uqa_execution/
set_operation.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Byte-bounded physical SQL set operations.
8
9use std::cmp::Ordering;
10use std::sync::Arc;
11
12use uqa_core::Value;
13use uqa_sql::ast::{ColumnType, SetOpKind};
14use uqa_sql::expr::RowLookup;
15
16#[cfg(test)]
17use uqa_sql::ResultRow;
18
19use crate::batch::DEFAULT_BATCH_SIZE;
20use crate::{
21    BackwardScanSupport, Batch, ExecError, ExecResult, ExpressionEvaluator, ExternalSort,
22    PhysicalOperator, PhysicalRow, PhysicalScanDirection, RowProjectionValue, RowSchema,
23    ScalarExpr, SortKey,
24};
25
26struct ColumnEvaluator;
27
28impl ExpressionEvaluator for ColumnEvaluator {
29    fn evaluate(&self, expression: &ScalarExpr, row: &dyn RowLookup) -> ExecResult<Value> {
30        let ScalarExpr::Column(column) = expression else {
31            return Err(ExecError::Other(
32                "set-operation sort key must be a column".into(),
33            ));
34        };
35        Ok(row.column(column).cloned().unwrap_or(Value::Null))
36    }
37}
38
39fn set_operation_types(left: &RowSchema, right: &RowSchema) -> ExecResult<Vec<Option<ColumnType>>> {
40    if left.len() != right.len() {
41        return Err(ExecError::Other(format!(
42            "set-operation inputs have different widths: {} and {}",
43            left.len(),
44            right.len()
45        )));
46    }
47    left.column_types()
48        .iter()
49        .zip(right.column_types())
50        .map(|(left, right)| match (left, right) {
51            (None, None) => Ok(None),
52            (Some(ty), None) | (None, Some(ty)) => Ok(Some(ty.clone())),
53            (Some(left), Some(right)) => uqa_execution_common_type(left, right).map(Some),
54        })
55        .collect()
56}
57
58fn uqa_execution_common_type(left: &ColumnType, right: &ColumnType) -> ExecResult<ColumnType> {
59    crate::common_type(left, right).map_err(ExecError::from)
60}
61
62fn coerce_set_value(
63    value: Value,
64    source_type: Option<&ColumnType>,
65    target_type: &ColumnType,
66) -> ExecResult<Value> {
67    let cast_target = match target_type {
68        ColumnType::Domain { base, .. } => base.as_ref(),
69        target => target,
70    };
71    let source_name = source_type.map(ColumnType::sql_name);
72    uqa_sql::expr::cast_value_from(&value, &cast_target.sql_name(), source_name.as_deref())
73        .map_err(ExecError::from)
74}
75
76struct AlignSchema<'a> {
77    child: Box<dyn PhysicalOperator + 'a>,
78    schema: RowSchema,
79    coercions: Vec<Option<ColumnType>>,
80}
81
82impl<'a> AlignSchema<'a> {
83    fn new(
84        child: Box<dyn PhysicalOperator + 'a>,
85        output: Vec<String>,
86        output_types: &[Option<ColumnType>],
87    ) -> ExecResult<Self> {
88        let source = child.schema().to_vec();
89        if source.len() != output.len() {
90            return Err(ExecError::Other(format!(
91                "set-operation inputs have different widths: {} and {}",
92                output.len(),
93                source.len()
94            )));
95        }
96        if output.len() != output_types.len() {
97            return Err(ExecError::Other(format!(
98                "set-operation output type width {} does not match input width {}",
99                output_types.len(),
100                output.len()
101            )));
102        }
103        let coercions = child
104            .row_schema()
105            .column_types()
106            .iter()
107            .zip(output_types)
108            .map(|(source, target)| {
109                target
110                    .as_ref()
111                    .filter(|target| source.as_ref() != Some(*target))
112                    .cloned()
113            })
114            .collect();
115        let schema = RowSchema::with_types(output, output_types.to_vec());
116        Ok(Self {
117            child,
118            schema,
119            coercions,
120        })
121    }
122}
123
124impl PhysicalOperator for AlignSchema<'_> {
125    fn row_schema(&self) -> &RowSchema {
126        &self.schema
127    }
128
129    fn backward_scan_support(&self) -> BackwardScanSupport {
130        self.child.backward_scan_support()
131    }
132
133    fn open(&mut self) -> ExecResult<()> {
134        self.child.open()
135    }
136
137    fn next(&mut self) -> ExecResult<Option<Batch>> {
138        let Some(batch) = self.child.next()? else {
139            return Ok(None);
140        };
141        self.align_batch(batch).map(Some)
142    }
143
144    fn next_direction(&mut self, direction: PhysicalScanDirection) -> ExecResult<Option<Batch>> {
145        let Some(batch) = self.child.next_direction(direction)? else {
146            return Ok(None);
147        };
148        self.align_batch(batch).map(Some)
149    }
150
151    fn rewind(&mut self) -> ExecResult<()> {
152        self.child.rewind()
153    }
154
155    fn close(&mut self) -> ExecResult<()> {
156        self.child.close()
157    }
158}
159
160impl AlignSchema<'_> {
161    fn align_batch(&self, batch: Batch) -> ExecResult<Batch> {
162        let identity_layout = batch.schema.physical_width() == self.schema.physical_width()
163            && (0..batch.schema.len())
164                .all(|position| batch.schema.slot(position) == Some(position));
165        if self.coercions.iter().all(Option::is_none) && identity_layout {
166            let rows = batch
167                .rows
168                .into_iter()
169                .map(PhysicalRow::without_lock_origins)
170                .collect();
171            return Ok(Batch::from_physical_rows(self.schema.clone(), rows));
172        }
173        if self.coercions.iter().all(Option::is_none) {
174            let slots = (0..batch.schema.len())
175                .map(|position| {
176                    batch.schema.slot(position).ok_or_else(|| {
177                        ExecError::Other(format!(
178                            "set-operation input column {position} has no physical slot"
179                        ))
180                    })
181                })
182                .collect::<ExecResult<Vec<_>>>()?;
183            let rows = batch
184                .rows
185                .into_iter()
186                .map(|row| row.project_slots(&slots).without_lock_origins())
187                .collect();
188            return Ok(Batch::from_physical_rows(self.schema.clone(), rows));
189        }
190        let mut rows = Vec::with_capacity(batch.rows.len());
191        for row in batch.rows {
192            let view = batch.schema.view(&row);
193            let values = self
194                .coercions
195                .iter()
196                .enumerate()
197                .map(|(position, target)| match target {
198                    Some(target) => coerce_set_value(
199                        view.value_at(position).cloned().unwrap_or(Value::Null),
200                        batch.schema.column_type(position),
201                        target,
202                    )
203                    .map(RowProjectionValue::Owned),
204                    None => Ok(batch.schema.slot(position).map_or(
205                        RowProjectionValue::Owned(Value::Null),
206                        RowProjectionValue::InputSlot,
207                    )),
208                })
209                .collect::<ExecResult<Vec<_>>>()?;
210            rows.push(row.project_with_values(values).without_lock_origins());
211        }
212        Ok(Batch::from_physical_rows(self.schema.clone(), rows))
213    }
214}
215
216struct RowGroup {
217    row: PhysicalRow,
218    count: usize,
219}
220
221struct RowCursor<'a> {
222    operator: Box<dyn PhysicalOperator + 'a>,
223    batch: std::vec::IntoIter<PhysicalRow>,
224    lookahead: Option<PhysicalRow>,
225    exhausted: bool,
226}
227
228impl<'a> RowCursor<'a> {
229    fn new(operator: Box<dyn PhysicalOperator + 'a>) -> Self {
230        Self {
231            operator,
232            batch: Vec::new().into_iter(),
233            lookahead: None,
234            exhausted: false,
235        }
236    }
237
238    fn open(&mut self) -> ExecResult<()> {
239        self.batch = Vec::new().into_iter();
240        self.lookahead = None;
241        self.exhausted = false;
242        self.operator.open()
243    }
244
245    fn next_row(&mut self) -> ExecResult<Option<PhysicalRow>> {
246        loop {
247            if let Some(row) = self.batch.next() {
248                return Ok(Some(row));
249            }
250            let Some(batch) = self.operator.next()? else {
251                self.exhausted = true;
252                return Ok(None);
253            };
254            self.batch = batch.rows.into_iter();
255        }
256    }
257
258    fn backward_scan_support(&self) -> BackwardScanSupport {
259        self.operator.backward_scan_support()
260    }
261
262    fn next_direction_row(
263        &mut self,
264        direction: PhysicalScanDirection,
265    ) -> ExecResult<Option<PhysicalRow>> {
266        if self.batch.len() != 0 || self.lookahead.is_some() {
267            return Err(ExecError::Other(
268                "set-operation input cannot mix batched and directional pulls".into(),
269            ));
270        }
271        let Some(batch) = self.operator.next_direction(direction)? else {
272            self.exhausted = true;
273            return Ok(None);
274        };
275        let mut rows = batch.rows.into_iter();
276        let row = rows.next();
277        if rows.next().is_some() {
278            return Err(ExecError::Other(
279                "directional set-operation input returned more than one row".into(),
280            ));
281        }
282        self.exhausted = row.is_none();
283        Ok(row)
284    }
285
286    fn rewind(&mut self) -> ExecResult<()> {
287        self.batch = Vec::new().into_iter();
288        self.lookahead = None;
289        self.exhausted = false;
290        self.operator.rewind()
291    }
292
293    fn take_group(&mut self, schema: &RowSchema) -> ExecResult<Option<RowGroup>> {
294        let first = match self.lookahead.take() {
295            Some(row) => row,
296            None => match self.next_row()? {
297                Some(row) => row,
298                None => return Ok(None),
299            },
300        };
301        let mut count = 1_usize;
302        while let Some(row) = self.next_row()? {
303            if compare_rows(&first, &row, schema) == Ordering::Equal {
304                count = count
305                    .checked_add(1)
306                    .ok_or_else(|| ExecError::Other("set-operation group count overflow".into()))?;
307            } else {
308                self.lookahead = Some(row);
309                break;
310            }
311        }
312        Ok(Some(RowGroup { row: first, count }))
313    }
314
315    fn close(&mut self) -> ExecResult<()> {
316        self.batch = Vec::new().into_iter();
317        self.lookahead = None;
318        self.exhausted = true;
319        self.operator.close()
320    }
321}
322
323#[derive(Clone, Copy)]
324enum DirectionalAppendPosition {
325    BeforeFirst,
326    Left,
327    Right,
328    AfterLast,
329}
330
331fn compare_rows(left: &PhysicalRow, right: &PhysicalRow, schema: &RowSchema) -> Ordering {
332    let left = schema.view(left);
333    let right = schema.view(right);
334    let null = Value::Null;
335    for position in 0..schema.len() {
336        let ordering = left
337            .value_at(position)
338            .unwrap_or(&null)
339            .cmp(right.value_at(position).unwrap_or(&null));
340        if ordering != Ordering::Equal {
341            return ordering;
342        }
343    }
344    Ordering::Equal
345}
346
347/// `UNION` / `INTERSECT` / `EXCEPT` physical operator.
348///
349/// `UNION ALL` streams its children without sorting. Other forms sort both
350/// inputs through byte-bounded external runs and merge adjacent multiplicity
351/// counts, avoiding whole-input materialisation and quadratic row scans.
352pub struct ExternalSetOperation<'a> {
353    left: RowCursor<'a>,
354    right: RowCursor<'a>,
355    kind: SetOpKind,
356    all: bool,
357    schema: RowSchema,
358    left_group: Option<RowGroup>,
359    right_group: Option<RowGroup>,
360    pending_row: Option<PhysicalRow>,
361    pending_count: usize,
362    union_all_left_done: bool,
363    incremental_union_all: bool,
364    directional_position: DirectionalAppendPosition,
365}
366
367impl<'a> ExternalSetOperation<'a> {
368    pub fn new(
369        left: Box<dyn PhysicalOperator + 'a>,
370        right: Box<dyn PhysicalOperator + 'a>,
371        kind: SetOpKind,
372        all: bool,
373        work_mem_bytes: usize,
374    ) -> ExecResult<Self> {
375        let output_types = set_operation_types(left.row_schema(), right.row_schema())?;
376        Self::new_with_types(left, right, kind, all, output_types, work_mem_bytes)
377    }
378
379    pub fn new_with_types(
380        left: Box<dyn PhysicalOperator + 'a>,
381        right: Box<dyn PhysicalOperator + 'a>,
382        kind: SetOpKind,
383        all: bool,
384        output_types: Vec<Option<ColumnType>>,
385        work_mem_bytes: usize,
386    ) -> ExecResult<Self> {
387        Self::new_with_types_and_mode(left, right, kind, all, output_types, work_mem_bytes, false)
388    }
389
390    /// Construct a set operation whose ordinary forward pulls remain one-row incremental when a scroll materialization boundary wraps the complete operation.
391    pub fn new_directional_with_types(
392        left: Box<dyn PhysicalOperator + 'a>,
393        right: Box<dyn PhysicalOperator + 'a>,
394        kind: SetOpKind,
395        all: bool,
396        output_types: Vec<Option<ColumnType>>,
397        work_mem_bytes: usize,
398    ) -> ExecResult<Self> {
399        Self::new_with_types_and_mode(left, right, kind, all, output_types, work_mem_bytes, true)
400    }
401
402    fn new_with_types_and_mode(
403        left: Box<dyn PhysicalOperator + 'a>,
404        right: Box<dyn PhysicalOperator + 'a>,
405        kind: SetOpKind,
406        all: bool,
407        output_types: Vec<Option<ColumnType>>,
408        work_mem_bytes: usize,
409        incremental_union_all: bool,
410    ) -> ExecResult<Self> {
411        let output = left.schema().to_vec();
412        let left: Box<dyn PhysicalOperator + 'a> =
413            Box::new(AlignSchema::new(left, output.clone(), &output_types)?);
414        let right: Box<dyn PhysicalOperator + 'a> =
415            Box::new(AlignSchema::new(right, output.clone(), &output_types)?);
416        let (left, right) = if matches!((kind, all), (SetOpKind::Union, true)) {
417            (left, right)
418        } else {
419            let keys = output
420                .iter()
421                .map(|column| SortKey {
422                    expr: ScalarExpr::Column(column.clone()),
423                    descending: false,
424                    nulls_first: Some(true),
425                })
426                .collect::<Vec<_>>();
427            let evaluator = Arc::new(ColumnEvaluator);
428            // Both merge inputs are live concurrently. Split the configured
429            // budget so their encoded in-memory runs cannot each claim it.
430            let per_input = (work_mem_bytes / 2).max(1);
431            (
432                Box::new(ExternalSort::new(
433                    left,
434                    keys.clone(),
435                    evaluator.clone(),
436                    None,
437                    per_input,
438                )) as Box<dyn PhysicalOperator + 'a>,
439                Box::new(ExternalSort::new(right, keys, evaluator, None, per_input))
440                    as Box<dyn PhysicalOperator + 'a>,
441            )
442        };
443        Ok(Self {
444            left: RowCursor::new(left),
445            right: RowCursor::new(right),
446            kind,
447            all,
448            schema: RowSchema::with_types(output, output_types),
449            left_group: None,
450            right_group: None,
451            pending_row: None,
452            pending_count: 0,
453            union_all_left_done: false,
454            incremental_union_all,
455            directional_position: DirectionalAppendPosition::BeforeFirst,
456        })
457    }
458
459    fn next_union_all(&mut self) -> ExecResult<Option<Batch>> {
460        let mut rows = Vec::with_capacity(DEFAULT_BATCH_SIZE);
461        while rows.len() < DEFAULT_BATCH_SIZE {
462            let next = if self.union_all_left_done {
463                self.right.next_row()?
464            } else if let Some(row) = self.left.next_row()? {
465                Some(row)
466            } else {
467                self.union_all_left_done = true;
468                self.right.next_row()?
469            };
470            let Some(row) = next else {
471                break;
472            };
473            rows.push(row.without_lock_origins());
474        }
475        if rows.is_empty() {
476            Ok(None)
477        } else {
478            Ok(Some(Batch::from_physical_rows(self.schema.clone(), rows)))
479        }
480    }
481
482    fn next_union_all_row(&mut self) -> ExecResult<Option<Batch>> {
483        let row = if self.union_all_left_done {
484            self.right.next_row()?
485        } else if let Some(row) = self.left.next_row()? {
486            Some(row)
487        } else {
488            self.union_all_left_done = true;
489            self.right.next_row()?
490        };
491        Ok(row.map(|row| self.directional_row_batch(row)))
492    }
493
494    fn directional_row_batch(&self, row: PhysicalRow) -> Batch {
495        Batch::from_physical_rows(self.schema.clone(), vec![row.without_lock_origins()])
496    }
497
498    fn next_union_all_direction(
499        &mut self,
500        direction: PhysicalScanDirection,
501    ) -> ExecResult<Option<Batch>> {
502        let mut position = match (self.directional_position, direction) {
503            (DirectionalAppendPosition::BeforeFirst, PhysicalScanDirection::Backward)
504            | (DirectionalAppendPosition::AfterLast, PhysicalScanDirection::Forward) => {
505                return Ok(None)
506            }
507            (DirectionalAppendPosition::BeforeFirst, PhysicalScanDirection::Forward)
508            | (DirectionalAppendPosition::Left, _) => DirectionalAppendPosition::Left,
509            (DirectionalAppendPosition::AfterLast, PhysicalScanDirection::Backward)
510            | (DirectionalAppendPosition::Right, _) => DirectionalAppendPosition::Right,
511        };
512        loop {
513            let row = match position {
514                DirectionalAppendPosition::Left => self.left.next_direction_row(direction)?,
515                DirectionalAppendPosition::Right => self.right.next_direction_row(direction)?,
516                DirectionalAppendPosition::BeforeFirst | DirectionalAppendPosition::AfterLast => {
517                    unreachable!()
518                }
519            };
520            if let Some(row) = row {
521                self.directional_position = position;
522                return Ok(Some(self.directional_row_batch(row)));
523            }
524            position = match (position, direction) {
525                (DirectionalAppendPosition::Left, PhysicalScanDirection::Forward) => {
526                    DirectionalAppendPosition::Right
527                }
528                (DirectionalAppendPosition::Right, PhysicalScanDirection::Backward) => {
529                    DirectionalAppendPosition::Left
530                }
531                (DirectionalAppendPosition::Left, PhysicalScanDirection::Backward) => {
532                    self.directional_position = DirectionalAppendPosition::BeforeFirst;
533                    return Ok(None);
534                }
535                (DirectionalAppendPosition::Right, PhysicalScanDirection::Forward) => {
536                    self.directional_position = DirectionalAppendPosition::AfterLast;
537                    return Ok(None);
538                }
539                (
540                    DirectionalAppendPosition::BeforeFirst | DirectionalAppendPosition::AfterLast,
541                    _,
542                ) => unreachable!(),
543            };
544        }
545    }
546
547    fn load_groups(&mut self) -> ExecResult<()> {
548        if self.left_group.is_none() && !self.left.exhausted {
549            self.left_group = self.left.take_group(&self.schema)?;
550        }
551        if self.right_group.is_none() && !self.right.exhausted {
552            self.right_group = self.right.take_group(&self.schema)?;
553        }
554        Ok(())
555    }
556
557    fn take_left_group(&mut self) -> ExecResult<RowGroup> {
558        self.left_group
559            .take()
560            .ok_or_else(|| ExecError::Other("set-operation selected a missing left group".into()))
561    }
562
563    fn take_right_group(&mut self) -> ExecResult<RowGroup> {
564        self.right_group
565            .take()
566            .ok_or_else(|| ExecError::Other("set-operation selected a missing right group".into()))
567    }
568
569    fn choose_group(&mut self) -> ExecResult<Option<(PhysicalRow, usize)>> {
570        self.load_groups()?;
571        let ordering = match (&self.left_group, &self.right_group) {
572            (Some(left), Some(right)) => Some(compare_rows(&left.row, &right.row, &self.schema)),
573            (Some(_), None) => Some(Ordering::Less),
574            (None, Some(_)) => Some(Ordering::Greater),
575            (None, None) => None,
576        };
577        let Some(ordering) = ordering else {
578            return Ok(None);
579        };
580
581        let selected = match (self.kind, ordering) {
582            (SetOpKind::Union, Ordering::Less) => {
583                let group = self.take_left_group()?;
584                Some((group.row, 1))
585            }
586            (SetOpKind::Union, Ordering::Greater) => {
587                let group = self.take_right_group()?;
588                Some((group.row, 1))
589            }
590            (SetOpKind::Union, Ordering::Equal) => {
591                let group = self.take_left_group()?;
592                self.right_group = None;
593                Some((group.row, 1))
594            }
595            (SetOpKind::Intersect, Ordering::Less) => {
596                self.left_group = None;
597                None
598            }
599            (SetOpKind::Intersect | SetOpKind::Except, Ordering::Greater) => {
600                self.right_group = None;
601                None
602            }
603            (SetOpKind::Intersect, Ordering::Equal) => {
604                let left = self.take_left_group()?;
605                let right = self.take_right_group()?;
606                Some((
607                    left.row,
608                    if self.all {
609                        left.count.min(right.count)
610                    } else {
611                        1
612                    },
613                ))
614            }
615            (SetOpKind::Except, Ordering::Less) => {
616                let left = self.take_left_group()?;
617                Some((left.row, if self.all { left.count } else { 1 }))
618            }
619            (SetOpKind::Except, Ordering::Equal) => {
620                let left = self.take_left_group()?;
621                let right = self.take_right_group()?;
622                let count = if self.all {
623                    left.count.saturating_sub(right.count)
624                } else {
625                    0
626                };
627                Some((left.row, count))
628            }
629        };
630        Ok(selected.filter(|(_, count)| *count > 0))
631    }
632}
633
634impl PhysicalOperator for ExternalSetOperation<'_> {
635    fn row_schema(&self) -> &RowSchema {
636        &self.schema
637    }
638
639    fn backward_scan_support(&self) -> BackwardScanSupport {
640        if matches!((self.kind, self.all), (SetOpKind::Union, true))
641            && self.left.backward_scan_support() == BackwardScanSupport::Native
642            && self.right.backward_scan_support() == BackwardScanSupport::Native
643        {
644            BackwardScanSupport::Native
645        } else {
646            BackwardScanSupport::Unsupported
647        }
648    }
649
650    fn open(&mut self) -> ExecResult<()> {
651        self.left_group = None;
652        self.right_group = None;
653        self.pending_row = None;
654        self.pending_count = 0;
655        self.union_all_left_done = false;
656        self.directional_position = DirectionalAppendPosition::BeforeFirst;
657        self.left.open()?;
658        self.right.open()
659    }
660
661    fn next(&mut self) -> ExecResult<Option<Batch>> {
662        if matches!((self.kind, self.all), (SetOpKind::Union, true)) {
663            return if self.incremental_union_all {
664                self.next_union_all_row()
665            } else {
666                self.next_union_all()
667            };
668        }
669        let mut rows = Vec::with_capacity(DEFAULT_BATCH_SIZE);
670        while rows.len() < DEFAULT_BATCH_SIZE {
671            if self.pending_count > 0 {
672                let row = self.pending_row.as_ref().ok_or_else(|| {
673                    ExecError::Other(
674                        "set-operation has pending multiplicity without a pending row".into(),
675                    )
676                })?;
677                rows.push(row.clone());
678                self.pending_count -= 1;
679                if self.pending_count == 0 {
680                    self.pending_row = None;
681                }
682                continue;
683            }
684            match self.choose_group()? {
685                Some((row, count)) => {
686                    self.pending_row = Some(row.without_lock_origins());
687                    self.pending_count = count;
688                }
689                None if self.left.exhausted && self.right.exhausted => break,
690                None => {}
691            }
692        }
693        if rows.is_empty() {
694            Ok(None)
695        } else {
696            Ok(Some(Batch::from_physical_rows(self.schema.clone(), rows)))
697        }
698    }
699
700    fn next_direction(&mut self, direction: PhysicalScanDirection) -> ExecResult<Option<Batch>> {
701        if self.backward_scan_support() == BackwardScanSupport::Native {
702            return self.next_union_all_direction(direction);
703        }
704        if direction == PhysicalScanDirection::Forward {
705            return if matches!((self.kind, self.all), (SetOpKind::Union, true)) {
706                self.next_union_all_row()
707            } else {
708                self.next()
709            };
710        }
711        Err(ExecError::Other(
712            "set operation does not support backwards scanning".into(),
713        ))
714    }
715
716    fn rewind(&mut self) -> ExecResult<()> {
717        if self.backward_scan_support() != BackwardScanSupport::Native {
718            return Err(ExecError::Other(
719                "set operation does not support rewind".into(),
720            ));
721        }
722        self.left.rewind()?;
723        self.right.rewind()?;
724        self.directional_position = DirectionalAppendPosition::BeforeFirst;
725        Ok(())
726    }
727
728    fn close(&mut self) -> ExecResult<()> {
729        self.left_group = None;
730        self.right_group = None;
731        self.pending_row = None;
732        self.pending_count = 0;
733        self.directional_position = DirectionalAppendPosition::BeforeFirst;
734        let left = self.left.close();
735        let right = self.right.close();
736        crate::physical::with_cleanup(left, right, "close right set-operation input")
737    }
738}
739
740#[cfg(test)]
741mod tests {
742    use super::*;
743    use crate::physical::run_to_rows;
744    use crate::scan::TableScan;
745
746    fn row(value: i64) -> ResultRow {
747        [("v".into(), Value::Int(value))].into_iter().collect()
748    }
749
750    fn execute(kind: SetOpKind, all: bool, left: &[i64], right: &[i64]) -> Vec<i64> {
751        let left = TableScan::from_rows(vec!["v".into()], left.iter().copied().map(row).collect());
752        let right = TableScan::from_rows(
753            vec!["other".into()],
754            right
755                .iter()
756                .copied()
757                .map(|value| [("other".into(), Value::Int(value))].into_iter().collect())
758                .collect(),
759        );
760        let mut set =
761            ExternalSetOperation::new(Box::new(left), Box::new(right), kind, all, 1).unwrap();
762        run_to_rows(&mut set)
763            .unwrap()
764            .1
765            .into_iter()
766            .map(|row| match row.get("v") {
767                Some(Value::Int(value)) => *value,
768                value => panic!("unexpected set value: {value:?}"),
769            })
770            .collect()
771    }
772
773    #[test]
774    fn external_set_semantics_include_bag_multiplicity() {
775        assert_eq!(
776            execute(SetOpKind::Union, false, &[2, 1, 1], &[3, 2]),
777            vec![1, 2, 3]
778        );
779        assert_eq!(
780            execute(SetOpKind::Union, true, &[2, 1, 1], &[3, 2]),
781            vec![2, 1, 1, 3, 2]
782        );
783        assert_eq!(
784            execute(SetOpKind::Intersect, true, &[1, 1, 1, 2], &[1, 1, 3]),
785            vec![1, 1]
786        );
787        assert_eq!(
788            execute(SetOpKind::Except, true, &[1, 1, 1, 2], &[1, 1, 3]),
789            vec![1, 2]
790        );
791        assert_eq!(
792            execute(SetOpKind::Except, false, &[1, 1, 2], &[1, 3]),
793            vec![2]
794        );
795    }
796
797    #[test]
798    fn set_inputs_compact_schema_only_projections_before_alignment() {
799        let left = TableScan::from_rows(
800            vec!["v".into(), "hidden".into()],
801            vec![[
802                ("v".into(), Value::Int(1)),
803                ("hidden".into(), Value::Int(99)),
804            ]
805            .into_iter()
806            .collect()],
807        );
808        let left: Box<dyn PhysicalOperator> = Box::new(crate::ColumnSelection::with_positions(
809            Box::new(left),
810            vec![("v".into(), 0)],
811        ));
812        let right = TableScan::from_rows(vec!["v".into()], vec![row(2)]);
813        let mut set =
814            ExternalSetOperation::new(left, Box::new(right), SetOpKind::Union, true, 1).unwrap();
815        let rows = run_to_rows(&mut set).unwrap().1;
816        assert_eq!(rows, vec![row(1), row(2)]);
817    }
818
819    fn directional_value(batch: Option<Batch>) -> Option<i64> {
820        let batch = batch?;
821        let row = batch.rows.first()?;
822        match batch.schema.view(row).value_at(0) {
823            Some(Value::Int(value)) => Some(*value),
824            value => panic!("unexpected directional set value: {value:?}"),
825        }
826    }
827
828    #[test]
829    fn union_all_scans_children_in_reverse_order_across_the_boundary() {
830        let left: Box<dyn PhysicalOperator> = Box::new(crate::ScrollMaterialize::new(Box::new(
831            TableScan::from_rows(vec!["v".into()], vec![row(1), row(2)]),
832        )));
833        let right: Box<dyn PhysicalOperator> = Box::new(crate::ScrollMaterialize::new(Box::new(
834            TableScan::from_rows(vec!["v".into()], vec![row(3), row(4)]),
835        )));
836        let mut set = ExternalSetOperation::new(left, right, SetOpKind::Union, true, 1).unwrap();
837        assert_eq!(set.backward_scan_support(), BackwardScanSupport::Native);
838        set.open().unwrap();
839        for (direction, expected) in [
840            (PhysicalScanDirection::Forward, 1),
841            (PhysicalScanDirection::Forward, 2),
842            (PhysicalScanDirection::Forward, 3),
843            (PhysicalScanDirection::Backward, 2),
844            (PhysicalScanDirection::Forward, 3),
845        ] {
846            assert_eq!(
847                directional_value(set.next_direction(direction).unwrap()),
848                Some(expected)
849            );
850        }
851        set.rewind().unwrap();
852        assert_eq!(
853            directional_value(set.next_direction(PhysicalScanDirection::Forward).unwrap()),
854            Some(1)
855        );
856        set.close().unwrap();
857    }
858
859    #[test]
860    fn directional_union_all_materialization_pulls_one_row_at_a_time() {
861        let left: Box<dyn PhysicalOperator> =
862            Box::new(TableScan::from_rows(vec!["v".into()], vec![row(1), row(2)]));
863        let right: Box<dyn PhysicalOperator> =
864            Box::new(TableScan::from_rows(vec!["v".into()], vec![row(3)]));
865        let mut set = ExternalSetOperation::new_directional_with_types(
866            left,
867            right,
868            SetOpKind::Union,
869            true,
870            vec![None],
871            1,
872        )
873        .unwrap();
874        assert_eq!(
875            set.backward_scan_support(),
876            BackwardScanSupport::Unsupported
877        );
878        set.open().unwrap();
879        for expected in [1, 2, 3] {
880            let batch = set.next().unwrap().unwrap();
881            assert_eq!(batch.rows.len(), 1);
882            assert_eq!(directional_value(Some(batch)), Some(expected));
883        }
884        assert!(set.next().unwrap().is_none());
885        set.close().unwrap();
886    }
887}
uqa_execution/set_operation.rs

uqa_execution/
set_operation.rs
