diff --git a/python/python/tests/test_expr.py b/python/python/tests/test_expr.py index 6aa78943e..0995a39cf 100644 --- a/python/python/tests/test_expr.py +++ b/python/python/tests/test_expr.py @@ -52,7 +52,7 @@ class TestExprConstruction: def test_func(self): e = func("lower", col("name")) assert isinstance(e, Expr) - assert e.to_sql() == "lower(name)" + assert e.to_sql() == "lower(`name`)" def test_func_unknown_raises(self): with pytest.raises(Exception): @@ -115,7 +115,7 @@ class TestExprOperators: def test_and_operator(self): e = (col("age") > lit(18)) & (col("status") == lit("active")) assert isinstance(e, Expr) - assert e.to_sql() == "((age > 18) AND (status = 'active'))" + assert e.to_sql() == "((age > 18) AND (`status` = 'active'))" def test_or_operator(self): e = (col("a") == lit(1)) | (col("b") == lit(2)) @@ -166,7 +166,7 @@ class TestExprOperators: def test_coerce_plain_str(self): e = col("name") == "alice" assert isinstance(e, Expr) - assert e.to_sql() == "(name = 'alice')" + assert e.to_sql() == "(`name` = 'alice')" def test_reflexive_comparisons(self): # 10 < col("age") swaps to col("age") > 10 @@ -198,53 +198,53 @@ class TestExprBytesLiteral: def test_bytes_equality_expr_sql(self): e = col("data") == lit(b"\xca\xfe") - assert e.to_sql() == "(data = X'CAFE')" + assert e.to_sql() == "(`data` = X'CAFE')" def test_bytes_ne_expr_sql(self): e = col("data") != lit(b"\xff") - assert e.to_sql() == "(data <> X'FF')" + assert e.to_sql() == "(`data` <> X'FF')" def test_bytes_compound_expr_sql(self): e = (col("data") == lit(b"\x01")) & (col("id") > lit(5)) - assert e.to_sql() == "((data = X'01') AND (id > 5))" + assert e.to_sql() == "((`data` = X'01') AND (id > 5))" def test_bytes_in_function_call(self): # Regression test: binary literals inside scalar function calls # used to fail because DataFusion's unparser does not support Binary # scalars. Now handled via a placeholder-substitution rewrite. e = func("contains", col("data"), lit(b"\xff")) - assert e.to_sql() == "contains(data, X'FF')" + assert e.to_sql() == "contains(`data`, X'FF')" def test_bytes_in_not(self): e = ~(col("data") == lit(b"\xff")) - assert e.to_sql() == "NOT (data = X'FF')" + assert e.to_sql() == "NOT (`data` = X'FF')" class TestExprStringMethods: def test_lower(self): e = col("name").lower() assert isinstance(e, Expr) - assert e.to_sql() == "lower(name)" + assert e.to_sql() == "lower(`name`)" def test_upper(self): e = col("name").upper() assert isinstance(e, Expr) - assert e.to_sql() == "upper(name)" + assert e.to_sql() == "upper(`name`)" def test_contains(self): e = col("text").contains(lit("hello")) assert isinstance(e, Expr) - assert e.to_sql() == "contains(text, 'hello')" + assert e.to_sql() == "contains(`text`, 'hello')" def test_contains_with_str_coerce(self): e = col("text").contains("hello") assert isinstance(e, Expr) - assert e.to_sql() == "contains(text, 'hello')" + assert e.to_sql() == "contains(`text`, 'hello')" def test_chained_lower_eq(self): e = col("name").lower() == lit("alice") assert isinstance(e, Expr) - assert e.to_sql() == "(lower(name) = 'alice')" + assert e.to_sql() == "(lower(`name`) = 'alice')" class TestExprCast: @@ -597,14 +597,14 @@ class TestExprIsin: def test_isin_strs(self): assert ( col("status").isin(["active", "pending"]).to_sql() - == "status IN ('active', 'pending')" + == "`status` IN ('active', 'pending')" ) def test_isin_coerces_and_mixes(self): assert col("id").isin([lit(1), 2]).to_sql() == "id IN (1, 2)" def test_isin_empty(self): - assert col("id").isin([]).to_sql() == "id IN ()" + assert col("id").isin([]).to_sql() == "false" def test_isin_filter(self, simple_table): result = simple_table.search().where(col("id").isin([1, 3, 5])).to_arrow() diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 2d47606fd..1a341eb75 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -8,6 +8,7 @@ import threading import warnings from concurrent.futures import ThreadPoolExecutor from datetime import date, datetime, timedelta +from decimal import Decimal from time import sleep from typing import List from unittest.mock import patch @@ -2225,6 +2226,61 @@ def test_update_expr_filter_literals(mem_db: DBConnection): assert table.to_arrow()["result"].to_pylist() == values +def test_update_expr_filter_preserves_typed_semantics(mem_db: DBConnection): + low = Decimal("1.234567890123456789") + high = Decimal("1.234567890123456790") + decimal_schema = pa.schema( + [("val", pa.decimal128(19, 18)), ("result", pa.string())] + ) + decimal_table = mem_db.create_table( + "update_expr_decimal", + pa.table( + {"val": [low, high], "result": ["old", "old"]}, + schema=decimal_schema, + ), + ) + predicate = col("val") < lit(high) + assert decimal_table.search().where(predicate).to_arrow().num_rows == 1 + result = decimal_table.update(where=predicate, values={"result": "new"}) + assert result.rows_updated == 1 + + keyword_table = mem_db.create_table( + "update_expr_keyword", [{"null": 1, "result": "old"}] + ) + predicate = col("null") == 1 + assert keyword_table.search().where(predicate).to_arrow().num_rows == 1 + result = keyword_table.update(where=predicate, values={"result": "new"}) + assert result.rows_updated == 1 + + empty_in_table = mem_db.create_table( + "update_expr_empty_in", [{"id": 1, "result": "old"}] + ) + predicate = col("id").isin([]) + assert empty_in_table.search().where(predicate).to_arrow().num_rows == 0 + result = empty_in_table.update(where=predicate, values={"result": "new"}) + assert result.rows_updated == 0 + + marker = "__lancedb_binary_placeholder_0__" + binary_schema = pa.schema( + [("payload", pa.binary()), ("text", pa.string()), ("result", pa.string())] + ) + binary_table = mem_db.create_table( + "update_expr_binary", + pa.table( + { + "payload": [b"\x01", b"\x02"], + "text": ["other", marker], + "result": ["old", "old"], + }, + schema=binary_schema, + ), + ) + predicate = (col("payload") == lit(b"\x01")) | (col("text") == marker) + assert binary_table.search().where(predicate).to_arrow().num_rows == 2 + result = binary_table.update(where=predicate, values={"result": "new"}) + assert result.rows_updated == 2 + + def test_update_types(mem_db: DBConnection): table = mem_db.create_table( "my_table", diff --git a/rust/lancedb/src/expr.rs b/rust/lancedb/src/expr.rs index 75cce443d..84174fc14 100644 --- a/rust/lancedb/src/expr.rs +++ b/rust/lancedb/src/expr.rs @@ -156,7 +156,7 @@ mod tests { use datafusion_common::ScalarValue; let expr = col("data").eq(lit(ScalarValue::Binary(Some(vec![0xca, 0xfe])))); let sql = expr_to_sql_string(&expr).unwrap(); - assert_eq!(sql, "(data = X'CAFE')"); + assert_eq!(sql, "(`data` = X'CAFE')"); } #[test] @@ -166,7 +166,7 @@ mod tests { let int_expr = col("id").gt(lit(5i64)); let combined = bin_expr.and(int_expr); let sql = expr_to_sql_string(&combined).unwrap(); - assert_eq!(sql, "((data = X'01') AND (id > 5))"); + assert_eq!(sql, "((`data` = X'01') AND (id > 5))"); } #[test] @@ -184,7 +184,7 @@ mod tests { // serialized correctly (regression test for placeholder rewrite path). let expr = contains(col("data"), lit(ScalarValue::Binary(Some(vec![0xff])))); let sql = expr_to_sql_string(&expr).unwrap(); - assert_eq!(sql, "contains(data, X'FF')"); + assert_eq!(sql, "contains(`data`, X'FF')"); } #[test] @@ -195,7 +195,7 @@ mod tests { .eq(lit(ScalarValue::Binary(Some(vec![0xab, 0xcd])))) .not(); let sql = expr_to_sql_string(&expr).unwrap(); - assert_eq!(sql, "NOT (data = X'ABCD')"); + assert_eq!(sql, "NOT (`data` = X'ABCD')"); } #[test] @@ -205,6 +205,45 @@ mod tests { assert!(sql.contains("IN"), "expected IN in: {}", sql); } + #[test] + fn test_empty_is_in() { + let expr = is_in(col("id"), vec![]); + assert_eq!(expr_to_sql_string(&expr).unwrap(), "false"); + } + + #[test] + fn test_keyword_identifier() { + let expr = col("null").eq(lit(1i64)); + assert_eq!(expr_to_sql_string(&expr).unwrap(), "(`null` = 1)"); + } + + #[test] + fn test_decimal_literal_preserves_type() { + use datafusion_common::ScalarValue; + + let expr = col("val").lt(lit(ScalarValue::Decimal128( + Some(1_234_567_890_123_456_790), + 19, + 18, + ))); + let sql = expr_to_sql_string(&expr).unwrap().replace(' ', ""); + assert_eq!(sql, "(val Option { - let needs_quote = identifier.chars().any(|c| c.is_ascii_uppercase()) - || !identifier - .chars() - .enumerate() - .all(|(i, c)| c == '_' || c.is_ascii_alphabetic() || (i > 0 && c.is_ascii_digit())); + let identifier_upper = identifier.to_ascii_uppercase(); + let needs_quote = + (identifier_upper != "ID" && ALL_KEYWORDS.contains(&identifier_upper.as_str())) + || identifier.chars().any(|c| c.is_ascii_uppercase()) + || !identifier.chars().enumerate().all(|(i, c)| { + c == '_' || c.is_ascii_alphabetic() || (i > 0 && c.is_ascii_digit()) + }); if needs_quote { Some('`') } else { None } } } @@ -39,24 +48,21 @@ fn bytes_to_hex_sql(bytes: &[u8]) -> String { format!("X'{hex}'") } -/// Returns true if *expr* contains a `Binary` or `LargeBinary` scalar literal -/// anywhere in its subtree. DataFusion's SQL unparser cannot serialize those -/// variants, so we route such expressions through a placeholder-substitution -/// path that emits SQL `X'...'` byte-string literals. -fn has_binary_literal(expr: &Expr) -> bool { - let mut found = false; +fn string_literals(expr: &Expr) -> HashSet { + let mut literals = HashSet::new(); let _ = expr.apply(&mut |e: &Expr| { - if matches!( - e, - Expr::Literal(ScalarValue::Binary(_) | ScalarValue::LargeBinary(_), _) - ) { - found = true; - Ok(TreeNodeRecursion::Stop) - } else { - Ok(TreeNodeRecursion::Continue) + if let Expr::Literal( + ScalarValue::Utf8(Some(value)) + | ScalarValue::LargeUtf8(Some(value)) + | ScalarValue::Utf8View(Some(value)), + _, + ) = e + { + literals.insert(value.clone()); } + Ok(TreeNodeRecursion::Continue) }); - found + literals } fn run_unparser(expr: &Expr) -> crate::Result { @@ -69,25 +75,28 @@ fn run_unparser(expr: &Expr) -> crate::Result { } pub fn expr_to_sql_string(expr: &Expr) -> crate::Result { - // Fast path: no binary literals — DataFusion's unparser handles everything. - if !has_binary_literal(expr) { - return run_unparser(expr); - } - - // Slow path: DataFusion's unparser cannot serialize `Binary`/`LargeBinary` - // scalars, so we rewrite each one to a unique string-literal placeholder, - // let the unparser do the rest of the work, then substitute the SQL - // `X'...'` byte-string literal back in. This keeps the operator/function - // serialization logic centralized in DataFusion and works for every - // expression node type the unparser supports. - let mut bindings: Vec> = Vec::new(); + // DataFusion's unparser needs a few adaptations before its SQL can be + // reparsed by Lance without changing the typed expression's semantics: + // + // * decimal literals need an explicit cast to preserve precision and scale; + // * an empty IN list is valid in DataFusion but invalid SQL; + // * binary literals are unsupported by the unparser and need placeholders. + let user_strings = string_literals(expr); + let mut binary_bindings: Vec<(String, Vec)> = Vec::new(); let rewritten = expr .clone() .transform(|e: Expr| match e { Expr::Literal(ScalarValue::Binary(Some(bytes)), m) | Expr::Literal(ScalarValue::LargeBinary(Some(bytes)), m) => { - let placeholder = format!("{}{}__", BINARY_PLACEHOLDER_PREFIX, bindings.len()); - bindings.push(bytes); + let mut placeholder = + format!("{}{}__", BINARY_PLACEHOLDER_PREFIX, binary_bindings.len()); + while user_strings + .iter() + .any(|value| value.contains(&placeholder)) + { + placeholder.push('_'); + } + binary_bindings.push((placeholder.clone(), bytes)); Ok(Transformed::yes(Expr::Literal( ScalarValue::Utf8(Some(placeholder)), m, @@ -97,6 +106,37 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result { | Expr::Literal(ScalarValue::LargeBinary(None), m) => { Ok(Transformed::yes(Expr::Literal(ScalarValue::Null, m))) } + Expr::Literal(ScalarValue::Decimal32(Some(value), precision, scale), m) => { + let value = Decimal32Type::format_decimal(value, precision, scale); + Ok(Transformed::yes(Expr::Cast(Cast::new( + Box::new(Expr::Literal(ScalarValue::Utf8(Some(value)), m)), + DataType::Decimal32(precision, scale), + )))) + } + Expr::Literal(ScalarValue::Decimal64(Some(value), precision, scale), m) => { + let value = Decimal64Type::format_decimal(value, precision, scale); + Ok(Transformed::yes(Expr::Cast(Cast::new( + Box::new(Expr::Literal(ScalarValue::Utf8(Some(value)), m)), + DataType::Decimal64(precision, scale), + )))) + } + Expr::Literal(ScalarValue::Decimal128(Some(value), precision, scale), m) => { + let value = Decimal128Type::format_decimal(value, precision, scale); + Ok(Transformed::yes(Expr::Cast(Cast::new( + Box::new(Expr::Literal(ScalarValue::Utf8(Some(value)), m)), + DataType::Decimal128(precision, scale), + )))) + } + Expr::Literal(ScalarValue::Decimal256(Some(value), precision, scale), m) => { + let value = Decimal256Type::format_decimal(value, precision, scale); + Ok(Transformed::yes(Expr::Cast(Cast::new( + Box::new(Expr::Literal(ScalarValue::Utf8(Some(value)), m)), + DataType::Decimal256(precision, scale), + )))) + } + Expr::InList(in_list) if in_list.list.is_empty() => Ok(Transformed::yes( + Expr::Literal(ScalarValue::Boolean(Some(in_list.negated)), None), + )), other => Ok(Transformed::no(other)), }) .map_err(|e| crate::Error::InvalidInput { @@ -105,11 +145,16 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result { .data; let mut sql = run_unparser(&rewritten)?; - for (i, bytes) in bindings.iter().enumerate() { - // The unparser quotes string literals with single quotes, so the - // placeholder appears as `'__lancedb_binary_placeholder___'`. - let quoted = format!("'{}{}__'", BINARY_PLACEHOLDER_PREFIX, i); - sql = sql.replace("ed, &bytes_to_hex_sql(bytes)); + for (placeholder, bytes) in binary_bindings { + // Each placeholder is unique and absent from every user string, so a + // single replacement cannot rewrite an unrelated literal. + let quoted = format!("'{placeholder}'"); + if !sql.contains("ed) { + return Err(crate::Error::InvalidInput { + message: "failed to bind binary literal while serializing expression".to_string(), + }); + } + sql = sql.replacen("ed, &bytes_to_hex_sql(&bytes), 1); } Ok(sql) }