diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 006ef8cb9..ff866d4db 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -5335,6 +5335,12 @@ def _infer_target_schema( # Use the most common length of the list as the dimensions dim = _modal_list_size(peeked.column(i)) + if dim is None: + # Every list entry is null (e.g. embeddings not computed yet), + # so the dimension is unknowable. Leave the column as a + # variable-length list instead of guessing a fixed size. + continue + # Determine target type based on value type if pa.types.is_floating(field.type.value_type): target_type = pa.list_(pa.float32(), dim) @@ -5374,9 +5380,14 @@ def _infer_target_schema( return schema, reader -def _modal_list_size(arr: Union[pa.ListArray, pa.ChunkedArray]) -> int: - # Use the most common length of the list as the dimensions - return pc.mode(pc.list_value_length(arr))[0].as_py()["mode"] +def _modal_list_size(arr: Union[pa.ListArray, pa.ChunkedArray]) -> Optional[int]: + # Use the most common length of the list as the dimensions. pyarrow's + # `mode` skips nulls, so an all-null column yields an empty result; return + # None in that case so the caller can keep a variable-length list. + modes = pc.mode(pc.list_value_length(arr)) + if len(modes) == 0: + return None + return modes[0].as_py()["mode"] def _infer_vector_dim(arr: Union[pa.Array, pa.ChunkedArray]) -> Optional[int]: diff --git a/python/python/tests/test_table.py b/python/python/tests/test_table.py index 7a4f5981c..827f2ef58 100644 --- a/python/python/tests/test_table.py +++ b/python/python/tests/test_table.py @@ -280,6 +280,30 @@ def test_create_table_infers_large_int_vectors(mem_db: DBConnection): assert vector_column.to_pylist() == [[0.0, 300.0]] +def test_create_table_all_null_vector_column(mem_db: DBConnection): + # A vector column can legitimately arrive with every value null, e.g. when + # the embeddings have not been computed yet. The dimension is unknowable in + # that case, so schema inference must not crash and should leave the column + # as a variable-length list rather than a fixed-size one. + data = pa.table( + { + "id": pa.array([1, 2]), + "vector": pa.array([None, None], type=pa.list_(pa.float32())), + } + ) + + table = mem_db.create_table("all_null_vector", data=data) + + vector_field = table.schema.field("vector") + assert pa.types.is_list(vector_field.type) + assert vector_field.type.value_type == pa.float32() + + result = table.to_arrow() + assert result.num_rows == 2 + assert result.column("vector").to_pylist() == [None, None] + assert table.count_rows() == 2 + + @pytest.mark.asyncio async def test_create_table_async_infers_large_int_vectors( mem_db_async: AsyncConnection,