fix(python): preserve typed expression round trips

python/python/lancedb/_lancedb.pyi

@@ -579,6 +579,7 @@ class PyQueryRequest:
     filter: Optional[Union[str, bytes]]
     full_text_search: Optional[FullTextQuery]
     select: Optional[Union[str, List[str]]]
+    select_source_columns: Optional[Dict[str, str]]
     fast_search: Optional[bool]
     with_row_id: Optional[bool]
     use_lsm: Optional[bool]

python/python/lancedb/expr.py

@@ -315,7 +315,7 @@ def func(name: str, *args: ExprLike) -> Expr:
     --------
     >>> from lancedb.expr import col, func
     >>> func("lower", col("name"))
-    Expr(lower(name))
+    Expr(lower(`name`))
     """
     inner_args = [_coerce(a)._inner for a in args]
     return Expr(expr_func(name, inner_args))

python/python/lancedb/query.py

@@ -2777,15 +2777,20 @@ 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
+        )
         self._blob_auto_row_id = blob_auto_row_id_for_scan(
             schema,
-            req.select,
+            projection,
             with_row_id=self._with_row_id,
         )
         if not self._blob_auto_row_id:
             self._blob_paths = ()
             return
-        self._blob_paths = tuple(blob_v2_projection_sources(schema, req.select).keys())
+        self._blob_paths = tuple(blob_v2_projection_sources(schema, projection).keys())
         self._inner.with_row_id()
 
     def select(self, columns: Union[List[str], dict[str, str]]) -> Self:

python/python/tests/test_blob.py

@@ -179,6 +179,20 @@ async def test_async_table_to_pandas_descriptions_mode_omits_row_id():
     assert set(descriptor.keys()) == {"kind", "position", "size", "blob_id", "blob_uri"}
 
 
+@pytest.mark.asyncio
+async def test_async_typed_blob_projection_preserves_source_column():
+    db = await lancedb.connect_async("memory:///typed_blob_projection")
+    schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("blob")])
+    table = await db.create_table("typed_blob_projection", schema=schema)
+    await table.add([{"id": 1, "blob": b"alpha"}])
+
+    hits = await table.query().select({"blob_alias": col("blob")}).to_arrow()
+
+    assert "_lance_row_id" in hits.schema.field("blob_alias").type.names
+    blobs = await table.fetch_blobs("blob", hits)
+    assert blobs.to_pylist() == [b"alpha"]
+
+
 def test_fetch_blobs_round_trip():
     table = _blob_table(
         "round_trip",

python/python/tests/test_table.py

@@ -2366,6 +2366,55 @@ def test_update_expr_filter_preserves_typed_semantics(mem_db: DBConnection):
     result = binary_table.update(where=predicate, values={"result": "new"})
     assert result.rows_updated == 2
 
+    nonfinite_table = mem_db.create_table(
+        "update_expr_nonfinite",
+        [{"x": 1.0, "result": "old"}, {"x": 2.0, "result": "old"}],
+    )
+    predicate = col("x") < float("inf")
+    assert nonfinite_table.search().where(predicate).to_arrow().num_rows == 2
+    result = nonfinite_table.update(where=predicate, values={"result": "new"})
+    assert result.rows_updated == 2
+
+    float16_table = mem_db.create_table(
+        "update_expr_float16",
+        [{"x": 1.0, "result": "old"}, {"x": 3.0, "result": "old"}],
+    )
+    predicate = col("x").cast(pa.float16()) < 2.0
+    assert float16_table.search().where(predicate).to_arrow().num_rows == 1
+    result = float16_table.update(where=predicate, values={"result": "new"})
+    assert result.rows_updated == 1
+
+    decimal256_schema = pa.schema(
+        [("val", pa.decimal256(40, 2)), ("result", pa.string())]
+    )
+    decimal256_table = mem_db.create_table(
+        "update_expr_decimal256",
+        pa.table(
+            {
+                "val": [Decimal("1.00"), Decimal("3.00")],
+                "result": ["old", "old"],
+            },
+            schema=decimal256_schema,
+        ),
+    )
+    predicate = col("val") < lit(Decimal("2.00")).cast(pa.decimal256(40, 2))
+    assert decimal256_table.search().where(predicate).to_arrow().num_rows == 1
+    result = decimal256_table.update(where=predicate, values={"result": "new"})
+    assert result.rows_updated == 1
+
+    binary_empty_table = mem_db.create_table(
+        "update_expr_binary_empty",
+        pa.table(
+            {"payload": [b"\x01", b"\x02"], "result": ["old", "old"]},
+            schema=pa.schema([("payload", pa.binary()), ("result", pa.string())]),
+        ),
+    )
+    predicate = (col("payload") == lit(b"\x01")).isin([])
+    assert binary_empty_table.search().where(predicate).to_arrow().num_rows == 0
+    assert predicate.to_sql() == "false"
+    result = binary_empty_table.update(where=predicate, values={"result": "new"})
+    assert result.rows_updated == 0
+
 
 def test_update_with_arrow_scalar(mem_db: DBConnection):
     schema = pa.schema({"id": pa.int64(), "vector": pa.list_(pa.float32(), 4)})

python/src/query.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
+use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -292,6 +293,7 @@ pub struct PyQueryRequest {
     pub filter: Option<PyQueryFilter>,
     pub full_text_search: Option<PyLanceDB<FtsQuery>>,
     pub select: PySelect,
+    pub select_source_columns: Option<HashMap<String, String>>,
     pub fast_search: Option<bool>,
     pub with_row_id: Option<bool>,
     pub use_lsm: Option<bool>,
@@ -322,6 +324,7 @@ impl From<AnyQuery> for PyQueryRequest {
                 full_text_search: query_request
                     .full_text_search
                     .map(|fts| PyLanceDB(fts.query)),
+                select_source_columns: PySelect::source_columns(&query_request.select),
                 select: PySelect(query_request.select),
                 fast_search: Some(query_request.fast_search),
                 with_row_id: Some(query_request.with_row_id),
@@ -347,6 +350,7 @@ impl From<AnyQuery> for PyQueryRequest {
                 offset: vector_query.base.offset,
                 filter: vector_query.base.filter.map(PyQueryFilter),
                 full_text_search: None,
+                select_source_columns: PySelect::source_columns(&vector_query.base.select),
                 select: PySelect(vector_query.base.select),
                 fast_search: Some(vector_query.base.fast_search),
                 with_row_id: Some(vector_query.base.with_row_id),
@@ -379,6 +383,25 @@ impl From<AnyQuery> for PyQueryRequest {
 #[derive(Clone)]
 pub struct PySelect(Select);
 
+impl PySelect {
+    fn source_columns(select: &Select) -> Option<HashMap<String, String>> {
+        match select {
+            Select::Expr(pairs) => Some(
+                pairs
+                    .iter()
+                    .filter_map(|(output, expr)| match expr {
+                        lancedb::expr::DfExpr::Column(column) if column.relation.is_none() => {
+                            Some((output.clone(), column.name.clone()))
+                        }
+                        _ => None,
+                    })
+                    .collect(),
+            ),
+            _ => None,
+        }
+    }
+}
+
 impl<'py> IntoPyObject<'py> for PySelect {
     type Target = PyAny;
     type Output = Bound<'py, Self::Target>;

rust/lancedb/src/expr.rs

@@ -211,6 +211,17 @@ mod tests {
         assert_eq!(expr_to_sql_string(&expr).unwrap(), "false");
     }
 
+    #[test]
+    fn test_empty_is_in_discards_binary_children() {
+        use datafusion_common::ScalarValue;
+
+        let expr = is_in(
+            col("payload").eq(lit(ScalarValue::Binary(Some(vec![0x01])))),
+            vec![],
+        );
+        assert_eq!(expr_to_sql_string(&expr).unwrap(), "false");
+    }
+
     #[test]
     fn test_keyword_identifier() {
         let expr = col("null").eq(lit(1i64));
@@ -226,8 +237,35 @@ mod tests {
             19,
             18,
         )));
-        let sql = expr_to_sql_string(&expr).unwrap().replace(' ', "");
-        assert_eq!(sql, "(val<CAST('1.234567890123456790'ASDECIMAL(19,18)))");
+        let sql = expr_to_sql_string(&expr).unwrap();
+        assert_eq!(
+            sql,
+            "(val < arrow_cast('1.234567890123456790', 'Decimal128(19, 18)'))"
+        );
+    }
+
+    #[test]
+    fn test_non_finite_float_literal_preserves_type() {
+        let expr = col("x").lt(lit(f64::INFINITY));
+        assert_eq!(
+            expr_to_sql_string(&expr).unwrap(),
+            "(x < arrow_cast('inf', 'Float64'))"
+        );
+    }
+
+    #[test]
+    fn test_cast_uses_arrow_type_name() {
+        let expr = expr_cast(col("x"), DataType::Float16).lt(lit(2.0));
+        assert_eq!(
+            expr_to_sql_string(&expr).unwrap(),
+            "(arrow_cast(x, 'Float16') < 2.0)"
+        );
+
+        let decimal = expr_cast(lit("2.00"), DataType::Decimal256(40, 2));
+        assert_eq!(
+            expr_to_sql_string(&decimal).unwrap(),
+            "arrow_cast('2.00', 'Decimal256(40, 2)')"
+        );
     }
 
     #[test]
@@ -244,6 +282,19 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_binary_placeholder_collision_search_is_linear() {
+        use datafusion_common::ScalarValue;
+
+        let collision_shaped = format!("__lancedb_binary_placeholder_0__{}", "_".repeat(64_000));
+        let expr = col("payload")
+            .eq(lit(ScalarValue::Binary(Some(vec![0x01]))))
+            .and(col("text").eq(lit(collision_shaped.clone())));
+        let sql = expr_to_sql_string(&expr).unwrap();
+        assert!(sql.contains("X'01'"));
+        assert!(sql.contains(&format!("'{collision_shaped}'")));
+    }
+
     #[test]
     fn test_multiple_binary_literals() {
         use datafusion_common::ScalarValue;

rust/lancedb/src/expr/sql.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The LanceDB Authors
 
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 
 use arrow_array::types::{
     Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType,
@@ -9,7 +9,10 @@ use arrow_array::types::{
 use arrow_schema::DataType;
 use datafusion_common::ScalarValue;
 use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
-use datafusion_expr::{Expr, expr::Cast};
+use datafusion_expr::Expr;
+use datafusion_functions::core::expr_fn::{
+    arrow_cast as datafusion_arrow_cast, arrow_try_cast as datafusion_arrow_try_cast,
+};
 use datafusion_sql::sqlparser::keywords::ALL_KEYWORDS;
 use datafusion_sql::unparser::{self, dialect::Dialect};
 
@@ -65,6 +68,113 @@ fn string_literals(expr: &Expr) -> HashSet<String> {
     literals
 }
 
+fn typed_string_literal(value: String, data_type: DataType) -> Expr {
+    datafusion_arrow_cast(
+        Expr::Literal(ScalarValue::Utf8(Some(value)), None),
+        Expr::Literal(ScalarValue::Utf8(Some(data_type.to_string())), None),
+    )
+}
+
+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<String>, next_id: &mut usize) -> String {
+    loop {
+        let placeholder = format!("{BINARY_PLACEHOLDER_PREFIX}{}__", *next_id);
+        *next_id += 1;
+        if !user_strings.contains(&placeholder) {
+            return placeholder;
+        }
+    }
+}
+
+fn bind_binary_literals(
+    sql: &str,
+    mut bindings: HashMap<String, Vec<u8>>,
+) -> crate::Result<String> {
+    let bytes = sql.as_bytes();
+    let mut output = Vec::with_capacity(bytes.len());
+    let mut index = 0;
+
+    // Walk SQL string tokens once. Placeholders are plain, unescaped string
+    // 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'\'' {
+            output.push(bytes[index]);
+            index += 1;
+            continue;
+        }
+
+        let literal_start = index;
+        index += 1;
+        let content_start = index;
+        let mut escaped = false;
+        let mut content_end = None;
+        while index < bytes.len() {
+            if bytes[index] == b'\'' {
+                if index + 1 < bytes.len() && bytes[index + 1] == b'\'' {
+                    escaped = true;
+                    index += 2;
+                } else {
+                    content_end = Some(index);
+                    index += 1;
+                    break;
+                }
+            } else {
+                index += 1;
+            }
+        }
+
+        let Some(content_end) = content_end else {
+            return Err(crate::Error::InvalidInput {
+                message: "unterminated string while binding binary literal".to_string(),
+            });
+        };
+
+        let placeholder = &sql[content_start..content_end];
+        if !escaped && let Some(value) = bindings.remove(placeholder) {
+            output.extend_from_slice(bytes_to_hex_sql(&value).as_bytes());
+        } else {
+            output.extend_from_slice(&bytes[literal_start..index]);
+        }
+    }
+
+    if !bindings.is_empty() {
+        return Err(crate::Error::InvalidInput {
+            message: "failed to bind binary literal while serializing expression".to_string(),
+        });
+    }
+
+    String::from_utf8(output).map_err(|e| crate::Error::InvalidInput {
+        message: format!("failed to bind binary literal: {e}"),
+    })
+}
+
 fn run_unparser(expr: &Expr) -> crate::Result<String> {
     let ast = unparser::Unparser::new(&LanceSqlDialect)
         .expr_to_sql(expr)
@@ -81,22 +191,30 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
     // * 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<u8>)> = Vec::new();
+    // Eliminate empty membership expressions before visiting their children.
+    // Otherwise a discarded binary child could leave behind a stale binding.
     let rewritten = expr
         .clone()
+        .transform(|e: Expr| match e {
+            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 {
+            message: format!("failed to rewrite expression: {e}"),
+        })?
+        .data;
+
+    let user_strings = string_literals(&rewritten);
+    let mut next_placeholder_id = 0;
+    let mut binary_bindings = HashMap::new();
+    let rewritten = rewritten
         .transform(|e: Expr| match e {
             Expr::Literal(ScalarValue::Binary(Some(bytes)), m)
             | Expr::Literal(ScalarValue::LargeBinary(Some(bytes)), m) => {
-                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));
+                let placeholder = next_binary_placeholder(&user_strings, &mut next_placeholder_id);
+                binary_bindings.insert(placeholder.clone(), bytes);
                 Ok(Transformed::yes(Expr::Literal(
                     ScalarValue::Utf8(Some(placeholder)),
                     m,
@@ -106,37 +224,61 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
             | Expr::Literal(ScalarValue::LargeBinary(None), m) => {
                 Ok(Transformed::yes(Expr::Literal(ScalarValue::Null, m)))
             }
-            Expr::Literal(ScalarValue::Decimal32(Some(value), precision, scale), 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)),
+                Ok(Transformed::yes(typed_string_literal(
+                    value,
                     DataType::Decimal32(precision, scale),
-                ))))
+                )))
             }
-            Expr::Literal(ScalarValue::Decimal64(Some(value), precision, scale), m) => {
+            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)),
+                Ok(Transformed::yes(typed_string_literal(
+                    value,
                     DataType::Decimal64(precision, scale),
-                ))))
+                )))
             }
-            Expr::Literal(ScalarValue::Decimal128(Some(value), precision, scale), m) => {
+            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)),
+                Ok(Transformed::yes(typed_string_literal(
+                    value,
                     DataType::Decimal128(precision, scale),
-                ))))
+                )))
             }
-            Expr::Literal(ScalarValue::Decimal256(Some(value), precision, scale), m) => {
+            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)),
+                Ok(Transformed::yes(typed_string_literal(
+                    value,
                     DataType::Decimal256(precision, scale),
-                ))))
+                )))
+            }
+            Expr::Literal(ScalarValue::Float16(Some(value)), _m) if !value.is_finite() => Ok(
+                Transformed::yes(typed_string_literal(value.to_string(), DataType::Float16)),
+            ),
+            Expr::Literal(ScalarValue::Float32(Some(value)), _m) if !value.is_finite() => Ok(
+                Transformed::yes(typed_string_literal(value.to_string(), DataType::Float32)),
+            ),
+            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::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 {
@@ -144,17 +286,10 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result<String> {
         })?
         .data;
 
-    let mut sql = run_unparser(&rewritten)?;
-    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(&quoted) {
-            return Err(crate::Error::InvalidInput {
-                message: "failed to bind binary literal while serializing expression".to_string(),
-            });
-        }
-        sql = sql.replacen(&quoted, &bytes_to_hex_sql(&bytes), 1);
+    let sql = run_unparser(&rewritten)?;
+    if binary_bindings.is_empty() {
+        Ok(sql)
+    } else {
+        bind_binary_literals(&sql, binary_bindings)
     }
-    Ok(sql)
 }