Add support for add_column with backfill

[jackye1995]

Data evolution like adding column with backfill is a critical feature in Lance, where users can append column data to an existing table very easily. Want to create this issue to track it, will post some thoughts about the design later.

[majin1102]

This seems dupliated with the issue : #15 ?

[jackye1995]

This is not nested column read or write. This is about the equivalent experience of https://lancedb.github.io/lance/integrations/ray/#parallel-column-merging

[majin1102]

This is not nested column read or write. This is about the equivalent experience of https://lancedb.github.io/lance/integrations/ray/#parallel-column-merging

Sorry I posted the wrong issue. It should be #15

I raised a issue to support addColumn and dropColumn to support sparksql scenarios(filling nulls or default values, this is the common usage for big data?). This has been implemented in our repo. I will upstream when I have time.

I have a suggestion. Can we describe this issue as merging column or appending column to distingish it with the default behavious of spark dataframe and SparkSQL(as well as other big data engines)?

[qidian99]

In our case, we need Spark-side feature engineering, and then only backfill new columns into Lance without rewriting blobs.

The process looks like:

1. Compute new column based on existing columns: e.g., col(c)=udf(a,b)

   In the simple case, we can do this using pyspark's dataframe api:

   df = df.withColumn("new_col", some_udf("col_a", "col_b")

   In other cases, to compute the new column(s) we may need to join other table or apply complex ETLs before materialization.

2. Backfill new column(s) without overwriting

   Now that we have computed new columns in Spark, we need to commit them into Lance. I can think of 2 possible safe commit strategies:
   A) Merge/join by key: Compute in Spark → materialize [key, new_col] → merge into Lance by key. Since the row alignment is enforced by join, no ordering constraints needed.
   B) Aligned stream: Compute in Spark in exact dataset order → stream only the new column(s) to Lance. This looks simpler but it requires strict global ordering, which is harder to guarantee in Spark at scale.

   In either 2 above approach, we need to:

   1. make sure new columns are aligned per fragment, which means padding nulls and apply the same DV if necessary.
   2. have a single committer that collects all metadata and do commit

This feature should be exposed in lance-spark via options like:
df.write.format("lance").option("backfill", "true").mode("append").save(uri)

or we may need some low level(fragment-level) semantics where we split tasks by fragments, add column to fragment, and commit new data files into fragments.

[fangbo]

@qidian99 I have a few idea to share for discussion. To align stream and write in parallel, we can extend spark by changing the physical plan.

1. Add _fragment_id and _row_addr in the spark table's schema.
2. Set _fragment_id as the shuffle distribution key. This ensures that records from the same Fragment are always written by the same Writer task.
3. Sort the records by _row_addr. This ensures that the new inserted columns can be aligned with exiting rows.

At commit stage in driver, we use Operation::Merge to add the new written columns to the Dataset.

[qidian99]

@qidian99 I have a few idea to share for discussion. To align stream and write in parallel, we can extend spark by changing the physical plan.

1. Add _fragment_id and _row_addr in the spark table's schema.
2. Set _fragment_id as the shuffle distribution key. This ensures that records from the same Fragment are always written by the same Writer task.
3. Sort the records by _row_addr. This ensures that the new inserted columns can be aligned with exiting rows.

At commit stage in driver, we use Operation::Merge to add the new written columns to the Dataset.

Yes there should be a new catalyst rule to do the rewrite, and the distribute by and ordering can be implemented in spark write as well.

1. Currently lance spark connector only exposes _rowid and _rowaddr during scan.
   1. _rowaddr can be added to schema.
   2. As for _fragment_id it's in the manifest, and can also be extract from _rowaddr -- which way is better?
2. When we process data from the same fragment, I'm not so sure about some details in data alignment:
   1. Let's say we have row addr=[0,5] to backfill to fragment 0, should we set the value for new column of row 1-4 to null?
   2. In the case when lance data type does not support null, should a exception be thrown?
   3. row 1-4 might have been deleted, or they simply do not match the backfill conditions -- do we need to distinguish between these cases?

[fangbo]

1. Currently lance spark connector only exposes _rowid and _rowaddr during scan.

   1. _rowaddr can be added to schema.
   2. As for _fragment_id it's in the manifest, and can also be extract from _rowaddr -- which way is better?

I think it is better to add a ConstantColumnVector containing _fragment_id like https://github.com/lancedb/lance-spark/pull/79/files#diff-f72ff3656c623491f0f40fb0394d05d03c2e5ffb37800bf407fceaa3b6f645b2R57

[fangbo]

This feature should be exposed in lance-spark via options like: df.write.format("lance").option("backfill", "true").mode("append").save(uri)

If write mode is use "append" , the exception will be thrown by spark. Such as

org.apache.spark.sql.AnalysisException: [INSERT_COLUMN_ARITY_MISMATCH.TOO_MANY_DATA_COLUMNS] Cannot write toadd_column_table, the reason is too many data columns: Table columns: id, text, _rowaddr, _fragid. Data columns: id, text, _rowaddr, _fragid, new_col`.

at org.apache.spark.sql.errors.QueryCompilationErrors$.cannotWriteTooManyColumnsToTableError(QueryCompilationErrors.scala:2114)
at org.apache.spark.sql.catalyst.analysis.TableOutputResolver$.resolveOutputColumns(TableOutputResolver.scala:52)
at org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveOutputRelation$$anonfun$apply$50.applyOrElse(Analyzer.scala:3424)
at org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveOutputRelation$$anonfun$apply$50.applyOrElse(Analyzer.scala:3420)
at org.apache.spark.sql.catalyst.plans.logical.AnalysisHelper.$anonfun$resolveOperatorsDownWithPruning$2(AnalysisHelper.scala:170)
at org.apache.spark.sql.catalyst.trees.CurrentOrigin$.withOrigin(origin.scala:76)

`

@qidian99 Do you have any suggestions about handling this mismatch in the number of fields ?

[qidian99]

This feature should be exposed in lance-spark via options like: df.write.format("lance").option("backfill", "true").mode("append").save(uri)

If write mode is use "append" , the exception will be thrown by spark. Such as

org.apache.spark.sql.AnalysisException: [INSERT_COLUMN_ARITY_MISMATCH.TOO_MANY_DATA_COLUMNS] Cannot write toadd_column_table, the reason is too many data columns: Table columns: id, text, _rowaddr, _fragid. Data columns: id, text, _rowaddr, _fragid, new_col`.

at org.apache.spark.sql.errors.QueryCompilationErrors$.cannotWriteTooManyColumnsToTableError(QueryCompilationErrors.scala:2114)
at org.apache.spark.sql.catalyst.analysis.TableOutputResolver$.resolveOutputColumns(TableOutputResolver.scala:52)
at org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveOutputRelation$$anonfun$apply$50.applyOrElse(Analyzer.scala:3424)
at org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveOutputRelation$$anonfun$apply$50.applyOrElse(Analyzer.scala:3420)
at org.apache.spark.sql.catalyst.plans.logical.AnalysisHelper.$anonfun$resolveOperatorsDownWithPruning$2(AnalysisHelper.scala:170)
at org.apache.spark.sql.catalyst.trees.CurrentOrigin$.withOrigin(origin.scala:76)

`

@qidian99 Do you have any suggestions about handling this mismatch in the number of fields ?

@fangbo I think there are 2 plausible approaches:

1. [implicit] schema evolution and column backfill in the same operation: need to bypass the arity check. The new analyzer rule should take effect before ResolveOutputRelation, where we match the AppendData plan, check that the relation is lance dataset, and align the columns before we proceed (need to get the backfill column name from hint or option)
2. [explicit] schema evolution and column backfill in different operations: in this case user should do alter table add column first, then execute column backfill

It depends on whether we want schema evolution to happen implicitly or explicitly.

[fangbo]

This feature should be exposed in lance-spark via options like: df.write.format("lance").option("backfill", "true").mode("append").save(uri)

If write mode is use "append" , the exception will be thrown by spark. Such as
org.apache.spark.sql.AnalysisException: [INSERT_COLUMN_ARITY_MISMATCH.TOO_MANY_DATA_COLUMNS] Cannot write toadd_column_table, the reason is too many data columns: Table columns: id, text, _rowaddr, _fragid. Data columns: id, text, _rowaddr, _fragid, new_col`.

at org.apache.spark.sql.errors.QueryCompilationErrors$.cannotWriteTooManyColumnsToTableError(QueryCompilationErrors.scala:2114)
at org.apache.spark.sql.catalyst.analysis.TableOutputResolver$.resolveOutputColumns(TableOutputResolver.scala:52)
at org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveOutputRelation$$anonfun$apply$50.applyOrElse(Analyzer.scala:3424)
at org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveOutputRelation$$anonfun$apply$50.applyOrElse(Analyzer.scala:3420)
at org.apache.spark.sql.catalyst.plans.logical.AnalysisHelper.$anonfun$resolveOperatorsDownWithPruning$2(AnalysisHelper.scala:170)
at org.apache.spark.sql.catalyst.trees.CurrentOrigin$.withOrigin(origin.scala:76)

`
@qidian99 Do you have any suggestions about handling this mismatch in the number of fields ?

@fangbo I think there are 2 plausible approaches:

1. [implicit] schema evolution and column backfill in the same operation: need to bypass the arity check. The new analyzer rule should take effect before ResolveOutputRelation, where we match the AppendData plan, check that the relation is lance dataset, and align the columns before we proceed (need to get the backfill column name from hint or option)
2. [explicit] schema evolution and column backfill in different operations: in this case user should do alter table add column first, then execute column backfill

It depends on whether we want schema evolution to happen implicitly or explicitly.

@qidian99 Thank you for your suggestion! I think that the first approach ([implicit] schema evolution and column backfill in the same operation) fits our usage habits better.

In addition, I tried to add a new Rule through an Extension, but at runtime ResolveOutputRelation still executes before this added customer Rule. Do you have any advice on that ?

[qidian99]

This feature should be exposed in lance-spark via options like: df.write.format("lance").option("backfill", "true").mode("append").save(uri)

If write mode is use "append" , the exception will be thrown by spark. Such as
org.apache.spark.sql.AnalysisException: [INSERT_COLUMN_ARITY_MISMATCH.TOO_MANY_DATA_COLUMNS] Cannot write toadd_column_table, the reason is too many data columns: Table columns: id, text, _rowaddr, _fragid. Data columns: id, text, _rowaddr, _fragid, new_col`.

at org.apache.spark.sql.errors.QueryCompilationErrors$.cannotWriteTooManyColumnsToTableError(QueryCompilationErrors.scala:2114)
at org.apache.spark.sql.catalyst.analysis.TableOutputResolver$.resolveOutputColumns(TableOutputResolver.scala:52)
at org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveOutputRelation$$anonfun$apply$50.applyOrElse(Analyzer.scala:3424)
at org.apache.spark.sql.catalyst.analysis.Analyzer$ResolveOutputRelation$$anonfun$apply$50.applyOrElse(Analyzer.scala:3420)
at org.apache.spark.sql.catalyst.plans.logical.AnalysisHelper.$anonfun$resolveOperatorsDownWithPruning$2(AnalysisHelper.scala:170)
at org.apache.spark.sql.catalyst.trees.CurrentOrigin$.withOrigin(origin.scala:76)

`
@qidian99 Do you have any suggestions about handling this mismatch in the number of fields ?

@fangbo I think there are 2 plausible approaches:

1. [implicit] schema evolution and column backfill in the same operation: need to bypass the arity check. The new analyzer rule should take effect before ResolveOutputRelation, where we match the AppendData plan, check that the relation is lance dataset, and align the columns before we proceed (need to get the backfill column name from hint or option)
2. [explicit] schema evolution and column backfill in different operations: in this case user should do alter table add column first, then execute column backfill

It depends on whether we want schema evolution to happen implicitly or explicitly.

@qidian99 Thank you for your suggestion! I think that the first approach ([implicit] schema evolution and column backfill in the same operation) fits our usage habits better.

In addition, I tried to add a new Rule through an Extension, but at runtime ResolveOutputRelation still executes before this added customer Rule. Do you have any advice on that ?

@fangbo I did some research, user rules are always appended to the end of resolution rules, and there's no way to disable/change the order of analyzer rules.

[fangbo]

@qidian99 I think we can't workaround Spark's limitation if we use "append" as write mode. I have another idea about using "overwrite". We can inject a Strategy which converts ReplaceTableAsSelect to a customized operator if write option has "backfill". In this operator, the schema can be aligned depending on the output's schema. This is a demo:

class AddColumnsAsSelectExec(
                                   override val catalog: TableCatalog,
                                   override val ident: Identifier,
                                   override val partitioning: Seq[Transform],
                                   override val query: LogicalPlan,
                                   override val tableSpec: TableSpec,
                                   override val writeOptions: Map[String, String],
                                   override val orCreate: Boolean,
                                   override val invalidateCache: (TableCatalog, Table, Identifier) => Unit)
  extends ReplaceTableAsSelectExec(catalog, ident, partitioning, query, tableSpec, writeOptions, orCreate, invalidateCache) {

  override protected def run(): Seq[InternalRow] = {
    if (!catalog.tableExists(ident)) {
      return super.run()
    }

    val table = catalog.loadTable(ident)
    invalidateCache(catalog, table, ident)

    table match {
      case lanceTable: LanceDataset => query.schema.foreach(field => lanceTable.addField(field))
      case _ => throw new UnsupportedOperationException("AddColumnsAsSelectExec only support LanceDataset")
    }

    writeToTable(catalog, table, writeOptions, ident, query)
  }
}

I am implementing a POC about this. I will submit a POC pull later and we can discuss on the code.

[jackye1995]

@fangbo if you are already doing a custom strategy, why not just introduce a completely new custom SQL parser and extension? (basically like what all the formats have done before Spark officially supported UPDATE/DELETE/MERGE)