PyVortex to use new scan executor (#2908)

Author: gatesn

pyvortex/pyproject.toml

@@ -8,7 +8,7 @@ dependencies = ["pyarrow>=17.0.0", "substrait>=0.23.0"]
 requires-python = ">= 3.10"
 
 [project.optional-dependencies]
-polars = ["polars>=1.24.0"]
+polars = ["polars>=1.27.0"]
 pandas = ["pandas>=2.2.0"]
 numpy = ["numpy>=1.26.0"]
 duckdb = ["duckdb>=1.1.2"]

pyvortex/python/vortex/file.py

@@ -30,11 +30,7 @@ def _io_source(
         if predicate is not None:
             predicate = polars_to_vortex(predicate)
 
-        reader = self.to_arrow(
-            projection=with_columns,
-            expr=predicate,
-            batch_size=batch_size or 8192,
-        )
+        reader = self.to_arrow(projection=with_columns, expr=predicate)
 
         for batch in reader:
             batch = pl.DataFrame._from_arrow(batch, rechunk=False)
@@ -46,7 +42,7 @@ def _io_source(
             data=pa.RecordBatch.from_arrays(
                 [pa.array([], type=field.type) for field in reader.schema],
                 schema=reader.schema,
-            )
+            ),
         )
 
     return register_io_source(_io_source, schema=schema)

pyvortex/python/vortex/polars_.py

@@ -25,7 +25,13 @@ def polars_to_vortex(expr: pl.Expr) -> ve.Expr:
     "LogicalOr": operator.or_,
 }
 
+
+def _unsupported(v, name: str):
+    raise ValueError(f"Unsupported Polars expression {name}: {v}")
+
+
 _LITERAL_TYPES = {
+    "Boolean": lambda v: vx.bool_(nullable=v is None),
     "Int": lambda v: vx.int_(64, nullable=v is None),
     "Int8": lambda v: vx.int_(8, nullable=v is None),
     "Int16": lambda v: vx.int_(16, nullable=v is None),
@@ -37,7 +43,6 @@ def polars_to_vortex(expr: pl.Expr) -> ve.Expr:
     "UInt64": lambda v: vx.uint(64, nullable=v is None),
     "Float32": lambda v: vx.float_(32, nullable=v is None),
     "Float64": lambda v: vx.float_(64, nullable=v is None),
-    "Boolean": lambda v: vx.bool_(nullable=v is None),
     "Null": lambda v: vx.null(),
     "String": lambda v: vx.utf8(nullable=v is None),
     "Binary": lambda v: vx.binary(nullable=v is None),
@@ -59,11 +64,28 @@ def _polars_to_vortex(expr: dict) -> ve.Expr:
     if "Column" in expr:
         return ve.column(expr["Column"])
 
+    # See https://github.com/pola-rs/polars/pull/21849)
+    if "Scalar" in expr:
+        dtype = expr["Scalar"]["dtype"]  # DType
+        value = expr["Scalar"]["value"]  # AnyValue
+
+        if "Null" in value:
+            value = None
+        elif "StringOwned" in value:
+            value = value["StringOwned"]
+        else:
+            raise ValueError(f"Unsupported Polars scalar value type {value}")
+
+        return ve.literal(_LITERAL_TYPES[dtype](value), value)
+
     if "Literal" in expr:
         expr = expr["Literal"]
 
         literal_type = next(iter(expr.keys()), None)
 
+        if literal_type == "Scalar":
+            return _polars_to_vortex(expr)
+
         # Special-case Series
         if literal_type == "Series":
             expr = pl.Expr.from_json(json.dumps({"Literal": expr}))
@@ -91,6 +113,12 @@ def _polars_to_vortex(expr: dict) -> ve.Expr:
             dtype = vx.ext("vortex.timestamp", vx.int_(64, nullable=value is None), metadata=metadata)
             return ve.literal(dtype, value)
 
+        # Unwrap 'Dyn' scalars, whose type hasn't been established yet.
+        # (post https://github.com/pola-rs/polars/pull/21849)
+        if literal_type == "Dyn":
+            expr = expr["Dyn"]
+            literal_type = next(iter(expr.keys()), None)
+
         if literal_type not in _LITERAL_TYPES:
             raise NotImplementedError(f"Unsupported Polars literal type: {literal_type}")
         value = expr[literal_type]
@@ -111,6 +139,11 @@ def _polars_to_vortex(expr: dict) -> ve.Expr:
 
                 # Vortex doesn't support is-in, so we need to construct a series of ORs?
 
+        if "StringExpr" in fn:
+            fn = fn["StringExpr"]
+            if "Contains" in fn:
+                raise ValueError("Unsupported Polars StringExpr.Contains")
+
         raise NotImplementedError(f"Unsupported Polars function: {fn}")
 
     raise NotImplementedError(f"Unsupported Polars expression: {expr}")

pyvortex/src/dataset.rs

@@ -49,7 +49,7 @@ pub async fn read_array_from_reader(
         scan = scan.with_row_indices(indices);
     }
 
-    let stream = scan.into_array_stream()?;
+    let stream = scan.spawn_tokio(TOKIO_RUNTIME.handle().clone())?;
     let dtype = stream.dtype().clone();
 
     let all_arrays = stream.try_collect::<Vec<_>>().await?;

pyvortex/src/file.rs

@@ -2,21 +2,20 @@ use std::sync::Arc;
 
 use arrow::array::RecordBatchReader;
 use arrow::pyarrow::IntoPyArrow;
-use futures::{SinkExt, StreamExt};
 use pyo3::exceptions::PyTypeError;
 use pyo3::prelude::*;
 use pyo3::types::PyList;
 use vortex::ToCanonical;
 use vortex::compute::try_cast;
 use vortex::dtype::Nullability::NonNullable;
 use vortex::dtype::{DType, PType};
-use vortex::error::{VortexExpect, vortex_err};
+use vortex::error::VortexError;
 use vortex::expr::{ExprRef, ident, select};
 use vortex::file::scan::SplitBy;
 use vortex::file::segments::MokaSegmentCache;
 use vortex::file::{VortexFile, VortexOpenOptions};
 use vortex::io::TokioFile;
-use vortex::stream::{ArrayStream, ArrayStreamAdapter, ArrayStreamExt};
+use vortex::stream::ArrayStreamExt;
 
 use crate::arrays::PyArrayRef;
 use crate::dataset::PyVortexDataset;
@@ -184,37 +183,23 @@ impl PyVortexFile {
         expr: Option<PyExpr>,
         batch_size: Option<usize>,
     ) -> PyResult<PyObject> {
-        let mut builder = slf
-            .get()
-            .vxf
-            .scan()?
-            .with_canonicalize(true)
-            .with_some_filter(expr.map(|e| e.into_inner()))
-            .with_projection(projection.map(|p| p.0).unwrap_or_else(ident));
-
-        if let Some(batch_size) = batch_size {
-            builder = builder.with_split_by(SplitBy::RowCount(batch_size));
-        }
-
-        let stream = ArrayStreamExt::boxed(builder.spawn_tokio(TOKIO_RUNTIME.handle().clone())?);
-        let dtype = stream.dtype().clone();
-
-        // The I/O of the array stream won't make progress unless it's polled. So we need to spawn it.
-        let (mut send, recv) = futures::channel::mpsc::unbounded();
-
-        TOKIO_RUNTIME
-            .block_on(TOKIO_RUNTIME.spawn(async move {
-                let mut stream = stream;
-                while let Some(batch) = stream.next().await {
-                    send.send(batch)
-                        .await
-                        .map_err(|e| vortex_err!("Send failed {}", e))
-                        .vortex_expect("send failed");
-                }
-            }))
-            .vortex_expect("failed to spawn stream");
-
-        let stream = ArrayStreamAdapter::new(dtype, recv);
+        let vxf = slf.get().vxf.clone();
+
+        let stream = slf.py().allow_threads(|| {
+            let mut builder = vxf
+                .scan()?
+                .with_canonicalize(true)
+                .with_some_filter(expr.map(|e| e.into_inner()))
+                .with_projection(projection.map(|p| p.0).unwrap_or_else(ident));
+
+            if let Some(batch_size) = batch_size {
+                builder = builder.with_split_by(SplitBy::RowCount(batch_size));
+            }
+
+            Ok::<_, VortexError>(ArrayStreamExt::boxed(
+                builder.spawn_tokio(TOKIO_RUNTIME.handle().clone())?,
+            ))
+        })?;
 
         let iter = ArrayStreamToIterator::new(stream);
         let rbr: Box<dyn RecordBatchReader + Send> =