feat: allow blob in write_fragments

[fecet]

Allow user use write_fragments for lance with blob storage class by returning a nullable list of blob ops

[fecet]

@westonpace Could you review it please? I'm wondering if the change in python/src/fragment.rs is valid? As my blob column in the balanced dataset become all None after I readed them by _take_rows.

[fecet]

I think the problem is it's not possible to modify id of schema on python side

[westonpace]

Thanks for working on this! Overall I think this looks like it is going in the right direction.

I have a few cleanup suggestions. Also, this will need a unit test before we can merge. test_commit_batch_append in test_dataset.py should be pretty close.

I think the problem is it's not possible to modify id of schema on python side

Shouldn't the id of the field remain the same? An append should only add new fragments, it should not modify the schema in any way.

If you can make a unit test (even if its a failing one) then I can debug further. At the moment it isn't clear to me exactly how you plan to use this so it is a bit hard to debug.

[westonpace]

Is there any reason we can't modify write_fragments above instead of introducing a new method?

It's ok for there to be breaking changes at the rust level because the public API is the write_fragments in fragment.py.

[westonpace]

How are these changes related? Are you doing a overwrite operation?

[westonpace]

Hmm...being able to send a transaction in directly seems like it introduces more complexity to the user than needed. Is it possible to modify https://github.com/lancedb/lance/blob/main/python/python/lance/dataset.py#L2494 (the Append operation) so that it can take in a second list of blob fragments:

    @dataclass
    class Append(BaseOperation):
        fragments: Iterable[FragmentMetadata]
        blob_fragments: Iterable[FragmentMetadata] = []

[fecet]

I have found that the current implementation still has some issues, but I can first outline the difficulties I am currently encountering:

I am using the code below for testing, and it currently needs to run on the modified version.

import shutil


def make_table(offset, num_rows, big_val):
    end = offset + num_rows
    values = pa.array([big_val for _ in range(num_rows)], pa.large_binary())
    idx = pa.array(range(offset, end), pa.uint64())
    table = pa.record_batch(
        [idx, values],
        schema=pa.schema(
            [
                pa.field("idx", pa.uint64()),
                pa.field(
                    "blobs",
                    pa.large_binary(),
                    metadata={
                        "lance-schema:storage-class": "blob",
                    },
                ),
            ]
        ),
    )
    return table


tbl = make_table(0, 10, b"0" * 1024 * 1024)
uri = "test1"
shutil.rmtree(uri, ignore_errors=True)

default_frags, blob_frags = lance.fragment.write_fragments(
    pa.Table.from_batches([tbl]),
    uri,
    with_blobs=True,
    enable_move_stable_row_ids=True,
)

blob_operation = lance.LanceOperation.Overwrite(tbl.schema, blob_frags)
operation = lance.LanceOperation.Overwrite(tbl.schema, default_frags)
transaction = lance.Transaction(
    operation=operation,
    blobs_op=blob_operation,
    read_version=0,
)

ds = lance.LanceDataset.commit(
    uri,
    transaction,
    enable_v2_manifest_paths=True,
)
ds._take_rows(range(10))

The problem lies in the Rust implementation where the operation._to_inner method is called. Additionally, the two fragments returned by write_fragments retain the column index information. Therefore, the Overwrite operation expects to correctly preserve the index in the schema.

blob_operation = lance.LanceOperation.Overwrite(tbl.schema, blob_frags)
blob_operation._to_inner()

this will return correct schema

Overwrite { fragments: [Fragment { id: 0, files: [DataFile { path: "823dcbd1-175e-4307-b726-d9fbc775e887.lance", fields: [1], column_indices: [0], file_major_version: 2, file_minor_version: 0 }], deletion_file: None, row_id_meta: None, physical_rows: Some(10) }], schema: Schema { fields: [Field { name: "blobs", id: 1, parent_id: -1, logical_type: LogicalType("large_binary"), metadata: {"lance-schema:storage-class": "blob"}, encoding: Some(VarBinary), nullable: true, children: [], dictionary: None, storage_class: Blob }], metadata: {} }, config_upsert_values: None }

I confirm it's correct by

lance.write_dataset(tbl, "test2", tbl.schema)

lance.LanceDataset("test2/_blobs").lance_schema

The key of problem is the id in schema: Schema { fields: [Field { name: "blobs", id: 1, this cannot be done on the python side, as the id will be relabeled from zero to the length of schema,

pub fn set_field_id(&mut self, max_existing_id: Option<i32>) {
        let schema_max_id = self.max_field_id().unwrap_or(-1);
        let max_existing_id = max_existing_id.unwrap_or(-1);
        let mut current_id = schema_max_id.max(max_existing_id) + 1;
        self.fields
            .iter_mut()
            .for_each(|f| f.set_id(-1, &mut current_id));
    }

blob_operation = lance.LanceOperation.Overwrite(pa.schema(
            [
                pa.field(
                    "blobs",
                    pa.large_binary(),
                    metadata={
                        "lance-schema:storage-class": "blob",
                    },
                ),
            ]
        ), blob_frags)
blob_operation._to_inner()

will return schema: Schema { fields: [Field { name: "blobs", id: 0, that's why I do changes in https://github.com/lancedb/lance/pull/3235/files/56cda28424b3934f9ca32b1060e8f47d340639d2#diff-31492cb9a22ac40a820049c18a44ac88f651a5d0d7115fc09386899ab95f0c10. But I'm agreed it's not that correct, that could be easier if we can provide a way to directly return the transaction and we can extract inner schema from there.

[westonpace]

So you are using Overwrite to create a new dataset? I agree that should work. Would it help if write_fragments also returned the schema? I think this is how we work around this problem in Fragment.merge_columns.

[fecet]

That should be a better way. So do you think we should change the definition in LanceOperation so that they can accept LanceSchame like merge?

schema: LanceSchema | pa.Schema

Append operation should be just ok but I haven't do any test

[wjones127]

That should be a better way. So do you think we should change the definition in LanceOperation so that they can accept LanceSchame like merge?

That's actually being done right now in #3240

[fecet]

That's actually being done right now in #3240

Yes the translating between rust and python for transaction often trips me, will this be merged soon?

[fecet]

@westonpace I rewrite the implement and add test for it, I think just return a transaction would be cleaner for end user as it can be convert to python easily.