Skip to content

Commit 1d20d1a

Browse files
mcheshkovmcheshkov
committed
fix(cubesql): Calculate proper limit and offset for CubeScan in nested limits case
1 parent 68a6c46 commit 1d20d1a

File tree

3 files changed

+305
-27
lines changed

3 files changed

+305
-27
lines changed

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

Lines changed: 77 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -240,8 +240,8 @@ impl RewriteRules for MemberRules {
240240
"?members",
241241
"?filters",
242242
"?orders",
243-
"?cube_fetch",
244-
"?offset",
243+
"?inner_fetch",
244+
"?inner_skip",
245245
"?split",
246246
"?can_pushdown_join",
247247
"CubeScanWrapped:false",
@@ -260,7 +260,14 @@ impl RewriteRules for MemberRules {
260260
"CubeScanWrapped:false",
261261
"?ungrouped",
262262
),
263-
self.push_down_limit("?skip", "?fetch", "?new_skip", "?new_fetch"),
263+
self.push_down_limit(
264+
"?skip",
265+
"?fetch",
266+
"?inner_skip",
267+
"?inner_fetch",
268+
"?new_skip",
269+
"?new_fetch",
270+
),
264271
),
265272
transforming_rewrite(
266273
"select-distinct-dimensions",
@@ -1694,47 +1701,90 @@ impl MemberRules {
16941701
&self,
16951702
skip_var: &'static str,
16961703
fetch_var: &'static str,
1704+
inner_skip_var: &'static str,
1705+
inner_fetch_var: &'static str,
16971706
new_skip_var: &'static str,
16981707
new_fetch_var: &'static str,
16991708
) -> impl Fn(&mut CubeEGraph, &mut Subst) -> bool {
17001709
let skip_var = var!(skip_var);
17011710
let fetch_var = var!(fetch_var);
1711+
let inner_skip_var = var!(inner_skip_var);
1712+
let inner_fetch_var = var!(inner_fetch_var);
17021713
let new_skip_var = var!(new_skip_var);
17031714
let new_fetch_var = var!(new_fetch_var);
17041715
move |egraph, subst| {
1716+
// This transform expects only single value in every (eclass, kind)
1717+
// No two different values of fetch or skip should ever get unified
1718+
17051719
let mut skip_value = None;
17061720
for skip in var_iter!(egraph[subst[skip_var]], LimitSkip) {
1707-
if skip.unwrap_or_default() > 0 {
1708-
skip_value = *skip;
1709-
break;
1710-
}
1721+
skip_value = *skip;
1722+
break;
17111723
}
17121724
let mut fetch_value = None;
17131725
for fetch in var_iter!(egraph[subst[fetch_var]], LimitFetch) {
1714-
if fetch.unwrap_or_default() > 0 {
1715-
fetch_value = *fetch;
1716-
break;
1717-
}
1726+
fetch_value = *fetch;
1727+
break;
1728+
}
1729+
// TODO support this case
1730+
if fetch_value == Some(0) {
1731+
// Broken and unsupported case for now
1732+
return false;
17181733
}
17191734

1720-
if skip_value.is_some() || fetch_value.is_some() {
1721-
subst.insert(
1722-
new_skip_var,
1723-
egraph.add(LogicalPlanLanguage::CubeScanOffset(CubeScanOffset(
1724-
skip_value,
1725-
))),
1726-
);
1727-
subst.insert(
1728-
new_fetch_var,
1729-
egraph.add(LogicalPlanLanguage::CubeScanLimit(CubeScanLimit(
1730-
fetch_value,
1731-
))),
1732-
);
1733-
1734-
return true;
1735+
let mut inner_skip_value = None;
1736+
for inner_skip in var_iter!(egraph[subst[inner_skip_var]], CubeScanOffset) {
1737+
inner_skip_value = *inner_skip;
1738+
break;
17351739
}
17361740

1737-
false
1741+
let mut inner_fetch_value = None;
1742+
for inner_fetch in var_iter!(egraph[subst[inner_fetch_var]], CubeScanLimit) {
1743+
inner_fetch_value = *inner_fetch;
1744+
break;
1745+
}
1746+
1747+
let new_skip = match (skip_value, inner_skip_value) {
1748+
(None, None) => None,
1749+
(Some(skip), None) | (None, Some(skip)) => Some(skip),
1750+
(Some(outer_skip), Some(inner_skip)) => Some(outer_skip + inner_skip),
1751+
};
1752+
// No need to set offset=0, it's same as no offset
1753+
let new_skip = if new_skip != Some(0) { new_skip } else { None };
1754+
let new_fetch = match (fetch_value, inner_fetch_value) {
1755+
(None, None) => None,
1756+
// Inner node have no limit, maybe just offset, result limit is same as for outer node
1757+
(Some(outer_fetch), None) => Some(outer_fetch),
1758+
// Outer node have no limit, but may have offset
1759+
// First, inner offset would apply
1760+
// Then inner node would limit rows
1761+
// Then outer offset would apply, which would yield no more than `inner_fetch - outer_skip` rows
1762+
(None, Some(inner_fetch)) => {
1763+
Some(inner_fetch.saturating_sub(skip_value.unwrap_or(0)))
1764+
}
1765+
// Both nodes have a limit
1766+
// First, inner offset would apply
1767+
// Then inner node would limit rows
1768+
// Then outer offset would apply, which would yield no more than `in_limit - out_offset` rows
1769+
// Then outer limit would apply, which would yield no more than minimal of two
1770+
(Some(outer_fetch), Some(inner_fetch)) => Some(usize::min(
1771+
inner_fetch.saturating_sub(skip_value.unwrap_or(0)),
1772+
outer_fetch,
1773+
)),
1774+
};
1775+
1776+
subst.insert(
1777+
new_skip_var,
1778+
egraph.add(LogicalPlanLanguage::CubeScanOffset(CubeScanOffset(
1779+
new_skip,
1780+
))),
1781+
);
1782+
subst.insert(
1783+
new_fetch_var,
1784+
egraph.add(LogicalPlanLanguage::CubeScanLimit(CubeScanLimit(new_fetch))),
1785+
);
1786+
1787+
true
17381788
}
17391789
}
17401790

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ pub mod test_cube_join;
3232
#[cfg(test)]
3333
pub mod test_cube_join_grouped;
3434
#[cfg(test)]
35+
pub mod test_cube_scan;
36+
#[cfg(test)]
3537
pub mod test_df_execution;
3638
#[cfg(test)]
3739
pub mod test_filters;

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

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

0 commit comments

Comments
 (0)