Repository: rapidsai/cudf
Pull request #18440: Fix auto-detection of compression type in host-side decompression
Author: login: shrshi
Labels: name: bug
description: Something isn't working
name: libcudf
description: Affects libcudf (C++/CUDA) code.
name: cuIO
description: cuIO issue
name: non-breaking
description: Non-breaking change
Created: 2025-04-04T23:05:07Z | Merged: 2025-04-22T17:37:10Z

## Description
Ensures that updated compression type is used for decompression instead of AUTO.

## Checklist
- [X] I am familiar with the [Contributing Guidelines](https://github.com/rapidsai/cudf/blob/HEAD/CONTRIBUTING.md).
- [X] New or existing tests cover these changes.
- [ ] The documentation is up to date with these changes.



Diff:
diff --git a/cpp/src/io/comp/uncomp.cpp b/cpp/src/io/comp/uncomp.cpp
index 5ca56dfa6b0..c0559c4cb76 100644
--- a/cpp/src/io/comp/uncomp.cpp
+++ b/cpp/src/io/comp/uncomp.cpp
@@ -485,6 +485,23 @@ source_properties get_source_properties(compression_type compression, host_span<
       if (compression != compression_type::AUTO) { break; }
       [[fallthrough]];
     }
+    case compression_type::ZSTD: {
+      auto const ret =
+        ZSTD_findDecompressedSize(reinterpret_cast<const void*>(src.data()), src.size());
+      uncomp_len = static_cast<size_t>(ret);
+      if (compression != compression_type::AUTO) {
+        CUDF_EXPECTS(ret != ZSTD_CONTENTSIZE_UNKNOWN,
+                     "Decompressed ZSTD size cannot be determined");
+        CUDF_EXPECTS(ret != ZSTD_CONTENTSIZE_ERROR, "Error determining decompressed ZSTD size");
+      } else if (ret != ZSTD_CONTENTSIZE_UNKNOWN && ret != ZSTD_CONTENTSIZE_ERROR) {
+        compression = compression_type::ZSTD;
+      }
+      if (compression != compression_type::AUTO) { break; }
+      [[fallthrough]];
+    }
+    // Snappy is detected after ZSTD since the
+    // magic characters checked at the beginning of the input buffer
+    // are valid for ZSTD compression as well.
     case compression_type::SNAPPY: {
       uncomp_len     = 0;
       auto cur       = src.begin();
@@ -511,20 +528,6 @@ source_properties get_source_properties(compression_type compression, host_span<
       if (compression != compression_type::AUTO) { break; }
       [[fallthrough]];
     }
-    case compression_type::ZSTD: {
-      auto const ret =
-        ZSTD_findDecompressedSize(reinterpret_cast<const void*>(src.data()), src.size());
-      uncomp_len = static_cast<size_t>(ret);
-      if (compression != compression_type::AUTO) {
-        CUDF_EXPECTS(ret != ZSTD_CONTENTSIZE_UNKNOWN,
-                     "Decompressed ZSTD size cannot be determined");
-        CUDF_EXPECTS(ret != ZSTD_CONTENTSIZE_ERROR, "Error determining decompressed ZSTD size");
-      } else if (ret != ZSTD_CONTENTSIZE_UNKNOWN && ret != ZSTD_CONTENTSIZE_ERROR) {
-        compression = compression_type::ZSTD;
-      }
-      if (compression != compression_type::AUTO) { break; }
-      [[fallthrough]];
-    }
     default: {
       uncomp_len = 0;
     }
@@ -568,13 +571,13 @@ std::vector<uint8_t> decompress(compression_type compression, host_span<uint8_t
                                      // ~4:1 compression for initial size
   }
 
-  if (compression == compression_type::GZIP) {
+  if (srcprops.compression == compression_type::GZIP) {
     // INFLATE
     std::vector<uint8_t> dst(srcprops.uncomp_len);
     decompress_gzip(src, dst);
     return dst;
   }
-  if (compression == compression_type::ZSTD) {
+  if (srcprops.compression == compression_type::ZSTD) {
     std::vector<uint8_t> dst(srcprops.uncomp_len);
     auto const decompressed_bytes = decompress_zstd(src, dst);
     CUDF_EXPECTS(decompressed_bytes == srcprops.uncomp_len,
@@ -583,12 +586,12 @@ std::vector<uint8_t> decompress(compression_type compression, host_span<uint8_t
     dst.resize(decompressed_bytes);
     return dst;
   }
-  if (compression == compression_type::ZIP) {
+  if (srcprops.compression == compression_type::ZIP) {
     std::vector<uint8_t> dst(srcprops.uncomp_len);
     cpu_inflate_vector(dst, srcprops.comp_data, srcprops.comp_len);
     return dst;
   }
-  if (compression == compression_type::BZIP2) {
+  if (srcprops.compression == compression_type::BZIP2) {
     size_t src_ofs = 0;
     size_t dst_ofs = 0;
     int bz_err     = 0;
@@ -612,7 +615,7 @@ std::vector<uint8_t> decompress(compression_type compression, host_span<uint8_t
     CUDF_EXPECTS(bz_err == 0, "Decompression: error in stream");
     return dst;
   }
-  if (compression == compression_type::SNAPPY) {
+  if (srcprops.compression == compression_type::SNAPPY) {
     std::vector<uint8_t> dst(srcprops.uncomp_len);
     decompress_snappy(src, dst);
     return dst;
diff --git a/cpp/src/io/functions.cpp b/cpp/src/io/functions.cpp
index d381bf077ad..3d7a88e7c33 100644
--- a/cpp/src/io/functions.cpp
+++ b/cpp/src/io/functions.cpp
@@ -48,7 +48,12 @@ compression_type infer_compression_type(compression_type compression, source_inf
 {
   if (compression != compression_type::AUTO) { return compression; }
 
-  if (info.type() != io_type::FILEPATH) { return compression_type::NONE; }
+  if (info.type() != io_type::FILEPATH) {
+    CUDF_LOG_WARN(
+      "Auto detection of compression type is supported only for file type buffers. For other "
+      "buffer types, AUTO compression type assumes uncompressed input.");
+    return compression_type::NONE;
+  }
 
   auto filepath = info.filepaths()[0];
 
@@ -68,6 +73,8 @@ compression_type infer_compression_type(compression_type compression, source_inf
   if (ext == "gz") { return compression_type::GZIP; }
   if (ext == "zip") { return compression_type::ZIP; }
   if (ext == "bz2") { return compression_type::BZIP2; }
+  if (ext == "zstd") { return compression_type::ZSTD; }
+  if (ext == "sz") { return compression_type::SNAPPY; }
   if (ext == "xz") { return compression_type::XZ; }
 
   return compression_type::NONE;
diff --git a/cpp/tests/io/comp/comp_test.cpp b/cpp/tests/io/comp/comp_test.cpp
index 75ca9911627..b8bb2143aef 100644
--- a/cpp/tests/io/comp/comp_test.cpp
+++ b/cpp/tests/io/comp/comp_test.cpp
@@ -47,18 +47,22 @@ enum class hw { CPU, GPU };
  * whose interface and setup is different for each codec.
  */
 template <typename Decompressor>
-struct DecompressTest : public cudf::test::BaseFixture, public testing::WithParamInterface<hw> {
+struct DecompressTest
+  : public cudf::test::BaseFixture,
+    public testing::WithParamInterface<std::tuple<hw, cudf::io::compression_type>> {
   [[nodiscard]] std::vector<uint8_t> vector_from_string(std::string const str) const
   {
     return {reinterpret_cast<uint8_t const*>(str.c_str()),
             reinterpret_cast<uint8_t const*>(str.c_str()) + strlen(str.c_str())};
   }
 
-  std::vector<uint8_t> Decompress(hw type,
+  std::vector<uint8_t> Decompress(std::tuple<hw, cudf::io::compression_type> type,
                                   cudf::host_span<uint8_t const> compressed,
                                   size_t uncompressed_size)
   {
-    if (type == hw::GPU) {
+    auto hw_type   = std::get<0>(GetParam());
+    auto comp_type = std::get<1>(GetParam());
+    if (hw_type == hw::GPU) {
       if constexpr (has_gpu_impl<Decompressor>::value) {
         return DeviceDecompress(compressed, uncompressed_size);
       } else {
@@ -66,7 +70,7 @@ struct DecompressTest : public cudf::test::BaseFixture, public testing::WithPara
       }
     } else {
       if constexpr (has_cpu_impl<Decompressor>::value) {
-        return HostDecompress(compressed, uncompressed_size);
+        return HostDecompress(comp_type, compressed, uncompressed_size);
       } else {
         CUDF_FAIL("Host decompression has not been implemented");
       }
@@ -103,10 +107,12 @@ struct DecompressTest : public cudf::test::BaseFixture, public testing::WithPara
     return decompressed;
   }
 
-  std::vector<uint8_t> HostDecompress(cudf::host_span<uint8_t const> compressed,
+  std::vector<uint8_t> HostDecompress(cudf::io::compression_type comp_type,
+                                      cudf::host_span<uint8_t const> compressed,
                                       size_t uncompressed_size)
   {
-    return static_cast<Decompressor*>(this)->host_dispatch(compressed, uncompressed_size);
+    return static_cast<Decompressor*>(this)->host_dispatch(
+      comp_type, compressed, uncompressed_size);
   }
 
   template <typename T, typename = void>
@@ -165,13 +171,17 @@ struct GzipDecompressTest : public DecompressTest<GzipDecompressTest> {
                                  cudf::get_default_stream());
   }
 
-  std::vector<uint8_t> host_dispatch(cudf::host_span<uint8_t const> compressed,
+  std::vector<uint8_t> host_dispatch(cudf::io::compression_type comp_type,
+                                     cudf::host_span<uint8_t const> compressed,
                                      size_t uncompressed_size)
   {
     CUDF_EXPECTS(uncompressed_size <= cudf::io::detail::get_uncompressed_size(
-                                        cudf::io::compression_type::GZIP, compressed),
+                                        cudf::io::compression_type::AUTO, compressed),
                  "Underestimating uncompressed size!");
-    return cudf::io::detail::decompress(cudf::io::compression_type::GZIP, compressed);
+    CUDF_EXPECTS(comp_type == cudf::io::compression_type::AUTO ||
+                   comp_type == cudf::io::compression_type::GZIP,
+                 "Invalid compression type");
+    return cudf::io::detail::decompress(comp_type, compressed);
   }
 };
 
@@ -179,13 +189,17 @@ struct GzipDecompressTest : public DecompressTest<GzipDecompressTest> {
  * @brief Derived fixture for GZIP decompression
  */
 struct ZstdDecompressTest : public DecompressTest<ZstdDecompressTest> {
-  std::vector<uint8_t> host_dispatch(cudf::host_span<uint8_t const> compressed,
+  std::vector<uint8_t> host_dispatch(cudf::io::compression_type comp_type,
+                                     cudf::host_span<uint8_t const> compressed,
                                      size_t uncompressed_size)
   {
     CUDF_EXPECTS(uncompressed_size <= cudf::io::detail::get_uncompressed_size(
-                                        cudf::io::compression_type::ZSTD, compressed),
+                                        cudf::io::compression_type::AUTO, compressed),
                  "Underestimating uncompressed size!");
-    return cudf::io::detail::decompress(cudf::io::compression_type::ZSTD, compressed);
+    CUDF_EXPECTS(comp_type == cudf::io::compression_type::AUTO ||
+                   comp_type == cudf::io::compression_type::ZSTD,
+                 "Invalid compression type");
+    return cudf::io::detail::decompress(comp_type, compressed);
   }
 };
 
@@ -200,13 +214,17 @@ struct SnappyDecompressTest : public DecompressTest<SnappyDecompressTest> {
     cudf::io::detail::gpu_unsnap(d_inf_in, d_inf_out, d_inf_stat, cudf::get_default_stream());
   }
 
-  std::vector<uint8_t> host_dispatch(cudf::host_span<uint8_t const> compressed,
+  std::vector<uint8_t> host_dispatch(cudf::io::compression_type comp_type,
+                                     cudf::host_span<uint8_t const> compressed,
                                      size_t uncompressed_size)
   {
     CUDF_EXPECTS(uncompressed_size <= cudf::io::detail::get_uncompressed_size(
-                                        cudf::io::compression_type::SNAPPY, compressed),
+                                        cudf::io::compression_type::AUTO, compressed),
                  "Underestimating uncompressed size!");
-    return cudf::io::detail::decompress(cudf::io::compression_type::SNAPPY, compressed);
+    CUDF_EXPECTS(comp_type == cudf::io::compression_type::AUTO ||
+                   comp_type == cudf::io::compression_type::SNAPPY,
+                 "Invalid compression type");
+    return cudf::io::detail::decompress(comp_type, compressed);
   }
 };
 
@@ -227,7 +245,9 @@ struct BrotliDecompressTest : public DecompressTest<BrotliDecompressTest> {
 
 INSTANTIATE_TEST_CASE_P(GzipDecompressTest,
                         GzipDecompressTest,
-                        ::testing::Values(hw::CPU, hw::GPU));
+                        ::testing::Combine(::testing::Values(hw::CPU, hw::GPU),
+                                           ::testing::Values(cudf::io::compression_type::AUTO,
+                                                             cudf::io::compression_type::GZIP)));
 
 TEST_P(GzipDecompressTest, HelloWorld)
 {
@@ -246,7 +266,9 @@ TEST_P(GzipDecompressTest, HelloWorld)
 
 INSTANTIATE_TEST_CASE_P(SnappyDecompressTest,
                         SnappyDecompressTest,
-                        ::testing::Values(hw::CPU, hw::GPU));
+                        ::testing::Combine(::testing::Values(hw::CPU, hw::GPU),
+                                           ::testing::Values(cudf::io::compression_type::AUTO,
+                                                             cudf::io::compression_type::SNAPPY)));
 
 TEST_P(SnappyDecompressTest, HelloWorld)
 {
@@ -276,7 +298,11 @@ TEST_P(SnappyDecompressTest, ShortLiteralAfterLongCopyAtStartup)
   EXPECT_EQ(output, input);
 }
 
-INSTANTIATE_TEST_CASE_P(BrotliDecompressTest, BrotliDecompressTest, ::testing::Values(hw::GPU));
+INSTANTIATE_TEST_CASE_P(
+  BrotliDecompressTest,
+  BrotliDecompressTest,
+  ::testing::Values(std::make_tuple(hw::GPU, cudf::io::compression_type::AUTO),
+                    std::make_tuple(hw::GPU, cudf::io::compression_type::BROTLI)));
 
 TEST_P(BrotliDecompressTest, HelloWorld)
 {
@@ -292,7 +318,11 @@ TEST_P(BrotliDecompressTest, HelloWorld)
   EXPECT_EQ(output, input);
 }
 
-INSTANTIATE_TEST_CASE_P(ZstdDecompressTest, ZstdDecompressTest, ::testing::Values(hw::CPU));
+INSTANTIATE_TEST_CASE_P(
+  ZstdDecompressTest,
+  ZstdDecompressTest,
+  ::testing::Values(std::make_tuple(hw::CPU, cudf::io::compression_type::AUTO),
+                    std::make_tuple(hw::CPU, cudf::io::compression_type::ZSTD)));
 
 TEST_P(ZstdDecompressTest, HelloWorld)
 {


Reviews:
user: login: bdice
body: Approving with some design suggestions for the future.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T02:06:38Z
user: login: vuule
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-12T02:17:09Z
user: login: vuule
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-14T20:50:02Z
user: login: shrshi
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T00:18:53Z
user: login: vuule
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T03:04:07Z

Review comments:
diff_hunk: @@ -485,6 +485,20 @@ source_properties get_source_properties(compression_type compression, host_span<
       if (compression != compression_type::AUTO) { break; }
       [[fallthrough]];
     }
+    case compression_type::ZSTD: {
path: cpp/src/io/comp/uncomp.cpp
user: login: bdice
body: I struggled to understand this logic from this diff, and had to expand the code a lot to get the full picture. Often, this lack of clarity in “local reasoning” can indicate unnecessary complexity. I wonder if there is a better way to express the AUTO behavior.

At a minimum, something like factoring the body of each `case` into a standalone helper function can make the overall control flow of the `switch` easier to see. I could see alternative approaches using std::optional (returning “maybe the detected sizes, or nullopt”) or using a std::variant holding one of the possible compression metadata values (offering more choices for what metadata might be returned by each compression type).

This can be follow-up work, but I think it is worth improving design of internals as we continue to invest in our compression code paths. The status quo of complexity and few people having a full understanding is partly driven by the difficulty of local reasoning.
created_at: 2025-04-07T01:56:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -68,6 +68,8 @@ compression_type infer_compression_type(compression_type compression, source_inf
   if (ext == "gz") { return compression_type::GZIP; }
   if (ext == "zip") { return compression_type::ZIP; }
   if (ext == "bz2") { return compression_type::BZIP2; }
+  if (ext == "zstd") { return compression_type::ZSTD; }
path: cpp/src/io/functions.cpp
user: login: bdice
body: We should optimize all our code paths for exiting early based on the distribution of most common compression types. If zstd and snappy are most popular, put them first. This function is cheap but other functions like the AUTO fallthrough logic above might not be, and a consistent ordering across all of our I/O code would be nice.
created_at: 2025-04-07T02:06:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -485,6 +485,20 @@ source_properties get_source_properties(compression_type compression, host_span<
       if (compression != compression_type::AUTO) { break; }
       [[fallthrough]];
     }
+    case compression_type::ZSTD: {
path: cpp/src/io/comp/uncomp.cpp
user: login: vuule
body: That's a good idea. 
I'd be fine with a separate PR with such refactoring; AFAIK, this PR is a part of a bug fix for an impactful issue.
created_at: 2025-04-12T02:17:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -169,9 +169,9 @@ struct GzipDecompressTest : public DecompressTest<GzipDecompressTest> {
                                      size_t uncompressed_size)
   {
     CUDF_EXPECTS(uncompressed_size <= cudf::io::detail::get_uncompressed_size(
-                                        cudf::io::compression_type::GZIP, compressed),
+                                        cudf::io::compression_type::AUTO, compressed),
path: cpp/tests/io/comp/comp_test.cpp
user: login: vuule
body: should we test both passing the concrete codec and AUTO?
created_at: 2025-04-14T20:31:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -485,6 +485,20 @@ source_properties get_source_properties(compression_type compression, host_span<
       if (compression != compression_type::AUTO) { break; }
       [[fallthrough]];
     }
+    case compression_type::ZSTD: {
+      auto const ret =
+        ZSTD_findDecompressedSize(reinterpret_cast<const void*>(src.data()), src.size());
+      uncomp_len = static_cast<size_t>(ret);
+      if (compression != compression_type::AUTO) {
+        CUDF_EXPECTS(ret != ZSTD_CONTENTSIZE_UNKNOWN,
+                     "Decompressed ZSTD size cannot be determined");
+        CUDF_EXPECTS(ret != ZSTD_CONTENTSIZE_ERROR, "Error determining decompressed ZSTD size");
+      } else if (ret != ZSTD_CONTENTSIZE_UNKNOWN && ret != ZSTD_CONTENTSIZE_ERROR) {
+        compression = compression_type::ZSTD;
+      }
+      if (compression != compression_type::AUTO) { break; }
+      [[fallthrough]];
+    }
path: cpp/src/io/comp/uncomp.cpp
user: login: vuule
body: please add a comment on why the snappy detection is performed last.
created_at: 2025-04-14T20:49:25Z
author_association: CONTRIBUTOR
diff_hunk: @@ -169,9 +169,9 @@ struct GzipDecompressTest : public DecompressTest<GzipDecompressTest> {
                                      size_t uncompressed_size)
   {
     CUDF_EXPECTS(uncompressed_size <= cudf::io::detail::get_uncompressed_size(
-                                        cudf::io::compression_type::GZIP, compressed),
+                                        cudf::io::compression_type::AUTO, compressed),
path: cpp/tests/io/comp/comp_test.cpp
user: login: shrshi
body: Good idea, I've added another parameter to the gtests
created_at: 2025-04-15T00:18:53Z
author_association: CONTRIBUTOR

PR comments:
user: login: copy-pr-bot[bot]
created_at: 2025-04-04T23:05:10Z
body: Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message [here](https://docs.gha-runners.nvidia.com/cpr/auto-sync).

author_association: NONE
user: login: shrshi
created_at: 2025-04-04T23:28:42Z
body: /ok to test
author_association: CONTRIBUTOR
user: login: shrshi
created_at: 2025-04-22T17:37:05Z
body: /merge
author_association: CONTRIBUTOR