diff --git a/Cargo.lock b/Cargo.lock index 7ce83cbb..4011549c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3375,6 +3375,7 @@ dependencies = [ "datafusion", "datafusion-ffi", "futures", + "half", "pyo3", "pyo3-build-config", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 6be38d06..408b1809 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,11 +24,12 @@ async-trait = "0.1" datafusion = { version = "54.0.0" } datafusion-ffi = { version = "54.0.0" } futures = { version = "0.3" } +half = "2.7" # `abi3-py310` builds against CPython's stable ABI, so a single wheel per -# platform works on all CPython >= 3.10 (matching `requires-python`). This -# lets the release workflow ship pre-built wheels for every interpreter -# without compiling per-version, avoiding local rebuilds on install. -pyo3 = { version = "0.28.0", features = ["extension-module", "abi3-py310"] } +# platform works on all CPython >= 3.10 (matching `requires-python`). Maturin +# enables `pyo3/extension-module` through pyproject.toml for wheel builds; it +# must stay disabled for ordinary Cargo test binaries so they link libpython. +pyo3 = { version = "0.28.0", features = ["abi3-py310"] } tokio = { version = "1.46.1", features = ["rt"] } diff --git a/pyproject.toml b/pyproject.toml index 8a35be4f..c47f2dfa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,6 +104,8 @@ module = [ "datafusion.*", "xarray.*", "pandas.*", + "cftime.*", + "xarray_sql._native", ] ignore_missing_imports = true diff --git a/src/lib.rs b/src/lib.rs index 63a5a6bd..1111a1b2 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -47,7 +47,7 @@ use std::fmt::Debug; use std::sync::Arc; use arrow::array::RecordBatch; -use arrow::datatypes::{DataType, Schema, SchemaRef, TimeUnit}; +use arrow::datatypes::{DataType, IntervalMonthDayNano, IntervalUnit, Schema, SchemaRef, TimeUnit}; use arrow::pyarrow::{FromPyArrow, ToPyArrow}; use async_stream::try_stream; use async_trait::async_trait; @@ -70,6 +70,7 @@ use datafusion::physical_plan::{ }; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::table_provider::FFI_TableProvider; +use half::f16; use pyo3::prelude::*; use pyo3::types::{PyCapsule, PyList}; @@ -77,50 +78,218 @@ use pyo3::types::{PyCapsule, PyList}; // 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. +/// Scalar value for orderable xarray dimension bounds. #[derive(Clone, Debug)] pub enum ScalarBound { - /// 64-bit integer (for integer coordinates) + Boolean(bool), + Int8(i8), + Int16(i16), + Int32(i32), Int64(i64), - /// 64-bit float (for lat/lon coordinates) + UInt8(u8), + UInt16(u16), + UInt32(u32), + UInt64(u64), + Float16(f16), + Float32(f32), Float64(f64), - /// Nanoseconds since Unix epoch (for datetime64[ns] coordinates) + Utf8(String), + Binary(Vec), + TimestampSecond(i64), + TimestampMillisecond(i64), + TimestampMicrosecond(i64), TimestampNanos(i64), + DurationSecond(i64), + DurationMillisecond(i64), + DurationMicrosecond(i64), + DurationNanosecond(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. - // Use checked_mul to avoid silent overflow in release builds; - // on overflow return None (conservative: include the partition). - (ScalarBound::TimestampNanos(a), ScalarValue::TimestampNanosecond(Some(b), _)) => { - Some(a.cmp(b)) - } - (ScalarBound::TimestampNanos(a), ScalarValue::TimestampMicrosecond(Some(b), _)) => { - b.checked_mul(1_000).map(|b_ns| a.cmp(&b_ns)) + if let (Some(a), Some(b)) = (integer_bound(self), integer_scalar(scalar)) { + return Some(compare_integers(a, b)); + } + if let (Some(a), Some(b)) = (float_bound(self), float_scalar(scalar)) { + return a.partial_cmp(&b); + } + if let (Some(a), Some(b)) = (timestamp_bound(self), timestamp_scalar(scalar)) { + return Some(a.cmp(&b)); + } + if let (Some(a), ScalarValue::IntervalMonthDayNano(Some(b))) = + (duration_bound(self), scalar) + { + if b.months != 0 { + return None; } - (ScalarBound::TimestampNanos(a), ScalarValue::TimestampMillisecond(Some(b), _)) => { - b.checked_mul(1_000_000).map(|b_ns| a.cmp(&b_ns)) + let a = i64::try_from(a).ok()?; + return Some(IntervalMonthDayNano::new(0, 0, a).cmp(b)); + } + if let (Some(a), Some(b)) = (duration_bound(self), duration_scalar(scalar)) { + return Some(a.cmp(&b)); + } + + match (self, scalar) { + (ScalarBound::Boolean(a), ScalarValue::Boolean(Some(b))) => Some(a.cmp(b)), + (ScalarBound::Utf8(a), ScalarValue::Utf8(Some(b))) + | (ScalarBound::Utf8(a), ScalarValue::Utf8View(Some(b))) + | (ScalarBound::Utf8(a), ScalarValue::LargeUtf8(Some(b))) => Some(a.cmp(b)), + (ScalarBound::Binary(a), ScalarValue::Binary(Some(b))) + | (ScalarBound::Binary(a), ScalarValue::BinaryView(Some(b))) + | (ScalarBound::Binary(a), ScalarValue::LargeBinary(Some(b))) + | (ScalarBound::Binary(a), ScalarValue::FixedSizeBinary(_, Some(b))) => Some(a.cmp(b)), + _ => None, + } + } + + fn compare_to_bound(&self, other: &ScalarBound) -> Option { + if let (Some(a), Some(b)) = (integer_bound(self), integer_bound(other)) { + return Some(compare_integers(a, b)); + } + if let (Some(a), Some(b)) = (float_bound(self), float_bound(other)) { + return a.partial_cmp(&b); + } + if let (Some(a), Some(b)) = (timestamp_bound(self), timestamp_bound(other)) { + return Some(a.cmp(&b)); + } + if let (Some(a), Some(b)) = (duration_bound(self), duration_bound(other)) { + return Some(a.cmp(&b)); + } + match (self, other) { + (ScalarBound::Boolean(a), ScalarBound::Boolean(b)) => Some(a.cmp(b)), + (ScalarBound::Utf8(a), ScalarBound::Utf8(b)) => Some(a.cmp(b)), + (ScalarBound::Binary(a), ScalarBound::Binary(b)) => Some(a.cmp(b)), + _ => None, + } + } +} + +#[derive(Clone, Copy)] +enum IntegerBound { + Signed(i128), + Unsigned(u128), +} + +fn compare_integers(a: IntegerBound, b: IntegerBound) -> std::cmp::Ordering { + match (a, b) { + (IntegerBound::Signed(a), IntegerBound::Signed(b)) => a.cmp(&b), + (IntegerBound::Unsigned(a), IntegerBound::Unsigned(b)) => a.cmp(&b), + (IntegerBound::Signed(a), IntegerBound::Unsigned(b)) => { + if a < 0 { + std::cmp::Ordering::Less + } else { + (a as u128).cmp(&b) } - (ScalarBound::TimestampNanos(a), ScalarValue::TimestampSecond(Some(b), _)) => { - b.checked_mul(1_000_000_000).map(|b_ns| a.cmp(&b_ns)) + } + (IntegerBound::Unsigned(a), IntegerBound::Signed(b)) => { + if b < 0 { + std::cmp::Ordering::Greater + } else { + a.cmp(&(b as u128)) } + } + } +} - // Incompatible types - _ => None, +fn integer_bound(bound: &ScalarBound) -> Option { + Some(match bound { + ScalarBound::Int8(v) => IntegerBound::Signed(*v as i128), + ScalarBound::Int16(v) => IntegerBound::Signed(*v as i128), + ScalarBound::Int32(v) => IntegerBound::Signed(*v as i128), + ScalarBound::Int64(v) => IntegerBound::Signed(*v as i128), + ScalarBound::UInt8(v) => IntegerBound::Unsigned(*v as u128), + ScalarBound::UInt16(v) => IntegerBound::Unsigned(*v as u128), + ScalarBound::UInt32(v) => IntegerBound::Unsigned(*v as u128), + ScalarBound::UInt64(v) => IntegerBound::Unsigned(*v as u128), + _ => return None, + }) +} + +fn integer_scalar(scalar: &ScalarValue) -> Option { + Some(match scalar { + ScalarValue::Int8(Some(v)) => IntegerBound::Signed(*v as i128), + ScalarValue::Int16(Some(v)) => IntegerBound::Signed(*v as i128), + ScalarValue::Int32(Some(v)) => IntegerBound::Signed(*v as i128), + ScalarValue::Int64(Some(v)) => IntegerBound::Signed(*v as i128), + ScalarValue::UInt8(Some(v)) => IntegerBound::Unsigned(*v as u128), + ScalarValue::UInt16(Some(v)) => IntegerBound::Unsigned(*v as u128), + ScalarValue::UInt32(Some(v)) => IntegerBound::Unsigned(*v as u128), + ScalarValue::UInt64(Some(v)) => IntegerBound::Unsigned(*v as u128), + _ => return None, + }) +} + +fn float_bound(bound: &ScalarBound) -> Option { + match bound { + ScalarBound::Float16(v) => Some(v.to_f64()), + ScalarBound::Float32(v) => Some(*v as f64), + ScalarBound::Float64(v) => Some(*v), + _ => None, + } +} + +fn float_scalar(scalar: &ScalarValue) -> Option { + match scalar { + ScalarValue::Float16(Some(v)) => Some(v.to_f64()), + ScalarValue::Float32(Some(v)) => Some(*v as f64), + ScalarValue::Float64(Some(v)) => Some(*v), + _ => None, + } +} + +fn temporal_nanos(value: i64, unit: TimeUnit) -> i128 { + let factor = match unit { + TimeUnit::Second => 1_000_000_000, + TimeUnit::Millisecond => 1_000_000, + TimeUnit::Microsecond => 1_000, + TimeUnit::Nanosecond => 1, + }; + i128::from(value) * factor +} + +fn timestamp_bound(bound: &ScalarBound) -> Option { + Some(match bound { + ScalarBound::TimestampSecond(v) => temporal_nanos(*v, TimeUnit::Second), + ScalarBound::TimestampMillisecond(v) => temporal_nanos(*v, TimeUnit::Millisecond), + ScalarBound::TimestampMicrosecond(v) => temporal_nanos(*v, TimeUnit::Microsecond), + ScalarBound::TimestampNanos(v) => temporal_nanos(*v, TimeUnit::Nanosecond), + _ => return None, + }) +} + +fn timestamp_scalar(scalar: &ScalarValue) -> Option { + Some(match scalar { + ScalarValue::TimestampSecond(Some(v), _) => temporal_nanos(*v, TimeUnit::Second), + ScalarValue::TimestampMillisecond(Some(v), _) => temporal_nanos(*v, TimeUnit::Millisecond), + ScalarValue::TimestampMicrosecond(Some(v), _) => temporal_nanos(*v, TimeUnit::Microsecond), + ScalarValue::TimestampNanosecond(Some(v), _) => temporal_nanos(*v, TimeUnit::Nanosecond), + _ => return None, + }) +} + +fn duration_bound(bound: &ScalarBound) -> Option { + Some(match bound { + ScalarBound::DurationSecond(v) => temporal_nanos(*v, TimeUnit::Second), + ScalarBound::DurationMillisecond(v) => temporal_nanos(*v, TimeUnit::Millisecond), + ScalarBound::DurationMicrosecond(v) => temporal_nanos(*v, TimeUnit::Microsecond), + ScalarBound::DurationNanosecond(v) => temporal_nanos(*v, TimeUnit::Nanosecond), + _ => return None, + }) +} + +fn duration_scalar(scalar: &ScalarValue) -> Option { + match scalar { + ScalarValue::DurationSecond(Some(v)) => Some(temporal_nanos(*v, TimeUnit::Second)), + ScalarValue::DurationMillisecond(Some(v)) => { + Some(temporal_nanos(*v, TimeUnit::Millisecond)) } + ScalarValue::DurationMicrosecond(Some(v)) => { + Some(temporal_nanos(*v, TimeUnit::Microsecond)) + } + ScalarValue::DurationNanosecond(Some(v)) => Some(temporal_nanos(*v, TimeUnit::Nanosecond)), + _ => None, } } @@ -129,9 +298,9 @@ impl ScalarBound { pub struct DimensionRange { /// The column name (dimension name from xarray) pub column_name: String, - /// Minimum value (inclusive) - first coordinate value in this partition + /// Minimum coordinate value in this partition (inclusive). pub min: ScalarBound, - /// Maximum value (inclusive) - last coordinate value in this partition + /// Maximum coordinate value in this partition (inclusive). pub max: ScalarBound, } @@ -223,6 +392,13 @@ impl PrunableStreamingTable { /// Conservative: returns false (include) if uncertain. fn filter_excludes_partition(&self, expr: &Expr, meta: &PartitionMetadata) -> bool { match expr { + Expr::Column(_) => self.boolean_filter_excludes(expr, true, meta), + Expr::IsTrue(inner) | Expr::IsNotFalse(inner) => { + self.boolean_filter_excludes(inner, true, meta) + } + Expr::IsFalse(inner) | Expr::IsNotTrue(inner) => { + self.boolean_filter_excludes(inner, false, meta) + } Expr::BinaryExpr(BinaryExpr { left, op, right }) => { // Handle AND/OR logic match op { @@ -240,7 +416,10 @@ impl PrunableStreamingTable { _ => self.comparison_excludes(left, op, right, meta), } } - Expr::Not(_) => { + Expr::Not(inner) => { + if self.pruning_column_name(inner).is_some() { + return self.boolean_filter_excludes(inner, false, meta); + } // NOT inverts the predicate. We cannot safely derive exclusion // from the inner result: if inner returns false (uncertain), // !false = true would incorrectly exclude the partition. @@ -266,10 +445,16 @@ impl PrunableStreamingTable { meta: &PartitionMetadata, ) -> bool { // Try to extract column and literal from either side - 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 + let (col_name, scalar, flipped) = if let (Some(column), Expr::Literal(scalar, _)) = + (self.pruning_column_name(left), right) + { + (column, scalar, false) + } else if let (Expr::Literal(scalar, _), Some(column)) = + (left, self.pruning_column_name(right)) + { + (column, scalar, true) + } else { + return false; }; // Get the dimension range for this column @@ -277,6 +462,9 @@ impl PrunableStreamingTable { Some(r) => r, None => return false, // Not a dimension column, can't prune }; + if !range_can_compare_scalar(range, scalar) { + return false; + } // Flip operator if literal was on left side let effective_op = if flipped { flip_operator(op) } else { *op }; @@ -349,9 +537,9 @@ impl PrunableStreamingTable { } // Extract column name - let col_name = match between.expr.as_ref() { - Expr::Column(c) => c.name.clone(), - _ => return false, + let col_name = match self.pruning_column_name(&between.expr) { + Some(name) => name, + None => return false, }; // Get dimension range @@ -365,6 +553,9 @@ impl PrunableStreamingTable { (Expr::Literal(l, _), Expr::Literal(h, _)) => (l, h), _ => return false, }; + if !range_can_compare_scalar(range, low) || !range_can_compare_scalar(range, high) { + return false; + } // Exclude if partition range doesn't overlap with [low, high] // No overlap if: partition.max < low OR partition.min > high @@ -388,9 +579,9 @@ impl PrunableStreamingTable { } // Extract column name - let col_name = match in_list.expr.as_ref() { - Expr::Column(c) => c.name.clone(), - _ => return false, + let col_name = match self.pruning_column_name(&in_list.expr) { + Some(name) => name, + None => return false, }; // Get dimension range @@ -402,16 +593,20 @@ impl PrunableStreamingTable { // 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 { + if !range_can_compare_scalar(range, scalar) { + return true; + } // Value is in range if: min <= value <= max - let above_min = matches!( + match ( 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 + ) { + (Some(min_cmp), Some(max_cmp)) => { + min_cmp != std::cmp::Ordering::Greater + && max_cmp != std::cmp::Ordering::Less + } + _ => true, + } } else { // Non-literal in list, be conservative true @@ -425,6 +620,11 @@ impl PrunableStreamingTable { /// Check if an expression is a filter on a dimension column. fn is_dimension_filter(&self, expr: &Expr) -> bool { match expr { + Expr::Column(_) => self.pruning_column_name(expr).is_some(), + Expr::IsTrue(inner) + | Expr::IsFalse(inner) + | Expr::IsNotTrue(inner) + | Expr::IsNotFalse(inner) => self.expr_references_dimension(inner), Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op { Operator::And | Operator::Or => { self.is_dimension_filter(left) || self.is_dimension_filter(right) @@ -440,11 +640,124 @@ impl PrunableStreamingTable { /// Check if an expression references a dimension column. fn expr_references_dimension(&self, expr: &Expr) -> bool { + self.pruning_column_name(expr).is_some() + } + + /// Return the dimension column behind an order-preserving expression. + fn pruning_column_name(&self, expr: &Expr) -> Option { match expr { - Expr::Column(c) => self.dimension_columns.contains(&c.name), - _ => false, + Expr::Column(column) if self.dimension_columns.contains(&column.name) => { + Some(column.name.clone()) + } + Expr::Cast(cast) => { + let column = self.pruning_column_name(&cast.expr)?; + let source = self.schema.field_with_name(&column).ok()?.data_type(); + order_preserving_cast(source, cast.field.data_type()).then_some(column) + } + Expr::TryCast(cast) => { + let column = self.pruning_column_name(&cast.expr)?; + let source = self.schema.field_with_name(&column).ok()?.data_type(); + order_preserving_cast(source, cast.field.data_type()).then_some(column) + } + _ => None, } } + + /// A bare boolean column means `column = TRUE`; `NOT column` is handled + /// by the caller as `column = FALSE`. + fn boolean_filter_excludes( + &self, + expr: &Expr, + expected: bool, + meta: &PartitionMetadata, + ) -> bool { + let Some(column) = self.pruning_column_name(expr) else { + return false; + }; + let Some(range) = meta.get_range(&column) else { + return false; + }; + let literal = ScalarValue::Boolean(Some(expected)); + if expected { + matches!( + range.max.compare_to_scalar(&literal), + Some(std::cmp::Ordering::Less) + ) + } else { + matches!( + range.min.compare_to_scalar(&literal), + Some(std::cmp::Ordering::Greater) + ) + } + } +} + +fn signed_integer_width(dtype: &DataType) -> Option { + match dtype { + DataType::Int8 => Some(8), + DataType::Int16 => Some(16), + DataType::Int32 => Some(32), + DataType::Int64 => Some(64), + _ => None, + } +} + +fn unsigned_integer_width(dtype: &DataType) -> Option { + match dtype { + DataType::UInt8 => Some(8), + DataType::UInt16 => Some(16), + DataType::UInt32 => Some(32), + DataType::UInt64 => Some(64), + _ => None, + } +} + +/// Whether a cast preserves ordering and every source value exactly. +fn order_preserving_cast(source: &DataType, target: &DataType) -> bool { + if source == target { + return true; + } + if let (Some(source_width), Some(target_width)) = + (signed_integer_width(source), signed_integer_width(target)) + { + return target_width >= source_width; + } + if let (Some(source_width), Some(target_width)) = ( + unsigned_integer_width(source), + unsigned_integer_width(target), + ) { + return target_width >= source_width; + } + if let (Some(source_width), Some(target_width)) = + (unsigned_integer_width(source), signed_integer_width(target)) + { + return target_width > source_width; + } + matches!( + (source, target), + (DataType::Float16, DataType::Float32 | DataType::Float64) + | (DataType::Float32, DataType::Float64) + | ( + DataType::Duration(_), + DataType::Interval(IntervalUnit::MonthDayNano) + ) + ) +} + +/// Duration-to-interval casts use an i64 nanosecond field in Arrow. If either +/// endpoint overflows it, retaining the partition is required: pruning it could +/// otherwise hide a row-level cast error (or a null from TRY_CAST). +fn range_can_compare_scalar(range: &DimensionRange, scalar: &ScalarValue) -> bool { + let ScalarValue::IntervalMonthDayNano(Some(interval)) = scalar else { + return true; + }; + if interval.months != 0 { + return false; + } + match (duration_bound(&range.min), duration_bound(&range.max)) { + (Some(min), Some(max)) => i64::try_from(min).is_ok() && i64::try_from(max).is_ok(), + _ => true, + } } /// Extension trait for partition streams that support column projection. @@ -479,18 +792,28 @@ fn flip_operator(op: &Operator) -> Operator { /// Convert a Python object to a ScalarBound using an explicit dtype tag. fn python_to_scalar_bound(obj: &Bound<'_, PyAny>, dtype_tag: &str) -> PyResult { match dtype_tag { - "timestamp_ns" => { - let val = obj.extract::()?; - Ok(ScalarBound::TimestampNanos(val)) - } - "float64" => { - let val = obj.extract::()?; - Ok(ScalarBound::Float64(val)) - } - "int64" => { - let val = obj.extract::()?; - Ok(ScalarBound::Int64(val)) - } + "bool" => Ok(ScalarBound::Boolean(obj.extract::()?)), + "int8" => Ok(ScalarBound::Int8(obj.extract::()?)), + "int16" => Ok(ScalarBound::Int16(obj.extract::()?)), + "int32" => Ok(ScalarBound::Int32(obj.extract::()?)), + "int64" => Ok(ScalarBound::Int64(obj.extract::()?)), + "uint8" => Ok(ScalarBound::UInt8(obj.extract::()?)), + "uint16" => Ok(ScalarBound::UInt16(obj.extract::()?)), + "uint32" => Ok(ScalarBound::UInt32(obj.extract::()?)), + "uint64" => Ok(ScalarBound::UInt64(obj.extract::()?)), + "float16" => Ok(ScalarBound::Float16(f16::from_f64(obj.extract::()?))), + "float32" => Ok(ScalarBound::Float32(obj.extract::()?)), + "float64" => Ok(ScalarBound::Float64(obj.extract::()?)), + "utf8" => Ok(ScalarBound::Utf8(obj.extract::()?)), + "binary" => Ok(ScalarBound::Binary(obj.extract::>()?)), + "timestamp_s" => Ok(ScalarBound::TimestampSecond(obj.extract::()?)), + "timestamp_ms" => Ok(ScalarBound::TimestampMillisecond(obj.extract::()?)), + "timestamp_us" => Ok(ScalarBound::TimestampMicrosecond(obj.extract::()?)), + "timestamp_ns" => Ok(ScalarBound::TimestampNanos(obj.extract::()?)), + "duration_s" => Ok(ScalarBound::DurationSecond(obj.extract::()?)), + "duration_ms" => Ok(ScalarBound::DurationMillisecond(obj.extract::()?)), + "duration_us" => Ok(ScalarBound::DurationMicrosecond(obj.extract::()?)), + "duration_ns" => Ok(ScalarBound::DurationNanosecond(obj.extract::()?)), _ => Err(pyo3::exceptions::PyTypeError::new_err(format!( "Unsupported dtype tag for partition bound: {dtype_tag}" ))), @@ -682,12 +1005,7 @@ fn sum_row_counts<'a>(metas: impl Iterator) -> Pre /// larger one. Returns `None` if the variants differ (never expected within a /// single dimension) so the caller can fall back to unknown. fn fold_bound(a: &ScalarBound, b: &ScalarBound, keep_min: bool) -> Option { - let ord = match (a, b) { - (ScalarBound::Int64(x), ScalarBound::Int64(y)) => x.partial_cmp(y), - (ScalarBound::Float64(x), ScalarBound::Float64(y)) => x.partial_cmp(y), - (ScalarBound::TimestampNanos(x), ScalarBound::TimestampNanos(y)) => x.partial_cmp(y), - _ => return None, - }?; + let ord = a.compare_to_bound(b)?; let take_a = if keep_min { ord != std::cmp::Ordering::Greater } else { @@ -703,24 +1021,37 @@ fn fold_bound(a: &ScalarBound, b: &ScalarBound, keep_min: bool) -> Option Option { match (bound, dtype) { + (ScalarBound::Boolean(v), DataType::Boolean) => Some(ScalarValue::Boolean(Some(*v))), + (ScalarBound::Int8(v), DataType::Int8) => Some(ScalarValue::Int8(Some(*v))), + (ScalarBound::Int16(v), DataType::Int16) => Some(ScalarValue::Int16(Some(*v))), + (ScalarBound::Int32(v), DataType::Int32) => Some(ScalarValue::Int32(Some(*v))), (ScalarBound::Int64(v), DataType::Int64) => Some(ScalarValue::Int64(Some(*v))), - (ScalarBound::Int64(v), DataType::Int32) => { - i32::try_from(*v).ok().map(|x| ScalarValue::Int32(Some(x))) - } + (ScalarBound::UInt8(v), DataType::UInt8) => Some(ScalarValue::UInt8(Some(*v))), + (ScalarBound::UInt16(v), DataType::UInt16) => Some(ScalarValue::UInt16(Some(*v))), + (ScalarBound::UInt32(v), DataType::UInt32) => Some(ScalarValue::UInt32(Some(*v))), + (ScalarBound::UInt64(v), DataType::UInt64) => Some(ScalarValue::UInt64(Some(*v))), + (ScalarBound::Float16(v), DataType::Float16) => Some(ScalarValue::Float16(Some(*v))), + (ScalarBound::Float32(v), DataType::Float32) => Some(ScalarValue::Float32(Some(*v))), (ScalarBound::Float64(v), DataType::Float64) => Some(ScalarValue::Float64(Some(*v))), - (ScalarBound::Float64(v), DataType::Float32) => Some(ScalarValue::Float32(Some(*v as f32))), - // Datetime coordinates arrive as nanoseconds (see `cftime.partition_bounds` - // and the datetime64[ns] path in `_block_metadata`). Map them onto the - // column's own timestamp unit, but only when the scaling is exact so a - // reported bound is never a rounded value. - (ScalarBound::TimestampNanos(v), DataType::Timestamp(unit, tz)) => { - let scaled = match unit { - TimeUnit::Nanosecond => Some(*v), - TimeUnit::Microsecond if v % 1_000 == 0 => Some(v / 1_000), - TimeUnit::Millisecond if v % 1_000_000 == 0 => Some(v / 1_000_000), - TimeUnit::Second if v % 1_000_000_000 == 0 => Some(v / 1_000_000_000), - _ => None, - }?; + (ScalarBound::Utf8(v), DataType::Utf8) => Some(ScalarValue::Utf8(Some(v.clone()))), + (ScalarBound::Utf8(v), DataType::Utf8View) => Some(ScalarValue::Utf8View(Some(v.clone()))), + (ScalarBound::Utf8(v), DataType::LargeUtf8) => { + Some(ScalarValue::LargeUtf8(Some(v.clone()))) + } + (ScalarBound::Binary(v), DataType::Binary) => Some(ScalarValue::Binary(Some(v.clone()))), + (ScalarBound::Binary(v), DataType::BinaryView) => { + Some(ScalarValue::BinaryView(Some(v.clone()))) + } + (ScalarBound::Binary(v), DataType::LargeBinary) => { + Some(ScalarValue::LargeBinary(Some(v.clone()))) + } + (ScalarBound::Binary(v), DataType::FixedSizeBinary(size)) + if usize::try_from(*size).ok() == Some(v.len()) => + { + Some(ScalarValue::FixedSizeBinary(*size, Some(v.clone()))) + } + (bound, DataType::Timestamp(unit, tz)) if timestamp_bound(bound).is_some() => { + let scaled = temporal_bound_in_unit(timestamp_bound(bound)?, unit)?; Some(match unit { TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(scaled), tz.clone()), TimeUnit::Microsecond => { @@ -732,10 +1063,32 @@ fn bound_to_scalar(bound: &ScalarBound, dtype: &DataType) -> Option TimeUnit::Second => ScalarValue::TimestampSecond(Some(scaled), tz.clone()), }) } + (bound, DataType::Duration(unit)) if duration_bound(bound).is_some() => { + let scaled = temporal_bound_in_unit(duration_bound(bound)?, unit)?; + Some(match unit { + TimeUnit::Second => ScalarValue::DurationSecond(Some(scaled)), + TimeUnit::Millisecond => ScalarValue::DurationMillisecond(Some(scaled)), + TimeUnit::Microsecond => ScalarValue::DurationMicrosecond(Some(scaled)), + TimeUnit::Nanosecond => ScalarValue::DurationNanosecond(Some(scaled)), + }) + } _ => None, } } +fn temporal_bound_in_unit(nanoseconds: i128, unit: &TimeUnit) -> Option { + let divisor = match unit { + TimeUnit::Second => 1_000_000_000, + TimeUnit::Millisecond => 1_000_000, + TimeUnit::Microsecond => 1_000, + TimeUnit::Nanosecond => 1, + }; + if nanoseconds % divisor != 0 { + return None; + } + i64::try_from(nanoseconds / divisor).ok() +} + /// Exact in-memory byte size of `num_rows` rows of `schema`, or `Absent` if any /// column is variable-width (e.g. Utf8) and cannot be sized from the row count /// alone. Our data model is dense fixed-width grids, so this is normally exact. @@ -777,18 +1130,31 @@ fn build_scan_statistics(output_schema: &Schema, metas: &[&PartitionMetadata]) - // with a representable bound; all such partitions share the same bound // variant, so the fold is well-defined. let mut folded: Option<(ScalarBound, ScalarBound)> = None; + let mut complete = !metas.is_empty(); for meta in metas { - if let Some(range) = meta.ranges.get(field.name()) { - folded = Some(match folded { - None => (range.min.clone(), range.max.clone()), - Some((lo, hi)) => ( - fold_bound(&lo, &range.min, true).unwrap_or(lo), - fold_bound(&hi, &range.max, false).unwrap_or(hi), - ), - }); - } + let Some(range) = meta.ranges.get(field.name()) else { + complete = false; + break; + }; + folded = match folded.take() { + None => Some((range.min.clone(), range.max.clone())), + Some((lo, hi)) => { + let Some(lo) = fold_bound(&lo, &range.min, true) else { + complete = false; + break; + }; + let Some(hi) = fold_bound(&hi, &range.max, false) else { + complete = false; + break; + }; + Some((lo, hi)) + } + }; } + if !complete { + continue; + } let Some((lo, hi)) = folded else { continue }; // This column is a coordinate axis: never null, so the null count is // exactly zero regardless of whether the bound maps to a ScalarValue. @@ -1173,6 +1539,12 @@ impl LazyArrowStreamTable { /// - ``metadata_dict`` is a ``dict[str, tuple[Any, Any, str]]`` /// mapping dimension name to ``(min, max, dtype_str)``; pass /// ``{}`` to skip pruning for a partition. + /// Supported tags are ``bool``, exact-width ``int*``, + /// ``uint*``, and ``float*`` tags, ``utf8``, ``binary``, + /// ``timestamp_s|ms|us|ns``, and + /// ``duration_s|ms|us|ns``. Temporal bounds are signed + /// integer counts in the tagged unit. Invalid tags are a + /// protocol error and raise ``TypeError``. /// - ``num_rows`` is the exact row count for the partition, so /// the scan reports exact ``Statistics`` to the optimizer. /// Generators are accepted, so partition state can be produced diff --git a/tests/test_df.py b/tests/test_df.py index 5185b573..099e3f73 100644 --- a/tests/test_df.py +++ b/tests/test_df.py @@ -1,4 +1,5 @@ import tracemalloc +from decimal import Decimal import numpy as np import pandas as pd @@ -612,11 +613,95 @@ def test_parse_schema_maps_object_string_coord_to_string(): assert _field_type(schema, "station") == pa.string() -def test_partition_metadata_skips_out_of_ns_datetime(): - # datetime64 coordinates outside the datetime64[ns] range (pre-1678 / - # post-2262) cannot be represented as int64 nanoseconds, so partition - # pruning must be skipped for that dimension rather than raising - # OverflowError. Registration must still succeed. +@pytest.mark.parametrize( + ("dtype", "values", "expected"), + [ + ("bool", [True, False, True], (False, True, "bool")), + ("int8", [-2, 3, 1], (-2, 3, "int8")), + ("int16", [-2, 3, 1], (-2, 3, "int16")), + ("int32", [-2, 3, 1], (-2, 3, "int32")), + ("int64", [-2, 3, 1], (-2, 3, "int64")), + ("uint8", [3, 0, 2], (0, 3, "uint8")), + ("uint16", [3, 0, 2], (0, 3, "uint16")), + ("uint32", [3, 0, 2], (0, 3, "uint32")), + ( + "uint64", + [2**63 + 1, 2**63 + 3, 2**63 + 2], + (2**63 + 1, 2**63 + 3, "uint64"), + ), + ("float16", [3.5, -1.5, 2.0], (-1.5, 3.5, "float16")), + ("float32", [3.5, -1.5, 2.0], (-1.5, 3.5, "float32")), + ("float64", [3.5, -1.5, 2.0], (-1.5, 3.5, "float64")), + ("U8", ["zulu", "alpha", "echo"], ("alpha", "zulu", "utf8")), + ("S8", [b"zulu", b"alpha", b"echo"], (b"alpha", b"zulu", "binary")), + ], +) +def test_partition_metadata_preserves_coordinate_dtype(dtype, values, expected): + coord = np.asarray(values, dtype=dtype) + ds = xr.Dataset({"v": (["x"], np.arange(len(coord)))}, coords={"x": coord}) + blocks = list(block_slices(ds, chunks={"x": len(coord)})) + + assert partition_metadata(ds, blocks)[0]["x"] == expected + + +@pytest.mark.parametrize("kind", ["datetime64", "timedelta64"]) +@pytest.mark.parametrize("unit", ["s", "ms", "us", "ns"]) +def test_partition_metadata_preserves_temporal_unit(kind, unit): + coord = np.asarray([-2, 3, 1], dtype=f"{kind}[{unit}]") + ds = xr.Dataset({"v": (["x"], np.arange(len(coord)))}, coords={"x": coord}) + blocks = list(block_slices(ds, chunks={"x": len(coord)})) + prefix = "timestamp" if kind == "datetime64" else "duration" + + assert partition_metadata(ds, blocks)[0]["x"] == ( + -2, + 3, + f"{prefix}_{unit}", + ) + + +@pytest.mark.parametrize( + ("categories", "arrow_type", "bounds"), + [ + (["zulu", "alpha", "echo"], pa.string(), ("alpha", "zulu", "utf8")), + ([30, 10, 20], pa.int64(), (10, 30, "int64")), + ], +) +def test_categorical_coordinate_uses_emitted_value_type( + categories, arrow_type, bounds +): + coord = pd.CategoricalIndex(categories, ordered=True, name="category") + ds = xr.Dataset( + {"v": (["category"], np.arange(len(coord)))}, + coords={"category": coord}, + ) + blocks = list(block_slices(ds, chunks={"category": len(coord)})) + + assert _field_type(_parse_schema(ds), "category") == arrow_type + assert partition_metadata(ds, blocks)[0]["category"] == bounds + + +@pytest.mark.parametrize( + "coord", + [ + np.asarray([1.0, np.nan]), + np.asarray(["valid", None], dtype=object), + np.asarray([Decimal("1.0"), Decimal("2.0")], dtype=object), + np.asarray( + [np.datetime64("NaT"), np.datetime64("2000-01-01")], + dtype="datetime64[ns]", + ), + ], +) +def test_partition_metadata_omits_unsafe_coordinate_bounds(coord): + ds = xr.Dataset({"v": (["x"], np.arange(len(coord)))}, coords={"x": coord}) + blocks = list(block_slices(ds, chunks={"x": len(coord)})) + + assert "x" not in partition_metadata(ds, blocks)[0] + + +def test_partition_metadata_preserves_out_of_ns_datetime_unit(): + # A microsecond coordinate outside the datetime64[ns] range remains + # representable and prunable when metadata retains its original unit. times = xr.date_range( "0001-01-01", periods=3, freq="100YS", use_cftime=True ).to_datetimeindex(time_unit="us", unsafe=True) @@ -625,11 +710,10 @@ def test_partition_metadata_skips_out_of_ns_datetime(): ) blocks = list(block_slices(ds, chunks={"time": 2})) - meta = partition_metadata(ds, blocks) # must not raise + meta = partition_metadata(ds, blocks) assert len(meta) == len(blocks) - # "time" is unpruneable here, so it is omitted from every partition. - assert all("time" not in m for m in meta) + assert all(m["time"][2] == "timestamp_us" for m in meta) def test_parse_schema_all_null_object_var_stays_null(): @@ -703,8 +787,8 @@ def test_string_dataset_round_trips_through_record_batch(): def test_partition_metadata_in_range_datetime_still_pruned(): - # Regression guard: ordinary datetimes must keep producing timestamp_ns - # bounds so filter pushdown still works after the overflow fix. + # Regression guard: ordinary datetimes retain their source unit and keep + # producing bounds so filter pushdown still works. times = pd.date_range("2000-01-01", periods=4, freq="D") ds = _ensure_default_indexes( xr.Dataset({"v": (["time"], np.arange(4.0))}, coords={"time": times}) @@ -714,9 +798,8 @@ def test_partition_metadata_in_range_datetime_still_pruned(): meta = partition_metadata(ds, blocks) assert all("time" in m for m in meta) - for m in meta: - _, _, tag = m["time"] - assert tag == "timestamp_ns" + expected_unit = np.datetime_data(ds.coords["time"].dtype)[0] + assert all(m["time"][2] == f"timestamp_{expected_unit}" for m in meta) class TestGroupVarsByDims: diff --git a/tests/test_duckdb_backend.py b/tests/test_duckdb_backend.py index 0a004502..26f9204b 100644 --- a/tests/test_duckdb_backend.py +++ b/tests/test_duckdb_backend.py @@ -359,7 +359,7 @@ def test_nan_coordinate_chunk_is_not_pruned(): def test_cftime_dataset_aggregates_under_projection(): - cftime = pytest.importorskip("cftime") + pytest.importorskip("cftime") times = xr.date_range( "2000-01-01", periods=6, calendar="360_day", use_cftime=True diff --git a/tests/test_reader.py b/tests/test_reader.py index 153a9f3b..9ef6d373 100644 --- a/tests/test_reader.py +++ b/tests/test_reader.py @@ -16,6 +16,7 @@ import threading import time +from decimal import Decimal import numpy as np import pandas as pd import pyarrow as pa @@ -1029,6 +1030,273 @@ def test_time_lt_filter_prunes_late_partitions(self, time_chunked_ds): count = result["cnt"].iloc[0] assert count == 125, f"Expected 125 rows, got {count}" + @pytest.mark.parametrize( + "dtype", + [ + "int8", + "int16", + "int32", + "int64", + "uint8", + "uint16", + "uint32", + "uint64", + ], + ) + def test_integer_coordinate_types_prune(self, dtype): + tracker = IterationTracker() + coord = np.arange(8, dtype=dtype) + ds = xr.Dataset({"value": (["x"], np.arange(8))}, coords={"x": coord}) + table = read_xarray_table( + ds, chunks={"x": 2}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + result = ctx.sql( + "SELECT x FROM test WHERE x >= 6 ORDER BY x" + ).to_pandas() + + assert result["x"].tolist() == [6, 7] + assert tracker.iteration_count == 1 + + @pytest.mark.parametrize("dtype", ["float32", "float64"]) + def test_float_coordinate_types_prune_with_natural_literal(self, dtype): + tracker = IterationTracker() + coord = np.arange(8, dtype=dtype) + ds = xr.Dataset({"value": (["x"], np.arange(8))}, coords={"x": coord}) + table = read_xarray_table( + ds, chunks={"x": 2}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + result = ctx.sql( + "SELECT x FROM test WHERE x >= 6.0 ORDER BY x" + ).to_pandas() + + assert result["x"].tolist() == [6.0, 7.0] + assert tracker.iteration_count == 1 + + @pytest.mark.parametrize( + ("coord", "predicate"), + [ + (np.asarray(list("abcdefgh")), "x >= 'g'"), + ( + np.asarray([bytes([value]) for value in range(97, 105)]), + "x >= X'67'", + ), + ], + ) + def test_text_and_binary_coordinate_types_prune(self, coord, predicate): + tracker = IterationTracker() + ds = xr.Dataset({"value": (["x"], np.arange(8))}, coords={"x": coord}) + table = read_xarray_table( + ds, chunks={"x": 2}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + result = ctx.sql( + f"SELECT x FROM test WHERE {predicate} ORDER BY x" + ).to_pandas() + + assert len(result) == 2 + assert tracker.iteration_count == 1 + + def test_categorical_string_coordinate_prunes(self): + tracker = IterationTracker() + coord = pd.CategoricalIndex(list("abcdefgh"), ordered=True, name="x") + ds = xr.Dataset({"value": (["x"], np.arange(8))}, coords={"x": coord}) + table = read_xarray_table( + ds, chunks={"x": 2}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + result = ctx.sql( + "SELECT x FROM test WHERE x >= 'g' ORDER BY x" + ).to_pandas() + + assert result["x"].tolist() == ["g", "h"] + assert tracker.iteration_count == 1 + + def test_boolean_optimizer_forms_prune(self): + coord = np.asarray([False] * 4 + [True] * 4) + ds = xr.Dataset( + {"value": (["flag"], np.arange(8))}, coords={"flag": coord} + ) + + for predicate, expected in ( + ("flag = TRUE", True), + ("flag = FALSE", False), + ("flag IS TRUE", True), + ("flag IS FALSE", False), + ): + tracker = IterationTracker() + table = read_xarray_table( + ds, chunks={"flag": 2}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + result = ctx.sql( + f"SELECT flag FROM test WHERE {predicate} ORDER BY flag" + ).to_pandas() + + assert result["flag"].tolist() == [expected] * 4 + assert tracker.iteration_count == 2 + + @pytest.mark.parametrize( + ("unit", "interval"), + [ + ("s", "6 seconds"), + ("ms", "6 milliseconds"), + ("us", "6 microseconds"), + ("ns", "6 nanoseconds"), + ], + ) + def test_duration_coordinate_units_prune_with_interval( + self, unit, interval + ): + tracker = IterationTracker() + coord = np.arange(8).astype(f"timedelta64[{unit}]") + ds = xr.Dataset( + {"value": (["lead_time"], np.arange(8))}, + coords={"lead_time": coord}, + ) + table = read_xarray_table( + ds, chunks={"lead_time": 2}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + result = ctx.sql( + "SELECT lead_time FROM test " + f"WHERE lead_time >= INTERVAL '{interval}' ORDER BY lead_time" + ).to_pandas() + + assert len(result) == 2 + assert tracker.iteration_count == 1 + + def test_negative_duration_literal_prunes(self): + tracker = IterationTracker() + coord = np.arange(-4, 4).astype("timedelta64[us]") + ds = xr.Dataset( + {"value": (["lead_time"], np.arange(8))}, + coords={"lead_time": coord}, + ) + table = read_xarray_table( + ds, chunks={"lead_time": 2}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + result = ctx.sql( + "SELECT lead_time FROM test " + "WHERE lead_time < INTERVAL '-2 microseconds' ORDER BY lead_time" + ).to_pandas() + + assert result["lead_time"].tolist() == [ + pd.Timedelta(microseconds=-4), + pd.Timedelta(microseconds=-3), + ] + assert tracker.iteration_count == 1 + + def test_invalid_metadata_tag_is_rejected(self): + schema = pa.schema([("x", pa.int64())]) + + with pytest.raises(TypeError, match="Unsupported dtype tag"): + LazyArrowStreamTable( + [(lambda: None, {"x": (0, 1, "not-a-dtype")}, 1)], + schema, + ) + + @pytest.mark.parametrize( + ("predicate", "expected", "expected_reads"), + [ + ("'g' <= x", ["g", "h"], 1), + ("x BETWEEN 'g' AND 'h'", ["g", "h"], 1), + ("x IN ('a', 'h')", ["a", "h"], 2), + ], + ) + def test_text_coordinate_filter_shapes_prune( + self, predicate, expected, expected_reads + ): + tracker = IterationTracker() + ds = xr.Dataset( + {"value": (["x"], np.arange(8))}, + coords={"x": np.asarray(list("abcdefgh"))}, + ) + table = read_xarray_table( + ds, chunks={"x": 2}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + result = ctx.sql( + f"SELECT x FROM test WHERE {predicate} ORDER BY x" + ).to_pandas() + + assert result["x"].tolist() == expected + assert tracker.iteration_count == expected_reads + + def test_calendar_month_interval_is_not_pruned(self): + tracker = IterationTracker() + coord = np.arange(8).astype("timedelta64[D]") + ds = xr.Dataset( + {"value": (["lead_time"], np.arange(8))}, + coords={"lead_time": coord}, + ) + table = read_xarray_table( + ds, chunks={"lead_time": 2}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + ctx.sql( + "SELECT lead_time FROM test WHERE lead_time < INTERVAL '1 month'" + ).collect() + + assert tracker.iteration_count == 4 + + def test_duration_cast_overflow_is_not_hidden_by_pruning(self): + tracker = IterationTracker() + coord = np.asarray([np.iinfo(np.int64).max], dtype="timedelta64[s]") + ds = xr.Dataset( + {"value": (["lead_time"], [1])}, coords={"lead_time": coord} + ) + table = read_xarray_table( + ds, chunks={"lead_time": 1}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + with pytest.raises(Exception, match="[Oo]verflow"): + ctx.sql( + "SELECT lead_time FROM test " + "WHERE lead_time < INTERVAL '1 second'" + ).collect() + + assert tracker.iteration_count == 1 + + def test_arrow_representable_unsupported_object_registers_without_pruning( + self, + ): + tracker = IterationTracker() + coord = np.asarray([Decimal("1.0"), Decimal("2.0")], dtype=object) + ds = xr.Dataset({"value": (["x"], [1, 2])}, coords={"x": coord}) + table = read_xarray_table( + ds, chunks={"x": 1}, _iteration_callback=tracker + ) + ctx = SessionContext() + ctx.register_table("test", table) + + result = ctx.sql("SELECT x FROM test ORDER BY x").to_pandas() + + assert result["x"].tolist() == [Decimal("1.0"), Decimal("2.0")] + assert tracker.iteration_count == 2 + def test_time_between_filter_prunes_outside_range(self, time_chunked_ds): """Query with BETWEEN should prune partitions outside the range.""" tracker = IterationTracker() diff --git a/tests/test_sql.py b/tests/test_sql.py index 66c6dd27..78042cfa 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -19,6 +19,37 @@ def test_sanity(air_dataset_small): assert all(col in result.columns for col in ["lat", "lon", "time", "air"]) +def test_timedelta_microseconds_register_filter_and_prune(monkeypatch): + """Duration coordinates use integer bounds and prune through interval SQL.""" + from xarray_sql import sql as sql_module + + blocks_seen = [] + original = sql_module.read_xarray_table + + def tracking_reader(ds, chunks=None, **kwargs): + kwargs["_iteration_callback"] = lambda block, projection_names: ( + blocks_seen.append(block) + ) + return original(ds, chunks, **kwargs) + + monkeypatch.setattr(sql_module, "read_xarray_table", tracking_reader) + lead_time = np.arange(-2, 6).astype("timedelta64[us]") + ds = xr.Dataset( + {"value": (["lead_time"], np.arange(8))}, + coords={"lead_time": lead_time}, + ) + ctx = XarrayContext() + + ctx.from_dataset("forecast", ds, chunks={"lead_time": 2}) + result = ctx.sql( + "SELECT value FROM forecast " + "WHERE lead_time >= INTERVAL '4 microseconds' ORDER BY value" + ).to_pandas() + + assert result["value"].tolist() == [6, 7] + assert len(blocks_seen) == 1 + + def test_aggregation_small(air_dataset_small): ctx = XarrayContext() ctx.from_dataset("air", air_dataset_small) diff --git a/tests/test_stats.py b/tests/test_stats.py index a1f199eb..619c5994 100644 --- a/tests/test_stats.py +++ b/tests/test_stats.py @@ -7,6 +7,7 @@ """ import numpy as np +import pytest import xarray as xr from xarray_sql import XarrayContext @@ -75,6 +76,53 @@ def test_dimension_column_min_max_in_scan_statistics(): assert "Min=Exact(Int64(0)) Max=Exact(Int64(4)) Null=Exact(0)" in plan +@pytest.mark.parametrize( + ("coord", "expected"), + [ + ( + np.arange(4, dtype="uint64") + 2**63, + "Min=Exact(UInt64(9223372036854775808)) " + "Max=Exact(UInt64(9223372036854775811)) Null=Exact(0)", + ), + ( + np.asarray(list("abcd")), + 'Min=Exact(Utf8("a")) Max=Exact(Utf8("d")) Null=Exact(0)', + ), + ( + np.arange(4).astype("timedelta64[us]"), + 'Min=Exact(DurationMicrosecond("0")) ' + 'Max=Exact(DurationMicrosecond("3")) Null=Exact(0)', + ), + ], +) +def test_coordinate_statistics_preserve_logical_type(coord, expected): + ds = xr.Dataset( + {"value": (["x"], np.arange(4))}, + coords={"x": coord}, + ) + ctx = XarrayContext() + ctx.from_dataset("values", ds, chunks={"x": 2}) + + plan = _explain(ctx, "SELECT x, value FROM values") + + assert expected in plan + + +def test_missing_partition_bound_removes_global_coordinate_statistics(): + ds = xr.Dataset( + {"value": (["x"], np.arange(4))}, + coords={"x": np.asarray([0.0, 1.0, np.nan, 3.0])}, + ) + ctx = XarrayContext() + ctx.from_dataset("values", ds, chunks={"x": 2}) + + plan = _explain(ctx, "SELECT x, value FROM values") + + assert "Min=Exact(Float64" not in plan + assert "Max=Exact(Float64" not in plan + assert "Null=Exact(0)" not in plan + + def test_count_star_answered_from_statistics(): """COUNT(*) returns the exact count from statistics (metadata only).""" ds = xr.Dataset( diff --git a/xarray_sql/cftime.py b/xarray_sql/cftime.py index 1f7c49ce..27a6c94c 100644 --- a/xarray_sql/cftime.py +++ b/xarray_sql/cftime.py @@ -20,6 +20,8 @@ from __future__ import annotations +from typing import cast + import numpy as np import pyarrow as pa import xarray as xr @@ -139,7 +141,7 @@ def to_microseconds(values) -> np.ndarray: units=DEFAULT_UNITS, calendar=values.ravel()[0].calendar, ) - return np.asarray(us, dtype=np.float64).astype(np.int64) + return cast(np.ndarray, np.asarray(us, dtype=np.float64).astype(np.int64)) def to_offsets(values, units: str, cal: str) -> np.ndarray: @@ -150,7 +152,7 @@ def to_offsets(values, units: str, cal: str) -> np.ndarray: import cftime as _cftime raw = _cftime.date2num(values.ravel(), units=units, calendar=cal) - return np.asarray(raw, dtype=np.float64).astype(np.int64) + return cast(np.ndarray, np.asarray(raw, dtype=np.float64).astype(np.int64)) def convert_for_field(values, field: pa.Field) -> np.ndarray: diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 83879f77..7ec1573c 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -431,7 +431,7 @@ def iter_record_batches( if name in ds.coords and name in ds.dims: k = dim_names.index(name) outer = int(np.prod(shape[:k])) - col = np.repeat(coord_values[name], strides[k]) + col: np.ndarray = np.repeat(coord_values[name], strides[k]) if outer > 1: col = np.tile(col, outer) full_arrays.append(_as_single_array(col, field.type)) @@ -487,6 +487,20 @@ def _arrow_type_for_object(values: np.ndarray) -> pa.DataType: return pa.array(np.asarray(values).ravel()).type +def _arrow_type_for_categorical(values: np.ndarray) -> pa.DataType: + """Return the Arrow value type emitted for a categorical coordinate. + + xarray preserves a pandas ``CategoricalDtype`` on the coordinate variable, + but ``.values`` exposes the category labels that are actually written to + Arrow. Infer from those labels so string categories become UTF-8 and + numeric categories retain their numeric width. + """ + values = np.asarray(values) + if values.dtype == np.dtype("O"): + return _arrow_type_for_object(values) + return pa.from_numpy_dtype(values.dtype) + + def _parse_schema(ds: xr.Dataset) -> pa.Schema: """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns. @@ -515,6 +529,13 @@ def _parse_schema(ds: xr.Dataset) -> pa.Schema: if cft.is_cftime_index(ds, coord_name): units, calendar = cft.encoding(ds, coord_name) columns.append(cft.arrow_field(coord_name, units, calendar)) + elif isinstance(coord_var.dtype, pd.CategoricalDtype): + columns.append( + pa.field( + coord_name, + _arrow_type_for_categorical(coord_var.values), + ) + ) elif coord_var.dtype == np.dtype("O"): # Object dtype that isn't cftime (e.g. string station names). arrow_type = _arrow_type_for_object(coord_var.values) @@ -549,6 +570,70 @@ def _parse_schema(ds: xr.Dataset) -> pa.Schema: PartitionBounds = dict[str, tuple[Any, Any, str]] +_INTEGER_BOUND_TAGS = { + ("i", 1): "int8", + ("i", 2): "int16", + ("i", 4): "int32", + ("i", 8): "int64", + ("u", 1): "uint8", + ("u", 2): "uint16", + ("u", 4): "uint32", + ("u", 8): "uint64", +} +_FLOAT_BOUND_TAGS = {2: "float16", 4: "float32", 8: "float64"} +_TEMPORAL_BOUND_UNITS = {"s", "ms", "us", "ns"} + + +def _coordinate_partition_bound( + coord_values: np.ndarray, +) -> tuple[Any, Any, str] | None: + """Return an exact, Rust-compatible bound for one coordinate slice. + + Unsupported, null-tainted, or unordered values return ``None``. Partition + pruning is optional, so omitting unsafe metadata is always preferable to + coercing a value or failing dataset registration. + """ + values = np.asarray(coord_values) + if values.size == 0: + return None + + kind = values.dtype.kind + + if kind in ("M", "m"): + if np.isnat(values).any(): + return None + unit, step = np.datetime_data(values.dtype) + if unit not in _TEMPORAL_BOUND_UNITS or step != 1: + return None + raw = values.astype(np.int64, copy=False) + prefix = "timestamp" if kind == "M" else "duration" + return int(raw.min()), int(raw.max()), f"{prefix}_{unit}" + + if kind == "b": + return bool(values.min()), bool(values.max()), "bool" + + if kind in ("i", "u"): + tag = _INTEGER_BOUND_TAGS.get((kind, values.dtype.itemsize)) + if tag is None: + return None + return int(values.min()), int(values.max()), tag + + if kind == "f": + tag = _FLOAT_BOUND_TAGS.get(values.dtype.itemsize) + if tag is None or np.isnan(values).any(): + return None + return float(values.min()), float(values.max()), tag + + if kind in ("U", "S", "O"): + python_values = values.ravel().tolist() + if all(isinstance(value, str) for value in python_values): + return min(python_values), max(python_values), "utf8" + if all(isinstance(value, bytes) for value in python_values): + return min(python_values), max(python_values), "binary" + + return None + + def _block_metadata( coord_arrays: dict, block: Block, @@ -586,39 +671,15 @@ def _block_metadata( if bounds is not None: ranges[str(dim)] = bounds continue - # String/object dtypes are not representable as ScalarBound - # (Int64/Float64/TimestampNanos) and numpy min/max ufuncs do not - # support them. Skip so pruning treats the dimension conservatively. - if coord_values.dtype.kind in ("U", "S", "O"): - continue - - # 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)): - # The Rust pruning layer only accepts int64 nanosecond bounds - # (ScalarBound::TimestampNanos). Dates outside the - # datetime64[ns] range (pre-1678 / post-2262) cannot be - # represented, so skip pruning for this dimension rather than - # raising -- registration still succeeds and the Rust pruner - # treats a missing dimension conservatively (never prunes on it). - try: - min_ns = int(pd.Timestamp(min_val).value) - max_ns = int(pd.Timestamp(max_val).value) - except (OverflowError, pd.errors.OutOfBoundsDatetime): - continue - ranges[str(dim)] = (min_ns, max_ns, "timestamp_ns") - elif hasattr(min_val, "item"): - min_val = min_val.item() - max_val = max_val.item() - dtype = "float64" if isinstance(min_val, float) else "int64" - ranges[str(dim)] = (min_val, max_val, dtype) - else: - dtype = "float64" if isinstance(min_val, float) else "int64" - ranges[str(dim)] = (min_val, max_val, dtype) + # Use actual min/max rather than first/last so non-monotonic axes are + # safe. Expected dtype/conversion failures disable pruning only for + # this dimension; they must not make table registration fail. + try: + bounds = _coordinate_partition_bound(coord_values) + except (OverflowError, TypeError, ValueError): + bounds = None + if bounds is not None: + ranges[str(dim)] = bounds return ranges @@ -638,10 +699,10 @@ def partition_metadata( List of dicts mapping dimension name to (min_value, max_value, dtype_str) tuples. - - For datetime64, values are nanoseconds since Unix epoch - (int64), dtype_str is "timestamp_ns" - - For numeric types, values are Python int or float, - dtype_str is "int64" or "float64" + Values preserve the coordinate's logical type and width. Temporal + values are signed integer counts in their Arrow unit; strings and + bytes use lexical bounds. Unsupported or null-tainted dimensions + are omitted so the Rust layer retains their partitions. Note: If a partition has an empty slice for a dimension, that dimension is diff --git a/xarray_sql/geometry.py b/xarray_sql/geometry.py index 6e3e6406..80e47d8c 100644 --- a/xarray_sql/geometry.py +++ b/xarray_sql/geometry.py @@ -81,7 +81,7 @@ def _wkb_points(x: np.ndarray, y: np.ndarray) -> pa.Array: "pa.binary() offsets are int32 and n * 21 bytes would " "overflow them. Use a smaller batch_size." ) - buf = np.empty((n, 21), dtype=np.uint8) + buf: np.ndarray = np.empty((n, 21), dtype=np.uint8) buf[:, 0] = 1 # little-endian byte order mark buf[:, 1:5] = np.array([1, 0, 0, 0], dtype=np.uint8) # WKB type 1: Point buf[:, 5:13] = x.view(np.uint8).reshape(n, 8) diff --git a/xarray_sql/lazyscan.py b/xarray_sql/lazyscan.py index cbf0551d..1f4b71c8 100644 --- a/xarray_sql/lazyscan.py +++ b/xarray_sql/lazyscan.py @@ -116,9 +116,12 @@ def distinct(self, column: str) -> np.ndarray: dim_only = self._df.select(col(f'"{column}"')).distinct() batches = [b.to_pyarrow() for b in dim_only.execute_stream()] if not batches: - return np.asarray([]) - return np.concatenate( - [b.column(0).to_numpy(zero_copy_only=False) for b in batches] + return cast(np.ndarray, np.asarray([])) + return cast( + np.ndarray, + np.concatenate( + [b.column(0).to_numpy(zero_copy_only=False) for b in batches] + ), ) def fetch( @@ -230,7 +233,10 @@ def distinct(self, column: str) -> np.ndarray: self._rel.project(duckdb.ColumnExpression(column)).distinct() ) ) - return np.asarray(table.column(0).to_numpy(zero_copy_only=False)) + return cast( + np.ndarray, + np.asarray(table.column(0).to_numpy(zero_copy_only=False)), + ) def fetch( self, specs: dict[str, DimSpec], columns: list[str] @@ -287,7 +293,7 @@ def distinct(self, column: str) -> np.ndarray: import polars as pl out = _collect_streaming(self._lf.select(pl.col(column).unique())) - return out.to_series().to_numpy() + return cast(np.ndarray, out.to_series().to_numpy()) def fetch( self, specs: dict[str, DimSpec], columns: list[str]