Skip to content

Commit b8ab1d9

Browse files
KKouldKKould
committed
refactor: decouple recursive CTEs from spill
1 parent 4acae59 commit b8ab1d9

12 files changed

Lines changed: 79 additions & 60 deletions

File tree

‎src/binder/parser.rs‎

Lines changed: 3 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,6 @@ use crate::planner::operator::alter_table::change_column::{DefaultChange, NotNul
3434
use crate::planner::operator::join::{JoinCondition, JoinOperator as LJoinOperator, JoinType};
3535
use crate::planner::operator::mark_apply::MarkApplyQuantifier;
3636
use crate::planner::operator::project::ProjectOperator;
37-
#[cfg(feature = "spill")]
3837
use crate::planner::operator::recursive_cte::{RecursiveCteOperator, RecursiveScanOperator};
3938
use crate::planner::operator::sort::SortField;
4039
use crate::planner::operator::Operator;
@@ -2860,7 +2859,6 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
28602859
})
28612860
}
28622861

2863-
#[cfg(feature = "spill")]
28642862
fn bind_recursive_cte_plan(
28652863
&mut self,
28662864
cte: &sqlparser::ast::Cte,
@@ -2970,18 +2968,10 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
29702968
let origin_step = self.context.step_now();
29712969
let cte_checkpoint = self.context.cte_checkpoint();
29722970
if let Some(with) = &query.with {
2973-
if with.recursive {
2974-
#[cfg(not(feature = "spill"))]
2971+
if with.recursive && with.cte_tables.len() != 1 {
29752972
return Err(DatabaseError::UnsupportedStmt(
2976-
"recursive CTEs require the spill feature".to_string(),
2973+
"only one recursive CTE is supported".to_string(),
29772974
));
2978-
2979-
#[cfg(feature = "spill")]
2980-
if with.cte_tables.len() != 1 {
2981-
return Err(DatabaseError::UnsupportedStmt(
2982-
"only one recursive CTE is supported".to_string(),
2983-
));
2984-
}
29852975
}
29862976

29872977
self.context.cte_depth += 1;
@@ -3002,12 +2992,7 @@ impl<'a, 'parent, T: Transaction, A: AsRef<[(&'static str, DataValue)]>> Binder<
30022992

30032993
let alias = sql_table_alias(cte.alias.clone());
30042994
let plan = if with.recursive {
3005-
#[cfg(feature = "spill")]
3006-
{
3007-
self.bind_recursive_cte_plan(cte, &alias, arena)?
3008-
}
3009-
#[cfg(not(feature = "spill"))]
3010-
unreachable!()
2995+
self.bind_recursive_cte_plan(cte, &alias, arena)?
30112996
} else {
30122997
self.bind_cte_plan(cte, &alias, arena)?
30132998
};

‎src/execution/dql/mod.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@ pub(crate) mod join;
2525
pub(crate) mod limit;
2626
pub(crate) mod mark_apply;
2727
pub(crate) mod projection;
28-
#[cfg(feature = "spill")]
2928
pub(crate) mod recursive_cte;
3029
pub(crate) mod scalar_apply;
3130
pub(crate) mod scalar_subquery;

‎src/execution/dql/recursive_cte.rs‎

Lines changed: 76 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
// limitations under the License.
1414

1515
use crate::errors::DatabaseError;
16+
#[cfg(feature = "spill")]
1617
use crate::execution::spill::{SpillReader, SpillVec};
1718
use crate::execution::{
1819
build_read, ExecArena, ExecId, ExecNode, ExecutionContext, ExecutorNode, ReadExecutor,
@@ -25,7 +26,10 @@ use std::mem;
2526

2627
pub(crate) enum RecursiveInput {
2728
One(Option<Tuple>),
29+
#[cfg(feature = "spill")]
2830
Many(SpillReader<Tuple>),
31+
#[cfg(not(feature = "spill"))]
32+
Many(std::vec::IntoIter<Tuple>),
2933
}
3034

3135
impl Iterator for RecursiveInput {
@@ -34,7 +38,10 @@ impl Iterator for RecursiveInput {
3438
fn next(&mut self) -> Option<Self::Item> {
3539
match self {
3640
Self::One(tuple) => tuple.take().map(Ok),
41+
#[cfg(feature = "spill")]
3742
Self::Many(rows) => rows.next(),
43+
#[cfg(not(feature = "spill"))]
44+
Self::Many(rows) => rows.next().map(Ok),
3845
}
3946
}
4047
}
@@ -47,8 +54,17 @@ enum RecursiveRows {
4754
tuple: Tuple,
4855
output_done: bool,
4956
},
57+
#[cfg(feature = "spill")]
5058
Writing(SpillVec<'static, Tuple>),
59+
#[cfg(feature = "spill")]
5160
Reading(SpillReader<Tuple>),
61+
#[cfg(not(feature = "spill"))]
62+
Writing(Vec<Tuple>),
63+
#[cfg(not(feature = "spill"))]
64+
Reading {
65+
rows: Vec<Tuple>,
66+
output_index: usize,
67+
},
5268
}
5369

5470
impl RecursiveRows {
@@ -65,14 +81,29 @@ impl RecursiveRows {
6581
output_done: false,
6682
} => {
6783
let first = mem::take(first);
68-
let mut rows = SpillVec::new();
69-
let _ = rows.push(first)?;
70-
let _ = rows.push(tuple)?;
71-
*self = Self::Writing(rows);
84+
#[cfg(feature = "spill")]
85+
{
86+
let mut rows = SpillVec::new();
87+
let _ = rows.push(first)?;
88+
let _ = rows.push(tuple)?;
89+
*self = Self::Writing(rows);
90+
}
91+
#[cfg(not(feature = "spill"))]
92+
{
93+
*self = Self::Writing(vec![first, tuple]);
94+
}
7295
}
96+
#[cfg(feature = "spill")]
7397
Self::Writing(rows) => {
7498
let _ = rows.push(tuple)?;
7599
}
100+
#[cfg(not(feature = "spill"))]
101+
Self::Writing(rows) => rows.push(tuple),
102+
#[cfg(feature = "spill")]
103+
Self::One { .. } | Self::Reading(_) => {
104+
unreachable!("cannot append to a finished recursive generation")
105+
}
106+
#[cfg(not(feature = "spill"))]
76107
Self::One { .. } | Self::Reading { .. } => {
77108
unreachable!("cannot append to a finished recursive generation")
78109
}
@@ -82,10 +113,16 @@ impl RecursiveRows {
82113

83114
fn finish(self) -> Result<Self, DatabaseError> {
84115
match self {
116+
#[cfg(feature = "spill")]
85117
Self::Writing(mut rows) => {
86118
let _ = rows.flush()?;
87119
Ok(Self::Reading(rows.into_iter()))
88120
}
121+
#[cfg(not(feature = "spill"))]
122+
Self::Writing(rows) => Ok(Self::Reading {
123+
rows,
124+
output_index: 0,
125+
}),
89126
rows => Ok(rows),
90127
}
91128
}
@@ -100,7 +137,14 @@ impl RecursiveRows {
100137
*output_done = true;
101138
Ok(Some(tuple.clone()))
102139
}
140+
#[cfg(feature = "spill")]
103141
Self::Reading(reader) => reader.next().transpose(),
142+
#[cfg(not(feature = "spill"))]
143+
Self::Reading { rows, output_index } => {
144+
let tuple = rows.get(*output_index).cloned();
145+
*output_index += usize::from(tuple.is_some());
146+
Ok(tuple)
147+
}
104148
Self::Writing(_) => unreachable!("recursive generation must be finished first"),
105149
}
106150
}
@@ -109,10 +153,13 @@ impl RecursiveRows {
109153
match self {
110154
Self::Empty => Ok(None),
111155
Self::One { tuple, .. } => Ok(Some(RecursiveInput::One(Some(tuple)))),
156+
#[cfg(feature = "spill")]
112157
Self::Reading(mut reader) => {
113158
reader.reset()?;
114159
Ok(Some(RecursiveInput::Many(reader)))
115160
}
161+
#[cfg(not(feature = "spill"))]
162+
Self::Reading { rows, .. } => Ok(Some(RecursiveInput::Many(rows.into_iter()))),
116163
Self::Writing(_) => unreachable!("recursive generation must be finished first"),
117164
}
118165
}
@@ -288,6 +335,7 @@ mod tests {
288335
use std::borrow::Cow;
289336
use tempfile::TempDir;
290337

338+
#[cfg(feature = "spill")]
291339
#[test]
292340
fn spilled_generation_is_written_once_and_replayed_for_scan() -> Result<(), DatabaseError> {
293341
let expected = (0..1100)
@@ -311,6 +359,30 @@ mod tests {
311359
Ok(())
312360
}
313361

362+
#[cfg(not(feature = "spill"))]
363+
#[test]
364+
fn memory_generation_is_replayed_for_scan() -> Result<(), DatabaseError> {
365+
let expected = (0..3)
366+
.map(|value| Tuple::new(None, vec![DataValue::Int32(value)]))
367+
.collect::<Vec<_>>();
368+
let mut rows = RecursiveRows::default();
369+
for tuple in expected.iter().cloned() {
370+
rows.push(tuple)?;
371+
}
372+
373+
let mut working = rows.finish()?;
374+
assert!(matches!(&working, RecursiveRows::Reading { .. }));
375+
let mut output = Vec::new();
376+
while let Some(tuple) = working.next_output()? {
377+
output.push(tuple);
378+
}
379+
assert_eq!(output, expected);
380+
381+
let scan = working.into_input()?.unwrap();
382+
assert_eq!(scan.collect::<Result<Vec<_>, _>>()?, expected);
383+
Ok(())
384+
}
385+
314386
#[test]
315387
fn empty_generation_has_no_recursive_input() -> Result<(), DatabaseError> {
316388
let rows = RecursiveRows::default();

‎src/execution/mod.rs‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,6 @@ use crate::execution::dql::index_scan::IndexScan;
5959
use crate::execution::dql::join::hash_join::HashJoin;
6060
use crate::execution::dql::limit::Limit;
6161
use crate::execution::dql::projection::Projection;
62-
#[cfg(feature = "spill")]
6362
use crate::execution::dql::recursive_cte::{RecursiveCte, RecursiveInput, RecursiveScan};
6463
use crate::execution::dql::scalar_subquery::ScalarSubquery;
6564
use crate::execution::dql::seq_scan::SeqScan;
@@ -199,9 +198,7 @@ pub(crate) enum ExecNode<'a, T: Transaction + 'a> {
199198
MarkApply(MarkApply),
200199
NestedLoopJoin(NestedLoopJoin),
201200
Projection(Projection),
202-
#[cfg(feature = "spill")]
203201
RecursiveCte(RecursiveCte<'a, T>),
204-
#[cfg(feature = "spill")]
205202
RecursiveScan(RecursiveScan),
206203
ScalarApply(ScalarApply),
207204
ScalarSubquery(ScalarSubquery),
@@ -321,11 +318,9 @@ impl<'a, T: Transaction + 'a> ExecNode<'a, T> {
321318
ExecNode::Projection(exec) => {
322319
<Projection as ExecutorNode<'a, T>>::next_tuple(exec, arena, plan_arena)
323320
}
324-
#[cfg(feature = "spill")]
325321
ExecNode::RecursiveCte(exec) => {
326322
<RecursiveCte<'a, T> as ExecutorNode<'a, T>>::next_tuple(exec, arena, plan_arena)
327323
}
328-
#[cfg(feature = "spill")]
329324
ExecNode::RecursiveScan(exec) => {
330325
<RecursiveScan as ExecutorNode<'a, T>>::next_tuple(exec, arena, plan_arena)
331326
}
@@ -391,7 +386,6 @@ pub(crate) struct ExecArena<'a, T: Transaction + 'a> {
391386
transaction: *mut T,
392387
runtime_probe_stack: Vec<RuntimeIndexProbe>,
393388
ddl_apply: Vec<DDLApply>,
394-
#[cfg(feature = "spill")]
395389
recursive_input: Option<RecursiveInput>,
396390
}
397391

@@ -441,7 +435,6 @@ impl<'a, T: Transaction + 'a> ExecArena<'a, T> {
441435
transaction: std::ptr::null_mut(),
442436
runtime_probe_stack: Vec::new(),
443437
ddl_apply: Vec::new(),
444-
#[cfg(feature = "spill")]
445438
recursive_input: None,
446439
}
447440
}
@@ -556,20 +549,17 @@ impl<'a, T: Transaction + 'a> ExecArena<'a, T> {
556549
self.runtime_probe_stack.len()
557550
}
558551

559-
#[cfg(feature = "spill")]
560552
pub(crate) fn set_recursive_input(&mut self, input: RecursiveInput) {
561553
debug_assert!(self.recursive_input.is_none());
562554
self.recursive_input = Some(input);
563555
}
564556

565-
#[cfg(feature = "spill")]
566557
pub(crate) fn take_recursive_input(&mut self) -> RecursiveInput {
567558
self.recursive_input
568559
.take()
569560
.expect("recursive input initialized")
570561
}
571562

572-
#[cfg(feature = "spill")]
573563
pub(crate) fn reset_for_rebuild(&mut self) {
574564
debug_assert!(self.runtime_probe_stack.is_empty());
575565
debug_assert!(self.ddl_apply.is_empty());
@@ -931,15 +921,13 @@ where
931921
cache,
932922
transaction,
933923
),
934-
#[cfg(feature = "spill")]
935924
Operator::RecursiveCte(_) => <RecursiveCte<'a, T> as ReadExecutor<'a, T>>::into_executor(
936925
childrens.pop_twins(),
937926
arena,
938927
plan_arena,
939928
cache,
940929
transaction,
941930
),
942-
#[cfg(feature = "spill")]
943931
Operator::RecursiveScan(op) => <RecursiveScan as ReadExecutor<'a, T>>::into_executor(
944932
op,
945933
arena,

‎src/optimizer/rule/implementation/mod.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -134,7 +134,6 @@ impl ImplementationRuleRootTag {
134134
| Operator::CreateView(_)
135135
| Operator::DropView(_)
136136
| Operator::DropIndex(_) => None,
137-
#[cfg(feature = "spill")]
138137
Operator::RecursiveCte(_) | Operator::RecursiveScan(_) => None,
139138
}
140139
}

‎src/optimizer/rule/normalization/column_pruning.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -662,7 +662,6 @@ impl ColumnPruning {
662662
}
663663
}
664664
}
665-
#[cfg(feature = "spill")]
666665
Operator::RecursiveCte(_) => {
667666
changed |= Self::apply_twins(
668667
required_columns,
@@ -678,7 +677,6 @@ impl ColumnPruning {
678677
Operator::Dummy | Operator::Values(_) | Operator::FunctionScan(_) => {
679678
outcome.removed_positions.truncate(output_start);
680679
}
681-
#[cfg(feature = "spill")]
682680
Operator::RecursiveScan(_) => {
683681
outcome.removed_positions.truncate(output_start);
684682
}

‎src/optimizer/rule/normalization/compilation_in_advance.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@ impl EvaluatorBind {
4444
| Operator::Union(_)
4545
| Operator::SetMembership(_)
4646
);
47-
#[cfg(feature = "spill")]
4847
let bind_right = bind_right || matches!(plan.operator, Operator::RecursiveCte(_));
4948
if bind_right {
5049
Self::_apply(right, arena)?;

‎src/optimizer/rule/normalization/mod.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,6 @@ impl NormalizationRuleRootTag {
135135
| Operator::Union(_)
136136
| Operator::SetMembership(_)
137137
| Operator::Window(_) => None,
138-
#[cfg(feature = "spill")]
139138
Operator::RecursiveCte(_) | Operator::RecursiveScan(_) => None,
140139
#[cfg(feature = "copy")]
141140
Operator::CopyFromFile(_) | Operator::CopyToFile(_) => None,

‎src/planner/mod.rs‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@ pub mod operator;
1717

1818
use crate::catalog::TableName;
1919
use crate::errors::DatabaseError;
20-
#[cfg(feature = "spill")]
2120
use crate::planner::operator::recursive_cte::{RecursiveCteOperator, RecursiveScanOperator};
2221
use crate::planner::operator::set_membership::SetMembershipOperator;
2322
use crate::planner::operator::union::UnionOperator;
@@ -242,7 +241,6 @@ impl LogicalPlan {
242241
left_schema_ref: schema_ref,
243242
..
244243
}) => schema_ref.clone(),
245-
#[cfg(feature = "spill")]
246244
Operator::RecursiveCte(RecursiveCteOperator { schema_ref })
247245
| Operator::RecursiveScan(RecursiveScanOperator { schema_ref }) => schema_ref.clone(),
248246
Operator::Dummy => Vec::new(),

0 commit comments

Comments
 (0)