/Users/andrewlamb/Software/datafusion/datafusion/physical-plan/src/placeholder_row.rs
Line | Count | Source (jump to first uncovered line) |
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 | | |
20 | | use std::any::Any; |
21 | | use std::sync::Arc; |
22 | | |
23 | | use super::{ |
24 | | common, DisplayAs, ExecutionMode, PlanProperties, SendableRecordBatchStream, |
25 | | Statistics, |
26 | | }; |
27 | | use crate::{memory::MemoryStream, DisplayFormatType, ExecutionPlan, Partitioning}; |
28 | | |
29 | | use arrow::array::{ArrayRef, NullArray}; |
30 | | use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; |
31 | | use arrow::record_batch::RecordBatch; |
32 | | use arrow_array::RecordBatchOptions; |
33 | | use datafusion_common::{internal_err, Result}; |
34 | | use datafusion_execution::TaskContext; |
35 | | use datafusion_physical_expr::EquivalenceProperties; |
36 | | |
37 | | use log::trace; |
38 | | |
39 | | /// Execution plan for empty relation with produce_one_row=true |
40 | | #[derive(Debug)] |
41 | | pub struct PlaceholderRowExec { |
42 | | /// The schema for the produced row |
43 | | schema: SchemaRef, |
44 | | /// Number of partitions |
45 | | partitions: usize, |
46 | | cache: PlanProperties, |
47 | | } |
48 | | |
49 | | impl PlaceholderRowExec { |
50 | | /// Create a new PlaceholderRowExec |
51 | 4 | pub fn new(schema: SchemaRef) -> Self { |
52 | 4 | let partitions = 1; |
53 | 4 | let cache = Self::compute_properties(Arc::clone(&schema), partitions); |
54 | 4 | PlaceholderRowExec { |
55 | 4 | schema, |
56 | 4 | partitions, |
57 | 4 | cache, |
58 | 4 | } |
59 | 4 | } |
60 | | |
61 | | /// Create a new PlaceholderRowExecPlaceholderRowExec with specified partition number |
62 | 1 | pub fn with_partitions(mut self, partitions: usize) -> Self { |
63 | 1 | self.partitions = partitions; |
64 | 1 | // Update output partitioning when updating partitions: |
65 | 1 | let output_partitioning = Self::output_partitioning_helper(self.partitions); |
66 | 1 | self.cache = self.cache.with_partitioning(output_partitioning); |
67 | 1 | self |
68 | 1 | } |
69 | | |
70 | 4 | fn data(&self) -> Result<Vec<RecordBatch>> { |
71 | 4 | Ok({ |
72 | 4 | let n_field = self.schema.fields.len(); |
73 | 4 | vec![RecordBatch::try_new_with_options( |
74 | 4 | Arc::new(Schema::new( |
75 | 4 | (0..n_field) |
76 | 52 | .map(|i| { |
77 | 52 | Field::new(format!("placeholder_{i}"), DataType::Null, true) |
78 | 52 | }) |
79 | 4 | .collect::<Fields>(), |
80 | 4 | )), |
81 | 4 | (0..n_field) |
82 | 52 | .map(|_i| { |
83 | 52 | let ret: ArrayRef = Arc::new(NullArray::new(1)); |
84 | 52 | ret |
85 | 52 | }) |
86 | 4 | .collect(), |
87 | 4 | // Even if column number is empty we can generate single row. |
88 | 4 | &RecordBatchOptions::new().with_row_count(Some(1)), |
89 | 4 | )?0 ] |
90 | | }) |
91 | 4 | } |
92 | | |
93 | 5 | fn output_partitioning_helper(n_partitions: usize) -> Partitioning { |
94 | 5 | Partitioning::UnknownPartitioning(n_partitions) |
95 | 5 | } |
96 | | |
97 | | /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. |
98 | 4 | fn compute_properties(schema: SchemaRef, n_partitions: usize) -> PlanProperties { |
99 | 4 | let eq_properties = EquivalenceProperties::new(schema); |
100 | 4 | // Get output partitioning: |
101 | 4 | let output_partitioning = Self::output_partitioning_helper(n_partitions); |
102 | 4 | |
103 | 4 | PlanProperties::new(eq_properties, output_partitioning, ExecutionMode::Bounded) |
104 | 4 | } |
105 | | } |
106 | | |
107 | | impl DisplayAs for PlaceholderRowExec { |
108 | 0 | fn fmt_as( |
109 | 0 | &self, |
110 | 0 | t: DisplayFormatType, |
111 | 0 | f: &mut std::fmt::Formatter, |
112 | 0 | ) -> std::fmt::Result { |
113 | 0 | match t { |
114 | | DisplayFormatType::Default | DisplayFormatType::Verbose => { |
115 | 0 | write!(f, "PlaceholderRowExec") |
116 | 0 | } |
117 | 0 | } |
118 | 0 | } |
119 | | } |
120 | | |
121 | | impl ExecutionPlan for PlaceholderRowExec { |
122 | 0 | fn name(&self) -> &'static str { |
123 | 0 | "PlaceholderRowExec" |
124 | 0 | } |
125 | | |
126 | | /// Return a reference to Any that can be used for downcasting |
127 | 0 | fn as_any(&self) -> &dyn Any { |
128 | 0 | self |
129 | 0 | } |
130 | | |
131 | 2 | fn properties(&self) -> &PlanProperties { |
132 | 2 | &self.cache |
133 | 2 | } |
134 | | |
135 | 2 | fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { |
136 | 2 | vec![] |
137 | 2 | } |
138 | | |
139 | 1 | fn with_new_children( |
140 | 1 | self: Arc<Self>, |
141 | 1 | _: Vec<Arc<dyn ExecutionPlan>>, |
142 | 1 | ) -> Result<Arc<dyn ExecutionPlan>> { |
143 | 1 | Ok(self) |
144 | 1 | } |
145 | | |
146 | 6 | fn execute( |
147 | 6 | &self, |
148 | 6 | partition: usize, |
149 | 6 | context: Arc<TaskContext>, |
150 | 6 | ) -> Result<SendableRecordBatchStream> { |
151 | 6 | trace!("Start PlaceholderRowExec::execute for partition {} of context session_id {} and task_id {:?}", partition, context.session_id(), context.task_id()0 ); |
152 | | |
153 | 6 | if partition >= self.partitions { |
154 | 2 | return internal_err!( |
155 | 2 | "PlaceholderRowExec invalid partition {} (expected less than {})", |
156 | 2 | partition, |
157 | 2 | self.partitions |
158 | 2 | ); |
159 | 4 | } |
160 | 4 | |
161 | 4 | Ok(Box::pin(MemoryStream::try_new( |
162 | 4 | self.data()?0 , |
163 | 4 | Arc::clone(&self.schema), |
164 | 4 | None, |
165 | 0 | )?)) |
166 | 6 | } |
167 | | |
168 | 0 | fn statistics(&self) -> Result<Statistics> { |
169 | 0 | let batch = self |
170 | 0 | .data() |
171 | 0 | .expect("Create single row placeholder RecordBatch should not fail"); |
172 | 0 | Ok(common::compute_record_batch_statistics( |
173 | 0 | &[batch], |
174 | 0 | &self.schema, |
175 | 0 | None, |
176 | 0 | )) |
177 | 0 | } |
178 | | } |
179 | | |
180 | | #[cfg(test)] |
181 | | mod tests { |
182 | | use super::*; |
183 | | use crate::{test, with_new_children_if_necessary}; |
184 | | |
185 | | #[test] |
186 | 1 | fn with_new_children() -> Result<()> { |
187 | 1 | let schema = test::aggr_test_schema(); |
188 | 1 | |
189 | 1 | let placeholder = Arc::new(PlaceholderRowExec::new(schema)); |
190 | | |
191 | 1 | let placeholder_2 = with_new_children_if_necessary( |
192 | 1 | Arc::clone(&placeholder) as Arc<dyn ExecutionPlan>, |
193 | 1 | vec![], |
194 | 1 | )?0 ; |
195 | 1 | assert_eq!(placeholder.schema(), placeholder_2.schema()); |
196 | | |
197 | 1 | let too_many_kids = vec![placeholder_2]; |
198 | 1 | assert!( |
199 | 1 | with_new_children_if_necessary(placeholder, too_many_kids).is_err(), |
200 | 0 | "expected error when providing list of kids" |
201 | | ); |
202 | 1 | Ok(()) |
203 | 1 | } |
204 | | |
205 | | #[tokio::test] |
206 | 1 | async fn invalid_execute() -> Result<()> { |
207 | 1 | let task_ctx = Arc::new(TaskContext::default()); |
208 | 1 | let schema = test::aggr_test_schema(); |
209 | 1 | let placeholder = PlaceholderRowExec::new(schema); |
210 | 1 | |
211 | 1 | // ask for the wrong partition |
212 | 1 | assert!(placeholder.execute(1, Arc::clone(&task_ctx)).is_err()); |
213 | 1 | assert!(placeholder.execute(20, task_ctx).is_err()); |
214 | 1 | Ok(()) |
215 | 1 | } |
216 | | |
217 | | #[tokio::test] |
218 | 1 | async fn produce_one_row() -> Result<()> { |
219 | 1 | let task_ctx = Arc::new(TaskContext::default()); |
220 | 1 | let schema = test::aggr_test_schema(); |
221 | 1 | let placeholder = PlaceholderRowExec::new(schema); |
222 | 1 | |
223 | 1 | let iter = placeholder.execute(0, task_ctx)?0 ; |
224 | 1 | let batches = common::collect(iter).await0 ?0 ; |
225 | 1 | |
226 | 1 | // should have one item |
227 | 1 | assert_eq!(batches.len(), 1); |
228 | 1 | |
229 | 1 | Ok(()) |
230 | 1 | } |
231 | | |
232 | | #[tokio::test] |
233 | 1 | async fn produce_one_row_multiple_partition() -> Result<()> { |
234 | 1 | let task_ctx = Arc::new(TaskContext::default()); |
235 | 1 | let schema = test::aggr_test_schema(); |
236 | 1 | let partitions = 3; |
237 | 1 | let placeholder = PlaceholderRowExec::new(schema).with_partitions(partitions); |
238 | 1 | |
239 | 3 | for n in 0..partitions1 { |
240 | 3 | let iter = placeholder.execute(n, Arc::clone(&task_ctx))?0 ; |
241 | 3 | let batches = common::collect(iter).await0 ?0 ; |
242 | 1 | |
243 | 1 | // should have one item |
244 | 3 | assert_eq!(batches.len(), 1); |
245 | 1 | } |
246 | 1 | |
247 | 1 | Ok(()) |
248 | 1 | } |
249 | | } |