Repository: risingwavelabs/risingwave
File: src/frontend/src/datafusion/expand.rs (Rust)

Code before:
    // `keep_columns` semantics in RW expand executor: for each subset, keep values on listed
    // positions and fill other expanded positions with NULLs.
    let input_len = (schema.fields().len() - 1) / 2;
    let subset_masks: Vec<Vec<bool>> = column_subsets
        .iter()
        .map(|subset| {
            let mut mask = vec![false; input_len];
            for &idx in subset {
                mask[idx] = true;
            }
            mask
        })
        .collect();

    #[for_await]
    for input_batch in input {
        let input_batch = input_batch?;
        let input_rows = input_batch.num_rows();
        let input_cols = input_batch.columns();

        for (flag, subset_mask) in subset_masks.iter().enumerate() {
            let flag = i64::try_from(flag)
                .map_err(|_| exec_datafusion_err!("ExpandExec: flag index overflows i64"))?;
            let flags: ArrayRef =
                Arc::new(Int64Array::from_iter(std::iter::repeat_n(flag, input_rows)));

            let mut output_cols = Vec::with_capacity(input_len * 2 + 1);
            for (idx, keep) in subset_mask.iter().enumerate() {
                if *keep {
                    output_cols.push(input_cols[idx].clone());
                } else {
                    output_cols.push(new_null_array(input_cols[idx].data_type(), input_rows));
                }
            }
            output_cols.extend(input_cols.iter().cloned());
            output_cols.push(flags);

            yield RecordBatch::try_new(schema.clone(), output_cols)?;
        }
    }
}

#[derive(Debug)]
pub struct ExpandPlanner;

#[async_trait]
impl ExtensionPlanner for ExpandPlanner {
    async fn plan_extension(
        &self,
        _planner: &dyn PhysicalPlanner,
        node: &dyn datafusion::logical_expr::UserDefinedLogicalNode,

Review comment (bug):
When constructing output columns, the code assumes `input_cols.len() == input_len` and indexes `input_cols[idx]` directly. Consider adding an explicit check (and returning `exec_err!`) so invariant violations fail gracefully instead of panicking, which is important for reliability/debuggability if the plan/schema ever becomes inconsistent.

Code after:
        .iter()
        .map(|subset| Bitmap::from_indices(input_len, subset))
        .collect();

    #[for_await]
    for input_batch in input {
        let input_batch = input_batch?;
        let input_rows = input_batch.num_rows();
        let input_cols = input_batch.columns();
        if input_cols.len() != input_len {
            return exec_err!("ExpandExec: unexpected number of input columns");
        }

        for (flag, subset_mask) in subset_masks.iter().enumerate() {
            let flag = i64::try_from(flag)
                .map_err(|_| exec_datafusion_err!("ExpandExec: flag index overflows i64"))?;
            let flags: ArrayRef =
                Arc::new(Int64Array::from_iter(std::iter::repeat_n(flag, input_rows)));

            let mut output_cols = Vec::with_capacity(input_len * 2 + 1);
            for (input_col, keep) in input_cols.iter().zip_eq_fast(subset_mask.iter()) {
                if keep {
                    output_cols.push(input_col.clone());
                } else {
                    output_cols.push(new_null_array(input_col.data_type(), input_rows));
                }
            }
            output_cols.extend(input_cols.iter().cloned());
            output_cols.push(flags);

            yield RecordBatch::try_new(schema.clone(), output_cols)?;
        }
    }
}

#[derive(Debug)]
pub struct ExpandPlanner;

#[async_trait]
impl ExtensionPlanner for ExpandPlanner {
    async fn plan_extension(
        &self,
        _planner: &dyn PhysicalPlanner,
        node: &dyn datafusion::logical_expr::UserDefinedLogicalNode,
        _logical_inputs: &[&LogicalPlan],
        physical_inputs: &[Arc<dyn ExecutionPlan>],
        _session_state: &SessionState,
    ) -> DFResult<Option<Arc<dyn ExecutionPlan>>> {
        let Some(expand) = node.as_any().downcast_ref::<Expand>() else {
            return Ok(None);
        };