Function execute_update_once

crates/integrations/datafusion/src/update.rs:88–160

Single attempt of UPDATE execution.

(
    ctx: &SQLContext,
    update: &Update,
    table: &Table,
)

/// Single attempt of UPDATE execution.
async fn execute_update_once(
    ctx: &SQLContext,
    update: &Update,
    table: &Table,
) -> DFResult<DataFrame> {
    // 1. Extract SET assignments
    let mut columns = Vec::new();
    let mut exprs = Vec::new();
    for assignment in &update.assignments {
        let col_name = match &assignment.target {
            AssignmentTarget::ColumnName(name) => name
                .0
                .last()
                .and_then(|p| p.as_ident())
                .map(|id| id.value.clone())
                .ok_or_else(|| {
                    DataFusionError::Plan(format!("Invalid column name in SET: {name}"))
                })?,
            AssignmentTarget::Tuple(_) => {
                return Err(DataFusionError::Plan(
                    "Tuple assignment in UPDATE SET is not supported".to_string(),
                ));
            }
        };
        columns.push(col_name);
        exprs.push(assignment.value.to_string());
    }

    // 2. Create DataEvolutionWriter (validates preconditions)
    let mut writer =
        DataEvolutionWriter::new(table, columns.clone()).map_err(to_datafusion_error)?;

    // 3. Query the target table directly with WHERE filter.
    let table_ref = update.table.to_string();

    let select_parts: Vec<String> =
        std::iter::once("\"_ROW_ID\"".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(),
    };

    let query_sql = format!("SELECT {select_clause} FROM {table_ref}{where_clause}");
    let batches = ctx.ctx().sql(&query_sql).await?.collect().await?;

    // 4. Project update columns (rename __upd_X → X)
    let total_count: u64 = batches.iter().map(|b| b.num_rows() as u64).sum();
    if total_count == 0 {
        return ok_result(ctx.ctx(), 0);
    }

    let update_batches = project_update_columns(&batches, &columns)?;