Skip to content

Commit a90f44a

Browse files
tkalkirilltkalkirill
authored
IGNITE-28813 Add the ability to add aggregate functions to Calcite via plugins (#13272)
1 parent d8b6b2d commit a90f44a

4 files changed

Lines changed: 156 additions & 10 deletions

File tree

‎modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/exp/agg/AccumulatorFactoryProvider.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package org.apache.ignite.internal.processors.query.calcite.exec.exp.agg;
19+
20+
import java.util.function.Supplier;
21+
import org.apache.calcite.plan.Context;
22+
import org.apache.calcite.rel.core.AggregateCall;
23+
import org.apache.calcite.tools.Frameworks;
24+
import org.apache.ignite.internal.processors.query.calcite.exec.ExecutionContext;
25+
import org.apache.ignite.plugin.PluginProvider;
26+
import org.jetbrains.annotations.Nullable;
27+
28+
/**
29+
* Factory that selects and creates an accumulator supplier for an aggregate call. Allows overriding standard aggregate
30+
* functions.
31+
*
32+
* <p>It can be set via {@link PluginProvider} when creating a configuration using
33+
* {@link PluginProvider#createComponent} via {@link Frameworks.ConfigBuilder#context(Context)}.</p>
34+
*/
35+
@FunctionalInterface
36+
public interface AccumulatorFactoryProvider {
37+
/** @return Accumulator supplier, {@code null} if no accumulator is required for this aggregate call. */
38+
@Nullable <Row> Supplier<Accumulator<Row>> factory(AggregateCall call, ExecutionContext<Row> ctx);
39+
}

‎modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/exp/agg/Accumulators.java‎

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,15 @@ private static <Row> Supplier<Accumulator<Row>> accumulatorFunctionFactory(
7171
) {
7272
RowHandler<Row> hnd = ctx.rowHandler();
7373

74+
AccumulatorFactoryProvider prov = ctx.unwrap(AccumulatorFactoryProvider.class);
75+
76+
if (prov != null) {
77+
Supplier<Accumulator<Row>> fac = prov.factory(call, ctx);
78+
79+
if (fac != null)
80+
return fac;
81+
}
82+
7483
switch (call.getAggregation().getName()) {
7584
case "COUNT":
7685
return () -> new LongCount<>(call, hnd);
@@ -280,21 +289,21 @@ private static <Row> Supplier<Accumulator<Row>> maxFactory(AggregateCall call, R
280289
}
281290

282291
/** */
283-
private abstract static class AbstractAccumulator<Row> implements Accumulator<Row> {
292+
public abstract static class AbstractAccumulator<Row> implements Accumulator<Row> {
284293
/** */
285294
private final RowHandler<Row> hnd;
286295

287296
/** */
288297
private final transient AggregateCall aggCall;
289298

290299
/** */
291-
AbstractAccumulator(AggregateCall aggCall, RowHandler<Row> hnd) {
300+
protected AbstractAccumulator(AggregateCall aggCall, RowHandler<Row> hnd) {
292301
this.aggCall = aggCall;
293302
this.hnd = hnd;
294303
}
295304

296305
/** */
297-
<T> T get(int idx, Row row) {
306+
protected <T> T get(int idx, Row row) {
298307
assert idx < arguments().size() : "idx=" + idx + "; arguments=" + arguments();
299308

300309
return (T)hnd.get(arguments().get(idx), row);
@@ -311,7 +320,7 @@ protected List<Integer> arguments() {
311320
}
312321

313322
/** */
314-
int columnCount(Row row) {
323+
protected int columnCount(Row row) {
315324
return hnd.columnCount(row);
316325
}
317326
}
@@ -1344,8 +1353,9 @@ public ListAggAccumulator(AggregateCall aggCall, RowHandler<Row> hnd) {
13441353
if (builder == null)
13451354
builder = new StringBuilder();
13461355

1347-
if (builder.length() != 0)
1356+
if (!builder.isEmpty())
13481357
builder.append(extractSeparator(row));
1358+
13491359
builder.append(val);
13501360
}
13511361

‎modules/calcite/src/test/java/org/apache/ignite/internal/processors/query/calcite/integration/OperatorsExtensionIntegrationTest.java‎

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,18 @@
1818

1919
import java.math.BigDecimal;
2020
import java.sql.Timestamp;
21+
import java.util.List;
22+
import java.util.function.Supplier;
2123
import com.google.common.collect.ImmutableList;
2224
import org.apache.calcite.adapter.enumerable.NullPolicy;
2325
import org.apache.calcite.avatica.util.TimeUnitRange;
2426
import org.apache.calcite.linq4j.tree.Expressions;
27+
import org.apache.calcite.plan.Contexts;
28+
import org.apache.calcite.rel.core.AggregateCall;
29+
import org.apache.calcite.rel.type.RelDataType;
2530
import org.apache.calcite.rex.RexBuilder;
2631
import org.apache.calcite.rex.RexNode;
32+
import org.apache.calcite.sql.SqlAggFunction;
2733
import org.apache.calcite.sql.SqlCall;
2834
import org.apache.calcite.sql.SqlFunction;
2935
import org.apache.calcite.sql.SqlFunctionCategory;
@@ -43,12 +49,19 @@
4349
import org.apache.calcite.tools.FrameworkConfig;
4450
import org.apache.calcite.tools.Frameworks;
4551
import org.apache.calcite.util.BuiltInMethod;
52+
import org.apache.calcite.util.Optionality;
4653
import org.apache.ignite.configuration.IgniteConfiguration;
4754
import org.apache.ignite.internal.processors.query.calcite.CalciteQueryProcessor;
55+
import org.apache.ignite.internal.processors.query.calcite.exec.ExecutionContext;
56+
import org.apache.ignite.internal.processors.query.calcite.exec.RowHandler;
4857
import org.apache.ignite.internal.processors.query.calcite.exec.exp.RexImpTable;
58+
import org.apache.ignite.internal.processors.query.calcite.exec.exp.agg.Accumulator;
59+
import org.apache.ignite.internal.processors.query.calcite.exec.exp.agg.AccumulatorFactoryProvider;
60+
import org.apache.ignite.internal.processors.query.calcite.exec.exp.agg.Accumulators;
4961
import org.apache.ignite.internal.processors.query.calcite.prepare.IgniteConvertletTable;
5062
import org.apache.ignite.internal.processors.query.calcite.prepare.IgniteSqlNodeRewriter;
5163
import org.apache.ignite.internal.processors.query.calcite.prepare.IgniteSqlValidator;
64+
import org.apache.ignite.internal.processors.query.calcite.type.IgniteTypeFactory;
5265
import org.apache.ignite.plugin.AbstractTestPluginProvider;
5366
import org.apache.ignite.plugin.PluginContext;
5467
import org.jetbrains.annotations.Nullable;
@@ -75,6 +88,9 @@ public class OperatorsExtensionIntegrationTest extends AbstractBasicIntegrationT
7588
.sqlValidatorConfig(
7689
((IgniteSqlValidator.Config)CalciteQueryProcessor.FRAMEWORK_CONFIG.getSqlValidatorConfig())
7790
.withSqlNodeRewriter(new SqlRewriter()))
91+
.context(Contexts.chain(
92+
CalciteQueryProcessor.FRAMEWORK_CONFIG.getContext(),
93+
Contexts.of(new AccumulatorFactoryProviderImpl())))
7894
.build();
7995

8096
return (T)cfg;
@@ -134,6 +150,14 @@ public void testOperatorsCallsInViews() {
134150
assertQuery("SELECT val_str from my_view").returns(new BigDecimal("0")).check();
135151
}
136152

153+
/** */
154+
@Test
155+
public void testCustomAggregateFunction() {
156+
assertQuery("SELECT TEST_SUM(x) FROM (VALUES (1), (2), (3)) t(x)")
157+
.returns(6L)
158+
.check();
159+
}
160+
137161
/** Rewrites LTRIM with 2 parameters. */
138162
public static SqlCall rewriteLtrim(SqlValidator validator, SqlCall call) {
139163
if (call.operandCount() != 2)
@@ -193,6 +217,9 @@ public static class OperatorTable extends ReflectiveSqlOperatorTable {
193217
OperandTypes.STRING_STRING,
194218
SqlFunctionCategory.STRING
195219
);
220+
221+
/** */
222+
public static final SqlAggFunction TEST_SUM = new SqlTestSumAggFunction();
196223
}
197224

198225
/** Extended convertlet table. */
@@ -229,4 +256,73 @@ private static class SqlRewriter implements IgniteSqlNodeRewriter {
229256
return node;
230257
}
231258
}
259+
260+
/** */
261+
private static class AccumulatorFactoryProviderImpl implements AccumulatorFactoryProvider {
262+
/** {@inheritDoc} */
263+
@Override public @Nullable <Row> Supplier<Accumulator<Row>> factory(AggregateCall call, ExecutionContext<Row> ctx) {
264+
if (call.getAggregation().getName().equals(OperatorTable.TEST_SUM.getName()))
265+
return () -> new TestSum<>(call, ctx.rowHandler());
266+
267+
return null;
268+
}
269+
}
270+
271+
/** */
272+
public static class SqlTestSumAggFunction extends SqlAggFunction {
273+
/** */
274+
public SqlTestSumAggFunction() {
275+
super(
276+
"TEST_SUM",
277+
null,
278+
SqlKind.SUM,
279+
ReturnTypes.AGG_SUM,
280+
null,
281+
OperandTypes.NUMERIC,
282+
SqlFunctionCategory.NUMERIC,
283+
false,
284+
false,
285+
Optionality.FORBIDDEN
286+
);
287+
}
288+
}
289+
290+
/** */
291+
private static class TestSum<Row> extends Accumulators.AbstractAccumulator<Row> {
292+
/** */
293+
private long sum;
294+
295+
/** */
296+
protected TestSum(AggregateCall aggCall, RowHandler<Row> hnd) {
297+
super(aggCall, hnd);
298+
}
299+
300+
/** {@inheritDoc} */
301+
@Override public void add(Row row) {
302+
Number val = get(0, row);
303+
304+
if (val != null)
305+
sum += val.longValue();
306+
}
307+
308+
/** {@inheritDoc} */
309+
@Override public void apply(Accumulator<Row> other) {
310+
sum += ((TestSum<Row>)other).sum;
311+
}
312+
313+
/** {@inheritDoc} */
314+
@Override public Object end() {
315+
return sum;
316+
}
317+
318+
/** {@inheritDoc} */
319+
@Override public List<RelDataType> argumentTypes(IgniteTypeFactory typeFactory) {
320+
return List.of(typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true));
321+
}
322+
323+
/** {@inheritDoc} */
324+
@Override public RelDataType returnType(IgniteTypeFactory typeFactory) {
325+
return typeFactory.createSqlType(org.apache.calcite.sql.type.SqlTypeName.BIGINT);
326+
}
327+
}
232328
}

‎modules/core/src/main/java/org/apache/ignite/internal/processors/query/QueryContext.java‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.util.Arrays;
2222
import java.util.List;
2323
import org.apache.ignite.internal.util.typedef.F;
24+
import org.jetbrains.annotations.Nullable;
2425

2526
/** */
2627
public final class QueryContext {
@@ -36,10 +37,10 @@ private QueryContext(Object[] params) {
3637
}
3738

3839
/**
39-
* Finds an instance of an interface implemented by this object,
40-
* or returns null if this object does not support that interface.
40+
* Finds an instance of an interface implemented by this object
41+
* or returns {@code null} if this object does not support that interface.
4142
*/
42-
public <C> C unwrap(Class<C> aClass) {
43+
public <C> @Nullable C unwrap(Class<C> aClass) {
4344
if (Object[].class == aClass)
4445
return aClass.cast(params);
4546

@@ -50,12 +51,12 @@ public <C> C unwrap(Class<C> aClass) {
5051
* @param params Context parameters.
5152
* @return Query context.
5253
*/
53-
public static QueryContext of(Object... params) {
54+
public static QueryContext of(@Nullable Object... params) {
5455
return !F.isEmpty(params) ? new QueryContext(build(null, params).toArray()) : new QueryContext(EMPTY);
5556
}
5657

5758
/** */
58-
private static List<Object> build(List<Object> dst, Object[] src) {
59+
private static List<Object> build(List<Object> dst, @Nullable Object[] src) {
5960
if (dst == null)
6061
dst = new ArrayList<>();
6162

0 commit comments

Comments
 (0)