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

Code before:
            self.schema.clone(),
            stream,
        )))
    }
}

// Imitate risingwave_[batch|stream] ExpandExecutor: one input scan, expand rows by subsets.
#[try_stream(ok = RecordBatch, error = datafusion_common::DataFusionError)]
async fn expand_stream(
    input: SendableRecordBatchStream,
    schema: ArrowSchemaRef,
    column_subsets: Vec<Vec<usize>>,
) {
    // `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)?;

Review comment (performance):
`input_len` is derived from the *output* schema via `(schema.fields().len() - 1) / 2`, and subset masks are built using it. If the physical input schema ever diverges (e.g., due to optimizer rewrites or future schema layout changes), this can lead to out-of-bounds panics when indexing `mask[idx]` / `input_cols[idx]`. Consider deriving `input_len` from `input_batch.num_columns()` (or `input.schema()`) and validating `schema.fields().len() == input_len * 2 + 1` up front, returning `exec_err!` on mismatch to avoid panics.
```suggestion

    #[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();
        let input_len = input_batch.num_columns();

        // Validate that the provided output schema is consistent with the physical input.
        let expected_fields = input_len * 2 + 1;
        let actual_fields = schema.fields().len();
        if actual_fields != expected_fields {
            Err(exec_datafusion_err!(
                "ExpandExec: schema mismatch, expected {} fields for input with {} columns, but got {}",
                expected_fields,
                input_len,
                actual_fields
            ))?;
        }

        // Build subset masks based on the physical input width, validating indices.
        let subset_masks: Vec<Vec<bool>> = column_subsets
            .iter()
            .map(|subset| {
                let mut mask = vec![false; input_len];
                for &idx in subset {
                    if idx >= input_len {
                        return Err(exec_datafusion_err!(
                            "ExpandExec: column subset index {} out of bounds for input with {} columns",
                            idx,
                            input_len
                        ));
                    }
                    mask[idx] = true;
                }
                Ok(mask)
            })
            .collect::<DFResult<_>>()?;
```

Code after:
        )))
    }
}

// imitate risingwave_batch_executors::executor::ExpandExecutor::do_execute
#[try_stream(ok = RecordBatch, error = datafusion_common::DataFusionError)]
async fn expand_stream(
    input: SendableRecordBatchStream,
    schema: ArrowSchemaRef,
    column_subsets: Vec<Vec<usize>>,
) {
    let input_len = (schema.fields().len() - 1) / 2;
    let subset_masks: Vec<Bitmap> = column_subsets
        .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;