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

Code before:
        }
    }

    #[try_stream(ok = RecordBatch, error = DataFusionError)]
    pub async fn execute_inner(self: Arc<Self>, chunk_size: usize, partition: usize) {
        let mut buffer = RecordBatchBuffer::new(chunk_size);
        let reader = self
            .table
            .reader_builder()
            .with_batch_size(chunk_size)
            .build();
        let task = self.tasks[partition].clone();
        let stream = reader
            .read(tokio_stream::once(Ok(task)).boxed())
            .await
            .map_err(to_datafusion_error)?;

        #[for_await]
        for (i, batch) in stream.enumerate() {
            let batch = batch.map_err(to_datafusion_error)?;
            let batch = append_metadata(
                batch,
                self.need_seq_num,
                self.need_file_path_and_pos,
                &self.tasks[partition],
                (i * chunk_size).try_into().unwrap(),
            )?;
            let batch = cast_batch(self.arrow_schema.clone(), batch)?;
            if let Some(batch) = buffer.add(batch)? {
                yield batch;
            }
        }

        if let Some(batch) = buffer.finish()? {
            yield batch;
        }
    }
}

struct RecordBatchBuffer {
    buffer: Vec<RecordBatch>,
    current_rows: usize,
    max_record_batch_rows: usize,
}

impl RecordBatchBuffer {
    fn new(max_record_batch_rows: usize) -> Self {
        Self {
            buffer: vec![],
            current_rows: 0,
            max_record_batch_rows,

Review comment (bug):
The position calculation assumes all batches have exactly `chunk_size` rows, but the last batch in a file could have fewer rows. This means `pos_start` could be incorrect for the last batch. Consider calculating `pos_start` based on the cumulative row count from previous batches rather than `i * chunk_size`. For example:
```rust
let mut cumulative_rows = 0i64;
for batch in stream {
    // ... 
    let batch = append_metadata(batch, ..., cumulative_rows)?;
    cumulative_rows += batch.num_rows() as i64;
}
```

Code after:
            }
        }
    }

    #[try_stream(ok = RecordBatch, error = DataFusionError)]
    pub async fn execute_inner(self: Arc<Self>, chunk_size: usize, partition: usize) {
        let reader = self
            .table
            .reader_builder()
            .with_batch_size(chunk_size)
            .build();

        for task in &self.tasks[partition] {
            let stream = reader
                .clone()
                .read(tokio_stream::once(Ok(task.clone())).boxed())
                .await
                .map_err(to_datafusion_error)?;
            let mut pos_start: i64 = 0;

            #[for_await]
            for batch in stream {
                let batch = batch.map_err(to_datafusion_error)?;
                let batch = append_metadata(
                    batch,
                    self.need_seq_num,
                    self.need_file_path_and_pos,
                    task,
                    pos_start,
                )?;
                let batch = cast_batch(self.arrow_schema.clone(), batch)?;
                pos_start += i64::try_from(batch.num_rows()).unwrap();
                yield batch;
            }
        }
    }
}

fn append_metadata(
    batch: RecordBatch,
    need_seq_num: bool,
    need_file_path_and_pos: bool,
    task: &FileScanTask,
    pos_start: i64,
) -> Result<RecordBatch, DataFusionError> {
    let mut columns = batch.columns().to_vec();
    let mut fields = batch.schema().fields().to_vec();

    if need_seq_num {
        let seq_num_array = Int64Array::from_value(task.sequence_number, batch.num_rows());
        columns.push(Arc::new(seq_num_array));