1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
//! `UNION`, `EXCEPT` and `INTERSECT`.
//!
//! The distinction that costs the most code here is `ALL`. Without it the operation is over sets
//! and the answer is a dedup. With it the operation is over multisets, and `EXCEPT ALL` of three
//! copies of a row minus one copy is two copies, not zero and not three. That rule is in the
//! standard, DuckDB implements it, and it is the sort of thing an implementation gets wrong once
//! and then nobody notices for a year because nothing in a normal query has duplicates in it.
//!
//! The output columns are the left side's, under the set operation's own table index. Both sides
//! were made type compatible by the binder, so nothing here casts anything.
//!
//! # Two pipelines and an edge
//!
//! This is the first operator with two inputs to move behind the push traits, and two inputs is two
//! pipelines. The right side ends in a [`Gather`](crate::gather::Gather), which holds its rows and
//! nothing else, and the left side ends here. The order is not a choice: every arm below needs the
//! whole right side before it can say anything about one left row, which is the dependency edge the
//! builder records on the pipeline. [`Query::run`](crate::Query::run) takes its order from that
//! edge, so the right side's pipeline has finalised before this one starts.
use std::sync::Mutex;
use rudb_common::{Error, Memory, Reservation, Result, Value};
use rudb_pipeline::{Lease, Progress, Sink};
use rudb_plan::SetOpKind;
use rudb_vector::Chunk;
use crate::buffer::Buffered;
use crate::gather::{self, Gathering, Rows};
use crate::key::{Key, RowMap, RowSet};
use crate::rows;
use crate::schema::Schema;
/// A set operation over two inputs.
#[derive(Debug)]
pub(crate) struct SetOp {
kind: SetOpKind,
all: bool,
schema: Schema,
memory: Memory,
/// The right side, filled by the pipeline this one depends on.
right: Rows,
/// The left side, as every instance gathered it.
left: Mutex<Vec<Vec<Value>>>,
/// What the left side is charged, given back once the finished chunks are charged instead.
charged: Mutex<Vec<Reservation>>,
/// What the finished chunks are charged, held for as long as they are readable.
held: Mutex<Reservation>,
out: Buffered,
}
impl SetOp {
/// The sink for the left side, and the source the answer comes out of.
///
/// `left` is the left input's schema, whose fields become the output's under `index`, and
/// `right` is the handle on the rows the other pipeline gathered.
pub(crate) fn new(
left: &Schema,
right: Rows,
kind: SetOpKind,
all: bool,
index: u32,
memory: &Memory,
) -> (Self, Buffered) {
let out = Buffered::new();
let setop = Self {
kind,
all,
schema: Schema::numbered(left.fields().to_vec(), index),
memory: memory.clone(),
right,
left: Mutex::new(Vec::new()),
charged: Mutex::new(Vec::new()),
held: Mutex::new(memory.reservation()),
out: out.clone(),
};
(setop, out)
}
/// What this operator produces, which is the left side's columns under its own table index.
pub(crate) fn schema(&self) -> &Schema {
&self.schema
}
}
impl Sink for SetOp {
type Local = Gathering;
fn local(&self) -> Gathering {
gather::gathering(&self.memory)
}
/// Not yet. Every arm of a set operation produces its rows in the order the left side arrived
/// in, so which instance got which morsel would show up in the answer.
fn parallel(&self) -> bool {
false
}
fn sink(&self, chunk: &Chunk, local: &mut Gathering) -> Result<Progress> {
gather::take(chunk, local)?;
Ok(Progress::More)
}
fn combine(&self, local: Gathering) -> Result<()> {
let (rows, charged) = gather::into_parts(local);
self.left.lock().map_err(poisoned)?.extend(rows);
self.charged.lock().map_err(poisoned)?.push(charged);
Ok(())
}
fn finalize(&self, _threads: &Lease<'_>) -> Result<()> {
// Both sides at once, which is what every arm below needs, and the counting tables on top
// of them. The tables are not charged separately, because a count per distinct row is
// bounded by the rows that are already charged and charging it twice would refuse a query
// that fits.
let left = std::mem::take(&mut *self.left.lock().map_err(poisoned)?);
let right = self.right.take()?;
let out = match (self.kind, self.all) {
(SetOpKind::Union, true) => {
let mut out = left;
out.extend(right);
out
}
(SetOpKind::Union, false) => {
let mut out = left;
out.extend(right);
deduplicated(out)
}
(SetOpKind::Except, true) => difference(left, &counts(&right)),
(SetOpKind::Except, false) => {
let held = counts(&right);
deduplicated(
left.into_iter().filter(|row| !held.contains_key(&Key(row.clone()))).collect(),
)
}
(SetOpKind::Intersect, true) => intersection(left, &counts(&right)),
(SetOpKind::Intersect, false) => {
let held = counts(&right);
deduplicated(
left.into_iter().filter(|row| held.contains_key(&Key(row.clone()))).collect(),
)
}
};
let mut held = self.held.lock().map_err(poisoned)?;
let chunks = rows::chunks(&self.schema.types(), &out, &mut held)?;
self.out.fill(chunks)?;
self.charged.lock().map_err(poisoned)?.clear();
Ok(())
}
}
fn poisoned<T>(_: T) -> Error {
Error::internal("a thread panicked while holding the rows a set operation gathered")
}
/// How many times each row appears.
fn counts(rows: &[Vec<Value>]) -> RowMap<usize> {
let mut held = RowMap::default();
for row in rows {
*held.entry(Key(row.clone())).or_insert(0) += 1;
}
held
}
/// The first occurrence of each row, in the order they arrived.
fn deduplicated(rows: Vec<Vec<Value>>) -> Vec<Vec<Value>> {
let mut seen = RowSet::default();
rows.into_iter().filter(|row| seen.insert(Key(row.clone()))).collect()
}
/// `EXCEPT ALL`: each left row survives unless a right row has already cancelled it.
fn difference(left: Vec<Vec<Value>>, right: &RowMap<usize>) -> Vec<Vec<Value>> {
let mut budget = right.clone();
let mut out = Vec::new();
for row in left {
match budget.get_mut(&Key(row.clone())) {
Some(remaining) if *remaining > 0 => *remaining -= 1,
_ => out.push(row),
}
}
out
}
/// `INTERSECT ALL`: a left row survives while the right side still has a copy to pair it with.
fn intersection(left: Vec<Vec<Value>>, right: &RowMap<usize>) -> Vec<Vec<Value>> {
let mut budget = right.clone();
let mut out = Vec::new();
for row in left {
if let Some(remaining) = budget.get_mut(&Key(row.clone())) {
if *remaining > 0 {
*remaining -= 1;
out.push(row);
}
}
}
out
}