[C++][Dataset] Apply stripe-selected ORC scanning with conservative null/overflow handling

Author: cbb330

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

@@ -642,6 +642,31 @@ class ORCFileReader::Impl {
                                              pool_);
   }
 
+  Result<std::shared_ptr<RecordBatchReader>> NextStripeReader(
+      int64_t batch_size, const std::vector<std::string>& include_names) {
+    if (current_row_ >= NumberOfRows()) {
+      return nullptr;
+    }
+
+    liborc::RowReaderOptions opts = DefaultRowReaderOptions();
+    if (!include_names.empty()) {
+      RETURN_NOT_OK(SelectNames(&opts, include_names));
+    }
+    StripeInformation stripe_info{0, 0, 0, 0};
+    RETURN_NOT_OK(SelectStripeWithRowNumber(&opts, current_row_, &stripe_info));
+    ARROW_ASSIGN_OR_RAISE(auto schema, ReadSchema(opts));
+    std::unique_ptr<liborc::RowReader> row_reader;
+
+    ORC_BEGIN_CATCH_NOT_OK
+    row_reader = reader_->createRowReader(opts);
+    row_reader->seekToRow(current_row_);
+    current_row_ = stripe_info.first_row_id + stripe_info.num_rows;
+    ORC_END_CATCH_NOT_OK
+
+    return std::make_shared<OrcStripeReader>(std::move(row_reader), schema, batch_size,
+                                             pool_);
+  }
+
   Result<std::shared_ptr<RecordBatchReader>> GetRecordBatchReader(
       int64_t batch_size, const std::vector<std::string>& include_names) {
     liborc::RowReaderOptions opts = DefaultRowReaderOptions();
@@ -746,6 +771,11 @@ Result<std::shared_ptr<RecordBatchReader>> ORCFileReader::NextStripeReader(
   return impl_->NextStripeReader(batch_size, include_indices);
 }
 
+Result<std::shared_ptr<RecordBatchReader>> ORCFileReader::NextStripeReader(
+    int64_t batch_size, const std::vector<std::string>& include_names) {
+  return impl_->NextStripeReader(batch_size, include_names);
+}
+
 int64_t ORCFileReader::NumberOfStripes() { return impl_->NumberOfStripes(); }
 
 int64_t ORCFileReader::NumberOfRows() { return impl_->NumberOfRows(); }

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

@@ -163,6 +163,22 @@ class ARROW_EXPORT ORCFileReader {
   Result<std::shared_ptr<RecordBatchReader>> NextStripeReader(
       int64_t batch_size, const std::vector<int>& include_indices);
 
+  /// \brief Get a stripe level record batch iterator.
+  ///
+  /// Each record batch will have up to `batch_size` rows.
+  /// NextStripeReader serves as a fine-grained alternative to ReadStripe
+  /// which may cause OOM issues by loading the whole stripe into memory.
+  ///
+  /// Note this will only read rows for the current stripe, not the entire
+  /// file.
+  ///
+  /// \param[in] batch_size the maximum number of rows in each record batch
+  /// \param[in] include_names the selected field names to read, if not empty
+  /// (otherwise all fields are read)
+  /// \return the stripe reader
+  Result<std::shared_ptr<RecordBatchReader>> NextStripeReader(
+      int64_t batch_size, const std::vector<std::string>& include_names);
+
   /// \brief Get a record batch iterator for the entire file.
   ///
   /// Each record batch will have up to `batch_size` rows.

cpp/src/arrow/dataset/file_orc.cc

@@ -31,6 +31,7 @@
 #include "arrow/util/logging.h"
 #include "arrow/util/string.h"
 #include "arrow/util/thread_pool.h"
+#include <limits>
 #include <numeric>
 
 namespace arrow {
@@ -87,6 +88,7 @@ class OrcScanTask {
                                               const FileFormat& format,
                                               const ScanOptions& scan_options,
                                               const std::shared_ptr<FileFragment>& fragment) {
+        ARROW_UNUSED(format);
         ARROW_ASSIGN_OR_RAISE(
             auto reader,
             OpenORCReader(source, std::make_shared<ScanOptions>(scan_options)));
@@ -123,23 +125,47 @@ class OrcScanTask {
           std::iota(stripe_indices.begin(), stripe_indices.end(), 0);
         }
 
-        // For this PR, we read all stripes but the infrastructure is in place
-        // A future PR can add GetRecordBatchReader overload with stripe_indices
-        std::shared_ptr<RecordBatchReader> record_batch_reader;
-        ARROW_ASSIGN_OR_RAISE(
-            record_batch_reader,
-            reader->GetRecordBatchReader(scan_options.batch_size, included_fields));
+        if (stripe_indices.empty()) {
+          return MakeEmptyIterator<std::shared_ptr<RecordBatch>>();
+        }
 
-        return RecordBatchIterator(Impl{std::move(record_batch_reader)});
+        return RecordBatchIterator(Impl{std::move(reader), std::move(included_fields),
+                                        std::move(stripe_indices), scan_options.batch_size,
+                                        0, nullptr});
       }
 
       Result<std::shared_ptr<RecordBatch>> Next() {
-        std::shared_ptr<RecordBatch> batch;
-        RETURN_NOT_OK(record_batch_reader_->ReadNext(&batch));
-        return batch;
+        while (true) {
+          if (current_stripe_reader_) {
+            std::shared_ptr<RecordBatch> batch;
+            RETURN_NOT_OK(current_stripe_reader_->ReadNext(&batch));
+            if (batch) {
+              return batch;
+            }
+            current_stripe_reader_.reset();
+            ++next_stripe_index_;
+            continue;
+          }
+
+          if (next_stripe_index_ >= stripe_indices_.size()) {
+            return IterationEnd<std::shared_ptr<RecordBatch>>();
+          }
+
+          const auto stripe = stripe_indices_[next_stripe_index_];
+          const auto stripe_info = reader_->GetStripeInformation(stripe);
+          RETURN_NOT_OK(reader_->Seek(stripe_info.first_row_id));
+          ARROW_ASSIGN_OR_RAISE(
+              current_stripe_reader_,
+              reader_->NextStripeReader(batch_size_, included_fields_));
+        }
       }
 
-      std::shared_ptr<RecordBatchReader> record_batch_reader_;
+      std::unique_ptr<adapters::orc::ORCFileReader> reader_;
+      std::vector<std::string> included_fields_;
+      std::vector<int> stripe_indices_;
+      int64_t batch_size_;
+      size_t next_stripe_index_;
+      std::shared_ptr<RecordBatchReader> current_stripe_reader_;
     };
 
     return Impl::Make(fragment_->source(),
@@ -319,10 +345,13 @@ Result<std::vector<compute::Expression>> OrcFileFragment::TestStripes(
 
   // Open reader if not already cached
   if (!cached_reader_) {
-    ARROW_ASSIGN_OR_RAISE(auto input,
-        arrow::io::RandomAccessFile::Open(source().path()));
-    ARROW_ASSIGN_OR_RAISE(cached_reader_,
-        adapters::orc::ORCFileReader::Open(input, arrow::default_memory_pool()));
+    auto lock = metadata_mutex_.Lock();
+    if (!cached_reader_) {
+      ARROW_ASSIGN_OR_RAISE(auto input, source().Open());
+      ARROW_ASSIGN_OR_RAISE(
+          cached_reader_,
+          adapters::orc::ORCFileReader::Open(std::move(input), arrow::default_memory_pool()));
+    }
   }
 
   // Process each field referenced in predicate (lazy evaluation)
@@ -349,8 +378,8 @@ Result<std::vector<compute::Expression>> OrcFileFragment::TestStripes(
     }
     statistics_expressions_complete_[field_index] = true;
 
-    // PR4 limitation: only support INT64
-    if (field->type()->id() != Type::INT64) {
+    // Support INT32 and INT64 types
+    if (field->type()->id() != Type::INT32 && field->type()->id() != Type::INT64) {
       continue;  // Unsupported type
     }
 
@@ -370,6 +399,16 @@ Result<std::vector<compute::Expression>> OrcFileFragment::TestStripes(
         continue;  // No statistics
       }
 
+      auto field_expr = compute::field_ref(field_ref);
+      const bool has_null = col_stats->hasNull();
+      const bool is_all_null = has_null && col_stats->getNumberOfValues() == 0;
+
+      if (is_all_null) {
+        // All values are null in this stripe for this column.
+        FoldingAnd(&statistics_expressions_[stripe_idx], compute::is_null(field_expr));
+        continue;
+      }
+
       const auto* int_stats =
           dynamic_cast<const liborc::IntegerColumnStatistics*>(col_stats);
       if (!int_stats || !int_stats->hasMinimum() || !int_stats->hasMaximum()) {
@@ -378,16 +417,28 @@ Result<std::vector<compute::Expression>> OrcFileFragment::TestStripes(
 
       int64_t min_value = int_stats->getMinimum();
       int64_t max_value = int_stats->getMaximum();
-      bool has_null = col_stats->hasNull();
 
       if (min_value > max_value) {
         continue;  // Invalid statistics
       }
 
-      // Build guarantee expression (from PR2 logic)
-      auto field_expr = compute::field_ref(field_ref);
-      auto min_scalar = std::make_shared<Int64Scalar>(min_value);
-      auto max_scalar = std::make_shared<Int64Scalar>(max_value);
+      // Build guarantee expression
+      std::shared_ptr<Scalar> min_scalar, max_scalar;
+
+      // Handle INT32 with overflow protection
+      if (field->type()->id() == Type::INT32) {
+        // Check for INT32 overflow
+        if (min_value < std::numeric_limits<int32_t>::min() ||
+            max_value > std::numeric_limits<int32_t>::max()) {
+          // Statistics overflow - skip predicate pushdown for safety
+          continue;
+        }
+        min_scalar = std::make_shared<Int32Scalar>(static_cast<int32_t>(min_value));
+        max_scalar = std::make_shared<Int32Scalar>(static_cast<int32_t>(max_value));
+      } else {
+        min_scalar = std::make_shared<Int64Scalar>(min_value);
+        max_scalar = std::make_shared<Int64Scalar>(max_value);
+      }
 
       auto min_expr = compute::greater_equal(field_expr, compute::literal(*min_scalar));
       auto max_expr = compute::less_equal(field_expr, compute::literal(*max_scalar));
@@ -421,6 +472,7 @@ Result<std::vector<compute::Expression>> OrcFileFragment::TestStripes(
 
 Result<std::vector<int>> OrcFileFragment::FilterStripes(
     const compute::Expression& predicate) {
+  RETURN_NOT_OK(EnsureMetadataCached());
 
   // Feature flag for disabling predicate pushdown
   if (auto env_var = arrow::internal::GetEnvVar("ARROW_ORC_DISABLE_PREDICATE_PUSHDOWN")) {

cpp/src/arrow/dataset/file_orc.h

@@ -82,7 +82,8 @@ class ARROW_DS_EXPORT OrcFileFragment : public FileFragment {
   /// \brief Filter stripes based on predicate using stripe statistics
   ///
   /// Returns indices of stripes where the predicate may be satisfied.
-  /// Currently supports INT64 columns with greater-than operator only.
+  /// Supports INT32/INT64 columns and conservative handling of missing or
+  /// unsupported statistics.
   ///
   /// \param predicate Arrow compute expression to evaluate
   /// \return Vector of stripe indices to read (0-based)