diff --git a/daft_lance/lance_merge_column.py b/daft_lance/lance_merge_column.py index 595bdd5..3086156 100644 --- a/daft_lance/lance_merge_column.py +++ b/daft_lance/lance_merge_column.py @@ -267,10 +267,18 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: rowaddrs = rowaddr_col.to_pylist() if hasattr(rowaddr_col, "to_pylist") else list(rowaddr_col) - # Build table of new columns + # Build table of new columns, preserving the Arrow type from the daft Series. + # pa.array(s.to_pylist()) loses type information: fixed_size_list[N] + # becomes list because Python floats are float64 and list structure is + # inferred from Python lists. Using s.to_arrow() avoids this type erasure. arrays = [] for s in data_cols: - arr = _pa.array(s.to_pylist() if hasattr(s, "to_pylist") else list(s)) + if hasattr(s, "to_arrow"): + arr = s.to_arrow() + if isinstance(arr, _pa.ChunkedArray): + arr = arr.combine_chunks() + else: + arr = _pa.array(list(s)) arrays.append(arr) tbl = _pa.table({name: arr for name, arr in zip(self.new_column_names, arrays)}) diff --git a/tests/io/lancedb/test_fast_path_merge.py b/tests/io/lancedb/test_fast_path_merge.py index 99d768b..9f69efb 100644 --- a/tests/io/lancedb/test_fast_path_merge.py +++ b/tests/io/lancedb/test_fast_path_merge.py @@ -615,3 +615,36 @@ def test_fast_path_check_does_not_set_result_cache(self, ds_path): assert getattr(df, "_result_cache", None) is None, ( "_result_cache was set — fast-path check used collect() instead of count_rows()" ) + + def test_fixed_size_list_float32_type_preserved(self, ds_path): + """Bug: pa.array(s.to_pylist()) erases fixed_size_list[N] → list. + + Python floats are float64 and list structure is inferred from nested Python lists, + so the Arrow type is lost. The fix uses s.to_arrow().combine_chunks() which + preserves the declared daft return type exactly. + """ + N = 8 + ds = create_dataset(ds_path, [{"id": [1, 2, 3]}]) + df = read_with_metadata(ds_path) + + @daft.func.batch(return_dtype=daft.DataType.fixed_size_list(daft.DataType.float32(), N)) + def _make_vec(ids): + import numpy as np + + return [np.array([float(i)] * N, dtype=np.float32) for i in ids.to_pylist()] + + df = df.with_column("embedding", _make_vec(daft.col("id"))) + ds2 = merge_columns_from_df(df, ds, ds_path) + + field = ds2.schema.field("embedding") + # Type must be preserved: fixed_size_list[N], NOT list + assert pa.types.is_fixed_size_list(field.type), f"Expected fixed_size_list, got {field.type}" + assert field.type.list_size == N + assert field.type.value_type == pa.float32(), f"Expected float32 values, got {field.type.value_type}" + + result = ds2.to_table().sort_by("id").to_pydict() + for i, emb in zip(result["id"], result["embedding"]): + assert emb is not None, f"Embedding for id={i} is null" + assert len(emb) == N + for v in emb: + assert pytest.approx(float(i), rel=1e-5) == v