datafusion_datasource/

sink.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.

//! Execution plan for writing data to [`DataSink`]s

use std::any::Any;
use std::fmt;
use std::fmt::Debug;
use std::sync::Arc;

use arrow::array::{ArrayRef, RecordBatch, UInt64Array};
use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use datafusion_common::tree_node::TreeNodeRecursion;
use datafusion_common::{Result, assert_eq_or_internal_err};
use datafusion_execution::TaskContext;
use datafusion_physical_expr::{Distribution, EquivalenceProperties, PhysicalExpr};
use datafusion_physical_expr_common::sort_expr::{LexRequirement, OrderingRequirements};
use datafusion_physical_plan::metrics::MetricsSet;
use datafusion_physical_plan::stream::RecordBatchStreamAdapter;
use datafusion_physical_plan::{
    ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan,
    ExecutionPlanProperties, InputDistributionRequirements, Partitioning, PlanProperties,
    ReplaceChildrenOptions, SendableRecordBatchStream, execute_input_stream,
};

use async_trait::async_trait;
use datafusion_physical_plan::execution_plan::{EvaluationType, SchedulingType};
use futures::StreamExt;

/// `DataSink` implements writing streams of [`RecordBatch`]es to
/// user defined destinations.
///
/// The `Display` impl is used to format the sink for explain plan
/// output.
#[async_trait]
pub trait DataSink: Any + DisplayAs + Debug + Send + Sync {
    /// Return a snapshot of the [MetricsSet] for this
    /// [DataSink].
    ///
    /// See [ExecutionPlan::metrics()] for more details
    fn metrics(&self) -> Option<MetricsSet> {
        None
    }

    /// Returns the sink schema
    fn schema(&self) -> &SchemaRef;

    // TODO add desired input ordering
    // How does this sink want its input ordered?

    /// Writes the data to the sink, returns the number of values written
    ///
    /// This method will be called exactly once during each DML
    /// statement. Thus prior to return, the sink should do any commit
    /// or rollback required.
    async fn write_all(
        &self,
        data: SendableRecordBatchStream,
        context: &Arc<TaskContext>,
    ) -> Result<u64>;

    /// Serialize this sink into a full protobuf plan node, if it knows how.
    ///
    /// Implementations can use `ctx` to encode the input plan, sink-specific
    /// expressions, and [`DataSinkExec::encode_sort_order`].
    ///
    /// Returning `Ok(None)` lets the caller try its extension codec instead.
    #[cfg(feature = "proto")]
    fn try_to_proto(
        &self,
        _exec: &DataSinkExec,
        _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
        Ok(None)
    }
}

impl dyn DataSink {
    /// Returns true if the inner type is `T`.
    pub fn is<T: DataSink>(&self) -> bool {
        (self as &dyn Any).is::<T>()
    }

    /// Returns a reference to the inner value as the type `T` if it is of that type.
    pub fn downcast_ref<T: DataSink>(&self) -> Option<&T> {
        (self as &dyn Any).downcast_ref()
    }
}

/// Execution plan for writing record batches to a [`DataSink`]
///
/// Returns a single row with the number of values written
#[derive(Clone)]
pub struct DataSinkExec {
    /// Input plan that produces the record batches to be written.
    input: Arc<dyn ExecutionPlan>,
    /// Sink to which to write
    sink: Arc<dyn DataSink>,
    /// Schema describing the structure of the output data.
    count_schema: SchemaRef,
    /// Optional required sort order for output data.
    sort_order: Option<LexRequirement>,
    cache: Arc<PlanProperties>,
}

impl Debug for DataSinkExec {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "DataSinkExec schema: {}", self.count_schema)
    }
}

impl DataSinkExec {
    /// Create a plan to write to `sink`
    /// Note: DataSinkExec requires its input to have a single partition.
    /// If the input has multiple partitions, the physical optimizer will
    /// automatically insert a Merge-related operator to merge them.
    /// If you construct PhysicalPlan without going through the physical optimizer,
    /// you must ensure that the input has a single partition.
    pub fn new(
        input: Arc<dyn ExecutionPlan>,
        sink: Arc<dyn DataSink>,
        sort_order: Option<LexRequirement>,
    ) -> Self {
        let count_schema = make_count_schema();
        let cache = Self::create_schema(&input, count_schema);
        Self {
            input,
            sink,
            count_schema: make_count_schema(),
            sort_order,
            cache: Arc::new(cache),
        }
    }

    /// Input execution plan
    pub fn input(&self) -> &Arc<dyn ExecutionPlan> {
        &self.input
    }

    /// Returns insert sink
    pub fn sink(&self) -> &dyn DataSink {
        self.sink.as_ref()
    }

    /// Optional sort order for output data
    pub fn sort_order(&self) -> &Option<LexRequirement> {
        &self.sort_order
    }

    /// Encode the optional sink ordering for a protobuf plan node.
    #[cfg(feature = "proto")]
    pub fn encode_sort_order(
        &self,
        ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection>>
    {
        use datafusion_physical_expr::PhysicalSortExpr;
        use datafusion_proto_models::protobuf;

        self.sort_order
            .as_ref()
            .map(|requirements| {
                requirements
                    .iter()
                    .map(|requirement| {
                        let expr: PhysicalSortExpr = requirement.to_owned().into();
                        Ok(protobuf::PhysicalSortExprNode {
                            expr: Some(Box::new(ctx.encode_expr(&expr.expr)?)),
                            asc: !expr.options.descending,
                            nulls_first: expr.options.nulls_first,
                        })
                    })
                    .collect::<Result<Vec<_>>>()
                    .map(|physical_sort_expr_nodes| {
                        protobuf::PhysicalSortExprNodeCollection {
                            physical_sort_expr_nodes,
                        }
                    })
            })
            .transpose()
    }

    /// Decode the optional sink ordering from a protobuf plan node.
    #[cfg(feature = "proto")]
    pub fn decode_sort_order(
        collection: Option<
            &datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection,
        >,
        ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>,
        schema: &Schema,
    ) -> Result<Option<LexRequirement>> {
        use arrow::compute::SortOptions;
        use datafusion_physical_expr::PhysicalSortExpr;

        let Some(collection) = collection else {
            return Ok(None);
        };
        let sort_exprs = collection
            .physical_sort_expr_nodes
            .iter()
            .map(|node| {
                let expr = node.expr.as_ref().ok_or_else(|| {
                    datafusion_common::internal_datafusion_err!(
                        "Unexpected empty physical expression"
                    )
                })?;
                Ok(PhysicalSortExpr {
                    expr: ctx.decode_expr(expr, schema)?,
                    options: SortOptions {
                        descending: !node.asc,
                        nulls_first: node.nulls_first,
                    },
                })
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(LexRequirement::new(sort_exprs.into_iter().map(Into::into)))
    }

    fn create_schema(
        input: &Arc<dyn ExecutionPlan>,
        schema: SchemaRef,
    ) -> PlanProperties {
        let eq_properties = EquivalenceProperties::new(schema);
        PlanProperties::new(
            eq_properties,
            Partitioning::UnknownPartitioning(1),
            input.pipeline_behavior(),
            input.boundedness(),
        )
        .with_scheduling_type(SchedulingType::Cooperative)
        .with_evaluation_type(EvaluationType::Eager)
    }
}

impl DisplayAs for DataSinkExec {
    fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
        match t {
            DisplayFormatType::Default | DisplayFormatType::Verbose => {
                write!(f, "DataSinkExec: sink=")?;
                self.sink.fmt_as(t, f)
            }
            DisplayFormatType::TreeRender => self.sink().fmt_as(t, f),
        }
    }
}

impl ExecutionPlan for DataSinkExec {
    fn name(&self) -> &'static str {
        "DataSinkExec"
    }

    /// Return a reference to Any that can be used for downcasting
    fn properties(&self) -> &Arc<PlanProperties> {
        &self.cache
    }

    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
        // DataSink is responsible for dynamically partitioning its
        // own input at execution time.
        vec![false]
    }

    fn required_input_distribution(&self) -> Vec<Distribution> {
        self.input_distribution_requirements().into_per_child()
    }

    fn input_distribution_requirements(&self) -> InputDistributionRequirements {
        // DataSink is responsible for dynamically partitioning its
        // own input at execution time, and so requires a single input partition.
        InputDistributionRequirements::new(vec![
            Distribution::SinglePartition;
            self.children().len()
        ])
    }

    fn required_input_ordering(&self) -> Vec<Option<OrderingRequirements>> {
        // The required input ordering is set externally (e.g. by a `ListingTable`).
        // Otherwise, there is no specific requirement (i.e. `sort_order` is `None`).
        vec![self.sort_order.as_ref().cloned().map(Into::into)]
    }

    fn maintains_input_order(&self) -> Vec<bool> {
        // Maintains ordering in the sense that the written file will reflect
        // the ordering of the input. For more context, see:
        //
        // https://github.com/apache/datafusion/pull/6354#discussion_r1195284178
        vec![true]
    }

    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
        vec![&self.input]
    }

    fn replace_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
        _: ReplaceChildrenOptions,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        Ok(Arc::new(Self::new(
            Arc::clone(&children[0]),
            Arc::clone(&self.sink),
            self.sort_order.clone(),
        )))
    }

    fn with_new_children(
        self: Arc<Self>,
        children: Vec<Arc<dyn ExecutionPlan>>,
    ) -> Result<Arc<dyn ExecutionPlan>> {
        self.replace_children(
            children,
            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
        )
    }

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

    /// Execute the plan and return a stream of `RecordBatch`es for
    /// the specified partition.
    fn execute(
        &self,
        partition: usize,
        context: Arc<TaskContext>,
    ) -> Result<SendableRecordBatchStream> {
        assert_eq_or_internal_err!(
            partition,
            0,
            "DataSinkExec can only be called on partition 0!"
        );
        let data = execute_input_stream(
            Arc::clone(&self.input),
            Arc::clone(self.sink.schema()),
            0,
            Arc::clone(&context),
        )?;

        let count_schema = Arc::clone(&self.count_schema);
        let sink = Arc::clone(&self.sink);

        let stream = futures::stream::once(async move {
            sink.write_all(data, &context).await.map(make_count_batch)
        })
        .boxed();

        Ok(Box::pin(RecordBatchStreamAdapter::new(
            count_schema,
            stream,
        )))
    }

    /// Returns the metrics of the underlying [DataSink]
    fn metrics(&self) -> Option<MetricsSet> {
        self.sink.metrics()
    }

    /// Delegates protobuf serialization to the underlying sink.
    #[cfg(feature = "proto")]
    fn try_to_proto(
        &self,
        ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>,
    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
        self.sink().try_to_proto(self, ctx)
    }
}

/// Create a output record batch with a count
///
/// ```text
/// +-------+,
/// | count |,
/// +-------+,
/// | 6     |,
/// +-------+,
/// ```
fn make_count_batch(count: u64) -> RecordBatch {
    let array = Arc::new(UInt64Array::from(vec![count])) as ArrayRef;

    RecordBatch::try_from_iter_with_nullable(vec![("count", array, false)]).unwrap()
}

fn make_count_schema() -> SchemaRef {
    // Define a schema.
    Arc::new(Schema::new(vec![Field::new(
        "count",
        DataType::UInt64,
        false,
    )]))
}