Add support for string/binary view to content defined chunking writer

Author: robert3005

cpp/src/parquet/chunker_internal.cc

@@ -21,6 +21,7 @@
 #include <cstdint>
 #include <iterator>
 #include <string>
+#include <string_view>
 #include <vector>
 
 #include "arrow/array.h"
@@ -349,6 +350,17 @@ class ContentDefinedChunker::Impl {
     });
   }
 
+  std::vector<Chunk> CalculateBinaryView(const int16_t* def_levels,
+                                         const int16_t* rep_levels, int64_t num_levels,
+                                         const ::arrow::Array& values) {
+    const auto& array = checked_cast<const ::arrow::BinaryViewArray&>(values);
+    return Calculate(def_levels, rep_levels, num_levels, [&](int64_t i) {
+      const std::string_view value = array.GetView(i);
+      Roll(reinterpret_cast<const uint8_t*>(value.data()),
+           static_cast<int64_t>(value.size()));
+    });
+  }
+
   std::vector<Chunk> GetChunks(const int16_t* def_levels, const int16_t* rep_levels,
                                int64_t num_levels, const ::arrow::Array& values) {
     auto handle_type = [&](auto&& type) -> std::vector<Chunk> {
@@ -380,6 +392,8 @@ class ContentDefinedChunker::Impl {
       } else if constexpr (::arrow::is_large_binary_like(ArrowType::type_id)) {
         return CalculateBinaryLike<::arrow::LargeBinaryArray>(def_levels, rep_levels,
                                                               num_levels, values);
+      } else if constexpr (::arrow::is_binary_view_like(ArrowType::type_id)) {
+        return CalculateBinaryView(def_levels, rep_levels, num_levels, values);
       } else if constexpr (::arrow::is_dictionary(ArrowType::type_id)) {
         return GetChunks(def_levels, rep_levels, num_levels,
                          *static_cast<const ::arrow::DictionaryArray&>(values).indices());

cpp/src/parquet/chunker_internal_test.cc

@@ -170,6 +170,10 @@ Result<std::shared_ptr<Array>> GenerateArray(const std::shared_ptr<Field>& field
                     std::string("bin_") + std::to_string(val))
       GENERATE_CASE(LARGE_BINARY, ::arrow::LargeBinaryBuilder,
                     std::string("bin_") + std::to_string(val))
+      GENERATE_CASE(STRING_VIEW, ::arrow::StringViewBuilder,
+                    std::string("str_") + std::to_string(val))
+      GENERATE_CASE(BINARY_VIEW, ::arrow::BinaryViewBuilder,
+                    std::string("bin_") + std::to_string(val))
     case ::arrow::Type::FIXED_SIZE_BINARY: {
       auto size =
           checked_cast<const ::arrow::FixedSizeBinaryType*>(type.get())->byte_width();
@@ -742,6 +746,13 @@ Result<int64_t> CalculateCdcSize(const std::shared_ptr<Array>& array, bool nulla
   } else if (::arrow::is_large_binary_like(type_id)) {
     auto binary_array = checked_cast<const ::arrow::LargeBinaryArray*>(array.get());
     result += binary_array->total_values_length();
+  } else if (::arrow::is_binary_view_like(type_id)) {
+    auto view_array = checked_cast<const ::arrow::BinaryViewArray*>(array.get());
+    for (int64_t i = 0; i < view_array->length(); i++) {
+      if (view_array->IsValid(i)) {
+        result += static_cast<int64_t>(view_array->GetView(i).size());
+      }
+    }
   } else {
     return Status::NotImplemented("CDC size calculation for type ",
                                   array->type()->ToString(), " is not implemented");
@@ -784,7 +795,8 @@ void AssertContentDefinedChunkSizes(const std::shared_ptr<::arrow::ChunkedArray>
     ASSERT_EQ(column_info.has_dictionary_page, expect_dictionary_page);
   }
 
-  if (::arrow::is_fixed_width(type_id) || ::arrow::is_base_binary_like(type_id)) {
+  if (::arrow::is_fixed_width(type_id) || ::arrow::is_base_binary_like(type_id) ||
+      ::arrow::is_binary_view_like(type_id)) {
     int64_t offset = 0;
 
     auto page_lengths = column_info.page_lengths;
@@ -1451,7 +1463,7 @@ INSTANTIATE_TEST_SUITE_P(
     testing::Values(
         CaseConfig{::arrow::boolean(), false}, CaseConfig{::arrow::int64(), true},
         // Binary-like
-        CaseConfig{::arrow::utf8(), false},
+        CaseConfig{::arrow::utf8(), false}, CaseConfig{::arrow::utf8_view(), true},
         CaseConfig{::arrow::fixed_size_binary(16), true},
         // Nested types
         CaseConfig{::arrow::list(::arrow::int32()), false},
@@ -1475,6 +1487,7 @@ INSTANTIATE_TEST_SUITE_P(
         // Binary-like
         CaseConfig{::arrow::utf8(), false}, CaseConfig{::arrow::binary(), true},
         CaseConfig{::arrow::fixed_size_binary(16), true},
+        CaseConfig{::arrow::utf8_view(), false}, CaseConfig{::arrow::binary_view(), true},
         // Temporal
         CaseConfig{::arrow::date32(), false},
         CaseConfig{::arrow::time32(::arrow::TimeUnit::MILLI), true},