diff --git a/cpp/src/arrow/adapters/orc/adapter.cc b/cpp/src/arrow/adapters/orc/adapter.cc index 51cca497485c..e63c16a3385c 100644 --- a/cpp/src/arrow/adapters/orc/adapter.cc +++ b/cpp/src/arrow/adapters/orc/adapter.cc @@ -156,6 +156,12 @@ class OrcStripeReader : public RecordBatchReader { return Status::OK(); } + if (schema_->num_fields() == 0) { + *out = RecordBatch::Make(schema_, batch->numElements, + std::vector>{}); + return Status::OK(); + } + ARROW_ASSIGN_OR_RAISE(builder, RecordBatchBuilder::Make(schema_, pool_, batch->numElements)); // The top-level type must be a struct to read into an arrow table @@ -501,6 +507,31 @@ class ORCFileReader::Impl { return Status::OK(); } + Result> GetStripeReader(int64_t stripe, + int64_t batch_size) { + return GetStripeReader(stripe, batch_size, DefaultRowReaderOptions()); + } + + Result> GetStripeReader( + int64_t stripe, int64_t batch_size, const std::vector& include_names) { + auto opts = DefaultRowReaderOptions(); + RETURN_NOT_OK(SelectNames(&opts, include_names)); + return GetStripeReader(stripe, batch_size, std::move(opts)); + } + + Result> GetStripeReader( + int64_t stripe, int64_t batch_size, liborc::RowReaderOptions opts) { + if (batch_size <= 0) { + return Status::Invalid("batch_size must be greater than zero"); + } + RETURN_NOT_OK(SelectStripe(&opts, stripe)); + std::unique_ptr row_reader; + ORC_CATCH_NOT_OK(row_reader = reader_->createRowReader(opts)); + ARROW_ASSIGN_OR_RAISE(auto schema, GetArrowSchema(row_reader->getSelectedType())); + return std::make_shared(std::move(row_reader), std::move(schema), + batch_size, pool_); + } + Result> NextStripeReader( int64_t batch_size, const std::vector& include_indices) { if (current_row_ >= NumberOfRows()) { @@ -615,6 +646,16 @@ Result> ORCFileReader::ReadStripe( Status ORCFileReader::Seek(int64_t row_number) { return impl_->Seek(row_number); } +Result> ORCFileReader::GetStripeReader( + int64_t stripe, int64_t batch_size) { + return impl_->GetStripeReader(stripe, batch_size); +} + +Result> ORCFileReader::GetStripeReader( + int64_t stripe, int64_t batch_size, const std::vector& include_names) { + return impl_->GetStripeReader(stripe, batch_size, include_names); +} + Result> ORCFileReader::NextStripeReader( int64_t batch_size) { return impl_->NextStripeReader(batch_size); diff --git a/cpp/src/arrow/adapters/orc/adapter.h b/cpp/src/arrow/adapters/orc/adapter.h index 4ffff81f355f..17a278bbbae0 100644 --- a/cpp/src/arrow/adapters/orc/adapter.h +++ b/cpp/src/arrow/adapters/orc/adapter.h @@ -135,6 +135,25 @@ class ARROW_EXPORT ORCFileReader { /// \param[in] row_number the rows number to seek Status Seek(int64_t row_number); + /// \brief Read a designated stripe incrementally without changing the Seek position. + /// + /// \param[in] stripe the stripe index + /// \param[in] batch_size the maximum number of rows per batch, which must be positive + /// \return an independent reader for all fields in the stripe + Result> GetStripeReader(int64_t stripe, + int64_t batch_size); + + /// \brief Read selected fields of a designated stripe incrementally. + /// + /// This does not change the Seek position. Field order follows the file schema. + /// An empty selection produces zero-column batches with the stripe's row count. + /// \param[in] stripe the stripe index + /// \param[in] batch_size the maximum number of rows per batch, which must be positive + /// \param[in] include_names the selected field names, including nested field paths + /// \return an independent reader for the selected fields in the stripe + Result> GetStripeReader( + int64_t stripe, int64_t batch_size, const std::vector& include_names); + /// \brief Get a stripe level record batch iterator. /// /// Each record batch will have up to `batch_size` rows. diff --git a/cpp/src/arrow/adapters/orc/adapter_test.cc b/cpp/src/arrow/adapters/orc/adapter_test.cc index 3cbb6d7828f1..8ef53c84d613 100644 --- a/cpp/src/arrow/adapters/orc/adapter_test.cc +++ b/cpp/src/arrow/adapters/orc/adapter_test.cc @@ -18,6 +18,7 @@ #include #include +#include #include #include @@ -30,6 +31,8 @@ #include "arrow/compute/cast.h" #include "arrow/io/interfaces.h" #include "arrow/io/memory.h" +#include "arrow/io/test_common.h" +#include "arrow/memory_pool.h" #include "arrow/status.h" #include "arrow/table.h" #include "arrow/testing/gtest_util.h" @@ -43,6 +46,7 @@ namespace liborc = orc; namespace arrow { +using internal::checked_cast; using internal::checked_pointer_cast; constexpr size_t kDefaultSmallMemStreamSize = 16384 * 5; // 80KB @@ -456,6 +460,193 @@ TEST(TestAdapterRead, ReadIntAndStringFileMultipleStripes) { EXPECT_EQ(num_rows / reader_batch_size, batches); } +int64_t StripeBatchValue(int64_t row, int column) { + return static_cast( + (static_cast(row + 1) * (0x123456789abcdefULL + column)) & + 0x7fffffffffffffffULL); +} + +std::shared_ptr MakeStripeBatchFile( + int64_t rows_per_stripe, int stripes, + liborc::CompressionKind compression = liborc::CompressionKind_NONE) { + MemoryOutputStream stream(64 * 1024 * 1024); + auto type = liborc::Type::buildTypeFromString( + "struct"); + liborc::WriterOptions options; + options.setCompression(compression); + options.setRowIndexStride(0); + options.setStripeSize(stripes == 1 ? 256 * 1024 * 1024 : 1); + auto writer = liborc::createWriter(*type, &stream, options); + auto batch = writer->createRowBatch(16384); + auto& root = checked_cast(*batch); + const int64_t total = rows_per_stripe * stripes; + for (int64_t offset = 0; offset < total; offset += 16384) { + root.numElements = std::min(16384, total - offset); + for (int col = 0; col < 4; ++col) { + auto& values = checked_cast(*root.fields[col]); + values.numElements = root.numElements; + for (uint64_t row = 0; row < root.numElements; ++row) { + values.data[row] = StripeBatchValue(offset + row, col); + } + } + writer->add(root); + } + writer->close(); + return Buffer::FromString(std::string(stream.getData(), stream.getLength())); +} + +TEST(TestAdapterRead, StripeBatchSelectionAndState) { + auto file = std::make_shared(MakeStripeBatchFile(16384, 4)); + ASSERT_OK_AND_ASSIGN(auto reader, + adapters::orc::ORCFileReader::Open(file, default_memory_pool())); + ASSERT_EQ(4, reader->NumberOfStripes()); + ASSERT_OK(reader->Seek(16384 + 7)); + ASSERT_OK_AND_ASSIGN(auto first, reader->GetStripeReader(2, 1000)); + ASSERT_OK_AND_ASSIGN(auto second, reader->GetStripeReader(0, 1)); + std::shared_ptr batch; + ASSERT_OK(second->ReadNext(&batch)); + ASSERT_EQ(StripeBatchValue(0, 0), + checked_cast(*batch->column(0)).Value(0)); + ASSERT_OK_AND_ASSIGN(auto ignored, reader->ReadStripe(3)); + int64_t rows = 0; + while (true) { + ASSERT_OK(first->ReadNext(&batch)); + if (!batch) break; + ASSERT_OK(batch->ValidateFull()); + ASSERT_LE(batch->num_rows(), 1000); + for (int64_t row = 0; row < batch->num_rows(); ++row) { + ASSERT_EQ(StripeBatchValue(2 * 16384 + rows + row, 0), + checked_cast(*batch->column(0)).Value(row)); + } + rows += batch->num_rows(); + } + ASSERT_EQ(16384, rows); + ASSERT_OK_AND_ASSIGN(auto cursor_reader, reader->NextStripeReader(3)); + ASSERT_OK(cursor_reader->ReadNext(&batch)); + ASSERT_EQ(StripeBatchValue(16384 + 7, 0), + checked_cast(*batch->column(0)).Value(0)); + reader.reset(); + ASSERT_OK(second->ReadNext(&batch)); + ASSERT_EQ(StripeBatchValue(1, 0), + checked_cast(*batch->column(0)).Value(0)); +} + +TEST(TestAdapterRead, StripeBatchProjectionAndArguments) { + auto file = std::make_shared(MakeStripeBatchFile(16384, 3)); + ASSERT_OK_AND_ASSIGN(auto reader, + adapters::orc::ORCFileReader::Open(file, default_memory_pool())); + ASSERT_RAISES(Invalid, reader->GetStripeReader(-1, 10)); + ASSERT_RAISES(Invalid, reader->GetStripeReader(3, 10)); + ASSERT_RAISES(Invalid, reader->GetStripeReader(0, 0)); + ASSERT_RAISES(Invalid, reader->GetStripeReader(0, -1)); + ASSERT_RAISES(IOError, reader->GetStripeReader(0, 10, {"missing"})); + ASSERT_OK_AND_ASSIGN(auto projected, + reader->GetStripeReader(1, 10000, {"col2", "col0", "col0"})); + ASSERT_EQ(projected->schema()->field_names(), + (std::vector{"col0", "col2"})); + std::shared_ptr batch; + ASSERT_OK(projected->ReadNext(&batch)); + ASSERT_EQ(10000, batch->num_rows()); + ASSERT_EQ(StripeBatchValue(16384, 2), + checked_cast(*batch->column(1)).Value(0)); + ASSERT_OK_AND_ASSIGN(auto zero, + reader->GetStripeReader(2, 1000, std::vector{})); + int64_t rows = 0; + while (true) { + ASSERT_OK(zero->ReadNext(&batch)); + if (!batch) break; + ASSERT_OK(batch->ValidateFull()); + ASSERT_EQ(0, batch->num_columns()); + ASSERT_GT(batch->num_rows(), 0); + ASSERT_LE(batch->num_rows(), 1000); + rows += batch->num_rows(); + } + ASSERT_EQ(16384, rows); + ProxyMemoryPool zero_pool(default_memory_pool()); + ASSERT_OK_AND_ASSIGN(auto zero_reader, + adapters::orc::ORCFileReader::Open(file, &zero_pool)); + ASSERT_OK_AND_ASSIGN(auto zero_batches, + zero_reader->GetStripeReader(0, 1000, std::vector{})); + ASSERT_OK(zero_batches->ReadNext(&batch)); + ASSERT_EQ(1000, batch->num_rows()); + ASSERT_EQ(0, zero_pool.max_memory()); // No full-column decode followed by projection. +} + +TEST(TestAdapterRead, StripeBatchReadsOnlySelectedRange) { + for (auto codec : {liborc::CompressionKind_NONE, liborc::CompressionKind_ZLIB}) { + auto file = std::make_shared(MakeStripeBatchFile(16384, 4, codec)); + std::shared_ptr tracked = + io::TrackedRandomAccessFile::Make(file.get()); + ASSERT_OK_AND_ASSIGN( + auto reader, adapters::orc::ORCFileReader::Open(tracked, default_memory_pool())); + ASSERT_EQ(4, reader->NumberOfStripes()); + tracked->ResetStats(); + for (int stripe : {2, 0}) { + ASSERT_OK_AND_ASSIGN(auto batches, reader->GetStripeReader(stripe, 17)); + std::shared_ptr batch; + ASSERT_OK(batches->ReadNext(&batch)); + ASSERT_EQ(17, batch->num_rows()); + ASSERT_GT(tracked->bytes_read(), 0); + auto reads = tracked->num_reads(); + batches.reset(); // Early stop must not read the rest or another stripe. + ASSERT_EQ(reads, tracked->num_reads()); + } + for (const auto& range : tracked->get_read_ranges()) { + for (int stripe : {1, 3}) { + auto unselected = reader->GetStripeInformation(stripe); + ASSERT_TRUE(range.offset + range.length <= unselected.offset || + range.offset >= unselected.offset + unselected.length); + } + } + } +} + +TEST(TestAdapterRead, StripeBatchMemoryBounded) { + constexpr int64_t batch_size = 256; + int64_t small_peak = 0; + for (int64_t rows : {262144, 1048576}) { + auto buffer = MakeStripeBatchFile(rows, 1); + int64_t full_peak; + { + ProxyMemoryPool pool(default_memory_pool()); + auto file = std::make_shared(buffer); + ASSERT_OK_AND_ASSIGN(auto reader, adapters::orc::ORCFileReader::Open(file, &pool)); + ASSERT_EQ(1, reader->NumberOfStripes()); + ASSERT_EQ(rows, reader->GetStripeInformation(0).num_rows); + ASSERT_OK_AND_ASSIGN(auto batch, reader->ReadStripe(0)); + ASSERT_EQ(rows, batch->num_rows()); + full_peak = pool.max_memory(); + } + ProxyMemoryPool pool(default_memory_pool()); + auto file = std::make_shared(buffer); + ASSERT_OK_AND_ASSIGN(auto reader, adapters::orc::ORCFileReader::Open(file, &pool)); + ASSERT_EQ(0, pool.bytes_allocated()); + ASSERT_OK_AND_ASSIGN(auto batches, reader->GetStripeReader(0, batch_size)); + ASSERT_EQ(0, pool.bytes_allocated()); // No whole-stripe Arrow allocation. + int64_t count = 0; + int64_t first_peak = 0; + while (true) { + std::shared_ptr batch; + ASSERT_OK(batches->ReadNext(&batch)); + if (!batch) break; + ASSERT_LE(batch->num_rows(), batch_size); + if (count == 0) first_peak = pool.max_memory(); + ASSERT_EQ(StripeBatchValue(count, 0), + checked_cast(*batch->column(0)).Value(0)); + count += batch->num_rows(); + } + ASSERT_EQ(rows, count); + ASSERT_EQ(0, pool.bytes_allocated()); + ASSERT_LT(first_peak, full_peak / 4); + ASSERT_LT(pool.max_memory(), full_peak / 4); + if (small_peak == 0) small_peak = pool.max_memory(); + ASSERT_LE(pool.max_memory(), 2 * small_peak); + std::cout << "ORC_BATCH_MEMORY rows=" << rows << " batch_size=" << batch_size + << " full_peak=" << full_peak << " first_peak=" << first_peak + << " streaming_peak=" << pool.max_memory() << std::endl; + } +} + TEST(TestAdapterRead, ReadCharAndVarcharType) { MemoryOutputStream mem_stream(kDefaultMemStreamSize); auto orc_type = liborc::Type::buildTypeFromString("struct"); diff --git a/docs/source/python/orc.rst b/docs/source/python/orc.rst index 8f3ae75e1013..75f2d591a7ae 100644 --- a/docs/source/python/orc.rst +++ b/docs/source/python/orc.rst @@ -167,6 +167,39 @@ concatenate them into a single table. You can read individual stripes with two: ["foo","bar","baz"] three: [true,false,true] +To consume a file in batches without first materializing an entire stripe as +Arrow data, use :meth:`~pyarrow.orc.ORCFile.iter_batches`: + +.. code-block:: python + + >>> batches = orc_file.iter_batches(batch_size=2, stripes=[0], columns=['one']) + >>> [batch.num_rows for batch in batches] + [2, 1] + +Omitting ``stripes`` reads the entire file in physical order. An explicit +sequence reads stripes in the supplied order, including repetitions; an empty +sequence reads nothing. Column selection follows ``read_stripe``. In particular, +``columns=[]`` returns zero-column batches with the correct row counts. + +Reading begins when the iterator is consumed. When stopping early, call +``close()`` to release its reader: + +.. code-block:: python + + >>> batches = orc_file.iter_batches(batch_size=2) + >>> try: + ... first_batch = next(batches) + ... finally: + ... batches.close() + +Closing the iterator does not close a file object supplied by the caller. +Returned batches remain valid. ``batch_size`` bounds the number of output rows, +not total process memory: ORC dictionaries, encoded streams, decompression +buffers, and individual wide rows can still require additional memory. +The Dataset interface also supports whole-file streaming; this method additionally +allows callers to select individual stripes. It does not perform predicate +pushdown or choose stripes from statistics. + We can write an ORC file using ``ORCWriter``: .. code-block:: python diff --git a/python/pyarrow/_orc.pxd b/python/pyarrow/_orc.pxd index aecbba317aec..3a865b8b3bef 100644 --- a/python/pyarrow/_orc.pxd +++ b/python/pyarrow/_orc.pxd @@ -25,7 +25,7 @@ from pyarrow.includes.common cimport * from pyarrow.includes.libarrow cimport (CArray, CSchema, CStatus, CResult, CTable, CMemoryPool, CKeyValueMetadata, - CRecordBatch, + CRecordBatch, CRecordBatchReader, CTable, CCompressionType, CRandomAccessFile, COutputStream, TimeUnit) @@ -103,6 +103,13 @@ cdef extern from "arrow/adapters/orc/adapter.h" \ CResult[shared_ptr[CRecordBatch]] ReadStripe( int64_t stripe, std_vector[c_string]) + CResult[shared_ptr[CRecordBatchReader]] GetStripeReader( + int64_t stripe, int64_t batch_size) + CResult[shared_ptr[CRecordBatchReader]] GetStripeReader( + int64_t stripe, int64_t batch_size, const std_vector[c_string]&) + CResult[shared_ptr[CRecordBatchReader]] GetRecordBatchReader( + int64_t batch_size, const std_vector[c_string]&) + CResult[shared_ptr[CTable]] Read() CResult[shared_ptr[CTable]] Read(std_vector[c_string]) diff --git a/python/pyarrow/_orc.pyx b/python/pyarrow/_orc.pyx index 1dd6848122c2..e30f6e87f232 100644 --- a/python/pyarrow/_orc.pyx +++ b/python/pyarrow/_orc.pyx @@ -239,6 +239,7 @@ cdef class ORCReader(_Weakrefable): object source CMemoryPool* allocator unique_ptr[ORCFileReader] reader + shared_ptr[CRandomAccessFile] source_handle def __cinit__(self, MemoryPool memory_pool=None): self.allocator = maybe_unbox_memory_pool(memory_pool) @@ -250,6 +251,7 @@ cdef class ORCReader(_Weakrefable): self.source = source get_reader(source, use_memory_map, &rd_handle) + self.source_handle = rd_handle with nogil: self.reader = move(GetResultValue( ORCFileReader.Open(rd_handle, self.allocator) @@ -364,6 +366,57 @@ cdef class ORCReader(_Weakrefable): return pyarrow_wrap_batch(sp_record_batch) + def iter_batches(self, int64_t batch_size, stripes=None, columns=None): + cdef: + shared_ptr[CRecordBatchReader] batch_reader + shared_ptr[CRecordBatch] batch + std_vector[c_string] names + int64_t stripe + + if batch_size <= 0: + raise ValueError("batch_size must be greater than zero") + if columns is not None: + names = [tobytes(name) for name in columns] + + # Keep self (and its source) alive for the generator's entire lifetime. + if stripes is None and (columns is None or len(columns) != 0): + stripe_indices = [None] + elif stripes is None: + stripe_indices = range(self.nstripes()) + else: + stripe_indices = stripes + + try: + for index in stripe_indices: + if index is None: + with nogil: + batch_reader = GetResultValue( + self.reader.get().GetRecordBatchReader(batch_size, names)) + else: + stripe = index + if columns is None: + with nogil: + batch_reader = GetResultValue( + self.reader.get().GetStripeReader(stripe, batch_size)) + else: + with nogil: + batch_reader = GetResultValue( + self.reader.get().GetStripeReader( + stripe, batch_size, names)) + while True: + if self.source_handle.get().closed(): + raise ValueError("I/O operation on closed file") + with nogil: + check_status(batch_reader.get().ReadNext(&batch)) + if batch.get() == NULL: + break + yield pyarrow_wrap_batch(batch) + batch.reset() + batch_reader.reset() + finally: + batch.reset() + batch_reader.reset() + def read(self, columns=None): cdef: shared_ptr[CTable] sp_table diff --git a/python/pyarrow/orc.py b/python/pyarrow/orc.py index 8da0a2de797e..d08334aabbc7 100644 --- a/python/pyarrow/orc.py +++ b/python/pyarrow/orc.py @@ -17,6 +17,7 @@ from numbers import Integral +import operator import warnings from pyarrow.lib import Table @@ -167,6 +168,56 @@ def read_stripe(self, n, columns=None): columns = self._select_names(columns) return self.reader.read_stripe(n, columns=columns) + def iter_batches(self, batch_size=65536, stripes=None, columns=None): + """Read streaming batches from the file or selected stripes. + + Parameters + ---------- + batch_size : int, default 65536 + Maximum number of rows per batch. Must be positive. + stripes : sequence of int, optional + Zero-based stripe indices, read in the supplied order. Repeated + indices repeat the corresponding rows. None reads the entire file; + an empty sequence produces no batches. + columns : list, optional + Column names, nested field paths, or top-level integer indices, + as in read_stripe. Output follows the file's column order. None + reads all columns; an empty list reads no columns but preserves + the number of rows. + + Yields + ------ + pyarrow.RecordBatch + A batch containing at most batch_size rows. Stripe boundaries may + produce smaller batches. + + Notes + ----- + Reading starts when the iterator is consumed. Call its close() method + to release the reader when stopping early. This does not close a source + supplied by the caller. Returned batches remain valid after closing. + + Batches are decoded incrementally, without first constructing an Arrow + batch for the entire stripe. batch_size limits rows, not bytes or total + process memory: ORC dictionaries and decompression buffers may still + depend on the stripe size. + """ + batch_size = operator.index(batch_size) + if batch_size > 2**63 - 1 or batch_size < -2**63: + raise OverflowError("batch_size does not fit in int64") + if batch_size <= 0: + raise ValueError("batch_size must be greater than zero") + if stripes is not None: + stripes = tuple(operator.index(stripe) for stripe in stripes) + for stripe in stripes: + if stripe > 2**63 - 1 or stripe < -2**63: + raise OverflowError("stripe index does not fit in int64") + if not 0 <= stripe < self.nstripes: + raise ValueError(f"Out of bounds stripe: {stripe}") + if columns is not None: + columns = list(self._select_names(list(columns))) + return self.reader.iter_batches(batch_size, stripes, columns) + def read(self, columns=None): """Read the whole file. diff --git a/python/pyarrow/tests/test_orc.py b/python/pyarrow/tests/test_orc.py index 27154a6f34f3..cf33ec6f0e19 100644 --- a/python/pyarrow/tests/test_orc.py +++ b/python/pyarrow/tests/test_orc.py @@ -17,10 +17,14 @@ import decimal import datetime +import gc +import io from pathlib import Path +import random import shutil import subprocess import sys +import weakref import pytest @@ -126,6 +130,14 @@ def check_example_file(orc_path, expected_df, need_fix=False): json_pos += len(batch) assert json_pos == orc_file.nrows + # Exercise both the full-file and selected-stripe native batch paths on + # examples covering the existing reader's supported ORC types. + for stripes in (None, list(range(orc_file.nstripes))): + batches = list(orc_file.iter_batches(batch_size=127, stripes=stripes)) + assert all(0 < batch.num_rows <= 127 for batch in batches) + actual = pa.Table.from_batches(batches, schema=table.schema) + assert actual.equals(table) + @pytest.mark.pandas @pytest.mark.parametrize('filename', [ @@ -697,3 +709,273 @@ def test_orc_writer_with_null_arrays(tempdir): table = pa.table({"int64": a, "utf8": b}) with pytest.raises(pa.ArrowNotImplementedError): orc.write_table(table, path) + + +@pytest.fixture(scope="module", params=[ + 'uncompressed', 'zlib', 'snappy', 'lz4', 'zstd']) +def batch_file(request): + from pyarrow import orc + + rows = 4 * 16384 + rng = random.Random(42) + values = [rng.getrandbits(62) for _ in range(rows)] + table = pa.table({'value': values, + 'other': [value // 3 for value in values]}) + sink = pa.BufferOutputStream() + orc.write_table(table, sink, stripe_size=1, batch_size=16384, + compression=request.param) + data = sink.getvalue() + reader = orc.ORCFile(data) + assert reader.nstripes == 4 + assert all(reader.read_stripe(i).num_rows == 16384 for i in range(4)) + return data, table + + +@pytest.mark.parametrize('batch_size', [1000, 16384, 20000, 65536]) +@pytest.mark.parametrize('stripes', [None, [2, 0, 2], []]) +def test_iter_batches_selection(batch_file, batch_size, stripes): + from pyarrow import orc + + data, table = batch_file + reader = orc.ORCFile(data) + batches = list(reader.iter_batches(batch_size, stripes=stripes)) + assert all(0 < batch.num_rows <= batch_size for batch in batches) + if stripes is None: + expected = table + else: + expected = pa.Table.from_batches( + [reader.read_stripe(i) for i in stripes], schema=table.schema) + actual = pa.Table.from_batches(batches, schema=table.schema) + assert actual.equals(expected) + for batch in batches: + batch.validate(full=True) + + +@pytest.mark.parametrize('stripes', [None, [2, 0]]) +@pytest.mark.parametrize('columns', [None, [], [1], ['other', 'value'], + ['other', 'other']]) +def test_iter_batches_projection(batch_file, stripes, columns): + from pyarrow import orc + + data, _ = batch_file + reader = orc.ORCFile(data) + batches = list(reader.iter_batches(1000, stripes, columns)) + expected_rows = reader.nrows if stripes is None else 2 * 16384 + assert sum(batch.num_rows for batch in batches) == expected_rows + assert all(0 < batch.num_rows <= 1000 for batch in batches) + names = reader._select_names(columns) + if columns == []: + assert all(batch.num_columns == 0 for batch in batches) + else: + expected = pa.Table.from_batches([ + reader.read_stripe(i, columns=names) + for i in (range(reader.nstripes) if stripes is None else stripes)]) + assert pa.Table.from_batches(batches).equals(expected) + + +def test_iter_batches_nested_and_nulls(): + from pyarrow import orc + + table = pa.table({ + 'struct': [{'middle': {'inner': 1}, 'other': 2}, None, + {'middle': {'inner': None}, 'other': 4}], + 'list-struct': [[{'inner1': 1, 'inner2': 2}], None, []], + 'map': pa.array([[('a', 1)], None, []], pa.map_(pa.string(), pa.int32())), + 'text': ['hello', None, 'world'], + 'binary': [b'abc', None, b'def'], + }) + sink = pa.BufferOutputStream() + orc.write_table(table, sink) + reader = orc.ORCFile(sink.getvalue()) + for columns in (None, ['struct.middle.inner'], ['list-struct.inner1'], + ['map'], [0, 2]): + expected = reader.read(columns=columns) + for stripes in (None, [0]): + batches = list(reader.iter_batches(1, stripes, columns)) + assert [batch.num_rows for batch in batches] == [1, 1, 1] + assert pa.Table.from_batches(batches).equals(expected) + assert sum(b.num_rows for b in reader.iter_batches()) == 3 + + +@pytest.mark.parametrize('batch_size,error', [ + (0, ValueError), (-1, ValueError), (1.5, TypeError), ('2', TypeError), + (None, TypeError), (2**63, OverflowError), (-2**63 - 1, OverflowError)]) +def test_iter_batches_invalid_batch_size(batch_file, batch_size, error): + from pyarrow import orc + + reader = orc.ORCFile(batch_file[0]) + with pytest.raises(error): + reader.iter_batches(batch_size) + + +@pytest.mark.parametrize('stripes,error', [ + ([-1], ValueError), ([4], ValueError), ([0, 4], ValueError), + ([1.2], TypeError), (['0'], TypeError), ([None], TypeError), + ([2**63], OverflowError), ([-2**63 - 1], OverflowError)]) +def test_iter_batches_invalid_stripes(batch_file, stripes, error): + from pyarrow import orc + + source = BatchTrackingFile(batch_file[0].to_pybytes()) + reader = orc.ORCFile(source) + source.reads.clear() + with pytest.raises(error): + reader.iter_batches(stripes=stripes) + assert source.reads == [] + + +def test_iter_batches_snapshot_and_state(batch_file): + from pyarrow import orc + + reader = orc.ORCFile(batch_file[0]) + stripes = [2, 0] + columns = ['value'] + first = reader.iter_batches(1000, stripes, columns) + stripes[:] = [1] + columns[:] = ['other'] + second = reader.iter_batches(7, [1]) + retained = next(first) + assert retained.equals(reader.read_stripe(2, ['value']).slice(0, 1000)) + assert next(second).equals(reader.read_stripe(1).slice(0, 7)) + reader.read() + assert next(first).equals(reader.read_stripe(2, ['value']).slice(1000, 1000)) + first.close() + second.close() + assert retained.equals(reader.read_stripe(2, ['value']).slice(0, 1000)) + with pytest.raises(ValueError): + reader.iter_batches(columns=[2]) + with pytest.raises(IOError): + next(reader.iter_batches(columns=['missing'])) + assert list(reader.iter_batches(stripes=[], columns=['missing'])) == [] + + +@pytest.mark.numpy +def test_iter_batches_numpy_integer_arguments(): + from pyarrow import orc + import numpy as np + + table = pa.table({'value': [0, 1, 2, 3]}) + sink = pa.BufferOutputStream() + orc.write_table(table, sink) + reader = orc.ORCFile(sink.getvalue()) + batches = list(reader.iter_batches(np.int64(2), [np.int64(0)])) + assert [batch.num_rows for batch in batches] == [2, 2] + assert pa.Table.from_batches(batches).equals(table) + + +class BatchTrackingFile(io.BytesIO): + def __init__(self, data): + super().__init__(data) + self.reads = [] + self.fail = False + + def read(self, size=-1): + if self.fail: + raise OSError('injected ORC read failure') + offset = self.tell() + data = super().read(size) + self.reads.append((offset, len(data))) + return data + + def readinto(self, buffer): + data = self.read(len(buffer)) + buffer[:len(data)] = data + return len(data) + + +@pytest.mark.parametrize('stripes', [None, [2, 0]], ids=['all', 'selected']) +def test_iter_batches_lazy_close_and_owner(batch_file, monkeypatch, stripes): + from pyarrow import orc + + source = BatchTrackingFile(batch_file[0].to_pybytes()) + file = orc.ORCFile(source) + source.reads.clear() + reader_ref = weakref.ref(file.reader) + source_ref = weakref.ref(source) + + for unstarted_stripes in (None, [2, 0], []): + unstarted = file.iter_batches(32, unstarted_stripes) + unstarted.close() + assert list(unstarted) == [] + assert source.reads == [] + assert not source.closed + del unstarted + + def fail_whole_stripe(*args, **kwargs): + pytest.fail('iter_batches must not call read_stripe') + + monkeypatch.setattr(orc.ORCFile, 'read_stripe', fail_whole_stripe) + iterator = file.iter_batches(32, stripes) + assert source.reads == [] + del file + gc.collect() + assert reader_ref() is not None + retained = next(iterator) + start = 0 if stripes is None else stripes[0] * 16384 + expected = batch_file[1].slice(start, 32) + assert pa.Table.from_batches([retained]).equals(expected) + assert source.reads + reads = list(source.reads) + iterator.close() + assert source.reads == reads + assert not source.closed + assert list(iterator) == [] + del iterator, source + gc.collect() + assert reader_ref() is None + assert source_ref() is None + assert pa.Table.from_batches([retained]).equals(expected) + + +def test_iter_batches_errors_after_output(batch_file): + from pyarrow import orc + + source = BatchTrackingFile(batch_file[0].to_pybytes()) + reader = orc.ORCFile(source) + iterator = reader.iter_batches(16384, [0, 2]) + first = next(iterator) + assert first.num_rows == 16384 + source.fail = True + with pytest.raises((OSError, pa.ArrowException), match='injected ORC read failure'): + next(iterator) + assert list(iterator) == [] + assert pa.Table.from_batches([first]).equals(batch_file[1].slice(0, 16384)) + source.fail = False + source.close() + with pytest.raises((OSError, ValueError, pa.ArrowException)): + next(reader.iter_batches(stripes=[2])) + + +@pytest.mark.parametrize('columns', [None, []]) +def test_iter_batches_closed_source_with_buffered_rows(batch_file, columns): + from pyarrow import orc + + source = BatchTrackingFile(batch_file[0].to_pybytes()) + reader = orc.ORCFile(source) + iterator = reader.iter_batches(1, [0], columns) + first = next(iterator) + source.close() + with pytest.raises(ValueError, match='closed file'): + next(iterator) + assert first.num_rows == 1 + assert list(iterator) == [] + + +def test_iter_batches_sources_and_empty(tmp_path, datadir): + from pyarrow import orc + + path = tmp_path / 'batches.orc' + table = pa.table({'a': [1, None, 3]}) + orc.write_table(table, path) + with path.open('rb') as file, pa.OSFile(str(path)) as native: + for source in (path, str(path), file, native, + pa.BufferReader(path.read_bytes())): + batches = list(orc.ORCFile(source).iter_batches(2, [0])) + assert pa.Table.from_batches(batches).equals(table) + empty = orc.ORCFile(datadir / 'TestOrcFile.emptyFile.orc') + assert list(empty.iter_batches()) == [] + assert list(empty.iter_batches(columns=[])) == [] + assert list(empty.iter_batches(stripes=[])) == [] + with pytest.raises(ValueError): + empty.iter_batches(stripes=[0]) + with pytest.raises((OSError, pa.ArrowException)): + orc.ORCFile(pa.BufferReader(b'ORC broken file')).iter_batches()