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>;