From 71a2a0c7200b7c8096daf99d26dd2f4c64782bfb Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 12:16:54 +0000 Subject: [PATCH] fix: preserve typed expression boundaries --- python/python/lancedb/query.py | 17 +++--- python/python/tests/test_blob.py | 44 ++++++++++++++++ python/python/tests/test_expr.py | 12 ++--- python/python/tests/test_table.py | 24 +++++++++ rust/lancedb/src/expr.rs | 26 +++++++++ rust/lancedb/src/expr/sql.rs | 87 +++++++++++++++---------------- 6 files changed, 152 insertions(+), 58 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index c3cf11a04..0003f510b 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -169,6 +169,12 @@ def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]: return {"columns": projection} +def _query_request_projection(req: "PyQueryRequest") -> QueryProjection: + if req.select_source_columns is not None: + return req.select_source_columns + return req.select + + def _scanner_kwargs_for_query( query: Query, blob_mode: BlobMode, @@ -2777,11 +2783,7 @@ class AsyncQueryBase(object): req = self._inner.to_query_request() schema = await self._table.schema() - projection = ( - req.select_source_columns - if req.select_source_columns is not None - else req.select - ) + projection = _query_request_projection(req) self._blob_auto_row_id = blob_auto_row_id_for_scan( schema, projection, @@ -3881,14 +3883,15 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): blob_paths: tuple[str, ...] = () if self._table is not None: schema = await self._table.schema() + projection = _query_request_projection(req) blob_auto_row_id = blob_auto_row_id_for_scan( schema, - req.select, + projection, with_row_id=self._with_row_id, ) if blob_auto_row_id: blob_paths = tuple( - blob_v2_projection_sources(schema, req.select).keys() + blob_v2_projection_sources(schema, projection).keys() ) self._blob_auto_row_id = blob_auto_row_id self._blob_paths = blob_paths diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 551fdbb4d..ded352f0f 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -430,6 +430,50 @@ async def test_blob_v2_hybrid_fetch_blobs_async(): assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"} +@pytest.mark.asyncio +async def test_async_hybrid_typed_blob_projection_preserves_source_column(): + db = await lancedb.connect_async("memory:///hybrid_typed_blob") + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("text", pa.utf8()), + pa.field("vector", pa.list_(pa.float32(), list_size=2)), + lancedb.blob("blob"), + ] + ) + table = await db.create_table("hybrid_typed_blob", schema=schema) + await table.add( + [ + { + "id": 1, + "text": "hello alpha", + "vector": [1.0, 0.0], + "blob": b"alpha", + }, + { + "id": 2, + "text": "hello beta", + "vector": [0.9, 0.1], + "blob": b"beta", + }, + ] + ) + await table.create_index("text", config=FTS(with_position=False)) + + hits = await ( + table.query() + .nearest_to([1.0, 0.0]) + .nearest_to_text("hello") + .select({"blob_alias": col("blob")}) + .limit(2) + .to_arrow() + ) + + assert "_lance_row_id" in hits.schema.field("blob_alias").type.names + blobs = await table.fetch_blobs("blob", hits) + assert {blobs[i].as_py() for i in range(len(blobs))} == {b"alpha", b"beta"} + + def test_blob_file_seek_read_and_read_range(): payload = _identifiable_payload(1024) table = _blob_table("seek_read", [{"id": 1, "image": payload}]) diff --git a/python/python/tests/test_expr.py b/python/python/tests/test_expr.py index 0995a39cf..dab7d8a3c 100644 --- a/python/python/tests/test_expr.py +++ b/python/python/tests/test_expr.py @@ -251,32 +251,32 @@ class TestExprCast: def test_cast_string(self): e = col("id").cast("string") assert isinstance(e, Expr) - assert e.to_sql() == "CAST(id AS VARCHAR)" + assert e.to_sql() == "arrow_cast(id, 'Utf8')" def test_cast_int32(self): e = col("score").cast("int32") assert isinstance(e, Expr) - assert e.to_sql() == "CAST(score AS INTEGER)" + assert e.to_sql() == "arrow_cast(score, 'Int32')" def test_cast_float64(self): e = col("val").cast("float64") assert isinstance(e, Expr) - assert e.to_sql() == "CAST(val AS DOUBLE)" + assert e.to_sql() == "arrow_cast(val, 'Float64')" def test_cast_pyarrow_type(self): e = col("score").cast(pa.int32()) assert isinstance(e, Expr) - assert e.to_sql() == "CAST(score AS INTEGER)" + assert e.to_sql() == "arrow_cast(score, 'Int32')" def test_cast_pyarrow_float64(self): e = col("val").cast(pa.float64()) assert isinstance(e, Expr) - assert e.to_sql() == "CAST(val AS DOUBLE)" + assert e.to_sql() == "arrow_cast(val, 'Float64')" def test_cast_pyarrow_string(self): e = col("id").cast(pa.string()) assert isinstance(e, Expr) - assert e.to_sql() == "CAST(id AS VARCHAR)" + assert e.to_sql() == "arrow_cast(id, 'Utf8')" def test_cast_pyarrow_and_string_equivalent(self): # pa.int32() and "int32" should produce equivalent SQL diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index a33ff473a..600556489 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -2384,6 +2384,30 @@ def test_update_expr_filter_preserves_typed_semantics(mem_db: DBConnection): result = float16_table.update(where=predicate, values={"result": "new"}) assert result.rows_updated == 1 + string_cast_table = mem_db.create_table( + "update_expr_string_cast", + [{"x": 1, "result": "old"}, {"x": 2, "result": "old"}], + ) + predicate = col("x").cast("string") == "1" + assert string_cast_table.search().where(predicate).to_arrow().num_rows == 1 + result = string_cast_table.update(where=predicate, values={"result": "new"}) + assert result.rows_updated == 1 + + quoted_identifier_schema = pa.schema( + [("payload", pa.binary()), ("odd'name", pa.int64()), ("result", pa.string())] + ) + quoted_identifier_table = mem_db.create_table( + "update_expr_quoted_identifier", + pa.table( + {"payload": [b"\x01"], "odd'name": [1], "result": ["old"]}, + schema=quoted_identifier_schema, + ), + ) + predicate = (col("payload") == lit(b"\x01")) & (col("odd'name") == 1) + assert quoted_identifier_table.search().where(predicate).to_arrow().num_rows == 1 + result = quoted_identifier_table.update(where=predicate, values={"result": "new"}) + assert result.rows_updated == 1 + decimal256_schema = pa.schema( [("val", pa.decimal256(40, 2)), ("result", pa.string())] ) diff --git a/rust/lancedb/src/expr.rs b/rust/lancedb/src/expr.rs index 6c2acdd01..1d489c5ba 100644 --- a/rust/lancedb/src/expr.rs +++ b/rust/lancedb/src/expr.rs @@ -255,6 +255,18 @@ mod tests { #[test] fn test_cast_uses_arrow_type_name() { + let string = expr_cast(col("x"), DataType::Utf8); + assert_eq!( + expr_to_sql_string(&string).unwrap(), + "arrow_cast(x, 'Utf8')" + ); + + let int32 = expr_cast(col("x"), DataType::Int32); + assert_eq!( + expr_to_sql_string(&int32).unwrap(), + "arrow_cast(x, 'Int32')" + ); + let expr = expr_cast(col("x"), DataType::Float16).lt(lit(2.0)); assert_eq!( expr_to_sql_string(&expr).unwrap(), @@ -282,6 +294,20 @@ mod tests { ); } + #[test] + fn test_binary_binding_skips_quoted_identifiers() { + use datafusion_common::ScalarValue; + + let expr = col("payload") + .eq(lit(ScalarValue::Binary(Some(vec![0x01])))) + .and(col("odd'name").eq(lit(1i64))) + .and(col("odd`'name").eq(lit(2i64))); + assert_eq!( + expr_to_sql_string(&expr).unwrap(), + "(((payload = X'01') AND (`odd'name` = 1)) AND (`odd``'name` = 2))" + ); + } + #[test] fn test_binary_placeholder_collision_search_is_linear() { use datafusion_common::ScalarValue; diff --git a/rust/lancedb/src/expr/sql.rs b/rust/lancedb/src/expr/sql.rs index 904f4d381..e8817fcd4 100644 --- a/rust/lancedb/src/expr/sql.rs +++ b/rust/lancedb/src/expr/sql.rs @@ -75,33 +75,6 @@ fn typed_string_literal(value: String, data_type: DataType) -> Expr { ) } -fn cast_requires_arrow_type_name(data_type: &DataType) -> bool { - !matches!( - data_type, - DataType::Boolean - | DataType::Int8 - | DataType::Int16 - | DataType::Int32 - | DataType::Int64 - | DataType::UInt8 - | DataType::UInt16 - | DataType::UInt32 - | DataType::UInt64 - | DataType::Float32 - | DataType::Float64 - | DataType::Timestamp(_, _) - | DataType::Date32 - | DataType::Date64 - | DataType::Interval(_) - | DataType::Utf8 - | DataType::LargeUtf8 - | DataType::Utf8View - | DataType::Decimal32(_, _) - | DataType::Decimal64(_, _) - | DataType::Decimal128(_, _) - ) -} - fn next_binary_placeholder(user_strings: &HashSet, next_id: &mut usize) -> String { loop { let placeholder = format!("{BINARY_PLACEHOLDER_PREFIX}{}__", *next_id); @@ -124,6 +97,33 @@ fn bind_binary_literals( // literals, so this remains linear even when user strings are large or // deliberately resemble the placeholder prefix. while index < bytes.len() { + if bytes[index] == b'`' { + let identifier_start = index; + index += 1; + let mut identifier_end = None; + while index < bytes.len() { + if bytes[index] == b'`' { + if index + 1 < bytes.len() && bytes[index + 1] == b'`' { + index += 2; + } else { + index += 1; + identifier_end = Some(index); + break; + } + } else { + index += 1; + } + } + + let Some(identifier_end) = identifier_end else { + return Err(crate::Error::InvalidInput { + message: "unterminated identifier while binding binary literal".to_string(), + }); + }; + output.extend_from_slice(&bytes[identifier_start..identifier_end]); + continue; + } + if bytes[index] != b'\'' { output.push(bytes[index]); index += 1; @@ -189,6 +189,7 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result { // reparsed by Lance without changing the typed expression's semantics: // // * decimal literals need an explicit cast to preserve precision and scale; + // * casts need exact Arrow type names rather than SQL type aliases; // * an empty IN list is valid in DataFusion but invalid SQL; // * binary literals are unsupported by the unparser and need placeholders. // Eliminate empty membership expressions before visiting their children. @@ -261,24 +262,20 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result { Expr::Literal(ScalarValue::Float64(Some(value)), _m) if !value.is_finite() => Ok( Transformed::yes(typed_string_literal(value.to_string(), DataType::Float64)), ), - Expr::Cast(cast) if cast_requires_arrow_type_name(cast.field.data_type()) => { - Ok(Transformed::yes(datafusion_arrow_cast( - *cast.expr, - Expr::Literal( - ScalarValue::Utf8(Some(cast.field.data_type().to_string())), - None, - ), - ))) - } - Expr::TryCast(cast) if cast_requires_arrow_type_name(cast.field.data_type()) => { - Ok(Transformed::yes(datafusion_arrow_try_cast( - *cast.expr, - Expr::Literal( - ScalarValue::Utf8(Some(cast.field.data_type().to_string())), - None, - ), - ))) - } + Expr::Cast(cast) => Ok(Transformed::yes(datafusion_arrow_cast( + *cast.expr, + Expr::Literal( + ScalarValue::Utf8(Some(cast.field.data_type().to_string())), + None, + ), + ))), + Expr::TryCast(cast) => Ok(Transformed::yes(datafusion_arrow_try_cast( + *cast.expr, + Expr::Literal( + ScalarValue::Utf8(Some(cast.field.data_type().to_string())), + None, + ), + ))), other => Ok(Transformed::no(other)), }) .map_err(|e| crate::Error::InvalidInput {