From 179f753a0b32e0250a74a4ad506ddc7286e3e83b Mon Sep 17 00:00:00 2001 From: jiangxt2 Date: Sat, 31 Oct 2026 17:58:07 +0800 Subject: [PATCH] GH-55326: [C++][Python] Preserve field identity in ORC Dataset scans Pass per-file top-level field indices to the ORC batch reader so literal dotted names cannot be confused with nested paths. Preserve the existing name-based API and add adapter and Dataset regression coverage. Signed-off-by: jiangxt2 --- cpp/src/arrow/adapters/orc/adapter.cc | 33 +++++++++++--- cpp/src/arrow/adapters/orc/adapter.h | 10 ++++ cpp/src/arrow/adapters/orc/adapter_test.cc | 48 ++++++++++++++++++++ cpp/src/arrow/dataset/file_orc.cc | 10 ++-- cpp/src/arrow/dataset/file_orc_test.cc | 20 ++++++++ python/pyarrow/tests/test_dataset.py | 53 ++++++++++++++++++++++ 6 files changed, 163 insertions(+), 11 deletions(-) diff --git a/cpp/src/arrow/adapters/orc/adapter.cc b/cpp/src/arrow/adapters/orc/adapter.cc index 51cca497485c..6928502be9b4 100644 --- a/cpp/src/arrow/adapters/orc/adapter.cc +++ b/cpp/src/arrow/adapters/orc/adapter.cc @@ -532,6 +532,28 @@ class ORCFileReader::Impl { if (!include_names.empty()) { RETURN_NOT_OK(SelectNames(&opts, include_names)); } + return MakeRecordBatchReader(batch_size, opts); + } + + Result> GetRecordBatchReaderByIndices( + int64_t batch_size, const std::vector& include_indices) { + ARROW_RETURN_IF(batch_size <= 0, + Status::Invalid("batch_size must be greater than zero")); + auto opts = DefaultRowReaderOptions(); + if (!include_indices.empty()) { + RETURN_NOT_OK(SelectIndices(&opts, include_indices)); + } + return MakeRecordBatchReader(batch_size, opts); + } + + Result> NextStripeReader(int64_t batch_size) { + std::vector empty_vec; + return NextStripeReader(batch_size, empty_vec); + } + + private: + Result> MakeRecordBatchReader( + int64_t batch_size, const liborc::RowReaderOptions& opts) { ARROW_ASSIGN_OR_RAISE(auto schema, ReadSchema(opts)); std::unique_ptr row_reader; @@ -543,12 +565,6 @@ class ORCFileReader::Impl { pool_); } - Result> NextStripeReader(int64_t batch_size) { - std::vector empty_vec; - return NextStripeReader(batch_size, empty_vec); - } - - private: MemoryPool* pool_; std::unique_ptr reader_; std::vector stripes_; @@ -625,6 +641,11 @@ Result> ORCFileReader::GetRecordBatchReader( return impl_->GetRecordBatchReader(batch_size, include_names); } +Result> ORCFileReader::GetRecordBatchReaderByIndices( + int64_t batch_size, const std::vector& include_indices) { + return impl_->GetRecordBatchReaderByIndices(batch_size, include_indices); +} + Result> ORCFileReader::NextStripeReader( int64_t batch_size, const std::vector& include_indices) { return impl_->NextStripeReader(batch_size, include_indices); diff --git a/cpp/src/arrow/adapters/orc/adapter.h b/cpp/src/arrow/adapters/orc/adapter.h index 4ffff81f355f..2742484ad054 100644 --- a/cpp/src/arrow/adapters/orc/adapter.h +++ b/cpp/src/arrow/adapters/orc/adapter.h @@ -174,6 +174,16 @@ class ARROW_EXPORT ORCFileReader { Result> GetRecordBatchReader( int64_t batch_size, const std::vector& include_names); + /// \brief Get an entire-file record batch iterator by top-level field index. + /// + /// Field order follows the file schema; repeated indices do not duplicate fields. + /// An empty selection reads all fields, as in GetRecordBatchReader(). + /// \param[in] batch_size the maximum number of rows per batch, which must be positive + /// \param[in] include_indices the zero-based top-level field indices to read + /// \return the record batch iterator + Result> GetRecordBatchReaderByIndices( + int64_t batch_size, const std::vector& include_indices); + /// \brief The number of stripes in the file int64_t NumberOfStripes(); diff --git a/cpp/src/arrow/adapters/orc/adapter_test.cc b/cpp/src/arrow/adapters/orc/adapter_test.cc index 3cbb6d7828f1..350be2ffb35e 100644 --- a/cpp/src/arrow/adapters/orc/adapter_test.cc +++ b/cpp/src/arrow/adapters/orc/adapter_test.cc @@ -20,6 +20,8 @@ #include #include +#include +#include #include "arrow/adapters/orc/adapter.h" #include "arrow/adapters/orc/util.h" @@ -30,6 +32,7 @@ #include "arrow/compute/cast.h" #include "arrow/io/interfaces.h" #include "arrow/io/memory.h" +#include "arrow/result.h" #include "arrow/status.h" #include "arrow/table.h" #include "arrow/testing/gtest_util.h" @@ -454,6 +457,51 @@ TEST(TestAdapterRead, ReadIntAndStringFileMultipleStripes) { } EXPECT_EQ(num_rows, stripe_row_count * stripe_count); EXPECT_EQ(num_rows / reader_batch_size, batches); + + ASSERT_OK_AND_ASSIGN(auto indexed, + reader->GetRecordBatchReaderByIndices(reader_batch_size, {0, 1})); + ASSERT_OK_AND_ASSIGN(auto indexed_table, indexed->ToTable()); + ASSERT_OK_AND_ASSIGN(auto expected_table, reader->Read()); + AssertTablesEqual(*expected_table, *indexed_table, false, false); +} + +TEST(TestAdapterRead, RecordBatchReaderByIndices) { + auto table = TableFromJSON( + schema({field("a", struct_({field("b", int64())})), field("a.b", int64()), + field("other", struct_({field("v", int64())}))}), + {R"([[{"b": 10}, 1, {"v": 101}], [null, 2, {"v": 102}], + [{"b": 30}, null, null]])"}); + ASSERT_OK_AND_ASSIGN(auto sink, io::BufferOutputStream::Create()); + ASSERT_OK_AND_ASSIGN(auto writer, adapters::orc::ORCFileWriter::Open(sink.get())); + ASSERT_OK(writer->Write(*table)); + ASSERT_OK(writer->Close()); + ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish()); + ASSERT_OK_AND_ASSIGN(auto reader, adapters::orc::ORCFileReader::Open( + std::make_shared(buffer), + default_memory_pool())); + auto check = [&](Result> maybe_stream, + const std::vector& columns) { + ASSERT_OK_AND_ASSIGN(auto stream, std::move(maybe_stream)); + ASSERT_OK_AND_ASSIGN(auto batches, stream->ToRecordBatches()); + for (const auto& batch : batches) { + ASSERT_LE(batch->num_rows(), 2); + } + ASSERT_OK_AND_ASSIGN(auto actual, Table::FromRecordBatches(batches)); + ASSERT_OK_AND_ASSIGN(auto expected, table->SelectColumns(columns)); + AssertTablesEqual(*expected, *actual, false, false); + }; + using Selection = std::pair, std::vector>; + for (const auto& [indices, columns] : + {Selection{{1}, {1}}, {{0}, {0}}, {{2, 1, 1}, {1, 2}}, {{}, {0, 1, 2}}}) { + check(reader->GetRecordBatchReaderByIndices(2, indices), columns); + } + check(reader->GetRecordBatchReader(2, {"a"}), {0}); + check(reader->GetRecordBatchReader(2, {"other.v"}), {2}); + check(reader->GetRecordBatchReader(2, {}), {0, 1, 2}); + ASSERT_RAISES(Invalid, reader->GetRecordBatchReaderByIndices(2, {-1})); + ASSERT_RAISES(IOError, reader->GetRecordBatchReaderByIndices(2, {3})); + ASSERT_RAISES(Invalid, reader->GetRecordBatchReaderByIndices(0, {1})); + ASSERT_RAISES(Invalid, reader->GetRecordBatchReaderByIndices(-1, {1})); } TEST(TestAdapterRead, ReadCharAndVarcharType) { diff --git a/cpp/src/arrow/dataset/file_orc.cc b/cpp/src/arrow/dataset/file_orc.cc index 1393df57f9d7..e0cb1373840f 100644 --- a/cpp/src/arrow/dataset/file_orc.cc +++ b/cpp/src/arrow/dataset/file_orc.cc @@ -76,19 +76,19 @@ class OrcScanTask { auto materialized_fields = scan_options.MaterializedFields(); // filter out virtual columns - std::vector included_fields; + std::vector included_indices; ARROW_ASSIGN_OR_RAISE(auto schema, reader->ReadSchema()); for (const auto& ref : materialized_fields) { ARROW_ASSIGN_OR_RAISE(auto match, ref.FindOneOrNone(*schema)); if (match.indices().empty()) continue; - included_fields.push_back(schema->field(match.indices()[0])->name()); + included_indices.push_back(match.indices()[0]); } std::shared_ptr record_batch_reader; - ARROW_ASSIGN_OR_RAISE( - record_batch_reader, - reader->GetRecordBatchReader(scan_options.batch_size, included_fields)); + ARROW_ASSIGN_OR_RAISE(record_batch_reader, + reader->GetRecordBatchReaderByIndices( + scan_options.batch_size, included_indices)); return RecordBatchIterator(Impl{std::move(record_batch_reader)}); } diff --git a/cpp/src/arrow/dataset/file_orc_test.cc b/cpp/src/arrow/dataset/file_orc_test.cc index 17be015de516..9057f67bce26 100644 --- a/cpp/src/arrow/dataset/file_orc_test.cc +++ b/cpp/src/arrow/dataset/file_orc_test.cc @@ -19,6 +19,7 @@ #include #include +#include #include "arrow/adapters/orc/adapter.h" #include "arrow/dataset/dataset_internal.h" @@ -75,6 +76,25 @@ TEST_P(TestOrcFileFormatScan, ScanRecordBatchReaderProjected) { TestScanProjecte TEST_P(TestOrcFileFormatScan, ScanRecordBatchReaderProjectedNested) { TestScanProjectedNested(); } +TEST_P(TestOrcFileFormatScan, ScanDottedFieldNameCollision) { + auto batch = RecordBatchFromJSON( + schema({field("a.b", int32()), field("a", struct_({field("b", int32())})), + field("other", int32())}), + R"([[1, {"b": 10}, 101], [2, {"b": 20}, 102]])"); + ASSERT_OK_AND_ASSIGN(auto expected, batch->SelectColumns({0, 2})); + for (const auto& indices : {std::vector{0, 1, 2}, {1, 0, 2}}) { + ASSERT_OK_AND_ASSIGN(auto reordered, batch->SelectColumns(indices)); + this->SetSchema(reordered->schema()->fields()); + this->Project({"a.b"}); + this->SetFilter(greater(field_ref("other"), literal(100))); + ASSERT_OK_AND_ASSIGN(auto reader, RecordBatchReader::Make({reordered})); + auto source = this->GetFileSource(reader.get()); + ASSERT_OK_AND_ASSIGN(auto batches, + PhysicalBatches(this->MakeFragment(*source)).ToVector()); + ASSERT_EQ(batches.size(), 1); + AssertBatchesEqual(*expected, *batches[0]); + } +} TEST_P(TestOrcFileFormatScan, ScanRecordBatchReaderProjectedMissingCols) { TestScanProjectedMissingCols(); } diff --git a/python/pyarrow/tests/test_dataset.py b/python/pyarrow/tests/test_dataset.py index 9ddfdfdc9c49..ab11f5896555 100644 --- a/python/pyarrow/tests/test_dataset.py +++ b/python/pyarrow/tests/test_dataset.py @@ -3432,6 +3432,59 @@ def test_orc_format(tempdir, dataset_reader): assert dataset_reader.count_rows(dataset, filter=ds.field("a") > 2) == 1 +@pytest.mark.orc +@pytest.mark.parametrize('literal_first', [True, False]) +@pytest.mark.parametrize('in_memory', [True, False]) +@pytest.mark.parametrize('empty', [True, False]) +def test_orc_dotted_field_identity( + tempdir, dataset_reader, literal_first, in_memory, empty): + from pyarrow import orc + table = pa.table({'a.b': [1, 2, None], + 'a': [{'b': 10}, None, {'b': 30}], + 'other': [101, 102, 103]}) + if not literal_first: + table = table.select(['a', 'a.b', 'other']) + if empty: + table = table.slice(0, 0) + output = pa.BufferOutputStream() if in_memory else str(tempdir / 'dotted.orc') + orc.write_table(table, output) + dataset = ( + ds.OrcFileFormat().make_fragment(pa.BufferReader(output.getvalue())) + if in_memory else ds.dataset(output, format='orc')) + + assert dataset_reader.to_table(dataset).equals(table) + projected = dataset_reader.to_table(dataset, columns=['a.b'], batch_size=1) + assert projected.equals(table.select(['a.b'])) + assert all(len(chunk) <= 1 for chunk in projected['a.b'].chunks) + predicate = ds.field('a.b') > 1 + assert dataset_reader.to_table( + dataset, columns=['other'], filter=predicate).equals( + table.filter(predicate).select(['other'])) + columns = {'literal': ds.field('a.b'), 'nested': ds.field(('a', 'b'))} + expected = pa.table({'literal': [1, 2, None], + 'nested': [10, None, 30]}).slice(0, table.num_rows) + assert dataset_reader.to_table(dataset, columns=columns).equals(expected) + assert dataset_reader.to_table(dataset, columns=[]).num_rows == table.num_rows + + +@pytest.mark.orc +def test_orc_dotted_field_schema_evolution(tempdir, dataset_reader): + from pyarrow import orc + table = pa.table({'a.b': [1, 2], 'a': [{'b': 10}, {'b': 20}]}) + paths = [] + for i, columns in enumerate([['a.b', 'a'], ['a', 'a.b'], ['a']]): + path = str(tempdir / f'part{i}.orc') + orc.write_table(table.select(columns), path) + paths.append(path) + dataset = ds.dataset(paths, format='orc', schema=table.schema) + missing = table.set_column(0, table.schema.field(0), pa.nulls(2, pa.int64())) + expected = pa.concat_tables([table, table, missing]) + assert dataset_reader.to_table(dataset).equals(expected) + assert dataset_reader.to_table( + dataset, columns=['a.b'], filter=ds.field('a.b') > 1).equals( + pa.table({'a.b': [2, 2]})) + + @pytest.mark.orc def test_orc_scan_options(tempdir, dataset_reader): from pyarrow import orc