## Coverage Report

#### Created: 2024-10-13 08:39

    /Users/andrewlamb/Software/datafusion/datafusion/physical-plan/src/memory.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
    //! Execution plan for reading in-memory batches of data
    19
    20
    use std::any::Any;
    21
    use std::fmt;
    22
    use std::sync::Arc;
    23
    use std::task::{Context, Poll};
    24
    25
    use super::expressions::PhysicalSortExpr;
    26
    use super::{
    27
        common, DisplayAs, DisplayFormatType, ExecutionMode, ExecutionPlan, Partitioning,
    28
        PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics,
    29
    };
    30
    31
    use arrow::datatypes::SchemaRef;
    32
    use arrow::record_batch::RecordBatch;
    33
    use datafusion_common::{internal_err, project_schema, Result};
    34
    use datafusion_execution::memory_pool::MemoryReservation;
    35
    use datafusion_execution::TaskContext;
    36
    use datafusion_physical_expr::{EquivalenceProperties, LexOrdering};
    37
    38
    use futures::Stream;
    39
    40
    /// Execution plan for reading in-memory batches of data
    41
    pub struct MemoryExec {
    42
        /// The partitions to query
    43
        partitions: Vec<Vec<RecordBatch>>,
    44
        /// Schema representing the data before projection
    45
        schema: SchemaRef,
    46
        /// Schema representing the data after the optional projection is applied
    47
        projected_schema: SchemaRef,
    48
        /// Optional projection
    49
        projection: Option<Vec<usize>>,
    50
        // Sort information: one or more equivalent orderings
    51
        sort_information: Vec<LexOrdering>,
    52
        cache: PlanProperties,
    53
        /// if partition sizes should be displayed
    54
        show_sizes: bool,
    55
    }
    56
    57
    impl fmt::Debug for MemoryExec {
    58
    0
        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
    59
    0
            f.debug_struct("MemoryExec")
    60
    0
                .field("partitions", &"[...]")
    61
    0
                .field("schema", &self.schema)
    62
    0
                .field("projection", &self.projection)
    63
    0
                .field("sort_information", &self.sort_information)
    64
    0
                .finish()
    65
    0
        }
    66
    }
    67
    68
    impl DisplayAs for MemoryExec {
    69
    6
        fn fmt_as(
    70
    6
            &self,
    71
    6
            t: DisplayFormatType,
    72
    6
            f: &mut std::fmt::Formatter,
    73
    6
        ) -> std::fmt::Result {
    74
    6
            match t {
    75
                DisplayFormatType::Default | DisplayFormatType::Verbose => {
    76
    6
                    let partition_sizes: Vec<_> =
    77
    6
                        self.partitions.iter().map(|b| b.len()).collect();
    78
    6
    79
    6
                    let output_ordering = self
    80
    6
                        .sort_information
    81
    6
                        .first()
    82
    6
                        .map(|output_ordering| {
    83
    3
                            format!(
    84
    3
                                ", output_ordering={}",
    85
    3
                                PhysicalSortExpr::format_list(output_ordering)
    86
    3
                            )
    87
    6
                        })
    88
    6
                        .unwrap_or_default();
    89
    6
    90
    6
                    if self.show_sizes {
    91
    6
                        write!(
    92
    6
                            f,
    93
    6
                            "MemoryExec: partitions={}, partition_sizes={partition_sizes:?}{output_ordering}",
    94
    6
                            partition_sizes.len(),
    95
    6
                        )
    96
                    } else {
    97
    0
                        write!(f, "MemoryExec: partitions={}", partition_sizes.len(),)
    98
                    }
    99
                }
    100
            }
    101
    6
        }
    102
    }
    103
    104
    impl ExecutionPlan for MemoryExec {
    105
    0
        fn name(&self) -> &'static str {
    106
    0
            "MemoryExec"
    107
    0
        }
    108
    109
        /// Return a reference to Any that can be used for downcasting
    110
    0
        fn as_any(&self) -> &dyn Any {
    111
    0
            self
    112
    0
        }
    113
    114
    26.6k
        fn properties(&self) -> &PlanProperties {
    115
    26.6k
            &self.cache
    116
    26.6k
        }
    117
    118
    6
        fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
    119
    6
            // this is a leaf node and has no children
    120
    6
            vec![]
    121
    6
        }
    122
    123
    0
        fn with_new_children(
    124
    0
            self: Arc<Self>,
    125
    0
            children: Vec<Arc<dyn ExecutionPlan>>,
    126
    0
        ) -> Result<Arc<dyn ExecutionPlan>> {
    127
    0
            // MemoryExec has no children
    128
    0
            if children.is_empty() {
    129
    0
                Ok(self)
    130
            } else {
    131
    0
                internal_err!("Children cannot be replaced in {self:?}")
    132
            }
    133
    0
        }
    134
    135
    2.58k
        fn execute(
    136
    2.58k
            &self,
    137
    2.58k
            partition: usize,
    138
    2.58k
            _context: Arc<TaskContext>,
    139
    2.58k
        ) -> Result<SendableRecordBatchStream> {
    140
    2.58k
            Ok(Box::pin(MemoryStream::try_new(
    141
    2.58k
                self.partitions[partition].clone(),
    142
    2.58k
                Arc::clone(&self.projected_schema),
    143
    2.58k
                self.projection.clone(),
    144
    2.58k
            )?0))
    145
    2.58k
        }
    146
    147
        /// We recompute the statistics dynamically from the arrow metadata as it is pretty cheap to do so
    148
    25
        fn statistics(&self) -> Result<Statistics> {
    149
    25
            Ok(common::compute_record_batch_statistics(
    150
    25
                &self.partitions,
    151
    25
                &self.schema,
    152
    25
                self.projection.clone(),
    153
    25
            ))
    154
    25
        }
    155
    }
    156
    157
    impl MemoryExec {
    158
        /// Create a new execution plan for reading in-memory record batches
    159
        /// The provided `schema` should not have the projection applied.
    160
    1.17k
        pub fn try_new(
    161
    1.17k
            partitions: &[Vec<RecordBatch>],
    162
    1.17k
            schema: SchemaRef,
    163
    1.17k
            projection: Option<Vec<usize>>,
    164
    1.17k
        ) -> Result<Self> {
    165
    1.17k
            let projected_schema = project_schema(&schema, projection.as_ref())?0;
    166
    1.17k
            let cache =
    167
    1.17k
                Self::compute_properties(Arc::clone(&projected_schema), &[], partitions);
    168
    1.17k
            Ok(Self {
    169
    1.17k
                partitions: partitions.to_vec(),
    170
    1.17k
                schema,
    171
    1.17k
                projected_schema,
    172
    1.17k
                projection,
    173
    1.17k
                sort_information: vec![],
    174
    1.17k
                cache,
    175
    1.17k
                show_sizes: true,
    176
    1.17k
            })
    177
    1.17k
        }
    178
    179
        /// set `show_sizes` to determine whether to display partition sizes
    180
    0
        pub fn with_show_sizes(mut self, show_sizes: bool) -> Self {
    181
    0
            self.show_sizes = show_sizes;
    182
    0
            self
    183
    0
        }
    184
    185
    0
        pub fn partitions(&self) -> &[Vec<RecordBatch>] {
    186
    0
            &self.partitions
    187
    0
        }
    188
    189
    0
        pub fn projection(&self) -> &Option<Vec<usize>> {
    190
    0
            &self.projection
    191
    0
        }
    192
    193
        /// A memory table can be ordered by multiple expressions simultaneously.
    194
        /// [`EquivalenceProperties`] keeps track of expressions that describe the
    195
        /// global ordering of the schema. These columns are not necessarily same; e.g.
    196
        /// ```text
    197
        /// ┌-------┐
    198
        /// | a | b |
    199
        /// |---|---|
    200
        /// | 1 | 9 |
    201
        /// | 2 | 8 |
    202
        /// | 3 | 7 |
    203
        /// | 5 | 5 |
    204
        /// └---┴---┘
    205
        /// ```
    206
        /// where both `a ASC` and `b DESC` can describe the table ordering. With
    207
        /// [`EquivalenceProperties`], we can keep track of these equivalences
    208
        /// and treat `a ASC` and `b DESC` as the same ordering requirement.
    209
    710
        pub fn with_sort_information(mut self, sort_information: Vec<LexOrdering>) -> Self {
    210
    710
            self.sort_information = sort_information;
    211
    710
    212
    710
            // We need to update equivalence properties when updating sort information.
    213
    710
            let eq_properties = EquivalenceProperties::new_with_orderings(
    214
    710
                self.schema(),
    215
    710
                &self.sort_information,
    216
    710
            );
    217
    710
            self.cache = self.cache.with_eq_properties(eq_properties);
    218
    710
            self
    219
    710
        }
    220
    221
    0
        pub fn original_schema(&self) -> SchemaRef {
    222
    0
            Arc::clone(&self.schema)
    223
    0
        }
    224
    225
        /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
    226
    1.17k
        fn compute_properties(
    227
    1.17k
            schema: SchemaRef,
    228
    1.17k
            orderings: &[LexOrdering],
    229
    1.17k
            partitions: &[Vec<RecordBatch>],
    230
    1.17k
        ) -> PlanProperties {
    231
    1.17k
            let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings);
    232
    1.17k
            PlanProperties::new(
    233
    1.17k
                eq_properties,                                       // Equivalence Properties
    234
    1.17k
                Partitioning::UnknownPartitioning(partitions.len()), // Output Partitioning
    235
    1.17k
                ExecutionMode::Bounded,                              // Execution Mode
    236
    1.17k
            )
    237
    1.17k
        }
    238
    }
    239
    240
    /// Iterator over batches
    241
    pub struct MemoryStream {
    242
        /// Vector of record batches
    243
        data: Vec<RecordBatch>,
    244
        /// Optional memory reservation bound to the data, freed on drop
    245
        reservation: Option<MemoryReservation>,
    246
        /// Schema representing the data
    247
        schema: SchemaRef,
    248
        /// Optional projection for which columns to load
    249
        projection: Option<Vec<usize>>,
    250
        /// Index into the data
    251
        index: usize,
    252
    }
    253
    254
    impl MemoryStream {
    255
        /// Create an iterator for a vector of record batches
    256
    2.59k
        pub fn try_new(
    257
    2.59k
            data: Vec<RecordBatch>,
    258
    2.59k
            schema: SchemaRef,
    259
    2.59k
            projection: Option<Vec<usize>>,
    260
    2.59k
        ) -> Result<Self> {
    261
    2.59k
            Ok(Self {
    262
    2.59k
                data,
    263
    2.59k
                reservation: None,
    264
    2.59k
                schema,
    265
    2.59k
                projection,
    266
    2.59k
                index: 0,
    267
    2.59k
            })
    268
    2.59k
        }
    269
    270
        /// Set the memory reservation for the data
    271
    1
        pub(super) fn with_reservation(mut self, reservation: MemoryReservation) -> Self {
    272
    1
            self.reservation = Some(reservation);
    273
    1
            self
    274
    1
        }
    275
    }
    276
    277
    impl Stream for MemoryStream {
    278
        type Item = Result<RecordBatch>;
    279
    280
    35.3k
        fn poll_next(
    281
    35.3k
            mut self: std::pin::Pin<&mut Self>,
    282
    35.3k
            _: &mut Context<'_>,
    283
    35.3k
        ) -> Poll<Option<Self::Item>> {
    284
    35.3k
            Poll::Ready(if self.index < self.data.len() {
    285
    32.7k
                self.index += 1;
    286
    32.7k
                let batch = &self.data[self.index - 1];
    287
    288
                // return just the columns requested
    289
    32.7k
                let batch = match self.projection.as_ref() {
    290
    0
                    Some(columns) => batch.project(columns)?,
    291
    32.7k
                    None => batch.clone(),
    292
                };
    293
    294
    32.7k
                Some(Ok(batch))
    295
            } else {
    296
    2.58k
                None
    297
            })
    298
    35.3k
        }
    299
    300
    0
        fn size_hint(&self) -> (usize, Option<usize>) {
    301
    0
            (self.data.len(), Some(self.data.len()))
    302
    0
        }
    303
    }
    304
    305
    impl RecordBatchStream for MemoryStream {
    306
        /// Get the schema
    307
    350
        fn schema(&self) -> SchemaRef {
    308
    350
            Arc::clone(&self.schema)
    309
    350
        }
    310
    }
    311
    312
    #[cfg(test)]
    313
    mod tests {
    314
        use std::sync::Arc;
    315
    316
        use crate::memory::MemoryExec;
    317
        use crate::ExecutionPlan;
    318
    319
        use arrow_schema::{DataType, Field, Schema, SortOptions};
    320
        use datafusion_physical_expr::expressions::col;
    321
        use datafusion_physical_expr::PhysicalSortExpr;
    322
    323
        #[test]
    324
    1
        fn test_memory_order_eq() -> datafusion_common::Result<()> {
    325
    1
            let schema = Arc::new(Schema::new(vec![
    326
    1
                Field::new("a", DataType::Int64, false),
    327
    1
                Field::new("b", DataType::Int64, false),
    328
    1
                Field::new("c", DataType::Int64, false),
    329
    1
            ]));
    330
    1
            let sort1 = vec![
    331
    1
                PhysicalSortExpr {
    332
    1
                    expr: col("a", &schema)?0,
    333
    1
                    options: SortOptions::default(),
    334
    1
                },
    335
    1
                PhysicalSortExpr {
    336
    1
                    expr: col("b", &schema)?0,
    337
    1
                    options: SortOptions::default(),
    338
                },
    339
            ];
    340
    1
            let sort2 = vec![PhysicalSortExpr {
    341
    1
                expr: col("c", &schema)?0,
    342
    1
                options: SortOptions::default(),
    343
    1
            }];
    344
    1
            let mut expected_output_order = vec![];
    345
    1
            expected_output_order.extend(sort1.clone());
    346
    1
            expected_output_order.extend(sort2.clone());
    347
    1
    348
    1
            let sort_information = vec![sort1.clone(), sort2.clone()];
    349
    1
            let mem_exec = MemoryExec::try_new(&[vec![]], schema, None)?0
    350
    1
                .with_sort_information(sort_information);
    351
    1
    352
    1
            assert_eq!(
    353
    1
                mem_exec.properties().output_ordering().unwrap(),
    354
    1
                expected_output_order
    355
    1
            );
    356
    1
            let eq_properties = mem_exec.properties().equivalence_properties();
    357
    1
            assert!(eq_properties.oeq_class().contains(&sort1));
    358
    1
            assert!(eq_properties.oeq_class().contains(&sort2));
    359
    1
            Ok(())
    360
    1
        }
    361
    }