Function execute_cow_update_inner

crates/integrations/datafusion/src/update.rs:223–294

(
    ctx: &SessionContext,
    columns: &[String],
    exprs: &[String],
    cow_table_name: &str,
    update: &Update,
    writer: &mut CopyOnWriteMergeWriter,
)

}

async fn execute_cow_update_inner(
    ctx: &SessionContext,
    columns: &[String],
    exprs: &[String],
    cow_table_name: &str,
    update: &Update,
    writer: &mut CopyOnWriteMergeWriter,
) -> DFResult<u64> {
    let select_parts: Vec<String> =
        std::iter::once("\"__paimon_file_idx\"".to_string())
            .chain(std::iter::once("\"__paimon_row_offset\"".to_string()))
            .chain(columns.iter().zip(exprs.iter()).map(|(col, expr)| {
                format!("{expr} AS {}", quote_identifier(&format!("__upd_{col}")))
            }))
            .collect();

    let select_clause = select_parts.join(", ");
    let where_clause = match &update.selection {
        Some(expr) => format!(" WHERE {expr}"),
        None => String::new(),
    };

    // Safety: where_clause comes from sqlparser AST to_string(), not raw user input.
    let query_sql = format!("SELECT {select_clause} FROM {cow_table_name}{where_clause}");
    let join_result = ctx.sql(&query_sql).await?.collect().await?;

    let mut update_value_batches: Vec<RecordBatch> = Vec::new();
    let mut batch_counter: usize = 0;
    let mut total_count: u64 = 0;

    for batch in &join_result {
        if batch.num_rows() == 0 {
            continue;
        }

        let (file_idx_col, row_offset_col) = extract_tracking_columns(batch)?;

        let mut upd_fields = Vec::new();
        let mut upd_columns: Vec<Arc<dyn Array>> = Vec::new();
        for col in columns {
            let prefixed = format!("__upd_{col}");
            let idx = batch.schema().index_of(&prefixed).map_err(|e| {
                DataFusionError::Internal(format!("Column {prefixed} not found: {e}"))
            })?;
            upd_fields.push(Field::new(
                col,
                batch.schema().field(idx).data_type().clone(),
                true,
            ));
            upd_columns.push(batch.column(idx).clone());
        }
        let upd_schema = Arc::new(Schema::new(upd_fields));
        let upd_batch = RecordBatch::try_new(upd_schema, upd_columns)?;

        let current_batch_idx = batch_counter;
        update_value_batches.push(upd_batch);
        batch_counter += 1;