Support for Zstandard decompression in ORC reader

cpp/src/io/comp/io_uncomp.hpp

@@ -40,7 +40,8 @@ std::vector<uint8_t> decompress(compression_type compression, host_span<uint8_t
 
 size_t decompress(compression_type compression,
                   host_span<uint8_t const> src,
-                  host_span<uint8_t> dst);
+                  host_span<uint8_t> dst,
+                  rmm::cuda_stream_view stream);
 
 /**
  * @brief GZIP header flags

cpp/src/io/comp/uncomp.cpp

@@ -15,10 +15,13 @@
  */
 
 #include "io_uncomp.hpp"
+#include "nvcomp_adapter.hpp"
 #include "unbz2.hpp"  // bz2 uncompress
 
+#include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/utilities/error.hpp>
 #include <cudf/utilities/span.hpp>
+#include <io/utilities/hostdevice_vector.hpp>
 
 #include <cuda_runtime.h>
 
@@ -498,14 +501,50 @@ size_t decompress_snappy(host_span<uint8_t const> src, host_span<uint8_t> dst)
   return uncompressed_size;
 }
 
+/**
+ * @brief ZSTD decompressor that uses nvcomp
+ */
+size_t decompress_zstd(host_span<uint8_t const> src,
+                       host_span<uint8_t> dst,
+                       rmm::cuda_stream_view stream)
+{
+  // Init device span of spans (source)
+  auto const d_src = cudf::detail::make_device_uvector_async(src, stream);
+  auto hd_srcs     = hostdevice_vector<device_span<uint8_t const>>(1, stream);
+  hd_srcs[0]       = d_src;
+  hd_srcs.host_to_device(stream);
+
+  // Init device span of spans (temporary destination)
+  auto d_dst   = rmm::device_uvector<uint8_t>(dst.size(), stream);
+  auto hd_dsts = hostdevice_vector<device_span<uint8_t>>(1, stream);
+  hd_dsts[0]   = d_dst;
+  hd_dsts.host_to_device(stream);
+
+  auto hd_stats                   = hostdevice_vector<decompress_status>(1, stream);
+  auto const max_uncomp_page_size = dst.size();
+  nvcomp::batched_decompress(
+    nvcomp::compression_type::ZSTD, hd_srcs, hd_dsts, hd_stats, max_uncomp_page_size, stream);
+
+  // Copy temporary output to `dst`
+  CUDF_CUDA_TRY(
+    cudaMemcpyAsync(dst.data(), d_dst.data(), dst.size(), cudaMemcpyDeviceToHost, stream.value()));
+
+  hd_stats.device_to_host(stream, true);
+  CUDF_EXPECTS(hd_stats[0].status == 0, "ZSTD decompression failed");
+
+  return hd_stats[0].bytes_written;
+}
+
 size_t decompress(compression_type compression,
                   host_span<uint8_t const> src,
-                  host_span<uint8_t> dst)
+                  host_span<uint8_t> dst,
+                  rmm::cuda_stream_view stream)
 {
   switch (compression) {
     case compression_type::GZIP: return decompress_gzip(src, dst);
     case compression_type::ZLIB: return decompress_zlib(src, dst);
     case compression_type::SNAPPY: return decompress_snappy(src, dst);
+    case compression_type::ZSTD: return decompress_zstd(src, dst, stream);
     default: CUDF_FAIL("Unsupported compression type");
   }
 }

cpp/src/io/functions.cpp

@@ -240,6 +240,7 @@ namespace detail_orc = cudf::io::detail::orc;
 
 raw_orc_statistics read_raw_orc_statistics(source_info const& src_info)
 {
+  auto stream = rmm::cuda_stream_default;
   // Get source to read statistics from
   std::unique_ptr<datasource> source;
   if (src_info.type() == io_type::FILEPATH) {
@@ -256,7 +257,7 @@ raw_orc_statistics read_raw_orc_statistics(source_info const& src_info)
     CUDF_FAIL("Unsupported source type");
   }
 
-  orc::metadata metadata(source.get());
+  orc::metadata metadata(source.get(), stream);
 
   // Initialize statistics to return
   raw_orc_statistics result;

cpp/src/io/orc/aggregate_orc_metadata.cpp

@@ -92,12 +92,13 @@ void add_column_to_mapping(std::map<size_type, std::vector<size_type>>& selected
 /**
  * @brief Create a metadata object from each element in the source vector
  */
-auto metadatas_from_sources(std::vector<std::unique_ptr<datasource>> const& sources)
+auto metadatas_from_sources(std::vector<std::unique_ptr<datasource>> const& sources,
+                            rmm::cuda_stream_view stream)
 {
   std::vector<metadata> metadatas;
   std::transform(
-    sources.cbegin(), sources.cend(), std::back_inserter(metadatas), [](auto const& source) {
-      return metadata(source.get());
+    sources.cbegin(), sources.cend(), std::back_inserter(metadatas), [stream](auto const& source) {
+      return metadata(source.get(), stream);
     });
   return metadatas;
 }
@@ -121,8 +122,8 @@ size_type aggregate_orc_metadata::calc_num_stripes() const
 }
 
 aggregate_orc_metadata::aggregate_orc_metadata(
-  std::vector<std::unique_ptr<datasource>> const& sources)
-  : per_file_metadata(metadatas_from_sources(sources)),
+  std::vector<std::unique_ptr<datasource>> const& sources, rmm::cuda_stream_view stream)
+  : per_file_metadata(metadatas_from_sources(sources, stream)),
     num_rows(calc_num_rows()),
     num_stripes(calc_num_stripes())
 {
@@ -152,7 +153,8 @@ aggregate_orc_metadata::aggregate_orc_metadata(
 std::vector<metadata::stripe_source_mapping> aggregate_orc_metadata::select_stripes(
   std::vector<std::vector<size_type>> const& user_specified_stripes,
   size_type& row_start,
-  size_type& row_count)
+  size_type& row_count,
+  rmm::cuda_stream_view stream)
 {
   std::vector<metadata::stripe_source_mapping> selected_stripes_mapping;
 
@@ -234,7 +236,7 @@ std::vector<metadata::stripe_source_mapping> aggregate_orc_metadata::select_stri
         const auto buffer =
           per_file_metadata[mapping.source_idx].source->host_read(sf_comp_offset, sf_comp_length);
         auto sf_data = per_file_metadata[mapping.source_idx].decompressor->decompress_blocks(
-          {buffer->data(), buffer->size()});
+          {buffer->data(), buffer->size()}, stream);
         ProtobufReader(sf_data.data(), sf_data.size())
           .read(per_file_metadata[mapping.source_idx].stripefooters[i]);
         mapping.stripe_info[i].second = &per_file_metadata[mapping.source_idx].stripefooters[i];

cpp/src/io/orc/aggregate_orc_metadata.hpp

@@ -65,7 +65,8 @@ class aggregate_orc_metadata {
   size_type const num_stripes;
   bool row_grp_idx_present{true};
 
-  aggregate_orc_metadata(std::vector<std::unique_ptr<datasource>> const& sources);
+  aggregate_orc_metadata(std::vector<std::unique_ptr<datasource>> const& sources,
+                         rmm::cuda_stream_view stream);
 
   [[nodiscard]] auto const& get_schema(int schema_idx) const
   {
@@ -116,7 +117,8 @@ class aggregate_orc_metadata {
   std::vector<metadata::stripe_source_mapping> select_stripes(
     std::vector<std::vector<size_type>> const& user_specified_stripes,
     size_type& row_start,
-    size_type& row_count);
+    size_type& row_count,
+    rmm::cuda_stream_view stream);
 
   /**
    * @brief Filters ORC file to a selection of columns, based on their paths in the file.

cpp/src/io/orc/orc.cpp

@@ -373,12 +373,16 @@ OrcDecompressor::OrcDecompressor(CompressionKind kind, uint32_t blockSize) : m_b
       break;
     case LZO: _compression = compression_type::LZO; break;
     case LZ4: _compression = compression_type::LZ4; break;
-    case ZSTD: _compression = compression_type::ZSTD; break;
+    case ZSTD:
+      m_log2MaxRatio = 11;
+      _compression   = compression_type::ZSTD;
+      break;
     default: CUDF_FAIL("Invalid compression type");
   }
 }
 
-host_span<uint8_t const> OrcDecompressor::decompress_blocks(host_span<uint8_t const> src)
+host_span<uint8_t const> OrcDecompressor::decompress_blocks(host_span<uint8_t const> src,
+                                                            rmm::cuda_stream_view stream)
 {
   // If uncompressed, just pass-through the input
   if (src.empty() or _compression == compression_type::NONE) { return src; }
@@ -419,7 +423,7 @@ host_span<uint8_t const> OrcDecompressor::decompress_blocks(host_span<uint8_t co
     } else {
       // Compressed block
       dst_length += decompress(
-        _compression, src.subspan(i, block_len), {m_buf.data() + dst_length, m_blockSize});
+        _compression, src.subspan(i, block_len), {m_buf.data() + dst_length, m_blockSize}, stream);
     }
     i += block_len;
   }
@@ -428,7 +432,7 @@ host_span<uint8_t const> OrcDecompressor::decompress_blocks(host_span<uint8_t co
   return m_buf;
 }
 
-metadata::metadata(datasource* const src) : source(src)
+metadata::metadata(datasource* const src, rmm::cuda_stream_view stream) : source(src)
 {
   const auto len         = source->size();
   const auto max_ps_size = std::min(len, static_cast<size_t>(256));
@@ -446,14 +450,14 @@ metadata::metadata(datasource* const src) : source(src)
 
   // Read compressed filefooter section
   buffer             = source->host_read(len - ps_length - 1 - ps.footerLength, ps.footerLength);
-  auto const ff_data = decompressor->decompress_blocks({buffer->data(), buffer->size()});
+  auto const ff_data = decompressor->decompress_blocks({buffer->data(), buffer->size()}, stream);
   ProtobufReader(ff_data.data(), ff_data.size()).read(ff);
   CUDF_EXPECTS(get_num_columns() > 0, "No columns found");
 
   // Read compressed metadata section
   buffer =
     source->host_read(len - ps_length - 1 - ps.footerLength - ps.metadataLength, ps.metadataLength);
-  auto const md_data = decompressor->decompress_blocks({buffer->data(), buffer->size()});
+  auto const md_data = decompressor->decompress_blocks({buffer->data(), buffer->size()}, stream);
   orc::ProtobufReader(md_data.data(), md_data.size()).read(md);
 
   init_parent_descriptors();

cpp/src/io/orc/orc.hpp

@@ -538,10 +538,12 @@ class OrcDecompressor {
    * @brief ORC block decompression
    *
    * @param src compressed data
+   * @param stream CUDA stream used for device memory operations and kernel launches
    *
    * @return decompressed data
    */
-  host_span<uint8_t const> decompress_blocks(host_span<uint8_t const> src);
+  host_span<uint8_t const> decompress_blocks(host_span<uint8_t const> src,
+                                             rmm::cuda_stream_view stream);
   [[nodiscard]] uint32_t GetLog2MaxCompressionRatio() const { return m_log2MaxRatio; }
   [[nodiscard]] uint32_t GetMaxUncompressedBlockSize(uint32_t block_len) const
   {
@@ -601,7 +603,7 @@ class metadata {
   };
 
  public:
-  explicit metadata(datasource* const src);
+  explicit metadata(datasource* const src, rmm::cuda_stream_view stream);
 
   [[nodiscard]] size_t get_total_rows() const { return ff.numberOfRows; }
   [[nodiscard]] int get_num_stripes() const { return ff.stripes.size(); }

cpp/src/io/orc/reader_impl.cu

@@ -384,6 +384,14 @@ rmm::device_buffer reader::impl::decompress_stripe_data(
           gpu_unsnap(inflate_in_view, inflate_out_view, inflate_stats, stream);
         }
         break;
+      case compression_type::ZSTD:
+        nvcomp::batched_decompress(nvcomp::compression_type::ZSTD,
+                                   inflate_in_view,
+                                   inflate_out_view,
+                                   inflate_stats,
+                                   max_uncomp_block_size,
+                                   stream);
+        break;
       default: CUDF_FAIL("Unexpected decompression dispatch"); break;
     }
     decompress_check(inflate_stats, any_block_failure.device_ptr(), stream);
@@ -859,10 +867,11 @@ void reader::impl::create_columns(std::vector<std::vector<column_buffer>>&& col_
 
 reader::impl::impl(std::vector<std::unique_ptr<datasource>>&& sources,
                    orc_reader_options const& options,
+                   rmm::cuda_stream_view stream,
                    rmm::mr::device_memory_resource* mr)
   : _mr(mr),
     _sources(std::move(sources)),
-    _metadata{_sources},
+    _metadata{_sources, stream},
     selected_columns{_metadata.select_columns(options.get_columns())}
 {
   // Override output timestamp resolution if requested
@@ -921,7 +930,7 @@ table_with_metadata reader::impl::read(size_type skip_rows,
     return {std::make_unique<table>(), std::move(out_metadata)};
 
   // Select only stripes required (aka row groups)
-  const auto selected_stripes = _metadata.select_stripes(stripes, skip_rows, num_rows);
+  const auto selected_stripes = _metadata.select_stripes(stripes, skip_rows, num_rows, stream);
 
   auto const tz_table = compute_timezone_table(selected_stripes, stream);
 
@@ -1287,7 +1296,7 @@ reader::reader(std::vector<std::unique_ptr<cudf::io::datasource>>&& sources,
                rmm::cuda_stream_view stream,
                rmm::mr::device_memory_resource* mr)
 {
-  _impl = std::make_unique<impl>(std::move(sources), options, mr);
+  _impl = std::make_unique<impl>(std::move(sources), options, stream, mr);
 }
 
 // Destructor within this translation unit

cpp/src/io/orc/reader_impl.hpp

@@ -84,6 +84,7 @@ class reader::impl {
    */
   explicit impl(std::vector<std::unique_ptr<datasource>>&& sources,
                 orc_reader_options const& options,
+                rmm::cuda_stream_view stream,
                 rmm::mr::device_memory_resource* mr);
 
   /**

python/cudf/cudf/tests/test_orc.py

@@ -1714,3 +1714,21 @@ def test_empty_columns():
 
     got_df = cudf.read_orc(buffer)
     assert_eq(expected, got_df)
+
+
+def test_orc_reader_zstd_compression(list_struct_buff):
+    expected = cudf.read_orc(list_struct_buff)
+    # save with ZSTD compression
+    buffer = BytesIO()
+    pyarrow_tbl = pyarrow.orc.ORCFile(list_struct_buff).read()
+    writer = pyarrow.orc.ORCWriter(buffer, compression="zstd")
+    writer.write(pyarrow_tbl)
+    writer.close()
+    try:
+        got = cudf.read_orc(buffer)
+        assert_eq(expected, got)
+    except RuntimeError as e:
+        if "Unsupported compression type" in str(e):
+            pytest.mark.xfail(reason="nvcomp build doesn't have zstd")
+        else:
+            raise e