Skip to content

Commit

Permalink
[feature](function) support orthogonal_bitmap_expr_calculate & ortho…
Browse files Browse the repository at this point in the history
…gonal_bitmap_expr_calculate_count for nereids (#44991) (#45012)

support orthogonal_bitmap_expr_calculate & orthogonal_bitmap_expr_calculate_count for nereids
  • Loading branch information
924060929 authored Dec 7, 2024
1 parent 358bcea commit 2543af8
Show file tree
Hide file tree
Showing 9 changed files with 271 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -228,7 +228,7 @@ struct AggOrthBitmapExprCalBaseData {
void init_add_key(const IColumn** columns, size_t row_num, int argument_size) {
if (first_init) {
DCHECK(argument_size > 1);
const auto& col = assert_cast<const ColVecData&>(*columns[2]);
const auto& col = assert_cast<const ColumnString&>(*columns[2]);
std::string expr = col.get_data_at(row_num).to_string();
bitmap_expr_cal.bitmap_calculation_init(expr);
first_init = false;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@
import org.apache.doris.nereids.trees.expressions.functions.agg.MultiDistinctSum;
import org.apache.doris.nereids.trees.expressions.functions.agg.MultiDistinctSum0;
import org.apache.doris.nereids.trees.expressions.functions.agg.Ndv;
import org.apache.doris.nereids.trees.expressions.functions.agg.OrthogonalBitmapExprCalculate;
import org.apache.doris.nereids.trees.expressions.functions.agg.OrthogonalBitmapExprCalculateCount;
import org.apache.doris.nereids.trees.expressions.functions.agg.OrthogonalBitmapIntersect;
import org.apache.doris.nereids.trees.expressions.functions.agg.OrthogonalBitmapIntersectCount;
import org.apache.doris.nereids.trees.expressions.functions.agg.OrthogonalBitmapUnionCount;
Expand Down Expand Up @@ -124,6 +126,8 @@ public class BuiltinAggregateFunctions implements FunctionHelper {
agg(MultiDistinctSum.class, "multi_distinct_sum"),
agg(MultiDistinctSum0.class, "multi_distinct_sum0"),
agg(Ndv.class, "approx_count_distinct", "ndv"),
agg(OrthogonalBitmapExprCalculate.class, "orthogonal_bitmap_expr_calculate"),
agg(OrthogonalBitmapExprCalculateCount.class, "orthogonal_bitmap_expr_calculate_count"),
agg(OrthogonalBitmapIntersect.class, "orthogonal_bitmap_intersect"),
agg(OrthogonalBitmapIntersectCount.class, "orthogonal_bitmap_intersect_count"),
agg(OrthogonalBitmapUnionCount.class, "orthogonal_bitmap_union_count"),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
import org.apache.doris.nereids.trees.expressions.functions.ExpressionTrait;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateParam;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregatePhase;
import org.apache.doris.nereids.trees.expressions.functions.agg.Count;
import org.apache.doris.nereids.trees.expressions.functions.agg.GroupConcat;
import org.apache.doris.nereids.trees.expressions.functions.agg.Max;
Expand Down Expand Up @@ -285,49 +286,58 @@ public List<Rule> buildRules() {
RuleType.ONE_PHASE_AGGREGATE_WITHOUT_DISTINCT.build(
basePattern
.when(agg -> agg.getDistinctArguments().isEmpty())
.when(agg -> agg.supportAggregatePhase(AggregatePhase.ONE))
.thenApplyMulti(ctx -> onePhaseAggregateWithoutDistinct(ctx.root, ctx.connectContext))
),
RuleType.TWO_PHASE_AGGREGATE_WITHOUT_DISTINCT.build(
basePattern
.when(agg -> agg.getDistinctArguments().isEmpty())
.when(agg -> agg.supportAggregatePhase(AggregatePhase.TWO))
.thenApplyMulti(ctx -> twoPhaseAggregateWithoutDistinct(ctx.root, ctx.connectContext))
),
// RuleType.TWO_PHASE_AGGREGATE_WITH_COUNT_DISTINCT_MULTI.build(
// basePattern
// .when(this::containsCountDistinctMultiExpr)
// .when(agg -> agg.supportAggregatePhase(AggregatePhase.TWO))
// .thenApplyMulti(ctx -> twoPhaseAggregateWithCountDistinctMulti(ctx.root, ctx.cascadesContext))
// ),
RuleType.THREE_PHASE_AGGREGATE_WITH_COUNT_DISTINCT_MULTI.build(
basePattern
.when(this::containsCountDistinctMultiExpr)
.when(agg -> agg.supportAggregatePhase(AggregatePhase.THREE))
.thenApplyMulti(ctx -> threePhaseAggregateWithCountDistinctMulti(ctx.root, ctx.cascadesContext))
),
RuleType.ONE_PHASE_AGGREGATE_SINGLE_DISTINCT_TO_MULTI.build(
basePattern
.when(agg -> agg.getDistinctArguments().size() == 1 && couldConvertToMulti(agg))
.when(agg -> agg.supportAggregatePhase(AggregatePhase.ONE))
.thenApplyMulti(ctx -> onePhaseAggregateWithMultiDistinct(ctx.root, ctx.connectContext))
),
RuleType.TWO_PHASE_AGGREGATE_SINGLE_DISTINCT_TO_MULTI.build(
basePattern
.when(agg -> agg.getDistinctArguments().size() == 1 && couldConvertToMulti(agg))
.when(agg -> agg.supportAggregatePhase(AggregatePhase.TWO))
.thenApplyMulti(ctx -> twoPhaseAggregateWithMultiDistinct(ctx.root, ctx.connectContext))
),
RuleType.TWO_PHASE_AGGREGATE_WITH_MULTI_DISTINCT.build(
basePattern
.when(agg -> agg.getDistinctArguments().size() > 1
&& !containsCountDistinctMultiExpr(agg)
&& couldConvertToMulti(agg))
.when(agg -> agg.supportAggregatePhase(AggregatePhase.TWO))
.thenApplyMulti(ctx -> twoPhaseAggregateWithMultiDistinct(ctx.root, ctx.connectContext))
),
// RuleType.TWO_PHASE_AGGREGATE_WITH_DISTINCT.build(
// basePattern
// .when(agg -> agg.getDistinctArguments().size() == 1)
// .when(agg -> agg.supportAggregatePhase(AggregatePhase.TWO))
// .thenApplyMulti(ctx -> twoPhaseAggregateWithDistinct(ctx.root, ctx.connectContext))
// ),
RuleType.THREE_PHASE_AGGREGATE_WITH_DISTINCT.build(
basePattern
.when(agg -> agg.getDistinctArguments().size() == 1)
.whenNot(agg -> agg.mustUseMultiDistinctAgg())
.when(agg -> agg.supportAggregatePhase(AggregatePhase.THREE))
.thenApplyMulti(ctx -> threePhaseAggregateWithDistinct(ctx.root, ctx.connectContext))
),
/*
Expand All @@ -352,6 +362,7 @@ && couldConvertToMulti(agg))
.when(agg -> agg.getDistinctArguments().size() == 1)
.when(agg -> agg.getGroupByExpressions().isEmpty())
.whenNot(agg -> agg.mustUseMultiDistinctAgg())
.when(agg -> agg.supportAggregatePhase(AggregatePhase.FOUR))
.thenApplyMulti(ctx -> {
Function<List<Expression>, RequireProperties> secondPhaseRequireDistinctHash =
groupByAndDistinct -> RequireProperties.of(
Expand Down Expand Up @@ -399,6 +410,7 @@ && couldConvertToMulti(agg))
}
return couldConvertToMulti(agg);
})
.when(agg -> agg.supportAggregatePhase(AggregatePhase.FOUR))
.thenApplyMulti(ctx -> {
Function<List<Expression>, RequireProperties> secondPhaseRequireGroupByAndDistinctHash =
groupByAndDistinct -> RequireProperties.of(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -131,6 +131,10 @@ public String toString() {
return getName() + "(" + (distinct ? "DISTINCT " : "") + args + ")";
}

public boolean supportAggregatePhase(AggregatePhase aggregatePhase) {
return true;
}

public List<Expression> getDistinctArguments() {
return distinct ? getArguments() : ImmutableList.of();
}
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

package org.apache.doris.nereids.trees.expressions.functions.agg;

/** AggregatePhase */
public enum AggregatePhase {
ONE, TWO, THREE, FOUR
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

package org.apache.doris.nereids.trees.expressions.functions.agg;

import org.apache.doris.catalog.FunctionSignature;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable;
import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
import org.apache.doris.nereids.trees.expressions.literal.VarcharLiteral;
import org.apache.doris.nereids.types.BitmapType;
import org.apache.doris.nereids.types.VarcharType;
import org.apache.doris.nereids.types.coercion.CharacterType;
import org.apache.doris.nereids.util.ExpressionUtils;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;

import java.util.List;

/** OrthogonalBitmapExprCalculate */
public class OrthogonalBitmapExprCalculate extends AggregateFunction
implements AlwaysNotNullable, OrthogonalBitmapFunction, ExplicitlyCastableSignature {

static final List<FunctionSignature> FUNCTION_SIGNATURES = ImmutableList.of(
FunctionSignature.ret(BitmapType.INSTANCE)
.varArgs(BitmapType.INSTANCE, VarcharType.SYSTEM_DEFAULT, VarcharType.SYSTEM_DEFAULT)
);

/**
* constructor with 3 arguments.
*/
public OrthogonalBitmapExprCalculate(
Expression bitmap, Expression filterColumn, VarcharLiteral inputString) {
super("orthogonal_bitmap_expr_calculate", ExpressionUtils.mergeArguments(bitmap, filterColumn, inputString));
}

/**
* constructor with 3 arguments.
*/
public OrthogonalBitmapExprCalculate(boolean distinct,
Expression bitmap, Expression filterColumn, VarcharLiteral inputString) {
super("orthogonal_bitmap_expr_calculate", distinct,
ExpressionUtils.mergeArguments(bitmap, filterColumn, inputString));
}

@Override
public boolean supportAggregatePhase(AggregatePhase aggregatePhase) {
return aggregatePhase == AggregatePhase.TWO;
}

@Override
public OrthogonalBitmapExprCalculate withDistinctAndChildren(boolean distinct, List<Expression> children) {
Preconditions.checkArgument(children.size() == 3
&& children.get(2).getDataType() instanceof CharacterType
&& children.get(2).getDataType() instanceof VarcharType);
return new OrthogonalBitmapExprCalculate(
distinct, children.get(0), children.get(1), (VarcharLiteral) children.get(2));
}

@Override
public List<FunctionSignature> getSignatures() {
return FUNCTION_SIGNATURES;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

package org.apache.doris.nereids.trees.expressions.functions.agg;

import org.apache.doris.catalog.FunctionSignature;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.functions.AlwaysNotNullable;
import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
import org.apache.doris.nereids.trees.expressions.literal.VarcharLiteral;
import org.apache.doris.nereids.types.BigIntType;
import org.apache.doris.nereids.types.BitmapType;
import org.apache.doris.nereids.types.VarcharType;
import org.apache.doris.nereids.types.coercion.CharacterType;
import org.apache.doris.nereids.util.ExpressionUtils;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;

import java.util.List;

/** OrthogonalBitmapExprCalculateCount */
public class OrthogonalBitmapExprCalculateCount extends AggregateFunction
implements AlwaysNotNullable, OrthogonalBitmapFunction, ExplicitlyCastableSignature {

static final List<FunctionSignature> FUNCTION_SIGNATURES = ImmutableList.of(
FunctionSignature.ret(BigIntType.INSTANCE)
.varArgs(BitmapType.INSTANCE, VarcharType.SYSTEM_DEFAULT, VarcharType.SYSTEM_DEFAULT)
);

/**
* constructor with 3 arguments.
*/
public OrthogonalBitmapExprCalculateCount(
Expression bitmap, Expression filterColumn, VarcharLiteral inputString) {
super("orthogonal_bitmap_expr_calculate_count",
ExpressionUtils.mergeArguments(bitmap, filterColumn, inputString));
}

/**
* constructor with 3 arguments.
*/
public OrthogonalBitmapExprCalculateCount(boolean distinct,
Expression bitmap, Expression filterColumn, VarcharLiteral inputString) {
super("orthogonal_bitmap_expr_calculate_count", distinct,
ExpressionUtils.mergeArguments(bitmap, filterColumn, inputString));
}

@Override
public boolean supportAggregatePhase(AggregatePhase aggregatePhase) {
return aggregatePhase == AggregatePhase.TWO;
}

@Override
public OrthogonalBitmapExprCalculateCount withDistinctAndChildren(boolean distinct, List<Expression> children) {
Preconditions.checkArgument(children.size() == 3
&& children.get(2).getDataType() instanceof CharacterType
&& children.get(2).getDataType() instanceof VarcharType);
return new OrthogonalBitmapExprCalculateCount(
distinct, children.get(0), children.get(1), (VarcharLiteral) children.get(2));
}

@Override
public List<FunctionSignature> getSignatures() {
return FUNCTION_SIGNATURES;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,8 @@
import org.apache.doris.nereids.trees.expressions.Slot;
import org.apache.doris.nereids.trees.expressions.SlotReference;
import org.apache.doris.nereids.trees.expressions.functions.ExpressionTrait;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
import org.apache.doris.nereids.trees.expressions.functions.agg.AggregatePhase;
import org.apache.doris.nereids.trees.expressions.functions.agg.Count;
import org.apache.doris.nereids.trees.expressions.functions.agg.Ndv;
import org.apache.doris.nereids.trees.plans.Plan;
Expand Down Expand Up @@ -402,4 +404,14 @@ public ImmutableSet<FdItem> computeFdItems() {
public void computeEqualSet(FunctionalDependencies.Builder fdBuilder) {
fdBuilder.addEqualSet(child().getLogicalProperties().getFunctionalDependencies());
}

/** supportAggregatePhase */
public boolean supportAggregatePhase(AggregatePhase aggregatePhase) {
for (AggregateFunction aggregateFunction : getAggregateFunctions()) {
if (!aggregateFunction.supportAggregatePhase(aggregatePhase)) {
return false;
}
}
return true;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.

suite("test_orthogonal_bitmap_expr_calculate") {
multi_sql """
drop table if exists test_orthogonal_bitmap_expr_calculate;
create table test_orthogonal_bitmap_expr_calculate(
id int,
tag int,
user_id bitmap bitmap_union
)
aggregate key(id, tag)
distributed by hash(id) buckets 1
properties(
'replication_num'='1'
);
insert into test_orthogonal_bitmap_expr_calculate values
(1, 100, bitmap_from_string('1,2,3,4,5')),
(1, 200, bitmap_from_string('3,4,5,6,7'));
set enable_fallback_to_original_planner=false;
"""

test {
sql """
select bitmap_to_string(orthogonal_bitmap_expr_calculate(user_id, tag, '(100&200)'))
from test_orthogonal_bitmap_expr_calculate
"""
result([['3,4,5']])
}

test {
sql """
select orthogonal_bitmap_expr_calculate_count(user_id, tag, '(100&200)')
from test_orthogonal_bitmap_expr_calculate
"""
result([[3L]])
}
}

0 comments on commit 2543af8

Please sign in to comment.