GH-43660: [C++] Add a CastingGenerator to Parquet Reader that applies required casts before slicing

cpp/src/arrow/dataset/CMakeLists.txt

@@ -241,3 +241,7 @@ endfunction()
 
 add_arrow_dataset_benchmark(file_benchmark)
 add_arrow_dataset_benchmark(scanner_benchmark)
+
+if(ARROW_PARQUET)
+  add_arrow_dataset_benchmark(parquet_scan_benchmark)
+endif()

cpp/src/arrow/dataset/file_parquet.cc

@@ -26,11 +26,17 @@
 
 #include "arrow/compute/cast.h"
 #include "arrow/compute/exec.h"
+#include "arrow/dataset/dataset.h"
 #include "arrow/dataset/dataset_internal.h"
 #include "arrow/dataset/parquet_encryption_config.h"
 #include "arrow/dataset/scanner.h"
 #include "arrow/filesystem/path_util.h"
+#include "arrow/memory_pool.h"
+#include "arrow/record_batch.h"
+#include "arrow/result.h"
 #include "arrow/table.h"
+#include "arrow/type.h"
+#include "arrow/type_fwd.h"
 #include "arrow/util/checked_cast.h"
 #include "arrow/util/future.h"
 #include "arrow/util/iterator.h"
@@ -565,6 +571,66 @@ Future<std::shared_ptr<parquet::arrow::FileReader>> ParquetFileFormat::GetReader
       });
 }
 
+struct CastingGenerator {
+  CastingGenerator(RecordBatchGenerator source, std::shared_ptr<Schema> final_schema,
+                   const std::unordered_set<std::string>& cols_to_skip,
+                   MemoryPool* pool = default_memory_pool())
+      : source_(source),
+        final_schema_(std::move(final_schema)),
+        cols_to_skip_(cols_to_skip),
+        exec_ctx(std::make_shared<compute::ExecContext>(pool)) {}
+
+  Future<std::shared_ptr<RecordBatch>> operator()() {
+    return this->source_().Then([this](const std::shared_ptr<RecordBatch>& next)
+                                    -> Result<std::shared_ptr<RecordBatch>> {
+      if (IsIterationEnd(next) || this->final_schema_ == nullptr) {
+        return next;
+      }
+      std::vector<std::shared_ptr<Array>> out_cols;
+      std::vector<std::shared_ptr<Field>> out_schema_fields;
+
+      bool changed = false;
+      for (const auto& field : this->final_schema_->fields()) {
+        FieldRef field_ref = FieldRef(field->name());
+        ARROW_ASSIGN_OR_RAISE(std::shared_ptr<Array> column,
+                              field_ref.GetOneOrNone(*next));
+        if (column) {
+          if (this->cols_to_skip_.count(field->name())) {
+            // Maintain the original input type.
+            out_schema_fields.emplace_back(field->WithType(column->type()));
+            out_cols.emplace_back(std::move(column));
+            continue;
+          }
+          if (!column->type()->Equals(field->type())) {
+            // Referenced field was present but didn't have the expected type.
+            ARROW_ASSIGN_OR_RAISE(
+                auto converted,
+                compute::Cast(column, field->type(), compute::CastOptions::Safe(),
+                              this->exec_ctx.get()));
+            column = converted.make_array();
+            changed = true;
+          }
+          out_schema_fields.emplace_back(field->Copy());
+          out_cols.emplace_back(std::move(column));
+        }
+      }
+
+      if (changed) {
+        return RecordBatch::Make(std::make_shared<Schema>(std::move(out_schema_fields),
+                                                          next->schema()->metadata()),
+                                 next->num_rows(), std::move(out_cols));
+      } else {
+        return next;
+      }
+    });
+  }
+
+  RecordBatchGenerator source_;
+  std::shared_ptr<Schema> final_schema_;
+  const std::unordered_set<std::string>& cols_to_skip_;
+  std::shared_ptr<compute::ExecContext> exec_ctx;
+};
+
 struct SlicingGenerator {
   SlicingGenerator(RecordBatchGenerator source, int64_t batch_size)
       : state(std::make_shared<State>(source, batch_size)) {}
@@ -627,6 +693,9 @@ Result<RecordBatchGenerator> ParquetFileFormat::ScanBatchesAsync(
       [this, options, parquet_fragment, pre_filtered,
        row_groups](const std::shared_ptr<parquet::arrow::FileReader>& reader) mutable
       -> Result<RecordBatchGenerator> {
+    // Since we already do the batching through the SlicingGenerator, we don't need the
+    // reader to batch its output.
+    reader->set_batch_size(std::numeric_limits<int64_t>::max());
     // Ensure that parquet_fragment has FileMetaData
     RETURN_NOT_OK(parquet_fragment->EnsureCompleteMetadata(reader.get()));
     if (!pre_filtered) {
@@ -649,8 +718,17 @@ Result<RecordBatchGenerator> ParquetFileFormat::ScanBatchesAsync(
     ARROW_ASSIGN_OR_RAISE(auto generator, reader->GetRecordBatchGenerator(
                                               reader, row_groups, column_projection,
                                               cpu_executor, rows_to_readahead));
+    // We need to skip casting the dictionary columns since the dataset_schema doesn't
+    // have the dictionary-encoding information. Parquet reader will return them with the
+    // dictionary type, which is what we eventually want.
+    const std::unordered_set<std::string>& dict_cols =
+        parquet_fragment->parquet_format_.reader_options.dict_columns;
+    // Casting before slicing is more efficient. Casts on slices might require wasteful
+    // allocations and computation.
+    RecordBatchGenerator casted = CastingGenerator(
+        std::move(generator), options->dataset_schema, dict_cols, options->pool);
     RecordBatchGenerator sliced =
-        SlicingGenerator(std::move(generator), options->batch_size);
+        SlicingGenerator(std::move(casted), options->batch_size);
     if (batch_readahead == 0) {
       return sliced;
     }

cpp/src/arrow/dataset/parquet_scan_benchmark.cc

@@ -0,0 +1,94 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/testing/gtest_util.h"
+#include "benchmark/benchmark.h"
+
+#include "arrow/api.h"
+#include "arrow/dataset/file_parquet.h"
+#include "arrow/io/memory.h"
+#include "parquet/arrow/writer.h"
+
+namespace arrow {
+namespace dataset {
+
+using parquet::arrow::WriteTable;
+
+Result<std::shared_ptr<Buffer>> WriteStringColParquetBuffer(int64_t nrows) {
+  auto schema = arrow::schema({arrow::field("my_string_col", arrow::utf8())});
+
+  arrow::StringBuilder builder;
+  for (int64_t i = 0; i < nrows; i++) {
+    ARROW_RETURN_NOT_OK(builder.Append("row_" + std::to_string(i)));
+  }
+  std::shared_ptr<arrow::Array> arr;
+  ARROW_RETURN_NOT_OK(builder.Finish(&arr));
+  auto table = arrow::Table::Make(schema, {arr});
+
+  ARROW_ASSIGN_OR_RAISE(auto sink, arrow::io::BufferOutputStream::Create());
+  ARROW_RETURN_NOT_OK(WriteTable(*table, arrow::default_memory_pool(), sink));
+  return sink->Finish();
+}
+
+static void ParquetScanToTableCastStrings(benchmark::State& state) {
+  // GH-43660: Scan parquet data including a String column using a dataset object with
+  // LargeString in schema.
+  size_t num_batches = state.range(0);
+  size_t batch_size = state.range(1);
+  size_t nrows = num_batches * batch_size;
+  auto format = std::make_shared<ParquetFileFormat>();
+
+  // Create a buffer with a single String column and wrap with FileFragment
+  ASSERT_OK_AND_ASSIGN(std::shared_ptr<Buffer> buffer,
+                       WriteStringColParquetBuffer(nrows));
+  auto buffer_reader = std::make_shared<arrow::io::BufferReader>(buffer);
+  FileSource source(buffer_reader, buffer->size());
+  ASSERT_OK_AND_ASSIGN(auto fragment, format->MakeFragment(source));
+  std::vector<std::shared_ptr<FileFragment>> fragments{fragment};
+
+  // Create a dataset from FileFragment and set schema to LargeString (require casting).
+  auto schema = arrow::schema({field("my_string_col", arrow::large_utf8())});
+  ASSERT_OK_AND_ASSIGN(auto dataset, FileSystemDataset::Make(
+                                         schema, compute::literal(true), format,
+                                         /*filesystem=*/nullptr, std::move(fragments)));
+
+  ASSERT_OK_AND_ASSIGN(auto builder, dataset->NewScan());
+  ASSERT_OK(builder->BatchSize(batch_size));
+  ASSERT_OK_AND_ASSIGN(auto scanner, builder->Finish());
+
+  for (auto _ : state) {
+    ASSERT_OK_AND_ASSIGN(auto table, scanner->ToTable());
+    benchmark::DoNotOptimize(table);
+  }
+
+  state.SetItemsProcessed(state.iterations() * nrows);
+}
+
+static void ParquetScanBenchmark_Customize(benchmark::internal::Benchmark* b) {
+  std::vector<int64_t> num_batches = {1000, 100, 10};
+  std::vector<int64_t> batch_sizes = {1000, 10000, 100000};
+
+  for (size_t i = 0; i < batch_sizes.size(); i++) {
+    b->Args({(num_batches[i]), batch_sizes[i]});
+  }
+  b->ArgNames({"num_batches", "batch_size"});
+}
+
+BENCHMARK(ParquetScanToTableCastStrings)->Apply(ParquetScanBenchmark_Customize);
+
+}  // namespace dataset
+}  // namespace arrow