datafusion_physical_plan/
recursive_query.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 recursive query plan
19
20use std::any::Any;
21use std::sync::Arc;
22use std::task::{Context, Poll};
23
24use super::work_table::{ReservedBatches, WorkTable};
25use crate::aggregates::group_values::{GroupValues, new_group_values};
26use crate::aggregates::order::GroupOrdering;
27use crate::common::project_plan_to_schema;
28use crate::execution_plan::{Boundedness, EmissionType, reset_plan_states};
29use crate::metrics::{
30    BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput,
31};
32use crate::{
33    ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties,
34    RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream,
35};
36use arrow::array::{BooleanArray, BooleanBuilder};
37use arrow::compute::filter_record_batch;
38use arrow::datatypes::SchemaRef;
39use arrow::record_batch::RecordBatch;
40use datafusion_common::tree_node::TreeNodeRecursion;
41use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
42use datafusion_common::{
43    Result, exec_datafusion_err, internal_datafusion_err, not_impl_err,
44};
45use datafusion_execution::TaskContext;
46use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation};
47use datafusion_physical_expr::PhysicalExpr;
48use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
49
50use futures::{Stream, StreamExt, ready};
51
52/// Recursive query execution plan.
53///
54/// This plan has two components: a base part (the static term) and
55/// a dynamic part (the recursive term). The execution will start from
56/// the base, and as long as the previous iteration produced at least
57/// a single new row (taking care of the distinction) the recursive
58/// part will be continuously executed.
59///
60/// Before each execution of the dynamic part, the rows from the previous
61/// iteration will be available in a "working table" (not a real table,
62/// can be only accessed using a continuance operation).
63///
64/// Note that there won't be any limit or checks applied to detect
65/// an infinite recursion, so it is up to the planner to ensure that
66/// it won't happen.
67#[derive(Debug, Clone)]
68pub struct RecursiveQueryExec {
69    /// Name of the query handler
70    name: String,
71    /// The working table of cte
72    work_table: Arc<WorkTable>,
73    /// The base part (static term)
74    static_term: Arc<dyn ExecutionPlan>,
75    /// The dynamic part (recursive term)
76    recursive_term: Arc<dyn ExecutionPlan>,
77    /// Distinction
78    is_distinct: bool,
79    /// Execution metrics
80    metrics: ExecutionPlanMetricsSet,
81    /// Cache holding plan properties like equivalences, output partitioning etc.
82    cache: Arc<PlanProperties>,
83}
84
85impl RecursiveQueryExec {
86    /// Create a new RecursiveQueryExec
87    pub fn try_new(
88        name: String,
89        output_schema: SchemaRef,
90        static_term: Arc<dyn ExecutionPlan>,
91        recursive_term: Arc<dyn ExecutionPlan>,
92        is_distinct: bool,
93    ) -> Result<Self> {
94        // Each recursive query needs its own work table
95        let work_table = Arc::new(WorkTable::new(name.clone()));
96        // Use the same work table for both the WorkTableExec and the recursive term
97        let static_term = project_plan_to_schema(static_term, &output_schema)?;
98        let recursive_term = assign_work_table(recursive_term, &work_table)?;
99        let recursive_term = project_plan_to_schema(recursive_term, &output_schema)?;
100        let cache = Self::compute_properties(output_schema);
101        Ok(RecursiveQueryExec {
102            name,
103            static_term,
104            recursive_term,
105            is_distinct,
106            work_table,
107            metrics: ExecutionPlanMetricsSet::new(),
108            cache: Arc::new(cache),
109        })
110    }
111
112    /// Ref to name
113    pub fn name(&self) -> &str {
114        &self.name
115    }
116
117    /// Ref to static term
118    pub fn static_term(&self) -> &Arc<dyn ExecutionPlan> {
119        &self.static_term
120    }
121
122    /// Ref to recursive term
123    pub fn recursive_term(&self) -> &Arc<dyn ExecutionPlan> {
124        &self.recursive_term
125    }
126
127    /// is distinct
128    pub fn is_distinct(&self) -> bool {
129        self.is_distinct
130    }
131
132    /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
133    fn compute_properties(schema: SchemaRef) -> PlanProperties {
134        let eq_properties = EquivalenceProperties::new(schema);
135
136        PlanProperties::new(
137            eq_properties,
138            Partitioning::UnknownPartitioning(1),
139            EmissionType::Incremental,
140            Boundedness::Bounded,
141        )
142    }
143}
144
145impl ExecutionPlan for RecursiveQueryExec {
146    fn name(&self) -> &'static str {
147        "RecursiveQueryExec"
148    }
149
150    fn properties(&self) -> &Arc<PlanProperties> {
151        &self.cache
152    }
153
154    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
155        vec![&self.static_term, &self.recursive_term]
156    }
157
158    fn apply_expressions(
159        &self,
160        _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
161    ) -> Result<TreeNodeRecursion> {
162        Ok(TreeNodeRecursion::Continue)
163    }
164
165    // TODO: control these hints and see whether we can
166    // infer some from the child plans (static/recursive terms).
167    fn maintains_input_order(&self) -> Vec<bool> {
168        vec![false, false]
169    }
170
171    fn benefits_from_input_partitioning(&self) -> Vec<bool> {
172        vec![false, false]
173    }
174
175    fn required_input_distribution(&self) -> Vec<crate::Distribution> {
176        self.input_distribution_requirements().into_per_child()
177    }
178
179    fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements {
180        crate::InputDistributionRequirements::new(vec![
181            crate::Distribution::SinglePartition,
182            crate::Distribution::SinglePartition,
183        ])
184    }
185
186    fn replace_children(
187        self: Arc<Self>,
188        children: Vec<Arc<dyn ExecutionPlan>>,
189        _: ReplaceChildrenOptions,
190    ) -> Result<Arc<dyn ExecutionPlan>> {
191        RecursiveQueryExec::try_new(
192            self.name.clone(),
193            self.schema(),
194            Arc::clone(&children[0]),
195            Arc::clone(&children[1]),
196            self.is_distinct,
197        )
198        .map(|e| Arc::new(e) as _)
199    }
200
201    fn with_new_children(
202        self: Arc<Self>,
203        children: Vec<Arc<dyn ExecutionPlan>>,
204    ) -> Result<Arc<dyn ExecutionPlan>> {
205        self.replace_children(
206            children,
207            ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute),
208        )
209    }
210
211    fn execute(
212        &self,
213        partition: usize,
214        context: Arc<TaskContext>,
215    ) -> Result<SendableRecordBatchStream> {
216        // TODO: we might be able to handle multiple partitions in the future.
217        if partition != 0 {
218            return Err(internal_datafusion_err!(
219                "RecursiveQueryExec got an invalid partition {partition} (expected 0)"
220            ));
221        }
222
223        let static_stream = self.static_term.execute(partition, Arc::clone(&context))?;
224        let baseline_metrics = BaselineMetrics::new(&self.metrics, partition);
225        Ok(Box::pin(RecursiveQueryStream::new(
226            context,
227            Arc::clone(&self.work_table),
228            Arc::clone(&self.recursive_term),
229            static_stream,
230            self.is_distinct,
231            baseline_metrics,
232        )?))
233    }
234
235    fn metrics(&self) -> Option<MetricsSet> {
236        Some(self.metrics.clone_inner())
237    }
238}
239
240impl DisplayAs for RecursiveQueryExec {
241    fn fmt_as(
242        &self,
243        t: DisplayFormatType,
244        f: &mut std::fmt::Formatter,
245    ) -> std::fmt::Result {
246        match t {
247            DisplayFormatType::Default | DisplayFormatType::Verbose => {
248                write!(
249                    f,
250                    "RecursiveQueryExec: name={}, is_distinct={}",
251                    self.name, self.is_distinct
252                )
253            }
254            DisplayFormatType::TreeRender => {
255                // TODO: collect info
256                write!(f, "")
257            }
258        }
259    }
260}
261
262/// The actual logic of the recursive queries happens during the streaming
263/// process. A simplified version of the algorithm is the following:
264///
265/// buffer = []
266///
267/// while batch := static_stream.next():
268///    buffer.push(batch)
269///    yield buffer
270///
271/// while buffer.len() > 0:
272///    sender, receiver = Channel()
273///    register_continuation(handle_name, receiver)
274///    sender.send(buffer.drain())
275///    recursive_stream = recursive_term.execute()
276///    while batch := recursive_stream.next():
277///        buffer.append(batch)
278///        yield buffer
279struct RecursiveQueryStream {
280    /// The context to be used for managing handlers & executing new tasks
281    task_context: Arc<TaskContext>,
282    /// The working table state, representing the self referencing cte table
283    work_table: Arc<WorkTable>,
284    /// The dynamic part (recursive term) as is (without being executed)
285    recursive_term: Arc<dyn ExecutionPlan>,
286    /// The static part (static term) as a stream. If the processing of this
287    /// part is completed, then it will be None.
288    static_stream: Option<SendableRecordBatchStream>,
289    /// The dynamic part (recursive term) as a stream. If the processing of this
290    /// part has not started yet, or has been completed, then it will be None.
291    recursive_stream: Option<SendableRecordBatchStream>,
292    /// The schema of the output.
293    schema: SchemaRef,
294    /// In-memory buffer for storing a copy of the current results. Will be
295    /// cleared after each iteration.
296    buffer: Vec<RecordBatch>,
297    /// Tracks the memory used by the buffer
298    reservation: MemoryReservation,
299    /// If the distinct flag is set, then we use this hash table to remove duplicates from result and work tables
300    distinct_deduplicator: Option<DistinctDeduplicator>,
301    /// Metrics.
302    baseline_metrics: BaselineMetrics,
303}
304
305impl RecursiveQueryStream {
306    /// Create a new recursive query stream
307    fn new(
308        task_context: Arc<TaskContext>,
309        work_table: Arc<WorkTable>,
310        recursive_term: Arc<dyn ExecutionPlan>,
311        static_stream: SendableRecordBatchStream,
312        is_distinct: bool,
313        baseline_metrics: BaselineMetrics,
314    ) -> Result<Self> {
315        let schema = static_stream.schema();
316        let reservation =
317            MemoryConsumer::new("RecursiveQuery").register(task_context.memory_pool());
318        let distinct_deduplicator = is_distinct
319            .then(|| DistinctDeduplicator::new(Arc::clone(&schema), &task_context))
320            .transpose()?;
321        Ok(Self {
322            task_context,
323            work_table,
324            recursive_term,
325            static_stream: Some(static_stream),
326            recursive_stream: None,
327            schema,
328            buffer: vec![],
329            reservation,
330            distinct_deduplicator,
331            baseline_metrics,
332        })
333    }
334
335    /// Push a clone of the given batch to the in memory buffer, and then return
336    /// a poll with it.
337    fn push_batch(
338        mut self: std::pin::Pin<&mut Self>,
339        mut batch: RecordBatch,
340    ) -> Poll<Option<Result<RecordBatch>>> {
341        let baseline_metrics = self.baseline_metrics.clone();
342
343        if let Some(deduplicator) = &mut self.distinct_deduplicator {
344            let _timer_guard = baseline_metrics.elapsed_compute().timer();
345            batch = deduplicator.deduplicate(&batch)?;
346        }
347
348        if let Err(e) = self.reservation.try_grow(batch.get_array_memory_size()) {
349            return Poll::Ready(Some(Err(e)));
350        }
351        self.buffer.push(batch.clone());
352        (&batch).record_output(&baseline_metrics);
353        Poll::Ready(Some(Ok(batch)))
354    }
355
356    /// Start polling for the next iteration, will be called either after the static term
357    /// is completed or another term is completed. It will follow the algorithm above on
358    /// to check whether the recursion has ended.
359    fn poll_next_iteration(
360        mut self: std::pin::Pin<&mut Self>,
361        cx: &mut Context<'_>,
362    ) -> Poll<Option<Result<RecordBatch>>> {
363        let total_length = self
364            .buffer
365            .iter()
366            .fold(0, |acc, batch| acc + batch.num_rows());
367
368        if total_length == 0 {
369            return Poll::Ready(None);
370        }
371
372        // Update the work table with the current buffer
373        let reserved_batches = ReservedBatches::new(
374            std::mem::take(&mut self.buffer),
375            self.reservation.take(),
376        );
377        self.work_table.update(reserved_batches);
378
379        // We always execute (and re-execute iteratively) the first partition.
380        // Downstream plans should not expect any partitioning.
381        let partition = 0;
382
383        let recursive_plan = reset_plan_states(Arc::clone(&self.recursive_term))?;
384        self.recursive_stream =
385            Some(recursive_plan.execute(partition, Arc::clone(&self.task_context))?);
386        self.poll_next(cx)
387    }
388}
389
390fn assign_work_table(
391    plan: Arc<dyn ExecutionPlan>,
392    work_table: &Arc<WorkTable>,
393) -> Result<Arc<dyn ExecutionPlan>> {
394    let mut work_table_refs = 0;
395    plan.transform_down(|plan| {
396        if let Some(new_plan) =
397            plan.with_new_state(Arc::clone(work_table) as Arc<dyn Any + Send + Sync>)
398        {
399            if work_table_refs > 0 {
400                not_impl_err!(
401                    "Multiple recursive references to the same CTE are not supported"
402                )
403            } else {
404                work_table_refs += 1;
405                Ok(Transformed::yes(new_plan))
406            }
407        } else {
408            Ok(Transformed::no(plan))
409        }
410    })
411    .data()
412}
413
414impl Stream for RecursiveQueryStream {
415    type Item = Result<RecordBatch>;
416
417    fn poll_next(
418        mut self: std::pin::Pin<&mut Self>,
419        cx: &mut Context<'_>,
420    ) -> Poll<Option<Self::Item>> {
421        if let Some(static_stream) = &mut self.static_stream {
422            // While the static term's stream is available, we'll be forwarding the batches from it (also
423            // saving them for the initial iteration of the recursive term).
424            let batch_result = ready!(static_stream.poll_next_unpin(cx));
425            match &batch_result {
426                None => {
427                    // Once this is done, we can start running the setup for the recursive term.
428                    self.static_stream = None;
429                    self.poll_next_iteration(cx)
430                }
431                Some(Ok(batch)) => self.push_batch(batch.clone()),
432                _ => Poll::Ready(batch_result),
433            }
434        } else if let Some(recursive_stream) = &mut self.recursive_stream {
435            let batch_result = ready!(recursive_stream.poll_next_unpin(cx));
436            match batch_result {
437                None => {
438                    self.recursive_stream = None;
439                    self.poll_next_iteration(cx)
440                }
441                Some(Ok(batch)) => self.push_batch(batch),
442                _ => Poll::Ready(batch_result),
443            }
444        } else {
445            Poll::Ready(None)
446        }
447    }
448}
449
450impl RecordBatchStream for RecursiveQueryStream {
451    /// Get the schema
452    fn schema(&self) -> SchemaRef {
453        Arc::clone(&self.schema)
454    }
455}
456
457/// Deduplicator based on a hash table.
458struct DistinctDeduplicator {
459    /// Grouped rows used for distinct
460    group_values: Box<dyn GroupValues>,
461    reservation: MemoryReservation,
462    intern_output_buffer: Vec<usize>,
463}
464
465impl DistinctDeduplicator {
466    fn new(schema: SchemaRef, task_context: &TaskContext) -> Result<Self> {
467        let group_values = new_group_values(schema, &GroupOrdering::None)?;
468        let reservation = MemoryConsumer::new("RecursiveQueryHashTable")
469            .register(task_context.memory_pool());
470        Ok(Self {
471            group_values,
472            reservation,
473            intern_output_buffer: Vec::new(),
474        })
475    }
476
477    /// Remove duplicated rows from the given batch, keeping a state between batches.
478    ///
479    /// We use a hash table to allocate new group ids for the new rows.
480    /// [`GroupValues`] allocate increasing group ids.
481    /// Hence, if groups (i.e., rows) are new, then they have ids >= length before interning, we keep them.
482    /// We also detect duplicates by enforcing that group ids are increasing.
483    fn deduplicate(&mut self, batch: &RecordBatch) -> Result<RecordBatch> {
484        let size_before = self.group_values.len();
485        let additional = batch.num_rows();
486        self.intern_output_buffer
487            .try_reserve(additional)
488            .map_err(|e| {
489                exec_datafusion_err!(
490                    "failed to reserve {additional} recursive query group ids: {e}"
491                )
492            })?;
493        self.group_values
494            .intern(batch.columns(), &mut self.intern_output_buffer)?;
495        let mask = new_groups_mask(&self.intern_output_buffer, size_before);
496        self.intern_output_buffer.clear();
497        // We update the reservation to reflect the new size of the hash table.
498        self.reservation.try_resize(self.group_values.size())?;
499        Ok(filter_record_batch(batch, &mask)?)
500    }
501}
502
503/// Return a mask, each element being true if, and only if, the element is greater than all previous elements and greater or equal than the provided max_already_seen_group_id
504fn new_groups_mask(
505    values: &[usize],
506    mut max_already_seen_group_id: usize,
507) -> BooleanArray {
508    let mut output = BooleanBuilder::with_capacity(values.len());
509    for value in values {
510        if *value >= max_already_seen_group_id {
511            output.append_value(true);
512            max_already_seen_group_id = *value + 1; // We want to be increasing
513        } else {
514            output.append_value(false);
515        }
516    }
517    output.finish()
518}
519
520#[cfg(test)]
521mod tests {
522    use super::*;
523    use crate::empty::EmptyExec;
524    use crate::projection::ProjectionExec;
525
526    use arrow::datatypes::{DataType, Field, Schema};
527
528    fn empty_exec(fields: Vec<Field>) -> Arc<dyn ExecutionPlan> {
529        Arc::new(EmptyExec::new(Arc::new(Schema::new(fields))))
530    }
531
532    #[test]
533    fn recursive_query_exec_projects_recursive_term_to_reconciled_schema() -> Result<()> {
534        let static_term = empty_exec(vec![Field::new("value", DataType::Int32, false)]);
535        let recursive_term =
536            empty_exec(vec![Field::new("value + Int32(1)", DataType::Int32, false)]);
537
538        let exec = RecursiveQueryExec::try_new(
539            "numbers".to_string(),
540            static_term.schema(),
541            Arc::clone(&static_term),
542            Arc::clone(&recursive_term),
543            false,
544        )?;
545
546        assert_eq!(exec.schema(), static_term.schema());
547        let projection = exec
548            .recursive_term()
549            .downcast_ref::<ProjectionExec>()
550            .expect("recursive term should be aligned with ProjectionExec");
551        assert!(Arc::ptr_eq(projection.input(), &recursive_term));
552        assert!(!projection.schema().field(0).is_nullable());
553        assert_eq!(projection.expr()[0].alias, "value");
554        Ok(())
555    }
556
557    #[test]
558    fn recursive_query_exec_reconciles_nullability() -> Result<()> {
559        let static_term = empty_exec(vec![Field::new("value", DataType::Int32, false)]);
560        let recursive_term =
561            empty_exec(vec![Field::new("value + Int32(1)", DataType::Int32, true)]);
562        let output_schema = Arc::new(Schema::new(vec![Field::new(
563            "value",
564            DataType::Int32,
565            true,
566        )]));
567
568        let exec = RecursiveQueryExec::try_new(
569            "numbers".to_string(),
570            Arc::clone(&output_schema),
571            static_term,
572            recursive_term,
573            false,
574        )?;
575
576        assert!(exec.schema().field(0).is_nullable());
577        assert!(exec.static_term().schema().field(0).is_nullable());
578        assert!(exec.recursive_term().schema().field(0).is_nullable());
579        Ok(())
580    }
581}
datafusion_physical_plan/recursive_query.rs

datafusion_physical_plan/
recursive_query.rs
