diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 9ba991322..f9265e783 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/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] diff --git a/python/python/lancedb/expr.py b/python/python/lancedb/expr.py index 752bb2fa0..a1d542e60 100644 --- a/python/python/lancedb/expr.py +++ b/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)) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 095a7b5ff..c3cf11a04 100644 --- a/python/python/lancedb/query.py +++ b/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: diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index fc8dd0296..551fdbb4d 100644 --- a/python/python/tests/test_blob.py +++ b/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", diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 62cc7f878..a33ff473a 100644 --- a/python/python/tests/test_table.py +++ b/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)}) diff --git a/python/src/query.rs b/python/src/query.rs index affbdf4fd..23c6d5cdd 100644 --- a/python/src/query.rs +++ b/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, pub full_text_search: Option>, pub select: PySelect, + pub select_source_columns: Option>, pub fast_search: Option, pub with_row_id: Option, pub use_lsm: Option, @@ -322,6 +324,7 @@ impl From 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 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 for PyQueryRequest { #[derive(Clone)] pub struct PySelect(Select); +impl PySelect { + fn source_columns(select: &Select) -> Option> { + 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>; diff --git a/rust/lancedb/src/expr.rs b/rust/lancedb/src/expr.rs index 84174fc14..6c2acdd01 100644 --- a/rust/lancedb/src/expr.rs +++ b/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 HashSet { 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, 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>, +) -> crate::Result { + 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 { let ast = unparser::Unparser::new(&LanceSqlDialect) .expr_to_sql(expr) @@ -81,22 +191,30 @@ pub fn expr_to_sql_string(expr: &Expr) -> crate::Result { // * 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(); + // 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 { | 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 { })? .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("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); + let sql = run_unparser(&rewritten)?; + if binary_bindings.is_empty() { + Ok(sql) + } else { + bind_binary_literals(&sql, binary_bindings) } - Ok(sql) }