enable zstd in parquet reader

vuule · May 13, 2022 · 34f0409 · 34f0409
1 parent df42642
commit 34f0409
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/cpp/src/io/parquet/reader_impl.cu b/cpp/src/io/parquet/reader_impl.cu
@@ -1108,7 +1108,8 @@ rmm::device_buffer reader::impl::decompress_page_data(
 
   std::array codecs{codec_stats{parquet::GZIP, 0, 0},
                     codec_stats{parquet::SNAPPY, 0, 0},
-                    codec_stats{parquet::BROTLI, 0, 0}};
+                    codec_stats{parquet::BROTLI, 0, 0},
+                    codec_stats{parquet::ZSTD, 0, 0}};
 
   auto is_codec_supported = [&codecs](int8_t codec) {
     if (codec == parquet::UNCOMPRESSED) return true;
@@ -1191,6 +1192,14 @@ rmm::device_buffer reader::impl::decompress_page_data(
           gpu_unsnap(d_comp_in, d_comp_out, d_comp_stats_view, stream);
         }
         break;
+      case parquet::ZSTD:
+        nvcomp::batched_decompress(nvcomp::compression_type::ZSTD,
+                                   d_comp_in,
+                                   d_comp_out,
+                                   d_comp_stats_view,
+                                   codec.max_decompressed_size,
+                                   stream);
+        break;
       case parquet::BROTLI:
         gpu_debrotli(d_comp_in,
                      d_comp_out,