From e26e72a64154b64110f6723c6139459073f6aea4 Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Sun, 4 Oct 2026 13:36:44 +0800 Subject: [PATCH 1/2] feat(python): support integer column indices in read_url projection `read_url`'s docstring and type stub both document `projection` as `list[str | int]` ("by their index or name"), but the implementation only accepted strings and raised `TypeError` on an integer, so the documented positional form never worked. Map an integer to the top-level field name at the matching position (threading the file dtype into `projection_from_python`), erroring on an out-of-range index or a non-struct file. String projection is unchanged. Signed-off-by: jackylee-ch --- vortex-python/src/dataset.rs | 34 +++++++++++++++++++++++++------- vortex-python/test/test_store.py | 23 +++++++++++++++++++++ 2 files changed, 50 insertions(+), 7 deletions(-) diff --git a/vortex-python/src/dataset.rs b/vortex-python/src/dataset.rs index 696ad86a7a0..561331b91c9 100644 --- a/vortex-python/src/dataset.rs +++ b/vortex-python/src/dataset.rs @@ -6,15 +6,18 @@ use std::sync::Arc; use arrow_array::RecordBatchReader; use arrow_schema::SchemaRef; use itertools::Itertools; +use pyo3::exceptions::PyIndexError; use pyo3::exceptions::PyTypeError; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; +use pyo3::types::PyInt; use pyo3::types::PyString; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::VortexSessionExecute; use vortex::array::arrays::PrimitiveArray; use vortex::array::iter::ArrayIteratorExt; +use vortex::dtype::DType; use vortex::dtype::FieldName; use vortex::dtype::FieldNames; use vortex::error::VortexResult; @@ -83,13 +86,30 @@ pub fn read_array_from_reader( scan.into_array_iter(&runtime)?.read_all() } -fn projection_from_python(columns: Option>>) -> PyResult { - fn field_from_pyany(field: &Bound) -> PyResult { - if field.clone().is_instance_of::() { +fn projection_from_python( + columns: Option>>, + dtype: &DType, +) -> PyResult { + fn field_from_pyany(field: &Bound, dtype: &DType) -> PyResult { + if field.is_instance_of::() { Ok(FieldName::from(field.cast::()?.to_str()?)) + } else if field.is_instance_of::() { + // Positional projection: map the index onto the top-level field name. + let DType::Struct(struct_dtype, _) = dtype else { + return Err(PyTypeError::new_err( + "projection: integer indices are only valid for a struct-typed file", + )); + }; + let index = field.extract::()?; + struct_dtype.field_name(index).cloned().ok_or_else(|| { + PyIndexError::new_err(format!( + "projection: column index {index} is out of range for {} columns", + struct_dtype.nfields() + )) + }) } else { Err(PyTypeError::new_err(format!( - "projection: expected list of strings or None, but found: {field}.", + "projection: expected a list of strings or integers or None, but found: {field}.", ))) } } @@ -99,7 +119,7 @@ fn projection_from_python(columns: Option>>) -> PyResult { let fields: Vec<_> = columns .iter() - .map(field_from_pyany) + .map(|field| field_from_pyany(field, dtype)) .collect::>()?; select(FieldNames::from(fields), root()) } @@ -148,7 +168,7 @@ impl PyVortexDataset { row_range: Option<(u64, u64)>, ) -> PyVortexResult { let vxf = self.vxf.clone(); - let projection = projection_from_python(columns)?; + let projection = projection_from_python(columns, vxf.dtype())?; let filter = filter_from_python(row_filter); let indices = indices.map(|i| i.into_inner()); @@ -187,7 +207,7 @@ impl PyVortexDataset { row_range: Option<(u64, u64)>, ) -> PyVortexResult> { let vxf = self_.vxf.clone(); - let projection = projection_from_python(columns)?; + let projection = projection_from_python(columns, vxf.dtype())?; let filter = filter_from_python(row_filter); let reader = self_.py().detach(move || { diff --git a/vortex-python/test/test_store.py b/vortex-python/test/test_store.py index d3909e5a4da..783214fcdc7 100644 --- a/vortex-python/test/test_store.py +++ b/vortex-python/test/test_store.py @@ -40,3 +40,26 @@ def test_store_roundtrip(tmp_path: Path) -> None: people = vx.io.read_url("people.vortex", store=local) assert people.to_pylist() == records.to_pylist() + + +def test_read_url_integer_projection(tmp_path: Path) -> None: + local = LocalStore(prefix=tmp_path) + records = vx.array([dict(name="Alice", salary=10), dict(name="Bob", salary=20)]) + vx.io.write(records, "people.vortex", store=local) + + # Columns are name (0) and salary (1); select salary by position. + by_index = vx.io.read_url("people.vortex", store=local, projection=[1]) + assert by_index.to_pylist() == [{"salary": 10}, {"salary": 20}] + + # Integer and name projection agree. + by_name = vx.io.read_url("people.vortex", store=local, projection=["salary"]) + assert by_index.to_pylist() == by_name.to_pylist() + + +def test_read_url_integer_projection_out_of_range(tmp_path: Path) -> None: + local = LocalStore(prefix=tmp_path) + records = vx.array([dict(name="Alice", salary=10)]) + vx.io.write(records, "people.vortex", store=local) + + with pytest.raises(IndexError): + vx.io.read_url("people.vortex", store=local, projection=[99]) From 04636b5087d784f261ee53ba7aa86830bfd702ba Mon Sep 17 00:00:00 2001 From: jackylee-ch Date: Sun, 4 Oct 2026 23:37:06 +0800 Subject: [PATCH 2/2] refactor(python): parse projection columns with a FromPyObject enum Replace the manual isinstance branching in projection_from_python with a FromPyObject-derived `ProjectionColumn` enum, as suggested in review. Signed-off-by: jackylee-ch --- vortex-python/src/dataset.rs | 65 +++++++++++++++++------------------- vortex-python/src/io.rs | 3 +- 2 files changed, 33 insertions(+), 35 deletions(-) diff --git a/vortex-python/src/dataset.rs b/vortex-python/src/dataset.rs index 561331b91c9..9dd356d397e 100644 --- a/vortex-python/src/dataset.rs +++ b/vortex-python/src/dataset.rs @@ -10,8 +10,6 @@ use pyo3::exceptions::PyIndexError; use pyo3::exceptions::PyTypeError; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use pyo3::types::PyInt; -use pyo3::types::PyString; use vortex::array::ArrayRef; use vortex::array::ExecutionCtx; use vortex::array::VortexSessionExecute; @@ -86,41 +84,40 @@ pub fn read_array_from_reader( scan.into_array_iter(&runtime)?.read_all() } +/// A projected column, selected either by name or by positional index. +#[derive(FromPyObject)] +pub enum ProjectionColumn { + Name(String), + Index(usize), +} + fn projection_from_python( - columns: Option>>, + columns: Option>, dtype: &DType, ) -> PyResult { - fn field_from_pyany(field: &Bound, dtype: &DType) -> PyResult { - if field.is_instance_of::() { - Ok(FieldName::from(field.cast::()?.to_str()?)) - } else if field.is_instance_of::() { - // Positional projection: map the index onto the top-level field name. - let DType::Struct(struct_dtype, _) = dtype else { - return Err(PyTypeError::new_err( - "projection: integer indices are only valid for a struct-typed file", - )); - }; - let index = field.extract::()?; - struct_dtype.field_name(index).cloned().ok_or_else(|| { - PyIndexError::new_err(format!( - "projection: column index {index} is out of range for {} columns", - struct_dtype.nfields() - )) - }) - } else { - Err(PyTypeError::new_err(format!( - "projection: expected a list of strings or integers or None, but found: {field}.", - ))) - } - } - Ok(match columns { None => root(), Some(columns) => { - let fields: Vec<_> = columns - .iter() - .map(|field| field_from_pyany(field, dtype)) - .collect::>()?; + let fields = columns + .into_iter() + .map(|column| match column { + ProjectionColumn::Name(name) => Ok(FieldName::from(name.as_str())), + ProjectionColumn::Index(index) => { + // Positional projection: map the index onto the top-level field name. + let DType::Struct(struct_dtype, _) = dtype else { + return Err(PyTypeError::new_err( + "projection: integer indices are only valid for a struct-typed file", + )); + }; + struct_dtype.field_name(index).cloned().ok_or_else(|| { + PyIndexError::new_err(format!( + "projection: column index {index} is out of range for {} columns", + struct_dtype.nfields() + )) + }) + } + }) + .collect::>>()?; select(FieldNames::from(fields), root()) } }) @@ -162,7 +159,7 @@ impl PyVortexDataset { pub(crate) fn to_array_inner<'py>( &self, py: Python<'py>, - columns: Option>>, + columns: Option>, row_filter: Option<&Bound<'py, PyExpr>>, indices: Option, row_range: Option<(u64, u64)>, @@ -190,7 +187,7 @@ impl PyVortexDataset { #[pyo3(signature = (*, columns = None, row_filter = None, indices = None, row_range = None))] pub fn to_array<'py>( self_: PyRef<'py, Self>, - columns: Option>>, + columns: Option>, row_filter: Option<&Bound<'py, PyExpr>>, indices: Option, row_range: Option<(u64, u64)>, @@ -201,7 +198,7 @@ impl PyVortexDataset { #[pyo3(signature = (*, columns = None, row_filter = None, split_by = None, row_range = None))] pub fn to_record_batch_reader( self_: PyRef, - columns: Option>>, + columns: Option>, row_filter: Option<&Bound<'_, PyExpr>>, split_by: Option, row_range: Option<(u64, u64)>, diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 794903369ee..77ab43f8142 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -34,6 +34,7 @@ use crate::arrow::FromPyArrow; use crate::classes::record_batch_reader_class; use crate::classes::table_class; use crate::current_runtime; +use crate::dataset::ProjectionColumn; use crate::dataset::PyVortexDataset; use crate::error::PyVortexResult; use crate::expr::PyExpr; @@ -125,7 +126,7 @@ pub fn read_url<'py>( py: Python<'py>, url: &str, store: Option>, - projection: Option>>, + projection: Option>, row_filter: Option<&Bound<'py, PyExpr>>, indices: Option, row_range: Option<(u64, u64)>,