Implement direct reads of Parquet RLE data into Arrow REE for string types

Author: lesterfan

cpp/src/arrow/compute/kernels/scalar_cast_internal.cc

@@ -209,6 +209,13 @@ Status UnpackDictionary(KernelContext* ctx, const ExecSpan& batch, ExecResult* o
   return Status::OK();
 }
 
+Status DecodeRunEndEncoded(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
+  Datum input(batch[0].array.ToArrayData());
+  ARROW_ASSIGN_OR_RAISE(Datum decoded, RunEndDecode(input, ctx->exec_context()));
+  out->value = std::move(decoded.array());
+  return Status::OK();
+}
+
 Status OutputAllNull(KernelContext* ctx, const ExecSpan& batch, ExecResult* out) {
   // TODO(wesm): there is no good reason to have to use ArrayData here, so we
   // should clean this up later. This is used in the dict<null>->null cast
@@ -286,6 +293,12 @@ static bool CanCastFromDictionary(Type::type type_id) {
           is_fixed_size_binary(type_id));
 }
 
+static bool CanCastFromRunEndEncoded(Type::type type_id) {
+  return (is_primitive(type_id) ||
+          is_base_binary_like(type_id) ||
+          is_fixed_size_binary(type_id));
+}
+
 void AddCommonCasts(Type::type out_type_id, OutputType out_ty, CastFunction* func) {
   // From null to this type
   ScalarKernel kernel;
@@ -303,6 +316,13 @@ void AddCommonCasts(Type::type out_type_id, OutputType out_ty, CastFunction* fun
                               MemAllocation::NO_PREALLOCATE));
   }
 
+  // From run-end-encoded to this type
+  if (CanCastFromRunEndEncoded(out_type_id)) {
+    DCHECK_OK(func->AddKernel(Type::RUN_END_ENCODED, {InputType(Type::RUN_END_ENCODED)}, out_ty,
+                              DecodeRunEndEncoded, NullHandling::COMPUTED_NO_PREALLOCATE,
+                              MemAllocation::NO_PREALLOCATE));
+  }
+
   // From extension type to this type
   DCHECK_OK(func->AddKernel(Type::EXTENSION, {InputType(Type::EXTENSION)}, out_ty,
                             CastFromExtension, NullHandling::COMPUTED_NO_PREALLOCATE,

cpp/src/arrow/dataset/file_parquet.cc

@@ -115,6 +115,10 @@ parquet::ArrowReaderProperties MakeArrowReaderProperties(
     auto column_index = metadata.schema()->ColumnIndex(name);
     properties.set_read_dictionary(column_index, true);
   }
+  for (const std::string& name : format.reader_options.ree_columns) {
+    auto column_index = metadata.schema()->ColumnIndex(name);
+    properties.set_read_ree(column_index, true);
+  }
   properties.set_coerce_int96_timestamp_unit(
       format.reader_options.coerce_int96_timestamp_unit);
   properties.set_binary_type(format.reader_options.binary_type);
@@ -445,6 +449,7 @@ bool ParquetFileFormat::Equals(const FileFormat& other) const {
 
   // FIXME implement comparison for decryption options
   return (reader_options.dict_columns == other_reader_options.dict_columns &&
+          reader_options.ree_columns == other_reader_options.ree_columns &&
           reader_options.coerce_int96_timestamp_unit ==
               other_reader_options.coerce_int96_timestamp_unit &&
           reader_options.binary_type == other_reader_options.binary_type &&

cpp/src/arrow/dataset/file_parquet.h

@@ -89,6 +89,7 @@ class ARROW_DS_EXPORT ParquetFileFormat : public FileFormat {
     ///
     /// @{
     std::unordered_set<std::string> dict_columns;
+    std::unordered_set<std::string> ree_columns;
     arrow::TimeUnit::type coerce_int96_timestamp_unit = arrow::TimeUnit::NANO;
     Type::type binary_type = Type::BINARY;
     Type::type list_type = Type::LIST;

cpp/src/arrow/dataset/file_parquet_test.cc

@@ -271,6 +271,17 @@ TEST_F(TestParquetFileFormat, InspectDictEncoded) {
   AssertSchemaEqual(*actual, expected_schema, /* check_metadata = */ false);
 }
 
+TEST_F(TestParquetFileFormat, InspectReeEncoded) {
+  auto reader = GetRecordBatchReader(schema({field("utf8", utf8())}));
+  auto source = GetFileSource(reader.get());
+
+  format_->reader_options.ree_columns = {"utf8"};
+  ASSERT_OK_AND_ASSIGN(auto actual, format_->Inspect(*source.get()));
+
+  Schema expected_schema({field("utf8", run_end_encoded(int32(), utf8()))});
+  AssertSchemaEqual(*actual, expected_schema, /* check_metadata = */ false);
+}
+
 TEST_F(TestParquetFileFormat, IsSupported) { TestIsSupported(); }
 
 TEST_F(TestParquetFileFormat, WriteRecordBatchReader) { TestWrite(); }
@@ -617,6 +628,25 @@ TEST_P(TestParquetFileFormatScan, ScanRecordBatchReaderDictEncoded) {
   }
   ASSERT_EQ(row_count, expected_rows());
 }
+TEST_P(TestParquetFileFormatScan, ScanRecordBatchReaderReeEncoded) {
+  auto reader = GetRecordBatchReader(schema({field("utf8", utf8())}));
+  auto source = GetFileSource(reader.get());
+
+  SetSchema(reader->schema()->fields());
+  SetFilter(literal(true));
+  format_->reader_options.ree_columns = {"utf8"};
+  ASSERT_OK_AND_ASSIGN(auto fragment, format_->MakeFragment(*source));
+
+  int64_t row_count = 0;
+  Schema expected_schema({field("utf8", run_end_encoded(int32(), utf8()))});
+
+  for (auto maybe_batch : PhysicalBatches(fragment)) {
+    ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
+    row_count += batch->num_rows();
+    AssertSchemaEqual(*batch->schema(), expected_schema, /* check_metadata = */ false);
+  }
+  ASSERT_EQ(row_count, expected_rows());
+}
 TEST_P(TestParquetFileFormatScan, ScanRecordBatchReaderPreBuffer) {
   auto reader = GetRecordBatchReader(schema({field("f64", float64())}));
   auto source = GetFileSource(reader.get());

cpp/src/arrow/util/rle_encoding_internal.h

@@ -426,6 +426,14 @@ class RleBitPackedDecoder {
   /// values.
   [[nodiscard]] bool Get(value_type* val);
 
+  // Get the next logical value and num_repeats within the specified batch_size.
+  [[nodiscard]] bool GetNextValueAndNumRepeats(value_type* val, int* num_repeats, int batch_size);
+
+  /// Like GetNextValueAndNumRepeats but add spacing for null entries.
+  [[nodiscard]] bool GetNextValueAndNumRepeatsSpaced(value_type* val, bool* is_null,
+                                             int* num_repeats, int batch_size,
+                                             const uint8_t* valid_bits, int64_t valid_bits_offset);
+
   /// Get a batch of values return the number of decoded elements.
   /// May write fewer elements to the output than requested if there are not enough values
   /// left or if an error occurred.
@@ -722,7 +730,6 @@ void RleBitPackedDecoder<T>::ParseWithCallable(Callable&& func) {
     auto OnBitPackedRun(BitPackedRun run) { return func(std::move(run)); }
     auto OnRleRun(RleRun run) { return func(std::move(run)); }
   } handler{std::move(func)};
-
   parser_.Parse(std::move(handler));
 }
 
@@ -731,6 +738,69 @@ bool RleBitPackedDecoder<T>::Get(value_type* val) {
   return GetBatch(val, 1) == 1;
 }
 
+template <typename T>
+bool RleBitPackedDecoder<T>::GetNextValueAndNumRepeats(value_type* val, int* num_repeats, int batch_size) {
+  using ControlFlow = RleBitPackedParser::ControlFlow;
+
+  if (ARROW_PREDICT_FALSE(run_remaining() > 0)) {
+    if (std::holds_alternative<BitPackedRunDecoder<value_type>>(decoder_)) {
+      auto& decoder = std::get<BitPackedRunDecoder<value_type>>(decoder_);
+      *num_repeats = 1;
+      return decoder.Get(val, value_bit_width_);
+    } else {
+      auto& decoder = std::get<RleRunDecoder<value_type>>(decoder_);
+      *num_repeats = std::min(decoder.remaining(), batch_size);
+      ARROW_DCHECK_EQ(decoder.Advance(*num_repeats, value_bit_width_), *num_repeats);
+      return decoder.Get(val, value_bit_width_);
+    }
+  }
+
+  bool read_new_value = false;
+
+  ParseWithCallable([&](auto run) {
+    if constexpr(std::is_same_v<decltype(run), BitPackedRun>) {
+      BitPackedRunDecoder<T> decoder(run, value_bit_width_);
+      read_new_value = decoder.Get(val, value_bit_width_);
+      *num_repeats = 1;
+      decoder_ = std::move(decoder);
+      return ControlFlow::Break;
+    }
+    else {
+      RleRunDecoder<T> decoder(run, value_bit_width_);
+      *num_repeats = std::min(decoder.remaining(), batch_size);
+      read_new_value = decoder.Get(val, value_bit_width_);
+      ARROW_DCHECK_EQ(decoder.Advance(*num_repeats, value_bit_width_), *num_repeats);
+      decoder_ = std::move(decoder);
+      return ControlFlow::Break;
+    }
+  });
+
+  return read_new_value;
+}
+
+template <typename T>
+bool RleBitPackedDecoder<T>::GetNextValueAndNumRepeatsSpaced(value_type* val, bool* is_null,
+                                                     int* num_repeats, int batch_size,
+                                                     const uint8_t* valid_bits, int64_t valid_bits_offset) {
+  arrow::internal::BitRunReader bit_reader(valid_bits, valid_bits_offset,
+                                           /*length=*/batch_size);
+  arrow::internal::BitRun valid_run = bit_reader.NextRun();
+  while (ARROW_PREDICT_FALSE(valid_run.length == 0)) {
+    valid_run = bit_reader.NextRun();
+  }
+  ARROW_DCHECK_GT(batch_size, 0);
+  ARROW_DCHECK_GT(valid_run.length, 0);
+  if (valid_run.set) {
+    return GetNextValueAndNumRepeats(
+        val, num_repeats,
+        static_cast<int>(std::min(valid_run.length, static_cast<int64_t>(batch_size))));
+  } else {
+    *is_null = true;
+    *num_repeats = static_cast<int>(valid_run.length);
+  }
+  return true;
+}
+
 template <typename T>
 auto RleBitPackedDecoder<T>::GetBatch(value_type* out, rle_size_t batch_size)
     -> rle_size_t {

cpp/src/parquet/CMakeLists.txt

@@ -406,7 +406,10 @@ add_parquet_test(arrow-reader-writer-test
                  SOURCES
                  arrow/arrow_reader_writer_test.cc
                  arrow/arrow_statistics_test.cc
-                 arrow/variant_test.cc)
+                 arrow/variant_test.cc
+                 EXTRA_LINK_LIBS
+                 arrow_compute_core_testing
+                 arrow_acero_testing)
 
 add_parquet_test(arrow-internals-test SOURCES arrow/path_internal_test.cc
                  arrow/reconstruct_internal_test.cc)