Skip to content

Commit 873c016

Browse files
mcheshkovmcheshkov
committed
fix(cubesql): Calculate proper limit and offset for CubeScan in nested limits case
1 parent 8fb14f0 commit 873c016

File tree

3 files changed

+283
-8
lines changed

3 files changed

+283
-8
lines changed

‎rust/cubesql/cubesql/src/compile/rewrite/rules/members.rs‎

Lines changed: 59 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -234,8 +234,8 @@ impl RewriteRules for MemberRules {
234234
"?members",
235235
"?filters",
236236
"?orders",
237-
"?cube_fetch",
238-
"?offset",
237+
"?inner_fetch",
238+
"?inner_skip",
239239
"?split",
240240
"?can_pushdown_join",
241241
"CubeScanWrapped:false",
@@ -254,7 +254,14 @@ impl RewriteRules for MemberRules {
254254
"CubeScanWrapped:false",
255255
"?ungrouped",
256256
),
257-
self.push_down_limit("?skip", "?fetch", "?new_skip", "?new_fetch"),
257+
self.push_down_limit(
258+
"?skip",
259+
"?fetch",
260+
"?inner_skip",
261+
"?inner_fetch",
262+
"?new_skip",
263+
"?new_fetch",
264+
),
258265
),
259266
// MOD function to binary expr
260267
transforming_rewrite_with_root(
@@ -1556,14 +1563,21 @@ impl MemberRules {
15561563
&self,
15571564
skip_var: &'static str,
15581565
fetch_var: &'static str,
1566+
inner_skip_var: &'static str,
1567+
inner_fetch_var: &'static str,
15591568
new_skip_var: &'static str,
15601569
new_fetch_var: &'static str,
15611570
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
15621571
let skip_var = var!(skip_var);
15631572
let fetch_var = var!(fetch_var);
1573+
let inner_skip_var = var!(inner_skip_var);
1574+
let inner_fetch_var = var!(inner_fetch_var);
15641575
let new_skip_var = var!(new_skip_var);
15651576
let new_fetch_var = var!(new_fetch_var);
15661577
move |egraph, subst| {
1578+
// This transform expects only single value in every (eclass, kind)
1579+
// No two different values of fetch or skip should ever get unified
1580+
15671581
let mut skip_value = None;
15681582
for skip in var_iter!(egraph[subst[skip_var]], LimitSkip) {
15691583
if skip.unwrap_or_default() > 0 {
@@ -1579,18 +1593,55 @@ impl MemberRules {
15791593
}
15801594
}
15811595

1582-
if skip_value.is_some() || fetch_value.is_some() {
1596+
let mut inner_skip_value = None;
1597+
for inner_skip in var_iter!(egraph[subst[inner_skip_var]], CubeScanOffset) {
1598+
inner_skip_value = *inner_skip;
1599+
break;
1600+
}
1601+
1602+
let mut inner_fetch_value = None;
1603+
for inner_fetch in var_iter!(egraph[subst[inner_fetch_var]], CubeScanLimit) {
1604+
inner_fetch_value = *inner_fetch;
1605+
break;
1606+
}
1607+
1608+
let new_skip = match (skip_value, inner_skip_value) {
1609+
(None, None) => None,
1610+
(Some(skip), None) | (None, Some(skip)) => Some(skip),
1611+
(Some(outer_skip), Some(inner_skip)) => Some(outer_skip + inner_skip),
1612+
};
1613+
let new_fetch = match (fetch_value, inner_fetch_value) {
1614+
(None, None) => None,
1615+
// Inner node have no limit, maybe just offset, result limit is same as for outer node
1616+
(Some(outer_fetch), None) => Some(outer_fetch),
1617+
// Outer node have no limit, but may have offset
1618+
// First, inner offset would apply
1619+
// Then inner node would limit rows
1620+
// Then outer offset would apply, which would yield no more than `inner_fetch - outer_skip` rows
1621+
(None, Some(inner_fetch)) => {
1622+
Some(inner_fetch.saturating_sub(skip_value.unwrap_or(0)))
1623+
}
1624+
// Both nodes have a limit
1625+
// First, inner offset would apply
1626+
// Then inner node would limit rows
1627+
// Then outer offset would apply, which would yield no more than `in_limit - out_offset` rows
1628+
// Then outer limit would apply, which would yield no more than minimal of two
1629+
(Some(outer_fetch), Some(inner_fetch)) => Some(usize::min(
1630+
inner_fetch.saturating_sub(skip_value.unwrap_or(0)),
1631+
outer_fetch,
1632+
)),
1633+
};
1634+
1635+
if new_skip.is_some() || new_fetch.is_some() {
15831636
subst.insert(
15841637
new_skip_var,
15851638
egraph.add(LogicalPlanLanguage::CubeScanOffset(CubeScanOffset(
1586-
skip_value,
1639+
new_skip,
15871640
))),
15881641
);
15891642
subst.insert(
15901643
new_fetch_var,
1591-
egraph.add(LogicalPlanLanguage::CubeScanLimit(CubeScanLimit(
1592-
fetch_value,
1593-
))),
1644+
egraph.add(LogicalPlanLanguage::CubeScanLimit(CubeScanLimit(new_fetch))),
15941645
);
15951646

15961647
return true;

‎rust/cubesql/cubesql/src/compile/test/mod.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ pub mod test_bi_workarounds;
3030
#[cfg(test)]
3131
pub mod test_cube_join;
3232
#[cfg(test)]
33+
pub mod test_cube_scan;
34+
#[cfg(test)]
3335
pub mod test_df_execution;
3436
#[cfg(test)]
3537
pub mod test_introspection;

‎rust/cubesql/cubesql/src/compile/test/test_cube_scan.rs‎

Lines changed: 222 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,222 @@
1+
use cubeclient::models::V1LoadRequestQuery;
2+
use pretty_assertions::assert_eq;
3+
4+
use crate::compile::{
5+
test::{convert_select_to_query_plan, init_testing_logger, utils::LogicalPlanTestUtils},
6+
DatabaseProtocol,
7+
};
8+
9+
#[tokio::test]
10+
async fn cubescan_limit_offset() {
11+
init_testing_logger();
12+
13+
let query_plan = convert_select_to_query_plan(
14+
// language=PostgreSQL
15+
r#"
16+
SELECT
17+
customer_gender
18+
FROM
19+
KibanaSampleDataEcommerce
20+
GROUP BY
21+
1
22+
LIMIT 2
23+
OFFSET 3
24+
"#
25+
.to_string(),
26+
DatabaseProtocol::PostgreSQL,
27+
)
28+
.await;
29+
30+
let logical_plan = query_plan.as_logical_plan();
31+
assert_eq!(
32+
logical_plan.find_cube_scan().request,
33+
V1LoadRequestQuery {
34+
measures: Some(vec![]),
35+
dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]),
36+
segments: Some(vec![]),
37+
order: Some(vec![]),
38+
limit: Some(2),
39+
offset: Some(3),
40+
..Default::default()
41+
}
42+
);
43+
}
44+
45+
#[tokio::test]
46+
async fn cubescan_limit_limit() {
47+
init_testing_logger();
48+
49+
let variants = vec![
50+
// language=PostgreSQL
51+
r#"
52+
SELECT
53+
customer_gender
54+
FROM (
55+
SELECT
56+
customer_gender
57+
FROM
58+
KibanaSampleDataEcommerce
59+
GROUP BY
60+
1
61+
LIMIT 3
62+
) scan
63+
LIMIT 2
64+
"#,
65+
// language=PostgreSQL
66+
r#"
67+
SELECT
68+
customer_gender
69+
FROM (
70+
SELECT
71+
customer_gender
72+
FROM
73+
KibanaSampleDataEcommerce
74+
GROUP BY
75+
1
76+
LIMIT 2
77+
) scan
78+
LIMIT 3
79+
"#,
80+
];
81+
82+
for variant in variants {
83+
let query_plan =
84+
convert_select_to_query_plan(variant.to_string(), DatabaseProtocol::PostgreSQL).await;
85+
86+
let logical_plan = query_plan.as_logical_plan();
87+
assert_eq!(
88+
logical_plan.find_cube_scan().request,
89+
V1LoadRequestQuery {
90+
measures: Some(vec![]),
91+
dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]),
92+
segments: Some(vec![]),
93+
order: Some(vec![]),
94+
limit: Some(2),
95+
..Default::default()
96+
}
97+
);
98+
}
99+
}
100+
101+
#[tokio::test]
102+
async fn cubescan_offset_offset() {
103+
init_testing_logger();
104+
105+
let variants = vec![
106+
// language=PostgreSQL
107+
r#"
108+
SELECT
109+
customer_gender
110+
FROM (
111+
SELECT
112+
customer_gender
113+
FROM
114+
KibanaSampleDataEcommerce
115+
GROUP BY
116+
1
117+
OFFSET 3
118+
) scan
119+
OFFSET 2
120+
"#,
121+
// language=PostgreSQL
122+
r#"
123+
SELECT
124+
customer_gender
125+
FROM (
126+
SELECT
127+
customer_gender
128+
FROM
129+
KibanaSampleDataEcommerce
130+
GROUP BY
131+
1
132+
OFFSET 2
133+
) scan
134+
OFFSET 3
135+
"#,
136+
];
137+
138+
for variant in variants {
139+
let query_plan =
140+
convert_select_to_query_plan(variant.to_string(), DatabaseProtocol::PostgreSQL).await;
141+
142+
let logical_plan = query_plan.as_logical_plan();
143+
assert_eq!(
144+
logical_plan.find_cube_scan().request,
145+
V1LoadRequestQuery {
146+
measures: Some(vec![]),
147+
dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]),
148+
segments: Some(vec![]),
149+
order: Some(vec![]),
150+
offset: Some(5),
151+
..Default::default()
152+
}
153+
);
154+
}
155+
}
156+
157+
#[tokio::test]
158+
async fn cubescan_limit_offset_limit_offset() {
159+
init_testing_logger();
160+
161+
let variants = vec![
162+
(
163+
// language=PostgreSQL
164+
r#"
165+
SELECT
166+
customer_gender
167+
FROM (
168+
SELECT
169+
customer_gender
170+
FROM
171+
KibanaSampleDataEcommerce
172+
GROUP BY
173+
1
174+
LIMIT 3
175+
OFFSET 3
176+
) scan
177+
LIMIT 2
178+
OFFSET 2
179+
"#,
180+
1,
181+
),
182+
(
183+
// language=PostgreSQL
184+
r#"
185+
SELECT
186+
customer_gender
187+
FROM (
188+
SELECT
189+
customer_gender
190+
FROM
191+
KibanaSampleDataEcommerce
192+
GROUP BY
193+
1
194+
LIMIT 10
195+
OFFSET 3
196+
) scan
197+
LIMIT 2
198+
OFFSET 2
199+
"#,
200+
2,
201+
),
202+
];
203+
204+
for (variant, limit) in variants {
205+
let query_plan =
206+
convert_select_to_query_plan(variant.to_string(), DatabaseProtocol::PostgreSQL).await;
207+
208+
let logical_plan = query_plan.as_logical_plan();
209+
assert_eq!(
210+
logical_plan.find_cube_scan().request,
211+
V1LoadRequestQuery {
212+
measures: Some(vec![]),
213+
dimensions: Some(vec!["KibanaSampleDataEcommerce.customer_gender".to_string()]),
214+
segments: Some(vec![]),
215+
order: Some(vec![]),
216+
limit: Some(limit),
217+
offset: Some(5),
218+
..Default::default()
219+
}
220+
);
221+
}
222+
}

0 commit comments

Comments
 (0)