From 0f271dba32ea8f25bbb8ef15df0c96aeb333a4f0 Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Mon, 9 Feb 2026 13:26:43 -0800 Subject: [PATCH 1/2] Xarray-specific filter predicate pushdown. Partition pruning: we can get statistics for the "primary keys" i.e. the dataset dims. --- Cargo.lock | 1 + Cargo.toml | 1 + src/lib.rs | 590 ++++++++++++++++++++++++++++++++++++-- xarray_sql/df.py | 46 +++ xarray_sql/reader.py | 34 ++- xarray_sql/reader_test.py | 280 ++++++++++++++++++ 6 files changed, 912 insertions(+), 40 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0b6258f1..db08ef5b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3443,6 +3443,7 @@ version = "0.2.0" dependencies = [ "arrow", "async-stream", + "async-trait", "datafusion", "datafusion-ffi", "futures", diff --git a/Cargo.toml b/Cargo.toml index d25f58da..33bf5810 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ exclude = [ [dependencies] arrow = { version = "57.2.0", features = ["pyarrow"] } async-stream = "0.3" +async-trait = "0.1" datafusion = { version = "51.0.0" } datafusion-ffi = { version = "51.0.0" } futures = { version = "0.3" } diff --git a/src/lib.rs b/src/lib.rs index 0a432e89..faf917ae 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -30,7 +30,21 @@ //! multiple cores. Due to a bug in DataFusion v51.0.0's `collect()` method, aggregation //! queries should use `to_arrow_table()` instead to ensure complete results. //! TODO(#107): Upgrading to the latest datafusion-python (52+) should fix this. +//! +//! ## Filter Pushdown (Partition Pruning) +//! +//! When partition metadata is provided, SQL filters on dimension columns (time, lat, lon) +//! automatically prune partitions that can't contain matching rows. For example: +//! +//! ```sql +//! SELECT * FROM air WHERE time > '2020-02-01' +//! ``` +//! +//! Will skip loading partitions whose time ranges are entirely before 2020-02-01. +//! Supported operators: `=`, `<`, `>`, `<=`, `>=`, `BETWEEN`, `IN`, `AND`, `OR`. +use std::any::Any; +use std::collections::HashMap; use std::ffi::CString; use std::fmt::Debug; use std::sync::Arc; @@ -39,18 +53,498 @@ use arrow::array::RecordBatch; use arrow::datatypes::SchemaRef; use arrow::pyarrow::FromPyArrow; use async_stream::try_stream; +use async_trait::async_trait; use datafusion::catalog::streaming::StreamingTable; -use datafusion::common::DataFusionError; +use datafusion::catalog::Session; +use datafusion::common::{DataFusionError, Result as DFResult, ScalarValue}; use datafusion::datasource::TableProvider; use datafusion::execution::TaskContext; +use datafusion::logical_expr::expr::InList; +use datafusion::logical_expr::{ + BinaryExpr, Expr, Operator, TableProviderFilterPushDown, TableType, +}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::streaming::PartitionStream; -use datafusion::physical_plan::SendableRecordBatchStream; +use datafusion::physical_plan::{ExecutionPlan, SendableRecordBatchStream}; use datafusion_ffi::table_provider::FFI_TableProvider; use pyo3::prelude::*; use pyo3::types::PyCapsule; use tokio::runtime::Handle; +// ============================================================================ +// Partition Metadata Types for Filter Pushdown +// ============================================================================ + +// TODO(alxmrs, Claude): Support every valid xarray coordinate type. +/// Scalar value for dimension bounds, supporting common xarray coordinate types. +#[derive(Clone, Debug)] +pub enum ScalarBound { + /// 64-bit integer (for integer coordinates) + Int64(i64), + /// 64-bit float (for lat/lon coordinates) + Float64(f64), + /// Nanoseconds since Unix epoch (for datetime64[ns] coordinates) + TimestampNanos(i64), +} + +impl ScalarBound { + /// Compare this bound with a DataFusion ScalarValue. + /// Returns None if types are incompatible. + fn compare_to_scalar(&self, scalar: &ScalarValue) -> Option { + match (self, scalar) { + // Integer comparisons + (ScalarBound::Int64(a), ScalarValue::Int64(Some(b))) => Some(a.cmp(b)), + (ScalarBound::Int64(a), ScalarValue::Int32(Some(b))) => Some(a.cmp(&(*b as i64))), + + // Float comparisons + (ScalarBound::Float64(a), ScalarValue::Float64(Some(b))) => a.partial_cmp(b), + (ScalarBound::Float64(a), ScalarValue::Float32(Some(b))) => a.partial_cmp(&(*b as f64)), + + // Timestamp comparisons - convert to nanoseconds + (ScalarBound::TimestampNanos(a), ScalarValue::TimestampNanosecond(Some(b), _)) => { + Some(a.cmp(b)) + } + (ScalarBound::TimestampNanos(a), ScalarValue::TimestampMicrosecond(Some(b), _)) => { + Some(a.cmp(&(b * 1_000))) + } + (ScalarBound::TimestampNanos(a), ScalarValue::TimestampMillisecond(Some(b), _)) => { + Some(a.cmp(&(b * 1_000_000))) + } + (ScalarBound::TimestampNanos(a), ScalarValue::TimestampSecond(Some(b), _)) => { + Some(a.cmp(&(b * 1_000_000_000))) + } + + // Incompatible types + _ => None, + } + } +} + +/// Range bounds for one dimension in a partition. +#[derive(Clone, Debug)] +pub struct DimensionRange { + /// The column name (dimension name from xarray) + pub column_name: String, + /// Minimum value (inclusive) - first coordinate value in this partition + pub min: ScalarBound, + /// Maximum value (inclusive) - last coordinate value in this partition + pub max: ScalarBound, +} + +/// Metadata for a single partition, used for filter-based pruning. +#[derive(Clone, Debug, Default)] +pub struct PartitionMetadata { + /// Dimension ranges for this partition, keyed by column name + pub ranges: HashMap, +} + +impl PartitionMetadata { + /// Get the range for a specific dimension column. + pub fn get_range(&self, column: &str) -> Option<&DimensionRange> { + self.ranges.get(column) + } +} + +// ============================================================================ +// Custom TableProvider with Filter Pushdown +// ============================================================================ + +/// A table provider that supports partition pruning via filter pushdown. +/// +/// This wraps partition streams with their metadata and implements +/// `TableProvider::supports_filters_pushdown` and partition pruning in `scan()`. +struct PrunableStreamingTable { + schema: SchemaRef, + /// Partition streams paired with their coordinate range metadata + partitions: Vec<(Arc, PartitionMetadata)>, + /// Set of column names that are dimension columns (eligible for pruning) + dimension_columns: std::collections::HashSet, +} + +impl PrunableStreamingTable { + fn new( + schema: SchemaRef, + partitions: Vec<(Arc, PartitionMetadata)>, + ) -> Self { + // Collect dimension column names from partition metadata + let dimension_columns: std::collections::HashSet = partitions + .first() + .map(|(_, meta)| meta.ranges.keys().cloned().collect()) + .unwrap_or_default(); + + Self { + schema, + partitions, + dimension_columns, + } + } + + /// Determine which partitions should be included based on filters. + /// Returns indices of partitions that may contain matching rows. + fn prune_partitions(&self, filters: &[Expr]) -> Vec { + self.partitions + .iter() + .enumerate() + .filter(|(_, (_, meta))| { + // Include partition unless a filter definitely excludes it + !filters + .iter() + .any(|f| self.filter_excludes_partition(f, meta)) + }) + .map(|(idx, _)| idx) + .collect() + } + + /// Returns true if this filter definitely excludes the partition. + /// Conservative: returns false (include) if uncertain. + fn filter_excludes_partition(&self, expr: &Expr, meta: &PartitionMetadata) -> bool { + match expr { + Expr::BinaryExpr(BinaryExpr { left, op, right }) => { + // Handle AND/OR logic + match op { + Operator::And => { + // For AND, exclude if either side excludes + self.filter_excludes_partition(left, meta) + || self.filter_excludes_partition(right, meta) + } + Operator::Or => { + // For OR, exclude only if both sides exclude + self.filter_excludes_partition(left, meta) + && self.filter_excludes_partition(right, meta) + } + // Handle comparison operators + _ => self.comparison_excludes(left, op, right, meta), + } + } + Expr::Not(inner) => { + // NOT inverts the logic, but we can't easily invert exclusion + // Be conservative and don't exclude + !self.filter_excludes_partition(inner, meta) + } + Expr::Between(between) => self.between_excludes(between, meta), + Expr::InList(in_list) => self.in_list_excludes(in_list, meta), + // Unknown expression type - be conservative + _ => false, + } + } + + /// Check if a comparison expression excludes this partition. + fn comparison_excludes( + &self, + left: &Expr, + op: &Operator, + right: &Expr, + meta: &PartitionMetadata, + ) -> bool { + // Try to extract column and literal from either side + let (col_name, scalar, flipped) = match (left.as_ref(), right.as_ref()) { + (Expr::Column(c), Expr::Literal(s, _)) => (c.name.clone(), s, false), + (Expr::Literal(s, _), Expr::Column(c)) => (c.name.clone(), s, true), + _ => return false, // Not a simple column-literal comparison + }; + + // Get the dimension range for this column + let range = match meta.get_range(&col_name) { + Some(r) => r, + None => return false, // Not a dimension column, can't prune + }; + + // Flip operator if literal was on left side + let effective_op = if flipped { + flip_operator(op) + } else { + op.clone() + }; + + // Determine if partition can be excluded based on operator + match effective_op { + // col > literal: exclude if max <= literal + Operator::Gt => matches!( + range.max.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Less | std::cmp::Ordering::Equal) + ), + // col >= literal: exclude if max < literal + Operator::GtEq => { + matches!( + range.max.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Less) + ) + } + // col < literal: exclude if min >= literal + Operator::Lt => matches!( + range.min.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Greater | std::cmp::Ordering::Equal) + ), + // col <= literal: exclude if min > literal + Operator::LtEq => { + matches!( + range.min.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Greater) + ) + } + // col = literal: exclude if literal outside [min, max] + Operator::Eq => { + let below_min = matches!( + range.min.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Greater) + ); + let above_max = matches!( + range.max.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Less) + ); + below_min || above_max + } + // col != literal: can't exclude (partition may have other values) + Operator::NotEq => false, + // Other operators: be conservative + _ => false, + } + } + + /// Check if a BETWEEN expression excludes this partition. + fn between_excludes( + &self, + between: &datafusion::logical_expr::Between, + meta: &PartitionMetadata, + ) -> bool { + if between.negated { + // NOT BETWEEN is complex, be conservative + return false; + } + + // Extract column name + let col_name = match between.expr.as_ref() { + Expr::Column(c) => c.name.clone(), + _ => return false, + }; + + // Get dimension range + let range = match meta.get_range(&col_name) { + Some(r) => r, + None => return false, + }; + + // Extract low and high bounds + let (low, high) = match (between.low.as_ref(), between.high.as_ref()) { + (Expr::Literal(l, _), Expr::Literal(h, _)) => (l, h), + _ => return false, + }; + + // Exclude if partition range doesn't overlap with [low, high] + // No overlap if: partition.max < low OR partition.min > high + let max_below_low = matches!( + range.max.compare_to_scalar(low), + Some(std::cmp::Ordering::Less) + ); + let min_above_high = matches!( + range.min.compare_to_scalar(high), + Some(std::cmp::Ordering::Greater) + ); + + max_below_low || min_above_high + } + + /// Check if an IN list expression excludes this partition. + fn in_list_excludes(&self, in_list: &InList, meta: &PartitionMetadata) -> bool { + if in_list.negated { + // NOT IN is complex, be conservative + return false; + } + + // Extract column name + let col_name = match in_list.expr.as_ref() { + Expr::Column(c) => c.name.clone(), + _ => return false, + }; + + // Get dimension range + let range = match meta.get_range(&col_name) { + Some(r) => r, + None => return false, + }; + + // Check if any value in the list could be in this partition's range + let any_in_range = in_list.list.iter().any(|expr| { + if let Expr::Literal(scalar, _) = expr { + // Value is in range if: min <= value <= max + let above_min = matches!( + range.min.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Less | std::cmp::Ordering::Equal) + ); + let below_max = matches!( + range.max.compare_to_scalar(scalar), + Some(std::cmp::Ordering::Greater | std::cmp::Ordering::Equal) + ); + above_min && below_max + } else { + // Non-literal in list, be conservative + true + } + }); + + // Exclude only if NO values could be in range + !any_in_range + } + + /// Check if an expression is a filter on a dimension column. + fn is_dimension_filter(&self, expr: &Expr) -> bool { + match expr { + Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op { + Operator::And | Operator::Or => { + self.is_dimension_filter(left) || self.is_dimension_filter(right) + } + _ => self.expr_references_dimension(left) || self.expr_references_dimension(right), + }, + Expr::Between(b) => self.expr_references_dimension(&b.expr), + Expr::InList(i) => self.expr_references_dimension(&i.expr), + Expr::Not(inner) => self.is_dimension_filter(inner), + _ => false, + } + } + + /// Check if an expression references a dimension column. + fn expr_references_dimension(&self, expr: &Expr) -> bool { + match expr { + Expr::Column(c) => self.dimension_columns.contains(&c.name), + _ => false, + } + } +} + +/// Flip a comparison operator (for when literal is on left side). +fn flip_operator(op: &Operator) -> Operator { + match op { + Operator::Lt => Operator::Gt, + Operator::LtEq => Operator::GtEq, + Operator::Gt => Operator::Lt, + Operator::GtEq => Operator::LtEq, + other => other.clone(), + } +} + +/// Convert a Python object to a ScalarBound. +fn python_to_scalar_bound(obj: &Bound<'_, PyAny>) -> PyResult { + // Try integer first (includes timestamps as nanoseconds) + if let Ok(val) = obj.extract::() { + // Check if this might be a timestamp (very large number) + // Python passes datetime64[ns] as nanoseconds since epoch + if val.abs() > 1_000_000_000_000_000 { + // Likely a nanosecond timestamp + return Ok(ScalarBound::TimestampNanos(val)); + } + return Ok(ScalarBound::Int64(val)); + } + + // Try float + if let Ok(val) = obj.extract::() { + return Ok(ScalarBound::Float64(val)); + } + + Err(pyo3::exceptions::PyTypeError::new_err(format!( + "Unsupported type for partition bound: {:?}", + obj.get_type().name() + ))) +} + +/// Convert Python partition metadata dict to Rust PartitionMetadata. +fn convert_python_metadata( + meta_dict: HashMap, Py)>, +) -> PyResult { + Python::attach(|py| { + let mut ranges = HashMap::new(); + for (dim_name, (min_obj, max_obj)) in meta_dict { + let min_bound = python_to_scalar_bound(min_obj.bind(py))?; + let max_bound = python_to_scalar_bound(max_obj.bind(py))?; + ranges.insert( + dim_name.clone(), + DimensionRange { + column_name: dim_name, + min: min_bound, + max: max_bound, + }, + ); + } + Ok(PartitionMetadata { ranges }) + }) +} + +impl Debug for PrunableStreamingTable { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PrunableStreamingTable") + .field("schema", &self.schema) + .field("num_partitions", &self.partitions.len()) + .field("dimension_columns", &self.dimension_columns) + .finish() + } +} + +#[async_trait] +impl TableProvider for PrunableStreamingTable { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> DFResult> { + // For dimension filters we can do exact pruning at partition level + // Return Inexact so DataFusion still applies row-level filtering + Ok(filters + .iter() + .map(|expr| { + if self.is_dimension_filter(expr) { + // We can prune partitions but not individual rows within + TableProviderFilterPushDown::Inexact + } else { + TableProviderFilterPushDown::Unsupported + } + }) + .collect()) + } + + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> DFResult> { + // Prune partitions based on filters + let included_indices = self.prune_partitions(filters); + + // Collect only the included partition streams + let included_partitions: Vec> = 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 + let empty_table = StreamingTable::try_new(Arc::clone(&self.schema), vec![])?; + return empty_table.scan(state, projection, &[], limit).await; + } + + // Create StreamingTable with the pruned partitions + let streaming = StreamingTable::try_new(Arc::clone(&self.schema), included_partitions)?; + + // Delegate to StreamingTable for actual execution + // Pass empty filters since we've already done partition pruning + // DataFusion will still apply row-level filtering + streaming.scan(state, projection, &[], limit).await + } +} + /// A partition stream that wraps a Python factory function that creates streams. /// /// The factory is called lazily on each `execute()` invocation, allowing @@ -147,6 +641,12 @@ impl PartitionStream for PyArrowStreamPartition { /// Each partition has its own factory function that is called on query execution /// to create a fresh stream, enabling true parallelism in DataFusion. /// +/// ## Filter Pushdown +/// +/// When partition metadata is provided, SQL filters on dimension columns (time, lat, lon, etc.) +/// will automatically prune partitions that can't contain matching rows. This dramatically +/// improves query performance for range queries on large datasets. +/// /// # Note /// /// Due to a bug in DataFusion v51.0.0's `collect()` method, use `to_arrow_table()` @@ -165,22 +665,28 @@ impl PartitionStream for PyArrowStreamPartition { /// lambda: pa.RecordBatchReader.from_batches(schema, batches_chunk_1), /// ] /// -/// # Wrap factories in lazy table - NO DATA LOADED -/// table = LazyArrowStreamTable(factories, schema) +/// # Partition metadata for filter pushdown (optional) +/// # Each dict maps dimension name to (min, max) coordinate values +/// metadata = [ +/// {'time': (0, 1000000000), 'lat': (-90.0, 0.0)}, # partition 0 +/// {'time': (1000000001, 2000000000), 'lat': (0.0, 90.0)}, # partition 1 +/// ] /// -/// # Register with DataFusion - STILL NO DATA LOADED +/// # Wrap factories in lazy table with metadata +/// table = LazyArrowStreamTable(factories, schema, metadata) +/// +/// # Register with DataFusion /// ctx = SessionContext() /// ctx.register_table("air", table) /// -/// # Data only loaded HERE during query execution -/// # Each partition runs in parallel with its own factory -/// # Use to_arrow_table() for aggregation queries -/// result = ctx.sql("SELECT AVG(air) FROM air").to_arrow_table() +/// # Queries with filters on dimension columns will prune partitions! +/// # This query might only read partition 1: +/// result = ctx.sql("SELECT AVG(air) FROM air WHERE lat > 0").to_arrow_table() /// ``` #[pyclass(name = "LazyArrowStreamTable")] struct LazyArrowStreamTable { - /// The underlying StreamingTable - table: Arc, + /// The underlying table provider with pruning support + table: Arc, } #[pymethods] @@ -194,12 +700,20 @@ impl LazyArrowStreamTable { /// Called on each query execution to create fresh streams. /// schema: A PyArrow Schema for the table. Required since the factories /// haven't been called yet. + /// partition_metadata: Optional list of dicts mapping dimension names to + /// (min, max) tuples. When provided, enables filter pushdown to + /// prune partitions based on SQL WHERE clauses. /// /// Raises: /// TypeError: If the schema is not a valid PyArrow Schema. - /// ValueError: If stream_factories is empty. + /// ValueError: If stream_factories is empty or metadata length doesn't match. #[new] - fn new(stream_factories: &Bound<'_, PyAny>, schema: &Bound<'_, PyAny>) -> PyResult { + #[pyo3(signature = (stream_factories, schema, partition_metadata=None))] + fn new( + stream_factories: &Bound<'_, PyAny>, + schema: &Bound<'_, PyAny>, + partition_metadata: Option<&Bound<'_, PyAny>>, + ) -> PyResult { // Convert the PyArrow schema to Arrow schema use arrow::datatypes::Schema; use arrow::pyarrow::FromPyArrow; @@ -222,21 +736,45 @@ impl LazyArrowStreamTable { )); } - // Create one partition per factory - let partitions: Vec> = factories + // Extract and convert partition metadata if provided + let metadata_list: Vec = if let Some(meta_py) = partition_metadata { + let meta_dicts: Vec, Py)>> = + meta_py.extract().map_err(|e| { + pyo3::exceptions::PyTypeError::new_err(format!( + "partition_metadata must be a list of dicts: {e}" + )) + })?; + + if meta_dicts.len() != factories.len() { + return Err(pyo3::exceptions::PyValueError::new_err(format!( + "partition_metadata length ({}) must match stream_factories length ({})", + meta_dicts.len(), + factories.len() + ))); + } + + meta_dicts + .into_iter() + .map(|meta_dict| convert_python_metadata(meta_dict)) + .collect::>>()? + } else { + // No metadata provided - create empty metadata for each partition + vec![PartitionMetadata::default(); factories.len()] + }; + + // Create partitions with their metadata + let partitions: Vec<(Arc, PartitionMetadata)> = factories .into_iter() - .map(|factory| { - Arc::new(PyArrowStreamPartition::new(factory, schema_ref.clone())) - as Arc + .zip(metadata_list.into_iter()) + .map(|(factory, meta)| { + let partition = Arc::new(PyArrowStreamPartition::new(factory, schema_ref.clone())) + as Arc; + (partition, meta) }) .collect(); - // Create the StreamingTable with multiple partitions - let table = StreamingTable::try_new(schema_ref, partitions).map_err(|e| { - pyo3::exceptions::PyRuntimeError::new_err(format!( - "Failed to create StreamingTable: {e}" - )) - })?; + // Create the PrunableStreamingTable + let table = PrunableStreamingTable::new(schema_ref, partitions); Ok(Self { table: Arc::new(table), @@ -257,9 +795,9 @@ impl LazyArrowStreamTable { // Try to get the current tokio runtime handle (available when called from DataFusion context) let runtime = Handle::try_current().ok(); - // Create FFI wrapper + // Create FFI wrapper with filter pushdown ENABLED let ffi_provider = FFI_TableProvider::new( - provider, false, // can_support_pushdown_filters + provider, true, // can_support_pushdown_filters = ENABLED runtime, ); diff --git a/xarray_sql/df.py b/xarray_sql/df.py index c799626f..a57f7201 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -170,3 +170,49 @@ def _parse_schema(ds) -> pa.Schema: columns.append(pa.field(var_name, pa_type)) return pa.schema(columns) + + +# Type alias for partition metadata: maps dimension name to (min, max) values +PartitionBounds = t.Dict[str, t.Tuple[t.Any, t.Any]] + + +def partition_metadata( + ds: xr.Dataset, blocks: t.List[Block] +) -> t.List[PartitionBounds]: + """Compute min/max coordinate values for each partition. + + This metadata enables filter pushdown: SQL queries with WHERE clauses + on dimension columns can prune partitions that can't contain matching rows. + + Args: + ds: The xarray Dataset containing coordinate values. + blocks: List of block slices from block_slices(). + + Returns: + List of dicts mapping dimension name to (min_value, max_value) tuples. + - For datetime64, values are nanoseconds since Unix epoch (int64) + - For numeric types, values are Python int or float + """ + metadata = [] + for block in blocks: + ranges: PartitionBounds = {} + for dim, slc in block.items(): + coord_values = ds.coords[dim].values[slc] + if len(coord_values) > 0: + min_val = coord_values[0] + max_val = coord_values[-1] + + # Convert numpy scalar types to Python native types + # This is required for PyO3 FFI conversion + if isinstance(min_val, (np.datetime64, pd.Timestamp)): + # Convert datetime to nanoseconds since epoch + min_val = int(pd.Timestamp(min_val).value) + max_val = int(pd.Timestamp(max_val).value) + elif hasattr(min_val, "item"): + # numpy scalar -> Python native + min_val = min_val.item() + max_val = max_val.item() + + ranges[str(dim)] = (min_val, max_val) + metadata.append(ranges) + return metadata diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py index 3b57559d..bb38133d 100644 --- a/xarray_sql/reader.py +++ b/xarray_sql/reader.py @@ -15,7 +15,7 @@ import pyarrow as pa import xarray as xr -from .df import Block, Chunks, block_slices, pivot, _parse_schema +from .df import Block, Chunks, block_slices, partition_metadata, pivot, _parse_schema if t.TYPE_CHECKING: from ._native import LazyArrowStreamTable @@ -182,18 +182,19 @@ def read_xarray_table( Each chunk becomes a separate partition, enabling DataFusion's parallel execution across multiple cores. - Note: - Due to a bug in DataFusion v51.0.0's collect() method, use - `to_arrow_table()` instead of `collect()` for aggregation queries - to ensure complete results:: + Filter Pushdown: + SQL queries with WHERE clauses on dimension columns (time, lat, lon, etc.) + automatically prune partitions that can't contain matching rows. For example: - # Correct - use to_arrow_table() - result = ctx.sql('SELECT lat, AVG(temp) FROM t GROUP BY lat').to_arrow_table() + # This query will skip loading partitions with time < '2020-02-01' + result = ctx.sql('SELECT * FROM air WHERE time > \"2020-02-01\"').to_arrow_table() - # May return partial results with collect() - result = ctx.sql('SELECT lat, AVG(temp) FROM t GROUP BY lat').collect() + Supported operators: =, <, >, <=, >=, BETWEEN, IN, AND, OR. - This should be fixed when we upgrade datafusion-python to 52 (#107). + Note: + Due to a bug in DataFusion v51.0.0's collect() method, use + `to_arrow_table()` instead of `collect()` for aggregation queries + to ensure complete results. This should be fixed in datafusion-python 52+. Args: ds: An xarray Dataset. All data_vars must share the same dimensions. @@ -217,16 +218,20 @@ def read_xarray_table( >>> ctx.register_table('air', table) >>> >>> # Data is only read here, during query execution + >>> # Filters on 'time' will prune partitions automatically! >>> result = ctx.sql('SELECT AVG(air) FROM air').to_arrow_table() - >>> # Can query again - each query creates a fresh stream - >>> result2 = ctx.sql('SELECT * FROM air LIMIT 10').to_arrow_table() """ from ._native import LazyArrowStreamTable # Get schema from dataset without creating a stream schema = _parse_schema(ds) - blocks = block_slices(ds, chunks) + # Compute block slices - need list for metadata computation + blocks = list(block_slices(ds, chunks)) + + # Compute partition metadata for filter pushdown + # This extracts min/max coordinate values per partition + metadata = partition_metadata(ds, blocks) # Create a factory function for each block (partition) # Each factory produces a RecordBatchReader for its specific chunk @@ -250,4 +255,5 @@ def make_stream() -> pa.RecordBatchReader: # Create one factory per block factories = [make_partition_factory(block) for block in blocks] - return LazyArrowStreamTable(factories, schema) + # Pass factories, schema, and metadata to Rust + return LazyArrowStreamTable(factories, schema, metadata) diff --git a/xarray_sql/reader_test.py b/xarray_sql/reader_test.py index 37304258..7b924713 100644 --- a/xarray_sql/reader_test.py +++ b/xarray_sql/reader_test.py @@ -942,3 +942,283 @@ def test_parallel_queries_independent(self, small_ds): assert ( tracker2.iteration_count == 2 ), f"Table2: expected 2 blocks, got {tracker2.iteration_count}" + + +class TestFilterPushdown: + """Tests for partition pruning via filter pushdown. + + These tests verify that SQL filters on dimension columns (time, lat, lon) + correctly prune partitions, reducing the number of partitions read. + """ + + @pytest.fixture + def time_chunked_ds(self): + """Dataset chunked by time for pruning tests.""" + np.random.seed(42) + # 100 days of data, chunked into 4 partitions of 25 days each + time = pd.date_range("2020-01-01", periods=100, freq="D") + lat = np.linspace(-90, 90, 5) + data = np.random.rand(100, 5).astype(np.float32) + + return xr.Dataset( + {"temperature": (["time", "lat"], data)}, + coords={"time": time, "lat": lat}, + ) + + def test_time_gt_filter_prunes_early_partitions(self, time_chunked_ds): + """Query with time > X should skip early partitions.""" + tracker = IterationTracker() + + # 4 partitions: days 0-24, 25-49, 50-74, 75-99 + # (Jan 1-25, Jan 26-Feb 19, Feb 20-Mar 15, Mar 16-Apr 9) + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query only last 25 days (Mar 16+) - should prune first 3 partitions + # 2020-03-16 is day 75 + result = ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE time >= '2020-03-16' + """ + ).to_arrow_table() + + # Should read only 1 partition (the last one) + assert ( + tracker.iteration_count == 1 + ), f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}" + + # Verify data correctness - 25 days * 5 lat = 125 rows + count = result.to_pandas()["cnt"].iloc[0] + assert count == 125, f"Expected 125 rows, got {count}" + + def test_time_lt_filter_prunes_late_partitions(self, time_chunked_ds): + """Query with time < X should skip late partitions.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query only first 25 days (< Jan 26) - should prune last 3 partitions + result = ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE time < '2020-01-26' + """ + ).to_arrow_table() + + # Should read only 1 partition (the first one) + assert ( + tracker.iteration_count == 1 + ), f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}" + + def test_time_between_filter_prunes_outside_range(self, time_chunked_ds): + """Query with BETWEEN should prune partitions outside the range.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query middle 50 days (Feb 1 - Mar 21) - should hit partitions 2 and 3 + result = ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE time BETWEEN '2020-02-01' AND '2020-03-21' + """ + ).to_arrow_table() + + # Should read 2-3 partitions (middle ones) + assert ( + tracker.iteration_count <= 3 + ), f"Expected at most 3 partitions, got {tracker.iteration_count}" + assert tracker.iteration_count >= 1, "Expected at least 1 partition" + + def test_lat_filter_prunes_partitions(self): + """Latitude filter should prune irrelevant partitions.""" + np.random.seed(42) + time = pd.date_range("2020-01-01", periods=10, freq="D") + # 100 lat values from -90 to 90 + lat = np.linspace(-90, 90, 100) + data = np.random.rand(10, 100).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat"], data)}, + coords={"time": time, "lat": lat}, + ) + + tracker = IterationTracker() + + # Chunk by latitude: 4 partitions (25 lat values each) + # Partition 0: lat -90 to ~-45 + # Partition 1: lat ~-45 to 0 + # Partition 2: lat 0 to ~45 + # Partition 3: lat ~45 to 90 + table = read_xarray_table( + ds, + chunks={"lat": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query southern hemisphere only (lat < 0) + result = ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE lat < 0 + """ + ).to_arrow_table() + + # Should read only ~2 partitions (southern hemisphere) + assert ( + tracker.iteration_count <= 3 + ), f"Expected at most 3 partitions for lat < 0, got {tracker.iteration_count}" + + def test_no_pruning_for_data_column_filters(self, time_chunked_ds): + """Filters on data columns (not dimensions) should not prune.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Filter on temperature (data column), not a dimension + result = ctx.sql( + """ + SELECT COUNT(*) FROM test WHERE temperature > 0.5 + """ + ).to_arrow_table() + + # All 4 partitions should be read (can't prune on data column) + assert ( + tracker.iteration_count == 4 + ), f"Expected 4 partitions (no pruning on data column), got {tracker.iteration_count}" + + def test_filter_correctness_preserved(self, time_chunked_ds): + """Verify filtered results are correct after pruning.""" + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Get count with filter + filtered = ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE time >= '2020-02-15' AND time <= '2020-03-15' + """ + ).to_arrow_table() + + # Manual calculation: Feb 15 (day 45) to Mar 15 (day 74) = 30 days + # 30 days * 5 lat values = 150 rows + count = filtered.to_pandas()["cnt"].iloc[0] + assert count == 150, f"Expected 150 rows, got {count}" + + def test_and_filter_combines_pruning(self, time_chunked_ds): + """AND filters should combine for maximum pruning.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Very narrow range that spans only 1 partition + result = ctx.sql( + """ + SELECT * FROM test + WHERE time >= '2020-03-20' AND time <= '2020-04-05' + """ + ).to_arrow_table() + + # Should read only 1 partition + assert ( + tracker.iteration_count == 1 + ), f"Expected 1 partition for narrow AND range, got {tracker.iteration_count}" + + def test_or_filter_is_conservative(self, time_chunked_ds): + """OR filters should include partitions matching either condition.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # First or last partition (OR condition) + result = ctx.sql( + """ + SELECT * FROM test + WHERE time < '2020-01-10' OR time > '2020-03-30' + """ + ).to_arrow_table() + + # Should read at least 2 partitions (first and last) + assert ( + tracker.iteration_count >= 2 + ), f"Expected at least 2 partitions for OR filter, got {tracker.iteration_count}" + + def test_empty_result_from_impossible_filter(self, time_chunked_ds): + """Filter that matches no data should return empty result.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query for dates outside the data range + result = ctx.sql( + """ + SELECT COUNT(*) as cnt FROM test + WHERE time > '2025-01-01' + """ + ).to_arrow_table() + + # Should read 0 partitions (all pruned) + assert ( + tracker.iteration_count == 0 + ), f"Expected 0 partitions for impossible filter, got {tracker.iteration_count}" + + # Result should be 0 rows + count = result.to_pandas()["cnt"].iloc[0] + assert count == 0, f"Expected 0 rows, got {count}" From 575f7dc77b293c4f0f2476bc81dc56fbd8a5020b Mon Sep 17 00:00:00 2001 From: Alex Merose Date: Mon, 9 Feb 2026 14:46:50 -0800 Subject: [PATCH 2/2] Fix clippy issues. --- src/lib.rs | 26 +++++++++++--------------- 1 file changed, 11 insertions(+), 15 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index faf917ae..ee695c8b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -237,7 +237,7 @@ impl PrunableStreamingTable { meta: &PartitionMetadata, ) -> bool { // Try to extract column and literal from either side - let (col_name, scalar, flipped) = match (left.as_ref(), right.as_ref()) { + let (col_name, scalar, flipped) = match (left, right) { (Expr::Column(c), Expr::Literal(s, _)) => (c.name.clone(), s, false), (Expr::Literal(s, _), Expr::Column(c)) => (c.name.clone(), s, true), _ => return false, // Not a simple column-literal comparison @@ -250,11 +250,7 @@ impl PrunableStreamingTable { }; // Flip operator if literal was on left side - let effective_op = if flipped { - flip_operator(op) - } else { - op.clone() - }; + let effective_op = if flipped { flip_operator(op) } else { *op }; // Determine if partition can be excluded based on operator match effective_op { @@ -418,7 +414,7 @@ fn flip_operator(op: &Operator) -> Operator { Operator::LtEq => Operator::GtEq, Operator::Gt => Operator::Lt, Operator::GtEq => Operator::LtEq, - other => other.clone(), + other => *other, } } @@ -738,12 +734,12 @@ impl LazyArrowStreamTable { // Extract and convert partition metadata if provided let metadata_list: Vec = if let Some(meta_py) = partition_metadata { - let meta_dicts: Vec, Py)>> = - meta_py.extract().map_err(|e| { - pyo3::exceptions::PyTypeError::new_err(format!( - "partition_metadata must be a list of dicts: {e}" - )) - })?; + type MetaDict = HashMap, Py)>; + let meta_dicts: Vec = meta_py.extract().map_err(|e| { + pyo3::exceptions::PyTypeError::new_err(format!( + "partition_metadata must be a list of dicts: {e}" + )) + })?; if meta_dicts.len() != factories.len() { return Err(pyo3::exceptions::PyValueError::new_err(format!( @@ -755,7 +751,7 @@ impl LazyArrowStreamTable { meta_dicts .into_iter() - .map(|meta_dict| convert_python_metadata(meta_dict)) + .map(convert_python_metadata) .collect::>>()? } else { // No metadata provided - create empty metadata for each partition @@ -765,7 +761,7 @@ impl LazyArrowStreamTable { // Create partitions with their metadata let partitions: Vec<(Arc, PartitionMetadata)> = factories .into_iter() - .zip(metadata_list.into_iter()) + .zip(metadata_list) .map(|(factory, meta)| { let partition = Arc::new(PyArrowStreamPartition::new(factory, schema_ref.clone())) as Arc;