Skip to content

Commit 1b0f10b

Browse files
wszhdshysKKould
andauthored
feat: TopK optimizer&planner&executor (#288)
* The topk algorithm is implemented to optimize the use of both order by and limit * Optimized logic * Optimized performance and added TOPK related tests. Improved the display of TOPK's plan * chore: codefmt * chore: codefmt * chore: codefmt --------- Co-authored-by: Kould <kould2333@gmail.com>
1 parent f3dfbef commit 1b0f10b

File tree

16 files changed

+931
-16
lines changed

16 files changed

+931
-16
lines changed

src/db.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,11 @@ impl<S: Storage> State<S> {
227227
NormalizationRuleImpl::CombineFilter,
228228
],
229229
)
230+
.batch(
231+
"TopK".to_string(),
232+
HepBatchStrategy::once_topdown(),
233+
vec![NormalizationRuleImpl::TopK],
234+
)
230235
.batch(
231236
"Expression Remapper".to_string(),
232237
HepBatchStrategy::once_topdown(),
@@ -249,6 +254,7 @@ impl<S: Storage> State<S> {
249254
ImplementationRuleImpl::IndexScan,
250255
ImplementationRuleImpl::FunctionScan,
251256
ImplementationRuleImpl::Sort,
257+
ImplementationRuleImpl::TopK,
252258
ImplementationRuleImpl::Values,
253259
// DML
254260
ImplementationRuleImpl::Analyze,

src/execution/dql/mod.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ pub(crate) mod seq_scan;
1313
pub(crate) mod show_table;
1414
pub(crate) mod show_view;
1515
pub(crate) mod sort;
16+
pub(crate) mod top_k;
1617
pub(crate) mod union;
1718
pub(crate) mod values;
1819

src/execution/dql/sort.rs

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ use std::pin::Pin;
1515
pub(crate) type BumpVec<'bump, T> = bumpalo::collections::Vec<'bump, T>;
1616

1717
#[derive(Clone)]
18-
pub(crate) struct NullableVec<'a, T>(BumpVec<'a, Option<T>>);
18+
pub(crate) struct NullableVec<'a, T>(pub(crate) BumpVec<'a, Option<T>>);
1919

2020
impl<'a, T> NullableVec<'a, T> {
2121
#[inline]
@@ -49,17 +49,31 @@ impl<'a, T> NullableVec<'a, T> {
4949
}
5050
}
5151

52-
struct RemappingIterator<'a> {
52+
pub struct RemappingIterator<'a> {
5353
pos: usize,
5454
tuples: NullableVec<'a, (usize, Tuple)>,
5555
indices: BumpVec<'a, usize>,
5656
}
5757

58+
impl RemappingIterator<'_> {
59+
pub fn new<'a>(
60+
pos: usize,
61+
tuples: NullableVec<'a, (usize, Tuple)>,
62+
indices: BumpVec<'a, usize>,
63+
) -> RemappingIterator<'a> {
64+
RemappingIterator {
65+
pos,
66+
tuples,
67+
indices,
68+
}
69+
}
70+
}
71+
5872
impl Iterator for RemappingIterator<'_> {
5973
type Item = Tuple;
6074

6175
fn next(&mut self) -> Option<Self::Item> {
62-
if self.pos > self.tuples.len() - 1 {
76+
if self.pos > self.indices.len() - 1 {
6377
return None;
6478
}
6579
let (_, tuple) = self.tuples.take(self.indices[self.pos]);
@@ -147,11 +161,7 @@ impl SortBy {
147161
}
148162
let indices = radix_sort(sort_keys, arena);
149163

150-
Ok(Box::new(RemappingIterator {
151-
pos: 0,
152-
tuples,
153-
indices,
154-
}))
164+
Ok(Box::new(RemappingIterator::new(0, tuples, indices)))
155165
}
156166
SortBy::Fast => {
157167
let fn_nulls_first = |nulls_first: bool| {
@@ -484,7 +494,7 @@ mod test {
484494
SortField {
485495
expr: ScalarExpression::Reference {
486496
expr: Box::new(ScalarExpression::Empty),
487-
pos: 0,
497+
pos: 1,
488498
},
489499
asc: asc_2,
490500
nulls_first: nulls_first_2,

0 commit comments

Comments
 (0)