Skip to content

Commit 9564db0

Browse files
chenzl25chenzl25
authored
feat(batch): Support index selection for sort aggregation with a descending ordering (risingwavelabs#8515)
1 parent 97b021d commit 9564db0

File tree

5 files changed

+164
-176
lines changed

5 files changed

+164
-176
lines changed

‎src/frontend/planner_test/tests/testdata/agg.yaml

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1230,3 +1230,23 @@
12301230
└─StreamHashAgg { group_key: [idx.col1, $expr1], aggs: [count] }
12311231
└─StreamProject { exprs: [idx.col1, idx.id, Vnode(idx.id) as $expr1] }
12321232
└─StreamTableScan { table: idx, columns: [idx.col1, idx.id], pk: [idx.id], dist: UpstreamHashShard(idx.id) }
1233+
- name: sort agg on an ascending index
1234+
sql: |
1235+
create table t (a int, b int);
1236+
create index idx_asc on t(a asc);
1237+
create index idx_desc on t(a desc);
1238+
select a, count(*) cnt from t group by a order by a asc;
1239+
batch_plan: |
1240+
BatchExchange { order: [idx_asc.a ASC], dist: Single }
1241+
└─BatchSortAgg { group_key: [idx_asc.a], aggs: [count] }
1242+
└─BatchScan { table: idx_asc, columns: [idx_asc.a], distribution: UpstreamHashShard(idx_asc.a) }
1243+
- name: sort agg on a descending index
1244+
sql: |
1245+
create table t (a int, b int);
1246+
create index idx_asc on t(a asc);
1247+
create index idx_desc on t(a desc);
1248+
select a, count(*) cnt from t group by a order by a desc;
1249+
batch_plan: |
1250+
BatchExchange { order: [idx_desc.a DESC], dist: Single }
1251+
└─BatchSortAgg { group_key: [idx_desc.a], aggs: [count] }
1252+
└─BatchScan { table: idx_desc, columns: [idx_desc.a], distribution: UpstreamHashShard(idx_desc.a) }

‎src/frontend/src/catalog/index_catalog.rs

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -122,12 +122,14 @@ impl IndexCatalog {
122122
self.index_table.columns.len() == self.primary_table.columns.len()
123123
}
124124

125-
/// a mapping maps column index of secondary index to column index of primary table
125+
/// A mapping maps the column index of the secondary index to the column index of the primary
126+
/// table.
126127
pub fn secondary_to_primary_mapping(&self) -> &BTreeMap<usize, usize> {
127128
&self.secondary_to_primary_mapping
128129
}
129130

130-
/// a mapping maps column index of primary table to column index of secondary index
131+
/// A mapping maps the column index of the primary table to the column index of the secondary
132+
/// index.
131133
pub fn primary_to_secondary_mapping(&self) -> &BTreeMap<usize, usize> {
132134
&self.primary_to_secondary_mapping
133135
}

‎src/frontend/src/optimizer/plan_node/logical_scan.rs

Lines changed: 66 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use fixedbitset::FixedBitSet;
2020
use itertools::Itertools;
2121
use risingwave_common::catalog::{ColumnDesc, Field, Schema, TableDesc};
2222
use risingwave_common::error::{ErrorCode, Result, RwError};
23-
use risingwave_common::util::sort_util::{ColumnOrder, OrderType};
23+
use risingwave_common::util::sort_util::ColumnOrder;
2424

2525
use super::generic::{GenericPlanNode, GenericPlanRef};
2626
use super::{
@@ -267,6 +267,65 @@ impl LogicalScan {
267267
self.i2o_col_mapping().rewrite_bitset(watermark_columns)
268268
}
269269

270+
/// Return indexes can satisfy the required order.
271+
pub fn indexes_satisfy_order(&self, required_order: &Order) -> Vec<&Rc<IndexCatalog>> {
272+
let output_col_map = self
273+
.output_col_idx()
274+
.iter()
275+
.cloned()
276+
.enumerate()
277+
.map(|(id, col)| (col, id))
278+
.collect::<BTreeMap<_, _>>();
279+
let unmatched_idx = output_col_map.len();
280+
self.indexes()
281+
.iter()
282+
.filter(|idx| {
283+
let s2p_mapping = idx.secondary_to_primary_mapping();
284+
Order {
285+
column_orders: idx
286+
.index_table
287+
.pk()
288+
.iter()
289+
.map(|idx_item| {
290+
ColumnOrder::new(
291+
*output_col_map
292+
.get(
293+
s2p_mapping
294+
.get(&idx_item.column_index)
295+
.expect("should be in s2p mapping"),
296+
)
297+
.unwrap_or(&unmatched_idx),
298+
idx_item.order_type,
299+
)
300+
})
301+
.collect(),
302+
}
303+
.satisfies(required_order)
304+
})
305+
.collect()
306+
}
307+
308+
/// If the index can cover the scan, transform it to the index scan.
309+
pub fn to_index_scan_if_index_covered(&self, index: &Rc<IndexCatalog>) -> Option<LogicalScan> {
310+
let p2s_mapping = index.primary_to_secondary_mapping();
311+
if self
312+
.required_col_idx()
313+
.iter()
314+
.all(|x| p2s_mapping.contains_key(x))
315+
{
316+
let index_scan = self.to_index_scan(
317+
&index.name,
318+
index.index_table.table_desc().into(),
319+
p2s_mapping,
320+
);
321+
Some(index_scan)
322+
} else {
323+
None
324+
}
325+
}
326+
327+
/// Prerequisite: the caller should guarantee that `primary_to_secondary_mapping` must cover the
328+
/// scan.
270329
pub fn to_index_scan(
271330
&self,
272331
index_name: &str,
@@ -581,32 +640,14 @@ impl LogicalScan {
581640
return None;
582641
}
583642

584-
let index = self.indexes().iter().find(|idx| {
585-
Order {
586-
column_orders: idx
587-
.index_item
588-
.iter()
589-
.map(|idx_item| ColumnOrder::new(idx_item.index, OrderType::ascending()))
590-
.collect(),
643+
let order_satisfied_index = self.indexes_satisfy_order(required_order);
644+
for index in order_satisfied_index {
645+
if let Some(index_scan) = self.to_index_scan_if_index_covered(index) {
646+
return Some(index_scan.to_batch());
591647
}
592-
.satisfies(required_order)
593-
})?;
594-
595-
let p2s_mapping = index.primary_to_secondary_mapping();
596-
if self
597-
.required_col_idx()
598-
.iter()
599-
.all(|x| p2s_mapping.contains_key(x))
600-
{
601-
let index_scan = self.to_index_scan(
602-
&index.name,
603-
index.index_table.table_desc().into(),
604-
p2s_mapping,
605-
);
606-
Some(index_scan.to_batch())
607-
} else {
608-
None
609648
}
649+
650+
None
610651
}
611652
}
612653

‎src/frontend/src/optimizer/rule/min_max_on_index_rule.rs

Lines changed: 49 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -57,13 +57,6 @@ impl Rule for MinMaxOnIndexRule {
5757
if !logical_scan.predicate().always_true() {
5858
return None;
5959
}
60-
let output_col_map = logical_scan
61-
.output_col_idx()
62-
.iter()
63-
.cloned()
64-
.enumerate()
65-
.map(|(id, col)| (col, id))
66-
.collect::<BTreeMap<_, _>>();
6760
let order = Order {
6861
column_orders: vec![ColumnOrder::new(
6962
calls.first()?.inputs.first()?.index(),
@@ -74,12 +67,10 @@ impl Rule for MinMaxOnIndexRule {
7467
},
7568
)],
7669
};
77-
if let Some(p) =
78-
self.try_on_index(logical_agg, logical_scan.clone(), &order, &output_col_map)
79-
{
70+
if let Some(p) = self.try_on_index(logical_agg, logical_scan.clone(), &order) {
8071
Some(p)
8172
} else {
82-
self.try_on_pk(logical_agg, logical_scan, &order, &output_col_map)
73+
self.try_on_pk(logical_agg, logical_scan, &order)
8374
}
8475
} else {
8576
None
@@ -96,93 +87,64 @@ impl MinMaxOnIndexRule {
9687
&self,
9788
logical_agg: &LogicalAgg,
9889
logical_scan: LogicalScan,
99-
order: &Order,
100-
output_col_map: &BTreeMap<usize, usize>,
90+
required_order: &Order,
10191
) -> Option<PlanRef> {
102-
let unmatched_idx = output_col_map.len();
103-
let index = logical_scan.indexes().iter().find(|idx| {
104-
let s2p_mapping = idx.secondary_to_primary_mapping();
105-
Order {
106-
column_orders: idx
107-
.index_table
108-
.pk()
109-
.iter()
110-
.map(|idx_item| {
111-
ColumnOrder::new(
112-
*output_col_map
113-
.get(
114-
s2p_mapping
115-
.get(&idx_item.column_index)
116-
.expect("should be in s2p mapping"),
117-
)
118-
.unwrap_or(&unmatched_idx),
119-
idx_item.order_type,
120-
)
121-
})
122-
.collect(),
92+
let order_satisfied_index = logical_scan.indexes_satisfy_order(required_order);
93+
for index in order_satisfied_index {
94+
if let Some(index_scan) = logical_scan.to_index_scan_if_index_covered(index) {
95+
let non_null_filter = LogicalFilter::create_with_expr(
96+
index_scan.into(),
97+
FunctionCall::new_unchecked(
98+
ExprType::IsNotNull,
99+
vec![ExprImpl::InputRef(Box::new(InputRef::new(
100+
0,
101+
logical_agg.schema().fields[0].data_type.clone(),
102+
)))],
103+
DataType::Boolean,
104+
)
105+
.into(),
106+
);
107+
108+
let limit = LogicalLimit::create(non_null_filter, 1, 0);
109+
110+
let formatting_agg = LogicalAgg::new(
111+
vec![PlanAggCall {
112+
agg_kind: logical_agg.agg_calls().first()?.agg_kind,
113+
return_type: logical_agg.schema().fields[0].data_type.clone(),
114+
inputs: vec![InputRef::new(
115+
0,
116+
logical_agg.schema().fields[0].data_type.clone(),
117+
)],
118+
order_by: vec![],
119+
distinct: false,
120+
filter: Condition {
121+
conjunctions: vec![],
122+
},
123+
}],
124+
vec![],
125+
limit,
126+
);
127+
128+
return Some(formatting_agg.into());
123129
}
124-
.satisfies(order)
125-
})?;
126-
127-
let p2s_mapping = index.primary_to_secondary_mapping();
130+
}
128131

129-
let index_scan = if logical_scan
130-
.required_col_idx()
131-
.iter()
132-
.all(|x| p2s_mapping.contains_key(x))
133-
{
134-
Some(logical_scan.to_index_scan(
135-
&index.name,
136-
index.index_table.table_desc().into(),
137-
p2s_mapping,
138-
))
139-
} else {
140-
None
141-
}?;
142-
143-
let non_null_filter = LogicalFilter::create_with_expr(
144-
index_scan.into(),
145-
FunctionCall::new_unchecked(
146-
ExprType::IsNotNull,
147-
vec![ExprImpl::InputRef(Box::new(InputRef::new(
148-
0,
149-
logical_agg.schema().fields[0].data_type.clone(),
150-
)))],
151-
DataType::Boolean,
152-
)
153-
.into(),
154-
);
155-
156-
let limit = LogicalLimit::create(non_null_filter, 1, 0);
157-
158-
let formatting_agg = LogicalAgg::new(
159-
vec![PlanAggCall {
160-
agg_kind: logical_agg.agg_calls().first()?.agg_kind,
161-
return_type: logical_agg.schema().fields[0].data_type.clone(),
162-
inputs: vec![InputRef::new(
163-
0,
164-
logical_agg.schema().fields[0].data_type.clone(),
165-
)],
166-
order_by: vec![],
167-
distinct: false,
168-
filter: Condition {
169-
conjunctions: vec![],
170-
},
171-
}],
172-
vec![],
173-
limit,
174-
);
175-
176-
Some(formatting_agg.into())
132+
None
177133
}
178134

179135
fn try_on_pk(
180136
&self,
181137
logical_agg: &LogicalAgg,
182138
logical_scan: LogicalScan,
183139
order: &Order,
184-
output_col_map: &BTreeMap<usize, usize>,
185140
) -> Option<PlanRef> {
141+
let output_col_map = logical_scan
142+
.output_col_idx()
143+
.iter()
144+
.cloned()
145+
.enumerate()
146+
.map(|(id, col)| (col, id))
147+
.collect::<BTreeMap<_, _>>();
186148
let unmatched_idx = output_col_map.len();
187149
let primary_key = logical_scan.primary_key();
188150
let primary_key_order = Order {

0 commit comments

Comments
 (0)