datafusion_physical_plan/
placeholder_row.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! EmptyRelation produce_one_row=true execution plan
19
20use std::sync::Arc;
21
22use crate::coop::cooperative;
23use crate::execution_plan::{Boundedness, EmissionType, SchedulingType};
24use crate::memory::MemoryStream;
25use crate::{
26    ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning,
27    PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, Statistics,
28    common,
29};
30
31use arrow::array::{ArrayRef, NullArray, RecordBatch, RecordBatchOptions};
32use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef};
33use datafusion_common::tree_node::TreeNodeRecursion;
34use datafusion_common::{Result, assert_or_internal_err};
35use datafusion_execution::TaskContext;
36use datafusion_physical_expr::EquivalenceProperties;
37use datafusion_physical_expr::PhysicalExpr;
38
39use crate::statistics::StatisticsArgs;
40use log::trace;
41
42/// Execution plan for empty relation with produce_one_row=true
43#[derive(Debug, Clone)]
44pub struct PlaceholderRowExec {
45    /// The schema for the produced row
46    schema: SchemaRef,
47    /// Number of partitions
48    partitions: usize,
49    cache: Arc<PlanProperties>,
50}
51
52impl PlaceholderRowExec {
53    /// Create a new PlaceholderRowExec
54    pub fn new(schema: SchemaRef) -> Self {
55        let partitions = 1;
56        let cache = Self::compute_properties(Arc::clone(&schema), partitions);
57        PlaceholderRowExec {
58            schema,
59            partitions,
60            cache: Arc::new(cache),
61        }
62    }
63
64    /// Create a new PlaceholderRowExecPlaceholderRowExec with specified partition number
65    pub fn with_partitions(mut self, partitions: usize) -> Self {
66        self.partitions = partitions;
67        // Update output partitioning when updating partitions:
68        let output_partitioning = Self::output_partitioning_helper(self.partitions);
69        Arc::make_mut(&mut self.cache).partitioning = output_partitioning;
70        self
71    }
72
73    fn data(&self) -> Result<Vec<RecordBatch>> {
74        Ok({
75            let n_field = self.schema.fields.len();
76            vec![RecordBatch::try_new_with_options(
77                Arc::new(Schema::new(
78                    (0..n_field)
79                        .map(|i| {
80                            Field::new(format!("placeholder_{i}"), DataType::Null, true)
81                        })
82                        .collect::<Fields>(),
83                )),
84                (0..n_field)
85                    .map(|_i| {
86                        let ret: ArrayRef = Arc::new(NullArray::new(1));
87                        ret
88                    })
89                    .collect(),
90                // Even if column number is empty we can generate single row.
91                &RecordBatchOptions::new().with_row_count(Some(1)),
92            )?]
93        })
94    }
95
96    fn output_partitioning_helper(n_partitions: usize) -> Partitioning {
97        Partitioning::UnknownPartitioning(n_partitions)
98    }
99
100    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
101    fn compute_properties(schema: SchemaRef, n_partitions: usize) -> PlanProperties {
102        PlanProperties::new(
103            EquivalenceProperties::new(schema),
104            Self::output_partitioning_helper(n_partitions),
105            EmissionType::Incremental,
106            Boundedness::Bounded,
107        )
108        .with_scheduling_type(SchedulingType::Cooperative)
109    }
110}
111
112impl DisplayAs for PlaceholderRowExec {
113    fn fmt_as(
114        &self,
115        t: DisplayFormatType,
116        f: &mut std::fmt::Formatter,
117    ) -> std::fmt::Result {
118        match t {
119            DisplayFormatType::Default | DisplayFormatType::Verbose => {
120                write!(f, "PlaceholderRowExec")
121            }
122
123            DisplayFormatType::TreeRender => Ok(()),
124        }
125    }
126}
127
128impl ExecutionPlan for PlaceholderRowExec {
129    fn name(&self) -> &'static str {
130        "PlaceholderRowExec"
131    }
132
133    /// Return a reference to Any that can be used for downcasting
134    fn properties(&self) -> &Arc<PlanProperties> {
135        &self.cache
136    }
137
138    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
139        vec![]
140    }
141
142    fn apply_expressions(
143        &self,
144        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
145    ) -> Result<TreeNodeRecursion> {
146        Ok(TreeNodeRecursion::Continue)
147    }
148
149    fn replace_children(
150        self: Arc<Self>,
151        _: Vec<Arc<dyn ExecutionPlan>>,
152        _: ReplaceChildrenOptions,
153    ) -> Result<Arc<dyn ExecutionPlan>> {
154        Ok(self)
155    }
156
157    fn with_new_children(
158        self: Arc<Self>,
159        children: Vec<Arc<dyn ExecutionPlan>>,
160    ) -> Result<Arc<dyn ExecutionPlan>> {
161        self.replace_children(
162            children,
163            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
164        )
165    }
166
167    fn execute(
168        &self,
169        partition: usize,
170        context: Arc<TaskContext>,
171    ) -> Result<SendableRecordBatchStream> {
172        trace!(
173            "Start PlaceholderRowExec::execute for partition {} of context session_id {} and task_id {:?}",
174            partition,
175            context.session_id(),
176            context.task_id()
177        );
178
179        assert_or_internal_err!(
180            partition < self.partitions,
181            "PlaceholderRowExec invalid partition {partition} (expected less than {})",
182            self.partitions
183        );
184
185        let ms = MemoryStream::try_new(self.data()?, Arc::clone(&self.schema), None)?;
186        Ok(Box::pin(cooperative(ms)))
187    }
188
189    fn statistics_from_inputs(
190        &self,
191        _input_stats: &[Arc<Statistics>],
192        args: &StatisticsArgs,
193    ) -> Result<Arc<Statistics>> {
194        let batches = self
195            .data()
196            .expect("Create single row placeholder RecordBatch should not fail");
197
198        let batches = match args.partition() {
199            Some(_) => vec![batches],
200            // entire plan
201            None => vec![batches; self.partitions],
202        };
203
204        Ok(Arc::new(common::compute_record_batch_statistics(
205            &batches,
206            &self.schema,
207            None,
208        )))
209    }
210
211    #[cfg(feature = "proto")]
212    fn try_to_proto(
213        &self,
214        _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>,
215    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalPlanNode>> {
216        use datafusion_proto_models::protobuf;
217        let schema = self.schema().as_ref().try_into()?;
218        Ok(Some(protobuf::PhysicalPlanNode {
219            physical_plan_type: Some(
220                protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow(
221                    protobuf::PlaceholderRowExecNode {
222                        schema: Some(schema),
223                        partitions: self
224                            .properties()
225                            .output_partitioning()
226                            .partition_count() as u32,
227                    },
228                ),
229            ),
230        }))
231    }
232}
233
234#[cfg(feature = "proto")]
235impl PlaceholderRowExec {
236    /// Reconstruct a [`PlaceholderRowExec`] from its protobuf representation.
237    pub fn try_from_proto(
238        node: &datafusion_proto_models::protobuf::PhysicalPlanNode,
239        _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>,
240    ) -> Result<Arc<dyn ExecutionPlan>> {
241        use datafusion_proto_models::protobuf;
242        let placeholder = crate::expect_plan_variant!(
243            node,
244            protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow,
245            "PlaceholderRowExec",
246        );
247        let schema = placeholder.schema.as_ref().ok_or_else(|| {
248            datafusion_common::internal_datafusion_err!(
249                "PlaceholderRowExec is missing required field 'schema'"
250            )
251        })?;
252        let schema = Arc::new(Schema::try_from(schema)?);
253        // A zero (absent) partition count comes from a plan encoded before the
254        // field existed, which always meant a single partition.
255        let partitions = placeholder.partitions.max(1) as usize;
256        Ok(Arc::new(
257            PlaceholderRowExec::new(schema).with_partitions(partitions),
258        ))
259    }
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265    use crate::{execution_plan::replace_children_if_necessary, test};
266
267    #[test]
268    fn replace_children() -> Result<()> {
269        let schema = test::aggr_test_schema();
270
271        let placeholder = Arc::new(PlaceholderRowExec::new(schema));
272
273        let placeholder_2 = replace_children_if_necessary(
274            Arc::clone(&placeholder) as Arc<dyn ExecutionPlan>,
275            vec![],
276        )?;
277        assert_eq!(placeholder.schema(), placeholder_2.schema());
278
279        let too_many_kids = vec![placeholder_2];
280        assert!(
281            replace_children_if_necessary(placeholder, too_many_kids).is_err(),
282            "expected error when providing list of kids"
283        );
284        Ok(())
285    }
286
287    #[tokio::test]
288    async fn invalid_execute() -> Result<()> {
289        let task_ctx = Arc::new(TaskContext::default());
290        let schema = test::aggr_test_schema();
291        let placeholder = PlaceholderRowExec::new(schema);
292
293        // Ask for the wrong partition
294        assert!(placeholder.execute(1, Arc::clone(&task_ctx)).is_err());
295        assert!(placeholder.execute(20, task_ctx).is_err());
296        Ok(())
297    }
298
299    #[tokio::test]
300    async fn produce_one_row() -> Result<()> {
301        let task_ctx = Arc::new(TaskContext::default());
302        let schema = test::aggr_test_schema();
303        let placeholder = PlaceholderRowExec::new(schema);
304
305        let iter = placeholder.execute(0, task_ctx)?;
306        let batches = common::collect(iter).await?;
307
308        // Should have one item
309        assert_eq!(batches.len(), 1);
310
311        Ok(())
312    }
313
314    #[tokio::test]
315    async fn produce_one_row_multiple_partition() -> Result<()> {
316        let task_ctx = Arc::new(TaskContext::default());
317        let schema = test::aggr_test_schema();
318        let partitions = 3;
319        let placeholder = PlaceholderRowExec::new(schema).with_partitions(partitions);
320
321        for n in 0..partitions {
322            let iter = placeholder.execute(n, Arc::clone(&task_ctx))?;
323            let batches = common::collect(iter).await?;
324
325            // Should have one item
326            assert_eq!(batches.len(), 1);
327        }
328
329        Ok(())
330    }
331}
datafusion_physical_plan/placeholder_row.rs

datafusion_physical_plan/
placeholder_row.rs
