diff --git a/daft_lance/lance_merge_column.py b/daft_lance/lance_merge_column.py index 3086156..8984120 100644 --- a/daft_lance/lance_merge_column.py +++ b/daft_lance/lance_merge_column.py @@ -314,8 +314,10 @@ def __call__(self, *cols: Any) -> list[dict[str, bytes]]: writer.write_batch(b) file_size = os.path.getsize(filepath) - # Determine field IDs for the new columns - next_fid = max(f.id() for f in self.lance_ds.lance_schema.fields()) + 1 + # Determine field IDs for the new columns. Lance's manifest-level + # max_field_id includes nested child fields and field IDs from dropped + # columns, so it is the correct high-water mark for dataset evolution. + next_fid = self.lance_ds.max_field_id + 1 # Stitch new data file into fragment metadata new_file_entry = { diff --git a/tests/io/lancedb/test_fast_path_merge.py b/tests/io/lancedb/test_fast_path_merge.py index 9f69efb..2947ad5 100644 --- a/tests/io/lancedb/test_fast_path_merge.py +++ b/tests/io/lancedb/test_fast_path_merge.py @@ -648,3 +648,36 @@ def _make_vec(ids): assert len(emb) == N for v in emb: assert pytest.approx(float(i), rel=1e-5) == v + + def test_next_fid_uses_manifest_max_field_id(self, ds_path): + """Bug: next_fid only scanned top-level lance_schema.fields(), missing child IDs. + + A struct column with M children occupies M+1 Lance field IDs (1 for parent + + M for children). With top-level-only scan: max_id = num_top_level_fields - 1, + next_fid can collide with a child field ID, and the new file's fragment metadata + maps to the wrong schema field → the new column reads back as null. + + Example layout for table with id(0) + meta struct(1, children a=2, b=3): + top-level scan → max=1, next_fid=2 (WRONG: collides with meta.a) + ds.max_field_id scan → max=3, next_fid=4 (CORRECT) + + The fix uses Lance's manifest-level max_field_id, which includes child IDs. + """ + struct_type = pa.struct([("a", pa.int64()), ("b", pa.utf8())]) + rows = [{"a": i, "b": f"s{i}"} for i in [1, 2, 3]] + tbl = pa.table( + { + "id": pa.array([1, 2, 3], type=pa.int64()), + "meta": pa.array(rows, type=struct_type), + } + ) + lance.write_dataset(tbl, ds_path) + ds = lance.dataset(ds_path) + + df = read_with_metadata(ds_path) + df = df.with_column("score", daft.col("id").cast(daft.DataType.int64()) * 10) + ds2 = merge_columns_from_df(df, ds, ds_path) + + result = ds2.to_table().sort_by("id").to_pydict() + assert result["score"] == [10, 20, 30], f"score is null or wrong (field ID collision): {result['score']}" + assert result["id"] == [1, 2, 3]