From d533a181000fe32fcc3705d77a36d061bf8bf96c Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 19:06:23 +0000 Subject: [PATCH] fix(python): materialize nested blob bytes in pandas --- python/python/lancedb/_blob.py | 62 +++++++++++++++---- python/python/tests/test_blob.py | 100 +++++++++++++++++++++++++++++++ 2 files changed, 151 insertions(+), 11 deletions(-) diff --git a/python/python/lancedb/_blob.py b/python/python/lancedb/_blob.py index dc4ed37df..88d2cccc5 100644 --- a/python/python/lancedb/_blob.py +++ b/python/python/lancedb/_blob.py @@ -110,7 +110,9 @@ def projection_includes_blob_column( if projection is None: return True for output, source in _iter_projection_pairs(projection): - if output in columns or source in columns: + if output in columns or any( + column == source or column.startswith(f"{source}.") for column in columns + ): return True return False @@ -122,14 +124,16 @@ def blob_v2_projection_sources( blob_columns = row_addressable_blob_v2_paths(schema) if not blob_columns: return {} - columns = set(blob_columns) if projection is None: return {column: column for column in blob_columns} - return { - output: source - for output, source in _iter_projection_pairs(projection) - if source in columns - } + sources = {} + for output, source in _iter_projection_pairs(projection): + for column in blob_columns: + if column == source: + sources[output] = column + elif column.startswith(f"{source}."): + sources[f"{output}{column[len(source) :]}"] = column + return sources def v2_projection_needs_row_id( @@ -180,7 +184,7 @@ async def replace_v2_blob_columns_with_bytes( fetch_blobs: FetchBlobsAsync, ) -> pa.Table: for output_name, source_name in blob_sources.items(): - if output_name not in tbl.column_names: + if not _has_column_path(tbl.schema, output_name): continue blobs = await fetch_blobs(source_name, tbl) tbl = _set_blob_column(tbl, output_name, blobs) @@ -193,7 +197,7 @@ def replace_v2_blob_columns_with_bytes_sync( fetch_blobs: FetchBlobsSync, ) -> pa.Table: for output_name, source_name in blob_sources.items(): - if output_name not in tbl.column_names: + if not _has_column_path(tbl.schema, output_name): continue blobs = fetch_blobs(source_name, tbl) tbl = _set_blob_column(tbl, output_name, blobs) @@ -287,9 +291,45 @@ def _iter_projection_pairs( 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: +def _has_column_path(schema: pa.Schema, path: str) -> bool: + if schema.get_field_index(path) >= 0: + return True + top_name, *rest = path.split(".") + index = schema.get_field_index(top_name) + if index < 0: + return False + field = schema.field(index) + for name in rest: + if not pa.types.is_struct(field.type): + return False + index = field.type.get_field_index(name) + if index < 0: + return False + field = field.type.field(index) + return True + + +def _set_blob_column( + tbl: pa.Table, output_name: str, blobs: pa.Array | pa.ChunkedArray +) -> pa.Table: index = tbl.schema.get_field_index(output_name) - return tbl.set_column(index, pa.field(output_name, blobs.type), [blobs]) + if index >= 0: + return tbl.set_column(index, pa.field(output_name, blobs.type), [blobs]) + + parent_path, leaf_name = output_name.rsplit(".", 1) + if isinstance(blobs, pa.ChunkedArray): + blobs = blobs.combine_chunks() + + def replace_blob(children: list, child_fields: list) -> None: + for i, field in enumerate(child_fields): + if field.name == leaf_name: + children[i] = blobs + child_fields[i] = pa.field( + leaf_name, blobs.type, nullable=field.nullable + ) + return + + return _transform_struct_column(tbl, parent_path, replace_blob) def _embed_row_id_in_column(tbl: pa.Table, path: str, row_ids: pa.Array) -> pa.Table: diff --git a/python/python/tests/test_blob.py b/python/python/tests/test_blob.py index 0d3a89c21..a3577b2d7 100644 --- a/python/python/tests/test_blob.py +++ b/python/python/tests/test_blob.py @@ -296,6 +296,106 @@ def test_blob_v2_projection_sources_use_typed_column_name(): } +def test_blob_v2_projection_sources_include_selected_struct(): + schema = pa.schema( + [ + pa.field( + "info", pa.struct([pa.field("name", pa.string()), lancedb.blob("blob")]) + ) + ] + ) + + assert blob_v2_projection_sources(schema, ["info"]) == {"info.blob": "info.blob"} + assert blob_v2_projection_sources(schema, {"renamed": "info"}) == { + "renamed.blob": "info.blob" + } + + +def _nested_blob_table(db, name): + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field( + "info", pa.struct([pa.field("name", pa.string()), lancedb.blob("blob")]) + ), + ] + ) + table = db.create_table(name, schema=schema) + info = pa.StructArray.from_arrays( + [ + pa.array(["first", "second", "third"]), + pa.array([b"hello", None, b"hidden"], type=pa.large_binary()), + ], + names=["name", "blob"], + mask=pa.array([False, False, True]), + ) + table.add(pa.table({"id": pa.array([0, 1, 2]), "info": info})) + return table + + +def test_nested_blob_to_pandas_bytes(tmp_path): + table = _nested_blob_table(lancedb.connect(tmp_path), "nested_pandas_bytes") + expected = [ + {"name": "first", "blob": b"hello"}, + {"name": "second", "blob": None}, + None, + ] + + assert table.to_pandas(blob_mode="bytes")["info"].tolist() == expected + assert table.search().to_pandas(blob_mode="bytes")["info"].tolist() == expected + assert ( + table.search().select(["info"]).to_pandas(blob_mode="bytes")["info"].tolist() + == expected + ) + assert ( + table.search() + .select({"renamed": "info"}) + .to_pandas(blob_mode="bytes")["renamed"] + .tolist() + == expected + ) + assert table.search().select({"b": "info.blob"}).to_pandas(blob_mode="bytes")[ + "b" + ].tolist() == [b"hello", None, None] + + +@pytest.mark.asyncio +async def test_nested_blob_to_pandas_bytes_async(tmp_path): + db = await lancedb.connect_async(tmp_path) + schema = pa.schema( + [ + pa.field("id", pa.int64()), + pa.field( + "info", pa.struct([pa.field("name", pa.string()), lancedb.blob("blob")]) + ), + ] + ) + table = await db.create_table("nested_pandas_bytes_async", schema=schema) + await table.add( + pa.table( + { + "id": pa.array([0, 1]), + "info": pa.StructArray.from_arrays( + [ + pa.array(["first", "second"]), + pa.array([b"hello", None], type=pa.large_binary()), + ], + names=["name", "blob"], + ), + } + ) + ) + expected = [ + {"name": "first", "blob": b"hello"}, + {"name": "second", "blob": None}, + ] + + assert (await table.to_pandas(blob_mode="bytes"))["info"].tolist() == expected + assert (await table.query().to_pandas(blob_mode="bytes"))[ + "info" + ].tolist() == expected + + def _legacy_v1_table(name): # Legacy v1 blob columns are only writable at file version <= 2.1. db = lancedb.connect(