add_columns dataset API

[wjones127]

We have an add_columns() API for fragments, and a merge() API for dataset. It would be nice to create a add_columns() API for dataset that is similar to the one for fragments.

It would also be very nice to add progress tracking as an option there, since this might be a long operation. For example, users might call this API to add a new embedding column to the dataset.

[wjones127]

We would also like to support a compute expression, so users can easily do a cast, for example.

The API could look like:

def add_columns(exprs_or_udf: Dict[str, str] | AddColumnUDF):
    ...

Then users could write expressions:

dataset.add_columns({"bloat16_vec": "cast(vector as array<bfloat16, 512>)"})

Or could specify a UDF using some decorator:

pool = ThreadPoolExecutor(8)

@add_columns_udf(
   reads_columns=['vector'],
   pool=pool,
)
def cast_bfloat16(batch: pa.RecordBatch):
    ...


dataset.add_columns(cast_bfloat16)

The pool parameter is important with UDFs to allow some parallism.

[wjones127]

Another requirement: for the UDF, it would be nice to provide some way for the data to be staged so that it can be resumed in case of a crash.

One way to do that in Python is use a sqlite file as a durable mapping from hash(input_data) to output_data. We could even be extra smart and keep track of which fragments have already been written and where. This could also help with cleanup.