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

Code before:
use crate::error::{Result as RwResult, RwError};
use crate::handler::RwPgResponse;
use crate::handler::util::{DataChunkToRowSetAdapter, to_pg_field};
use crate::optimizer::BatchOptimizedLogicalPlanRoot;
use crate::scheduler::SchedulerError;
use crate::session::SessionImpl;
use crate::utils::DropGuard;

mod memory_ctx;
mod query_planner;

#[derive(Clone)]
pub struct DfBatchQueryPlanResult {
    pub(crate) plan: Arc<LogicalPlan>,
    pub(crate) schema: RwSchema,
    pub(crate) stmt_type: StatementType,
}

#[derive(Debug, Error)]
pub enum GenDataFusionPlanError {
    #[error("Missing Iceberg Scan")]
    MissingIcebergScan,
    #[error("Unsupported Plan Node")]
    UnsupportedPlanNode,
    #[error("Generating plan error: {0}")]
    Generating(RwError),
}

pub fn try_gen_datafusion_plan(
    optimized_logical: &BatchOptimizedLogicalPlanRoot,
) -> Result<Arc<LogicalPlan>, GenDataFusionPlanError> {
    use crate::optimizer::DataFusionExecuteCheckerExt;

    let check_result = optimized_logical.plan.check_for_datafusion();
    if !check_result.have_iceberg_scan {
        return Err(GenDataFusionPlanError::MissingIcebergScan);
    }
    if !check_result.supported {
        return Err(GenDataFusionPlanError::UnsupportedPlanNode);
    }

    let plan = optimized_logical
        .gen_datafusion_logical_plan()
        .map_err(GenDataFusionPlanError::Generating)?;
    Ok(plan)
}

pub fn create_datafusion_context(session: &SessionImpl) -> RwResult<DFSessionContext> {
    let df_config = create_config(session);
    let memory_pool = Arc::new(RwMemoryPool::new(session.env().mem_context()));
    let runtime = RuntimeEnvBuilder::new()

Review comment (suggestion):
`GenDataFusionPlanError::Generating(RwError)` is formatted into the message but is not marked as an error source. Since the caller logs with `err.as_report()`, the underlying `RwError` chain may be lost. Consider marking the wrapped error as `#[source]` (or making the variant `#[error(transparent)]` + `#[from]`) so report formatting includes the full causal chain.
```suggestion
    Generating(#[source] RwError),
```

Code after:
use crate::error::{Result as RwResult, RwError};
use crate::handler::RwPgResponse;
use crate::handler::util::{DataChunkToRowSetAdapter, to_pg_field};
use crate::optimizer::BatchOptimizedLogicalPlanRoot;
use crate::scheduler::SchedulerError;
use crate::session::SessionImpl;
use crate::utils::DropGuard;

mod memory_ctx;
mod query_planner;

#[derive(Clone)]
pub struct DfBatchQueryPlanResult {
    pub(crate) plan: Arc<LogicalPlan>,
    pub(crate) schema: RwSchema,
    pub(crate) stmt_type: StatementType,
}

#[derive(Debug, Error)]
pub enum GenDataFusionPlanError {
    #[error("Missing Iceberg Scan")]
    MissingIcebergScan,
    #[error("Unsupported Plan Node")]
    UnsupportedPlanNode,
    #[error("Generating plan error: {0}")]
    Generating(#[source] RwError),
}

pub fn try_gen_datafusion_plan(
    optimized_logical: &BatchOptimizedLogicalPlanRoot,
) -> Result<Arc<LogicalPlan>, GenDataFusionPlanError> {
    use crate::optimizer::DataFusionExecuteCheckerExt;

    let check_result = optimized_logical.plan.check_for_datafusion();
    if !check_result.have_iceberg_scan {
        return Err(GenDataFusionPlanError::MissingIcebergScan);
    }
    if !check_result.supported {
        return Err(GenDataFusionPlanError::UnsupportedPlanNode);
    }

    let plan = optimized_logical
        .gen_datafusion_logical_plan()
        .map_err(GenDataFusionPlanError::Generating)?;
    Ok(plan)
}

pub fn create_datafusion_context(session: &SessionImpl) -> RwResult<DFSessionContext> {
    let df_config = create_config(session);
    let memory_pool = Arc::new(RwMemoryPool::new(session.env().mem_context()));
    let runtime = RuntimeEnvBuilder::new()