diff --git a/vortex-python/python/vortex/_lib/file.pyi b/vortex-python/python/vortex/_lib/file.pyi index 901e0323eec..fd15eafa44a 100644 --- a/vortex-python/python/vortex/_lib/file.pyi +++ b/vortex-python/python/vortex/_lib/file.pyi @@ -51,6 +51,7 @@ class VortexFile: *, expr: Expr | None = None, limit: int | None = None, + indices: Array | None = None, batch_size: int | None = None, schema: pa.Schema | None = None, ) -> pa.RecordBatchReader: ... diff --git a/vortex-python/python/vortex/file.py b/vortex-python/python/vortex/file.py index a78a59ead1c..050f64e53b6 100644 --- a/vortex-python/python/vortex/file.py +++ b/vortex-python/python/vortex/file.py @@ -208,6 +208,7 @@ def to_arrow( *, limit: int | None = None, expr: Expr | None = None, + indices: Array | None = None, batch_size: int | None = None, schema: pa.Schema | None = None, ) -> RecordBatchReader: @@ -220,6 +221,8 @@ def to_arrow( from the file) or an explicit list of desired columns. expr : :class:`vortex.Expr` | None The predicate used to filter rows. The filter columns need not appear in the projection. + indices : :class:`vortex.Array` | None + The indices of the rows to read. Must be strictly increasing and non-null. batch_size : :class:`int` | None The number of rows to read per chunk. schema : :class:`pyarrow.Schema` | None @@ -227,7 +230,9 @@ def to_arrow( Use ``pyarrow.binary()`` for ``BinaryArray`` fields. """ - return self._file.to_arrow(projection, expr=expr, limit=limit, batch_size=batch_size, schema=schema) + return self._file.to_arrow( + projection, expr=expr, limit=limit, indices=indices, batch_size=batch_size, schema=schema + ) def to_dataset(self) -> VortexDataset: """Scan the Vortex file using the :class:`pyarrow.dataset.Dataset` API.""" diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index 3912f840d8b..4b362cd37e8 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -201,16 +201,20 @@ impl PyVortexFile { }) } - #[pyo3(signature = (projection = None, *, expr = None, limit = None, batch_size = None, schema = None))] + #[pyo3(signature = (projection = None, *, expr = None, limit = None, indices = None, batch_size = None, schema = None))] fn to_arrow( slf: Bound, projection: Option, expr: Option, limit: Option, + indices: Option, batch_size: Option, schema: Option<&Bound>, ) -> PyVortexResult> { let vxf = &slf.get().vxf; + let projection = projection.map(|p| p.0); + let expr = expr.map(|e| e.into_inner()); + let indices = row_indices(slf.py(), indices)?; let schema = schema .map(|schema| Schema::from_pyarrow(&schema.as_borrowed())) .transpose()? @@ -218,26 +222,7 @@ impl PyVortexFile { // Building the reader is lazy and cheap, so it runs without releasing the GIL. The scan // runs as pyarrow pulls batches, and pyarrow releases the GIL while it does. - let filter = expr - .map(|e| e.into_inner().bind(vxf.dtype())?.optimize_recursive()) - .transpose()?; - let projection = projection - .map(|p| p.0) - .unwrap_or_else(root) - .bind(vxf.dtype())? - .optimize_recursive()?; - let mut builder = vxf - .scan()? - .with_some_filter(filter) - .with_projection(projection); - - if let Some(limit) = limit { - builder = builder.with_limit(limit); - } - - if let Some(batch_size) = batch_size { - builder = builder.with_split_by(SplitBy::RowCount(batch_size)); - } + let builder = scan_builder(vxf, projection, expr, limit, indices, batch_size)?; let schema = match schema { Some(schema) => schema, diff --git a/vortex-python/test/test_file_arrow_indices.py b/vortex-python/test/test_file_arrow_indices.py new file mode 100644 index 00000000000..51937682607 --- /dev/null +++ b/vortex-python/test/test_file_arrow_indices.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +from pathlib import Path + +import pyarrow as pa +import pytest + +import vortex as vx +import vortex.expr as ve + + +@pytest.fixture +def indexed_file(tmp_path: Path) -> tuple[vx.VortexFile, pa.Table]: + table = pa.table( + { + "id": list(range(9)), + "payload": ["first", None, "long string with 数据", "", "four", "five", None, "seven", "last"], + } + ) + path = tmp_path / "indexed.vortex" + vx.io.write(vx.compress(vx.array(table)), str(path)) + return vx.open(str(path)), table + + +@pytest.mark.parametrize("indices", [[], [1], [0, 2, 5, 8]]) +@pytest.mark.parametrize("batch_size", [1, 3]) +def test_arrow_indices_preserve_values_and_requested_schema(indexed_file, indices, batch_size): + file, table = indexed_file + indices = pa.array(indices, type=pa.uint64()) + schema = pa.schema([("payload", pa.string())]) + reader = file.to_arrow(["payload"], indices=vx.array(indices), batch_size=batch_size, schema=schema) + batches = list(reader) + scan_batches = list(file.scan(["payload"], indices=vx.array(indices), batch_size=batch_size)) + assert [batch.num_rows for batch in batches] == [len(batch) for batch in scan_batches] + result = pa.Table.from_batches(batches, schema=reader.schema) + assert result.equals(table.select(["payload"]).take(indices)) + + +def test_arrow_indices_apply_filter(indexed_file): + file, table = indexed_file + reader = file.to_arrow( + ["payload"], + indices=vx.array([0, 2, 5, 8]), + expr=ve.column("id") >= 4, + schema=pa.schema([("payload", pa.string())]), + ) + assert reader.read_all().equals(table.select(["payload"]).take(pa.array([5, 8]))) + + +def test_arrow_indices_apply_limit(indexed_file): + file, table = indexed_file + reader = file.to_arrow( + ["payload"], + indices=vx.array([0, 2, 5, 8]), + limit=1, + schema=pa.schema([("payload", pa.string())]), + ) + assert reader.read_all().equals(table.select(["payload"]).slice(0, 1)) + + +def test_arrow_indices_use_scan_filter_limit_validation(indexed_file): + file, _ = indexed_file + reader = file.to_arrow(indices=vx.array([0, 2, 5, 8]), expr=ve.column("id") >= 4, limit=1) + with pytest.raises(pa.ArrowInvalid, match="doesn't support scans with both a filter and a limit"): + reader.read_all() + + +@pytest.mark.parametrize("indices", [[2, 1], [1, 1], [None, 1]]) +def test_arrow_indices_use_scan_validation(indexed_file, indices): + file, _ = indexed_file + indices = vx.array(pa.array(indices, type=pa.uint64())) + with pytest.raises(RuntimeError): + file.to_arrow(indices=indices)