Merge pull request #11879 from abellina/fix_read_parquet_ztd_temp_size

Fixes bug in temporary decompression space estimation before calling nvcomp
rapidsai · Oct 8, 2022 · 17868b7 · 17868b7
2 parents 4c4bce9 + 61e2499
commit 17868b7
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/cpp/src/io/comp/nvcomp_adapter.cpp b/cpp/src/io/comp/nvcomp_adapter.cpp
@@ -147,12 +147,15 @@ size_t batched_decompress_temp_size(compression_type compression,
                                     size_t max_uncomp_chunk_size,
                                     size_t max_total_uncomp_size)
 {
-  size_t temp_size = 0;
-  auto const nvcomp_status =
-    batched_decompress_get_temp_size_ex(
-      compression, num_chunks, max_uncomp_chunk_size, &temp_size, max_total_uncomp_size)
-      .value_or(batched_decompress_get_temp_size(
-        compression, num_chunks, max_uncomp_chunk_size, &temp_size));
+  size_t temp_size   = 0;
+  auto nvcomp_status = batched_decompress_get_temp_size_ex(
+    compression, num_chunks, max_uncomp_chunk_size, &temp_size, max_total_uncomp_size);
+
+  if (nvcomp_status.value_or(nvcompStatus_t::nvcompErrorInternal) !=
+      nvcompStatus_t::nvcompSuccess) {
+    nvcomp_status =
+      batched_decompress_get_temp_size(compression, num_chunks, max_uncomp_chunk_size, &temp_size);
+  }
 
   CUDF_EXPECTS(nvcomp_status == nvcompStatus_t::nvcompSuccess,
                "Unable to get scratch size for decompression");