datafusion_physical_plan/
test.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//! Utilities for testing datafusion-physical-plan
19
20use std::collections::HashMap;
21use std::fmt;
22use std::fmt::{Debug, Formatter};
23use std::pin::Pin;
24use std::sync::Arc;
25use std::task::Context;
26
27use crate::common;
28use crate::execution_plan::{Boundedness, EmissionType};
29use crate::memory::MemoryStream;
30use crate::metrics::MetricsSet;
31use crate::statistics::StatisticsArgs;
32use crate::stream::RecordBatchStreamAdapter;
33use crate::streaming::PartitionStream;
34use crate::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions};
35use crate::{DisplayAs, DisplayFormatType, PlanProperties};
36
37use arrow::array::{Array, ArrayRef, Int32Array, RecordBatch};
38use arrow_schema::{DataType, Field, Schema, SchemaRef};
39use datafusion_common::tree_node::TreeNodeRecursion;
40use datafusion_common::{
41    Result, Statistics, assert_or_internal_err, config::ConfigOptions, project_schema,
42};
43use datafusion_execution::{SendableRecordBatchStream, TaskContext};
44use datafusion_physical_expr::equivalence::{
45    OrderingEquivalenceClass, ProjectionMapping,
46};
47use datafusion_physical_expr::expressions::Column;
48use datafusion_physical_expr::utils::collect_columns;
49use datafusion_physical_expr::{
50    EquivalenceProperties, LexOrdering, Partitioning, PhysicalExpr,
51};
52
53use futures::{Future, FutureExt};
54
55pub mod exec;
56
57/// `TestMemoryExec` is a mock equivalent to [`MemorySourceConfig`] with [`ExecutionPlan`] implemented for testing.
58/// i.e. It has some but not all the functionality of [`MemorySourceConfig`].
59/// This implements an in-memory DataSource rather than explicitly implementing a trait.
60/// It is implemented in this manner to keep relevant unit tests in place
61/// while avoiding circular dependencies between `datafusion-physical-plan` and `datafusion-datasource`.
62///
63/// [`MemorySourceConfig`]: https://github.com/apache/datafusion/tree/main/datafusion/datasource/src/memory.rs
64#[derive(Clone, Debug)]
65pub struct TestMemoryExec {
66    /// The partitions to query
67    partitions: Vec<Vec<RecordBatch>>,
68    /// Schema representing the data before projection
69    schema: SchemaRef,
70    /// Schema representing the data after the optional projection is applied
71    projected_schema: SchemaRef,
72    /// Optional projection
73    projection: Option<Vec<usize>>,
74    /// Sort information: one or more equivalent orderings
75    sort_information: Vec<LexOrdering>,
76    /// if partition sizes should be displayed
77    show_sizes: bool,
78    /// The maximum number of records to read from this plan. If `None`,
79    /// all records after filtering are returned.
80    fetch: Option<usize>,
81    cache: Arc<PlanProperties>,
82}
83
84impl DisplayAs for TestMemoryExec {
85    fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result {
86        write!(f, "DataSourceExec: ")?;
87        match t {
88            DisplayFormatType::Default | DisplayFormatType::Verbose => {
89                let partition_sizes: Vec<_> =
90                    self.partitions.iter().map(|b| b.len()).collect();
91
92                let output_ordering = self
93                    .sort_information
94                    .first()
95                    .map(|output_ordering| format!(", output_ordering={output_ordering}"))
96                    .unwrap_or_default();
97
98                let eq_properties = self.eq_properties();
99                let constraints = eq_properties.constraints();
100                let constraints = if constraints.is_empty() {
101                    String::new()
102                } else {
103                    format!(", {constraints}")
104                };
105
106                let limit = self
107                    .fetch
108                    .map_or(String::new(), |limit| format!(", fetch={limit}"));
109                if self.show_sizes {
110                    write!(
111                        f,
112                        "partitions={}, partition_sizes={partition_sizes:?}{limit}{output_ordering}{constraints}",
113                        partition_sizes.len(),
114                    )
115                } else {
116                    write!(
117                        f,
118                        "partitions={}{limit}{output_ordering}{constraints}",
119                        partition_sizes.len(),
120                    )
121                }
122            }
123            DisplayFormatType::TreeRender => {
124                // TODO: collect info
125                write!(f, "")
126            }
127        }
128    }
129}
130
131impl ExecutionPlan for TestMemoryExec {
132    fn name(&self) -> &'static str {
133        "DataSourceExec"
134    }
135
136    fn properties(&self) -> &Arc<PlanProperties> {
137        &self.cache
138    }
139
140    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
141        Vec::new()
142    }
143
144    fn apply_expressions(
145        &self,
146        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
147    ) -> Result<TreeNodeRecursion> {
148        Ok(TreeNodeRecursion::Continue)
149    }
150
151    fn replace_children(
152        self: Arc<Self>,
153        _: Vec<Arc<dyn ExecutionPlan>>,
154        _: ReplaceChildrenOptions,
155    ) -> Result<Arc<dyn ExecutionPlan>> {
156        Ok(self)
157    }
158
159    fn with_new_children(
160        self: Arc<Self>,
161        children: Vec<Arc<dyn ExecutionPlan>>,
162    ) -> Result<Arc<dyn ExecutionPlan>> {
163        self.replace_children(
164            children,
165            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
166        )
167    }
168
169    fn repartitioned(
170        &self,
171        _target_partitions: usize,
172        _config: &ConfigOptions,
173    ) -> Result<Option<Arc<dyn ExecutionPlan>>> {
174        unimplemented!()
175    }
176
177    fn execute(
178        &self,
179        partition: usize,
180        context: Arc<TaskContext>,
181    ) -> Result<SendableRecordBatchStream> {
182        self.open(partition, context)
183    }
184
185    fn metrics(&self) -> Option<MetricsSet> {
186        unimplemented!()
187    }
188
189    fn statistics_from_inputs(
190        &self,
191        _input_stats: &[Arc<Statistics>],
192        args: &StatisticsArgs,
193    ) -> Result<Arc<Statistics>> {
194        if args.partition().is_some() {
195            Ok(Arc::new(Statistics::new_unknown(&self.schema)))
196        } else {
197            Ok(Arc::new(self.statistics_inner()?))
198        }
199    }
200
201    fn fetch(&self) -> Option<usize> {
202        self.fetch
203    }
204}
205
206impl TestMemoryExec {
207    fn open(
208        &self,
209        partition: usize,
210        _context: Arc<TaskContext>,
211    ) -> Result<SendableRecordBatchStream> {
212        Ok(Box::pin(
213            MemoryStream::try_new(
214                self.partitions[partition].clone(),
215                Arc::clone(&self.projected_schema),
216                self.projection.clone(),
217            )?
218            .with_fetch(self.fetch),
219        ))
220    }
221
222    fn compute_properties(&self) -> PlanProperties {
223        PlanProperties::new(
224            self.eq_properties(),
225            self.output_partitioning(),
226            EmissionType::Incremental,
227            Boundedness::Bounded,
228        )
229    }
230
231    fn output_partitioning(&self) -> Partitioning {
232        Partitioning::UnknownPartitioning(self.partitions.len())
233    }
234
235    fn eq_properties(&self) -> EquivalenceProperties {
236        EquivalenceProperties::new_with_orderings(
237            Arc::clone(&self.projected_schema),
238            self.sort_information.clone(),
239        )
240    }
241
242    fn statistics_inner(&self) -> Result<Statistics> {
243        Ok(common::compute_record_batch_statistics(
244            &self.partitions,
245            &self.schema,
246            self.projection.clone(),
247        ))
248    }
249
250    pub fn try_new(
251        partitions: &[Vec<RecordBatch>],
252        schema: SchemaRef,
253        projection: Option<Vec<usize>>,
254    ) -> Result<Self> {
255        let projected_schema = project_schema(&schema, projection.as_ref())?;
256        Ok(Self {
257            partitions: partitions.to_vec(),
258            schema,
259            cache: Arc::new(PlanProperties::new(
260                EquivalenceProperties::new_with_orderings(
261                    Arc::clone(&projected_schema),
262                    Vec::<LexOrdering>::new(),
263                ),
264                Partitioning::UnknownPartitioning(partitions.len()),
265                EmissionType::Incremental,
266                Boundedness::Bounded,
267            )),
268            projected_schema,
269            projection,
270            sort_information: vec![],
271            show_sizes: true,
272            fetch: None,
273        })
274    }
275
276    /// Create a new `DataSourceExec` Equivalent plan for reading in-memory record batches
277    /// The provided `schema` should not have the projection applied.
278    pub fn try_new_exec(
279        partitions: &[Vec<RecordBatch>],
280        schema: SchemaRef,
281        projection: Option<Vec<usize>>,
282    ) -> Result<Arc<TestMemoryExec>> {
283        let mut source = Self::try_new(partitions, schema, projection)?;
284        let cache = source.compute_properties();
285        source.cache = Arc::new(cache);
286        Ok(Arc::new(source))
287    }
288
289    // Equivalent of `DataSourceExec::new`
290    pub fn update_cache(source: &Arc<TestMemoryExec>) -> TestMemoryExec {
291        let cache = source.compute_properties();
292        let mut source = (**source).clone();
293        source.cache = Arc::new(cache);
294        source
295    }
296
297    /// Set the limit of the files
298    pub fn with_limit(mut self, limit: Option<usize>) -> Self {
299        self.fetch = limit;
300        self
301    }
302
303    /// Ref to partitions
304    pub fn partitions(&self) -> &[Vec<RecordBatch>] {
305        &self.partitions
306    }
307
308    /// Ref to projection
309    pub fn projection(&self) -> &Option<Vec<usize>> {
310        &self.projection
311    }
312
313    /// Ref to sort information
314    pub fn sort_information(&self) -> &[LexOrdering] {
315        &self.sort_information
316    }
317
318    /// refer to `try_with_sort_information` at MemorySourceConfig for more information.
319    /// <https://github.com/apache/datafusion/tree/main/datafusion/datasource/src/memory.rs>
320    pub fn try_with_sort_information(
321        mut self,
322        mut sort_information: Vec<LexOrdering>,
323    ) -> Result<Self> {
324        // All sort expressions must refer to the original schema
325        let fields = self.schema.fields();
326        let ambiguous_column = sort_information
327            .iter()
328            .flat_map(|ordering| ordering.clone())
329            .flat_map(|expr| collect_columns(&expr.expr))
330            .find(|col| {
331                fields
332                    .get(col.index())
333                    .map(|field| field.name() != col.name())
334                    .unwrap_or(true)
335            });
336        assert_or_internal_err!(
337            ambiguous_column.is_none(),
338            "Column {:?} is not found in the original schema of the TestMemoryExec",
339            ambiguous_column.as_ref().unwrap()
340        );
341
342        // If there is a projection on the source, we also need to project orderings
343        if let Some(projection) = &self.projection {
344            let base_schema = self.original_schema();
345            let proj_exprs = projection.iter().map(|idx| {
346                let name = base_schema.field(*idx).name();
347                (Arc::new(Column::new(name, *idx)) as _, name.to_string())
348            });
349            let projection_mapping =
350                ProjectionMapping::try_new(proj_exprs, &base_schema)?;
351            let base_eqp = EquivalenceProperties::new_with_orderings(
352                Arc::clone(&base_schema),
353                sort_information,
354            );
355            let proj_eqp =
356                base_eqp.project(&projection_mapping, Arc::clone(&self.projected_schema));
357            let oeq_class: OrderingEquivalenceClass = proj_eqp.into();
358            sort_information = oeq_class.into();
359        }
360
361        self.sort_information = sort_information;
362        self.cache = Arc::new(self.compute_properties());
363        Ok(self)
364    }
365
366    /// Arc clone of ref to original schema
367    pub fn original_schema(&self) -> SchemaRef {
368        Arc::clone(&self.schema)
369    }
370}
371
372/// Asserts that given future is pending.
373pub fn assert_is_pending<'a, T>(fut: &mut Pin<Box<dyn Future<Output = T> + Send + 'a>>) {
374    let waker = futures::task::noop_waker();
375    let mut cx = Context::from_waker(&waker);
376    let poll = fut.poll_unpin(&mut cx);
377
378    assert!(poll.is_pending());
379}
380
381/// Get the schema for the aggregate_test_* csv files
382pub fn aggr_test_schema() -> SchemaRef {
383    let mut f1 = Field::new("c1", DataType::Utf8, false);
384    f1.set_metadata(HashMap::from_iter(vec![("testing".into(), "test".into())]));
385    let schema = Schema::new(vec![
386        f1,
387        Field::new("c2", DataType::UInt32, false),
388        Field::new("c3", DataType::Int8, false),
389        Field::new("c4", DataType::Int16, false),
390        Field::new("c5", DataType::Int32, false),
391        Field::new("c6", DataType::Int64, false),
392        Field::new("c7", DataType::UInt8, false),
393        Field::new("c8", DataType::UInt16, false),
394        Field::new("c9", DataType::UInt32, false),
395        Field::new("c10", DataType::UInt64, false),
396        Field::new("c11", DataType::Float32, false),
397        Field::new("c12", DataType::Float64, false),
398        Field::new("c13", DataType::Utf8, false),
399    ]);
400
401    Arc::new(schema)
402}
403
404/// Returns record batch with 3 columns of i32 in memory
405pub fn build_table_i32(
406    a: (&str, &Vec<i32>),
407    b: (&str, &Vec<i32>),
408    c: (&str, &Vec<i32>),
409) -> RecordBatch {
410    let schema = Schema::new(vec![
411        Field::new(a.0, DataType::Int32, false),
412        Field::new(b.0, DataType::Int32, false),
413        Field::new(c.0, DataType::Int32, false),
414    ]);
415
416    RecordBatch::try_new(
417        Arc::new(schema),
418        vec![
419            Arc::new(Int32Array::from(a.1.clone())),
420            Arc::new(Int32Array::from(b.1.clone())),
421            Arc::new(Int32Array::from(c.1.clone())),
422        ],
423    )
424    .unwrap()
425}
426
427/// Returns record batch with 2 columns of i32 in memory
428pub fn build_table_i32_two_cols(
429    a: (&str, &Vec<i32>),
430    b: (&str, &Vec<i32>),
431) -> RecordBatch {
432    let schema = Schema::new(vec![
433        Field::new(a.0, DataType::Int32, false),
434        Field::new(b.0, DataType::Int32, false),
435    ]);
436
437    RecordBatch::try_new(
438        Arc::new(schema),
439        vec![
440            Arc::new(Int32Array::from(a.1.clone())),
441            Arc::new(Int32Array::from(b.1.clone())),
442        ],
443    )
444    .unwrap()
445}
446
447/// Returns memory table scan wrapped around record batch with 3 columns of i32
448pub fn build_table_scan_i32(
449    a: (&str, &Vec<i32>),
450    b: (&str, &Vec<i32>),
451    c: (&str, &Vec<i32>),
452) -> Arc<dyn ExecutionPlan> {
453    let batch = build_table_i32(a, b, c);
454    let schema = batch.schema();
455    TestMemoryExec::try_new_exec(&[vec![batch]], schema, None).unwrap()
456}
457
458/// Return a RecordBatch with a single Int32 array with values (0..sz) in a field named "i"
459pub fn make_partition(sz: i32) -> RecordBatch {
460    let seq_start = 0;
461    let seq_end = sz;
462    let values = (seq_start..seq_end).collect::<Vec<_>>();
463    let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, true)]));
464    let arr = Arc::new(Int32Array::from(values));
465    let arr = arr as ArrayRef;
466
467    RecordBatch::try_new(schema, vec![arr]).unwrap()
468}
469
470pub fn make_partition_utf8(sz: i32) -> RecordBatch {
471    let seq_start = 0;
472    let seq_end = sz;
473    let values = (seq_start..seq_end)
474        .map(|i| format!("test_long_string_that_is_roughly_42_bytes_{i}"))
475        .collect::<Vec<_>>();
476    let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Utf8, true)]));
477    let mut string_array = arrow::array::StringArray::from(values);
478    string_array.shrink_to_fit();
479    let arr = Arc::new(string_array);
480    let arr = arr as ArrayRef;
481
482    RecordBatch::try_new(schema, vec![arr]).unwrap()
483}
484
485/// Returns a `DataSourceExec` that scans `partitions` of 100 batches each
486pub fn scan_partitioned(partitions: usize) -> Arc<dyn ExecutionPlan> {
487    Arc::new(mem_exec(partitions))
488}
489
490pub fn scan_partitioned_utf8(partitions: usize) -> Arc<dyn ExecutionPlan> {
491    Arc::new(mem_exec_utf8(partitions))
492}
493
494/// Returns a `DataSourceExec` that scans `partitions` of 100 batches each
495pub fn mem_exec(partitions: usize) -> TestMemoryExec {
496    let data: Vec<Vec<_>> = (0..partitions).map(|_| vec![make_partition(100)]).collect();
497
498    let schema = data[0][0].schema();
499    let projection = None;
500
501    TestMemoryExec::try_new(&data, schema, projection).unwrap()
502}
503
504pub fn mem_exec_utf8(partitions: usize) -> TestMemoryExec {
505    let data: Vec<Vec<_>> = (0..partitions)
506        .map(|_| vec![make_partition_utf8(100)])
507        .collect();
508
509    let schema = data[0][0].schema();
510    let projection = None;
511
512    TestMemoryExec::try_new(&data, schema, projection).unwrap()
513}
514
515// Construct a stream partition for test purposes
516#[derive(Debug)]
517pub struct TestPartitionStream {
518    pub schema: SchemaRef,
519    pub batches: Vec<RecordBatch>,
520}
521
522impl TestPartitionStream {
523    /// Create a new stream partition with the provided batches
524    pub fn new_with_batches(batches: Vec<RecordBatch>) -> Self {
525        let schema = batches[0].schema();
526        Self { schema, batches }
527    }
528}
529impl PartitionStream for TestPartitionStream {
530    fn schema(&self) -> &SchemaRef {
531        &self.schema
532    }
533    fn execute(&self, _ctx: Arc<TaskContext>) -> SendableRecordBatchStream {
534        let stream = futures::stream::iter(self.batches.clone().into_iter().map(Ok));
535        Box::pin(RecordBatchStreamAdapter::new(
536            Arc::clone(&self.schema),
537            stream,
538        ))
539    }
540}
541
542#[cfg(test)]
543macro_rules! assert_join_metrics {
544    ($metrics:expr, $expected_rows:expr) => {
545        assert_eq!($metrics.output_rows().unwrap(), $expected_rows);
546
547        let elapsed_compute = $metrics
548            .elapsed_compute()
549            .expect("did not find elapsed_compute metric");
550        let join_time = $metrics
551            .sum_by_name("join_time")
552            .expect("did not find join_time metric")
553            .as_usize();
554        let build_time = $metrics
555            .sum_by_name("build_time")
556            .expect("did not find build_time metric")
557            .as_usize();
558        // ensure join_time and build_time are considered in elapsed_compute
559        assert!(
560            join_time + build_time <= elapsed_compute,
561            "join_time ({}) + build_time ({}) = {} was <= elapsed_compute = {}",
562            join_time,
563            build_time,
564            join_time + build_time,
565            elapsed_compute
566        );
567    };
568}
569#[cfg(test)]
570pub(crate) use assert_join_metrics;
datafusion_physical_plan/test.rs

datafusion_physical_plan/
test.rs
