Use nvCOMP for ZLIB decompression in ORC reader (#11024)

Issue #11023 Adds `DEFLATE` compression type to the nvCOMP adapter. ORC reader uses the adapter when nvCOMP experimental integrations are enabled (for now). Otherwise behavior is unchanged and the internal implementation is still used. No tests as no visible behavior change is expected. Pending: Performance impact Authors: - Vukasin Milovanovic (https://github.com/vuule) Approvers: - Nghia Truong (https://github.com/ttnghia) - David Wendt (https://github.com/davidwendt) - Jim Brennan (https://github.com/jbrennan333) - Bradley Dice (https://github.com/bdice) URL: #11024
rapidsai · Jun 2, 2022 · 447028f · 447028f
1 parent 9301235
commit 447028f
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 3 deletions.
diff --git a/cpp/src/io/comp/nvcomp_adapter.cpp b/cpp/src/io/comp/nvcomp_adapter.cpp
@@ -29,6 +29,14 @@
 #define NVCOMP_HAS_ZSTD 0
 #endif
 
+#define NVCOMP_DEFLATE_HEADER <nvcomp/deflate.h>
+#if __has_include(NVCOMP_DEFLATE_HEADER)
+#include NVCOMP_DEFLATE_HEADER
+#define NVCOMP_HAS_DEFLATE 1
+#else
+#define NVCOMP_HAS_DEFLATE 0
+#endif
+
 namespace cudf::io::nvcomp {
 
 template <typename... Args>
@@ -40,6 +48,10 @@ auto batched_decompress_get_temp_size(compression_type compression, Args&&... ar
 #if NVCOMP_HAS_ZSTD
     case compression_type::ZSTD:
       return nvcompBatchedZstdDecompressGetTempSize(std::forward<Args>(args)...);
+#endif
+#if NVCOMP_HAS_DEFLATE
+    case compression_type::DEFLATE:
+      return nvcompBatchedDeflateDecompressGetTempSize(std::forward<Args>(args)...);
 #endif
     default: CUDF_FAIL("Unsupported compression type");
   }
@@ -54,6 +66,10 @@ auto batched_decompress_async(compression_type compression, Args&&... args)
 #if NVCOMP_HAS_ZSTD
     case compression_type::ZSTD:
       return nvcompBatchedZstdDecompressAsync(std::forward<Args>(args)...);
+#endif
+#if NVCOMP_HAS_DEFLATE
+    case compression_type::DEFLATE:
+      return nvcompBatchedDeflateDecompressAsync(std::forward<Args>(args)...);
 #endif
     default: CUDF_FAIL("Unsupported compression type");
   }

diff --git a/cpp/src/io/comp/nvcomp_adapter.hpp b/cpp/src/io/comp/nvcomp_adapter.hpp
@@ -24,7 +24,7 @@
 
 namespace cudf::io::nvcomp {
 
-enum class compression_type { SNAPPY, ZSTD };
+enum class compression_type { SNAPPY, ZSTD, DEFLATE };
 
 /**
  * @brief Device batch decompression of given type.

diff --git a/cpp/src/io/orc/reader_impl.cu b/cpp/src/io/orc/reader_impl.cu
@@ -369,8 +369,17 @@ rmm::device_buffer reader::impl::decompress_stripe_data(
     device_span<device_span<uint8_t>> inflate_out_view{inflate_out.data(), num_compressed_blocks};
     switch (decompressor.compression()) {
       case compression_type::ZLIB:
-        gpuinflate(
-          inflate_in_view, inflate_out_view, inflate_stats, gzip_header_included::NO, stream);
+        if (nvcomp_integration::is_all_enabled()) {
+          nvcomp::batched_decompress(nvcomp::compression_type::DEFLATE,
+                                     inflate_in_view,
+                                     inflate_out_view,
+                                     inflate_stats,
+                                     max_uncomp_block_size,
+                                     stream);
+        } else {
+          gpuinflate(
+            inflate_in_view, inflate_out_view, inflate_stats, gzip_header_included::NO, stream);
+        }
         break;
       case compression_type::SNAPPY:
         if (nvcomp_integration::is_stable_enabled()) {