From d11808b44670787b67a307044933283233415cab Mon Sep 17 00:00:00 2001 From: Jonas Dedden Date: Thu, 27 Aug 2026 17:27:16 +0200 Subject: [PATCH] fix(substrait): decode deprecated timestamp literals as microseconds --- python/python/tests/test_filter.py | 23 +++++ rust/lance-datafusion/src/substrait.rs | 130 +++++++++++++++++++++---- 2 files changed, 134 insertions(+), 19 deletions(-) diff --git a/python/python/tests/test_filter.py b/python/python/tests/test_filter.py index 16bd3c0061d..e2d4cbe0121 100644 --- a/python/python/tests/test_filter.py +++ b/python/python/tests/test_filter.py @@ -8,6 +8,7 @@ from datetime import date, datetime, timedelta from decimal import Decimal from pathlib import Path +from zoneinfo import ZoneInfo import lance import numpy as np @@ -108,6 +109,28 @@ def test_sql_predicates(dataset): assert dataset.to_table(filter=expr).num_rows == expected_num_rows +@pytest.mark.parametrize("unit", ["s", "ms", "us"]) +@pytest.mark.parametrize("timezone", [None, "UTC", "America/New_York"]) +def test_timestamp_pyarrow_predicates(tmp_path: Path, unit: str, timezone: str | None): + # PyArrow filters reach Lance as Substrait, where the timestamp literal used to be + # decoded in the wrong unit. + tz = ZoneInfo(timezone) if timezone else None + start = datetime(2021, 1, 1, tzinfo=tz) + ts_type = pa.timestamp(unit, timezone) + table = pa.table( + {"ts": pa.array([start + timedelta(hours=i) for i in range(100)], ts_type)} + ) + dataset = lance.write_dataset(table, tmp_path / f"{unit}_{timezone}") + + cutoff = pa.scalar(start + timedelta(hours=50), ts_type) + for expr in [ + pc.field("ts") > cutoff, + pc.field("ts") < cutoff, + pc.field("ts") == cutoff, + ]: + assert dataset.to_table(filter=expr) == table.filter(expr) + + def test_sql_current_date(tmp_path: Path): table = pa.table( {"date": pa.array([date(2020, 1, 1), date(2020, 1, 2)], type=pa.date32())} diff --git a/rust/lance-datafusion/src/substrait.rs b/rust/lance-datafusion/src/substrait.rs index 9f14fb5cdc0..f3dc9b4d3df 100644 --- a/rust/lance-datafusion/src/substrait.rs +++ b/rust/lance-datafusion/src/substrait.rs @@ -13,8 +13,9 @@ use datafusion_substrait::logical_plan::consumer::{ use datafusion_substrait::substrait::proto::{ AggregateRel, Expression, ExpressionReference, ExtendedExpression, NamedStruct, Plan, Type, expression::{ - RexType, + Literal, RexType, field_reference::{ReferenceType, RootType}, + literal::{LiteralType, PrecisionTimestamp}, reference_segment, }, expression_reference::ExprType, @@ -280,14 +281,47 @@ fn missing_field(what: &str) -> Error { )) } -fn remap_expr_references(expr: &mut Expression, mapping: &HashMap) -> Result<()> { +/// Substrait's deprecated `timestamp`/`timestamp_tz` literals are always microseconds, but +/// DataFusion takes their unit from `type_variation_reference` and reads the default 0 as +/// seconds. PyArrow emits exactly that, so filters silently matched the wrong rows. +/// +/// Rewrite them into the `precision_timestamp` forms, which state the unit. +#[allow(deprecated)] +fn normalize_deprecated_timestamp_literal(lit: &mut Literal) { + let precision = match lit.type_variation_reference { + // 0 is Substrait's default reference (microseconds); 1..=3 are DataFusion's ms/us/ns. + 0 | 2 => 6, + 1 => 3, + 3 => 9, + _ => return, + }; + let replacement = match lit.literal_type { + Some(LiteralType::Timestamp(value)) => { + LiteralType::PrecisionTimestamp(PrecisionTimestamp { precision, value }) + } + Some(LiteralType::TimestampTz(value)) => { + LiteralType::PrecisionTimestampTz(PrecisionTimestamp { precision, value }) + } + _ => return, + }; + lit.literal_type = Some(replacement); + lit.type_variation_reference = 0; +} + +/// Reject operators we cannot push down, normalize ambiguous literals, and remap field +/// references onto the schema `remove_extension_types` left behind. +fn normalize_expr(expr: &mut Expression, mapping: &HashMap) -> Result<()> { match expr .rex_type .as_mut() .ok_or_else(|| missing_field("expression"))? { + RexType::Literal(lit) => { + normalize_deprecated_timestamp_literal(lit); + Ok(()) + } // Simple, no field references possible - RexType::Literal(_) | RexType::Nested(_) | RexType::DynamicParameter(_) => Ok(()), + RexType::Nested(_) | RexType::DynamicParameter(_) => Ok(()), // Enum literals are deprecated in Substrait and should only appear in older plans. #[allow(deprecated)] RexType::Enum(_) => Ok(()), @@ -302,7 +336,7 @@ fn remap_expr_references(expr: &mut Expression, mapping: &HashMap) RexType::ScalarFunction(func) => { #[allow(deprecated)] for arg in &mut func.args { - remap_expr_references(arg, mapping)?; + normalize_expr(arg, mapping)?; } for arg in &mut func.arguments { match arg @@ -310,7 +344,7 @@ fn remap_expr_references(expr: &mut Expression, mapping: &HashMap) .as_mut() .ok_or_else(|| missing_field("function argument"))? { - ArgType::Value(expr) => remap_expr_references(expr, mapping)?, + ArgType::Value(expr) => normalize_expr(expr, mapping)?, ArgType::Enum(_) | ArgType::Type(_) => {} } } @@ -318,7 +352,7 @@ fn remap_expr_references(expr: &mut Expression, mapping: &HashMap) } RexType::IfThen(ifthen) => { for (i, clause) in ifthen.ifs.iter_mut().enumerate() { - remap_expr_references( + normalize_expr( clause .r#if .as_mut() @@ -326,7 +360,7 @@ fn remap_expr_references(expr: &mut Expression, mapping: &HashMap) mapping, )?; match clause.then.as_mut() { - Some(then) => remap_expr_references(then, mapping)?, + Some(then) => normalize_expr(then, mapping)?, // Only the leading clause may omit `then`, in which case its condition is // the case expression being matched against. None if i == 0 => {} @@ -334,26 +368,26 @@ fn remap_expr_references(expr: &mut Expression, mapping: &HashMap) } } if let Some(otherwise) = ifthen.r#else.as_mut() { - remap_expr_references(otherwise, mapping)?; + normalize_expr(otherwise, mapping)?; } Ok(()) } RexType::SwitchExpression(switch) => { for clause in switch.ifs.iter_mut() { if let Some(then) = clause.then.as_mut() { - remap_expr_references(then, mapping)?; + normalize_expr(then, mapping)?; } } if let Some(otherwise) = switch.r#else.as_mut() { - remap_expr_references(otherwise, mapping)?; + normalize_expr(otherwise, mapping)?; } Ok(()) } RexType::SingularOrList(orlist) => { for opt in orlist.options.iter_mut() { - remap_expr_references(opt, mapping)?; + normalize_expr(opt, mapping)?; } - remap_expr_references( + normalize_expr( orlist .value .as_mut() @@ -365,16 +399,16 @@ fn remap_expr_references(expr: &mut Expression, mapping: &HashMap) RexType::MultiOrList(orlist) => { for opt in orlist.options.iter_mut() { for field in opt.fields.iter_mut() { - remap_expr_references(field, mapping)?; + normalize_expr(field, mapping)?; } } for val in orlist.value.iter_mut() { - remap_expr_references(val, mapping)?; + normalize_expr(val, mapping)?; } Ok(()) } RexType::Cast(cast) => { - remap_expr_references( + normalize_expr( cast.input .as_mut() .ok_or_else(|| missing_field("cast input"))?, @@ -473,9 +507,10 @@ pub async fn parse_substrait( let (substrait_schema, _, index_mapping) = remove_extension_types(envelope.base_schema.as_ref().unwrap(), input_schema.clone())?; - // Always walk the expression: this also rejects operators we cannot push down. When no - // fields were removed the mapping is the identity, so the remap itself is a no-op. - remap_expr_references(&mut expr, &index_mapping)?; + // Always walk the expression: this also rejects operators we cannot push down and + // normalizes literals. When no fields were removed the mapping is the identity, so the + // remap itself is a no-op. + normalize_expr(&mut expr, &index_mapping)?; substrait_schema } else { @@ -698,7 +733,7 @@ async fn parse_measures( mod tests { use std::sync::Arc; - use arrow_schema::{DataType, Field, Schema}; + use arrow_schema::{DataType, Field, Schema, TimeUnit}; use datafusion::{ execution::SessionState, logical_expr::{BinaryExpr, Case, Operator}, @@ -724,6 +759,7 @@ mod tests { r#type::{Boolean, I32, Kind, Nullability, Struct}, }; use prost::Message; + use rstest::rstest; use crate::substrait::{encode_substrait, parse_substrait}; @@ -907,6 +943,62 @@ mod tests { assert_eq!(expr, Expr::Column(Column::new_unqualified("x"))); } + /// The deprecated literal is always microseconds, but DataFusion reads the default + /// variation reference as seconds. + #[rstest] + #[case::default_reference(0, TimeUnit::Microsecond)] + #[case::milli_reference(1, TimeUnit::Millisecond)] + #[case::micro_reference(2, TimeUnit::Microsecond)] + #[case::nano_reference(3, TimeUnit::Nanosecond)] + #[tokio::test] + async fn test_deprecated_timestamp_literal_units( + #[case] type_variation_reference: u32, + #[case] expected_unit: TimeUnit, + ) { + const MICROS: i64 = 1_704_247_200_000_000; + + #[allow(deprecated)] + let expr = parse_unpruned_expr(RexType::Literal(Literal { + nullable: false, + type_variation_reference, + literal_type: Some(LiteralType::Timestamp(MICROS)), + })) + .await + .unwrap(); + + let expected = match expected_unit { + TimeUnit::Millisecond => ScalarValue::TimestampMillisecond(Some(MICROS), None), + TimeUnit::Microsecond => ScalarValue::TimestampMicrosecond(Some(MICROS), None), + TimeUnit::Nanosecond => ScalarValue::TimestampNanosecond(Some(MICROS), None), + other => panic!("unexpected time unit {other:?}"), + }; + assert_eq!(expr, Expr::Literal(expected, None)); + } + + /// DataFusion has no consumer branch for the deprecated `timestamp_tz`, so it failed the + /// filter outright rather than answering wrongly. + #[tokio::test] + async fn test_deprecated_timestamp_tz_literal() { + const MICROS: i64 = 1_704_247_200_000_000; + + #[allow(deprecated)] + let expr = parse_unpruned_expr(RexType::Literal(Literal { + nullable: false, + type_variation_reference: 0, + literal_type: Some(LiteralType::TimestampTz(MICROS)), + })) + .await + .unwrap(); + + assert_eq!( + expr, + Expr::Literal( + ScalarValue::TimestampMicrosecond(Some(MICROS), Some("UTC".into())), + None + ) + ); + } + /// Optional message fields that a producer may legitimately omit must not panic the walker. #[tokio::test] async fn test_unpruned_if_then_with_omitted_optional_fields() {