Skip to content

Commit

Permalink
[BugFix] Fix RewriteUnnestBitmapRule rewrite bug (backport #55168) (#…
Browse files Browse the repository at this point in the history 
…55188)

Co-authored-by: shuming.li <[email protected]>
  • Loading branch information
@mergify @LiShuMing
mergify[bot] and LiShuMing authored Jan 17, 2025
1 parent 396f81c commit fe0b9c2
Show file tree
Hide file tree
Showing 4 changed files with 105 additions and 7 deletions.
20 changes: 13 additions & 7 deletions ...rc/main/java/com/starrocks/sql/optimizer/rule/transformation/RewriteUnnestBitmapRule.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
// limitations under the License.
package com.starrocks.sql.optimizer.rule.transformation;

import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.starrocks.analysis.Expr;
import com.starrocks.catalog.Function;
Expand Down Expand Up @@ -85,14 +86,22 @@ public List<OptExpression> transform(OptExpression input, OptimizerContext conte
return false;
}).map(entry -> new Pair<>(entry.getKey(), entry.getValue())).findFirst().get();

// if bitmap_to_array's output will be used by upper nodes, it's not safe to rewrite
// if bitmap_to_array's output are used by upper nodes, it's not safe to rewrite
if (originalTableFunctionOperator.getOuterColRefs().contains(bitmapToArray.first)) {
return Lists.newArrayList();
}

ColumnRefOperator bitmapColumn = bitmapToArray.second.getColumnRefs().get(0);
columnRefMap.remove(bitmapToArray.first);
columnRefMap.putIfAbsent(bitmapColumn, bitmapColumn);
Preconditions.checkArgument(bitmapToArray.second.getChildren().size() == 1);
// use bitmapToArrayArg to replace bitmapColumn
// eg:
// with r1 as (select b1 as b2 from test_agg),
// r2 as (select sub_bitmap(b1, 0, 10) as b2 from test_agg),
// r3 as (select bitmap_and(t0.b2, t1.b2) as b2 from r1 t0 join r2 t1)
// select unnest as r1 from test_agg, unnest(bitmap_to_array(b2)) order by r1;
// use bitmap_and(t0.b2, t1.b2) as b2 to replace bitmap_to_array(b2)
final ScalarOperator bitmapToArrayArg = bitmapToArray.second.getChild(0);
final ColumnRefOperator bitmapColumn = bitmapToArray.first;
columnRefMap.put(bitmapColumn, bitmapToArrayArg);

TableFunction unnestBitmapFn =
(TableFunction) Expr.getBuiltinFunction(FunctionSet.UNNEST_BITMAP, new Type[] {Type.BITMAP},
Expand All @@ -107,6 +116,3 @@ public List<OptExpression> transform(OptExpression input, OptimizerContext conte
return Lists.newArrayList(OptExpression.create(newTableFunctionOperator, input.inputAt(0)));
}
}



13 changes: 13 additions & 0 deletions fe/fe-core/src/test/java/com/starrocks/sql/plan/TableFunctionTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -276,4 +276,17 @@ public void testRewrite() throws Exception {
assertContains(plan, "tableFunctionName: unnest_bitmap");
assertNotContains(plan, "bitmap_to_array");
}

@Test
public void testUnnesetBitmapToArrayToUnnestBitmapRewrite() throws Exception {
String sql = "with r1 as (select b1 as b2 from test_agg),\n" +
"\t r2 as (select sub_bitmap(b1, 0, 10) as b2 from test_agg),\n" +
"\t r3 as (select bitmap_and(t0.b2, t1.b2) as b2 from r1 t0 join r2 t1)\n" +
"select unnest as r1 from r3, unnest(bitmap_to_array(b2)) order by r1;";
String plan = getFragmentPlan(sql);
PlanTestBase.assertContains(plan, "5:Project\n" +
" | <slot 28> : bitmap_and(10: b1, 25: sub_bitmap)");
PlanTestBase.assertContains(plan, "tableFunctionName: unnest_bitmap");
PlanTestBase.assertNotContains(plan, "bitmap_to_array");
}
}
50 changes: 50 additions & 0 deletions test/sql/test_bitmap_functions/R/test_unnest_bitmap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,4 +42,54 @@ select c1, unnest_bitmap as c3 from t1, unnest_bitmap(c2) order by c1 desc, c3 d
4 18
4 17
4 16
-- !result
set enable_rewrite_unnest_bitmap_to_array=true;
-- result:
-- !result
CREATE TABLE test_tags (
c1 varchar(65533) NOT NULL,
tag_name varchar(65533) NOT NULL,
tag_value varchar(65533) NOT NULL,
rb bitmap NOT NULL
) ENGINE=OLAP
PRIMARY KEY(c1, tag_name, tag_value)
PARTITION BY (c1)
DISTRIBUTED BY HASH(tag_name, tag_value)
PROPERTIES (
"replication_num" = "1"
);
-- result:
-- !result
insert into test_tags(c1, tag_name, tag_value, rb) SELECT '20250114050000','a','57',bitmap_from_string("57");
-- result:
-- !result
insert into test_tags(c1, tag_name, tag_value, rb) SELECT '20250114050000','a','a',bitmap_from_string("57,22253296,29101576,43027104");
-- result:
-- !result
WITH result AS (SELECT rb FROM test_tags),
page_result AS (SELECT sub_bitmap(bitmap_union ( rb ), 0, 20 ) rb FROM result),
all_tag AS (SELECT tag_name, tag_value, c1, bitmap_and (t0.rb, t1.rb) rb2 FROM test_tags t0 JOIN page_result t1 WHERE bitmap_has_any ( t0.rb, t1.rb))
SELECT tag_name, bitmap_to_string(rb2) as row_id_ FROM all_tag order by 1, 2;
-- result:
a 57
a 57,22253296,29101576,43027104
-- !result
WITH result AS (SELECT rb FROM test_tags),
page_result AS (SELECT sub_bitmap(bitmap_union ( rb ), 0, 20 ) rb FROM result),
all_tag AS (SELECT tag_name, tag_value, c1, bitmap_and (t0.rb, t1.rb) rb2 FROM test_tags t0 JOIN page_result t1 WHERE bitmap_has_any ( t0.rb, t1.rb))
SELECT tag_name, bitmap_to_array(rb2) as row_id_ FROM all_tag order by 1, 2;
-- result:
a [57]
a [57,22253296,29101576,43027104]
-- !result
WITH result AS (SELECT rb FROM test_tags),
page_result AS (SELECT sub_bitmap(bitmap_union ( rb ), 0, 20 ) rb FROM result),
all_tag AS (SELECT tag_name, tag_value, c1, bitmap_and (t0.rb, t1.rb) rb2 FROM test_tags t0 JOIN page_result t1 WHERE bitmap_has_any ( t0.rb, t1.rb))
SELECT tag_name, unnest as row_id_ FROM all_tag, unnest (bitmap_to_array(rb2)) order by 1, 2;
-- result:
a 57
a 57
a 22253296
a 29101576
a 43027104
-- !result
29 changes: 29 additions & 0 deletions test/sql/test_bitmap_functions/T/test_unnest_bitmap
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,3 +16,32 @@ insert into t1 select 5, null;
select sum(c1), sum(unnest_bitmap) from t1, unnest_bitmap(c2);
select c1, unnest_bitmap as c3 from t1, unnest_bitmap(c2) order by c1 asc, c3 asc limit 5;
select c1, unnest_bitmap as c3 from t1, unnest_bitmap(c2) order by c1 desc, c3 desc limit 5;

set enable_rewrite_unnest_bitmap_to_array=true;
CREATE TABLE test_tags (
c1 varchar(65533) NOT NULL,
tag_name varchar(65533) NOT NULL,
tag_value varchar(65533) NOT NULL,
rb bitmap NOT NULL
) ENGINE=OLAP
PRIMARY KEY(c1, tag_name, tag_value)
PARTITION BY (c1)
DISTRIBUTED BY HASH(tag_name, tag_value)
PROPERTIES (
"replication_num" = "1"
);
insert into test_tags(c1, tag_name, tag_value, rb) SELECT '20250114050000','a','57',bitmap_from_string("57");
insert into test_tags(c1, tag_name, tag_value, rb) SELECT '20250114050000','a','a',bitmap_from_string("57,22253296,29101576,43027104");

WITH result AS (SELECT rb FROM test_tags),
page_result AS (SELECT sub_bitmap(bitmap_union ( rb ), 0, 20 ) rb FROM result),
all_tag AS (SELECT tag_name, tag_value, c1, bitmap_and (t0.rb, t1.rb) rb2 FROM test_tags t0 JOIN page_result t1 WHERE bitmap_has_any ( t0.rb, t1.rb))
SELECT tag_name, bitmap_to_string(rb2) as row_id_ FROM all_tag order by 1, 2;
WITH result AS (SELECT rb FROM test_tags),
page_result AS (SELECT sub_bitmap(bitmap_union ( rb ), 0, 20 ) rb FROM result),
all_tag AS (SELECT tag_name, tag_value, c1, bitmap_and (t0.rb, t1.rb) rb2 FROM test_tags t0 JOIN page_result t1 WHERE bitmap_has_any ( t0.rb, t1.rb))
SELECT tag_name, bitmap_to_array(rb2) as row_id_ FROM all_tag order by 1, 2;
WITH result AS (SELECT rb FROM test_tags),
page_result AS (SELECT sub_bitmap(bitmap_union ( rb ), 0, 20 ) rb FROM result),
all_tag AS (SELECT tag_name, tag_value, c1, bitmap_and (t0.rb, t1.rb) rb2 FROM test_tags t0 JOIN page_result t1 WHERE bitmap_has_any ( t0.rb, t1.rb))
SELECT tag_name, unnest as row_id_ FROM all_tag, unnest (bitmap_to_array(rb2)) order by 1, 2;

0 comments on commit fe0b9c2

Please sign in to comment.