This is an automated email from the ASF dual-hosted git repository. shuwenwei pushed a commit to branch flink-iotdb-table-connector in repository https://gitbox.apache.org/repos/asf/iotdb-extras.git
commit 7d8469e98e35745db378e2208754bf8e722a94c5 Author: shuwenwei <[email protected]> AuthorDate: Mon Sep 21 15:40:59 2026 +0800 feat: support aggregate pushdown --- .../flink/source/pushdown/AggregateSpec.java | 52 ++++++ .../pushdown/IoTDBAggregatePushDownUtils.java | 192 +++++++++++++++++++ .../iotdb/relational/flink/utils/IoTDBUtils.java | 33 ++++ .../iotdb/relational/flink/source/IoTDBSource.java | 11 +- .../flink/source/IoTDBSourceEnumerator.java | 25 ++- .../table/IoTDBRelationalDynamicTableSource.java | 51 ++++- .../IoTDBRelationalSourcePushDownPlannerTest.java | 205 ++++++++++++++++++++- 7 files changed, 547 insertions(+), 22 deletions(-) diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/AggregateSpec.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/AggregateSpec.java new file mode 100644 index 0000000..d2b4644 --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/AggregateSpec.java @@ -0,0 +1,52 @@ +/* + * 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.iotdb.relational.flink.source.pushdown; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +/** + * Serializable description of a pushed-down aggregate query. + * + * <p>{@code selectExpressions} are the rendered grouping keys followed by the aggregate + * expressions; {@code groupByExpressions} are the rendered grouping keys. + */ +public class AggregateSpec implements Serializable { + + private static final long serialVersionUID = 1L; + + private final List<String> selectExpressions; + private final List<String> groupByExpressions; + + public AggregateSpec(List<String> selectExpressions, List<String> groupByExpressions) { + this.selectExpressions = Collections.unmodifiableList(new ArrayList<>(selectExpressions)); + this.groupByExpressions = Collections.unmodifiableList(new ArrayList<>(groupByExpressions)); + } + + public List<String> getSelectExpressions() { + return selectExpressions; + } + + public List<String> getGroupByExpressions() { + return groupByExpressions; + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/IoTDBAggregatePushDownUtils.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/IoTDBAggregatePushDownUtils.java new file mode 100644 index 0000000..0abb7b3 --- /dev/null +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/source/pushdown/IoTDBAggregatePushDownUtils.java @@ -0,0 +1,192 @@ +/* + * 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.iotdb.relational.flink.source.pushdown; + +import org.apache.iotdb.relational.flink.utils.IoTDBUtils; + +import org.apache.flink.table.expressions.AggregateExpression; +import org.apache.flink.table.expressions.FieldReferenceExpression; +import org.apache.flink.table.types.DataType; +import org.apache.tsfile.enums.TSDataType; + +import java.util.ArrayList; +import java.util.List; + +/** + * Translates Flink aggregate pushdown information into a serializable {@link AggregateSpec}. + * + * <p>The pushdown is all-or-nothing: any unsupported aggregate, argument or grouping makes this + * return {@code null} so the whole aggregation stays in Flink. + * + * <p>Each aggregate is classified and translated in a single pass from its function class name; + * grouping columns and aggregate arguments are rendered through {@link IoTDBExpressionVisitor}. + */ +public final class IoTDBAggregatePushDownUtils { + + private IoTDBAggregatePushDownUtils() {} + + public static AggregateSpec translate( + List<int[]> groupingSets, + List<AggregateExpression> aggregateExpressions, + DataType sourceRowDataType, + DataType producedDataType) { + if (groupingSets == null + || groupingSets.size() != 1 + || aggregateExpressions == null + || aggregateExpressions.isEmpty()) { + return null; + } + + final List<String> sourceFieldNames; + final List<DataType> producedFieldTypes; + try { + sourceFieldNames = DataType.getFieldNames(sourceRowDataType); + producedFieldTypes = DataType.getFieldDataTypes(producedDataType); + } catch (RuntimeException e) { + return null; + } + + int[] grouping = groupingSets.get(0); + if (grouping == null) { + return null; + } + + List<String> groupByExpressions = new ArrayList<>(grouping.length); + List<String> selectExpressions = new ArrayList<>(grouping.length + aggregateExpressions.size()); + for (int index : grouping) { + if (index < 0 || index >= sourceFieldNames.size()) { + return null; + } + String column = IoTDBUtils.quoteIdentifier(sourceFieldNames.get(index)); + groupByExpressions.add(column); + selectExpressions.add(column); + } + + IoTDBExpressionVisitor visitor = new IoTDBExpressionVisitor(); + for (int i = 0; i < aggregateExpressions.size(); i++) { + int producedIndex = grouping.length + i; + if (producedIndex >= producedFieldTypes.size()) { + return null; + } + String sql = + translateAggregate( + aggregateExpressions.get(i), producedFieldTypes.get(producedIndex), visitor); + if (sql == null) { + return null; + } + selectExpressions.add(sql); + } + + return new AggregateSpec(selectExpressions, groupByExpressions); + } + + private static String translateAggregate( + AggregateExpression aggregate, DataType producedType, IoTDBExpressionVisitor visitor) { + if (!isSupportedAggregate(aggregate)) { + return null; + } + + Class<?> functionClass = aggregate.getFunctionDefinition().getClass(); + String simpleName = functionClass == null ? null : functionClass.getSimpleName(); + TSDataType outType = toTsDataType(producedType); + if (simpleName == null || outType == null) { + return null; + } + + List<FieldReferenceExpression> args = aggregate.getArgs(); + int argCount = args == null ? 0 : args.size(); + + String expression; + if (simpleName.endsWith("Count1AggFunction")) { + if (argCount != 0 || !isCountType(outType)) { + return null; + } + expression = "COUNT(*)"; + } else if (simpleName.endsWith("CountAggFunction")) { + if (argCount != 1 || !isCountType(outType)) { + return null; + } + String argSql = args.get(0).accept(visitor); + if (argSql == null || toTsDataType(args.get(0).getOutputDataType()) == null) { + return null; + } + expression = "COUNT(" + argSql + ")"; + } else if (simpleName.endsWith("Sum0AggFunction") || simpleName.endsWith("SumAggFunction")) { + if (argCount != 1 || !isNumeric(outType)) { + return null; + } + String argSql = args.get(0).accept(visitor); + TSDataType argType = argSql == null ? null : toTsDataType(args.get(0).getOutputDataType()); + if (argType == null || !isNumeric(argType)) { + return null; + } + expression = "SUM(" + argSql + ")"; + } else if (simpleName.endsWith("MaxAggFunction") || simpleName.endsWith("MinAggFunction")) { + if (argCount != 1) { + return null; + } + String argSql = args.get(0).accept(visitor); + if (argSql == null || toTsDataType(args.get(0).getOutputDataType()) == null) { + return null; + } + String functionName = simpleName.endsWith("MaxAggFunction") ? "MAX" : "MIN"; + expression = functionName + "(" + argSql + ")"; + } else { + return null; + } + + return "CAST(" + expression + " AS " + outType.name() + ")"; + } + + private static boolean isSupportedAggregate(AggregateExpression aggregate) { + return aggregate != null + && !aggregate.isDistinct() + && !aggregate.isApproximate() + && !aggregate.isIgnoreNulls() + && !aggregate.getFilterExpression().isPresent(); + } + + private static boolean isNumeric(TSDataType dataType) { + switch (dataType) { + case INT32: + case INT64: + case FLOAT: + case DOUBLE: + return true; + default: + return false; + } + } + + private static boolean isCountType(TSDataType dataType) { + return dataType == TSDataType.INT32 || dataType == TSDataType.INT64; + } + + private static TSDataType toTsDataType(DataType dataType) { + if (dataType == null) { + return null; + } + try { + return IoTDBUtils.toIoTDBDataType(dataType); + } catch (RuntimeException e) { + return null; + } + } +} diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/utils/IoTDBUtils.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/utils/IoTDBUtils.java index 908dd59..f3c9237 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/utils/IoTDBUtils.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-base/src/main/java/org/apache/iotdb/relational/flink/utils/IoTDBUtils.java @@ -245,6 +245,39 @@ public final class IoTDBUtils { return sql.toString(); } + /** + * Builds a table-model aggregation query. + * + * @param table IoTDB table name + * @param selectExpressions already-rendered SELECT expressions (grouping columns and aggregates) + * @param filterQueries already-rendered IoTDB predicate fragments + * @param groupByExpressions already-rendered GROUP BY expressions, or {@code null} for a global + * aggregation + * @return IoTDB SELECT SQL + */ + public static String buildAggregateQuery( + String table, + List<String> selectExpressions, + List<String> filterQueries, + List<String> groupByExpressions) { + if (selectExpressions == null || selectExpressions.isEmpty()) { + throw new IllegalArgumentException("IoTDB aggregate query requires at least one column."); + } + + StringBuilder sql = + new StringBuilder("SELECT ") + .append(String.join(", ", selectExpressions)) + .append(" FROM ") + .append(quoteIdentifier(table)); + if (filterQueries != null && !filterQueries.isEmpty()) { + sql.append(" WHERE ").append(String.join(" AND ", filterQueries)); + } + if (groupByExpressions != null && !groupByExpressions.isEmpty()) { + sql.append(" GROUP BY ").append(String.join(", ", groupByExpressions)); + } + return sql.toString(); + } + private static void validateColumnExists( String columnName, String optionName, Map<String, TSDataType> dataTypesByColumn) { if (!dataTypesByColumn.containsKey(columnName)) { diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSource.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSource.java index ea4abc3..94754e4 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSource.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSource.java @@ -23,6 +23,7 @@ import org.apache.iotdb.relational.flink.cfg.IoTDBOptions; import org.apache.iotdb.relational.flink.source.deserializer.IoTDBDeserializationSchema; import org.apache.iotdb.relational.flink.source.enumerator.IoTDBSourceEnumeratorState; import org.apache.iotdb.relational.flink.source.enumerator.IoTDBSourceEnumeratorStateSerializer; +import org.apache.iotdb.relational.flink.source.pushdown.AggregateSpec; import org.apache.iotdb.relational.flink.source.split.IoTDBSourceSplit; import org.apache.iotdb.relational.flink.source.split.IoTDBSourceSplitSerializer; @@ -54,18 +55,21 @@ public class IoTDBSource<OUT> implements Source<OUT, IoTDBSourceSplit, IoTDBSour private final IoTDBDeserializationSchema<OUT> deserializer; private final List<String> filterQueries; private final long limit; + private final AggregateSpec aggregateSpec; public IoTDBSource( IoTDBOptions options, DataType rowDataType, IoTDBDeserializationSchema<OUT> deserializer, List<String> filterQueries, - long limit) { + long limit, + AggregateSpec aggregateSpec) { this.options = options; this.rowDataType = rowDataType; this.deserializer = deserializer; this.filterQueries = filterQueries == null ? new ArrayList<>() : new ArrayList<>(filterQueries); this.limit = limit; + this.aggregateSpec = aggregateSpec; } @Override @@ -81,14 +85,15 @@ public class IoTDBSource<OUT> implements Source<OUT, IoTDBSourceSplit, IoTDBSour @Override public SplitEnumerator<IoTDBSourceSplit, IoTDBSourceEnumeratorState> createEnumerator( SplitEnumeratorContext<IoTDBSourceSplit> enumContext) { - return new IoTDBSourceEnumerator(enumContext, options, rowDataType, filterQueries, limit); + return new IoTDBSourceEnumerator( + enumContext, options, rowDataType, filterQueries, limit, aggregateSpec); } @Override public SplitEnumerator<IoTDBSourceSplit, IoTDBSourceEnumeratorState> restoreEnumerator( SplitEnumeratorContext<IoTDBSourceSplit> enumContext, IoTDBSourceEnumeratorState checkpoint) { return new IoTDBSourceEnumerator( - enumContext, options, rowDataType, filterQueries, limit, checkpoint); + enumContext, options, rowDataType, filterQueries, limit, aggregateSpec, checkpoint); } @Override diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSourceEnumerator.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSourceEnumerator.java index 29ea548..99bf3fb 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSourceEnumerator.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/source/IoTDBSourceEnumerator.java @@ -21,6 +21,7 @@ package org.apache.iotdb.relational.flink.source; import org.apache.iotdb.relational.flink.cfg.IoTDBOptions; import org.apache.iotdb.relational.flink.source.enumerator.IoTDBSourceEnumeratorState; +import org.apache.iotdb.relational.flink.source.pushdown.AggregateSpec; import org.apache.iotdb.relational.flink.source.split.IoTDBSourceSplit; import org.apache.iotdb.relational.flink.utils.IoTDBUtils; @@ -48,6 +49,7 @@ public class IoTDBSourceEnumerator private final DataType rowDataType; private final List<String> filterQueries; private final long limit; + private final AggregateSpec aggregateSpec; private final Deque<IoTDBSourceSplit> pendingSplits = new ArrayDeque<>(); private final Deque<Integer> readersAwaitingSplit = new ArrayDeque<>(); private final Set<Integer> assignedReaders = new HashSet<>(); @@ -60,8 +62,9 @@ public class IoTDBSourceEnumerator IoTDBOptions options, DataType rowDataType, List<String> filterQueries, - long limit) { - this(context, options, rowDataType, filterQueries, limit, null); + long limit, + AggregateSpec aggregateSpec) { + this(context, options, rowDataType, filterQueries, limit, aggregateSpec, null); } public IoTDBSourceEnumerator( @@ -70,12 +73,14 @@ public class IoTDBSourceEnumerator DataType rowDataType, List<String> filterQueries, long limit, + AggregateSpec aggregateSpec, @Nullable IoTDBSourceEnumeratorState checkpoint) { this.context = context; this.options = options; this.rowDataType = rowDataType; this.filterQueries = filterQueries == null ? new ArrayList<>() : new ArrayList<>(filterQueries); this.limit = limit; + this.aggregateSpec = aggregateSpec; if (checkpoint != null) { pendingSplits.addAll(checkpoint.getRemainingSplits()); allSplitsCreated = true; @@ -145,9 +150,17 @@ public class IoTDBSourceEnumerator private IoTDBSourceSplit createSingleSplit() { String splitId = UUID.randomUUID().toString(); - String sql = - IoTDBUtils.buildSelectQuery( - options.getTable(), rowDataType, filterQueries, limit); - return new IoTDBSourceSplit(splitId, options.getDatabase(), options.getTable(), sql); + return new IoTDBSourceSplit(splitId, options.getDatabase(), options.getTable(), buildSql()); + } + + private String buildSql() { + if (aggregateSpec != null) { + return IoTDBUtils.buildAggregateQuery( + options.getTable(), + aggregateSpec.getSelectExpressions(), + filterQueries, + aggregateSpec.getGroupByExpressions()); + } + return IoTDBUtils.buildSelectQuery(options.getTable(), rowDataType, filterQueries, limit); } } diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSource.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSource.java index 6699168..f9d43c9 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSource.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/main/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalDynamicTableSource.java @@ -22,6 +22,8 @@ package org.apache.iotdb.relational.flink.table; import org.apache.iotdb.relational.flink.cfg.IoTDBOptions; import org.apache.iotdb.relational.flink.source.IoTDBSource; import org.apache.iotdb.relational.flink.source.deserializer.RowDataDeserializationSchema; +import org.apache.iotdb.relational.flink.source.pushdown.AggregateSpec; +import org.apache.iotdb.relational.flink.source.pushdown.IoTDBAggregatePushDownUtils; import org.apache.iotdb.relational.flink.source.pushdown.IoTDBExpressionVisitor; import org.apache.iotdb.relational.flink.utils.IoTDBUtils; @@ -30,9 +32,11 @@ import org.apache.flink.table.connector.ChangelogMode; import org.apache.flink.table.connector.source.DynamicTableSource; import org.apache.flink.table.connector.source.ScanTableSource; import org.apache.flink.table.connector.source.SourceProvider; +import org.apache.flink.table.connector.source.abilities.SupportsAggregatePushDown; import org.apache.flink.table.connector.source.abilities.SupportsFilterPushDown; import org.apache.flink.table.connector.source.abilities.SupportsLimitPushDown; import org.apache.flink.table.connector.source.abilities.SupportsProjectionPushDown; +import org.apache.flink.table.expressions.AggregateExpression; import org.apache.flink.table.expressions.ResolvedExpression; import org.apache.flink.table.types.DataType; @@ -50,13 +54,15 @@ public class IoTDBRelationalDynamicTableSource implements ScanTableSource, SupportsFilterPushDown, SupportsLimitPushDown, - SupportsProjectionPushDown { + SupportsProjectionPushDown, + SupportsAggregatePushDown { private final IoTDBOptions options; private final ResolvedSchema schema; private DataType physicalRowDataType; private final List<String> resolvedFilterQueries = new ArrayList<>(); private long limit = -1L; + private AggregateSpec aggregateSpec; public IoTDBRelationalDynamicTableSource(IoTDBOptions options, ResolvedSchema schema) { this.options = options; @@ -77,7 +83,8 @@ public class IoTDBRelationalDynamicTableSource physicalRowDataType, new RowDataDeserializationSchema(physicalRowDataType), resolvedFilterQueries, - limit)); + limit, + aggregateSpec)); } @Override @@ -87,7 +94,9 @@ public class IoTDBRelationalDynamicTableSource @Override public void applyProjection(int[][] projectedFields, DataType producedDataType) { - this.physicalRowDataType = producedDataType; + if (aggregateSpec == null) { + this.physicalRowDataType = producedDataType; + } } @Override @@ -111,9 +120,29 @@ public class IoTDBRelationalDynamicTableSource return Result.of(acceptedFilters, remainingFilters); } + @Override + public boolean applyAggregates( + List<int[]> groupingSets, + List<AggregateExpression> aggregateExpressions, + DataType producedDataType) { + // Grouping and argument indices refer to the scan's current row type, which is the row type + // after any projection that has already been pushed into this source. + AggregateSpec spec = + IoTDBAggregatePushDownUtils.translate( + groupingSets, aggregateExpressions, physicalRowDataType, producedDataType); + if (spec == null) { + return false; + } + this.aggregateSpec = spec; + this.physicalRowDataType = producedDataType; + return true; + } + @Override public void applyLimit(long limit) { - this.limit = limit; + if (aggregateSpec == null) { + this.limit = limit; + } } @Override @@ -122,6 +151,7 @@ public class IoTDBRelationalDynamicTableSource copy.physicalRowDataType = physicalRowDataType; copy.resolvedFilterQueries.addAll(resolvedFilterQueries); copy.limit = limit; + copy.aggregateSpec = aggregateSpec; return copy; } @@ -132,13 +162,24 @@ public class IoTDBRelationalDynamicTableSource /** * Builds the IoTDB query that this source would execute. Exposed for tests so the pushed-down - * projection, filters and limit can be verified without executing any query. + * projection, filters, limit and aggregation can be verified without executing any query. */ String buildQuery() { + if (aggregateSpec != null) { + return IoTDBUtils.buildAggregateQuery( + options.getTable(), + aggregateSpec.getSelectExpressions(), + resolvedFilterQueries, + aggregateSpec.getGroupByExpressions()); + } return IoTDBUtils.buildSelectQuery( options.getTable(), physicalRowDataType, resolvedFilterQueries, limit); } + AggregateSpec getAggregateSpec() { + return aggregateSpec; + } + List<String> getResolvedFilterQueries() { return resolvedFilterQueries; } diff --git a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/test/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalSourcePushDownPlannerTest.java b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/test/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalSourcePushDownPlannerTest.java index 2e53e0f..497a85f 100644 --- a/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/test/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalSourcePushDownPlannerTest.java +++ b/connectors/flink-iotdb-table-connector/flink-iotdb-table-connector-flink1/src/test/java/org/apache/iotdb/relational/flink/table/IoTDBRelationalSourcePushDownPlannerTest.java @@ -122,8 +122,7 @@ public class IoTDBRelationalSourcePushDownPlannerTest { optimize("SELECT device_id FROM iotdb_t WHERE temperature + humidity > 30.0E0"); assertEquals( - "SELECT \"device_id\" FROM \"sensor\" " - + "WHERE ((\"temperature\" + \"humidity\") > 30.0)", + "SELECT \"device_id\" FROM \"sensor\" " + "WHERE ((\"temperature\" + \"humidity\") > 30.0)", source.buildQuery()); } @@ -174,8 +173,7 @@ public class IoTDBRelationalSourcePushDownPlannerTest { optimize("SELECT device_id FROM iotdb_t WHERE temperature <> 0.0E0"); assertEquals( - "SELECT \"device_id\" FROM \"sensor\" WHERE (\"temperature\" <> 0.0)", - source.buildQuery()); + "SELECT \"device_id\" FROM \"sensor\" WHERE (\"temperature\" <> 0.0)", source.buildQuery()); } @Test @@ -238,8 +236,7 @@ public class IoTDBRelationalSourcePushDownPlannerTest { assertEquals( Collections.singletonList("(\"device_id\" = 'd1')"), source.getResolvedFilterQueries()); assertEquals( - "SELECT \"temperature\" FROM \"sensor\" WHERE (\"device_id\" = 'd1')", - source.buildQuery()); + "SELECT \"temperature\" FROM \"sensor\" WHERE (\"device_id\" = 'd1')", source.buildQuery()); } @Test @@ -293,8 +290,7 @@ public class IoTDBRelationalSourcePushDownPlannerTest { optimize("SELECT device_id FROM iotdb_t WHERE `time` < NOW()"); assertEquals( - "SELECT \"device_id\" FROM \"sensor\" WHERE (\"time\" < now())", - source.buildQuery()); + "SELECT \"device_id\" FROM \"sensor\" WHERE (\"time\" < now())", source.buildQuery()); } @Test @@ -306,9 +302,202 @@ public class IoTDBRelationalSourcePushDownPlannerTest { assertEquals("SELECT \"device_id\" FROM \"sensor\"", source.buildQuery()); } + @Test + public void testGlobalCountStarIsNotPushedDown() { + // Flink 1.17 feeds a global COUNT(*) through a constant Calc, and the pushdown rule only + // accepts field-projection Calcs, so it is intentionally left to Flink. + IoTDBRelationalDynamicTableSource source = optimize("SELECT COUNT(*) FROM iotdb_t"); + + assertEquals("SELECT \"time\" FROM \"sensor\"", source.buildQuery()); + } + + @Test + public void testCountColumnPushDown() { + IoTDBRelationalDynamicTableSource source = optimize("SELECT COUNT(temperature) FROM iotdb_t"); + + assertNotNull(source.getAggregateSpec()); + assertEquals( + "SELECT CAST(COUNT(\"temperature\") AS INT64) FROM \"sensor\"", source.buildQuery()); + } + + @Test + public void testSumPushDown() { + IoTDBRelationalDynamicTableSource source = optimize("SELECT SUM(temperature) FROM iotdb_t"); + + assertNotNull(source.getAggregateSpec()); + assertEquals( + "SELECT CAST(SUM(\"temperature\") AS DOUBLE) FROM \"sensor\"", source.buildQuery()); + } + + @Test + public void testAvgDecomposedPushDown() { + // AVG is decomposed by the planner into SUM0 + COUNT and both must be pushed down. + IoTDBRelationalDynamicTableSource source = optimize("SELECT AVG(temperature) FROM iotdb_t"); + + assertNotNull(source.getAggregateSpec()); + assertEquals( + "SELECT CAST(SUM(\"temperature\") AS DOUBLE), CAST(COUNT(\"temperature\") AS INT64) " + + "FROM \"sensor\"", + source.buildQuery()); + } + + @Test + public void testMaxAndMinPushDown() { + IoTDBRelationalDynamicTableSource source = + optimize("SELECT MAX(temperature), MIN(temperature) FROM iotdb_t"); + + assertNotNull(source.getAggregateSpec()); + assertEquals( + "SELECT CAST(MAX(\"temperature\") AS DOUBLE), CAST(MIN(\"temperature\") AS DOUBLE) " + + "FROM \"sensor\"", + source.buildQuery()); + } + + @Test + public void testMinStringPushDown() { + IoTDBRelationalDynamicTableSource source = optimize("SELECT MIN(device_id) FROM iotdb_t"); + + assertNotNull(source.getAggregateSpec()); + assertEquals("SELECT CAST(MIN(\"device_id\") AS STRING) FROM \"sensor\"", source.buildQuery()); + } + + @Test + public void testGroupByCountPushDown() { + IoTDBRelationalDynamicTableSource source = + optimize("SELECT device_id, COUNT(*) FROM iotdb_t GROUP BY device_id"); + + assertNotNull(source.getAggregateSpec()); + assertEquals( + "SELECT \"device_id\", CAST(COUNT(*) AS INT64) FROM \"sensor\" GROUP BY \"device_id\"", + source.buildQuery()); + } + + @Test + public void testGroupBySumPushDown() { + IoTDBRelationalDynamicTableSource source = + optimize("SELECT device_id, SUM(temperature) FROM iotdb_t GROUP BY device_id"); + + assertNotNull(source.getAggregateSpec()); + assertEquals( + "SELECT \"device_id\", CAST(SUM(\"temperature\") AS DOUBLE) FROM \"sensor\" " + + "GROUP BY \"device_id\"", + source.buildQuery()); + } + + @Test + public void testAggregateWithFilterPushDown() { + IoTDBRelationalDynamicTableSource source = + optimize("SELECT SUM(temperature) FROM iotdb_t WHERE temperature > 30.0E0"); + + assertNotNull(source.getAggregateSpec()); + assertEquals( + Collections.singletonList("(\"temperature\" > 30.0)"), source.getResolvedFilterQueries()); + assertEquals( + "SELECT CAST(SUM(\"temperature\") AS DOUBLE) FROM \"sensor\" " + + "WHERE (\"temperature\" > 30.0)", + source.buildQuery()); + } + + @Test + public void testMultipleAggregatesPushDown() { + IoTDBRelationalDynamicTableSource source = + optimize( + "SELECT device_id, COUNT(*), SUM(temperature), MAX(humidity), MIN(temperature) " + + "FROM iotdb_t GROUP BY device_id"); + + assertNotNull(source.getAggregateSpec()); + assertEquals( + "SELECT \"device_id\", CAST(COUNT(*) AS INT64), CAST(SUM(\"temperature\") AS DOUBLE), " + + "CAST(MAX(\"humidity\") AS DOUBLE), CAST(MIN(\"temperature\") AS DOUBLE) " + + "FROM \"sensor\" GROUP BY \"device_id\"", + source.buildQuery()); + } + + @Test + public void testGroupByMultipleColumnsPushDown() { + IoTDBRelationalDynamicTableSource source = + optimize( + "SELECT device_id, temperature, COUNT(*) FROM iotdb_t " + + "GROUP BY device_id, temperature"); + + assertNotNull(source.getAggregateSpec()); + assertEquals( + "SELECT \"device_id\", \"temperature\", CAST(COUNT(*) AS INT64) FROM \"sensor\" " + + "GROUP BY \"device_id\", \"temperature\"", + source.buildQuery()); + } + + @Test + public void testAggregateWithFilterAndGroupByPushDown() { + IoTDBRelationalDynamicTableSource source = + optimize( + "SELECT device_id, AVG(temperature) FROM iotdb_t " + + "WHERE humidity > 10.0E0 GROUP BY device_id"); + + assertNotNull(source.getAggregateSpec()); + assertEquals( + "SELECT \"device_id\", CAST(SUM(\"temperature\") AS DOUBLE), " + + "CAST(COUNT(\"temperature\") AS INT64) FROM \"sensor\" " + + "WHERE (\"humidity\" > 10.0) GROUP BY \"device_id\"", + source.buildQuery()); + } + + @Test + public void testSumOfExpressionIsNotPushedDown() { + // sum(a + b): the argument is an expression, which Flink evaluates in a Calc before the + // local aggregate, so the pushdown rule leaves it to Flink. + IoTDBRelationalDynamicTableSource source = + optimize("SELECT SUM(temperature + humidity) FROM iotdb_t"); + + assertEquals("SELECT \"temperature\", \"humidity\" FROM \"sensor\"", source.buildQuery()); + } + + @Test + public void testSumOfScalingExpressionIsNotPushedDown() { + IoTDBRelationalDynamicTableSource source = + optimize("SELECT SUM(temperature * 2.0E0) FROM iotdb_t"); + + assertEquals("SELECT \"temperature\" FROM \"sensor\"", source.buildQuery()); + } + + @Test + public void testAggregateOfFunctionIsNotPushedDown() { + IoTDBRelationalDynamicTableSource source = + optimize("SELECT MAX(LOWER(device_id)) FROM iotdb_t"); + + assertEquals("SELECT \"device_id\" FROM \"sensor\"", source.buildQuery()); + } + + @Test + public void testGroupByExpressionIsNotPushedDown() { + IoTDBRelationalDynamicTableSource source = + optimize("SELECT device_id || 'x', COUNT(*) FROM iotdb_t GROUP BY device_id || 'x'"); + + assertEquals("SELECT \"device_id\" FROM \"sensor\"", source.buildQuery()); + } + + @Test + public void testCountDistinctIsNotPushedDown() { + IoTDBRelationalDynamicTableSource source = + optimize("SELECT COUNT(DISTINCT device_id) FROM iotdb_t"); + + assertEquals("SELECT \"device_id\" FROM \"sensor\"", source.buildQuery()); + } + + @Test + public void testUnsupportedAggregateIsNotPushedDown() { + IoTDBRelationalDynamicTableSource source = + optimize("SELECT STDDEV_POP(temperature) FROM iotdb_t"); + + assertEquals("SELECT \"temperature\" FROM \"sensor\"", source.buildQuery()); + } + private static IoTDBRelationalDynamicTableSource optimize(String query) { TableEnvironmentImpl tableEnvironment = (TableEnvironmentImpl) TableEnvironment.create(EnvironmentSettings.inBatchMode()); + // Aggregate pushdown is opt-in and needs a local (partial) aggregate to be generated. + tableEnvironment.getConfig().set("table.optimizer.source.aggregate-pushdown-enabled", "true"); + tableEnvironment.getConfig().set("table.optimizer.agg-phase-strategy", "TWO_PHASE"); tableEnvironment.executeSql(DDL); Table table = tableEnvironment.sqlQuery(query);
