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    Batch, ExecError, ExecResult, ExpressionEvaluator, ExternalSort, PhysicalOperator, PhysicalRow,
22    RowProjectionValue, RowSchema, ScalarExpr, SortKey,
23};
24
25struct ColumnEvaluator;
26
27impl ExpressionEvaluator for ColumnEvaluator {
28    fn evaluate(&self, expression: &ScalarExpr, row: &dyn RowLookup) -> ExecResult<Value> {
29        let ScalarExpr::Column(column) = expression else {
30            return Err(ExecError::Other(
31                "set-operation sort key must be a column".into(),
32            ));
33        };
34        Ok(row.column(column).cloned().unwrap_or(Value::Null))
35    }
36}
37
38fn set_operation_types(left: &RowSchema, right: &RowSchema) -> ExecResult<Vec<Option<ColumnType>>> {
39    if left.len() != right.len() {
40        return Err(ExecError::Other(format!(
41            "set-operation inputs have different widths: {} and {}",
42            left.len(),
43            right.len()
44        )));
45    }
46    left.column_types()
47        .iter()
48        .zip(right.column_types())
49        .map(|(left, right)| match (left, right) {
50            (None, None) => Ok(None),
51            (Some(ty), None) | (None, Some(ty)) => Ok(Some(ty.clone())),
52            (Some(left), Some(right)) => uqa_execution_common_type(left, right).map(Some),
53        })
54        .collect()
55}
56
57fn uqa_execution_common_type(left: &ColumnType, right: &ColumnType) -> ExecResult<ColumnType> {
58    crate::common_type(left, right).map_err(ExecError::from)
59}
60
61fn coerce_set_value(
62    value: Value,
63    source_type: Option<&ColumnType>,
64    target_type: &ColumnType,
65) -> ExecResult<Value> {
66    let cast_target = match target_type {
67        ColumnType::Domain { base, .. } => base.as_ref(),
68        target => target,
69    };
70    let source_name = source_type.map(ColumnType::sql_name);
71    uqa_sql::expr::cast_value_from(&value, &cast_target.sql_name(), source_name.as_deref())
72        .map_err(ExecError::from)
73}
74
75struct AlignSchema<'a> {
76    child: Box<dyn PhysicalOperator + 'a>,
77    schema: RowSchema,
78    coercions: Vec<Option<ColumnType>>,
79}
80
81impl<'a> AlignSchema<'a> {
82    fn new(
83        child: Box<dyn PhysicalOperator + 'a>,
84        output: Vec<String>,
85        output_types: &[Option<ColumnType>],
86    ) -> ExecResult<Self> {
87        let source = child.schema().to_vec();
88        if source.len() != output.len() {
89            return Err(ExecError::Other(format!(
90                "set-operation inputs have different widths: {} and {}",
91                output.len(),
92                source.len()
93            )));
94        }
95        if output.len() != output_types.len() {
96            return Err(ExecError::Other(format!(
97                "set-operation output type width {} does not match input width {}",
98                output_types.len(),
99                output.len()
100            )));
101        }
102        let coercions = child
103            .row_schema()
104            .column_types()
105            .iter()
106            .zip(output_types)
107            .map(|(source, target)| {
108                target
109                    .as_ref()
110                    .filter(|target| source.as_ref() != Some(*target))
111                    .cloned()
112            })
113            .collect();
114        let schema = RowSchema::with_types(output, output_types.to_vec());
115        Ok(Self {
116            child,
117            schema,
118            coercions,
119        })
120    }
121}
122
123impl PhysicalOperator for AlignSchema<'_> {
124    fn row_schema(&self) -> &RowSchema {
125        &self.schema
126    }
127
128    fn open(&mut self) -> ExecResult<()> {
129        self.child.open()
130    }
131
132    fn next(&mut self) -> ExecResult<Option<Batch>> {
133        let Some(batch) = self.child.next()? else {
134            return Ok(None);
135        };
136        let identity_layout = batch.schema.physical_width() == self.schema.physical_width()
137            && (0..batch.schema.len())
138                .all(|position| batch.schema.slot(position) == Some(position));
139        if self.coercions.iter().all(Option::is_none) && identity_layout {
140            let rows = batch
141                .rows
142                .into_iter()
143                .map(PhysicalRow::without_lock_origins)
144                .collect();
145            return Ok(Some(Batch::from_physical_rows(self.schema.clone(), rows)));
146        }
147        if self.coercions.iter().all(Option::is_none) {
148            let slots = (0..batch.schema.len())
149                .map(|position| {
150                    batch.schema.slot(position).ok_or_else(|| {
151                        ExecError::Other(format!(
152                            "set-operation input column {position} has no physical slot"
153                        ))
154                    })
155                })
156                .collect::<ExecResult<Vec<_>>>()?;
157            let rows = batch
158                .rows
159                .into_iter()
160                .map(|row| row.project_slots(&slots).without_lock_origins())
161                .collect();
162            return Ok(Some(Batch::from_physical_rows(self.schema.clone(), rows)));
163        }
164        let mut rows = Vec::with_capacity(batch.rows.len());
165        for row in batch.rows {
166            let view = batch.schema.view(&row);
167            let values = self
168                .coercions
169                .iter()
170                .enumerate()
171                .map(|(position, target)| match target {
172                    Some(target) => coerce_set_value(
173                        view.value_at(position).cloned().unwrap_or(Value::Null),
174                        batch.schema.column_type(position),
175                        target,
176                    )
177                    .map(RowProjectionValue::Owned),
178                    None => Ok(batch.schema.slot(position).map_or(
179                        RowProjectionValue::Owned(Value::Null),
180                        RowProjectionValue::InputSlot,
181                    )),
182                })
183                .collect::<ExecResult<Vec<_>>>()?;
184            rows.push(row.project_with_values(values).without_lock_origins());
185        }
186        Ok(Some(Batch::from_physical_rows(self.schema.clone(), rows)))
187    }
188
189    fn close(&mut self) -> ExecResult<()> {
190        self.child.close()
191    }
192}
193
194struct RowGroup {
195    row: PhysicalRow,
196    count: usize,
197}
198
199struct RowCursor<'a> {
200    operator: Box<dyn PhysicalOperator + 'a>,
201    batch: std::vec::IntoIter<PhysicalRow>,
202    lookahead: Option<PhysicalRow>,
203    exhausted: bool,
204}
205
206impl<'a> RowCursor<'a> {
207    fn new(operator: Box<dyn PhysicalOperator + 'a>) -> Self {
208        Self {
209            operator,
210            batch: Vec::new().into_iter(),
211            lookahead: None,
212            exhausted: false,
213        }
214    }
215
216    fn open(&mut self) -> ExecResult<()> {
217        self.batch = Vec::new().into_iter();
218        self.lookahead = None;
219        self.exhausted = false;
220        self.operator.open()
221    }
222
223    fn next_row(&mut self) -> ExecResult<Option<PhysicalRow>> {
224        loop {
225            if let Some(row) = self.batch.next() {
226                return Ok(Some(row));
227            }
228            let Some(batch) = self.operator.next()? else {
229                self.exhausted = true;
230                return Ok(None);
231            };
232            self.batch = batch.rows.into_iter();
233        }
234    }
235
236    fn take_group(&mut self, schema: &RowSchema) -> ExecResult<Option<RowGroup>> {
237        let first = match self.lookahead.take() {
238            Some(row) => row,
239            None => match self.next_row()? {
240                Some(row) => row,
241                None => return Ok(None),
242            },
243        };
244        let mut count = 1_usize;
245        while let Some(row) = self.next_row()? {
246            if compare_rows(&first, &row, schema) == Ordering::Equal {
247                count = count
248                    .checked_add(1)
249                    .ok_or_else(|| ExecError::Other("set-operation group count overflow".into()))?;
250            } else {
251                self.lookahead = Some(row);
252                break;
253            }
254        }
255        Ok(Some(RowGroup { row: first, count }))
256    }
257
258    fn close(&mut self) -> ExecResult<()> {
259        self.batch = Vec::new().into_iter();
260        self.lookahead = None;
261        self.exhausted = true;
262        self.operator.close()
263    }
264}
265
266fn compare_rows(left: &PhysicalRow, right: &PhysicalRow, schema: &RowSchema) -> Ordering {
267    let left = schema.view(left);
268    let right = schema.view(right);
269    let null = Value::Null;
270    for position in 0..schema.len() {
271        let ordering = left
272            .value_at(position)
273            .unwrap_or(&null)
274            .cmp(right.value_at(position).unwrap_or(&null));
275        if ordering != Ordering::Equal {
276            return ordering;
277        }
278    }
279    Ordering::Equal
280}
281
282/// `UNION` / `INTERSECT` / `EXCEPT` physical operator.
283///
284/// `UNION ALL` streams its children without sorting. Other forms sort both
285/// inputs through byte-bounded external runs and merge adjacent multiplicity
286/// counts, avoiding whole-input materialisation and quadratic row scans.
287pub struct ExternalSetOperation<'a> {
288    left: RowCursor<'a>,
289    right: RowCursor<'a>,
290    kind: SetOpKind,
291    all: bool,
292    schema: RowSchema,
293    left_group: Option<RowGroup>,
294    right_group: Option<RowGroup>,
295    pending_row: Option<PhysicalRow>,
296    pending_count: usize,
297    union_all_left_done: bool,
298}
299
300impl<'a> ExternalSetOperation<'a> {
301    pub fn new(
302        left: Box<dyn PhysicalOperator + 'a>,
303        right: Box<dyn PhysicalOperator + 'a>,
304        kind: SetOpKind,
305        all: bool,
306        work_mem_bytes: usize,
307    ) -> ExecResult<Self> {
308        let output_types = set_operation_types(left.row_schema(), right.row_schema())?;
309        Self::new_with_types(left, right, kind, all, output_types, work_mem_bytes)
310    }
311
312    pub fn new_with_types(
313        left: Box<dyn PhysicalOperator + 'a>,
314        right: Box<dyn PhysicalOperator + 'a>,
315        kind: SetOpKind,
316        all: bool,
317        output_types: Vec<Option<ColumnType>>,
318        work_mem_bytes: usize,
319    ) -> ExecResult<Self> {
320        let output = left.schema().to_vec();
321        let left: Box<dyn PhysicalOperator + 'a> =
322            Box::new(AlignSchema::new(left, output.clone(), &output_types)?);
323        let right: Box<dyn PhysicalOperator + 'a> =
324            Box::new(AlignSchema::new(right, output.clone(), &output_types)?);
325        let (left, right) = if matches!((kind, all), (SetOpKind::Union, true)) {
326            (left, right)
327        } else {
328            let keys = output
329                .iter()
330                .map(|column| SortKey {
331                    expr: ScalarExpr::Column(column.clone()),
332                    descending: false,
333                    nulls_first: Some(true),
334                })
335                .collect::<Vec<_>>();
336            let evaluator = Arc::new(ColumnEvaluator);
337            // Both merge inputs are live concurrently. Split the configured
338            // budget so their encoded in-memory runs cannot each claim it.
339            let per_input = (work_mem_bytes / 2).max(1);
340            (
341                Box::new(ExternalSort::new(
342                    left,
343                    keys.clone(),
344                    evaluator.clone(),
345                    None,
346                    per_input,
347                )) as Box<dyn PhysicalOperator + 'a>,
348                Box::new(ExternalSort::new(right, keys, evaluator, None, per_input))
349                    as Box<dyn PhysicalOperator + 'a>,
350            )
351        };
352        Ok(Self {
353            left: RowCursor::new(left),
354            right: RowCursor::new(right),
355            kind,
356            all,
357            schema: RowSchema::with_types(output, output_types),
358            left_group: None,
359            right_group: None,
360            pending_row: None,
361            pending_count: 0,
362            union_all_left_done: false,
363        })
364    }
365
366    fn next_union_all(&mut self) -> ExecResult<Option<Batch>> {
367        let mut rows = Vec::with_capacity(DEFAULT_BATCH_SIZE);
368        while rows.len() < DEFAULT_BATCH_SIZE {
369            let next = if self.union_all_left_done {
370                self.right.next_row()?
371            } else if let Some(row) = self.left.next_row()? {
372                Some(row)
373            } else {
374                self.union_all_left_done = true;
375                self.right.next_row()?
376            };
377            let Some(row) = next else {
378                break;
379            };
380            rows.push(row.without_lock_origins());
381        }
382        if rows.is_empty() {
383            Ok(None)
384        } else {
385            Ok(Some(Batch::from_physical_rows(self.schema.clone(), rows)))
386        }
387    }
388
389    fn load_groups(&mut self) -> ExecResult<()> {
390        if self.left_group.is_none() && !self.left.exhausted {
391            self.left_group = self.left.take_group(&self.schema)?;
392        }
393        if self.right_group.is_none() && !self.right.exhausted {
394            self.right_group = self.right.take_group(&self.schema)?;
395        }
396        Ok(())
397    }
398
399    fn take_left_group(&mut self) -> ExecResult<RowGroup> {
400        self.left_group
401            .take()
402            .ok_or_else(|| ExecError::Other("set-operation selected a missing left group".into()))
403    }
404
405    fn take_right_group(&mut self) -> ExecResult<RowGroup> {
406        self.right_group
407            .take()
408            .ok_or_else(|| ExecError::Other("set-operation selected a missing right group".into()))
409    }
410
411    fn choose_group(&mut self) -> ExecResult<Option<(PhysicalRow, usize)>> {
412        self.load_groups()?;
413        let ordering = match (&self.left_group, &self.right_group) {
414            (Some(left), Some(right)) => Some(compare_rows(&left.row, &right.row, &self.schema)),
415            (Some(_), None) => Some(Ordering::Less),
416            (None, Some(_)) => Some(Ordering::Greater),
417            (None, None) => None,
418        };
419        let Some(ordering) = ordering else {
420            return Ok(None);
421        };
422
423        let selected = match (self.kind, ordering) {
424            (SetOpKind::Union, Ordering::Less) => {
425                let group = self.take_left_group()?;
426                Some((group.row, 1))
427            }
428            (SetOpKind::Union, Ordering::Greater) => {
429                let group = self.take_right_group()?;
430                Some((group.row, 1))
431            }
432            (SetOpKind::Union, Ordering::Equal) => {
433                let group = self.take_left_group()?;
434                self.right_group = None;
435                Some((group.row, 1))
436            }
437            (SetOpKind::Intersect, Ordering::Less) => {
438                self.left_group = None;
439                None
440            }
441            (SetOpKind::Intersect | SetOpKind::Except, Ordering::Greater) => {
442                self.right_group = None;
443                None
444            }
445            (SetOpKind::Intersect, Ordering::Equal) => {
446                let left = self.take_left_group()?;
447                let right = self.take_right_group()?;
448                Some((
449                    left.row,
450                    if self.all {
451                        left.count.min(right.count)
452                    } else {
453                        1
454                    },
455                ))
456            }
457            (SetOpKind::Except, Ordering::Less) => {
458                let left = self.take_left_group()?;
459                Some((left.row, if self.all { left.count } else { 1 }))
460            }
461            (SetOpKind::Except, Ordering::Equal) => {
462                let left = self.take_left_group()?;
463                let right = self.take_right_group()?;
464                let count = if self.all {
465                    left.count.saturating_sub(right.count)
466                } else {
467                    0
468                };
469                Some((left.row, count))
470            }
471        };
472        Ok(selected.filter(|(_, count)| *count > 0))
473    }
474}
475
476impl PhysicalOperator for ExternalSetOperation<'_> {
477    fn row_schema(&self) -> &RowSchema {
478        &self.schema
479    }
480
481    fn open(&mut self) -> ExecResult<()> {
482        self.left_group = None;
483        self.right_group = None;
484        self.pending_row = None;
485        self.pending_count = 0;
486        self.union_all_left_done = false;
487        self.left.open()?;
488        self.right.open()
489    }
490
491    fn next(&mut self) -> ExecResult<Option<Batch>> {
492        if matches!((self.kind, self.all), (SetOpKind::Union, true)) {
493            return self.next_union_all();
494        }
495        let mut rows = Vec::with_capacity(DEFAULT_BATCH_SIZE);
496        while rows.len() < DEFAULT_BATCH_SIZE {
497            if self.pending_count > 0 {
498                let row = self.pending_row.as_ref().ok_or_else(|| {
499                    ExecError::Other(
500                        "set-operation has pending multiplicity without a pending row".into(),
501                    )
502                })?;
503                rows.push(row.clone());
504                self.pending_count -= 1;
505                if self.pending_count == 0 {
506                    self.pending_row = None;
507                }
508                continue;
509            }
510            match self.choose_group()? {
511                Some((row, count)) => {
512                    self.pending_row = Some(row.without_lock_origins());
513                    self.pending_count = count;
514                }
515                None if self.left.exhausted && self.right.exhausted => break,
516                None => {}
517            }
518        }
519        if rows.is_empty() {
520            Ok(None)
521        } else {
522            Ok(Some(Batch::from_physical_rows(self.schema.clone(), rows)))
523        }
524    }
525
526    fn close(&mut self) -> ExecResult<()> {
527        self.left_group = None;
528        self.right_group = None;
529        self.pending_row = None;
530        self.pending_count = 0;
531        let left = self.left.close();
532        let right = self.right.close();
533        crate::physical::with_cleanup(left, right, "close right set-operation input")
534    }
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540    use crate::physical::run_to_rows;
541    use crate::scan::TableScan;
542
543    fn row(value: i64) -> ResultRow {
544        [("v".into(), Value::Int(value))].into_iter().collect()
545    }
546
547    fn execute(kind: SetOpKind, all: bool, left: &[i64], right: &[i64]) -> Vec<i64> {
548        let left = TableScan::from_rows(vec!["v".into()], left.iter().copied().map(row).collect());
549        let right = TableScan::from_rows(
550            vec!["other".into()],
551            right
552                .iter()
553                .copied()
554                .map(|value| [("other".into(), Value::Int(value))].into_iter().collect())
555                .collect(),
556        );
557        let mut set =
558            ExternalSetOperation::new(Box::new(left), Box::new(right), kind, all, 1).unwrap();
559        run_to_rows(&mut set)
560            .unwrap()
561            .1
562            .into_iter()
563            .map(|row| match row.get("v") {
564                Some(Value::Int(value)) => *value,
565                value => panic!("unexpected set value: {value:?}"),
566            })
567            .collect()
568    }
569
570    #[test]
571    fn external_set_semantics_include_bag_multiplicity() {
572        assert_eq!(
573            execute(SetOpKind::Union, false, &[2, 1, 1], &[3, 2]),
574            vec![1, 2, 3]
575        );
576        assert_eq!(
577            execute(SetOpKind::Union, true, &[2, 1, 1], &[3, 2]),
578            vec![2, 1, 1, 3, 2]
579        );
580        assert_eq!(
581            execute(SetOpKind::Intersect, true, &[1, 1, 1, 2], &[1, 1, 3]),
582            vec![1, 1]
583        );
584        assert_eq!(
585            execute(SetOpKind::Except, true, &[1, 1, 1, 2], &[1, 1, 3]),
586            vec![1, 2]
587        );
588        assert_eq!(
589            execute(SetOpKind::Except, false, &[1, 1, 2], &[1, 3]),
590            vec![2]
591        );
592    }
593
594    #[test]
595    fn set_inputs_compact_schema_only_projections_before_alignment() {
596        let left = TableScan::from_rows(
597            vec!["v".into(), "hidden".into()],
598            vec![[
599                ("v".into(), Value::Int(1)),
600                ("hidden".into(), Value::Int(99)),
601            ]
602            .into_iter()
603            .collect()],
604        );
605        let left: Box<dyn PhysicalOperator> = Box::new(crate::ColumnSelection::with_positions(
606            Box::new(left),
607            vec![("v".into(), 0)],
608        ));
609        let right = TableScan::from_rows(vec!["v".into()], vec![row(2)]);
610        let mut set =
611            ExternalSetOperation::new(left, Box::new(right), SetOpKind::Union, true, 1).unwrap();
612        let rows = run_to_rows(&mut set).unwrap().1;
613        assert_eq!(rows, vec![row(1), row(2)]);
614    }
615}
uqa_execution/set_operation.rs

uqa_execution/
set_operation.rs
