native-engine/datafusion-ext-plans/src/union_exec.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use std::{any::Any, fmt::Formatter, sync::Arc};

use arrow::array::{RecordBatch, RecordBatchOptions};
use arrow_schema::SchemaRef;
use datafusion::{
    common::{Result, Statistics},
    execution::{SendableRecordBatchStream, TaskContext},
    physical_expr::{EquivalenceProperties, Partitioning::UnknownPartitioning},
    physical_plan::{
        DisplayAs, DisplayFormatType, EmptyRecordBatchStream, ExecutionPlan, PlanProperties,
        execution_plan::{Boundedness, EmissionType},
        metrics::{ExecutionPlanMetricsSet, MetricsSet},
    },
};
use datafusion_ext_commons::arrow::cast::cast;
use futures_util::StreamExt;
use once_cell::sync::OnceCell;

use crate::common::execution_context::ExecutionContext;

#[derive(Debug)]
pub struct UnionExec {
    inputs: Vec<UnionInput>,
    schema: SchemaRef,
    num_partitions: usize,
    cur_partition: usize,
    metrics: ExecutionPlanMetricsSet,
    props: OnceCell<Arc<PlanProperties>>,
}

#[derive(Debug, Clone)]
pub struct UnionInput(pub Arc<dyn ExecutionPlan>, pub usize);

impl UnionExec {
    pub fn new(
        inputs: Vec<UnionInput>,
        schema: SchemaRef,
        num_partitions: usize,
        cur_partition: usize,
    ) -> Self {
        UnionExec {
            inputs,
            schema,
            num_partitions,
            cur_partition,
            metrics: ExecutionPlanMetricsSet::new(),
            props: OnceCell::new(),
        }
    }
}

impl DisplayAs for UnionExec {
    fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result {
        write!(f, "UnionExec")
    }
}

impl ExecutionPlan for UnionExec {
    fn name(&self) -> &str {
        "UnionExec"
    }

    fn schema(&self) -> SchemaRef {
        self.schema.clone()
    }

    fn properties(&self) -> &Arc<PlanProperties> {
        self.props.get_or_init(|| {
            Arc::new(PlanProperties::new(
                EquivalenceProperties::new(self.schema()),
                UnknownPartitioning(self.num_partitions),
                EmissionType::Both,
                Boundedness::Bounded,
            ))
        })
    }

    fn apply_expressions(
        &self,
        _f: &mut dyn FnMut(
            &Arc<dyn datafusion::physical_expr::PhysicalExpr>,
        ) -> Result<datafusion::common::tree_node::TreeNodeRecursion>,
    ) -> Result<datafusion::common::tree_node::TreeNodeRecursion> {
        Ok(datafusion::common::tree_node::TreeNodeRecursion::Continue)
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        self.inputs.iter().map(|input| &input.0).collect()
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        if children.len() != self.inputs.len() {
            return datafusion_ext_commons::df_execution_err!(
                "UnionExec expected {} children, got {}",
                self.inputs.len(),
                children.len()
            );
        }
        Ok(Arc::new(Self::new(
            children
                .into_iter()
                .zip(&self.inputs)
                .map(|(child, input)| UnionInput(child, input.1))
                .collect(),
            self.schema.clone(),
            self.num_partitions,
            self.cur_partition,
        )))
    }

    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        if partition != self.cur_partition {
            return Ok(Box::pin(EmptyRecordBatchStream::new(self.schema())));
        }
        let exec_ctx = ExecutionContext::new(context, partition, self.schema(), &self.metrics);
        let input_streams = self
            .inputs
            .iter()
            .map(|UnionInput(input, partition)| {
                let in_exec_ctx = ExecutionContext::new(
                    exec_ctx.task_ctx().clone(),
                    *partition,
                    input.schema(),
                    &self.metrics,
                );
                in_exec_ctx.execute_with_input_stats(input)
            })
            .collect::<Result<Vec<_>>>()?;

        Ok(exec_ctx
            .clone()
            .output_with_sender("Union", move |sender| async move {
                let output_schema = exec_ctx.output_schema();
                let output_fields = output_schema.fields();

                // iterate all input streams
                for mut input in input_streams {
                    // cast data type if input schema not matching output schema
                    let input_schema = input.schema();
                    let input_fields = input_schema.fields();
                    let columns_need_cast = input_fields
                        .iter()
                        .zip(output_fields)
                        .map(|(input_field, output_field)| input_field != output_field)
                        .collect::<Vec<_>>();
                    let need_cast = columns_need_cast.contains(&true);

                    while let Some(mut batch) = input.next().await.transpose()? {
                        if need_cast {
                            let mut casted_cols = Vec::with_capacity(batch.num_columns());
                            for ((mut col, &col_need_cast), output_field) in batch
                                .columns()
                                .iter()
                                .cloned()
                                .zip(&columns_need_cast)
                                .zip(output_schema.fields())
                            {
                                if col_need_cast {
                                    col = cast(&col, output_field.data_type())?;
                                }
                                casted_cols.push(col);
                            }
                            batch = RecordBatch::try_new_with_options(
                                output_schema.clone(),
                                casted_cols,
                                &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())),
                            )?;
                        }
                        exec_ctx.baseline_metrics().record_output(batch.num_rows());
                        sender.send(batch).await;
                    }
                }
                Ok(())
            }))
    }

    fn metrics(&self) -> Option<MetricsSet> {
        Some(self.metrics.clone_inner())
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use arrow::{
        array::Int32Array,
        datatypes::{DataType, Field, Schema},
        record_batch::RecordBatch,
    };
    use datafusion::{
        assert_batches_eq,
        common::Result,
        physical_plan::{ExecutionPlan, ExecutionPlanProperties, common, test::TestMemoryExec},
        prelude::SessionContext,
    };

    use super::{UnionExec, UnionInput};

    fn input(values: &[i32]) -> Result<Arc<dyn ExecutionPlan>> {
        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Int32,
            false,
        )]));
        let partitions = values
            .iter()
            .map(|value| {
                Ok(vec![RecordBatch::try_new(
                    schema.clone(),
                    vec![Arc::new(Int32Array::from(vec![*value]))],
                )?])
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(Arc::new(TestMemoryExec::try_new(
            &partitions,
            schema,
            None,
        )?))
    }

    #[tokio::test]
    async fn replacing_children_preserves_union_partition_mapping() -> Result<()> {
        let left = input(&[10, 11])?;
        let right = input(&[20, 21, 22])?;
        let union: Arc<dyn ExecutionPlan> = Arc::new(UnionExec::new(
            vec![UnionInput(left.clone(), 1), UnionInput(right, 2)],
            left.schema(),
            8,
            7,
        ));

        let rebuilt = union.with_new_children(vec![input(&[30, 31])?, input(&[40, 41, 42])?])?;
        assert_eq!(rebuilt.output_partitioning().partition_count(), 8);
        let ctx = SessionContext::new();
        let batches = common::collect(rebuilt.execute(7, ctx.task_ctx())?).await?;
        assert_batches_eq!(
            [
                "+-------+",
                "| value |",
                "+-------+",
                "| 31    |",
                "| 42    |",
                "+-------+"
            ],
            &batches
        );
        assert!(
            common::collect(rebuilt.execute(6, ctx.task_ctx())?)
                .await?
                .is_empty()
        );
        Ok(())
    }
    #[test]
    fn replacing_children_rejects_changed_input_count() -> Result<()> {
        let input = input(&[10, 11])?;
        let union: Arc<dyn ExecutionPlan> = Arc::new(UnionExec::new(
            vec![UnionInput(input.clone(), 1)],
            input.schema(),
            2,
            1,
        ));
        assert!(union.clone().with_new_children(vec![]).is_err());
        assert!(union.with_new_children(vec![input.clone(), input]).is_err());
        Ok(())
    }
}