datafusion_physical_plan/
work_table.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//! Defines the work table query plan
19
20use std::any::Any;
21use std::sync::{Arc, Mutex};
22
23use crate::coop::cooperative;
24use crate::execution_plan::{Boundedness, EmissionType, SchedulingType};
25use crate::memory::MemoryStream;
26use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet};
27use crate::{
28    ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
29    ReplaceChildrenOptions, SendableRecordBatchStream, Statistics,
30};
31
32use crate::statistics::StatisticsArgs;
33use arrow::datatypes::SchemaRef;
34use arrow::record_batch::RecordBatch;
35use datafusion_common::tree_node::TreeNodeRecursion;
36use datafusion_common::{Result, assert_eq_or_internal_err, internal_datafusion_err};
37use datafusion_execution::TaskContext;
38use datafusion_execution::memory_pool::MemoryReservation;
39use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr};
40
41/// A vector of record batches with a memory reservation.
42#[derive(Debug)]
43pub(super) struct ReservedBatches {
44    batches: Vec<RecordBatch>,
45    reservation: MemoryReservation,
46}
47
48impl ReservedBatches {
49    pub(super) fn new(batches: Vec<RecordBatch>, reservation: MemoryReservation) -> Self {
50        ReservedBatches {
51            batches,
52            reservation,
53        }
54    }
55}
56
57/// The name is from PostgreSQL's terminology.
58/// See <https://wiki.postgresql.org/wiki/CTEReadme#How_Recursion_Works>
59/// This table serves as a mirror or buffer between each iteration of a recursive query.
60#[derive(Debug)]
61pub struct WorkTable {
62    batches: Mutex<Option<ReservedBatches>>,
63    name: String,
64}
65
66impl WorkTable {
67    /// Create a new work table.
68    pub(super) fn new(name: String) -> Self {
69        Self {
70            batches: Mutex::new(None),
71            name,
72        }
73    }
74
75    /// Take the previously written batches from the work table.
76    /// This will be called by the [`WorkTableExec`] when it is executed.
77    fn take(&self) -> Result<ReservedBatches> {
78        self.batches
79            .lock()
80            .unwrap()
81            .take()
82            .ok_or_else(|| internal_datafusion_err!("Unexpected empty work table"))
83    }
84
85    /// Update the results of a recursive query iteration to the work table.
86    pub(super) fn update(&self, batches: ReservedBatches) {
87        self.batches.lock().unwrap().replace(batches);
88    }
89}
90
91/// A temporary "working table" operation where the input data will be
92/// taken from the named handle during the execution and will be re-published
93/// as is (kind of like a mirror).
94///
95/// Most notably used in the implementation of recursive queries where the
96/// underlying relation does not exist yet but the data will come as the previous
97/// term is evaluated. This table will be used such that the recursive plan
98/// will register a receiver in the task context and this plan will use that
99/// receiver to get the data and stream it back up so that the batches are available
100/// in the next iteration.
101#[derive(Clone, Debug)]
102pub struct WorkTableExec {
103    /// Name of the relation handler
104    name: String,
105    /// The schema of the stream
106    schema: SchemaRef,
107    /// Projection to apply to build the output stream from the recursion state
108    projection: Option<Vec<usize>>,
109    /// The work table
110    work_table: Arc<WorkTable>,
111    /// Execution metrics
112    metrics: ExecutionPlanMetricsSet,
113    /// Cache holding plan properties like equivalences, output partitioning etc.
114    cache: Arc<PlanProperties>,
115}
116
117impl WorkTableExec {
118    /// Create a new execution plan for a worktable exec.
119    pub fn new(
120        name: String,
121        mut schema: SchemaRef,
122        projection: Option<Vec<usize>>,
123    ) -> Result<Self> {
124        if let Some(projection) = &projection {
125            schema = Arc::new(schema.project(projection)?);
126        }
127        let cache = Self::compute_properties(Arc::clone(&schema));
128        Ok(Self {
129            name: name.clone(),
130            schema,
131            projection,
132            work_table: Arc::new(WorkTable::new(name)),
133            metrics: ExecutionPlanMetricsSet::new(),
134            cache: Arc::new(cache),
135        })
136    }
137
138    /// Ref to name
139    pub fn name(&self) -> &str {
140        &self.name
141    }
142
143    /// Arc clone of ref to schema
144    pub fn schema(&self) -> SchemaRef {
145        Arc::clone(&self.schema)
146    }
147
148    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
149    fn compute_properties(schema: SchemaRef) -> PlanProperties {
150        PlanProperties::new(
151            EquivalenceProperties::new(schema),
152            Partitioning::UnknownPartitioning(1),
153            EmissionType::Incremental,
154            Boundedness::Bounded,
155        )
156        .with_scheduling_type(SchedulingType::Cooperative)
157    }
158}
159
160impl DisplayAs for WorkTableExec {
161    fn fmt_as(
162        &self,
163        t: DisplayFormatType,
164        f: &mut std::fmt::Formatter,
165    ) -> std::fmt::Result {
166        match t {
167            DisplayFormatType::Default | DisplayFormatType::Verbose => {
168                write!(f, "WorkTableExec: name={}", self.name)
169            }
170            DisplayFormatType::TreeRender => {
171                write!(f, "name={}", self.name)
172            }
173        }
174    }
175}
176
177impl ExecutionPlan for WorkTableExec {
178    fn name(&self) -> &'static str {
179        "WorkTableExec"
180    }
181
182    fn properties(&self) -> &Arc<PlanProperties> {
183        &self.cache
184    }
185
186    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
187        vec![]
188    }
189
190    fn apply_expressions(
191        &self,
192        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
193    ) -> Result<TreeNodeRecursion> {
194        Ok(TreeNodeRecursion::Continue)
195    }
196
197    fn replace_children(
198        self: Arc<Self>,
199        _: Vec<Arc<dyn ExecutionPlan>>,
200        _: ReplaceChildrenOptions,
201    ) -> Result<Arc<dyn ExecutionPlan>> {
202        Ok(Arc::clone(&self) as Arc<dyn ExecutionPlan>)
203    }
204
205    fn with_new_children(
206        self: Arc<Self>,
207        children: Vec<Arc<dyn ExecutionPlan>>,
208    ) -> Result<Arc<dyn ExecutionPlan>> {
209        self.replace_children(
210            children,
211            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
212        )
213    }
214
215    /// Stream the batches that were written to the work table.
216    fn execute(
217        &self,
218        partition: usize,
219        _context: Arc<TaskContext>,
220    ) -> Result<SendableRecordBatchStream> {
221        // WorkTable streams must be the plan base.
222        assert_eq_or_internal_err!(
223            partition,
224            0,
225            "WorkTableExec got an invalid partition {partition} (expected 0)"
226        );
227        let ReservedBatches {
228            mut batches,
229            reservation,
230        } = self.work_table.take()?;
231        if let Some(projection) = &self.projection {
232            // We apply the projection
233            // TODO: it would be better to apply it as soon as possible and not only here
234            // TODO: an aggressive projection makes the memory reservation smaller, even if we do not edit it
235            batches = batches
236                .into_iter()
237                .map(|b| b.project(projection))
238                .collect::<Result<Vec<_>, _>>()?;
239        }
240
241        let stream = MemoryStream::try_new(batches, Arc::clone(&self.schema), None)?
242            .with_reservation(reservation);
243        Ok(Box::pin(cooperative(stream)))
244    }
245
246    fn metrics(&self) -> Option<MetricsSet> {
247        Some(self.metrics.clone_inner())
248    }
249
250    fn statistics_from_inputs(
251        &self,
252        _input_stats: &[Arc<Statistics>],
253        _args: &StatisticsArgs,
254    ) -> Result<Arc<Statistics>> {
255        Ok(Arc::new(Statistics::new_unknown(&self.schema())))
256    }
257
258    /// Injects run-time state into this `WorkTableExec`.
259    ///
260    /// The only state this node currently understands is an [`Arc<WorkTable>`].
261    /// If `state` can be down-cast to that type, a new `WorkTableExec` backed
262    /// by the provided work table is returned.  Otherwise `None` is returned
263    /// so that callers can attempt to propagate the state further down the
264    /// execution plan tree.
265    fn with_new_state(
266        &self,
267        state: Arc<dyn Any + Send + Sync>,
268    ) -> Option<Arc<dyn ExecutionPlan>> {
269        // Down-cast to the expected state type; propagate `None` on failure
270        let work_table = state.downcast::<WorkTable>().ok()?;
271
272        if work_table.name != self.name {
273            return None; // Different table
274        }
275
276        Some(Arc::new(Self {
277            name: self.name.clone(),
278            schema: Arc::clone(&self.schema),
279            projection: self.projection.clone(),
280            metrics: ExecutionPlanMetricsSet::new(),
281            work_table,
282            cache: Arc::clone(&self.cache),
283        }))
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290    use arrow::array::{ArrayRef, Int16Array, Int32Array, Int64Array};
291    use arrow_schema::{DataType, Field, Schema};
292    use datafusion_execution::memory_pool::{MemoryConsumer, UnboundedMemoryPool};
293    use futures::StreamExt;
294
295    #[test]
296    fn test_work_table() {
297        let work_table = WorkTable::new("test".into());
298        // Can't take from empty work_table
299        assert!(work_table.take().is_err());
300
301        let pool = Arc::new(UnboundedMemoryPool::default()) as _;
302        let reservation = MemoryConsumer::new("test_work_table").register(&pool);
303
304        // Update batch to work_table
305        let array: ArrayRef = Arc::new((0..5).collect::<Int32Array>());
306        let batch = RecordBatch::try_from_iter(vec![("col", array)]).unwrap();
307        reservation.try_grow(100).unwrap();
308        work_table.update(ReservedBatches::new(vec![batch.clone()], reservation));
309        // Take from work_table
310        let reserved_batches = work_table.take().unwrap();
311        assert_eq!(reserved_batches.batches, vec![batch.clone()]);
312
313        // Consume the batch by the MemoryStream
314        let memory_stream =
315            MemoryStream::try_new(reserved_batches.batches, batch.schema(), None)
316                .unwrap()
317                .with_reservation(reserved_batches.reservation);
318
319        // Should still be reserved
320        assert_eq!(pool.reserved(), 100);
321
322        // The reservation should be freed after drop the memory_stream
323        drop(memory_stream);
324        assert_eq!(pool.reserved(), 0);
325    }
326
327    #[tokio::test]
328    async fn test_work_table_exec() {
329        let schema = Arc::new(Schema::new(vec![
330            Field::new("a", DataType::Int64, false),
331            Field::new("b", DataType::Int32, false),
332            Field::new("c", DataType::Int16, false),
333        ]));
334        let work_table_exec =
335            WorkTableExec::new("wt".into(), Arc::clone(&schema), Some(vec![2, 1]))
336                .unwrap();
337
338        // We inject the work table
339        let work_table = Arc::new(WorkTable::new("wt".into()));
340        let work_table_exec = work_table_exec
341            .with_new_state(Arc::clone(&work_table) as _)
342            .unwrap();
343
344        // We update the work table
345        let pool = Arc::new(UnboundedMemoryPool::default()) as _;
346        let reservation = MemoryConsumer::new("test_work_table").register(&pool);
347        let batch = RecordBatch::try_new(
348            Arc::clone(&schema),
349            vec![
350                Arc::new(Int64Array::from(vec![1, 2, 3, 4, 5])),
351                Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])),
352                Arc::new(Int16Array::from(vec![1, 2, 3, 4, 5])),
353            ],
354        )
355        .unwrap();
356        work_table.update(ReservedBatches::new(vec![batch], reservation));
357
358        // We get back the batch from the work table
359        let returned_batch = work_table_exec
360            .execute(0, Arc::new(TaskContext::default()))
361            .unwrap()
362            .next()
363            .await
364            .unwrap()
365            .unwrap();
366        assert_eq!(
367            returned_batch,
368            RecordBatch::try_from_iter(vec![
369                ("c", Arc::new(Int16Array::from(vec![1, 2, 3, 4, 5])) as _),
370                ("b", Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])) as _),
371            ])
372            .unwrap()
373        );
374    }
375}
datafusion_physical_plan/work_table.rs

datafusion_physical_plan/
work_table.rs
