fix(substrait): decode deprecated timestamp literals as microseconds

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())}

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<usize, usize>) -> 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<usize, usize>) -> 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,58 +336,58 @@ fn remap_expr_references(expr: &mut Expression, mapping: &HashMap<usize, usize>)
         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
                     .arg_type
                     .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(_) => {}
                 }
             }
             Ok(())
         }
         RexType::IfThen(ifthen) => {
             for (i, clause) in ifthen.ifs.iter_mut().enumerate() {
-                remap_expr_references(
+                normalize_expr(
                     clause
                         .r#if
                         .as_mut()
                         .ok_or_else(|| missing_field("if clause condition"))?,
                     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 => {}
                     None => return Err(missing_field("if clause result")),
                 }
             }
             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<usize, usize>)
         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() {