GH-55326: [C++][Python] Preserve field identity in ORC Dataset scans

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<std::shared_ptr<RecordBatchReader>> GetRecordBatchReaderByIndices(
+      int64_t batch_size, const std::vector<int>& 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<std::shared_ptr<RecordBatchReader>> NextStripeReader(int64_t batch_size) {
+    std::vector<int> empty_vec;
+    return NextStripeReader(batch_size, empty_vec);
+  }
+
+ private:
+  Result<std::shared_ptr<RecordBatchReader>> MakeRecordBatchReader(
+      int64_t batch_size, const liborc::RowReaderOptions& opts) {
     ARROW_ASSIGN_OR_RAISE(auto schema, ReadSchema(opts));
     std::unique_ptr<liborc::RowReader> row_reader;
 
@@ -543,12 +565,6 @@ class ORCFileReader::Impl {
                                              pool_);
   }
 
-  Result<std::shared_ptr<RecordBatchReader>> NextStripeReader(int64_t batch_size) {
-    std::vector<int> empty_vec;
-    return NextStripeReader(batch_size, empty_vec);
-  }
-
- private:
   MemoryPool* pool_;
   std::unique_ptr<liborc::Reader> reader_;
   std::vector<StripeInformation> stripes_;
@@ -625,6 +641,11 @@ Result<std::shared_ptr<RecordBatchReader>> ORCFileReader::GetRecordBatchReader(
   return impl_->GetRecordBatchReader(batch_size, include_names);
 }
 
+Result<std::shared_ptr<RecordBatchReader>> ORCFileReader::GetRecordBatchReaderByIndices(
+    int64_t batch_size, const std::vector<int>& include_indices) {
+  return impl_->GetRecordBatchReaderByIndices(batch_size, include_indices);
+}
+
 Result<std::shared_ptr<RecordBatchReader>> ORCFileReader::NextStripeReader(
     int64_t batch_size, const std::vector<int>& include_indices) {
   return impl_->NextStripeReader(batch_size, include_indices);

cpp/src/arrow/adapters/orc/adapter.h

@@ -174,6 +174,16 @@ class ARROW_EXPORT ORCFileReader {
   Result<std::shared_ptr<RecordBatchReader>> GetRecordBatchReader(
       int64_t batch_size, const std::vector<std::string>& 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<std::shared_ptr<RecordBatchReader>> GetRecordBatchReaderByIndices(
+      int64_t batch_size, const std::vector<int>& include_indices);
+
   /// \brief The number of stripes in the file
   int64_t NumberOfStripes();
 

cpp/src/arrow/adapters/orc/adapter_test.cc

@@ -20,6 +20,8 @@
 
 #include <orc/OrcFile.hh>
 #include <string>
+#include <utility>
+#include <vector>
 
 #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,56 @@ 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 input = TableFromJSON(
+      schema({field("a", struct_({field("b", int64())})),
+              field("a.b", int64(), true, key_value_metadata({"role"}, {"literal"})),
+              field("other", struct_({field("v", int64())}))}),
+      {R"([[{"b": 10}, 1, {"v": 101}], [null, 2, {"v": 102}],
+          [{"b": 30}, null, null]])"});
+  for (const auto& table : {input, input->Slice(0, 0)}) {
+    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<io::BufferReader>(buffer),
+                                          default_memory_pool()));
+    auto check = [&](Result<std::shared_ptr<RecordBatchReader>> maybe_stream,
+                     const std::vector<int>& columns) {
+      ASSERT_OK_AND_ASSIGN(auto stream, std::move(maybe_stream));
+      ASSERT_OK_AND_ASSIGN(auto expected, table->SelectColumns(columns));
+      AssertSchemaEqual(expected->schema(), stream->schema(), /*check_metadata=*/true);
+      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(stream->schema(), batches));
+      AssertTablesEqual(*expected, *actual, false, false);
+    };
+    using Selection = std::pair<std::vector<int>, std::vector<int>>;
+    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) {

cpp/src/arrow/dataset/file_orc.cc

@@ -18,6 +18,8 @@
 #include "arrow/dataset/file_orc.h"
 
 #include <memory>
+#include <utility>
+#include <vector>
 
 #include "arrow/adapters/orc/adapter.h"
 #include "arrow/dataset/dataset_internal.h"
@@ -58,6 +60,19 @@ Result<std::unique_ptr<arrow::adapters::orc::ORCFileReader>> OpenORCReader(
   return reader;
 }
 
+// Preserve each name component when mapping logical positions to file fields.
+FieldRef ToUniversalRef(const FieldPath& path, const Schema& schema) {
+  std::vector<FieldRef> refs;
+  refs.reserve(path.indices().size());
+  const FieldVector* fields = &schema.fields();
+  for (int index : path) {
+    const auto& field = (*fields)[index];
+    refs.emplace_back(field->name());
+    fields = &field->type()->fields();
+  }
+  return FieldRef(std::move(refs));
+}
+
 /// \brief A ScanTask backed by an ORC file.
 class OrcScanTask {
  public:
@@ -76,19 +91,25 @@ class OrcScanTask {
 
         auto materialized_fields = scan_options.MaterializedFields();
         // filter out virtual columns
-        std::vector<std::string> included_fields;
+        std::vector<int> included_indices;
         ARROW_ASSIGN_OR_RAISE(auto schema, reader->ReadSchema());
+        const auto& dataset_schema = *scan_options.dataset_schema;
+        const bool same_schema = dataset_schema.Equals(*schema);
         for (const auto& ref : materialized_fields) {
-          ARROW_ASSIGN_OR_RAISE(auto match, ref.FindOneOrNone(*schema));
+          ARROW_ASSIGN_OR_RAISE(auto path, ref.FindOneOrNone(dataset_schema));
+          if (path.empty()) continue;
+          auto file_ref =
+              same_schema ? FieldRef(path) : ToUniversalRef(path, dataset_schema);
+          ARROW_ASSIGN_OR_RAISE(auto match, file_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<RecordBatchReader> 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)});
       }

cpp/src/arrow/dataset/file_orc_test.cc

@@ -19,6 +19,7 @@
 
 #include <memory>
 #include <utility>
+#include <vector>
 
 #include "arrow/adapters/orc/adapter.h"
 #include "arrow/dataset/dataset_internal.h"
@@ -75,6 +76,30 @@ 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]])");
+  this->SetSchema(
+      {batch->schema()->field(1), batch->schema()->field(0), batch->schema()->field(2)});
+  using Projection = std::pair<FieldRef, int>;
+  for (const auto& indices : {std::vector<int>{0, 1, 2}, {1, 0, 2}}) {
+    ASSERT_OK_AND_ASSIGN(auto reordered, batch->SelectColumns(indices));
+    for (const auto& [ref, field_index] :
+         {Projection{FieldRef("a.b"), 0}, {FieldRef(0), 1}, {FieldRef(0, 0), 1}}) {
+      this->Project({field_ref(ref)}, {"selected"});
+      this->SetFilter(greater(field_ref("other"), literal(100)));
+      ASSERT_OK_AND_ASSIGN(auto expected, batch->SelectColumns({field_index, 2}));
+      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();
 }

python/pyarrow/tests/test_dataset.py

@@ -3432,6 +3432,93 @@ 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
+@pytest.mark.parametrize('nested', [False, True])
+@pytest.mark.parametrize('use_filter', [False, True])
+def test_orc_positional_field_identity(tempdir, dataset_reader, nested, use_filter):
+    from pyarrow import orc
+    table = pa.table({'a': [{'b': 10}, None, {'b': None}], 'a.b': [1, 2, 3]})
+    path = str(tempdir / 'positional.orc')
+    orc.write_table(table.select(['a.b', 'a']), path)
+    dataset = ds.dataset(path, format='orc', schema=table.schema)
+    reference = ds.field((0, 0)) if nested else ds.field(0)
+    if use_filter:
+        actual = dataset_reader.to_table(
+            dataset, columns=['a.b'], filter=reference.is_valid())
+        expected = pa.table({'a.b': [1] if nested else [1, 3]})
+    else:
+        actual = dataset_reader.to_table(dataset, columns={'selected': reference})
+        values = table['a'].flatten()[0] if nested else table['a']
+        expected = pa.table({'selected': values})
+    assert actual.equals(expected)
+
+
+@pytest.mark.orc
+def test_orc_duplicate_field_positional_projection(dataset_reader):
+    from pyarrow import orc
+    table = pa.Table.from_arrays([[1, 2], [3, 4]], names=['f', 'f'])
+    sink = pa.BufferOutputStream()
+    orc.write_table(table, sink)
+    fragment = ds.OrcFileFormat().make_fragment(pa.BufferReader(sink.getvalue()))
+    for index in range(2):
+        actual = dataset_reader.to_table(
+            fragment, schema=table.schema, columns={'selected': ds.field(index)})
+        assert actual.equals(table.select([index]).rename_columns(['selected']))
+
+
+@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