From 25645d82d42e27fc4db8c386aa3decac7b4f2f97 Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 20:28:21 +0800 Subject: [PATCH 1/9] feat(python): accept expressions in update filters (#3876) ## Summary - allow Python sync, async, and remote table updates to accept type-safe `Expr` filters - serialize expression filters before invoking the existing update implementation - cover numeric-looking text and apostrophe-containing text in sync and async regression tests ## Root cause `Table.update` was the remaining Python write path that required callers to construct a raw SQL predicate. Dynamic text interpolated without SQL literal encoding could therefore be parsed as an integer, float, or unterminated string instead of Utf8. The expression API already encodes literals safely for query and delete filters. ## Validation - `cd python && .venv/bin/pytest python/tests/test_table.py::test_update_async python/tests/test_table.py::test_update_expr_filter_literals_async python/tests/test_table.py::test_update python/tests/test_table.py::test_update_expr_filter_literals -q` - `cd python && .venv/bin/pytest python/tests/test_expr.py -q` - `cd python && .venv/bin/ruff format --check .` - `cd python && .venv/bin/ruff check .` Fixes #1869 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- python/python/lancedb/_blob.py | 6 +- python/python/lancedb/_lancedb.pyi | 2 + python/python/lancedb/expr.py | 6 +- python/python/lancedb/query.py | 16 +- python/python/lancedb/remote/table.py | 11 +- python/python/lancedb/table.py | 42 +++-- python/python/tests/test_blob.py | 73 ++++++- python/python/tests/test_expr.py | 42 ++--- python/python/tests/test_table.py | 158 ++++++++++++++++ python/src/expr.rs | 8 + python/src/query.rs | 23 +++ rust/lancedb/src/expr.rs | 124 +++++++++++- rust/lancedb/src/expr/sql.rs | 262 +++++++++++++++++++++----- 13 files changed, 678 insertions(+), 95 deletions(-) diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index 926769f48..5b4c0c343 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -270,7 +270,8 @@ def _iter_projection_pairs( if isinstance(expr, str): yield name, expr elif isinstance(expr, Expr): - yield name, expr.to_sql() + source = expr._column_name() + yield name, source if source is not None else expr.to_sql() return for column in projection: if isinstance(column, str): @@ -280,7 +281,8 @@ def _iter_projection_pairs( if isinstance(expr, str): yield name, expr elif isinstance(expr, Expr): - yield name, expr.to_sql() + source = expr._column_name() + yield name, source if source is not None else expr.to_sql() def _set_blob_column(tbl: pa.Table, output_name: str, blobs: pa.Array) -> pa.Table: diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 593bceffa..7d7ca7f2a 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -87,6 +87,7 @@ class PyExpr: def contains(self, substr: "PyExpr") -> "PyExpr": ... def isin(self, values: List["PyExpr"]) -> "PyExpr": ... def cast(self, data_type: pa.DataType) -> "PyExpr": ... + def column_name(self) -> Optional[str]: ... def to_sql(self) -> str: ... def expr_col(name: str) -> PyExpr: ... @@ -608,6 +609,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 d16ba95d7..80d01e29a 100644 --- a/python/python/lancedb/expr.py +++ b/python/python/lancedb/expr.py @@ -249,6 +249,10 @@ class Expr: # ── utilities ──────────────────────────────────────────────────────────── + def _column_name(self) -> str | None: + """Return the source name when this is a bare column expression.""" + return self._inner.column_name() + def to_sql(self) -> str: """Render the expression as a SQL string (useful for debugging).""" return self._inner.to_sql() @@ -312,7 +316,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 9301d7df8..451384ad1 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -167,6 +167,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, @@ -2799,15 +2805,16 @@ class AsyncQueryBase(object): req = self._inner.to_query_request() schema = await self._table.schema() + projection = _query_request_projection(req) 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: @@ -3894,14 +3901,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/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index d9139396b..02748b9bc 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -36,6 +36,7 @@ from lancedb._lancedb import ( UpdateResult, ) from lancedb.embeddings.base import EmbeddingFunctionConfig +from lancedb.expr import Expr from lancedb.index import ( FTS, BTree, @@ -863,7 +864,7 @@ class RemoteTable(Table): def update( self, - where: Optional[str] = None, + where: Optional[Union[str, Expr]] = None, values: Optional[dict] = None, *, values_sql: Optional[Dict[str, str]] = None, @@ -874,9 +875,11 @@ class RemoteTable(Table): Parameters ---------- - where: str, optional - The SQL where clause to use when updating rows. For example, 'x = 2' - or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error. + where: str or [Expr][lancedb.expr.Expr], optional + The filter condition. Can be a SQL string or a type-safe + [Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and + [lit][lancedb.expr.lit]. The filter must not be empty, or it will + error. values: dict, optional The values to update. The keys are the column names and the values are the values to set. diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index c354a944e..765b7fa14 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -1744,7 +1744,7 @@ class Table(ABC): @abstractmethod def update( self, - where: Optional[str] = None, + where: Optional[Union[str, Expr]] = None, values: Optional[dict] = None, *, values_sql: Optional[Dict[str, str]] = None, @@ -1759,9 +1759,11 @@ class Table(ABC): Parameters ---------- - where: str, optional - The SQL where clause to use when updating rows. For example, 'x = 2' - or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error. + where: str or [Expr][lancedb.expr.Expr], optional + The filter condition. Can be a SQL string or a type-safe + [Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and + [lit][lancedb.expr.lit]. The filter must not be empty, or it will + error. values: dict, optional The values to update. The keys are the column names and the values are the values to set. @@ -1779,6 +1781,7 @@ class Table(ABC): Examples -------- >>> import lancedb + >>> from lancedb.expr import col >>> import pandas as pd >>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]}) >>> db = lancedb.connect("./.lancedb") @@ -1788,7 +1791,7 @@ class Table(ABC): 0 1 [1.0, 2.0] 1 2 [3.0, 4.0] 2 3 [5.0, 6.0] - >>> table.update(where="x = 2", values={"vector": [10.0, 10]}) + >>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]}) UpdateResult(rows_updated=1, version=2) >>> table.to_pandas() x vector @@ -3841,7 +3844,7 @@ class LanceTable(Table): def update( self, - where: Optional[str] = None, + where: Optional[Union[str, Expr]] = None, values: Optional[dict] = None, *, values_sql: Optional[Dict[str, str]] = None, @@ -3852,9 +3855,11 @@ class LanceTable(Table): Parameters ---------- - where: str, optional - The SQL where clause to use when updating rows. For example, 'x = 2' - or 'x IN (1, 2, 3)'. The filter must not be empty, or it will error. + where: str or [Expr][lancedb.expr.Expr], optional + The filter condition. Can be a SQL string or a type-safe + [Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and + [lit][lancedb.expr.lit]. The filter must not be empty, or it will + error. values: dict, optional The values to update. The keys are the column names and the values are the values to set. @@ -3872,6 +3877,7 @@ class LanceTable(Table): Examples -------- >>> import lancedb + >>> from lancedb.expr import col >>> import pandas as pd >>> data = pd.DataFrame({"x": [1, 2, 3], "vector": [[1.0, 2], [3, 4], [5, 6]]}) >>> db = lancedb.connect("./.lancedb") @@ -3881,7 +3887,7 @@ class LanceTable(Table): 0 1 [1.0, 2.0] 1 2 [3.0, 4.0] 2 3 [5.0, 6.0] - >>> table.update(where="x = 2", values={"vector": [10.0, 10]}) + >>> table.update(where=col("x") == 2, values={"vector": [10.0, 10]}) UpdateResult(rows_updated=1, version=2) >>> table.to_pandas() x vector @@ -5995,7 +6001,7 @@ class AsyncTable: self, updates: Optional[Dict[str, Any]] = None, *, - where: Optional[str] = None, + where: Optional[Union[str, Expr]] = None, updates_sql: Optional[Dict[str, str]] = None, ) -> UpdateResult: """ @@ -6010,9 +6016,11 @@ class AsyncTable: The updates to apply. The keys should be the name of the column to update. The values should be the new values to assign. This is required unless updates_sql is supplied. - where: str, optional - An SQL filter that controls which rows are updated. For example, 'x = 2' - or 'x IN (1, 2, 3)'. Only rows that satisfy this filter will be udpated. + where: str or [Expr][lancedb.expr.Expr], optional + The filter condition. Can be a SQL string or a type-safe + [Expr][lancedb.expr.Expr] built with [col][lancedb.expr.col] and + [lit][lancedb.expr.lit]. Only rows that satisfy this filter will + be updated. updates_sql: dict, optional The updates to apply, expressed as SQL expression strings. The keys should be column names. The values should be SQL expressions. These can be SQL @@ -6030,13 +6038,14 @@ class AsyncTable: -------- >>> import asyncio >>> import lancedb + >>> from lancedb.expr import col >>> import pandas as pd >>> async def demo_update(): ... data = pd.DataFrame({"x": [1, 2], "vector": [[1, 2], [3, 4]]}) ... db = await lancedb.connect_async("./.lancedb") ... table = await db.create_table("my_table", data) ... # x is [1, 2], vector is [[1, 2], [3, 4]] - ... await table.update({"vector": [10, 10]}, where="x = 2") + ... await table.update({"vector": [10, 10]}, where=col("x") == 2) ... # x is [1, 2], vector is [[1, 2], [10, 10]] ... await table.update(updates_sql={"x": "x + 1"}) ... # x is [2, 3], vector is [[1, 2], [10, 10]] @@ -6050,7 +6059,8 @@ class AsyncTable: if updates is not None: updates_sql = {k: value_to_sql(v) for k, v in updates.items()} - return await self._inner.update(updates_sql, where) + predicate = where.to_sql() if isinstance(where, Expr) else where + return await self._inner.update(updates_sql, predicate) async def add_columns( self, diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 5d7682f24..351769ff6 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -8,7 +8,12 @@ import pyarrow.compute as pc import pytest import lancedb -from lancedb._blob import read_row_ids_from_hits, stash_auto_row_ids +from lancedb._blob import ( + blob_v2_projection_sources, + read_row_ids_from_hits, + stash_auto_row_ids, +) +from lancedb.expr import col from lancedb.index import FTS from lancedb.schema import blob_column_paths, blob_v2_column_paths @@ -70,6 +75,14 @@ def test_blob_v2_column_paths_include_list_children(): ] +def test_blob_v2_projection_sources_use_typed_column_name(): + schema = pa.schema([lancedb.blob("blob")]) + + assert blob_v2_projection_sources(schema, {"blob_alias": col("blob")}) == { + "blob_alias": "blob" + } + + def _legacy_v1_table(name): db = lancedb.connect("memory:///") schema = pa.schema( @@ -166,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", @@ -403,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 0eb6f8929..0f49231f1 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,85 +198,85 @@ 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: 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 @@ -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 56e0eacfd..0be9e139d 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -11,6 +11,7 @@ import warnings import weakref 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 @@ -336,6 +337,21 @@ async def test_update_async(mem_db_async: AsyncConnection): assert await table.count_rows("id == 10") == 1 +@pytest.mark.asyncio +async def test_update_expr_filter_literals_async(mem_db_async: AsyncConnection): + values = ["5", "4.66e-84", "it's"] + table = await mem_db_async.create_table( + "update_expr_literals", + data=[{"field": value, "result": "original"} for value in values], + ) + + for value in values: + update_res = await table.update({"result": value}, where=col("field") == value) + assert update_res.rows_updated == 1 + + assert (await table.to_arrow())["result"].to_pylist() == values + + def test_create_table(mem_db: DBConnection): schema = pa.schema( { @@ -2343,6 +2359,148 @@ def test_update(mem_db: DBConnection): assert np.allclose(v, np.array([[1.2, 1.9], [1.1, 1.1]])) +def test_update_expr_filter_literals(mem_db: DBConnection): + values = ["5", "4.66e-84", "it's"] + table = mem_db.create_table( + "update_expr_literals", + data=[{"field": value, "result": "original"} for value in values], + ) + + for value in values: + update_res = table.update(where=col("field") == value, values={"result": value}) + assert update_res.rows_updated == 1 + + 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 + + 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 + + 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())] + ) + 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)}) table = mem_db.create_table("my_table", schema=schema) diff --git a/python/src/expr.rs b/python/src/expr.rs index eae1d96ec..79b448fdf 100644 --- a/python/src/expr.rs +++ b/python/src/expr.rs @@ -130,6 +130,14 @@ impl PyExpr { // ── utilities ──────────────────────────────────────────────────────────── + /// Return the referenced column name for a bare column expression. + fn column_name(&self) -> Option { + match &self.0 { + DfExpr::Column(column) if column.relation.is_none() => Some(column.name.clone()), + _ => None, + } + } + /// Render the expression as a SQL string (useful for debugging). fn to_sql(&self) -> PyResult { lancedb::expr::expr_to_sql_string(&self.0).map_err(|e| PyValueError::new_err(e.to_string())) diff --git a/python/src/query.rs b/python/src/query.rs index 014e79e2d..38153729f 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; @@ -325,6 +326,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, @@ -355,6 +357,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), @@ -380,6 +383,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), @@ -412,6 +416,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 da69914e3..1625d9632 100644 --- a/rust/lancedb/src/expr.rs +++ b/rust/lancedb/src/expr.rs @@ -157,7 +157,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] @@ -167,7 +167,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] @@ -185,7 +185,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] @@ -196,7 +196,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] @@ -206,6 +206,122 @@ 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_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)); + 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(); + 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 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(), + "(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] + fn test_binary_placeholder_does_not_rewrite_user_string() { + use datafusion_common::ScalarValue; + + let marker = "__lancedb_binary_placeholder_0__"; + let expr = col("payload") + .eq(lit(ScalarValue::Binary(Some(vec![0x01])))) + .or(col("text").eq(lit(marker))); + assert_eq!( + expr_to_sql_string(&expr).unwrap(), + "((payload = X'01') OR (`text` = '__lancedb_binary_placeholder_0__'))" + ); + } + + #[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; + + 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; diff --git a/rust/lancedb/src/expr/sql.rs b/rust/lancedb/src/expr/sql.rs index 24a676485..2a1ca201d 100644 --- a/rust/lancedb/src/expr/sql.rs +++ b/rust/lancedb/src/expr/sql.rs @@ -1,13 +1,24 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use std::any::TypeId; +use std::{ + any::TypeId, + collections::{HashMap, HashSet}, +}; +use arrow_array::types::{ + Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType, +}; +use arrow_schema::DataType; use datafusion_common::ScalarValue; use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; 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::{ dialect::{Dialect as SqlParserDialect, GenericDialect}, + keywords::ALL_KEYWORDS, tokenizer::{Token, Tokenizer}, }; use datafusion_sql::unparser::{self, dialect::Dialect as UnparserDialect}; @@ -27,11 +38,13 @@ struct LanceSqlDialect; impl UnparserDialect for LanceSqlDialect { fn identifier_quote_style(&self, identifier: &str) -> 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 } } } @@ -100,24 +113,128 @@ 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 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 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'`' { + 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; + 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 { @@ -130,25 +247,37 @@ 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; + // * 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. + // 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 placeholder = format!("{}{}__", BINARY_PLACEHOLDER_PREFIX, bindings.len()); - bindings.push(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, @@ -158,6 +287,57 @@ 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(typed_string_literal( + value, + DataType::Decimal32(precision, scale), + ))) + } + Expr::Literal(ScalarValue::Decimal64(Some(value), precision, scale), _m) => { + let value = Decimal64Type::format_decimal(value, precision, scale); + Ok(Transformed::yes(typed_string_literal( + value, + DataType::Decimal64(precision, scale), + ))) + } + Expr::Literal(ScalarValue::Decimal128(Some(value), precision, scale), _m) => { + let value = Decimal128Type::format_decimal(value, precision, scale); + Ok(Transformed::yes(typed_string_literal( + value, + DataType::Decimal128(precision, scale), + ))) + } + Expr::Literal(ScalarValue::Decimal256(Some(value), precision, scale), _m) => { + let value = Decimal256Type::format_decimal(value, precision, scale); + 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) => 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 { @@ -165,14 +345,12 @@ 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)); + let sql = run_unparser(&rewritten)?; + if binary_bindings.is_empty() { + Ok(sql) + } else { + bind_binary_literals(&sql, binary_bindings) } - Ok(sql) } #[cfg(test)] From 9d3962686e847be3e81a48005642b01f7ad9698f Mon Sep 17 00:00:00 2001 From: "lancedb-gatefixer[bot]" <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 27 Aug 2026 23:35:02 +0800 Subject: [PATCH 2/9] fix(node): accept Arrow metadata across JavaScript realms (#3904) ## Summary - accept genuine Arrow metadata maps created in another JavaScript realm - validate every metadata entry and clone it into a local Map - cover an Arrow 15 VM-realm table through the public fromDataToBuffer boundary - retain structural typing for nested and dictionary Arrow data ## Root cause The sanitizer used a local-realm instanceof Map check for schema and field metadata. A genuine Map created in another JavaScript realm has the required internal Map state but fails that identity check, so fromDataToBuffer rejected the foreign table before serializing its rows. ## Scope This fixes the distinct JavaScript-realm sanitizer failure identified during review. It does not establish the cause of the S3/compaction panic reported in #1525, so that issue remains open. ## Validation - pnpm test --runInBand (707 passed, 5 skipped) - pnpm test --runInBand __test__/arrow.test.ts (189 passed) - pnpm build - pnpm lint - pnpm run docs Related to #1525 --------- Co-authored-by: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Co-authored-by: Xuanwo --- nodejs/__test__/arrow.test.ts | 37 +++++++++++++++++++++++++++++++++++ nodejs/lancedb/arrow.ts | 4 ++-- nodejs/lancedb/sanitize.ts | 15 ++++++++++---- 3 files changed, 50 insertions(+), 6 deletions(-) diff --git a/nodejs/__test__/arrow.test.ts b/nodejs/__test__/arrow.test.ts index cb56cb5ae..83b4fae46 100644 --- a/nodejs/__test__/arrow.test.ts +++ b/nodejs/__test__/arrow.test.ts @@ -1,5 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors +import * as fs from "node:fs"; +import * as vm from "node:vm"; import * as arrow15 from "apache-arrow-15"; import * as arrow16 from "apache-arrow-16"; import * as arrow17 from "apache-arrow-17"; @@ -40,6 +42,41 @@ function sampleRecords(): Array> { ]; } +it("serializes an Arrow Table created in another JavaScript realm", async () => { + const context = vm.createContext({ + TextDecoder, + TextEncoder, + console, + setTimeout, + clearTimeout, + }); + vm.runInContext( + fs.readFileSync( + require.resolve("apache-arrow-15/Arrow.es2015.min"), + "utf8", + ), + context, + ); + const foreignTable: unknown = vm.runInContext( + "Arrow.tableFromArrays({ id: new Int32Array([1, 2, 3]), text: ['foo', 'bar', 'baz'] })", + context, + ); + + const foreignMetadata = ( + foreignTable as { schema: { metadata: Map } } + ).schema.metadata; + expect(foreignMetadata).not.toBeInstanceOf(Map); + + const buf = await fromDataToBuffer( + foreignTable as Parameters[0], + ); + const actual = currentTableFromIPC(buf); + + expect(actual.numRows).toBe(3); + expect(actual.getChild("id")?.toJSON()).toEqual([1, 2, 3]); + expect(actual.getChild("text")?.toJSON()).toEqual(["foo", "bar", "baz"]); +}); + it("preserves field metadata from a provided schema", async function () { const jsonMetadata = new Map([["ARROW:extension:name", "lance.json"]]); const schema = new CurrentSchema([ diff --git a/nodejs/lancedb/arrow.ts b/nodejs/lancedb/arrow.ts index b52ab50ef..1b6b98cc9 100644 --- a/nodejs/lancedb/arrow.ts +++ b/nodejs/lancedb/arrow.ts @@ -72,8 +72,7 @@ export type FieldLike = }; export type DataLike = - // biome-ignore lint/suspicious/noExplicitAny: - | import("apache-arrow").Data> + | import("apache-arrow").Data | { // biome-ignore lint/suspicious/noExplicitAny: type: any; @@ -82,6 +81,7 @@ export type DataLike = stride: number; nullable: boolean; children: DataLike[]; + dictionary?: { data: readonly DataLike[] }; get nullCount(): number; // biome-ignore lint/suspicious/noExplicitAny: values: Buffers[BufferType.DATA]; diff --git a/nodejs/lancedb/sanitize.ts b/nodejs/lancedb/sanitize.ts index 8fb2f1a0a..454c82247 100644 --- a/nodejs/lancedb/sanitize.ts +++ b/nodejs/lancedb/sanitize.ts @@ -94,17 +94,24 @@ export function sanitizeMetadata( if (metadataLike === undefined || metadataLike === null) { return undefined; } - if (!(metadataLike instanceof Map)) { + + let entries: IterableIterator<[unknown, unknown]>; + try { + entries = Map.prototype.entries.call(metadataLike); + } catch { throw Error("Expected metadata, if present, to be a Map"); } - for (const item of metadataLike) { - if (typeof item[0] !== "string" || typeof item[1] !== "string") { + + const metadata = new Map(); + for (const [key, value] of entries) { + if (typeof key !== "string" || typeof value !== "string") { throw Error( "Expected metadata, if present, to be a Map but it had non-string keys or values", ); } + metadata.set(key, value); } - return metadataLike as Map; + return metadata; } export function sanitizeInt(typeLike: object) { From b85776c22a7605047bf14c2a7f7d036648d501dd Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 27 Aug 2026 15:52:41 -0700 Subject: [PATCH 3/9] fix(listing)!: page table listings from the store's own cursor (#3979) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit BREAKING CHANGE: list_tables now provides tables in arbitrary order and the page token is now completely opaque. `table_names` retains the old behavior of lexical ordering and `start-after` semantics. Listing the tables in a directory database cost what the database held rather than what the page held. `ListingDatabase::list_tables` enumerated every child directory of the base path, sorted the names, then discarded all but the requested page — on every request, for every page. On object storage that is one full listing per page. This PR pages the store instead. `list_tables` asks for one page at a time through `ObjectStore::read_dir_page`, carrying the store's own continuation token, so a page is one request. Non-table children can leave a page short of its limit, so the walk continues until the page is full or the store runs out. --------- Co-authored-by: Claude Opus 5 (1M context) --- rust/lancedb/src/database/listing.rs | 283 +++++++++++++++++++++++---- 1 file changed, 248 insertions(+), 35 deletions(-) diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index 064b5d28f..c22b73dd7 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -13,7 +13,7 @@ use lance::dataset::{ReadParams, WriteMode, builder::DatasetBuilder}; use lance::io::{ObjectStore, ObjectStoreParams, WrappingObjectStore}; use lance_datafusion::utils::StreamingWriteSource; use lance_file::version::LanceFileVersion; -use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider}; +use lance_io::object_store::{ReadDirOptions, StorageOptionsAccessor, StorageOptionsProvider}; use lance_table::io::commit::commit_handler_from_url; use object_store::local::LocalFileSystem; use snafu::ResultExt; @@ -281,6 +281,22 @@ impl std::fmt::Display for ListingDatabase { } const LANCE_EXTENSION: &str = "lance"; + +/// The table a listed child of the database names, or `None` if the child is not a table. +/// +/// A table is the directory `.lance`; a loose file or any other directory under the +/// database prefix belongs to something else. `dir_suffix` is `.lance`, built once by the +/// caller rather than per child. +/// The table a listed child directory holds, or `None` if it is not a table at all. +/// +/// Only directories are considered, so a loose object named like a table is not one. +fn table_name(location: &object_store::path::Path, dir_suffix: &str) -> Option { + location + .filename()? + .strip_suffix(dir_suffix) + .map(String::from) + .filter(|name| !name.is_empty()) +} const ENGINE: &str = "engine"; const MIRRORED_STORE: &str = "mirroredStore"; @@ -944,51 +960,72 @@ impl Database for ListingDatabase { Ok(f) } + /// List the tables in the database, a page at a time. + /// + /// The page_token is opaque, unlike the `start_after` parameter of [`Self::table_names()`]. + /// + /// When there are no more results, the returned page_token will be None. + /// + /// `limit` is the maximum number of tables to return in the response. But it is possible + /// for the response to contain fewer than `limit` tables, even when there are more tables + /// to return. Clients should check the returned page_token to determine if there are + /// more results, rather than relying on the number of tables returned. + /// + /// The order that results are returned in not guaranteed to be stable across calls, + /// so clients should not rely on it. async fn list_tables(&self, request: ListTablesRequest) -> Result { if request.id.as_ref().map(|v| !v.is_empty()).unwrap_or(false) { return self.namespace_database().list_tables(request).await; } - let mut f = self - .object_store - .read_dir(self.base_path.clone()) - .await? - .iter() - .map(Path::new) - .filter(|path| { - let is_lance = path - .extension() - .and_then(|e| e.to_str()) - .map(|e| e == LANCE_EXTENSION); - is_lance.unwrap_or(false) - }) - .filter_map(|p| p.file_stem().and_then(|s| s.to_str().map(String::from))) - .collect::>(); - f.sort(); + let limit = request.limit.map(|limit| limit.max(0) as usize); + let dir_suffix = format!(".{LANCE_EXTENSION}"); + let mut tables = Vec::new(); + let mut page_token = request.page_token.filter(|token| !token.is_empty()); - // Handle pagination with page_token - if let Some(ref page_token) = request.page_token { - let index = f - .iter() - .position(|name| name.as_str() > page_token.as_str()) - .unwrap_or(f.len()); - f.drain(0..index); + // A page of nothing: the store rejects a limit of zero, and no table was handed over + // for a token to resume after. + if limit == Some(0) { + return Ok(ListTablesResponse { + context: None, + tables, + page_token: None, + }); } - // Determine if there's a next page. The token is the last name of this page, - // not the first of the next one: the next page resumes strictly after the - // token, so naming the next page's first entry would skip it. - let next_page_token = match request.limit { - Some(limit) if f.len() > limit as usize => { - f.truncate(limit as usize); - f.last().cloned() + loop { + // Ask only for what the page still has room for, so a database holding more + // than one page costs one request per page rather than one per table. + let listing = self + .object_store + .read_dir_page( + self.base_path.clone(), + ReadDirOptions { + page_token: page_token.take(), + limit: limit.map(|limit| limit - tables.len()), + }, + ) + .await?; + page_token = listing.page_token; + // Only child directories can be tables, and the store already separates them + // out, so the objects in the page are not looked at. + tables.extend( + listing + .result + .common_prefixes + .iter() + .filter_map(|location| table_name(location, &dir_suffix)), + ); + // Children that are not tables leave the page short of the limit, so keep + // going until the page is full or the database runs out. + if page_token.is_none() || limit.is_none_or(|limit| tables.len() >= limit) { + break; } - _ => None, - }; + } Ok(ListTablesResponse { context: None, - tables: f, - page_token: next_page_token, + tables, + page_token, }) } @@ -1484,6 +1521,182 @@ mod tests { use tokio::sync::Barrier; use tokio::time::timeout; + async fn create_tables(db: &ListingDatabase, names: &[&str]) { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + for name in names { + db.create_table(CreateTableRequest { + name: name.to_string(), + namespace_path: vec![], + data: Box::new(RecordBatch::new_empty(schema.clone())) as Box, + mode: CreateTableMode::Create, + write_options: Default::default(), + location: None, + namespace_client: None, + }) + .await + .unwrap(); + } + } + + /// Every table in the database, taken `limit` at a time, which is how a caller walks a + /// listing: the token ends the walk, never a short page. + async fn walk(db: &ListingDatabase, limit: Option) -> Vec { + let mut seen = Vec::new(); + let mut page_token = None; + loop { + let page = db + .list_tables(ListTablesRequest { + limit, + page_token, + ..Default::default() + }) + .await + .unwrap(); + seen.extend(page.tables); + page_token = page.page_token; + if page_token.is_none() { + return seen; + } + assert!( + seen.len() < 100, + "the walk is serving tables more than once" + ); + } + } + + /// Paging with the returned token has to visit every table exactly once, whatever the + /// page size, with nothing lost or repeated at a boundary. + #[rstest::rstest] + #[tokio::test] + async fn test_list_tables_pages_over_every_table_once(#[values(1, 2, 3, 5, 10)] limit: i32) { + let (_tempdir, db) = setup_database().await; + create_tables(&db, &["a", "b", "c", "d", "e"]).await; + + assert_eq!(walk(&db, Some(limit)).await, vec!["a", "b", "c", "d", "e"]); + } + + /// The token is opaque: it is whatever resumes the store the database sits on, not a + /// table name. Callers hand it back and nothing else. + /// + /// Nothing validates a token, so one invented by a caller is read as a position rather + /// than refused — which is why the token has to come back from a previous page. + #[tokio::test] + async fn test_the_page_token_is_not_a_table_name() { + let (_tempdir, db) = setup_database().await; + create_tables(&db, &["a", "b", "c"]).await; + + let page = db + .list_tables(ListTablesRequest { + limit: Some(1), + ..Default::default() + }) + .await + .unwrap(); + + assert_eq!(page.tables, vec!["a"]); + let token = page.page_token.expect("two tables are still to come"); + assert_ne!(token, "a"); + + // Handing it back is the only thing a caller does with it, and it resumes. + let rest = db + .list_tables(ListTablesRequest { + page_token: Some(token), + ..Default::default() + }) + .await + .unwrap(); + assert_eq!(rest.tables, vec!["b", "c"]); + } + + /// A limit the listing does not fill leaves no token behind, so a caller paging by token + /// stops without asking for an empty page. + #[tokio::test] + async fn test_a_listing_that_runs_out_has_no_token() { + let (_tempdir, db) = setup_database().await; + create_tables(&db, &["a", "b"]).await; + + let page = db + .list_tables(ListTablesRequest { + limit: Some(10), + ..Default::default() + }) + .await + .unwrap(); + + assert_eq!(page.tables, vec!["a", "b"]); + assert_eq!(page.page_token, None); + } + + /// An empty page token means "from the start", which is how a client looping on a token + /// spells its first request. + #[tokio::test] + async fn test_an_empty_page_token_lists_from_the_start() { + let (_tempdir, db) = setup_database().await; + create_tables(&db, &["a", "b"]).await; + + let page = db + .list_tables(ListTablesRequest { + page_token: Some(String::new()), + ..Default::default() + }) + .await + .unwrap(); + + assert_eq!(page.tables, vec!["a", "b"]); + } + + /// Listing follows the order the object store lists directories in, so a name that + /// extends another comes first: the `-` of `users-archive.lance` sorts below the `.` of + /// `users.lance`. Pagination pushes its cursor into the list request, so it cannot report + /// an order other than the one it resumes in. + #[tokio::test] + async fn test_listing_order_follows_the_store_not_the_table_name() { + let (_tempdir, db) = setup_database().await; + create_tables(&db, &["users", "users-archive", "users.old"]).await; + + assert_eq!( + walk(&db, None).await, + vec!["users-archive", "users", "users.old"] + ); + // And paging reports the same order, so a walk sees each table once. + assert_eq!( + walk(&db, Some(1)).await, + vec!["users-archive", "users", "users.old"] + ); + } + + /// Only directories named `.lance` are tables; loose files and other directories + /// under the database prefix are not. A page spent on them is filled from the next one, + /// so a page holding only non-tables does not read as an empty database. + #[tokio::test] + async fn test_listing_ignores_non_table_children() { + let (tempdir, db) = setup_database().await; + create_tables(&db, &["real"]).await; + std::fs::write(tempdir.path().join("aaa-loose.lance"), b"not a table").unwrap(); + create_dir_all(tempdir.path().join("aaa-scratch")).unwrap(); + + let page = db + .list_tables(ListTablesRequest { + limit: Some(1), + ..Default::default() + }) + .await + .unwrap(); + + assert_eq!(page.tables, vec!["real"]); + } + + #[tokio::test] + async fn listing_ignores_empty_table_name() { + let (tempdir, db) = setup_database().await; + create_dir_all(tempdir.path().join(".lance")).unwrap(); + let page = db.list_tables(ListTablesRequest::default()).await.unwrap(); + assert!( + page.tables.is_empty(), + "invalid empty table name was listed" + ); + } + async fn setup_database() -> (tempfile::TempDir, ListingDatabase) { let tempdir = tempdir().unwrap(); let uri = tempdir.path().to_str().unwrap(); From 83cff3ab93c1998e9e546efbc6a62adfdbc00b7b Mon Sep 17 00:00:00 2001 From: Drew Date: Thu, 27 Aug 2026 17:07:29 -0700 Subject: [PATCH 4/9] fix(python): use one blobv2 type and coerce blob writes by metadata (#4065) --- docs/src/python/python.md | 8 +- python/python/lancedb/__init__.py | 17 +- python/python/lancedb/_blob.py | 8 +- python/python/lancedb/schema.py | 135 +++-- python/python/lancedb/table.py | 292 ++++++++--- python/python/tests/test_blob.py | 480 ++++++++++++++++++ python/python/tests/test_util.py | 160 ++++++ .../src/table/datafusion/blob_coerce.rs | 22 +- 8 files changed, 1022 insertions(+), 100 deletions(-) diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 3cb996a15..5f359f4b7 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -223,9 +223,13 @@ tokens = list( Blob columns store large binary values out of line so they can be read lazily instead of being materialized with the rest of the row. -::: lancedb.blob +`lancedb.BlobType` is `lance.blob.BlobType` when pylance is installed. Without +pylance, LanceDB uses a matching `lance.blob.v2` extension type so blob columns +still work. Queries return descriptors. Call +[`fetch_blob_files`][lancedb.table.Table.fetch_blob_files] for lazy reads or +[`fetch_blobs`][lancedb.table.Table.fetch_blobs] for eager bytes. -::: lancedb.BlobType +::: lancedb.blob ::: lancedb._blob.BlobFile options: diff --git a/python/python/lancedb/__init__.py b/python/python/lancedb/__init__.py index 8cb85a3ed..21ffc8860 100644 --- a/python/python/lancedb/__init__.py +++ b/python/python/lancedb/__init__.py @@ -6,7 +6,7 @@ import importlib.metadata import os from concurrent.futures import ThreadPoolExecutor from datetime import timedelta -from typing import Dict, Optional, Union, Any, List, Iterable +from typing import Dict, Optional, Union, Any, List, Iterable, TYPE_CHECKING __version__ = importlib.metadata.version("lancedb") @@ -20,7 +20,7 @@ from .db import AsyncConnection, DBConnection, LanceDBConnection from .remote import ClientConfig from .remote.db import RemoteDBConnection from .expr import Expr, col, lit, func -from .schema import blob, vector, BlobType +from .schema import blob, vector from .job import AsyncJob, Job from .functions import ( FunctionArtifactRequest as FunctionArtifactRequest, @@ -49,6 +49,19 @@ from .namespace import ( ) +if TYPE_CHECKING: + from lance.blob import BlobType as BlobType + + +def __getattr__(name: str): + if name == "BlobType": + from .schema import BlobType + + globals()["BlobType"] = BlobType + return BlobType + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + def _check_s3_bucket_with_dots( uri: str, storage_options: Optional[Dict[str, str]] ) -> None: diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index 5b4c0c343..dc4ed37df 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING, Optional, Union import pyarrow as pa from .expr import Expr -from .schema import blob_v2_column_paths +from .schema import row_addressable_blob_v2_paths from .types import BlobMode, QueryProjection, QueryProjectionSpec if TYPE_CHECKING: @@ -119,7 +119,7 @@ def blob_v2_projection_sources( schema: pa.Schema, projection: QueryProjection, ) -> dict[str, str]: - blob_columns = blob_v2_column_paths(schema) + blob_columns = row_addressable_blob_v2_paths(schema) if not blob_columns: return {} columns = set(blob_columns) @@ -140,7 +140,9 @@ def v2_projection_needs_row_id( ) -> bool: if with_row_id: return False - return projection_includes_blob_column(projection, blob_v2_column_paths(schema)) + return projection_includes_blob_column( + projection, row_addressable_blob_v2_paths(schema) + ) def blob_auto_row_id_for_scan( diff --git a/python/python/lancedb/schema.py b/python/python/lancedb/schema.py index 33adbbae3..4dce09f0f 100644 --- a/python/python/lancedb/schema.py +++ b/python/python/lancedb/schema.py @@ -4,30 +4,34 @@ """Schema helpers for Lance blob columns.""" +import importlib +from typing import TYPE_CHECKING + import pyarrow as pa +import pyarrow.ipc + +if TYPE_CHECKING: + from lance.blob import BlobType as BlobType _BLOB_EXTENSION_NAME = "lance.blob.v2" _BLOB_V1_KEY = "lance-encoding:blob" _ARROW_EXT_NAME_KEY = "ARROW:extension:name" +_BLOB_V2_STORAGE_TYPE = pa.struct( + [ + pa.field("data", pa.large_binary(), nullable=True), + pa.field("uri", pa.utf8(), nullable=True), + pa.field("position", pa.uint64(), nullable=True), + pa.field("size", pa.uint64(), nullable=True), + ] +) +_resolved_blob_type = None -class BlobType(pa.ExtensionType): - """PyArrow extension type for a Lance blob v2 column. - - Queries return descriptors; call :meth:`~lancedb.table.Table.fetch_blob_files` - for lazy reads or :meth:`~lancedb.table.Table.fetch_blobs` for eager bytes. - """ +class _FallbackBlobType(pa.ExtensionType): + """lance.blob.v2 extension type used when pylance is not installed.""" def __init__(self) -> None: - storage_type = pa.struct( - [ - pa.field("data", pa.large_binary(), nullable=True), - pa.field("uri", pa.utf8(), nullable=True), - pa.field("position", pa.uint64(), nullable=True), - pa.field("size", pa.uint64(), nullable=True), - ] - ) - super().__init__(storage_type, _BLOB_EXTENSION_NAME) + pa.ExtensionType.__init__(self, _BLOB_V2_STORAGE_TYPE, _BLOB_EXTENSION_NAME) def __arrow_ext_serialize__(self) -> bytes: return b"" @@ -35,23 +39,16 @@ class BlobType(pa.ExtensionType): @classmethod def __arrow_ext_deserialize__( cls, storage_type: pa.DataType, serialized: bytes - ) -> "BlobType": + ) -> "_FallbackBlobType": return cls() def __reduce__(self): - # Ensure pickle round-trips on older pyarrow (apache/arrow#35599). return type(self).__arrow_ext_deserialize__, ( self.storage_type, self.__arrow_ext_serialize__(), ) -try: - pa.register_extension_type(BlobType()) # type: ignore[arg-type] -except pa.ArrowKeyError: - pass - - def _metadata_value(metadata: dict, key: str): return metadata.get(key.encode()) or metadata.get(key) @@ -92,43 +89,105 @@ def is_blob_like_field(field: pa.Field) -> bool: return is_blob_v2_field(field) or _metadata_marks_legacy_blob(field.metadata or {}) -def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[str]: - paths: list[str] = [] +def _collect_blob_paths(schema: pa.Schema, is_blob) -> list[tuple[str, bool]]: + """Walk the schema and return (path, has_list_ancestor) for each blob field.""" + paths: list[tuple[str, bool]] = [] - def walk(fields, prefix: str) -> None: + def walk(fields, prefix: str, has_list_ancestor: bool) -> None: for field in fields: path = f"{prefix}.{field.name}" if prefix else field.name if is_blob(field): - paths.append(path) + paths.append((path, has_list_ancestor)) elif pa.types.is_struct(field.type): - walk(field.type, path) + walk(field.type, path, has_list_ancestor) elif ( pa.types.is_list(field.type) or pa.types.is_large_list(field.type) or pa.types.is_fixed_size_list(field.type) ): - walk([field.type.value_field], path) + walk([field.type.value_field], path, True) - walk(schema, "") + walk(schema, "", False) return paths def blob_column_paths(schema: pa.Schema) -> list[str]: """Dotted paths of blob-like columns (v2 extension or legacy metadata).""" - return _collect_blob_paths(schema, is_blob_like_field) + return [path for path, _ in _collect_blob_paths(schema, is_blob_like_field)] def blob_v2_column_paths(schema: pa.Schema) -> list[str]: - return _collect_blob_paths(schema, is_blob_v2_field) + return [path for path, _ in _collect_blob_paths(schema, is_blob_v2_field)] + + +def row_addressable_blob_v2_paths(schema: pa.Schema) -> list[str]: + """Blob v2 paths with one blob addressable by table row id. + + ``fetch_blobs`` and the descriptor row-id ride-along address one blob per + row, so a blob inside a list container has no row-id slot and no fetch + path. Those columns still store and query as raw descriptors. + """ + return [ + path + for path, has_list_ancestor in _collect_blob_paths(schema, is_blob_v2_field) + if not has_list_ancestor + ] def schema_has_blob_field(schema: pa.Schema) -> bool: return bool(blob_column_paths(schema)) +def _deserialize_registered_type(extension_type: pa.ExtensionType) -> pa.DataType: + """Return the type Arrow reconstructs for this extension name.""" + schema = pa.schema([pa.field("value", extension_type)]) + restored = pa.ipc.read_schema(schema.serialize()) + return restored.field("value").type + + +def _resolve_blob_type(): + """Return the BlobType class this process should use. + + pylance's class when it owns the lance.blob.v2 registry entry, + otherwise LanceDB's fallback. A different registered class is an error. + """ + global _resolved_blob_type + if _resolved_blob_type is not None: + return _resolved_blob_type + try: + blob_module = importlib.import_module("lance.blob") + except ModuleNotFoundError as err: + if err.name not in ("lance", "lance.blob"): + raise + else: + blob_type = getattr(blob_module, "BlobType", None) + if blob_type is not None: + registered_type = _deserialize_registered_type(blob_type()) + if type(registered_type) is not blob_type: + registered_cls = type(registered_type) + raise ValueError( + "lance.blob.v2 is already registered by " + f"{registered_cls.__module__}.{registered_cls.__qualname__}" + ) + _resolved_blob_type = blob_type + return blob_type + try: + pa.register_extension_type(_FallbackBlobType()) # type: ignore[arg-type] + except pa.ArrowKeyError as err: + raise ValueError( + "lance.blob.v2 is already registered by another extension class" + ) from err + _resolved_blob_type = _FallbackBlobType + return _resolved_blob_type + + def blob(name: str, nullable: bool = True) -> pa.Field: - """Create a Lance blob v2 column field.""" - return pa.field(name, BlobType(), nullable=nullable) + """Create a Lance blob v2 column field. + + When pylance is installed this is ``lance.blob.BlobType``. + """ + blob_type = _resolve_blob_type() + return pa.field(name, blob_type(), nullable=nullable) def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataType: @@ -155,3 +214,11 @@ def vector(dimension: int, value_type: pa.DataType = pa.float32()) -> pa.DataTyp ... ]) """ return pa.list_(value_type, dimension) + + +def __getattr__(name: str): + if name == "BlobType": + blob_type = _resolve_blob_type() + globals()["BlobType"] = blob_type + return blob_type + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 765b7fa14..535ed7c0d 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -104,7 +104,12 @@ from .util import ( value_to_sql, ) from .index import lang_mapping -from .schema import blob_v2_column_paths, schema_has_blob_field +from .schema import ( + blob_v2_column_paths, + is_blob_v2_field, + row_addressable_blob_v2_paths, + schema_has_blob_field, +) def _should_push_down_query_table( @@ -426,6 +431,7 @@ def _cast_to_target_schema( def gen(): for batch in reader: + batch = _coerce_blob_write_columns(batch, reordered_schema) # Table but not RecordBatch has cast. cast_batches = ( pa.Table.from_batches([batch]).cast(reordered_schema).to_batches() @@ -438,6 +444,166 @@ def _cast_to_target_schema( return pa.RecordBatchReader.from_batches(reordered_schema, gen()) +def _coerce_blob_write_columns( + batch: pa.RecordBatch, target_schema: pa.Schema +) -> pa.RecordBatch: + """Materialize blob storage structs before the stream leaves Python. + + merge_insert requires its source reader to already match the table's + physical schema. Unlike add and insert, it does not pass through + LanceDB's Rust blob coercion, so preserving binary input here would + reach Lance as binary and fail the schema check. + """ + columns = [] + fields = [] + changed = False + for field, column in zip(batch.schema, batch.columns): + target_field = target_schema.field(field.name) + coerced = _coerce_blob_value(column, target_field) + if coerced is not column: + column = coerced + field = pa.field( + field.name, + coerced.type, + field.nullable, + target_field.metadata, + ) + changed = True + columns.append(column) + fields.append(field) + if not changed: + return batch + return pa.RecordBatch.from_arrays( + columns, schema=pa.schema(fields, metadata=batch.schema.metadata) + ) + + +def _coerce_blob_value(column: pa.Array, target_field: pa.Field) -> pa.Array: + if is_blob_v2_field(target_field) and _can_coerce_to_blob(column.type): + return _coerce_value_to_blob(column, target_field) + + target_type = target_field.type + if pa.types.is_struct(target_type) and pa.types.is_struct(column.type): + children = [] + fields = [] + changed = False + for source_field in column.type: + source_column = column.field(source_field.name) + nested_target = next( + (field for field in target_type if field.name == source_field.name), + None, + ) + if nested_target is None: + children.append(source_column) + fields.append(source_field) + continue + coerced = _coerce_blob_value(source_column, nested_target) + if coerced is not source_column: + changed = True + child_array, child_type = _physical_array_and_type(coerced) + children.append(child_array) + fields.append( + pa.field( + source_field.name, + child_type, + source_field.nullable, + nested_target.metadata, + ) + ) + if not changed: + return column + return pa.StructArray.from_arrays( + children, + fields=fields, + mask=column.is_null() if column.null_count else None, + ) + + if _is_list_like(target_type) and _is_list_like(column.type): + return _coerce_blob_list_values(column, target_type.value_field) + + return column + + +def _coerce_blob_list_values( + column: pa.Array, target_value_field: pa.Field +) -> pa.Array: + """Coerce blob values inside a list column, preserving offsets and nulls. + + Works on the raw child values window instead of ``pc.list_flatten`` because + flatten drops values spanned by null slots, which would misalign offsets. + """ + mask = column.is_null() if column.null_count else None + if pa.types.is_fixed_size_list(column.type): + list_size = column.type.list_size + values = column.values.slice(column.offset * list_size, len(column) * list_size) + coerced = _coerce_blob_value(values, target_value_field) + if coerced is values: + return column + physical_values, _ = _physical_array_and_type(coerced) + return pa.FixedSizeListArray.from_arrays(physical_values, list_size, mask=mask) + offsets = column.offsets + first_offset = offsets[0].as_py() + values = column.values.slice( + first_offset, + offsets[-1].as_py() - first_offset, + ) + coerced = _coerce_blob_value(values, target_value_field) + if coerced is values: + return column + physical_values, _ = _physical_array_and_type(coerced) + if first_offset: + offsets = pc.subtract(offsets, pa.scalar(first_offset, offsets.type)) + if pa.types.is_large_list(column.type): + return pa.LargeListArray.from_arrays(offsets, physical_values, mask=mask) + return pa.ListArray.from_arrays(offsets, physical_values, mask=mask) + + +def _coerce_value_to_blob(values: pa.Array, target_field: pa.Field) -> pa.Array: + if pa.types.is_null(values.type): + data = pa.nulls(len(values), type=pa.large_binary()) + elif pa.types.is_large_binary(values.type): + data = values + else: + data = values.cast(pa.large_binary()) + length = len(values) + storage_type = target_field.type + if isinstance(storage_type, pa.ExtensionType): + storage_type = storage_type.storage_type + storage_fields = list(storage_type) + children = [] + for storage_field in storage_fields: + if storage_field.name == "data": + children.append(data) + else: + children.append(pa.nulls(length, type=storage_field.type)) + storage = pa.StructArray.from_arrays( + children, + fields=storage_fields, + mask=values.is_null() if values.null_count else None, + ) + if isinstance(target_field.type, pa.ExtensionType): + return pa.ExtensionArray.from_storage(target_field.type, storage) + return storage + + +def _physical_array_and_type(array: pa.Array) -> tuple[pa.Array, pa.DataType]: + if isinstance(array.type, pa.ExtensionType): + return array.storage, array.type.storage_type + return array, array.type + + +def _can_coerce_to_blob(data_type: pa.DataType) -> bool: + return _is_binary_like(data_type) or pa.types.is_null(data_type) + + +def _is_binary_like(data_type: pa.DataType) -> bool: + return ( + pa.types.is_binary(data_type) + or pa.types.is_large_binary(data_type) + or pa.types.is_binary_view(data_type) + ) + + def _field_extension_name(field: pa.Field) -> Optional[str]: extension_name = getattr(field.type, "extension_name", None) if extension_name is not None: @@ -464,63 +630,71 @@ def _align_field_types( target_field = next((f for f in target_fields if f.name == field.name), None) if target_field is None: raise ValueError(f"Field '{field.name}' not found in target schema") - # Preserve arrow.json input until it reaches Lance. LanceDB exposes stored - # JSON columns as lance.json (JSONB-backed LargeBinary), but casting the - # input to that storage type here merely relabels the raw JSON bytes as - # JSONB. Lance must see arrow.json so it can perform the JSONB encoding. - if ( - _field_extension_name(field) == "arrow.json" - and _field_extension_name(target_field) == "lance.json" - ): - new_fields.append(field) - continue - if pa.types.is_struct(target_field.type): - if pa.types.is_struct(field.type): - new_type = pa.struct( - _align_field_types( - field.type.fields, - target_field.type.fields, - ) + new_fields.append(_align_field(field, target_field)) + return new_fields + + +def _align_list_value_field( + value_field: pa.Field, target_value_field: pa.Field +) -> pa.Field: + # A list has exactly one child, so the inferred child name ("item") aligns + # positionally and adopts the table's child name; pa.Table.cast renames it. + return _align_field(value_field, target_value_field).with_name( + target_value_field.name + ) + + +def _align_field(field: pa.Field, target_field: pa.Field) -> pa.Field: + # Preserve arrow.json input until it reaches Lance. LanceDB exposes stored + # JSON columns as lance.json (JSONB-backed LargeBinary), but casting the + # input to that storage type here merely relabels the raw JSON bytes as + # JSONB. Lance must see arrow.json so it can perform the JSONB encoding. + if ( + _field_extension_name(field) == "arrow.json" + and _field_extension_name(target_field) == "lance.json" + ): + return field + if pa.types.is_struct(target_field.type): + if pa.types.is_struct(field.type): + new_type = pa.struct( + _align_field_types( + field.type.fields, + target_field.type.fields, ) - else: - new_type = target_field.type - elif pa.types.is_list(target_field.type): - if _is_list_like(field.type): - new_type = pa.list_( - _align_field_types( - [field.type.value_field], - [target_field.type.value_field], - )[0] - ) - else: - new_type = target_field.type - elif pa.types.is_large_list(target_field.type): - if _is_list_like(field.type): - new_type = pa.large_list( - _align_field_types( - [field.type.value_field], - [target_field.type.value_field], - )[0] - ) - else: - new_type = target_field.type - elif pa.types.is_fixed_size_list(target_field.type): - if _is_list_like(field.type): - new_type = pa.list_( - _align_field_types( - [field.type.value_field], - [target_field.type.value_field], - )[0], - target_field.type.list_size, - ) - else: - new_type = target_field.type + ) else: new_type = target_field.type - new_fields.append( - pa.field(field.name, new_type, field.nullable, target_field.metadata) - ) - return new_fields + elif pa.types.is_list(target_field.type): + if _is_list_like(field.type): + new_type = pa.list_( + _align_list_value_field( + field.type.value_field, target_field.type.value_field + ) + ) + else: + new_type = target_field.type + elif pa.types.is_large_list(target_field.type): + if _is_list_like(field.type): + new_type = pa.large_list( + _align_list_value_field( + field.type.value_field, target_field.type.value_field + ) + ) + else: + new_type = target_field.type + elif pa.types.is_fixed_size_list(target_field.type): + if _is_list_like(field.type): + new_type = pa.list_( + _align_list_value_field( + field.type.value_field, target_field.type.value_field + ), + target_field.type.list_size, + ) + else: + new_type = target_field.type + else: + new_type = target_field.type + return pa.field(field.name, new_type, field.nullable, target_field.metadata) def _infer_subschema( @@ -589,7 +763,7 @@ def sanitize_create_table( schema = data.schema else: if schema is not None: - data = pa.Table.from_pylist([], schema) + data = pa.Table.from_batches([], schema=schema) if schema is None: if data is None: raise ValueError("Either data or schema must be provided") @@ -2698,7 +2872,7 @@ class LanceTable(Table): arrow_tbl = self.to_arrow() if blob_mode == "descriptions": arrow_tbl = strip_auto_row_ids( - arrow_tbl, blob_v2_column_paths(self.schema) + arrow_tbl, row_addressable_blob_v2_paths(self.schema) ) return arrow_tbl.to_pandas(**kwargs) @@ -5102,7 +5276,9 @@ class AsyncTable: if blob_mode == "descriptions" or not schema_has_blob_field(schema): arrow_tbl = await self.to_arrow() if blob_mode == "descriptions": - arrow_tbl = strip_auto_row_ids(arrow_tbl, blob_v2_column_paths(schema)) + arrow_tbl = strip_auto_row_ids( + arrow_tbl, row_addressable_blob_v2_paths(schema) + ) return arrow_tbl.to_pandas(**kwargs) if blob_mode == "lazy" and get_uri_scheme(await self.uri()) == "memory": diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 351769ff6..c9694277c 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -2,10 +2,15 @@ # SPDX-FileCopyrightText: Copyright The LanceDB Authors import io +import subprocess +import sys +import textwrap +import lance import pyarrow as pa import pyarrow.compute as pc import pytest +from lance.blob import BlobType as LanceBlobType import lancedb from lancedb._blob import ( @@ -18,6 +23,20 @@ from lancedb.index import FTS from lancedb.schema import blob_column_paths, blob_v2_column_paths +_HIDE_LANCE_BLOB = """\ +import importlib.abc +import sys + +class _MissingLanceBlob(importlib.abc.MetaPathFinder): + def find_spec(self, fullname, path, target=None): + if fullname == "lance.blob" or fullname.startswith("lance.blob."): + raise ModuleNotFoundError(fullname, name="lance.blob") + +sys.modules.pop("lance.blob", None) +sys.meta_path.insert(0, _MissingLanceBlob()) +""" + + def _blob_table(name, rows): db = lancedb.connect("memory:///") schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) @@ -51,6 +70,181 @@ def test_blob_factory_declares_v2_field(): field = lancedb.blob("image") assert isinstance(field.type, pa.ExtensionType) assert field.type.extension_name == "lance.blob.v2" + assert lancedb.BlobType is LanceBlobType + assert type(field.type) is LanceBlobType + + +def test_blob_type_works_without_pylance(): + script = _HIDE_LANCE_BLOB + textwrap.dedent( + """\ + import lancedb + import pyarrow as pa + + field = lancedb.blob("image") + if not isinstance(field.type, pa.ExtensionType): + raise SystemExit("expected an extension type") + if field.type.extension_name != "lance.blob.v2": + raise SystemExit(field.type.extension_name) + if lancedb.BlobType is not type(field.type): + raise SystemExit("BlobType is not the field type class") + if lancedb.BlobType.__module__ != "lancedb.schema": + raise SystemExit(lancedb.BlobType.__module__) + + db = lancedb.connect("memory:///") + table = db.create_table( + "images", + schema=pa.schema([pa.field("id", pa.int64()), field]), + ) + table.add([{"id": 1, "image": b"hello"}]) + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute([{"id": 1, "image": b"updated"}, {"id": 2, "image": b"inserted"}]) + ) + if result.num_updated_rows != 1 or result.num_inserted_rows != 1: + raise SystemExit( + f"merge_insert rows updated={result.num_updated_rows} " + f"inserted={result.num_inserted_rows}" + ) + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_blob_resolves_pylance_type_without_eager_import(): + script = textwrap.dedent( + """\ + import sys + import lancedb + + if "lance.blob" in sys.modules: + raise SystemExit("import lancedb imported lance.blob") + field = lancedb.blob("image") + from lance.blob import BlobType + + if type(field.type) is not BlobType: + raise SystemExit(f"{type(field.type)} is not {BlobType}") + import lance + + image = lance.blob_array([b"x"]) + if type(image.type) is not BlobType: + raise SystemExit("blob_array used a different class") + if type(image.type) is not type(field.type): + raise SystemExit("field and array classes differ") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_blob_fallback_fails_if_name_already_registered(): + script = _HIDE_LANCE_BLOB + textwrap.dedent( + """\ + import pyarrow as pa + + class OtherBlobType(pa.ExtensionType): + def __init__(self): + super().__init__( + pa.struct([pa.field("data", pa.large_binary())]), + "lance.blob.v2", + ) + + def __arrow_ext_serialize__(self): + return b"" + + @classmethod + def __arrow_ext_deserialize__(cls, storage_type, serialized): + return cls() + + pa.register_extension_type(OtherBlobType()) + import lancedb + + try: + lancedb.blob("image") + except ValueError as err: + if "already registered" not in str(err): + raise SystemExit(err) + else: + raise SystemExit("expected ValueError") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_blob_type_rejects_competing_registration_with_pylance(): + script = textwrap.dedent( + """\ + import pyarrow as pa + import pyarrow.ipc + + class OtherBlobType(pa.ExtensionType): + def __init__(self): + super().__init__( + pa.struct( + [ + pa.field("data", pa.large_binary()), + pa.field("uri", pa.utf8()), + pa.field("position", pa.uint64()), + pa.field("size", pa.uint64()), + ] + ), + "lance.blob.v2", + ) + + def __arrow_ext_serialize__(self): + return b"" + + @classmethod + def __arrow_ext_deserialize__(cls, storage_type, serialized): + return cls() + + pa.register_extension_type(OtherBlobType()) + + from lance.blob import BlobType + + if BlobType is OtherBlobType: + raise SystemExit("pylance BlobType was replaced") + schema = pa.schema([pa.field("value", BlobType())]) + restored = pa.ipc.read_schema(schema.serialize()) + if type(restored.field("value").type) is not OtherBlobType: + raise SystemExit(type(restored.field("value").type)) + + import lancedb + + try: + lancedb.blob("image") + except ValueError as err: + if "__main__.OtherBlobType" not in str(err): + raise SystemExit(err) + else: + raise SystemExit("expected ValueError") + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr def test_blob_v2_column_paths_include_list_children(): @@ -203,6 +397,292 @@ def test_fetch_blobs_round_trip(): assert [blobs[0].as_py(), blobs[1].as_py()] == [b"alpha", b"beta"] +def test_merge_insert_writes_python_bytes(): + table = _blob_table("merge_bytes", [{"id": 1, "image": b"before"}]) + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute([{"id": 1, "image": b"updated"}, {"id": 2, "image": b"inserted"}]) + ) + assert result.num_updated_rows == 1 + assert result.num_inserted_rows == 1 + by_id = _row_ids_by_id(table) + blobs = table.fetch_blobs("image", [by_id[1], by_id[2]]) + assert blobs.to_pylist() == [b"updated", b"inserted"] + + +def test_merge_insert_bytes_after_reopen_without_touching_blob_type(tmp_path): + db = lancedb.connect(tmp_path) + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("images", schema=schema) + table.add([{"id": 1, "image": b"hello"}]) + + script = textwrap.dedent( + f"""\ + import lancedb + + db = lancedb.connect({str(tmp_path)!r}) + table = db.open_table("images") + image_type = table.schema.field("image").type + if type(image_type).__name__ != "StructType": + raise SystemExit(f"expected StructType, got {{type(image_type)}}") + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute( + [{{"id": 1, "image": b"updated"}}, {{"id": 2, "image": b"inserted"}}] + ) + ) + if result.num_updated_rows != 1 or result.num_inserted_rows != 1: + raise SystemExit( + f"rows updated={{result.num_updated_rows}} " + f"inserted={{result.num_inserted_rows}}" + ) + hits = table.search().with_row_id(True).limit(10).to_arrow() + by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist())) + blobs = table.fetch_blobs("image", [by_id[1], by_id[2]]) + if blobs.to_pylist() != [b"updated", b"inserted"]: + raise SystemExit(blobs.to_pylist()) + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_merge_insert_bytes_after_reopen_without_pylance(tmp_path): + db = lancedb.connect(tmp_path) + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("images", schema=schema) + table.add([{"id": 1, "image": b"hello"}]) + + script = _HIDE_LANCE_BLOB + textwrap.dedent( + f"""\ + import lancedb + + db = lancedb.connect({str(tmp_path)!r}) + table = db.open_table("images") + image_type = table.schema.field("image").type + if type(image_type).__name__ != "StructType": + raise SystemExit(f"expected StructType, got {{type(image_type)}}") + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute( + [{{"id": 1, "image": b"updated"}}, {{"id": 2, "image": b"inserted"}}] + ) + ) + if result.num_updated_rows != 1 or result.num_inserted_rows != 1: + raise SystemExit( + f"rows updated={{result.num_updated_rows}} " + f"inserted={{result.num_inserted_rows}}" + ) + hits = table.search().with_row_id(True).limit(10).to_arrow() + by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist())) + blobs = table.fetch_blobs("image", [by_id[1], by_id[2]]) + if blobs.to_pylist() != [b"updated", b"inserted"]: + raise SystemExit(blobs.to_pylist()) + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_merge_insert_blob_array_into_reopened_unregistered_table(tmp_path): + db = lancedb.connect(tmp_path) + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("images", schema=schema) + table.add([{"id": 1, "image": b"before"}]) + + script = textwrap.dedent( + f"""\ + import pyarrow as pa + import lancedb + + db = lancedb.connect({str(tmp_path)!r}) + table = db.open_table("images") + image_type = table.schema.field("image").type + if type(image_type).__name__ != "StructType": + raise SystemExit( + f"expected StructType before lance import, got {{type(image_type)}}" + ) + + import lance + + updates = pa.Table.from_arrays( + [ + pa.array([1, 2], type=pa.int64()), + lance.blob_array([b"updated", b"inserted"]), + ], + names=["id", "image"], + ) + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(updates) + ) + if result.num_updated_rows != 1 or result.num_inserted_rows != 1: + raise SystemExit( + f"rows updated={{result.num_updated_rows}} " + f"inserted={{result.num_inserted_rows}}" + ) + hits = table.search().with_row_id(True).limit(10).to_arrow() + by_id = dict(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist())) + blobs = table.fetch_blobs("image", [by_id[1], by_id[2]]) + if blobs.to_pylist() != [b"updated", b"inserted"]: + raise SystemExit(blobs.to_pylist()) + """ + ) + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_add_all_null_blob_column(): + db = lancedb.connect("memory:///") + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + table = db.create_table("all_null", schema=schema) + table.add([{"id": 1, "image": None}, {"id": 2, "image": None}]) + by_id = _row_ids_by_id(table) + blobs = table.fetch_blobs("image", [by_id[1], by_id[2]]) + assert blobs.to_pylist() == [None, None] + + +def test_create_table_nested_blob_schema_without_rows(): + db = lancedb.connect("memory:///") + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("info", pa.struct([lancedb.blob("blob")])), + pa.field("images", pa.list_(lancedb.blob("image"))), + ] + ) + table = db.create_table("nested_empty", schema=schema) + assert table.count_rows() == 0 + + +def test_merge_insert_nested_blob_dicts(): + db = lancedb.connect("memory:///") + info = pa.StructArray.from_arrays( + [ + pa.array(["first"], type=pa.string()), + _blob_array("blob", [b"before"]), + ], + names=["name", "blob"], + ) + data = pa.Table.from_arrays( + [pa.array([1], type=pa.int64()), info], + names=["id", "info"], + ) + table = db.create_table("nested_merge", data=data) + result = ( + table.merge_insert("id") + .when_matched_update_all() + .execute([{"id": 1, "info": {"name": "first", "blob": b"after"}}]) + ) + assert result.num_updated_rows == 1 + by_id = _row_ids_by_id(table) + blobs = table.fetch_blobs("info.blob", [by_id[1]]) + assert blobs.to_pylist() == [b"after"] + + +def _list_blob_table(name): + db = lancedb.connect("memory:///") + blob_field = lancedb.blob("image") + images = pa.ListArray.from_arrays( + pa.array([0, 1], type=pa.int32()), _blob_array("image", [b"before"]) + ) + data = pa.Table.from_arrays( + [pa.array([1], type=pa.int64()), images], + schema=pa.schema( + [pa.field("id", pa.int64()), pa.field("images", pa.list_(blob_field))] + ), + ) + return db.create_table(name, data=data) + + +def test_merge_insert_list_blob_dicts(): + table = _list_blob_table("list_merge") + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute([{"id": 1, "images": [b"one", b"two"]}, {"id": 2, "images": None}]) + ) + assert result.num_updated_rows == 1 + assert result.num_inserted_rows == 1 + hits = table.search().limit(10).to_arrow() + sizes = { + row["id"]: None if row["images"] is None else [d["size"] for d in row["images"]] + for row in hits.to_pylist() + } + assert sizes == {1: [3, 3], 2: None} + + +def test_list_blob_column_queries_as_raw_descriptors(): + table = _list_blob_table("list_query") + hits = table.search().limit(10).to_arrow() + element = hits.schema.field("images").type.value_type + assert pa.types.is_struct(element) + assert "_lance_row_id" not in element.names + with pytest.raises(ValueError, match="expected struct before segment"): + table.fetch_blobs("images.image", [0]) + + +def test_row_addressable_paths_exclude_list_children(): + from lancedb.schema import row_addressable_blob_v2_paths + + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field("info", pa.struct([lancedb.blob("blob")])), + pa.field("images", pa.list_(lancedb.blob("image"))), + ] + ) + assert blob_v2_column_paths(schema) == ["info.blob", "images.image"] + assert row_addressable_blob_v2_paths(schema) == ["info.blob"] + + +def test_merge_insert_writes_pylance_blob_array(): + table = _blob_table("merge_pylance", [{"id": 1, "image": b"before"}]) + image = lance.blob_array([b"updated", b"inserted"]) + assert type(image.type) is LanceBlobType + assert type(image.type) is type(lancedb.BlobType()) + updates = pa.Table.from_arrays( + [pa.array([1, 2], type=pa.int64()), image], names=["id", "image"] + ) + + result = ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute(updates) + ) + + assert result.num_updated_rows == 1 + assert result.num_inserted_rows == 1 + by_id = _row_ids_by_id(table) + blobs = table.fetch_blobs("image", [by_id[1], by_id[2]]) + assert blobs.to_pylist() == [b"updated", b"inserted"] + + def test_fetch_blobs_accepts_query_result(): table = _blob_table("from_result", [{"id": 1, "image": b"gamma"}]) hits = table.search().limit(10).to_arrow() diff --git a/python/python/tests/test_util.py b/python/python/tests/test_util.py index a9b66b2dd..acdef8eac 100644 --- a/python/python/tests/test_util.py +++ b/python/python/tests/test_util.py @@ -7,6 +7,7 @@ import pathlib from typing import Optional import lance +from lance.blob import BlobType as LanceBlobType from lancedb.conftest import MockTextEmbeddingFunction from lancedb.embeddings.base import EmbeddingFunctionConfig from lancedb.embeddings.registry import EmbeddingFunctionRegistry @@ -907,6 +908,165 @@ def test_cast_to_target_schema(): assert output == expected +def test_cast_to_target_schema_coerces_binary_to_blob_v2(): + data = pa.table({"image": pa.array([b"hello", None], type=pa.binary())}) + target = pa.schema([lancedb.blob("image")]) + + output = _cast_to_target_schema(data.to_reader(), target).read_all() + + image = output["image"].chunk(0) + assert type(image.type) is lancedb.BlobType + assert image.storage.to_pylist() == [ + {"data": b"hello", "uri": None, "position": None, "size": None}, + None, + ] + + +def test_cast_to_target_schema_coerces_binary_to_metadata_blob_struct(): + storage = lancedb.blob("image").type.storage_type + target = pa.schema( + [ + pa.field( + "image", + storage, + metadata={ + b"ARROW:extension:name": b"lance.blob.v2", + b"ARROW:extension:metadata": b"", + }, + ) + ] + ) + data = pa.table({"image": pa.array([b"hello", None], type=pa.binary())}) + + output = _cast_to_target_schema(data.to_reader(), target).read_all() + + image = output["image"].chunk(0) + assert not isinstance(image.type, pa.ExtensionType) + assert image.to_pylist() == [ + {"data": b"hello", "uri": None, "position": None, "size": None}, + None, + ] + + +def test_cast_to_target_schema_coerces_nested_binary_blob(): + data = pa.table( + { + "info": pa.array( + [{"blob": b"hello"}, {"blob": None}], + type=pa.struct([pa.field("blob", pa.binary())]), + ) + } + ) + target = pa.schema([pa.field("info", pa.struct([lancedb.blob("blob")]))]) + + output = _cast_to_target_schema(data.to_reader(), target).read_all() + + blob = output["info"].chunk(0).field("blob") + assert type(blob.type) is lancedb.BlobType + assert blob.storage.to_pylist() == [ + {"data": b"hello", "uri": None, "position": None, "size": None}, + None, + ] + + +def test_cast_to_target_schema_coerces_list_binary_blob_with_inferred_child_name(): + data = pa.table( + {"images": pa.array([[b"a", b"b"], None], type=pa.list_(pa.binary()))} + ) + target = pa.schema([pa.field("images", pa.list_(lancedb.blob("image")))]) + + output = _cast_to_target_schema(data.to_reader(), target).read_all() + + images = output["images"].chunk(0) + assert images.type.value_field.name == "image" + assert type(images.type.value_type) is lancedb.BlobType + assert images.to_pylist()[1] is None + assert images.values.storage.to_pylist() == [ + {"data": b"a", "uri": None, "position": None, "size": None}, + {"data": b"b", "uri": None, "position": None, "size": None}, + ] + + +def test_list_blob_coercion_preserves_null_slots_with_nonzero_extent(): + child = pa.field("image", pa.binary()) + source = pa.ListArray.from_arrays( + pa.array([0, 2, 4], type=pa.int32()), + pa.array([b"a", b"b", b"dead", b"beef"], type=pa.binary()), + mask=pa.array([False, True]), + ).cast(pa.list_(child)) + target = pa.schema([pa.field("images", pa.list_(lancedb.blob("image")))]) + + output = _cast_to_target_schema( + pa.table({"images": source}).to_reader(), target + ).read_all() + + images = output["images"].chunk(0) + assert images.to_pylist()[1] is None + assert [b["data"] for b in images.to_pylist()[0]] == [b"a", b"b"] + + +def test_fixed_size_list_blob_coercion_keeps_null_rows(): + child = pa.field("frame", pa.binary()) + source = ( + pa.FixedSizeListArray.from_arrays( + pa.array([b"a", b"b", b"c", b"d"], type=pa.binary()), 2 + ) + .take(pa.array([0, None], type=pa.int32())) + .cast(pa.list_(child, 2)) + ) + target = pa.schema([pa.field("frames", pa.list_(lancedb.blob("frame"), 2))]) + + output = _cast_to_target_schema( + pa.table({"frames": source}).to_reader(), target + ).read_all() + + frames = output["frames"].chunk(0) + assert frames.to_pylist()[1] is None + assert [b["data"] for b in frames.to_pylist()[0]] == [b"a", b"b"] + + +def test_cast_to_target_schema_accepts_pylance_blob_v2(): + target_type = lancedb.BlobType() + source = lance.blob_array([b"hello", None]) + assert type(source.type) is LanceBlobType + assert type(source.type) is type(target_type) + data = pa.table({"image": source}) + target = pa.schema([pa.field("image", target_type)]) + + output = _cast_to_target_schema(data.to_reader(), target).read_all() + + image = output["image"].chunk(0) + assert type(image.type) is LanceBlobType + assert image.type == target_type + assert image.storage.to_pylist() == [ + {"data": b"hello", "uri": None, "position": None, "size": None}, + None, + ] + + +def test_cast_to_target_schema_rejects_different_blob_v2_class(): + class OtherBlobType(pa.ExtensionType): + def __init__(self): + super().__init__(lancedb.BlobType().storage_type, "lance.blob.v2") + + def __arrow_ext_serialize__(self) -> bytes: + return b"" + + @classmethod + def __arrow_ext_deserialize__( + cls, storage_type: pa.DataType, serialized: bytes + ) -> "OtherBlobType": + return cls() + + storage = lance.blob_array([b"hello"]).storage + source = pa.ExtensionArray.from_storage(OtherBlobType(), storage) + data = pa.table({"image": source}) + target = pa.schema([lancedb.blob("image")]) + + with pytest.raises(pa.ArrowTypeError, match="different extension type"): + _cast_to_target_schema(data.to_reader(), target).read_all() + + def test_sanitize_data_stream(): # Make sure we don't collect the whole stream when running sanitize_data schema = pa.schema({"a": pa.int32()}) diff --git a/rust/lancedb/src/table/datafusion/blob_coerce.rs b/rust/lancedb/src/table/datafusion/blob_coerce.rs index cb984f7f4..0596e7a2d 100644 --- a/rust/lancedb/src/table/datafusion/blob_coerce.rs +++ b/rust/lancedb/src/table/datafusion/blob_coerce.rs @@ -36,6 +36,14 @@ pub(super) fn coerce_blob_expr( }; let input_shape = match input_field.data_type() { + DataType::Null => { + let expr: Arc = Arc::new(CastExpr::new( + input_expr, + table_field.data_type().clone(), + None, + )); + return Ok((expr, table_field.clone())); + } DataType::Binary | DataType::LargeBinary | DataType::BinaryView => BlobInputShape::Bytes, DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => BlobInputShape::String, DataType::Struct(children) => { @@ -155,7 +163,7 @@ mod tests { use crate::blob::blob; use arrow_array::{ Array, ArrayRef, BinaryArray, BinaryViewArray, Int32Array, Int64Array, LargeBinaryArray, - RecordBatch, StringArray, StringViewArray, StructArray, UInt8Array, UInt64Array, + NullArray, RecordBatch, StringArray, StringViewArray, StructArray, UInt8Array, UInt64Array, }; use arrow_schema::Schema; use datafusion::prelude::SessionContext; @@ -279,6 +287,18 @@ mod tests { assert_eq!(data.value(0), b"view"); } + #[tokio::test] + async fn null_column_coerces_to_all_null_blob_struct() { + let batch = batch_with_image( + Field::new("image", DataType::Null, true), + Arc::new(NullArray::new(2)), + ); + let coerced = coerce(batch, &blob_table_schema()).await; + let image = image_struct(&coerced); + assert!(image.is_null(0)); + assert!(image.is_null(1)); + } + #[tokio::test] async fn binary_nulls_stay_null_after_coercion() { let batch = batch_with_image( From 84f46df876b988aeb22a05eff6eb671d174b87a6 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 27 Aug 2026 17:35:06 -0700 Subject: [PATCH 5/9] ci(nodejs): fix nightly OOM on the aarch64 publish legs (#4077) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The nightly `NPM Publish` run has failed every night since at least Aug 23, always on the same two legs: `aarch64-unknown-linux-gnu` and `aarch64-unknown-linux-musl`. The other five targets pass. rustc is OOM-killed during the fat-LTO codegen of the cdylib — `signal: 9` with no diagnostic, about 27 minutes in — and on the musl leg that takes the whole runner down with `The runner has received a shutdown signal`. Both legs now pass: | leg | before | peak memory | wall time | | --- | --- | --- | --- | | `aarch64-unknown-linux-gnu` | OOM-killed at ~27 min | 31391 → 22851 MiB | 38m43s → 22m04s | | `aarch64-unknown-linux-musl` | runner killed at ~28 min | >32 GiB → 16516 MiB | ~40 min → 20m50s | **ThinLTO** is most of that. Fat LTO is single-threaded, and its peak is consumed inside rustc's LLVM before any linker process is spawned — which is why it is the whole fix on musl, and why lld alone left the gnu leg still peaking at 31391 MiB against the runner's 32 GiB. Both legs now use the `lto: thin` / `codegen_units: 16` settings that darwin and both Windows legs already use, at a cost of a few percent runtime performance. **lld** covers the rest, on the gnu leg. arm64 Linux otherwise links through GNU `ld` where x86_64 already defaults to `rust-lld`, which is why only the arm64 legs hit this at all; on a comparable arm64 build (`lancedb/sophon#7313`) it cut the largest single linker process from 7.0 to 4.0 GiB and wall time by 35%. The flags live in a small wrapper script used as the linker rather than in `-C link-arg`, because the per-target rustflags variable does not reach every unit that links: dependency crates linking a dylib (`crc-fast`, `lance-arrow`) were invoked as bare `clang`, which targets the x86_64 host and fails with `Relocations in generic ELF (EM: 183)`. Separately, and affecting five legs rather than two: the three ThinLTO targets exported `CARGO_PROFILE_RELEASE_LTO` and `CARGO_PROFILE_RELEASE_CODEGEN_UNITS` from `pre_build`, which runs inside the build step — after the cache step. `Swatinem/rust-cache` computes its key when the action runs, before any step, so step-local values are invisible to it. The result is a loop that never converges: the key never changes, so restores are exact hits, an exact hit makes the post-run save a no-op, and cargo invalidates the restored artifacts anyway because the flags differ. Those legs have been rebuilding cold on every run. Both values move to job-level `env:` ahead of the cache step, driven by new `lto:`/`codegen_units:` matrix fields, and are forwarded into the containers with `-e` since `docker run` inherits nothing. Every leg's cache key shifts once as a result, so expect one cold rebuild. A `Report peak memory` step is added so whether these legs fit is a number rather than an inference from whether the runner survived. It produced the figures above. ## Not included Moving these legs to native arm64 runners. It would retire the zig cross path, the `AT_HWCAP2` workaround and the `TARGET_CC` override, and arm64 runners are billed roughly 37% below x64 at equal core count — but the `lts-debian-aarch64` image exists to link against the manylinux2014 sysroot's glibc 2.17, and building natively on ubuntu-24.04 would raise the minimum glibc for every published aarch64 binary. That is a user-facing decision, not a CI cleanup. Dropping these legs to smaller runners, which is where the real cost saving is — larger runners are billed even on public repos. On these numbers it is not available yet: musl at 16516 MiB is about 130 MiB over what a 16 GB standard runner has. Worth revisiting as a follow-up. ## Testing Cargo's rustflags precedence was checked locally rather than taken from the docs, since getting it wrong would silently change the published binaries. With a throwaway crate carrying both a `target.'cfg(all())'` and a per-target rustflags table: setting `RUSTFLAGS` discards both, and setting it to the empty string discards them too. That rules out routing the linker flag through a job-level `RUSTFLAGS`, because `env:` keys cannot be conditionally omitted and every other leg would then silently lose the `target-cpu`/`target-feature` settings in `.cargo/config.toml` — `+avx2` on x86_64 and `-crt-static` on aarch64-musl. --------- Co-authored-by: Claude Opus 5 (1M context) --- .github/workflows/npm-publish.yml | 102 ++++++++++++++++++++++-------- 1 file changed, 74 insertions(+), 28 deletions(-) diff --git a/.github/workflows/npm-publish.yml b/.github/workflows/npm-publish.yml index 72ef5ad13..ee6906e00 100644 --- a/.github/workflows/npm-publish.yml +++ b/.github/workflows/npm-publish.yml @@ -40,40 +40,31 @@ jobs: - target: aarch64-apple-darwin host: macos-latest features: fp16kernels + # Fat LTO was ~111 of this job's ~113 minutes. + lto: thin + codegen_units: 16 pre_build: |- brew install protobuf - # Fat LTO (the workspace default in .cargo/config.toml) is - # single-threaded and is the peak-memory step of the build. On - # this runner it accounted for ~111 of the job's ~113 minutes, - # making it the critical path of the entire publish pipeline. - # ThinLTO parallelizes it across the runner's cores, for a few - # percent of runtime performance. - export CARGO_PROFILE_RELEASE_LTO=thin - export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16 - target: x86_64-pc-windows-msvc host: windows-2025 features: "," + # The lower peak also keeps this on the standard 4-core runner. + lto: thin + codegen_units: 16 pre_build: |- choco install --no-progress protoc ninja nasm tail -n 1000 /c/ProgramData/chocolatey/logs/chocolatey.log # There is an issue where choco doesn't add nasm to the path export PATH="$PATH:/c/Program Files/NASM" nasm -v - # See the ThinLTO note on aarch64-apple-darwin above. Keeping - # peak memory down is also what lets this run on the standard - # 4-core runner: the 8-core larger runner was only needed to - # stop fat LTO from OOMing rustc-LLVM. - export CARGO_PROFILE_RELEASE_LTO=thin - export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16 - target: aarch64-pc-windows-msvc host: windows-2025 features: "," + lto: thin + codegen_units: 16 pre_build: |- choco install --no-progress protoc rustup target add aarch64-pc-windows-msvc - # See the ThinLTO note on aarch64-apple-darwin above. - export CARGO_PROFILE_RELEASE_LTO=thin - export CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16 - target: x86_64-unknown-linux-gnu host: ubuntu-latest features: fp16kernels @@ -103,6 +94,14 @@ jobs: # https://github.com/napi-rs/napi-rs/blob/main/debian-aarch64.Dockerfile docker: ghcr.io/napi-rs/napi-rs/nodejs-rust:lts-debian-aarch64 features: "fp16kernels" + # Fat LTO OOM-killed rustc every nightly; even with lld it peaked + # at 31391 MiB of the runner's 32 GiB. + lto: thin + codegen_units: 16 + # arm64 Linux links through GNU `ld` where x86_64 defaults to + # `rust-lld`, which is why only arm64 OOM'd. lld cut the largest + # linker process 7.0 -> 4.0 GiB (lancedb/sophon#7313). + linker: /tmp/aarch64-lld-clang pre_build: |- set -e && apt-get update && @@ -112,9 +111,30 @@ jobs: # AT_HWCAP2 (added in Linux 3.17). Define it for aws-lc-sys. export CFLAGS="$CFLAGS -DAT_HWCAP2=26" && rustup target add aarch64-unknown-linux-gnu + # Not `&&`-chained: in dash, errexit does not fire for a + # non-final command in an `&&` list, so failures were ignored. + # + # A wrapper rather than `-C link-arg` because the per-target + # rustflags variable does not reach every unit that links, while + # the linker variable does. `clang` because GCC silently ignores + # `-fuse-ld=lld` unless built with lld support. Two echoes + # because printf's newline escape gets rewritten to `;` between + # here and the container. + echo '#!/bin/sh' > /tmp/aarch64-lld-clang + echo 'exec clang --target=aarch64-unknown-linux-gnu --sysroot=/usr/aarch64-unknown-linux-gnu/aarch64-unknown-linux-gnu/sysroot --gcc-toolchain=/usr/aarch64-unknown-linux-gnu -fuse-ld=lld "$@"' >> /tmp/aarch64-lld-clang + chmod 0755 /tmp/aarch64-lld-clang + # Fail now, not at the cdylib link ~30 minutes later. Linking at + # all also proves lld resolved; clang errors out when it cannot. + echo 'int main(void){return 0;}' > /tmp/probe.c + /tmp/aarch64-lld-clang /tmp/probe.c -o /tmp/probe + readelf -h /tmp/probe | grep AArch64 - target: aarch64-unknown-linux-musl host: ubuntu-2404-8x-x64 features: "," + # Fat LTO took the whole runner down. lld cannot help: it died + # inside rustc's LLVM, before any linker was spawned. + lto: thin + codegen_units: 16 pre_build: |- set -e && sudo apt-get update && @@ -123,6 +143,19 @@ jobs: export EXTRA_ARGS="-x" name: build - ${{ matrix.settings.target }} runs-on: ${{ matrix.settings.host }} + # On the job, not exported from `pre_build`: `Swatinem/rust-cache` hashes + # `CARGO_*` into its cache key before any step runs, so a step-local export + # leaves the key unchanged while cargo still rebuilds cold. The ThinLTO + # legs had been doing that every run. + # + # Not `RUSTFLAGS`: setting it, even to "", discards every config-file + # rustflag, silently dropping .cargo/config.toml's `target-cpu` and + # `target-feature` from the published binaries. + env: + CARGO_PROFILE_RELEASE_LTO: ${{ matrix.settings.lto || 'fat' }} + CARGO_PROFILE_RELEASE_CODEGEN_UNITS: ${{ matrix.settings.codegen_units || '1' }} + # Empty elsewhere: a per-target variable is only read for that triple. + CARGO_TARGET_AARCH64_UNKNOWN_LINUX_GNU_LINKER: ${{ matrix.settings.linker }} defaults: run: working-directory: nodejs @@ -169,19 +202,15 @@ jobs: # creating ref). The nightly cadence also keeps entries inside # GitHub's 7-day eviction window, which a tag-only trigger would not. save-if: ${{ github.ref == 'refs/heads/main' }} - # Docker builds can use rust-cache too. `target/` already lives on the - # host because the whole workspace is bind-mounted into the container, and - # rust-cache's prune and save run host-side, so they can manage it -- which - # is what keeps the entry to dependency artifacts rather than a multi-GB - # copy of everything. + # Docker builds can use rust-cache too: the workspace is bind-mounted, so + # `target/` lives on the host and rust-cache's prune keeps the entry + # small. # # Two differences from the native builds. The container's CARGO_HOME is - # bind-mounted from `.cargo-cache` rather than the host's ~/.cargo, so that - # has to be cached explicitly. And the key is derived from the *host* rustc - # version, which is not the compiler that produced these artifacts; that is - # safe because cargo fingerprints the real compiler and rebuilds on a - # mismatch, it just means a base-image toolchain bump costs one cold build - # instead of invalidating the key. + # bind-mounted from `.cargo-cache` rather than ~/.cargo, so that is cached + # explicitly. And the key uses the *host* rustc version, not the compiler + # that built these artifacts -- safe, since cargo fingerprints the real + # one; a base-image bump just costs one cold build. - name: Cache cargo (docker builds) uses: Swatinem/rust-cache@v2 if: ${{ matrix.settings.docker }} @@ -210,9 +239,14 @@ jobs: # cache step above saves. Previously the registry mounts pointed at # `.cargo/...`, a path nothing cached, so the container re-downloaded # the whole crate registry on every run. + # + # `docker run` inherits nothing; `-e NAME` carries the job's `env:` in. options: "--user 0:0 -v ${{ github.workspace }}/.cargo-cache/git/db:/usr/local/cargo/git/db \ -v ${{ github.workspace }}/.cargo-cache/registry/cache:/usr/local/cargo/registry/cache \ -v ${{ github.workspace }}/.cargo-cache/registry/index:/usr/local/cargo/registry/index \ + -e CARGO_PROFILE_RELEASE_LTO \ + -e CARGO_PROFILE_RELEASE_CODEGEN_UNITS \ + -e CARGO_TARGET_AARCH64_UNKNOWN_LINUX_GNU_LINKER \ -v ${{ github.workspace }}:/build -w /build/nodejs" run: | set -e @@ -256,6 +290,18 @@ jobs: if: always() run: df -h shell: bash + - name: Report peak memory + if: always() && runner.os == 'Linux' + shell: bash + run: | + peak=$(find /sys/fs/cgroup -name memory.peak -readable \ + -exec cat {} + 2>/dev/null | sort -n | tail -1) + if [ -n "$peak" ]; then + echo "peak memory: $((peak / 1024 / 1024)) MiB" + else + echo "peak memory: unavailable (no readable cgroup v2 memory.peak)" + fi + free -g || true - name: Upload artifact uses: actions/upload-artifact@v7 with: From 6c8aa22704690bf9875ef446897c43e864bfc502 Mon Sep 17 00:00:00 2001 From: Wyatt Alt Date: Thu, 27 Aug 2026 17:47:27 -0700 Subject: [PATCH 6/9] feat: let a computed-column batch read its own earlier declarations (#4072) `add_columns().computed()` accepted several columns in one call but bound each against the table's schema as it stood before the call, so `a` and `b = a + 1` had to be two commits. A server staging declarations behind other schema work has no atomic way to do that, and a caller reading the builder's plural signature reasonably expects the batch to be one. Each accepted column now joins the schema the next one resolves against, so the batch is planned and committed as one. Order is the dependency order; reading ahead is still an unknown column. `validate_declarations` exposes the schema-level checks -- the Function-binding guard and the planning -- without a commit, for callers that must reject before earlier work in the same request lands; LSM state is table state and stays a commit-time check. Refresh order matters for a dependent column: `b = coalesce(a, 0)` refreshed before `a` would bake zeros from `a`'s placeholder null, and the fill-once contract keeps them. Refresh now refuses, naming the input, while a computed input still has rows a refresh of it would fill -- the same probe refresh already uses to detect a no-op. Otherwise it is one snapshot and one commit, as before; a concurrent append is not in the commit and waits for the next refresh. Refreshing dependencies on the caller's behalf was considered and rejected: it is not how materialized views or our own backfill scheduler behave, and it needs multi-commit fencing that an explicit per-row fill marker would make unnecessary. --- rust/lancedb/src/table/computed_columns.rs | 105 +++++++++++- rust/lancedb/src/table/refresh.rs | 190 +++++++++++++++++++-- 2 files changed, 270 insertions(+), 25 deletions(-) diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 2b95cb34c..0f89ca612 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -22,7 +22,7 @@ use std::collections::{BTreeSet, HashMap}; use std::sync::Arc; -use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, SchemaRef}; +use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema, SchemaRef}; use datafusion_common::tree_node::TreeNode; use datafusion_physical_plan::PhysicalExpr; use lance::dataset::NewColumnTransform; @@ -1273,6 +1273,11 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< /// refresh time: that the expression parses, that every column it reads /// exists, and that the target name is free. A declaration that survives this /// is one a refresh can always act on. +/// +/// Each accepted column joins the schema the next one resolves against, so a +/// batch may declare `a` and then `b = a + 1` in one commit. Refresh order +/// then matters, and refresh enforces it: `b` is refused while `a` still has +/// unfilled rows. pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result> { if columns.is_empty() { return Err(Error::InvalidInput { @@ -1280,11 +1285,11 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result = Vec::with_capacity(columns.len()); for (name, expression) in columns { - if schema.field_with_name(name).is_ok() || declared.contains(&name.as_str()) { + if schema.field_with_name(name).is_ok() { return Err(Error::ColumnAlreadyExists { name: name.clone() }); } @@ -1292,16 +1297,50 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result(), + schema.metadata().clone(), + )); + fields.push(field); } Ok(fields) } +/// Run the schema-level checks of +/// [`AddColumnsBuilder::computed`](super::AddColumnsBuilder::computed) against +/// `schema` without committing: the Function-binding guard and the planning of +/// every declaration. For callers that stage declarations behind other work +/// and need those rejections before any of it lands. +/// +/// Only the schema is consulted. Declaring also refuses a table with an LSM +/// write spec or retained SSTables; that is table state, checked at commit. +/// +/// ``` +/// # use std::sync::Arc; +/// # use arrow_schema::{DataType, Field, Schema}; +/// use lancedb::table::computed_columns::validate_declarations; +/// +/// let schema = Arc::new(Schema::new(vec![Field::new("x", DataType::Int32, false)])); +/// let declarations = vec![ +/// ("a".to_string(), "x + 1".to_string()), +/// ("b".to_string(), "a * 2".to_string()), +/// ]; +/// assert!(validate_declarations(schema.clone(), &declarations).is_ok()); +/// assert!(validate_declarations(schema, &[("c".into(), "random()".into())]).is_err()); +/// ``` +pub fn validate_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result<()> { + ensure_no_function_bindings_for_mutation(schema.as_ref(), "schema evolution")?; + plan(schema, columns).map(drop) +} + /// Build the transform that declares `columns` against `schema`. /// /// An all-null column is how a binding with no values yet is carried into a @@ -1340,6 +1379,22 @@ pub(super) async fn add_foreign_kind(table: &crate::Table, name: &str, kind: &st #[cfg(test)] mod tests { + /// The gate's reproducer: the validator applies the same schema-level + /// guard declaring does, so a staging caller is refused before it commits + /// anything else. + #[test] + fn test_validate_declarations_matches_schema_admission_barriers() { + let schema = Arc::new(ArrowSchema::new_with_metadata( + vec![ArrowField::new("x", DataType::Int32, true)], + HashMap::from([( + FUNCTION_BINDINGS_META_KEY.to_string(), + "not valid binding metadata".to_string(), + )]), + )); + let declarations = vec![("a".to_string(), "x + 1".to_string())]; + assert!(super::validate_declarations(schema, &declarations).is_err()); + } + #[test] fn output_arrow_type_grammar_matches_the_shared_golden() { let golden: serde_json::Value = serde_json::from_str(include_str!( @@ -1582,6 +1637,40 @@ mod tests { assert!(declared(&table).await.is_empty()); } + /// A batch may build on itself: one commit, and the later entry's inputs + /// name the earlier one. + #[tokio::test] + async fn test_a_declaration_may_read_one_declared_before_it() { + let table = table_with_ints("chain").await; + let before = table.version().await.unwrap(); + add_computed( + &table, + &[("a".into(), "x + 1".into()), ("b".into(), "a * 2".into())], + ) + .await + .unwrap(); + assert_eq!(table.version().await.unwrap(), before + 1); + let declared = declared(&table).await; + assert_eq!(declared[1].name, "b"); + assert_eq!(declared[1].inputs, vec!["a".to_string()]); + + // Order is the dependency order; reading ahead is still unknown. + let err = add_computed( + &table, + &[("c".into(), "d + 1".into()), ("d".into(), "x + 1".into())], + ) + .await + .unwrap_err(); + assert!(matches!(err, Error::InvalidExpression { column, .. } if column == "c")); + assert!( + validate_declarations( + table.schema().await.unwrap(), + &[("e".into(), "random()".into())] + ) + .is_err() + ); + } + /// A column added by an ordinary transform is materialized, not bound, so /// it carries no declaration to report. #[tokio::test] diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index 35f883411..bc2cc38d1 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -7,6 +7,16 @@ //! therefore idempotent and does not observe input mutation -- once a row is //! filled, changing what the expression reads leaves the stored result alone. //! +//! A column's computed inputs are filled first -- the dependency graph is +//! walked once, each reachable column filled once in dependency order, each +//! fill its own commit. Every fill in the pass, the requested column's +//! included, covers only the fragments of the snapshot the pass started +//! from: a commit may rebase over a concurrent append, and the fragment that +//! admits carries placeholder nulls no earlier fill covered, so it waits for +//! a later refresh rather than being read as values. Two concurrent fills of +//! one input collide on its field in lance's conflict check, so a dependent +//! fill can only commit over inputs that were durable when it read them. +//! //! Two passes per fragment. The first scans only the unfilled live rows and //! evaluates the expression over them, which yields the exact fill count and //! decides whether the fragment is staged at all -- a fragment where nothing @@ -41,7 +51,8 @@ use crate::{Error, Result}; /// The result of refreshing a computed column. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] pub struct RefreshColumnResult { - /// Rows that had a value computed. + /// Rows that had a value computed, in the requested column only; inputs + /// filled on its behalf are not counted. #[serde(default)] pub rows_filled: u64, /// The commit version associated with the operation. @@ -52,6 +63,7 @@ pub struct RefreshColumnResult { struct RefreshExecution { result: RefreshColumnResult, source_version: u64, + published_version: Option, } /// Internal implementation of the refresh logic. @@ -74,7 +86,12 @@ async fn execute_refresh_column_with_source( let expression = declared_expression(&dataset, column)?; let schema = Arc::new(ArrowSchema::from(dataset.schema())); - let bound = Arc::new(super::computed_columns::bind(schema, column, &expression)?); + let bound = Arc::new(super::computed_columns::bind( + schema.clone(), + column, + &expression, + )?); + ensure_inputs_filled(&dataset, &schema, column, &bound).await?; let field = dataset .schema() .field(column) @@ -100,25 +117,25 @@ async fn execute_refresh_column_with_source( replacements.push(fragment.write_columns(values, &column_schema).await?); } + let source_version = dataset.version().version; if replacements.is_empty() { - let source_version = dataset.version().version; return Ok(RefreshExecution { result: RefreshColumnResult { rows_filled: 0, version: source_version, }, source_version, + published_version: None, }); } - let read_version = dataset.version().version; // The dataset's own session, so registrations and caches survive the // commit being installed on the handle. let session = dataset.session(); let new_dataset = Dataset::commit( WriteDestination::Dataset(dataset.clone()), Operation::DataReplacement { replacements }, - Some(read_version), + Some(source_version), None, None, session, @@ -133,10 +150,52 @@ async fn execute_refresh_column_with_source( rows_filled, version, }, - source_version: read_version, + source_version, + published_version: Some(version), }) } +/// Refuse while a computed input still has rows a refresh of it would fill: +/// read now, its placeholder null would be evaluated as a value and kept. +async fn ensure_inputs_filled( + dataset: &Dataset, + schema: &Arc, + column: &str, + bound: &BoundExpression, +) -> Result<()> { + for input in &bound.roots { + let Some(declaration) = schema + .field_with_name(input) + .ok() + .and_then(computed_column_from_field) + else { + continue; + }; + let ComputedColumnKind::Sql { expression } = &declaration.kind else { + return Err(Error::NotSupported { + message: format!( + "computed column '{column}' reads '{input}', whose fill state this \ + refresh cannot check; refresh '{input}' first" + ), + }); + }; + let input_bound = super::computed_columns::bind(schema.clone(), input, expression)?; + let mut unfilled = 0u64; + for fragment in dataset.get_fragments() { + unfilled += count_fragment_gains(dataset, &fragment, &input_bound, input).await?; + } + if unfilled > 0 { + return Err(Error::InvalidInput { + message: format!( + "computed column '{column}' reads '{input}', which has {unfilled} unfilled \ + rows; refresh '{input}' first" + ), + }); + } + } + Ok(()) +} + /// Run the refresh as a [`Job`] in this process. pub(crate) async fn execute_refresh_column_async( table: &NativeTable, @@ -160,8 +219,7 @@ pub(crate) async fn execute_refresh_column_async( rows_failed: 0, rows_remaining: 0, source_version: execution.source_version, - published_version: (execution.result.rows_filled > 0) - .then_some(execution.result.version), + published_version: execution.published_version, }) }))) } @@ -384,7 +442,8 @@ mod tests { .version) } - async fn read(table: &Table, column: &str) -> Vec> { + async fn read(table: &Table, column: &str) -> Vec> { + use arrow_array::{Array, Int64Array}; let batches = table .query() .select(Select::columns(&[column])) @@ -394,15 +453,19 @@ mod tests { .try_collect::>() .await .unwrap(); - let mut values: Vec> = batches + let mut values: Vec> = batches .iter() .flat_map(|batch| { - batch[column] - .as_any() - .downcast_ref::() - .unwrap() - .iter() - .collect::>() + let array = &batch[column]; + match array.as_any().downcast_ref::() { + Some(ints) => ints.iter().map(|v| v.map(i64::from)).collect::>(), + None => array + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .collect::>(), + } }) .collect(); values.sort(); @@ -414,6 +477,98 @@ mod tests { table.add(batch).execute().await.unwrap(); } + /// The gate's reproducer: `b = coalesce(a, 0)` refreshed before `a` + /// must not bake zeros from `a`'s placeholder null. It is refused, and + /// names the input, until `a` is filled -- after every append too. + #[tokio::test] + async fn test_dependent_refresh_refuses_an_unfilled_input() { + let table = table_with("dependent_refresh_order", vec![1, 2, 3]).await; + table + .add_columns() + .computed("a", "x + 1") + .computed("b", "coalesce(a, 0)") + .execute() + .await + .unwrap(); + + let err = table.refresh_column("b").await.unwrap_err(); + assert!( + matches!(&err, Error::InvalidInput { message } if message.contains("refresh 'a' first")), + "{err}" + ); + assert_eq!(read(&table, "b").await, vec![None, None, None]); + + assert_eq!(table.refresh_column("a").await.unwrap().rows_filled, 3); + assert_eq!(table.refresh_column("b").await.unwrap().rows_filled, 3); + assert_eq!(read(&table, "b").await, vec![Some(2), Some(3), Some(4)]); + + append(&table, vec![10]).await; + assert!(table.refresh_column("b").await.is_err()); + table.refresh_column("a").await.unwrap(); + assert_eq!(table.refresh_column("b").await.unwrap().rows_filled, 1); + assert_eq!( + table.count_rows(Some("b = 0".to_string())).await.unwrap(), + 0 + ); + } + + /// Names that need quoting, and a nested input, survive the trip through + /// declaration metadata and the dependency check: the recorded inputs + /// are matched by name, never re-parsed as SQL. + #[tokio::test] + async fn test_dependent_refresh_handles_awkward_column_names() { + use arrow_array::{Int32Array, StructArray}; + use arrow_schema::{DataType, Field, Fields}; + + let conn = connect("memory://").execute().await.unwrap(); + let age_fields = Fields::from(vec![Field::new("age", DataType::Int32, true)]); + let meta = StructArray::new( + age_fields.clone(), + vec![Arc::new(Int32Array::from(vec![10, 20])) as _], + None, + ); + let schema = Arc::new(arrow_schema::Schema::new(vec![ + Field::new("camelCase", DataType::Int32, true), + Field::new("with-hyphen", DataType::Int32, true), + Field::new("meta", DataType::Struct(age_fields), true), + ])); + let batch = arrow_array::RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2])) as _, + Arc::new(Int32Array::from(vec![100, 200])) as _, + Arc::new(meta) as _, + ], + ) + .unwrap(); + let table = conn + .create_table("awkward_names", batch) + .execute() + .await + .unwrap(); + + table + .add_columns() + .computed("y", "`camelCase` * 2") + .computed("z", "coalesce(y, 0) + `with-hyphen` + meta.age") + .execute() + .await + .unwrap(); + let z = crate::table::computed_columns::computed_columns( + table.schema().await.unwrap().as_ref(), + ) + .into_iter() + .find(|c| c.name == "z") + .unwrap(); + assert_eq!(z.inputs, vec!["meta.age", "with-hyphen", "y"]); + + let err = table.refresh_column("z").await.unwrap_err(); + assert!(err.to_string().contains("refresh 'y' first"), "{err}"); + assert_eq!(table.refresh_column("y").await.unwrap().rows_filled, 2); + assert_eq!(table.refresh_column("z").await.unwrap().rows_filled, 2); + assert_eq!(read(&table, "z").await, vec![Some(112), Some(224)]); + } + #[tokio::test] async fn test_refresh_fills_a_declared_column() { let table = table_with("refresh_fills", vec![1, 2, 3]).await; @@ -651,7 +806,8 @@ mod tests { let read_back = read(&table, "doubled").await; assert_eq!(read_back.len(), 20_000); - let mut expected: Vec> = values.iter().map(|v| Some(v * 2)).collect(); + let mut expected: Vec> = + values.iter().map(|v| Some(i64::from(v * 2))).collect(); expected.sort(); assert_eq!(read_back, expected); } From c94d9a2a166eca9606c8308d6d25946c44f1b6a5 Mon Sep 17 00:00:00 2001 From: Lance Release Date: Fri, 28 Aug 2026 00:51:15 +0000 Subject: [PATCH 7/9] =?UTF-8?q?Bump=20version:=200.38.0-beta.11=20?= =?UTF-8?q?=E2=86=92=200.38.0-beta.12?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .bumpversion.toml | 2 +- Cargo.lock | 6 +++--- docs/src/java/java.md | 2 +- java/lancedb-core/pom.xml | 2 +- java/pom.xml | 2 +- nodejs/Cargo.toml | 2 +- nodejs/npm/darwin-arm64/package.json | 2 +- nodejs/npm/linux-arm64-gnu/package.json | 2 +- nodejs/npm/linux-arm64-musl/package.json | 2 +- nodejs/npm/linux-x64-gnu/package.json | 2 +- nodejs/npm/linux-x64-musl/package.json | 2 +- nodejs/npm/win32-arm64-msvc/package.json | 2 +- nodejs/npm/win32-x64-msvc/package.json | 2 +- nodejs/package-lock.json | 4 ++-- nodejs/package.json | 2 +- python/Cargo.toml | 2 +- rust/lancedb/Cargo.toml | 2 +- 17 files changed, 20 insertions(+), 20 deletions(-) diff --git a/.bumpversion.toml b/.bumpversion.toml index 2e0b78bf3..3d57dc0fe 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "0.38.0-beta.11" +current_version = "0.38.0-beta.12" parse = """(?x) (?P0|[1-9]\\d*)\\. (?P0|[1-9]\\d*)\\. diff --git a/Cargo.lock b/Cargo.lock index 822689a0c..f12988efa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5402,7 +5402,7 @@ dependencies = [ [[package]] name = "lancedb" -version = "0.38.0-beta.11" +version = "0.38.0-beta.12" dependencies = [ "ahash", "anyhow", @@ -5490,7 +5490,7 @@ dependencies = [ [[package]] name = "lancedb-nodejs" -version = "0.38.0-beta.11" +version = "0.38.0-beta.12" dependencies = [ "arrow-array", "arrow-buffer", @@ -5515,7 +5515,7 @@ dependencies = [ [[package]] name = "lancedb-python" -version = "0.38.0-beta.11" +version = "0.38.0-beta.12" dependencies = [ "arrow", "async-trait", diff --git a/docs/src/java/java.md b/docs/src/java/java.md index 1ce012522..e19880e29 100644 --- a/docs/src/java/java.md +++ b/docs/src/java/java.md @@ -14,7 +14,7 @@ Add the following dependency to your `pom.xml`: com.lancedb lancedb-core - 0.38.0-beta.11 + 0.38.0-beta.12 ``` diff --git a/java/lancedb-core/pom.xml b/java/lancedb-core/pom.xml index c6acc7dfe..3864ed127 100644 --- a/java/lancedb-core/pom.xml +++ b/java/lancedb-core/pom.xml @@ -8,7 +8,7 @@ com.lancedb lancedb-parent - 0.38.0-beta.11 + 0.38.0-beta.12 ../pom.xml diff --git a/java/pom.xml b/java/pom.xml index 0b85b69df..b3521f9c6 100644 --- a/java/pom.xml +++ b/java/pom.xml @@ -6,7 +6,7 @@ com.lancedb lancedb-parent - 0.38.0-beta.11 + 0.38.0-beta.12 pom ${project.artifactId} LanceDB Java SDK Parent POM diff --git a/nodejs/Cargo.toml b/nodejs/Cargo.toml index 6496c6384..c3b69424f 100644 --- a/nodejs/Cargo.toml +++ b/nodejs/Cargo.toml @@ -1,7 +1,7 @@ [package] name = "lancedb-nodejs" edition.workspace = true -version = "0.38.0-beta.11" +version = "0.38.0-beta.12" publish = false license.workspace = true description.workspace = true diff --git a/nodejs/npm/darwin-arm64/package.json b/nodejs/npm/darwin-arm64/package.json index 38b3db7d7..2b8d43c3d 100644 --- a/nodejs/npm/darwin-arm64/package.json +++ b/nodejs/npm/darwin-arm64/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-darwin-arm64", - "version": "0.38.0-beta.11", + "version": "0.38.0-beta.12", "os": ["darwin"], "cpu": ["arm64"], "main": "lancedb.darwin-arm64.node", diff --git a/nodejs/npm/linux-arm64-gnu/package.json b/nodejs/npm/linux-arm64-gnu/package.json index a256cb1ee..6656fdd5c 100644 --- a/nodejs/npm/linux-arm64-gnu/package.json +++ b/nodejs/npm/linux-arm64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-gnu", - "version": "0.38.0-beta.11", + "version": "0.38.0-beta.12", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-gnu.node", diff --git a/nodejs/npm/linux-arm64-musl/package.json b/nodejs/npm/linux-arm64-musl/package.json index 567b785b0..f8f3e151f 100644 --- a/nodejs/npm/linux-arm64-musl/package.json +++ b/nodejs/npm/linux-arm64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-arm64-musl", - "version": "0.38.0-beta.11", + "version": "0.38.0-beta.12", "os": ["linux"], "cpu": ["arm64"], "main": "lancedb.linux-arm64-musl.node", diff --git a/nodejs/npm/linux-x64-gnu/package.json b/nodejs/npm/linux-x64-gnu/package.json index 4443a2748..20efa860a 100644 --- a/nodejs/npm/linux-x64-gnu/package.json +++ b/nodejs/npm/linux-x64-gnu/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-gnu", - "version": "0.38.0-beta.11", + "version": "0.38.0-beta.12", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-gnu.node", diff --git a/nodejs/npm/linux-x64-musl/package.json b/nodejs/npm/linux-x64-musl/package.json index 5c0710d56..4b735a687 100644 --- a/nodejs/npm/linux-x64-musl/package.json +++ b/nodejs/npm/linux-x64-musl/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-linux-x64-musl", - "version": "0.38.0-beta.11", + "version": "0.38.0-beta.12", "os": ["linux"], "cpu": ["x64"], "main": "lancedb.linux-x64-musl.node", diff --git a/nodejs/npm/win32-arm64-msvc/package.json b/nodejs/npm/win32-arm64-msvc/package.json index 648c985f1..35fee5ee0 100644 --- a/nodejs/npm/win32-arm64-msvc/package.json +++ b/nodejs/npm/win32-arm64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-arm64-msvc", - "version": "0.38.0-beta.11", + "version": "0.38.0-beta.12", "os": [ "win32" ], diff --git a/nodejs/npm/win32-x64-msvc/package.json b/nodejs/npm/win32-x64-msvc/package.json index 1f3bfdeb8..925211fbe 100644 --- a/nodejs/npm/win32-x64-msvc/package.json +++ b/nodejs/npm/win32-x64-msvc/package.json @@ -1,6 +1,6 @@ { "name": "@lancedb/lancedb-win32-x64-msvc", - "version": "0.38.0-beta.11", + "version": "0.38.0-beta.12", "os": ["win32"], "cpu": ["x64"], "main": "lancedb.win32-x64-msvc.node", diff --git a/nodejs/package-lock.json b/nodejs/package-lock.json index d46f08628..b996b0810 100644 --- a/nodejs/package-lock.json +++ b/nodejs/package-lock.json @@ -1,12 +1,12 @@ { "name": "@lancedb/lancedb", - "version": "0.38.0-beta.11", + "version": "0.38.0-beta.12", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "@lancedb/lancedb", - "version": "0.38.0-beta.11", + "version": "0.38.0-beta.12", "cpu": [ "x64", "arm64" diff --git a/nodejs/package.json b/nodejs/package.json index 665e2a522..19c7f4d32 100644 --- a/nodejs/package.json +++ b/nodejs/package.json @@ -11,7 +11,7 @@ "ann" ], "private": false, - "version": "0.38.0-beta.11", + "version": "0.38.0-beta.12", "main": "dist/index.js", "exports": { ".": "./dist/index.js", diff --git a/python/Cargo.toml b/python/Cargo.toml index b97fad0ed..0ee561977 100644 --- a/python/Cargo.toml +++ b/python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb-python" -version = "0.38.0-beta.11" +version = "0.38.0-beta.12" publish = false edition.workspace = true description = "Python bindings for LanceDB" diff --git a/rust/lancedb/Cargo.toml b/rust/lancedb/Cargo.toml index a71e3c948..881a5017e 100644 --- a/rust/lancedb/Cargo.toml +++ b/rust/lancedb/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "lancedb" -version = "0.38.0-beta.11" +version = "0.38.0-beta.12" edition.workspace = true description = "LanceDB: A serverless, low-latency vector database for AI applications" license.workspace = true From 6ab3b9eb30d449d61ae8e1ca76744babb9052af5 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Thu, 27 Aug 2026 19:06:07 -0700 Subject: [PATCH 8/9] ci: upgrade chacha20 to 0.10.2 (#4078) The pinned version was yanked due to UB in some SIMD kernels. Upgrading. --- Cargo.lock | 4 ++-- deny.toml | 12 +----------- 2 files changed, 3 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f12988efa..cff38a304 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1597,9 +1597,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "chacha20" -version = "0.10.0" +version = "0.10.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6f8d983286843e49675a4b7a2d174efe136dc93a18d69130dd18198a6c167601" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if 1.0.4", "cpufeatures 0.3.0", diff --git a/deny.toml b/deny.toml index 3672321d0..20510adc2 100644 --- a/deny.toml +++ b/deny.toml @@ -131,18 +131,13 @@ allow = [ "BSD-3-Clause", "ISC", "Unicode-3.0", - "Unicode-DFS-2016", "Zlib", "CC0-1.0", "MPL-2.0", "BSL-1.0", - "OpenSSL", # 0BSD ("BSD Zero Clause") is effectively public domain — no attribution # required. Pulled in by `mock_instant`. "0BSD", - # bzip2-1.0.6 is the permissive upstream bzip2 license (BSD-like). Pulled - # in by `libbz2-rs-sys`, the pure-Rust bzip2 implementation. - "bzip2-1.0.6", # CDLA-Permissive-2.0 is a permissive data license used by `webpki-roots` # for the Mozilla CA root bundle. Data-only, distribution-compatible. "CDLA-Permissive-2.0", @@ -150,12 +145,7 @@ allow = [ confidence-threshold = 0.8 # Per-crate license exceptions: allow a license for a specific crate only, # rather than globally via the `allow` list above. -exceptions = [ - # CDDL-1.0 (copyleft) is pulled in only as a dev/profiling dependency via - # `inferno` -> `pprof` -> `lance-testing`; it is a test dependency that we - # do not distribute, so scope the allowance to `inferno` alone. - { allow = ["CDDL-1.0"], crate = "inferno" }, -] +exceptions = [] # Crates whose license cannot be determined from Cargo metadata but whose # license we've manually confirmed from upstream. Keep this list minimal. [[licenses.clarify]] From 0559108fa94b29cc5db4d390d8b37794b9fcc41c Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Fri, 28 Aug 2026 17:23:49 +0800 Subject: [PATCH 9/9] feat: support blob computed column refresh (#4081) Computed-column planning currently sees Blob v2 storage descriptors, so expressions cannot consume payload bytes or preserve Blob semantics in their outputs. A computed declaration now derives its output field from its expression. A direct projection of a Blob v2 field inherits the source field's Blob metadata; other expressions retain their ordinary Arrow-inferred type. Declarations remain ordered, so the same rule applies across chained projections. Refresh materializes referenced Blob inputs as `LargeBinary` payload bytes and publishes inherited Blob outputs through Lance's Blob conversion path. Remote requests remain within the shared namespace contract as `{name, computed}`; the server planner is being updated in tandem to implement the same Blob-aware planning semantics, and remote enablement must be aligned with that server rollout. The existing null-as-unfilled contract remains unchanged. Row-level freshness and cell flags remain follow-up work. --- python/python/lancedb/remote/table.py | 10 +- python/python/lancedb/table.py | 15 +- python/python/tests/test_table.py | 23 + rust/lancedb/src/remote/table.rs | 8 +- rust/lancedb/src/table.rs | 4 +- rust/lancedb/src/table/computed_columns.rs | 289 ++++++++++-- rust/lancedb/src/table/refresh.rs | 522 ++++++++++++++++++++- 7 files changed, 821 insertions(+), 50 deletions(-) diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 02748b9bc..55014a423 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -67,7 +67,15 @@ from ..query import ( LanceTakeQueryBuilder, LanceVectorQueryBuilder, ) -from ..table import AsyncTable, BlobMode, Branches, IndexStatistics, Query, Table, Tags +from ..table import ( + AsyncTable, + BlobMode, + Branches, + IndexStatistics, + Query, + Table, + Tags, +) from ..types import BaseTokenizerType diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 535ed7c0d..36c5b727d 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2165,9 +2165,11 @@ class Table(ABC): Function columns are supported only on LanceDB Cloud and Enterprise. computed: Dict[str, str], optional - A map of column name to a SQL expression defining the column. The - column's type and inputs are derived from the expression, so no - data type is supplied. + A mapping from output column names to SQL expressions derives each + output field from its expression. A direct projection of a Blob v2 + field inherits Blob v2 semantics; other expressions derive their + ordinary Arrow type. Mapping order is declaration and dependency + order. Unlike ``transforms``, the expression is stored rather than evaluated now: the column is committed with no values, and rows get @@ -6268,8 +6270,11 @@ class AsyncTable: Function columns are supported only on LanceDB Cloud and Enterprise. computed: Dict[str, str], optional - A map of column name to a SQL expression defining the column. The - column's type and inputs are derived from the expression. + A mapping from output column names to SQL expressions derives each + output field from its expression. A direct projection of a Blob v2 + field inherits Blob v2 semantics; other expressions derive their + ordinary Arrow type. Mapping order is declaration and dependency + order. Unlike ``transforms``, the expression is stored rather than evaluated now: the column is committed with no values, and rows get diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 0be9e139d..fbdfac5d8 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -4087,6 +4087,29 @@ def test_computed_column_rejects_transforms_and_computed_together(tmp_path): table.add_columns({"a": "x + 1"}, computed={"b": "x * 2"}) +def test_computed_column_blob_projection_inherits_semantics(tmp_path): + schema = pa.schema([pa.field("id", pa.int64()), lancedb.blob("image")]) + db = lancedb.connect(tmp_path) + table = db.create_table("computed_column_blob", schema=schema) + table.add( + [ + {"id": 1, "image": b"hello"}, + {"id": 2, "image": b""}, + {"id": 3, "image": None}, + ] + ) + + table.add_columns(computed={"image_copy": "image", "second_copy": "image_copy"}) + assert table.refresh_column("image_copy").rows_filled == 2 + assert table.refresh_column("second_copy").rows_filled == 2 + assert table.blob_columns() == ["image", "image_copy", "second_copy"] + + hits = table.search().with_row_id(True).limit(10).to_arrow() + rows = sorted(zip(hits["id"].to_pylist(), hits["_rowid"].to_pylist())) + copied = table.fetch_blobs("second_copy", [row_id for _, row_id in rows]) + assert copied.to_pylist() == [b"hello", b"", None] + + @pytest.mark.asyncio async def test_computed_column_async(tmp_path): db = await lancedb.connect_async(tmp_path) diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 1afc2615a..57d2dc47d 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -3180,8 +3180,8 @@ impl BaseTable for RemoteTable { self.schema().await?.as_ref(), "schema evolution", )?; - // The server plans the declaration: expression validation, type - // inference and the persisted binding all happen there. + // The server plans the declaration against its table schema, including + // Blob v2 semantics inherited by a direct field projection. let entries = columns .iter() .map( @@ -7388,8 +7388,8 @@ mod tests { assert_eq!(result.version, if old_server { 0 } else { 43 }); } - /// A declaration is sent as `{name, computed}` entries for the server to - /// plan; the client never types the expression itself. + /// A declaration is sent as `{name, computed}` for the server to plan; the + /// client never types the expression itself. #[tokio::test] async fn test_add_computed_columns_sends_the_expression() { let table = Table::new_with_handler("my_table", |request| match request.url().path() { diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 8436657ca..70b86f3cf 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -750,8 +750,8 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { /// Declare computed columns, each defined by a SQL expression. /// /// Where the declaration is planned depends on the backend: a local table - /// validates and types the expression itself, a remote one sends the text - /// for the server to plan. + /// validates and types the expression itself, while a remote one sends the + /// expression for the server to plan. async fn add_computed_columns( &self, _columns: &[(String, String)], diff --git a/rust/lancedb/src/table/computed_columns.rs b/rust/lancedb/src/table/computed_columns.rs index 0f89ca612..4e8d8211e 100644 --- a/rust/lancedb/src/table/computed_columns.rs +++ b/rust/lancedb/src/table/computed_columns.rs @@ -9,29 +9,35 @@ //! refresh fills the rows. //! //! The rule is tagged by kind ([`ComputedColumnKind`]) because kinds differ in -//! where the column's type and inputs come from. A SQL expression is -//! self-describing -- both are derived from the expression, so a caller writes -//! neither -- while a kind resolved through a registry cannot be typed without -//! consulting it. Registered Functions use an exact remote version plus a -//! schema-level Function binding; unknown newer kinds remain readable and fail -//! closed before mutation. +//! where the column's type and inputs come from. A SQL expression determines +//! its inputs and physical result type. A direct projection of a Blob v2 field +//! also inherits that field's semantic type while execution continues to use +//! `LargeBinary`. A kind resolved through a registry cannot be typed without +//! consulting it. +//! Registered Functions use an exact remote version plus a schema-level +//! Function binding; unknown newer kinds remain readable and fail closed +//! before mutation. //! //! [`computed_columns`] and [`computed_column_from_field`] read declarations //! back off a schema. -use std::collections::{BTreeSet, HashMap}; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::sync::Arc; use arrow_schema::{DataType, Field as ArrowField, Fields, Schema as ArrowSchema, SchemaRef}; -use datafusion_common::tree_node::TreeNode; +use datafusion_common::{ScalarValue, tree_node::TreeNode}; +use datafusion_expr::Expr; use datafusion_physical_plan::PhysicalExpr; use lance::dataset::NewColumnTransform; +use lance_arrow::FieldExt; +use lance_core::datatypes::{BLOB_V2_DESC_FIELD, format_field_path_minimal, parse_field_path}; use lance_datafusion::planner::Planner; use lance_namespace::models::{JsonArrowDataType, JsonArrowField, JsonArrowSchema}; use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::function::{FunctionApplication, FunctionBinding}; +use crate::utils::resolve_arrow_field_path; use crate::{Error, Result}; /// Field metadata key marking a column as computed. The value is `"true"`. @@ -1106,15 +1112,20 @@ pub(crate) fn ensure_no_foreign_declarations<'a>( fields: impl IntoIterator>, ) -> Result<()> { for field in fields { - if field.metadata().keys().any(|k| is_declaration_key(k)) { - return Err(Error::InvalidInput { - message: format!( - "field '{}' carries computed-column metadata; declare computed columns \ - with add_columns().computed()", - field.name() - ), - }); - } + ensure_no_foreign_declaration(field)?; + } + Ok(()) +} + +fn ensure_no_foreign_declaration(field: &ArrowField) -> Result<()> { + if field.metadata().keys().any(|k| is_declaration_key(k)) { + return Err(Error::InvalidInput { + message: format!( + "field '{}' carries computed-column metadata; declare computed columns \ + with add_columns().computed()", + field.name() + ), + }); } Ok(()) } @@ -1162,15 +1173,154 @@ pub(crate) struct BoundExpression { /// The columns the expression names, as written; nested inputs keep /// their dotted path. pub inputs: Vec, - /// The top-level columns evaluation reads, in [`Self::read_schema`] - /// order. A nested input appears through its root. + /// The top-level columns evaluation reads, in physical-expression order. + /// A nested input appears through its root. pub roots: Vec, - /// The projected schema evaluation runs against. - pub read_schema: SchemaRef, /// The compiled expression. pub physical: Arc, /// The type the expression yields. pub data_type: DataType, + /// Blob v2 leaves the scan must materialize as `LargeBinary`. + pub blob_paths: Vec, + /// A directly projected Blob v2 field whose semantics the output inherits. + projected_blob_field: Option, +} + +fn is_direct_field_projection(expr: &Expr) -> bool { + match expr { + Expr::Column(_) => true, + Expr::ScalarFunction(function) + if function.name() == "get_field" && function.args.len() == 2 => + { + is_direct_field_projection(&function.args[0]) + && matches!( + &function.args[1], + Expr::Literal(ScalarValue::Utf8(Some(_)), _) + ) + } + _ => false, + } +} + +fn projected_blob_field(schema: &ArrowSchema, expr: &Expr) -> Result> { + if !is_direct_field_projection(expr) { + return Ok(None); + } + let paths = Planner::column_names_in_expr(expr); + let [path] = paths.as_slice() else { + return Ok(None); + }; + let (_, field) = resolve_arrow_field_path(schema, path)?; + Ok(field.is_blob_v2().then_some(field)) +} + +fn collect_blob_paths(field: &ArrowField, parent: &[String], paths: &mut Vec>) { + let mut path = parent.to_vec(); + path.push(field.name().clone()); + if field.is_blob_v2() { + paths.push(path); + return; + } + match field.data_type() { + DataType::Struct(children) => { + for child in children { + collect_blob_paths(child, &path, paths); + } + } + DataType::List(child) + | DataType::LargeList(child) + | DataType::FixedSizeList(child, _) + | DataType::Map(child, _) => collect_blob_paths(child, &path, paths), + _ => {} + } +} + +fn schema_blob_paths(schema: &ArrowSchema) -> Vec> { + let mut paths = Vec::new(); + for field in schema.fields() { + collect_blob_paths(field, &[], &mut paths); + } + paths +} + +fn transform_blob_field( + field: &ArrowField, + parent: &[String], + materialized: &HashSet>, +) -> ArrowField { + let mut path = parent.to_vec(); + path.push(field.name().clone()); + if field.is_blob_v2() { + if materialized.contains(&path) { + return ArrowField::new(field.name(), DataType::LargeBinary, field.is_nullable()); + } + return ArrowField::new( + field.name(), + BLOB_V2_DESC_FIELD.data_type().clone(), + field.is_nullable(), + ) + .with_metadata(BLOB_V2_DESC_FIELD.metadata().clone()); + } + + let data_type = match field.data_type() { + DataType::Struct(children) => DataType::Struct( + children + .iter() + .map(|child| Arc::new(transform_blob_field(child, &path, materialized))) + .collect(), + ), + DataType::List(child) => { + DataType::List(Arc::new(transform_blob_field(child, &path, materialized))) + } + DataType::LargeList(child) => { + DataType::LargeList(Arc::new(transform_blob_field(child, &path, materialized))) + } + DataType::FixedSizeList(child, size) => DataType::FixedSizeList( + Arc::new(transform_blob_field(child, &path, materialized)), + *size, + ), + DataType::Map(child, sorted) => DataType::Map( + Arc::new(transform_blob_field(child, &path, materialized)), + *sorted, + ), + _ => return field.clone(), + }; + ArrowField::new(field.name(), data_type, field.is_nullable()) + .with_metadata(field.metadata().clone()) +} + +fn blob_runtime_schema(schema: &ArrowSchema, materialized: &HashSet>) -> SchemaRef { + Arc::new(ArrowSchema::new_with_metadata( + schema + .fields() + .iter() + .map(|field| Arc::new(transform_blob_field(field, &[], materialized))) + .collect::(), + schema.metadata().clone(), + )) +} + +fn referenced_blob_paths(schema: &ArrowSchema, inputs: &[String]) -> Result>> { + let input_paths = inputs + .iter() + .map(|input| { + parse_field_path(input).map_err(|error| Error::InvalidInput { + message: format!("invalid computed-column input path '{input}': {error}"), + }) + }) + .collect::>>()?; + Ok(schema_blob_paths(schema) + .into_iter() + .filter(|blob_path| { + input_paths.iter().any(|input_path| { + input_path.len() <= blob_path.len() + && input_path + .iter() + .zip(blob_path) + .all(|(input, blob)| input == blob) + }) + }) + .collect()) } /// Parse, resolve and compile `expression` against `schema`. @@ -1185,10 +1335,18 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< message, }; - let planner = Planner::new(schema.clone()); + // Blob v2 is a semantic type whose runtime expression ABI is + // `LargeBinary`. Parse against that ABI first so a direct Blob reference + // is not mistaken for its storage descriptor struct. + let all_blob_paths = schema_blob_paths(schema.as_ref()) + .into_iter() + .collect::>(); + let parsing_schema = blob_runtime_schema(schema.as_ref(), &all_blob_paths); + let planner = Planner::new(parsing_schema); let parsed = planner .parse_expr(expression) .map_err(|e| invalid(e.to_string()))?; + let projected_blob_field = projected_blob_field(schema.as_ref(), &parsed)?; // A declaration is evaluated more than once -- staging and writing are // separate passes, and a refresh years later replays the same text -- so @@ -1218,13 +1376,19 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< inputs.sort(); inputs.dedup(); + let blob_paths = referenced_blob_paths(schema.as_ref(), &inputs)?; + let runtime_schema = blob_runtime_schema( + schema.as_ref(), + &blob_paths.iter().cloned().collect::>(), + ); + // A nested input is recorded by its path but read through its root // column; Schema::index_of resolves top-level names only. Resolved here // rather than left to the planner so an unknown column names itself in // the error instead of surfacing as a plan failure. let mut indices = Vec::with_capacity(inputs.len()); for input in &inputs { - let index = schema + let index = runtime_schema .index_of(root(input)) .map_err(|_| invalid(format!("unknown column '{input}'")))?; if !indices.contains(&index) { @@ -1237,7 +1401,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< // compiles the expression has to be built on the projected schema // evaluation will actually read. let read_schema = Arc::new( - schema + runtime_schema .project(&indices) .map_err(|e| invalid(e.to_string()))?, ); @@ -1247,7 +1411,8 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< .map(|field| field.name().clone()) .collect(); - let optimized = planner + let runtime_planner = Planner::new(runtime_schema); + let optimized = runtime_planner .optimize_expr(parsed) .map_err(|e| invalid(e.to_string()))?; let physical = Planner::new(read_schema.clone()) @@ -1260,9 +1425,16 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< Ok(BoundExpression { inputs, roots, - read_schema, physical, data_type, + blob_paths: blob_paths + .iter() + .map(|path| { + let segments = path.iter().map(String::as_str).collect::>(); + format_field_path_minimal(&segments) + }) + .collect(), + projected_blob_field, }) } @@ -1278,7 +1450,7 @@ pub(crate) fn bind(schema: SchemaRef, column: &str, expression: &str) -> Result< /// batch may declare `a` and then `b = a + 1` in one commit. Refresh order /// then matters, and refresh enforces it: `b` is refused while `a` still has /// unfilled rows. -pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result> { +fn plan_declarations(schema: SchemaRef, columns: &[(String, String)]) -> Result> { if columns.is_empty() { return Err(Error::InvalidInput { message: "at least one computed column is required".into(), @@ -1290,15 +1462,28 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result { + let mut metadata = source.metadata().clone(); + metadata.retain(|key, _| !is_declaration_key(key)); + metadata.extend(computed_metadata); + source + .with_name(name) + .with_nullable(true) + .with_metadata(metadata) + } + None => ArrowField::new(name, bound.data_type, true).with_metadata(computed_metadata), + }; schema = Arc::new(ArrowSchema::new_with_metadata( schema .fields() @@ -1314,6 +1499,10 @@ pub(crate) fn plan(schema: SchemaRef, columns: &[(String, String)]) -> Result Result> { + plan_declarations(schema, columns) +} + /// Run the schema-level checks of /// [`AddColumnsBuilder::computed`](super::AddColumnsBuilder::computed) against /// `schema` without committing: the Function-binding guard and the planning of @@ -1352,7 +1541,7 @@ pub(crate) fn declare( schema: SchemaRef, columns: &[(String, String)], ) -> Result { - let fields = plan(schema, columns)?; + let fields = plan_declarations(schema, columns)?; Ok(NewColumnTransform::AllNulls(Arc::new(ArrowSchema::new( fields, )))) @@ -1478,6 +1667,44 @@ mod tests { ); } + #[test] + fn test_direct_blob_projection_inherits_semantics() { + let schema = Arc::new(ArrowSchema::new(vec![crate::blob("image", false)])); + let fields = plan( + schema, + &[ + ("first".to_string(), "image".to_string()), + ("second".to_string(), "first".to_string()), + ], + ) + .unwrap(); + + for field in &fields { + assert!(field.is_blob_v2()); + assert!(field.is_nullable()); + } + assert_eq!( + fields[1] + .metadata() + .get(EXPRESSION_META_KEY) + .map(String::as_str), + Some("first") + ); + } + + #[test] + fn test_blob_expression_transformation_does_not_inherit_semantics() { + let schema = Arc::new(ArrowSchema::new(vec![crate::blob("image", true)])); + let fields = plan( + schema, + &[("payload".to_string(), "coalesce(image, image)".to_string())], + ) + .unwrap(); + + assert!(!fields[0].is_blob_v2()); + assert_eq!(fields[0].data_type(), &DataType::LargeBinary); + } + /// The binding reaches the schema only if `AllNulls` carries per-field /// metadata through the commit. The whole representation rests on it. #[tokio::test] diff --git a/rust/lancedb/src/table/refresh.rs b/rust/lancedb/src/table/refresh.rs index bc2cc38d1..511fce8ff 100644 --- a/rust/lancedb/src/table/refresh.rs +++ b/rust/lancedb/src/table/refresh.rs @@ -29,10 +29,14 @@ //! inputs masked to null first, so a poison value in a row nobody is filling //! cannot fail the refresh. +use std::collections::HashSet; use std::sync::Arc; -use arrow_array::{ArrayRef, BooleanArray, RecordBatch, RecordBatchOptions}; -use arrow_schema::Schema as ArrowSchema; +use arrow_array::{ + Array, ArrayRef, BooleanArray, LargeBinaryArray, RecordBatch, RecordBatchOptions, StructArray, + new_null_array, +}; +use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema}; use datafusion_expr::ColumnarValue; use futures::{Stream, StreamExt, TryStreamExt}; use lance::Dataset; @@ -40,7 +44,7 @@ use lance::dataset::WriteDestination; use lance::dataset::fragment::FileFragment; use lance::dataset::transaction::Operation; use lance_core::ROW_ID; -use lance_core::datatypes::Schema as LanceSchema; +use lance_core::datatypes::{BlobHandling, Schema as LanceSchema}; use serde::{Deserialize, Serialize}; use super::computed_columns::{BoundExpression, ComputedColumnKind, computed_column_from_field}; @@ -104,6 +108,7 @@ async fn execute_refresh_column_with_source( fields: vec![field.clone()], metadata: Default::default(), }; + let output_is_blob = field.is_blob_v2(); let mut rows_filled = 0u64; let mut replacements = Vec::new(); @@ -113,7 +118,8 @@ async fn execute_refresh_column_with_source( continue; } rows_filled += gained; - let values = fill_stream(&dataset, &fragment, bound.clone(), column).await?; + let values = + fill_stream(&dataset, &fragment, bound.clone(), column, output_is_blob).await?; replacements.push(fragment.write_columns(values, &column_schema).await?); } @@ -294,12 +300,15 @@ fn evaluation_batch( mask_out: Option<&BooleanArray>, ) -> lance_core::Result { let mut columns = Vec::with_capacity(bound.roots.len()); + let mut fields = Vec::with_capacity(bound.roots.len()); for name in &bound.roots { - let column = batch.column_by_name(name).ok_or_else(|| { + let index = batch.schema_ref().index_of(name).map_err(|_| { lance_core::Error::invalid_input(format!( "refreshing a computed column read no {name} column" )) })?; + let column = batch.column(index); + fields.push(batch.schema_ref().field(index).clone()); // Rows outside the mask must not reach the expression: a value in a // deleted or already-filled row can be one it would choke on. columns.push(match mask_out { @@ -308,7 +317,7 @@ fn evaluation_batch( }); } Ok(RecordBatch::try_new_with_options( - bound.read_schema.clone(), + Arc::new(ArrowSchema::new(fields)), columns, &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), )?) @@ -329,6 +338,99 @@ fn evaluate(bound: &BoundExpression, batch: &RecordBatch) -> lance_core::Result< } } +fn materialized_blob_ids(schema: &LanceSchema, paths: &[String]) -> Result> { + paths + .iter() + .map(|path| { + let field = schema + .resolve(path) + .and_then(|fields| fields.last().copied()) + .ok_or_else(|| Error::InvalidInput { + message: format!("computed Blob input '{path}' no longer exists"), + })?; + if !field.is_blob_v2() { + return Err(Error::InvalidInput { + message: format!("computed Blob input '{path}' is no longer Blob v2"), + }); + } + u32::try_from(field.id).map_err(|_| Error::InvalidInput { + message: format!( + "computed Blob input '{path}' has invalid field id {}", + field.id + ), + }) + }) + .collect() +} + +fn configure_blob_inputs( + scanner: &mut lance::dataset::scanner::Scanner, + schema: &LanceSchema, + bound: &BoundExpression, + extra_blob_id: Option, +) -> Result<()> { + let mut ids = materialized_blob_ids(schema, &bound.blob_paths)?; + ids.extend(extra_blob_id); + scanner.blob_handling(BlobHandling::SomeBlobsBinary(ids)); + Ok(()) +} + +fn blob_array_from_binary( + array: &ArrayRef, + target_field: &ArrowField, +) -> lance_core::Result { + let values = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + lance_core::Error::invalid_input(format!( + "a Blob v2 computed output produced {}, expected LargeBinary", + array.data_type() + )) + })?; + let mut builder = lance::blob::BlobArrayBuilder::new(values.len()); + for index in 0..values.len() { + if values.is_null(index) { + builder.push_null()?; + } else { + builder.push_bytes(values.value(index))?; + } + } + let minimal = builder.finish()?; + let minimal = minimal + .as_any() + .downcast_ref::() + .ok_or_else(|| lance_core::Error::internal("Blob builder returned a non-struct array"))?; + let DataType::Struct(target_fields) = target_field.data_type() else { + return Err(lance_core::Error::invalid_input(format!( + "Blob v2 output field '{}' has non-struct type {}", + target_field.name(), + target_field.data_type() + ))); + }; + let columns = target_fields + .iter() + .map(|field| match field.name().as_str() { + "data" | "uri" => minimal + .column_by_name(field.name()) + .cloned() + .ok_or_else(|| { + lance_core::Error::internal(format!("Blob builder omitted '{}'", field.name())) + }), + "position" | "size" => Ok(new_null_array(field.data_type(), minimal.len())), + name => Err(lance_core::Error::invalid_input(format!( + "Blob v2 output field '{}' has unsupported logical child '{name}'", + target_field.name() + ))), + }) + .collect::>>()?; + Ok(Arc::new(StructArray::try_new( + target_fields.clone(), + columns, + minimal.nulls().cloned(), + )?)) +} + /// How many rows of one fragment would gain a value. /// /// Scans only the unfilled live rows -- deleted rows never reach the @@ -347,6 +449,7 @@ async fn count_fragment_gains( .with_row_id() .filter(&format!("{} IS NULL", quote_identifier(column)))? .project(&bound.roots)?; + configure_blob_inputs(&mut scanner, dataset.schema(), bound, None)?; let mut gained = 0u64; let mut batches = scanner.try_into_stream().await?; @@ -368,6 +471,7 @@ async fn fill_stream( fragment: &FileFragment, bound: Arc, column: &str, + output_is_blob: bool, ) -> Result> + Send + use<>> { let mut projection: Vec = bound.roots.clone(); projection.push(column.to_string()); @@ -377,6 +481,20 @@ async fn fill_stream( .with_row_id() .include_deleted_rows() .project(&projection)?; + let output_blob_id = output_is_blob + .then(|| { + dataset + .schema() + .field(column) + .and_then(|field| u32::try_from(field.id).ok()) + }) + .flatten(); + configure_blob_inputs( + &mut scanner, + dataset.schema(), + bound.as_ref(), + output_blob_id, + )?; let projected = Arc::new(ArrowSchema::new(vec![ ArrowSchema::from(dataset.schema()) @@ -412,6 +530,11 @@ async fn fill_stream( let computed = evaluate(&bound, &evaluation_batch(&batch, &bound, Some(&keep))?)?; let merged = arrow_select::zip::zip(&fill, &computed, existing)?; + let merged = if output_is_blob { + blob_array_from_binary(&merged, projected.field(0))? + } else { + merged + }; Ok(RecordBatch::try_new(projected.clone(), vec![merged])?) })) } @@ -420,8 +543,12 @@ async fn fill_stream( mod tests { use std::sync::Arc; - use arrow_array::{Int32Array, record_batch}; + use arrow_array::{ + Array, ArrayRef, Int32Array, LargeBinaryArray, RecordBatch, StructArray, record_batch, + }; + use arrow_schema::Field as ArrowField; use futures::TryStreamExt; + use lance_core::ROW_ID; use crate::connect; use crate::query::{ExecutableQuery, QueryBase, Select}; @@ -477,6 +604,25 @@ mod tests { table.add(batch).execute().await.unwrap(); } + #[test] + fn test_blob_output_matches_complete_logical_field() { + let values: ArrayRef = Arc::new(LargeBinaryArray::from(vec![ + Some(b"hello".as_slice()), + None, + ])); + let field = ArrowField::new( + "image", + lance_core::datatypes::BLOB_V2_LOGICAL_TYPE.clone(), + true, + ); + + let output = super::blob_array_from_binary(&values, &field).unwrap(); + assert_eq!(output.data_type(), field.data_type()); + let output = output.as_any().downcast_ref::().unwrap(); + assert_eq!(output.column_by_name("position").unwrap().null_count(), 2); + assert_eq!(output.column_by_name("size").unwrap().null_count(), 2); + } + /// The gate's reproducer: `b = coalesce(a, 0)` refreshed before `a` /// must not bake zeros from `a`'s placeholder null. It is refused, and /// names the input, until `a` is filled -- after every append too. @@ -1164,4 +1310,366 @@ mod tests { let err = table.refresh_column("embedding").await.unwrap_err(); assert!(matches!(err, Error::NotSupported { message } if message.contains("udf"))); } + + fn blob_batch(ids: Vec, payloads: Vec>) -> RecordBatch { + use arrow_array::Int32Array; + use arrow_schema::{Field, Schema}; + + let mut builder = lance::blob::BlobArrayBuilder::new(payloads.len()); + for payload in payloads { + match payload { + Some(payload) => builder.push_bytes(payload).unwrap(), + None => builder.push_null().unwrap(), + } + } + RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", arrow_schema::DataType::Int32, false), + crate::blob("image", true), + ])), + vec![Arc::new(Int32Array::from(ids)), builder.finish().unwrap()], + ) + .unwrap() + } + + async fn create_blob_table(path: &std::path::Path, batch: RecordBatch) -> Table { + let conn = connect(path.to_str().unwrap()).execute().await.unwrap(); + conn.create_table("blobs", batch).execute().await.unwrap() + } + + #[tokio::test] + async fn test_refresh_inherits_and_publishes_blob_output() { + use arrow_array::UInt64Array; + use lance_arrow::{ + BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, BLOB_INLINE_SIZE_THRESHOLD_META_KEY, + }; + use lance_core::datatypes::BlobKind; + + use crate::table::schema_evolution::FieldMetadataUpdate; + + let tmp = tempfile::tempdir().unwrap(); + let table = create_blob_table( + tmp.path(), + blob_batch( + vec![1, 2, 3, 4], + vec![Some(b"hello"), Some(b"ab"), Some(b""), None], + ), + ) + .await; + table + .add_columns() + .computed("image_copy", "image") + .execute() + .await + .unwrap(); + table + .update_field_metadata(&[FieldMetadataUpdate::new("image_copy") + .set(BLOB_INLINE_SIZE_THRESHOLD_META_KEY, "1") + .set(BLOB_DEDICATED_SIZE_THRESHOLD_META_KEY, "4")]) + .await + .unwrap(); + + let first_refresh = table.refresh_column("image_copy").await.unwrap(); + assert_eq!(first_refresh.rows_filled, 3); + assert_eq!( + table.blob_columns().await.unwrap(), + vec!["image".to_string(), "image_copy".to_string()] + ); + + let batches = table + .query() + .with_row_id() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let batch = arrow_select::concat::concat_batches(&batches[0].schema(), &batches).unwrap(); + assert!( + batch + .column_by_name("image_copy") + .unwrap() + .as_any() + .is::() + ); + let row_ids = batch + .column_by_name(ROW_ID) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec(); + let original = table.fetch_blobs("image", &row_ids).await.unwrap(); + let copied = table.fetch_blobs("image_copy", &row_ids).await.unwrap(); + assert_eq!(original, copied); + let ids = batch + .column_by_name("id") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + let files = table + .fetch_blob_files("image_copy", &row_ids) + .await + .unwrap(); + let mut layouts = ids + .values() + .iter() + .copied() + .zip(files) + .map(|(id, file)| (id, file.and_then(|file| file.kind()))) + .collect::>(); + layouts.sort_by_key(|(id, _)| *id); + assert_eq!( + layouts, + vec![ + (1, Some(BlobKind::Dedicated)), + (2, Some(BlobKind::Packed)), + (3, Some(BlobKind::Inline)), + (4, None), + ] + ); + + table + .add(blob_batch(vec![5], vec![Some(b"appended")])) + .execute() + .await + .unwrap(); + table + .optimize(crate::table::OptimizeAction::Compact { + options: crate::table::CompactionOptions::default(), + remap_options: None, + }) + .await + .unwrap(); + assert_eq!( + table + .refresh_column("image_copy") + .await + .unwrap() + .rows_filled, + 1 + ); + assert_eq!( + table + .refresh_column("image_copy") + .await + .unwrap() + .rows_filled, + 0 + ); + + table.checkout(first_refresh.version).await.unwrap(); + assert_eq!(table.count_rows(None).await.unwrap(), 4); + assert_eq!( + table.blob_columns().await.unwrap(), + vec!["image".to_string(), "image_copy".to_string()] + ); + table.checkout_latest().await.unwrap(); + } + + #[tokio::test] + async fn test_refresh_inherits_nested_struct_blob_input() { + use arrow_array::{Int32Array, StructArray, UInt64Array}; + use arrow_schema::{DataType, Field, Fields, Schema}; + + let tmp = tempfile::tempdir().unwrap(); + let mut blob_builder = lance::blob::BlobArrayBuilder::new(2); + blob_builder.push_bytes(b"nested").unwrap(); + blob_builder.push_null().unwrap(); + let blob_field = crate::blob("image", true); + let metadata_fields = Fields::from(vec![blob_field.clone()]); + let metadata = StructArray::new( + metadata_fields.clone(), + vec![blob_builder.finish().unwrap()], + None, + ); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("metadata", DataType::Struct(metadata_fields), true), + ])), + vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(metadata)], + ) + .unwrap(); + let table = create_blob_table(tmp.path(), batch).await; + table + .add_columns() + .computed("payload_copy", "metadata.image") + .execute() + .await + .unwrap(); + + assert_eq!( + table + .refresh_column("payload_copy") + .await + .unwrap() + .rows_filled, + 1 + ); + assert_eq!( + table.blob_columns().await.unwrap(), + vec!["metadata.image".to_string(), "payload_copy".to_string()] + ); + let batches = table + .query() + .with_row_id() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let row_ids = batches[0] + .column_by_name(ROW_ID) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values(); + let payloads = table.fetch_blobs("payload_copy", row_ids).await.unwrap(); + assert_eq!(payloads.value(0), b"nested"); + assert!(payloads.is_null(1)); + } + + #[tokio::test] + async fn test_refresh_preserves_list_shape_when_materializing_blob_input() { + use arrow_array::{Int32Array, ListArray}; + use arrow_buffer::{OffsetBuffer, ScalarBuffer}; + use arrow_schema::{DataType, Field, Schema}; + + let tmp = tempfile::tempdir().unwrap(); + let mut blob_builder = lance::blob::BlobArrayBuilder::new(3); + blob_builder.push_bytes(b"a").unwrap(); + blob_builder.push_bytes(b"bb").unwrap(); + blob_builder.push_null().unwrap(); + let item = Arc::new(crate::blob("item", true)); + let images = ListArray::new( + item.clone(), + OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 3])), + blob_builder.finish().unwrap(), + None, + ); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("images", DataType::List(item), true), + ])), + vec![Arc::new(Int32Array::from(vec![1, 2])), Arc::new(images)], + ) + .unwrap(); + let table = create_blob_table(tmp.path(), batch).await; + table + .add_columns() + .computed("image_payloads", "images") + .execute() + .await + .unwrap(); + + assert_eq!( + table + .refresh_column("image_payloads") + .await + .unwrap() + .rows_filled, + 2 + ); + let batches = table + .query() + .select(Select::columns(&["image_payloads"])) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let output = batches[0] + .column_by_name("image_payloads") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(output.value_offsets(), &[0, 2, 3]); + assert!(output.values().as_any().is::()); + } + + #[tokio::test] + async fn test_refresh_inherits_external_blob_input() { + use arrow_array::{Int32Array, StringArray, UInt64Array}; + use arrow_schema::{DataType, Field, Schema}; + + let tmp = tempfile::tempdir().unwrap(); + let payload = b"external-payload"; + let path = tmp.path().join("payload.bin"); + std::fs::write(&path, payload).unwrap(); + let uri = url::Url::from_file_path(path).unwrap().to_string(); + let conn = connect(tmp.path().join("db").to_str().unwrap()) + .execute() + .await + .unwrap(); + let table = conn + .create_empty_table( + "external", + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + crate::blob("image", true), + ])), + ) + .execute() + .await + .unwrap(); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("image", DataType::Utf8, true), + ])), + vec![ + Arc::new(Int32Array::from(vec![1])), + Arc::new(StringArray::from(vec![Some(uri)])), + ], + ) + .unwrap(); + table + .add(batch) + .allow_external_blob_outside_bases(true) + .execute() + .await + .unwrap(); + table + .add_columns() + .computed("payload_copy", "image") + .execute() + .await + .unwrap(); + + assert_eq!( + table + .refresh_column("payload_copy") + .await + .unwrap() + .rows_filled, + 1 + ); + let batches = table + .query() + .with_row_id() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let row_ids = batches[0] + .column_by_name(ROW_ID) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .values(); + let payloads = table.fetch_blobs("payload_copy", row_ids).await.unwrap(); + assert_eq!(payloads.value(0), payload); + } }