Column Projection Pushdown (#119)

Only access/pivot the columns from the Xarray dataset that are actually
used by the query.

For a `SELECT temperature FROM ds` query on a 250-variable ERA5 dataset,
this reduces per-partition data loading from ~38 GB to ~1.5 GB (the
temperature array plus coordinate columns).

Fixes #103, Fixes #138.

---------

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: alxmrs

src/lib.rs

@@ -50,7 +50,7 @@ use std::fmt::Debug;
 use std::sync::Arc;
 
 use arrow::array::RecordBatch;
-use arrow::datatypes::SchemaRef;
+use arrow::datatypes::{Schema, SchemaRef};
 use arrow::pyarrow::FromPyArrow;
 use async_stream::try_stream;
 use async_trait::async_trait;
@@ -68,7 +68,7 @@ use datafusion::physical_plan::streaming::PartitionStream;
 use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream};
 use datafusion_ffi::table_provider::FFI_TableProvider;
 use pyo3::prelude::*;
-use pyo3::types::PyCapsule;
+use pyo3::types::{PyCapsule, PyList};
 use tokio::runtime::Handle;
 
 // ============================================================================
@@ -157,16 +157,18 @@ impl PartitionMetadata {
 /// `TableProvider::supports_filters_pushdown` and partition pruning in `scan()`.
 struct PrunableStreamingTable {
     schema: SchemaRef,
-    /// Partition streams paired with their coordinate range metadata
-    partitions: Vec<(Arc<dyn PartitionStream>, PartitionMetadata)>,
+    /// Partition streams paired with their coordinate range metadata.
+    /// Stored behind the `ProjectableStream` trait so `PrunableStreamingTable`
+    /// is not coupled to `PyArrowStreamPartition`.
+    partitions: Vec<(Arc<dyn ProjectableStream>, PartitionMetadata)>,
     /// Set of column names that are dimension columns (eligible for pruning)
     dimension_columns: HashSet<String>,
 }
 
 impl PrunableStreamingTable {
     fn new(
         schema: SchemaRef,
-        partitions: Vec<(Arc<dyn PartitionStream>, PartitionMetadata)>,
+        partitions: Vec<(Arc<dyn ProjectableStream>, PartitionMetadata)>,
     ) -> Self {
         // Collect dimension column names from the first partition that has
         // non-empty metadata. All partitions share the same dimension names,
@@ -435,6 +437,24 @@ impl PrunableStreamingTable {
     }
 }
 
+/// Extension trait for partition streams that support column projection.
+///
+/// Implemented by `PyArrowStreamPartition` so that `PrunableStreamingTable`
+/// can push projections to Python factories without coupling to the concrete type.
+/// Any new stream implementation (e.g. for non-Python backends) can implement this
+/// trait and be used with `PrunableStreamingTable` directly.
+trait ProjectableStream: PartitionStream + Debug {
+    /// Return a new stream that emits only the specified columns.
+    fn clone_with_projection(
+        &self,
+        projection: Arc<[String]>,
+        projected_schema: SchemaRef,
+    ) -> Arc<dyn PartitionStream>;
+
+    /// Clone this stream as a generic `PartitionStream` Arc.
+    fn clone_as_stream(&self) -> Arc<dyn PartitionStream>;
+}
+
 /// Flip a comparison operator (for when literal is on left side).
 fn flip_operator(op: &Operator) -> Operator {
     match op {
@@ -545,49 +565,137 @@ impl TableProvider for PrunableStreamingTable {
         // Prune partitions based on filters
         let included_indices = self.prune_partitions(filters);
 
-        // Collect only the included partition streams
-        let included_partitions: Vec<Arc<dyn PartitionStream>> = included_indices
-            .iter()
-            .map(|&idx| Arc::clone(&self.partitions[idx].0))
-            .collect();
-
-        // Handle empty case - create an empty streaming table
-        if included_partitions.is_empty() {
-            // Create a streaming table with no partitions
-            // DataFusion will return empty result
+        // Handle empty case — all partitions pruned, return empty plan
+        if included_indices.is_empty() {
             let empty_table = StreamingTable::try_new(Arc::clone(&self.schema), vec![])?;
             return empty_table.scan(state, projection, filters, limit).await;
         }
 
-        // Create StreamingTable with the pruned partitions
-        let streaming = StreamingTable::try_new(Arc::clone(&self.schema), included_partitions)?;
+        // Determine whether to push projection down to the Python factory.
+        //
+        // We push when the projection includes at least one data variable
+        // (non-dimension column), because xarray can selectively load only
+        // the requested data arrays while dimension coordinates are always
+        // available via xarray's coordinate system.
+        //
+        // We do NOT push when:
+        //   - projection is None (load everything — factory receives None)
+        //   - projection is Some([]) (COUNT(*) — let StreamingTable handle)
+        //   - projection contains only dimension columns (ds.to_dataframe()
+        //     needs at least one data variable; dimensions are always loaded)
+        let push_projection = match projection {
+            Some(indices) if !indices.is_empty() => indices
+                .iter()
+                .any(|&i| !self.dimension_columns.contains(self.schema.field(i).name())),
+            _ => false,
+        };
 
-        // Delegate to StreamingTable for actual execution.
-        // We report Inexact, so DataFusion still applies row-level filtering.
-        streaming.scan(state, projection, filters, limit).await
+        if push_projection {
+            let indices = projection.unwrap();
+
+            // Build the projected schema (only the requested fields)
+            let proj_fields: Vec<_> = indices
+                .iter()
+                .map(|&i| self.schema.field(i).clone())
+                .collect();
+            let projected_schema = Arc::new(Schema::new(proj_fields));
+
+            // Collect the requested column names to send to the factory.
+            // Stored in an Arc so each clone_with_projection call shares the
+            // same allocation via an atomic refcount increment (no N Vec copies).
+            let proj_col_names: Arc<[String]> = indices
+                .iter()
+                .map(|&i| self.schema.field(i).name().to_string())
+                .collect::<Vec<_>>()
+                .into();
+
+            // Clone each pruned partition with the projection baked in.
+            // The factory will receive proj_col_names and load only those vars.
+            let projected_partitions: Vec<Arc<dyn PartitionStream>> = included_indices
+                .iter()
+                .map(|&idx| {
+                    self.partitions[idx].0.clone_with_projection(
+                        Arc::clone(&proj_col_names),
+                        Arc::clone(&projected_schema),
+                    )
+                })
+                .collect();
+
+            // StreamingTable already has the projected schema — pass None for
+            // projection so it doesn't wrap the stream in a redundant ProjectionExec.
+            let streaming = StreamingTable::try_new(projected_schema, projected_partitions)?;
+            streaming.scan(state, None, filters, limit).await
+        } else {
+            // No projection pushdown — factory is called with None (loads all
+            // columns). StreamingTable applies projection via ProjectionExec.
+            let included_partitions: Vec<Arc<dyn PartitionStream>> = included_indices
+                .iter()
+                .map(|&idx| self.partitions[idx].0.clone_as_stream())
+                .collect();
+            let streaming = StreamingTable::try_new(Arc::clone(&self.schema), included_partitions)?;
+            streaming.scan(state, projection, filters, limit).await
+        }
     }
 }
 
 /// A partition stream that wraps a Python factory function that creates streams.
 ///
 /// The factory is called lazily on each `execute()` invocation, allowing
 /// the same table to be queried multiple times.
+///
+/// When `projection` is set, the factory is called with that list of column
+/// names so that xarray only loads the requested data variables rather than
+/// materializing every variable in the dataset.
 struct PyArrowStreamPartition {
     schema: SchemaRef,
-    /// A Python callable (factory) that returns a fresh stream implementing `__arrow_c_stream__`.
-    /// Called on each execute() to create a new stream.
-    stream_factory: Py<PyAny>,
+    /// A Python callable (factory) that returns a fresh stream.
+    /// Signature: `make_stream(projection_names: Optional[List[str]]) -> RecordBatchReader`
+    ///
+    /// Wrapped in `Arc` so `ProjectableStream::clone_with_projection` can share
+    /// the same Python object across projected partitions without acquiring the
+    /// GIL — only an atomic reference-count increment is needed.
+    stream_factory: Arc<Py<PyAny>>,
+    /// Column names to pass to the factory. `None` means load all columns.
+    /// Stored as `Arc<[String]>` so multiple projected clones share one allocation.
+    projection: Option<Arc<[String]>>,
 }
 
 impl PyArrowStreamPartition {
     fn new(stream_factory: Py<PyAny>, schema: SchemaRef) -> Self {
         Self {
             schema,
-            stream_factory,
+            stream_factory: Arc::new(stream_factory),
+            projection: None,
         }
     }
 }
 
+impl ProjectableStream for PyArrowStreamPartition {
+    /// Return a new partition that emits only the given columns.
+    ///
+    /// Clones the factory `Arc` (atomic refcount increment, no GIL) so the
+    /// same Python callable is shared across all projected partitions.
+    fn clone_with_projection(
+        &self,
+        projection: Arc<[String]>,
+        projected_schema: SchemaRef,
+    ) -> Arc<dyn PartitionStream> {
+        Arc::new(Self {
+            schema: projected_schema,
+            stream_factory: Arc::clone(&self.stream_factory),
+            projection: Some(projection),
+        })
+    }
+
+    fn clone_as_stream(&self) -> Arc<dyn PartitionStream> {
+        Arc::new(Self {
+            schema: Arc::clone(&self.schema),
+            stream_factory: Arc::clone(&self.stream_factory),
+            projection: self.projection.clone(),
+        })
+    }
+}
+
 impl Debug for PyArrowStreamPartition {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("PyArrowStreamPartition")
@@ -604,8 +712,9 @@ impl PartitionStream for PyArrowStreamPartition {
     fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
         let schema = Arc::clone(&self.schema);
 
-        // Clone the factory with the GIL held
-        let factory = Python::attach(|py| self.stream_factory.clone_ref(py));
+        // Clone the factory Arc (no GIL needed) and the projection list.
+        let factory = Arc::clone(&self.stream_factory);
+        let projection = self.projection.clone();
 
         // Create a lazy stream using try_stream! macro.
         // The GIL is acquired only for the duration of each Python call
@@ -614,9 +723,22 @@ impl PartitionStream for PyArrowStreamPartition {
         // executor a chance to poll other partition streams (which can then
         // acquire the GIL and make progress in parallel).
         let batch_stream = try_stream! {
-            // Call factory to get the PyArrow RecordBatchReader
+            // Call factory with the projection argument.
+            // `projection` is either a Python list of column names or None
+            // (load all columns).  The factory always receives exactly one arg
+            // so it can distinguish "no projection" from "empty projection".
             let reader: Py<PyAny> = Python::attach(|py| {
-                factory.call0(py).map_err(|e| {
+                let proj_arg = match &projection {
+                    Some(cols) => PyList::new(py, cols.iter().map(|s| s.as_str()))
+                        .map_err(|e| {
+                            DataFusionError::Execution(format!(
+                                "Failed to build projection list: {e}"
+                            ))
+                        })?
+                        .into_any(),
+                    None => py.None().into_bound(py).into_any(),
+                };
+                factory.call1(py, (proj_arg,)).map_err(|e| {
                     DataFusionError::Execution(format!("Failed to call stream factory: {e}"))
                 })
             })?;
@@ -748,7 +870,9 @@ impl LazyArrowStreamTable {
         // eliminating the per-partition Python::attach() calls of the old
         // three-list approach.  Python can release each block dict, factory
         // closure, and metadata dict as soon as Rust has ingested them.
-        let mut partition_list: Vec<(Arc<dyn PartitionStream>, PartitionMetadata)> = Vec::new();
+        // Stored as Arc<dyn ProjectableStream> so PrunableStreamingTable
+        // is decoupled from PyArrowStreamPartition.
+        let mut partition_list: Vec<(Arc<dyn ProjectableStream>, PartitionMetadata)> = Vec::new();
         for item_result in partitions.try_iter()? {
             let item = item_result?;
             let (factory_obj, meta_obj): (Py<PyAny>, Py<PyAny>) = item.extract().map_err(|e| {
@@ -757,8 +881,8 @@ impl LazyArrowStreamTable {
                 ))
             })?;
             let meta = convert_python_metadata_from_bound(meta_obj.bind(partitions.py()))?;
-            let partition = Arc::new(PyArrowStreamPartition::new(factory_obj, schema_ref.clone()))
-                as Arc<dyn PartitionStream>;
+            let partition: Arc<dyn ProjectableStream> =
+                Arc::new(PyArrowStreamPartition::new(factory_obj, schema_ref.clone()));
             partition_list.push((partition, meta));
         }
 

xarray_sql/df.py

@@ -324,11 +324,11 @@ def _block_metadata(coord_arrays: dict, block: Block) -> PartitionBounds:
   for dim, slc in block.items():
     coord_values = coord_arrays[str(dim)][slc]
     if len(coord_values) > 0:
-      first, last = coord_values[0], coord_values[-1]
-      if first <= last:
-        min_val, max_val = first, last
-      else:
-        min_val, max_val = last, first
+      # Use actual min/max rather than first/last so that non-monotonic
+      # coordinate axes (e.g. descending latitude 90→-90) are handled
+      # correctly.  np.min/max work for both numeric and datetime64 arrays.
+      min_val = coord_values.min()
+      max_val = coord_values.max()
 
       if isinstance(min_val, (np.datetime64, pd.Timestamp)):
         min_val = int(pd.Timestamp(min_val).value)

xarray_sql/reader.py

@@ -63,7 +63,9 @@ def __init__(
       chunks: Chunks = None,
       *,
       batch_size: int = DEFAULT_BATCH_SIZE,
-      _iteration_callback: Callable[[Block], None] | None = None,
+      _iteration_callback: (
+          Callable[[Block, list[str] | None], None] | None
+      ) = None,
   ):
     """Initialize the lazy reader.
 
@@ -105,9 +107,10 @@ def _generate_batches(self) -> Iterator[pa.RecordBatch]:
     emitted as one or more RecordBatches of at most self._batch_size rows.
     """
     for block in block_slices(self._ds, self._chunks):
-      # Call the iteration callback if provided (for testing)
+      # Call the iteration callback if provided (for testing).
+      # XarrayRecordBatchReader has no projection concept, so always passes None.
       if self._iteration_callback is not None:
-        self._iteration_callback(block)
+        self._iteration_callback(block, None)
 
       yield from iter_record_batches(
           self._ds.isel(block), self._schema, self._batch_size
@@ -187,7 +190,9 @@ def read_xarray_table(
     chunks: Chunks = None,
     *,
     batch_size: int = DEFAULT_BATCH_SIZE,
-    _iteration_callback: Callable[[Block], None] | None = None,
+    _iteration_callback: (
+        Callable[[Block, list[str] | None], None] | None
+    ) = None,
 ) -> "LazyArrowStreamTable":
   """Create a lazy DataFusion table from an xarray Dataset.
 
@@ -248,14 +253,38 @@ def read_xarray_table(
   # Zarr-backed datasets (e.g. ARCO-ERA5 on GCS).
   coord_arrays = {str(dim): ds.coords[dim].values for dim in ds.dims}
 
+  # Determine which column names are data variables (not dimension coordinates).
+  # Used by the factory to skip loading unrequested variables.
+  data_var_names = set(ds.data_vars.keys())
+
   def make_partition_factory(
       block: Block,
-  ) -> Callable[[], pa.RecordBatchReader]:
-    def make_stream() -> pa.RecordBatchReader:
+  ) -> Callable[[list[str] | None], pa.RecordBatchReader]:
+    def make_stream(
+        projection_names: list[str] | None,
+    ) -> pa.RecordBatchReader:
       if _iteration_callback is not None:
-        _iteration_callback(block)
+        _iteration_callback(block, projection_names)
+
+      if projection_names is not None:
+        # Restrict to the data variables mentioned in the projection.
+        # Dimension coordinates come along automatically via coords.
+        data_vars_needed = [c for c in projection_names if c in data_var_names]
+        if data_vars_needed:
+          ds_block = ds[data_vars_needed].isel(block)
+        else:
+          # Only dimension coords requested — drop all data vars to avoid
+          # loading them unnecessarily (e.g. for queries like SELECT lat, lon).
+          ds_block = ds[[]].isel(block)
+        batch_schema = pa.schema(
+            [schema.field(name) for name in projection_names]
+        )
+      else:
+        ds_block = ds.isel(block)
+        batch_schema = schema
+
       return pa.RecordBatchReader.from_batches(
-          schema, iter_record_batches(ds.isel(block), schema, batch_size)
+          batch_schema, iter_record_batches(ds_block, batch_schema, batch_size)
       )
 
     return make_stream