Use nvcomp's snappy decompressor in avro reader

conda/environments/cudf_dev_cuda11.0.yml

@@ -19,6 +19,7 @@ dependencies:
   - pandas>=1.0,<1.4.0dev0
   - pyarrow=5.0.0=*cuda
   - fastavro>=0.22.9
+  - python-snappy>=0.6.0
   - notebook>=0.5.0
   - cython>=0.29,<0.30
   - fsspec>=0.6.0

conda/environments/cudf_dev_cuda11.2.yml

@@ -19,6 +19,7 @@ dependencies:
   - pandas>=1.0,<1.4.0dev0
   - pyarrow=5.0.0=*cuda
   - fastavro>=0.22.9
+  - python-snappy>=0.6.0
   - notebook>=0.5.0
   - cython>=0.29,<0.30
   - fsspec>=0.6.0

cpp/src/io/avro/reader_impl.cu

@@ -16,6 +16,7 @@
 
 #include "avro.h"
 #include "avro_gpu.h"
+#include "thrust/iterator/transform_output_iterator.h"
 
 #include <io/comp/gpuinflate.h>
 #include <io/utilities/column_buffer.hpp>
@@ -30,9 +31,13 @@
 #include <cudf/utilities/span.hpp>
 #include <cudf/utilities/traits.hpp>
 
+#include <numeric>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/device_uvector.hpp>
+#include <rmm/exec_policy.hpp>
+
+#include <nvcomp/snappy.h>
 
 #include <memory>
 #include <string>
@@ -151,99 +156,172 @@ rmm::device_buffer decompress_data(datasource& source,
                                    rmm::device_buffer const& comp_block_data,
                                    rmm::cuda_stream_view stream)
 {
-  size_t uncompressed_data_size = 0;
+  if (meta.codec == "deflate") {
+    size_t uncompressed_data_size = 0;
 
-  auto inflate_in  = hostdevice_vector<gpu_inflate_input_s>(meta.block_list.size());
-  auto inflate_out = hostdevice_vector<gpu_inflate_status_s>(meta.block_list.size());
+    auto inflate_in  = hostdevice_vector<gpu_inflate_input_s>(meta.block_list.size());
+    auto inflate_out = hostdevice_vector<gpu_inflate_status_s>(meta.block_list.size());
 
-  if (meta.codec == "deflate") {
     // Guess an initial maximum uncompressed block size
     uint32_t initial_blk_len = (meta.max_block_size * 2 + 0xfff) & ~0xfff;
     uncompressed_data_size   = initial_blk_len * meta.block_list.size();
     for (size_t i = 0; i < inflate_in.size(); ++i) {
       inflate_in[i].dstSize = initial_blk_len;
     }
-  } else if (meta.codec == "snappy") {
-    // Extract the uncompressed length from the snappy stream
-    for (size_t i = 0; i < meta.block_list.size(); i++) {
-      auto const buffer  = source.host_read(meta.block_list[i].offset, 4);
-      uint8_t const* blk = buffer->data();
-      uint32_t blk_len   = blk[0];
-      if (blk_len > 0x7f) {
-        blk_len = (blk_len & 0x7f) | (blk[1] << 7);
-        if (blk_len > 0x3fff) {
-          blk_len = (blk_len & 0x3fff) | (blk[2] << 14);
-          if (blk_len > 0x1fffff) { blk_len = (blk_len & 0x1fffff) | (blk[3] << 21); }
-        }
-      }
-      inflate_in[i].dstSize = blk_len;
-      uncompressed_data_size += blk_len;
-    }
-  } else {
-    CUDF_FAIL("Unsupported compression codec\n");
-  }
 
-  rmm::device_buffer decomp_block_data(uncompressed_data_size, stream);
+    rmm::device_buffer decomp_block_data(uncompressed_data_size, stream);
 
-  auto const base_offset = meta.block_list[0].offset;
-  for (size_t i = 0, dst_pos = 0; i < meta.block_list.size(); i++) {
-    auto const src_pos = meta.block_list[i].offset - base_offset;
+    auto const base_offset = meta.block_list[0].offset;
+    for (size_t i = 0, dst_pos = 0; i < meta.block_list.size(); i++) {
+      auto const src_pos = meta.block_list[i].offset - base_offset;
 
-    inflate_in[i].srcDevice = static_cast<uint8_t const*>(comp_block_data.data()) + src_pos;
-    inflate_in[i].srcSize   = meta.block_list[i].size;
-    inflate_in[i].dstDevice = static_cast<uint8_t*>(decomp_block_data.data()) + dst_pos;
+      inflate_in[i].srcDevice = static_cast<uint8_t const*>(comp_block_data.data()) + src_pos;
+      inflate_in[i].srcSize   = meta.block_list[i].size;
+      inflate_in[i].dstDevice = static_cast<uint8_t*>(decomp_block_data.data()) + dst_pos;
 
-    // Update blocks offsets & sizes to refer to uncompressed data
-    meta.block_list[i].offset = dst_pos;
-    meta.block_list[i].size   = static_cast<uint32_t>(inflate_in[i].dstSize);
-    dst_pos += meta.block_list[i].size;
-  }
+      // Update blocks offsets & sizes to refer to uncompressed data
+      meta.block_list[i].offset = dst_pos;
+      meta.block_list[i].size   = static_cast<uint32_t>(inflate_in[i].dstSize);
+      dst_pos += meta.block_list[i].size;
+    }
 
-  for (int loop_cnt = 0; loop_cnt < 2; loop_cnt++) {
-    inflate_in.host_to_device(stream);
-    CUDA_TRY(
-      cudaMemsetAsync(inflate_out.device_ptr(), 0, inflate_out.memory_size(), stream.value()));
-    if (meta.codec == "deflate") {
+    for (int loop_cnt = 0; loop_cnt < 2; loop_cnt++) {
+      inflate_in.host_to_device(stream);
+      CUDA_TRY(
+        cudaMemsetAsync(inflate_out.device_ptr(), 0, inflate_out.memory_size(), stream.value()));
       CUDA_TRY(gpuinflate(
         inflate_in.device_ptr(), inflate_out.device_ptr(), inflate_in.size(), 0, stream));
-    } else if (meta.codec == "snappy") {
-      CUDA_TRY(
-        gpu_unsnap(inflate_in.device_ptr(), inflate_out.device_ptr(), inflate_in.size(), stream));
-    } else {
-      CUDF_FAIL("Unsupported compression codec\n");
-    }
-    inflate_out.device_to_host(stream, true);
-
-    // Check if larger output is required, as it's not known ahead of time
-    if (meta.codec == "deflate" && !loop_cnt) {
-      size_t actual_uncompressed_size = 0;
-      for (size_t i = 0; i < meta.block_list.size(); i++) {
-        // If error status is 1 (buffer too small), the `bytes_written` field
-        // is actually contains the uncompressed data size
-        if (inflate_out[i].status == 1 && inflate_out[i].bytes_written > inflate_in[i].dstSize) {
-          inflate_in[i].dstSize = inflate_out[i].bytes_written;
+      inflate_out.device_to_host(stream, true);
+
+      // Check if larger output is required, as it's not known ahead of time
+      if (loop_cnt == 0) {
+        size_t actual_uncompressed_size = 0;
+        for (size_t i = 0; i < meta.block_list.size(); i++) {
+          // If error status is 1 (buffer too small), the `bytes_written` field
+          // is actually contains the uncompressed data size
+          if (inflate_out[i].status == 1 && inflate_out[i].bytes_written > inflate_in[i].dstSize) {
+            inflate_in[i].dstSize = inflate_out[i].bytes_written;
+          }
+          actual_uncompressed_size += inflate_in[i].dstSize;
         }
-        actual_uncompressed_size += inflate_in[i].dstSize;
-      }
-      if (actual_uncompressed_size > uncompressed_data_size) {
-        decomp_block_data.resize(actual_uncompressed_size, stream);
-        for (size_t i = 0, dst_pos = 0; i < meta.block_list.size(); i++) {
-          auto dst_base           = static_cast<uint8_t*>(decomp_block_data.data());
-          inflate_in[i].dstDevice = dst_base + dst_pos;
-
-          meta.block_list[i].offset = dst_pos;
-          meta.block_list[i].size   = static_cast<uint32_t>(inflate_in[i].dstSize);
-          dst_pos += meta.block_list[i].size;
+        if (actual_uncompressed_size > uncompressed_data_size) {
+          decomp_block_data.resize(actual_uncompressed_size, stream);
+          for (size_t i = 0, dst_pos = 0; i < meta.block_list.size(); i++) {
+            auto dst_base           = static_cast<uint8_t*>(decomp_block_data.data());
+            inflate_in[i].dstDevice = dst_base + dst_pos;
+
+            meta.block_list[i].offset = dst_pos;
+            meta.block_list[i].size   = static_cast<uint32_t>(inflate_in[i].dstSize);
+            dst_pos += meta.block_list[i].size;
+          }
+        } else {
+          break;
         }
-      } else {
-        break;
       }
-    } else {
-      break;
     }
-  }
 
-  return decomp_block_data;
+    return decomp_block_data;
+  } else if (meta.codec == "snappy") {
+    size_t const num_blocks = meta.block_list.size();
+
+    // comp_block_data contains contents of the avro file starting from the first block, excluding
+    // file header. meta.block_list[i].offset refers to offset of block i in the file, including
+    // file header.
+    // Find ptrs to each compressed block in comp_block_data by removing header offset.
+    hostdevice_vector<void const*> compressed_data_ptrs(num_blocks, stream);
+    std::transform(meta.block_list.begin(),
+                   meta.block_list.end(),
+                   compressed_data_ptrs.host_ptr(),
+                   [&](auto const& block) {
+                     return static_cast<std::byte const*>(comp_block_data.data()) +
+                            (block.offset - meta.block_list[0].offset);
+                   });
+    compressed_data_ptrs.host_to_device(stream);
+
+    hostdevice_vector<size_t> compressed_data_sizes(num_blocks, stream);
+    std::transform(meta.block_list.begin(),
+                   meta.block_list.end(),
+                   compressed_data_sizes.host_ptr(),
+                   [](auto const& block) { return block.size; });
+    compressed_data_sizes.host_to_device(stream);
+
+    hostdevice_vector<size_t> uncompressed_data_sizes(num_blocks, stream);
+    nvcompStatus_t status =
+      nvcompBatchedSnappyGetDecompressSizeAsync(compressed_data_ptrs.device_ptr(),
+                                                compressed_data_sizes.device_ptr(),
+                                                uncompressed_data_sizes.device_ptr(),
+                                                num_blocks,
+                                                stream.value());
+    CUDF_EXPECTS(status == nvcompStatus_t::nvcompSuccess,
+                 "Unable to get uncompressed sizes for snappy compressed blocks");
+    uncompressed_data_sizes.device_to_host(stream, true);
+
+    size_t const uncompressed_data_size =
+      std::reduce(uncompressed_data_sizes.begin(), uncompressed_data_sizes.end());
+    size_t const max_uncomp_block_size = std::reduce(
+      uncompressed_data_sizes.begin(), uncompressed_data_sizes.end(), 0, thrust::maximum<size_t>());
+
+    size_t temp_size;
+    status =
+      nvcompBatchedSnappyDecompressGetTempSize(num_blocks, max_uncomp_block_size, &temp_size);
+    CUDF_EXPECTS(status == nvcompStatus_t::nvcompSuccess,
+                 "Unable to get scratch size for snappy decompression");
+
+    rmm::device_buffer scratch(temp_size, stream);
+    rmm::device_buffer decomp_block_data(uncompressed_data_size, stream);
+    rmm::device_uvector<void*> uncompressed_data_ptrs(num_blocks, stream);
+    hostdevice_vector<size_t> uncompressed_data_offsets(num_blocks, stream);
+
+    std::exclusive_scan(uncompressed_data_sizes.begin(),
+                        uncompressed_data_sizes.end(),
+                        uncompressed_data_offsets.begin(),
+                        0);
+    uncompressed_data_offsets.host_to_device(stream);
+
+    thrust::tabulate(rmm::exec_policy(),
+                     uncompressed_data_ptrs.begin(),
+                     uncompressed_data_ptrs.end(),
+                     [off  = uncompressed_data_offsets.device_ptr(),
+                      data = static_cast<std::byte*>(decomp_block_data.data())] __device__(int i) {
+                       return data + off[i];
+                     });
+
+    rmm::device_uvector<size_t> actual_uncompressed_data_sizes(num_blocks, stream);
+    rmm::device_uvector<nvcompStatus_t> statuses(num_blocks, stream);
+
+    status = nvcompBatchedSnappyDecompressAsync(compressed_data_ptrs.device_ptr(),
+                                                compressed_data_sizes.device_ptr(),
+                                                uncompressed_data_sizes.device_ptr(),
+                                                actual_uncompressed_data_sizes.data(),
+                                                num_blocks,
+                                                scratch.data(),
+                                                scratch.size(),
+                                                uncompressed_data_ptrs.data(),
+                                                statuses.data(),
+                                                stream);
+    CUDF_EXPECTS(status == nvcompStatus_t::nvcompSuccess, "unable to perform snappy decompression");
+
+    CUDF_EXPECTS(thrust::equal(rmm::exec_policy(stream),
+                               uncompressed_data_sizes.d_begin(),
+                               uncompressed_data_sizes.d_end(),
+                               actual_uncompressed_data_sizes.begin()),
+                 "Mismatch in expected and actual decompressed size during snappy decompression");
+    CUDF_EXPECTS(thrust::equal(rmm::exec_policy(stream),
+                               statuses.begin(),
+                               statuses.end(),
+                               thrust::make_constant_iterator(nvcompStatus_t::nvcompSuccess)),
+                 "Error during snappy decompression");
+
+    // Update blocks offsets & sizes to refer to uncompressed data
+    for (size_t i = 0; i < num_blocks; i++) {
+      meta.block_list[i].offset = uncompressed_data_offsets[i];
+      meta.block_list[i].size   = uncompressed_data_sizes[i];
+    }
+
+    return decomp_block_data;
+  } else {
+    CUDF_FAIL("Unsupported compression codec\n");
+  }
 }
 
 std::vector<column_buffer> decode_data(metadata& meta,

cpp/src/io/utilities/hostdevice_vector.hpp

@@ -86,6 +86,10 @@ class hostdevice_vector {
 
   T& operator[](size_t i) const { return h_data[i]; }
   T* host_ptr(size_t offset = 0) const { return h_data + offset; }
+  T* begin() { return h_data; }
+  T* end() { return h_data + num_elements; }
+  T* d_begin() { return static_cast<T*>(d_data.data()); }
+  T* d_end() { return static_cast<T*>(d_data.data()) + num_elements; }
   T* device_ptr(size_t offset = 0) { return reinterpret_cast<T*>(d_data.data()) + offset; }
   T const* device_ptr(size_t offset = 0) const
   {

python/cudf/cudf/tests/test_avro_reader_fastavro_integration.py

@@ -19,6 +19,7 @@
 
 import cudf
 from cudf.testing._utils import assert_eq
+from cudf.testing.dataset_generator import rand_dataframe
 
 
 def cudf_from_avro_util(schema, records):
@@ -207,3 +208,39 @@ def test_can_parse_no_schema():
     actual = cudf_from_avro_util(schema_root, records)
     expected = cudf.DataFrame()
     assert_eq(expected, actual)
+
+
+@pytest.mark.parametrize("rows", [0, 1, 10, 100000])
+@pytest.mark.parametrize("codec", ["null", "deflate", "snappy"])
+def test_avro_compression(rows, codec):
+    schema = {
+        "name": "root",
+        "type": "record",
+        "fields": [
+            {"name": "0", "type": "int"},
+            {"name": "1", "type": "string"},
+        ],
+    }
+
+    df = rand_dataframe(
+        [
+            {"dtype": "int32", "null_frequency": 0, "cardinality": 1000},
+            {
+                "dtype": "str",
+                "null_frequency": 0,
+                "cardinality": 100,
+                "max_string_length": 10,
+            },
+        ],
+        rows,
+    )
+    expected_df = cudf.DataFrame.from_arrow(df)
+
+    records = df.to_pandas().to_dict(orient="records")
+
+    buffer = io.BytesIO()
+    fastavro.writer(buffer, schema, records, codec=codec)
+    buffer.seek(0)
+    got_df = cudf.read_avro(buffer)
+
+    assert_eq(expected_df, got_df)

python/cudf/requirements/cuda-11.0/dev_requirements.txt

@@ -10,6 +10,7 @@ cmake-setuptools>=0.1.3
 cython>=0.29,<0.30
 dlpack
 fastavro>=0.22.9
+python-snappy>=0.6.0
 fsspec>=0.6.0
 hypothesis
 mimesis<4.1

python/cudf/requirements/cuda-11.2/dev_requirements.txt

@@ -10,6 +10,7 @@ cmake-setuptools>=0.1.3
 cython>=0.29,<0.30
 dlpack
 fastavro>=0.22.9
+python-snappy>=0.6.0
 fsspec>=0.6.0
 hypothesis
 mimesis<4.1

python/cudf/setup.py

@@ -21,7 +21,6 @@
 install_requires = [
     "numba>=0.53.1",
     "Cython>=0.29,<0.30",
-    "fastavro>=0.22.9",
     "fsspec>=0.6.0",
     "numpy",
     "pandas>=1.0,<1.4.0dev0",
@@ -38,6 +37,8 @@
         "pytest-benchmark",
         "pytest-xdist",
         "hypothesis" "mimesis",
+        "fastavro>=0.22.9",
+        "python-snappy>=0.6.0",
         "pyorc",
         "msgpack",
         "transformers",