diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index bb10e2c17..e71169624 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -1,4 +1,5 @@ -from datetime import datetime, timedelta +from datetime import date, datetime, timedelta +from decimal import Decimal from typing import Dict, List, Optional, Tuple, Any, TypedDict, Union, Literal import pyarrow as pa @@ -53,7 +54,9 @@ class PyExpr: def to_sql(self) -> str: ... def expr_col(name: str) -> PyExpr: ... -def expr_lit(value: Union[bool, int, float, str, bytes]) -> PyExpr: ... +def expr_lit( + value: Union[bool, int, float, str, bytes, date, datetime, Decimal], +) -> PyExpr: ... def expr_func(name: str, args: List[PyExpr]) -> PyExpr: ... class Session: diff --git a/python/python/lancedb/expr.py b/python/python/lancedb/expr.py index c28bc208a..e8b2d63a4 100644 --- a/python/python/lancedb/expr.py +++ b/python/python/lancedb/expr.py @@ -19,6 +19,8 @@ operators:: from __future__ import annotations +from datetime import date, datetime +from decimal import Decimal from typing import Iterable, Union import pyarrow as pa @@ -63,7 +65,7 @@ def _coerce(value: "ExprLike") -> "Expr": # Type alias used in annotations. -ExprLike = Union["Expr", bool, int, float, str, bytes] +ExprLike = Union["Expr", bool, int, float, str, bytes, date, datetime, Decimal] class Expr: @@ -118,10 +120,18 @@ class Expr: """Logical AND (``expr_a & expr_b``).""" return Expr(self._inner.and_(_coerce(other)._inner)) + def __rand__(self, other: ExprLike) -> "Expr": + """Right-hand logical AND (``True & expr``).""" + return Expr(_coerce(other)._inner.and_(self._inner)) + def __or__(self, other: "Expr") -> "Expr": """Logical OR (``expr_a | expr_b``).""" return Expr(self._inner.or_(_coerce(other)._inner)) + def __ror__(self, other: ExprLike) -> "Expr": + """Right-hand logical OR (``False | expr``).""" + return Expr(_coerce(other)._inner.or_(self._inner)) + def __invert__(self) -> "Expr": """Logical NOT (``~expr``).""" return Expr(self._inner.not_()) @@ -266,13 +276,14 @@ def col(name: str) -> Expr: return Expr(expr_col(name)) -def lit(value: Union[bool, int, float, str, bytes]) -> Expr: +def lit(value: Union[bool, int, float, str, bytes, date, datetime, Decimal]) -> Expr: """Create a literal (constant) value expression. Parameters ---------- value: - A Python ``bool``, ``int``, ``float``, ``str``, or ``bytes``. + A Python ``bool``, ``int``, ``float``, ``str``, ``bytes``, ``date``, + ``datetime``, or ``Decimal``. Examples -------- @@ -280,6 +291,9 @@ def lit(value: Union[bool, int, float, str, bytes]) -> Expr: >>> col("price") * lit(1.1) Expr((price * 1.1)) """ + if not isinstance(value, (bool, int, float, str, bytes, date, datetime, Decimal)): + raise TypeError(f"Unsupported literal type: {type(value).__name__}") + return Expr(expr_lit(value)) diff --git a/python/tests/test_expr.py b/python/python/tests/test_expr.py similarity index 76% rename from python/tests/test_expr.py rename to python/python/tests/test_expr.py index 7858995e7..6aa78943e 100644 --- a/python/tests/test_expr.py +++ b/python/python/tests/test_expr.py @@ -3,10 +3,14 @@ """Tests for the type-safe expression builder API.""" -import pytest +from datetime import date, datetime, timedelta, timezone +from decimal import Decimal + import pyarrow as pa +import pytest + import lancedb -from lancedb.expr import Expr, col, lit, func +from lancedb.expr import Expr, col, func, lit # ── unit tests for Expr construction ───────────────────────────────────────── @@ -54,6 +58,28 @@ class TestExprConstruction: with pytest.raises(Exception): func("not_a_real_function", col("x")) + def test_lit_date(self): + e = lit(date(2024, 1, 1)) + assert isinstance(e, Expr) + + def test_lit_datetime(self): + # Naive datetime + e = lit(datetime(2024, 1, 1, 10, 0)) + assert isinstance(e, Expr) + + def test_lit_datetime_tz(self): + # Timezone-aware datetime + tz = timezone(timedelta(hours=5)) + dt = datetime(2024, 1, 1, 10, 0, tzinfo=tz) + e = lit(dt) + assert isinstance(e, Expr) + + def test_lit_decimal_precision(self): + # High precision Decimal that would be rounded if converted to float + d = Decimal("1.234567890123456789") + e = lit(d) + assert isinstance(e, Expr) + class TestExprOperators: def test_eq_operator(self): @@ -142,6 +168,20 @@ class TestExprOperators: assert isinstance(e, Expr) assert e.to_sql() == "(name = 'alice')" + def test_reflexive_comparisons(self): + # 10 < col("age") swaps to col("age") > 10 + assert (10 < col("age")).to_sql() == "(age > 10)" + assert (10 <= col("age")).to_sql() == "(age >= 10)" + assert (10 > col("age")).to_sql() == "(age < 10)" + assert (10 >= col("age")).to_sql() == "(age <= 10)" + assert (10 == col("age")).to_sql() == "(age = 10)" + assert (10 != col("age")).to_sql() == "(age <> 10)" + + def test_reflexive_logical(self): + # True & Expr calls Expr.__rand__(True) + assert (True & (col("age") > 18)).to_sql() == "(true AND (age > 18))" + assert (False | (col("age") > 18)).to_sql() == "(false OR (age > 18))" + class TestExprBytesLiteral: def test_bytes_to_sql(self): @@ -282,6 +322,40 @@ class TestExprRepr: {e: 1} +class TestExprReflexive: + def test_reflexive_eq(self): + e = 1 == col("x") + assert isinstance(e, Expr) + assert e.to_sql() == "(x = 1)" + + def test_reflexive_ne(self): + e = 1 != col("x") + assert isinstance(e, Expr) + assert e.to_sql() == "(x <> 1)" + + def test_reflexive_lt(self): + # 1 < x => (x > 1) + e = 1 < col("x") + assert isinstance(e, Expr) + assert e.to_sql() == "(x > 1)" + + def test_reflexive_gt(self): + # 1 > x => (x < 1) + e = 1 > col("x") + assert isinstance(e, Expr) + assert e.to_sql() == "(x < 1)" + + def test_reflexive_and(self): + e = True & col("active") + assert isinstance(e, Expr) + assert e.to_sql() == "(true AND active)" + + def test_reflexive_or(self): + e = False | col("inactive") + assert isinstance(e, Expr) + assert e.to_sql() == "(false OR inactive)" + + # ── integration tests: end-to-end query against a real table ───────────────── @@ -432,6 +506,72 @@ class TestColNamingIntegration: assert sorted(result["upper_name"].to_pylist()) == ["ALICE", "BOB", "CHARLIE"] +@pytest.fixture +def type_check_table(tmp_path): + """Fixture that creates a table with Date32 and Decimal128 columns.""" + db = lancedb.connect(str(tmp_path)) + schema = pa.schema( + [ + ("date", pa.date32()), + ("decimal", pa.decimal128(10, 2)), + ("binary", pa.binary()), + ] + ) + data = pa.table( + { + "date": [date(2024, 1, 1), date(2024, 1, 2)], + "decimal": [Decimal("10.50"), Decimal("20.75")], + "binary": [b"\x01", b"\x02"], + }, + schema=schema, + ) + return db.create_table("extended_types", data) + + +class TestExtendedTypeIntegration: + """Integration tests verifying that typed literals work correctly in filters.""" + + def test_date_integration(self, type_check_table): + """Verify that Date32 literals are correctly parsed and filtered.""" + result = ( + type_check_table.search() + .where(col("date") == lit(date(2024, 1, 1))) + .to_arrow() + ) + assert result.num_rows == 1 + assert result["date"][0].as_py() == date(2024, 1, 1) + + def test_decimal_integration(self, tmp_path): + """A Decimal literal must retain full 128-bit precision in a filter. + + 1.234567890123456789 and 1.234567890123456790 differ only in the last + digit and are indistinguishable once truncated to f64. The filter + therefore returns the single expected row only if ``lit(Decimal)`` + produces a true ``Decimal128`` scalar rather than being coerced to f64. + """ + low = Decimal("1.234567890123456789") + high = Decimal("1.234567890123456790") + + db = lancedb.connect(str(tmp_path / "decimal_precision")) + schema = pa.schema([("val", pa.decimal128(19, 18))]) + table = db.create_table( + "decimal_precision", + pa.table({"val": [low, high]}, schema=schema), + ) + + result = table.search().where(col("val") < lit(high)).to_arrow() + assert result.num_rows == 1 + assert result["val"][0].as_py() == low + + def test_binary_integration(self, type_check_table): + """Verify that Binary literals are correctly filtered.""" + result = ( + type_check_table.search().where(col("binary") == lit(b"\x01")).to_arrow() + ) + assert result.num_rows == 1 + assert result["binary"][0].as_py() == b"\x01" + + # ── bytes / binary column integration tests ─────────────────────────────────── diff --git a/python/src/expr.rs b/python/src/expr.rs index d3e9ea5aa..5d5ceb268 100644 --- a/python/src/expr.rs +++ b/python/src/expr.rs @@ -7,12 +7,14 @@ //! build type-safe filter / projection expressions that map directly to //! DataFusion [`Expr`] nodes, bypassing SQL string parsing. +use std::ops::{Add, Div, Mul, Not, Sub}; + use arrow::{datatypes::DataType, pyarrow::PyArrowType}; use datafusion_common::ScalarValue; use lancedb::expr::{ DfExpr, col as ldb_col, contains, expr_cast, is_in, lit as df_lit, lower, upper, }; -use pyo3::types::PyBytes; +use pyo3::types::{PyBytes, PyDate, PyDateTime}; use pyo3::{Bound, PyAny, PyResult, exceptions::PyValueError, prelude::*, pyfunction}; /// A type-safe DataFusion expression. @@ -63,30 +65,30 @@ impl PyExpr { Self(self.0.clone().or(other.0.clone())) } + /// Logical NOT. fn not_(&self) -> Self { - use std::ops::Not; Self(self.0.clone().not()) } // ── arithmetic ─────────────────────────────────────────────────────────── + /// Add expressions. fn add(&self, other: &Self) -> Self { - use std::ops::Add; Self(self.0.clone().add(other.0.clone())) } + /// Subtract expressions. fn sub(&self, other: &Self) -> Self { - use std::ops::Sub; Self(self.0.clone().sub(other.0.clone())) } + /// Multiply expressions. fn mul(&self, other: &Self) -> Self { - use std::ops::Mul; Self(self.0.clone().mul(other.0.clone())) } + /// Divide expressions. fn div(&self, other: &Self) -> Self { - use std::ops::Div; Self(self.0.clone().div(other.0.clone())) } @@ -153,7 +155,8 @@ pub fn expr_col(name: &str) -> PyExpr { /// Create a literal value expression. /// -/// Supported Python types: `bool`, `int`, `float`, `str`, `bytes`. +/// Supported Python types: `bool`, `int`, `float`, `str`, `bytes`, `date`, +/// `datetime`, `Decimal`. #[pyfunction] pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult { // bool must be checked before int because bool is a subclass of int in Python @@ -163,6 +166,19 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult { if let Ok(i) = value.extract::() { return Ok(PyExpr(df_lit(i))); } + // Decimal must be checked before f64: Python's Decimal implements __float__, + // so value.extract::() would succeed and silently truncate the value to + // f64, losing precision. Build a Decimal128 scalar to preserve it instead. + if value.get_type().name()? == "Decimal" { + let s = value.call_method0("__str__")?.extract::()?; + // Parse the decimal string into an i128 value, precision, and scale. + let (val, precision, scale) = parse_decimal(&s)?; + return Ok(PyExpr(df_lit(ScalarValue::Decimal128( + Some(val), + precision, + scale, + )))); + } if let Ok(f) = value.extract::() { return Ok(PyExpr(df_lit(f))); } @@ -173,12 +189,48 @@ pub fn expr_lit(value: Bound<'_, PyAny>) -> PyResult { let bytes = value.extract::>()?; return Ok(PyExpr(df_lit(ScalarValue::Binary(Some(bytes))))); } + + // datetime.datetime is a subclass of datetime.date, so it must be checked first. + if let Ok(dt) = value.downcast::() { + let ts: f64 = dt.call_method0("timestamp")?.extract()?; + let micros = (ts * 1_000_000.0).round() as i64; + return Ok(PyExpr(df_lit(ScalarValue::TimestampMicrosecond( + Some(micros), + None, + )))); + } + if let Ok(d) = value.downcast::() { + let ordinal: i32 = d.call_method0("toordinal")?.extract()?; + let days = ordinal - 719163; // Unix epoch is 1970-01-01 + return Ok(PyExpr(df_lit(ScalarValue::Date32(Some(days))))); + } + Err(PyValueError::new_err(format!( - "unsupported literal type: {}. Supported: bool, int, float, str, bytes", + "unsupported literal type: {}. Supported: bool, int, float, str, bytes, date, datetime, Decimal", value.get_type().name()? ))) } +fn parse_decimal(s: &str) -> PyResult<(i128, u8, i8)> { + let s = s.trim(); + let dot_pos = s.find('.'); + let scale = if let Some(pos) = dot_pos { + (s.len() - pos - 1) as i8 + } else { + 0 + }; + + let digits = s.replace('.', ""); + let val = digits + .parse::() + .map_err(|e| PyValueError::new_err(format!("failed to parse decimal digits: {}", e)))?; + + // Precision is total number of digits + let precision = digits.trim_start_matches('-').len() as u8; + + Ok((val, precision, scale)) +} + /// Call an arbitrary registered SQL function by name. /// /// See `lancedb::expr::func` for the list of supported function names.