You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@ignite.apache.org by am...@apache.org on 2024/01/30 09:33:41 UTC
(ignite-3) branch main updated: IGNITE-21277 Sql: Partition pruning (#3105)
This is an automated email from the ASF dual-hosted git repository.
amashenkov pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/ignite-3.git
The following commit(s) were added to refs/heads/main by this push:
new b932c7f5c0 IGNITE-21277 Sql: Partition pruning (#3105)
b932c7f5c0 is described below
commit b932c7f5c0f56e645440372ba190fd52164570c0
Author: Max Zhuravkov <sh...@gmail.com>
AuthorDate: Tue Jan 30 11:33:35 2024 +0200
IGNITE-21277 Sql: Partition pruning (#3105)
---
.../prepare/pruning/PartitionPruningColumns.java | 71 +++
.../prepare/pruning/PartitionPruningMetadata.java | 58 +++
.../pruning/PartitionPruningMetadataExtractor.java | 568 +++++++++++++++++++++
.../planner/PartitionPruningExtractorSelfTest.java | 389 ++++++++++++++
.../planner/PartitionPruningMetadataTest.java | 473 +++++++++++++++++
.../sql/engine/planner/PartitionPruningTest.java | 137 +++++
6 files changed, 1696 insertions(+)
diff --git a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningColumns.java b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningColumns.java
new file mode 100644
index 0000000000..7ba1a44e9a
--- /dev/null
+++ b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningColumns.java
@@ -0,0 +1,71 @@
+/*
+ * 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.ignite.internal.sql.engine.prepare.pruning;
+
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.stream.Collectors;
+import org.apache.calcite.rex.RexNode;
+import org.apache.ignite.internal.tostring.IgniteToStringInclude;
+import org.apache.ignite.internal.tostring.S;
+import org.jetbrains.annotations.TestOnly;
+
+/**
+ * Metadata for partition pruning. Collection of column values for a particular relation/physical operator. Example:
+ * <pre>
+ * Condition: c1 = 10 AND c2 = 42 OR c1 = 77 AND c2 = 173
+ * Colocation keys: c1, c2
+ * Columns: [c1=10, c2=42], [c1=77, c2=173]
+ * </pre>
+ *
+ * @see PartitionPruningMetadataExtractor
+ */
+public class PartitionPruningColumns {
+
+ @IgniteToStringInclude
+ private final List<Int2ObjectMap<RexNode>> columns;
+
+ /** Constructor. */
+ public PartitionPruningColumns(List<Int2ObjectMap<RexNode>> columns) {
+ this.columns = Collections.unmodifiableList(columns);
+ }
+
+ /** A list of column values. */
+ public List<Int2ObjectMap<RexNode>> columns() {
+ return columns;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public String toString() {
+ return S.toString(this);
+ }
+
+ /** Returns column values in canonical form. E.g. {@code [1=2, 0=3]} becomes {@code [0=3, 1=2]} */
+ @TestOnly
+ public static List<List<Map.Entry<Integer, RexNode>>> canonicalForm(PartitionPruningColumns columns) {
+ return columns.columns.stream()
+ .map(cols -> cols.int2ObjectEntrySet().stream().map(e -> Map.entry(e.getIntKey(), e.getValue()))
+ .sorted(Entry.comparingByKey())
+ .collect(Collectors.toList()))
+ .collect(Collectors.toList());
+ }
+}
diff --git a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningMetadata.java b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningMetadata.java
new file mode 100644
index 0000000000..0fdd51c9f2
--- /dev/null
+++ b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningMetadata.java
@@ -0,0 +1,58 @@
+/*
+ * 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.ignite.internal.sql.engine.prepare.pruning;
+
+import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
+import it.unimi.dsi.fastutil.longs.Long2ObjectMaps;
+import org.apache.ignite.internal.tostring.IgniteToStringInclude;
+import org.apache.ignite.internal.tostring.S;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Partition pruning metadata includes partition pruning information for each operator of a physical plan that supports partition pruning.
+ *
+ * @see PartitionPruningColumns
+ * @see PartitionPruningMetadataExtractor
+ */
+public class PartitionPruningMetadata {
+
+ @IgniteToStringInclude
+ private final Long2ObjectMap<PartitionPruningColumns> data;
+
+ /** Constructor. */
+ public PartitionPruningMetadata(Long2ObjectMap<PartitionPruningColumns> data) {
+ this.data = Long2ObjectMaps.unmodifiable(data);
+ }
+
+ /** Return metadata for operator with the given sourceId, or {@code null} if there is no metadata for it. */
+ @Nullable
+ public PartitionPruningColumns get(long sourceId) {
+ return data.get(sourceId);
+ }
+
+ /** A map of operators with their metadata. */
+ public Long2ObjectMap<PartitionPruningColumns> data() {
+ return data;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public String toString() {
+ return S.toString(this);
+ }
+}
diff --git a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningMetadataExtractor.java b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningMetadataExtractor.java
new file mode 100644
index 0000000000..8b637ef957
--- /dev/null
+++ b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/pruning/PartitionPruningMetadataExtractor.java
@@ -0,0 +1,568 @@
+/*
+ * 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.ignite.internal.sql.engine.prepare.pruning;
+
+import it.unimi.dsi.fastutil.ints.Int2ObjectArrayMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap;
+import it.unimi.dsi.fastutil.ints.Int2ObjectMap.Entry;
+import it.unimi.dsi.fastutil.ints.IntArrayList;
+import it.unimi.dsi.fastutil.ints.IntList;
+import it.unimi.dsi.fastutil.longs.Long2ObjectArrayMap;
+import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexDynamicParam;
+import org.apache.calcite.rex.RexLiteral;
+import org.apache.calcite.rex.RexLocalRef;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexUtil;
+import org.apache.calcite.sql.SqlKind;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.ignite.internal.sql.engine.prepare.IgniteRelShuttle;
+import org.apache.ignite.internal.sql.engine.rel.IgniteIndexScan;
+import org.apache.ignite.internal.sql.engine.rel.IgniteRel;
+import org.apache.ignite.internal.sql.engine.rel.IgniteTableScan;
+import org.apache.ignite.internal.sql.engine.schema.IgniteTable;
+import org.apache.ignite.internal.sql.engine.trait.IgniteDistribution;
+import org.apache.ignite.internal.tostring.S;
+import org.jetbrains.annotations.Nullable;
+import org.jetbrains.annotations.VisibleForTesting;
+
+/**
+ * Extracts partition pruning metadata from a physical plan. Examples:
+ * <pre>
+ * SELECT * FROM t WHERE pk = 10
+ * =>
+ * t = [ [pk=10] ]
+ *
+ * SELECT * FROM t WHERE pk = ?
+ * =>
+ * t = [ [pk = ?0] ]
+ *
+ * SELECT * FROM t WHERE pk = 10 AND col = 42
+ * t = [ [pk=10] ]
+ *
+ * SELECT * FROM t WHERE pk = 10 OR pk = 42
+ * =>
+ * t = [ [pk=10], [pk=42] ]
+ *
+ * SELECT * FROM t WHERE pk = 10 OR col = 42
+ * =>
+ * []
+ *
+ * SELECT * FROM t WHERE colo_key1 = 10 AND colo_key2 = 20
+ * =>
+ * t = [ [colo_key1=10], [colo_key2=20] ]
+ *
+ * SELECT * FROM t WHERE colo_key1 = 10 AND colo_key2 = 20 OR colo_key1 = 40 AND colo_key2 = 30
+ * =>
+ * t = [ [colo_key1=10, colo_key2=20], [colo_key1=20, colo_key2=30] ]
+ * </pre>
+ */
+public class PartitionPruningMetadataExtractor extends IgniteRelShuttle {
+
+ private final Long2ObjectMap<PartitionPruningColumns> result = new Long2ObjectArrayMap<>();
+
+ /**
+ * Extracts partition pruning metadata from the given physical plan.
+ * This method traverses a physical plan and attempts to extract metadata from operators that support it
+ * and includes such metadata is that complete (e.g. values of all colocation keys present in a scan predicate).
+ *
+ * @param rel Physical plan.
+ * @return Partition pruning metadata.
+ */
+ public PartitionPruningMetadata go(IgniteRel rel) {
+ rel.accept(this);
+
+ return new PartitionPruningMetadata(result);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public IgniteRel visit(IgniteIndexScan rel) {
+ RexNode condition = rel.condition();
+
+ IgniteTable table = rel.getTable().unwrap(IgniteTable.class);
+ assert table != null : "No table";
+
+ extractFromTable(rel.sourceId(), table, condition, rel.getCluster().getRexBuilder());
+
+ return super.visit(rel);
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public IgniteRel visit(IgniteTableScan rel) {
+ RexNode condition = rel.condition();
+
+ IgniteTable table = rel.getTable().unwrap(IgniteTable.class);
+ assert table != null : "No table";
+
+ extractFromTable(rel.sourceId(), table, condition, rel.getCluster().getRexBuilder());
+
+ return rel;
+ }
+
+ private void extractFromTable(long sourceId, IgniteTable table, @Nullable RexNode condition, RexBuilder rexBuilder) {
+ if (condition == null) {
+ return;
+ }
+
+ IgniteDistribution distribution = table.distribution();
+ if (!distribution.function().affinity()) {
+ return;
+ }
+ IntArrayList keysList = new IntArrayList(distribution.getKeys().size());
+ for (Integer key : distribution.getKeys()) {
+ keysList.add(key.intValue());
+ }
+
+ PartitionPruningColumns metadata = extractMetadata(keysList, condition, rexBuilder);
+
+ if (metadata != null) {
+ result.put(sourceId, metadata);
+ }
+ }
+
+ /** Extracts values of colocated columns from the given condition. */
+ @VisibleForTesting
+ @Nullable
+ public static PartitionPruningColumns extractMetadata(IntList keys, RexNode condition, RexBuilder rexBuilder) {
+ Result res = extractMetadata(condition, keys, rexBuilder, false);
+
+ // Both unknown condition and additional condition can not be used to extract metadata.
+ if (res == Result.UNKNOWN || res == Result.RESTRICT) {
+ return null;
+ }
+
+ PruningColumnSets columnSets;
+
+ if (res instanceof PruningColumnSet) {
+ PruningColumnSet columnSet = (PruningColumnSet) res;
+ columnSets = new PruningColumnSets();
+ columnSets.candidates.add(columnSet);
+ } else {
+ columnSets = (PruningColumnSets) res;
+ }
+
+ for (PruningColumnSet columnSet : columnSets.candidates) {
+
+ if (!columnSet.columns.keySet().containsAll(keys)) {
+ return null;
+ }
+ }
+
+ List<Int2ObjectMap<RexNode>> result = new ArrayList<>(columnSets.candidates.size());
+
+ for (PruningColumnSet columnSet : columnSets.candidates) {
+ result.add(columnSet.columns);
+ }
+
+ return new PartitionPruningColumns(result);
+ }
+
+ private static Result extractMetadata(RexNode node, IntList keys, RexBuilder rexBuilder, boolean negate) {
+
+ if (isColocationKey(node, keys)) {
+ // a standalone <bool_col> ref.
+ if (node.getType().getSqlTypeName() == SqlTypeName.BOOLEAN) {
+ RexLocalRef ref = (RexLocalRef) node;
+ return new PruningColumnSet(ref.getIndex(), rexBuilder.makeLiteral(!negate));
+ }
+ } else if (node.isA(SqlKind.LOCAL_REF)) {
+ return Result.RESTRICT;
+ }
+
+ List<RexNode> operands;
+ if (node instanceof RexCall) {
+ operands = ((RexCall) node).getOperands();
+ } else {
+ return Result.UNKNOWN;
+ }
+
+ switch (node.getKind()) {
+ case IS_NOT_DISTINCT_FROM:
+ case EQUALS: {
+ // NOT (c1 = <val>) -> c1 != <val> and we can not use a value.
+ // But NOT (c1 != <val>) -> c1 = <val>, and it can be used.
+ RexNode lhs;
+ RexNode rhs;
+
+ if (operands.get(0).isA(SqlKind.LOCAL_REF)) {
+ lhs = operands.get(0);
+ rhs = operands.get(1);
+ } else {
+ lhs = operands.get(1);
+ rhs = operands.get(0);
+ }
+
+ if (isColocationKey(lhs, keys) && isValueExpr(rhs)) {
+ if (negate) {
+ return Result.UNKNOWN;
+ } else {
+ RexLocalRef column = (RexLocalRef) lhs;
+ return new PruningColumnSet(column.getIndex(), rhs);
+ }
+ } else if (lhs.isA(SqlKind.LOCAL_REF) && isValueExpr(rhs)) {
+ // some column = <val> - preserve, in case of AND it can be ignored.
+ return Result.RESTRICT;
+ } else {
+ // Not a simple expression.
+ return Result.UNKNOWN;
+ }
+ }
+ case NOT_EQUALS:
+ case IS_DISTINCT_FROM: {
+ RexNode lhs;
+ RexNode rhs;
+
+ if (operands.get(0).isA(SqlKind.LOCAL_REF)) {
+ lhs = operands.get(0);
+ rhs = operands.get(1);
+ } else {
+ lhs = operands.get(1);
+ rhs = operands.get(0);
+ }
+
+ if (isColocationKey(lhs, keys) && isValueExpr(rhs)) {
+ // NOT(colo_key != <val>) => colo_key = <val>
+ if (negate) {
+ RexLocalRef column = (RexLocalRef) lhs;
+ return new PruningColumnSet(column.getIndex(), rhs);
+ } else {
+ return Result.UNKNOWN;
+ }
+ } else if (lhs.isA(SqlKind.LOCAL_REF) && isValueExpr(rhs)) {
+ // some column != <val> - preserve, in case of AND it can be ignored.
+ return Result.RESTRICT;
+
+ } else {
+ // Not a simple expression.
+ return Result.UNKNOWN;
+ }
+ }
+ case OR: {
+ PruningColumnSets res = new PruningColumnSets();
+
+ for (RexNode operand : operands) {
+ Result child = extractMetadata(operand, keys, rexBuilder, negate);
+
+ // In case of OR: we can not ignore additional condition,
+ // because OR increases the search space.
+ if (child == Result.UNKNOWN || child == Result.RESTRICT) {
+ return Result.UNKNOWN;
+ }
+
+ res.add(child);
+ }
+
+ return res;
+ }
+ case AND: {
+ PruningColumnSets res = new PruningColumnSets();
+
+ for (RexNode operand : operands) {
+ Result child = extractMetadata(operand, keys, rexBuilder, negate);
+
+ if (child == Result.UNKNOWN) {
+ return Result.UNKNOWN;
+ }
+
+ // In case of AND: we can ignore additional condition, because
+ // such condition only narrows the search space.
+ if (child == Result.RESTRICT) {
+ continue;
+ }
+
+ res.combine(child);
+
+ if (res.conflict) {
+ return Result.UNKNOWN;
+ }
+ }
+
+ return res;
+ }
+ case SEARCH: {
+ RexNode expandedSearch = RexUtil.expandSearch(rexBuilder, null, node);
+ assert !expandedSearch.isA(SqlKind.SEARCH) : "Search operation is not expanded: " + node;
+
+ return extractMetadata(expandedSearch, keys, rexBuilder, false);
+ }
+ case NOT: {
+ if (isColocationKey(operands.get(0), keys)) {
+ RexLocalRef column = (RexLocalRef) operands.get(0);
+
+ return new PruningColumnSet(column.getIndex(), rexBuilder.makeLiteral(negate));
+ } else {
+ return extractMetadata(operands.get(0), keys, rexBuilder, !negate);
+ }
+ }
+ case IS_NULL:
+ case IS_NOT_NULL: {
+ RexNode operand = operands.get(0);
+
+ if (operand.isA(SqlKind.LOCAL_REF)) {
+ return Result.RESTRICT;
+ } else {
+ // If we reach this branch with a colocated key which is never null,
+ // then IS_NOT_NULL is always true / IS_NULL is always false so there is something wrong here.
+ return Result.UNKNOWN;
+ }
+ }
+ case IS_FALSE:
+ case IS_TRUE: {
+ RexNode operand = operands.get(0);
+
+ if (isColocationKey(operand, keys)) {
+ RexLocalRef ref = (RexLocalRef) operand;
+
+ boolean value;
+ if (negate) {
+ // NOT (col IS FALSE) => col IS TRUE => col = true
+ value = node.getKind() == SqlKind.IS_FALSE;
+ } else {
+ // NOT (col IS TRUE) => col IS FALSE => col = false
+ value = node.getKind() == SqlKind.IS_TRUE;
+ }
+
+ return new PruningColumnSet(ref.getIndex(), rexBuilder.makeLiteral(value));
+ } else if (operand.isA(SqlKind.LOCAL_REF)) {
+ return Result.RESTRICT;
+ } else {
+ return Result.UNKNOWN;
+ }
+ }
+ case IS_NOT_FALSE:
+ case IS_NOT_TRUE: {
+ // IS_NOT_TRUE is used by case/when expression rewriter.
+ // IS_NOT_FALSE is added for symmetry.
+
+ boolean value;
+ if (negate) {
+ value = node.getKind() != SqlKind.IS_NOT_FALSE;
+ } else {
+ value = node.getKind() != SqlKind.IS_NOT_TRUE;
+ }
+
+ if (isColocationKey(operands.get(0), keys)) {
+ RexLocalRef column = (RexLocalRef) operands.get(0);
+ return new PruningColumnSet(column.getIndex(), rexBuilder.makeLiteral(value));
+ } else {
+ return extractMetadata(operands.get(0), keys, rexBuilder, !negate);
+ }
+ }
+ default: {
+ if (node.isA(SqlKind.BINARY_COMPARISON)) {
+ // Convert binary comparision operations to Result::RESTRICT
+ RexNode lhs;
+ RexNode rhs;
+
+ if (operands.get(0).isA(SqlKind.LOCAL_REF)) {
+ lhs = operands.get(0);
+ rhs = operands.get(1);
+ } else {
+ lhs = operands.get(1);
+ rhs = operands.get(0);
+ }
+ if (isColocationKey(lhs, keys) && isValueExpr(rhs)) {
+ // We can not extract values from expressions such as colo_key > 10
+ return Result.UNKNOWN;
+ } else if (lhs.isA(SqlKind.LOCAL_REF) && isValueExpr(rhs)) {
+ // We can use non_colo_key > 10 to narrow the search space.
+ return Result.RESTRICT;
+ } else {
+ return Result.UNKNOWN;
+ }
+ } else {
+ return Result.UNKNOWN;
+ }
+ }
+ }
+ }
+
+ private static boolean isColocationKey(RexNode node, IntList keys) {
+ if (node instanceof RexLocalRef) {
+ RexLocalRef localRef = (RexLocalRef) node;
+ return keys.contains(localRef.getIndex());
+ } else {
+ return false;
+ }
+ }
+
+ private static boolean isValueExpr(RexNode node) {
+ return node instanceof RexLiteral || node instanceof RexDynamicParam;
+ }
+
+ /** Intermediate result of extracting partition pruning metadata. */
+ private abstract static class Result {
+
+ /** Represents an unknown result, information can not used. */
+ private static final Result UNKNOWN = new Result() {
+ @Override
+ public String toString() {
+ return "<unknown>";
+ }
+ };
+
+ /** Represents additional condition, such condition can be ignored by AND operator, since they narrow the search space. */
+ private static final Result RESTRICT = new Result() {
+ @Override
+ public String toString() {
+ return "<restrict>";
+ }
+ };
+ }
+
+ /** A set of colocation key columns with their values. */
+ private static class PruningColumnSet extends Result {
+
+ private final Int2ObjectMap<RexNode> columns;
+
+ PruningColumnSet(Int2ObjectMap<RexNode> columns) {
+ this.columns = columns;
+ }
+
+ PruningColumnSet(int column, RexNode value) {
+ columns = new Int2ObjectArrayMap<>();
+ columns.put(column, value);
+ }
+
+ @Override
+ public String toString() {
+ return columns.toString();
+ }
+ }
+
+ /** A collection of colocation key column sets. */
+ private static class PruningColumnSets extends Result {
+
+ private final List<PruningColumnSet> candidates = new ArrayList<>();
+
+ private boolean conflict;
+
+ /** Adds the given result to this collection of sets. Argument can be either PruningColumnSet or PruningColumnSets. */
+ void add(Result res) {
+ // Add (OR) simply adds another candidate
+ //
+ // lhs : [c1 = 1, c2 = 2]
+ // rhs : [c3 = 3]
+ // result : [c1 = 1, c2 = 2], [c3 = 3]
+ //
+ // lhs : [c1 = 1, c2 = 2]
+ // rhs : [a = 2]
+ // result : [c1 = 1, c2 = 2], [a = 2]
+ //
+ // lhs : [c1 = 1, c2 = 2]
+ // rhs : [c1 = 2, c2 = 3]
+ // result : [c1 = 1, c2 = 2], [c1 = 2, c2 = 3]
+
+
+ if (res instanceof PruningColumnSet) {
+ PruningColumnSet columnSet = (PruningColumnSet) res;
+ candidates.add(columnSet);
+ } else {
+ PruningColumnSets columnSets = (PruningColumnSets) res;
+ candidates.addAll(columnSets.candidates);
+ }
+ }
+
+ /** Combines this column sets with the given result. Argument can be either PruningColumnSet or PruningColumnSets. */
+ void combine(Result res) {
+ // Combine (AND) merges each existing column sets with each candidate column set an for each of those
+ // produces new result
+
+ // lhs : [c1 = 1, c2 = 2]
+ // rhs : [c3 = 3]
+ // result : [c1 = 1, c2 = 2, c3 = 3]
+
+ // lhs : [c1 = 1, c2 = 2]
+ // rhs : [c3 = 3], [c4 = 4]
+ // result : [c1 = 1, c2 = 2, c3 = 3], [c1 = 1, c2 = 2, c4 = 4]
+
+ // ADD also does not allow to have the same columns with different values:
+
+ // lhs : [c1 = 1, c2 = 2]
+ // rhs : [c1 = 2]
+ // result : conflict -> c1 # c1 can't be equal to both 1 and 2
+
+ // lhs : [c1 = 1, c2 = 2]
+ // rhs : [c1 = 2, c2 = 3]
+ // result : conflict -> c1, c2
+
+ if (candidates.isEmpty()) {
+ if (res instanceof PruningColumnSet) {
+ PruningColumnSet columnSet = (PruningColumnSet) res;
+ candidates.add(columnSet);
+ } else {
+ PruningColumnSets other = (PruningColumnSets) res;
+ candidates.addAll(other.candidates);
+ }
+ } else {
+
+ if (conflict) {
+ return;
+ }
+
+ PruningColumnSets other;
+ if (res instanceof PruningColumnSet) {
+ other = new PruningColumnSets();
+ other.candidates.add((PruningColumnSet) res);
+ } else {
+ other = (PruningColumnSets) res;
+ }
+
+ List<PruningColumnSet> newOutput = new ArrayList<>();
+
+ for (PruningColumnSet candidate : other.candidates) {
+ for (PruningColumnSet val : candidates) {
+ for (Entry<RexNode> ckv : candidate.columns.int2ObjectEntrySet()) {
+ if (!val.columns.containsKey(ckv.getIntKey())) {
+ continue;
+ }
+
+ Object existing = val.columns.get(ckv.getIntKey());
+ if (!Objects.equals(existing, ckv.getValue())) {
+ conflict = true;
+ return;
+ }
+ }
+
+ Int2ObjectMap<RexNode> newValue = new Int2ObjectArrayMap<>(val.columns.size() + candidate.columns.size());
+ newValue.putAll(val.columns);
+ newValue.putAll(candidate.columns);
+
+ PruningColumnSet newColSet = new PruningColumnSet(newValue);
+ newOutput.add(newColSet);
+ }
+ }
+
+ candidates.clear();
+ candidates.addAll(newOutput);
+ }
+ }
+
+ @Override
+ public String toString() {
+ return S.toString(this);
+ }
+ }
+}
diff --git a/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/planner/PartitionPruningExtractorSelfTest.java b/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/planner/PartitionPruningExtractorSelfTest.java
new file mode 100644
index 0000000000..2eab8dc6e1
--- /dev/null
+++ b/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/planner/PartitionPruningExtractorSelfTest.java
@@ -0,0 +1,389 @@
+/*
+ * 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.ignite.internal.sql.engine.planner;
+
+import static org.apache.ignite.internal.lang.IgniteStringFormatter.format;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import it.unimi.dsi.fastutil.ints.IntImmutableList;
+import it.unimi.dsi.fastutil.ints.IntList;
+import java.util.Collections;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.SqlOperator;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
+import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.ignite.internal.sql.engine.prepare.pruning.PartitionPruningColumns;
+import org.apache.ignite.internal.sql.engine.prepare.pruning.PartitionPruningMetadataExtractor;
+import org.apache.ignite.internal.sql.engine.util.Commons;
+import org.apache.ignite.internal.testframework.BaseIgniteAbstractTest;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+/**
+ * Tests basic operations that extract partition pruning metadata using non-optimized expressions. For tests against optimized expressions
+ * see {@link PartitionPruningMetadataTest}.
+ */
+public class PartitionPruningExtractorSelfTest extends BaseIgniteAbstractTest {
+
+ private static final RelDataTypeFactory TYPE_FACTORY = Commons.typeFactory();
+
+ private static final RexBuilder rexBuilder = new RexBuilder(TYPE_FACTORY);
+
+ enum TestCase {
+ SIMPLE_1a(
+ // col0 = 1 => [0=1] colo keys = [0]
+ eq(intCol(0), lit(1)),
+ List.of(0), List.of("[0=1]")
+ ),
+ SIMPLE_1b(
+ // 1 = col0 same result as col1 = 1
+ eq(lit(1), intCol(0)),
+ List.of(0), List.of("[0=1]")
+ ),
+ SIMPLE_2c(
+ // another colo key.
+ eq(intCol(0), lit(1)),
+ List.of(1), List.of()
+ ),
+ SIMPLE_2d(
+ // col1 key is missing
+ eq(intCol(0), lit(1)),
+ List.of(0, 1), List.of()
+ ),
+
+ // OR
+
+ OR_1a(
+ // col0 = 1 OR col0 = 2 => two alternatives
+ or(
+ eq(intCol(0), lit(1)),
+ eq(intCol(0), lit(2))
+ ),
+ List.of(0),
+ List.of("[0=1]", "[0=2]")
+ ),
+
+ OR_1b(
+ // col0 = 1 OR col1 = 1 => no result both as both cols are colo keys.
+ or(
+ eq(intCol(0), lit(1)),
+ eq(intCol(1), lit(2))
+ ),
+ List.of(0, 1),
+ List.of()
+ ),
+
+ OR_1c(
+ // col0 = 1 OR col1 = 2 => no result col1 is not a colo key
+ or(
+ eq(intCol(0), lit(1)),
+ eq(intCol(1), lit(2))
+ ),
+ List.of(0),
+ List.of()
+ ),
+
+ AND_1a(
+ // col0 = 1 AND col1 = 2 => [0=1, 1=2]
+ and(
+ eq(intCol(0), lit(1)),
+ eq(intCol(1), lit(2))
+ ),
+ List.of(0, 1),
+ List.of("[0=1, 1=2]")
+ ),
+ AND_1b(
+ // col0 = 1 AND col0 = 2 => contradiction => no result
+ and(
+ eq(intCol(0), lit(1)),
+ eq(intCol(0), lit(2))
+ ),
+ List.of(0),
+ List.of()
+ ),
+ AND_1c(
+ // col0 = 1 AND col1 = 2 => [0=1] since col1 is not a colo key
+ and(
+ eq(intCol(0), lit(1)),
+ eq(intCol(1), lit(2))
+ ),
+ List.of(0),
+ List.of("[0=1]")
+ ),
+ AND_1d(
+ // col0 = 1 AND col1 = 2 AND col0 = 1 => [0=1, 1=2]
+ and(
+ eq(intCol(0), lit(1)),
+ eq(intCol(1), lit(2)), eq(intCol(0), lit(1))
+ ),
+ List.of(0, 1),
+ List.of("[0=1, 1=2]")
+ ),
+
+ DYNAMIC_PARAM_1a(
+ // col0 = ?
+ eq(intCol(0), intDynParam(0)),
+ List.of(0),
+ List.of("[0=?0]")
+ ),
+
+ DYNAMIC_PARAM_1b(
+ // col0 = 1 OR col0 = ?
+ or(
+ eq(intCol(0), lit(1)),
+ eq(intCol(0), intDynParam(0))
+ ),
+ List.of(0),
+ List.of("[0=1]", "[0=?0]")
+ ),
+
+ DYNAMIC_PARAM_1c(
+ // col0 = 1 AND col0 = ? => contradiction.
+ and(
+ eq(intCol(0), lit(1)),
+ eq(intCol(0), intDynParam(0))
+ ),
+ List.of(0),
+ List.of()
+ ),
+
+ COMPOUND_1a(
+ // col0 = 1 AND col2 = 2 OR col0=3 AND col2=4
+ or(
+ and(
+ eq(intCol(0), lit(1)),
+ eq(intCol(1), lit(2))
+ ),
+ and(
+ eq(intCol(0), lit(3)),
+ eq(intCol(1), lit(4))
+ )
+ ),
+ List.of(0, 1), List.of("[0=1, 1=2]", "[0=3, 1=4]")
+ ),
+ COMPOUND_1b(
+ // col0 = 1 AND col2 = 2 OR col0=3 => []
+ or(
+ and(
+ eq(intCol(0), lit(1)),
+ eq(intCol(1), lit(2))
+ ),
+ eq(intCol(0), lit(3))
+ ),
+ List.of(0, 1),
+ List.of()
+ ),
+ COMPOUND_1c(
+ // col0 = 1 AND col2 = 2 OR col0=3 => [col=1], [col0=3] since col0 is the only colo key.
+ or(
+ and(
+ eq(intCol(0), lit(1)),
+ eq(intCol(1), lit(2))
+ ),
+ eq(intCol(0), lit(3))
+ ),
+ List.of(0),
+ List.of("[0=1]", "[0=3]")
+ ),
+
+ // NEGATE
+
+ NEGATE_1a(
+ // NOT(col0 = 1) => col0 != 1 => []
+ not(eq(intCol(0), lit(1))),
+ List.of(0),
+ List.of()
+ ),
+ NEGATE_1b(
+ // NOT(col0 != 1) => col0 = 1 => [0=1]
+ not(notEq(intCol(0), lit(1))),
+ List.of(0),
+ List.of("[0=1]")
+ ),
+ NEGATE_1c(
+ // IS_NOT_TRUE(col0 = 1) => col0 != 1 => []
+ call(SqlStdOperatorTable.IS_NOT_TRUE, eq(intCol(0), lit(1))),
+ List.of(0),
+ List.of()
+ ),
+ NEGATE_1d(
+ // IS_NOT_TRUE(col0 != 1) => col0 = 1 => [0=1]
+ call(SqlStdOperatorTable.IS_NOT_TRUE, notEq(intCol(0), lit(1))),
+ List.of(0),
+ List.of("[0=1]")
+ ),
+
+ // bools
+ NEGATE_2a(
+ // NOT(bool col0) => [col0=false]
+ not(boolCol(0)), List.of(0), List.of("[0=false]")
+ ),
+ NEGATE_2b(
+ // NOT(NOT(bool col0)) => [col0=true]
+ not(not(boolCol(0))),
+ List.of(0), List.of("[0=true]")
+ ),
+ NEGATE_2c(
+ // NOT(IS_TRUE(bool col0)) => [col0=false]
+ not(call(SqlStdOperatorTable.IS_TRUE, boolCol(0))),
+ List.of(0), List.of("[0=false]")
+ ),
+ NEGATE_1e(
+ // NOT(IS_FALSE(bool col0)) => [col0=true]
+ not(call(SqlStdOperatorTable.IS_FALSE, boolCol(0))),
+ List.of(0), List.of("[0=true]")
+ ),
+ NEGATE_2f(
+ // IS_NOT_TRUE(bool col0) => [col0=false]
+ call(SqlStdOperatorTable.IS_NOT_TRUE, boolCol(0)),
+ List.of(0), List.of("[0=false]")
+ ),
+ NEGATE_2g(
+ // IS_NOT_FALSE(bool col0) => [col0=true]
+ call(SqlStdOperatorTable.IS_NOT_FALSE, boolCol(0)),
+ List.of(0), List.of("[0=true]")
+ ),
+ NEGATE_2h(
+ // NOT(IS_NOT_TRUE(bool col0)) => [col0=true]
+ not(call(SqlStdOperatorTable.IS_NOT_TRUE, boolCol(0))),
+ List.of(0), List.of("[0=true]")
+ ),
+ NEGATE_2i(
+ // NOT(IS_NOT_FALSE(bool col0)) => [col0=false]
+ not(call(SqlStdOperatorTable.IS_NOT_FALSE, boolCol(0))),
+ List.of(0), List.of("[0=false]")
+ ),
+
+ IS_NULL_1a(
+ // col0 IS NULL
+ call(SqlStdOperatorTable.IS_NULL, intCol(0)),
+ List.of(0), List.of()
+ ),
+ IS_NULL_1b(
+ // col0 IS NOT NULL
+ call(SqlStdOperatorTable.IS_NULL, intCol(0)),
+ List.of(0), List.of()
+ )
+
+ ;
+
+ final RexNode condition;
+
+ final List<Integer> keys;
+
+ final List<String> expected;
+
+ TestCase(RexNode condition, List<Integer> keys, List<String> expected) {
+ this.condition = condition;
+ this.keys = keys;
+ this.expected = expected;
+ }
+
+ @Override
+ public String toString() {
+ return condition + " " + keys + " -> " + expected;
+ }
+ }
+
+ @ParameterizedTest
+ @EnumSource(TestCase.class)
+ public void test(TestCase test) {
+ expectMetadata(test.condition, test.keys, test.expected);
+ }
+
+ private void expectMetadata(RexNode condition, List<Integer> keys, List<String> expectedMetadata) {
+ log.info("Condition: {}", condition);
+ log.info("Keys: {}", keys);
+ log.info("Expected metadata: {}", expectedMetadata);
+
+ IntList keyList = IntImmutableList.toList(keys.stream().mapToInt(Integer::intValue));
+ PartitionPruningColumns actualMetadata = PartitionPruningMetadataExtractor.extractMetadata(keyList, condition, rexBuilder);
+
+ List<String> actual;
+
+ if (actualMetadata == null) {
+ actual = Collections.emptyList();
+ } else {
+ // [1 = RexExpr1, 2 = RexExpr2] => [1 = "expr1-as-string", 2 = "expr2-as-string" ]
+ actual = PartitionPruningColumns.canonicalForm(actualMetadata).stream()
+ .map(String::valueOf)
+ .collect(Collectors.toList());
+ }
+
+ log.info("Actual metadata: {}", actual);
+
+ assertEquals(expectedMetadata, actual);
+ }
+
+ private static RexNode intCol(int i) {
+ return rexBuilder.makeLocalRef(TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER), i);
+ }
+
+ private static RexNode boolCol(int i) {
+ return rexBuilder.makeLocalRef(TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER), i);
+ }
+
+ private static RexNode lit(Object value) {
+ if (value instanceof Boolean) {
+ return rexBuilder.makeLiteral(value, type(SqlTypeName.BOOLEAN));
+ } else if (value instanceof Integer) {
+ return rexBuilder.makeLiteral(value, type(SqlTypeName.INTEGER));
+ } else if (value == null) {
+ return rexBuilder.makeLiteral(null, type(SqlTypeName.NULL));
+ } else {
+ throw new IllegalArgumentException(format("Unexpected literal: <{}> of type: {}", value, value.getClass()));
+ }
+ }
+
+ private static RexNode intDynParam(int i) {
+ return rexBuilder.makeDynamicParam(TYPE_FACTORY.createSqlType(SqlTypeName.INTEGER), i);
+ }
+
+ private static RexNode or(RexNode... exprs) {
+ return rexBuilder.makeCall(SqlStdOperatorTable.OR, exprs);
+ }
+
+ private static RexNode and(RexNode... exprs) {
+ return rexBuilder.makeCall(SqlStdOperatorTable.AND, exprs);
+ }
+
+ private static RexNode eq(RexNode lhs, RexNode rhs) {
+ return rexBuilder.makeCall(SqlStdOperatorTable.EQUALS, lhs, rhs);
+ }
+
+ private static RexNode not(RexNode expr) {
+ return rexBuilder.makeCall(SqlStdOperatorTable.NOT, expr);
+ }
+
+ private static RexNode notEq(RexNode lhs, RexNode rhs) {
+ return rexBuilder.makeCall(SqlStdOperatorTable.NOT_EQUALS, lhs, rhs);
+ }
+
+ private static RexNode call(SqlOperator op, RexNode... exprs) {
+ return rexBuilder.makeCall(op, exprs);
+ }
+
+ private static RelDataType type(SqlTypeName typeName) {
+ return TYPE_FACTORY.createSqlType(typeName);
+ }
+}
diff --git a/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/planner/PartitionPruningMetadataTest.java b/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/planner/PartitionPruningMetadataTest.java
new file mode 100644
index 0000000000..2e6af22748
--- /dev/null
+++ b/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/planner/PartitionPruningMetadataTest.java
@@ -0,0 +1,473 @@
+/*
+ * 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.ignite.internal.sql.engine.planner;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.apache.calcite.plan.RelOptUtil;
+import org.apache.calcite.sql.SqlExplainFormat;
+import org.apache.calcite.sql.SqlExplainLevel;
+import org.apache.ignite.internal.sql.engine.framework.TestBuilders;
+import org.apache.ignite.internal.sql.engine.prepare.pruning.PartitionPruningColumns;
+import org.apache.ignite.internal.sql.engine.prepare.pruning.PartitionPruningMetadata;
+import org.apache.ignite.internal.sql.engine.prepare.pruning.PartitionPruningMetadataExtractor;
+import org.apache.ignite.internal.sql.engine.rel.IgniteRel;
+import org.apache.ignite.internal.sql.engine.schema.ColumnDescriptor;
+import org.apache.ignite.internal.sql.engine.schema.IgniteSchema;
+import org.apache.ignite.internal.sql.engine.schema.IgniteTable;
+import org.apache.ignite.internal.sql.engine.schema.TableDescriptor;
+import org.apache.ignite.internal.sql.engine.trait.IgniteDistributions;
+import org.apache.ignite.internal.type.NativeTypes;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+
+/**
+ * Tests for {@link PartitionPruningMetadataExtractor} against optimized expressions.
+ */
+public class PartitionPruningMetadataTest extends AbstractPlannerTest {
+
+ private static final IgniteSchema TABLE_C1 = createSchema(TestBuilders.table().name("T")
+ .addKeyColumn("C1", NativeTypes.INT32)
+ .addColumn("C2", NativeTypes.INT32, false)
+ .distribution(IgniteDistributions.affinity(List.of(0), 1, 2))
+ .build());
+
+ private static final IgniteSchema TABLE_C1_NULLABLE_C2 = createSchema(TestBuilders.table().name("T")
+ .addKeyColumn("C1", NativeTypes.INT32)
+ .addColumn("C2", NativeTypes.INT32, true)
+ .distribution(IgniteDistributions.affinity(List.of(0), 1, 2))
+ .build());
+
+ private static final IgniteSchema TABLE_C1_C2 = createSchema(TestBuilders.table().name("T")
+ .addKeyColumn("C1", NativeTypes.INT32)
+ .addKeyColumn("C2", NativeTypes.INT32)
+ .addColumn("C3", NativeTypes.INT32, false)
+ .distribution(IgniteDistributions.affinity(List.of(0, 1), 1, 2))
+ .build());
+
+ private static final IgniteSchema TABLE_C1_C2_C3 = createSchema(TestBuilders.table().name("T")
+ .addKeyColumn("C1", NativeTypes.INT32)
+ .addKeyColumn("C2", NativeTypes.INT32)
+ .addKeyColumn("C3", NativeTypes.INT32)
+ .addColumn("C4", NativeTypes.INT32, false)
+ .distribution(IgniteDistributions.affinity(List.of(0, 1, 2), 1, 2))
+ .build());
+
+ private static final IgniteSchema TABLE_BOOL_C1 = createSchema(TestBuilders.table().name("T")
+ .addKeyColumn("C1", NativeTypes.BOOLEAN)
+ .addColumn("C2", NativeTypes.INT32, false)
+ .distribution(IgniteDistributions.affinity(List.of(0), 1, 2))
+ .build());
+
+ private static final IgniteSchema TABLE_BOOL_C1_C3 = createSchema(TestBuilders.table().name("T")
+ .addKeyColumn("C1", NativeTypes.BOOLEAN)
+ .addKeyColumn("C2", NativeTypes.BOOLEAN)
+ .addColumn("C3", NativeTypes.INT32, false)
+ .distribution(IgniteDistributions.affinity(List.of(0, 1), 1, 2))
+ .build());
+
+ private static final IgniteSchema TABLE_C1_BOOLS = createSchema(TestBuilders.table().name("T")
+ .addKeyColumn("C1", NativeTypes.INT32)
+ .addColumn("C2", NativeTypes.BOOLEAN, false)
+ .distribution(IgniteDistributions.affinity(List.of(0), 1, 2))
+ .build());
+
+ private static final IgniteSchema TABLE_ALL_BOOLS_C1 = createSchema(TestBuilders.table().name("T")
+ .addKeyColumn("C1", NativeTypes.BOOLEAN)
+ .addColumn("C2", NativeTypes.BOOLEAN, false)
+ .distribution(IgniteDistributions.affinity(List.of(0), 1, 2))
+ .build());
+
+ /** Basic test cases for partition pruning metadata extractor. */
+ @ParameterizedTest
+ @EnumSource(TestCaseBasic.class)
+ public void testBasic(TestCaseBasic testCaseSimple) {
+ checkPruningMetadata(testCaseSimple.data);
+ }
+
+ enum TestCaseBasic {
+ // SELECT * FROM t WHERE ... is omitted
+
+ SIMPLE_1a("c1 = 42", TABLE_C1, "[c1=42]"),
+ SIMPLE_1b("42 = c1", TABLE_C1, "[c1=42]"),
+ SIMPLE_1c("c1 = ?", TABLE_C1, "[c1=?0]"),
+ SIMPLE_1d("? = c1", TABLE_C1, "[c1=?0]"),
+ SIMPLE_1e("c1 = '42'::INTEGER", TABLE_C1, "[c1=42]"),
+
+ // OR
+
+ OR_1a("c1 = 1 OR c1 = 2", TABLE_C1, "[c1=1]", "[c1=2]"),
+ OR_1b("c1 = 1 OR c1 = 1", TABLE_C1, "[c1=1]"),
+
+ OR_1c("c1 = 42 OR c1 = ?", TABLE_C1, "[c1=42]", "[c1=?0]"),
+ OR_1d("c1 = 42 OR c1 = ? OR c1 = ?", TABLE_C1, "[c1=42]", "[c1=?0]", "[c1=?1]"),
+
+ OR_1e("c1 = 1 OR c2 = 2", TABLE_C1),
+ OR_1f("c1 = 1 OR (c2 = 2 OR c1 = 3)", TABLE_C1),
+
+ OR_2a("c1 = 1 OR c1 > 10", TABLE_C1),
+
+ OR_BOOL_EXPR_1a("c1 = 42 OR true", TABLE_C1),
+ OR_BOOL_EXPR_1b("c1 = 42 OR false", TABLE_C1, "[c1=42]"),
+ // filters=[=($t0, 42)], because c2 is not nullable
+ OR_BOOL_EXPR_1c("c1 = 42 OR c2 IS NULL", TABLE_C1, "[c1=42]"),
+ // filters=[true]
+ OR_BOOL_EXPR_1d("c1 = 42 OR c2 IS NOT NULL", TABLE_C1),
+ // filters=[OR(=($t0, 42), $t1)]
+ OR_BOOL_EXPR_1e("c1 = 42 OR c2 IS TRUE", TABLE_C1_BOOLS),
+ // filters=[OR(=($t0, 42), NOT($t1))]
+ OR_BOOL_EXPR_1f("c1 = 42 OR c2 IS NOT TRUE", TABLE_C1_BOOLS),
+ // filters=[OR(=($t0, 42), NOT($t1))]
+ OR_BOOL_EXPR_1g("c1 = 42 OR c2 IS FALSE", TABLE_C1_BOOLS),
+ // filters=[OR(=($t0, 42), $t1)]
+ OR_BOOL_EXPR_1h("c1 = 42 OR c2 IS NOT FALSE", TABLE_C1_BOOLS),
+
+ // AND
+
+ AND_1a("c1 = 42 AND c2 = 3", TABLE_C1, "[c1=42]"),
+ AND_1b("c2 = 3 AND c1 = 42", TABLE_C1, "[c1=42]"),
+
+ AND_1c("c1 = 42 AND c2 = 3", TABLE_C1, "[c1=42]"),
+ AND_1d("c2 = 3 AND c1 = 42", TABLE_C1, "[c1=42]"),
+
+ AND_2a("c1 = 42 AND c2 = 3", TABLE_C1_C2, "[c1=42, c2=3]"),
+ AND_2b("c1 = 42 AND c2 = 3 AND c3 = 100", TABLE_C1_C2, "[c1=42, c2=3]"),
+
+ // Column has different values -> conflict
+ AND_3a("c1 = 1 AND c2 = 2 AND c1 = 2", TABLE_C1_C2),
+ AND_3b("c1 = 1 AND c2 = 2 AND c1 = 2 and c2 = 3", TABLE_C1_C2),
+ AND_3c("c1 = 1 AND c2 = 2 AND c1 = 2 or c3 = 3", TABLE_C1_C2),
+
+ // Column might have different values.
+ // The query probably contains an error c1 = constant and a dynamic parameter.
+ AND_3d("c1 = 42 AND c1 = ?", TABLE_C1),
+ AND_3e("c1 = 42 AND c1 = ? AND c2 = ?", TABLE_C1_C2),
+
+ AND_4a("c1 = 1 AND c2 = 2 AND c3 = 3 and c4 = 4", TABLE_C1_C2_C3, "[c1=1, c2=2, c3=3]"),
+ AND_4b("c1 = 1 AND c2 = 2 AND c3 = 3 OR c3 = 4 AND c2 = 5 AND c1 = 6", TABLE_C1_C2_C3, "[c1=1, c2=2, c3=3]", "[c1=6, c2=5, c3=4]"),
+
+ // Some colocation key columns are missing
+ MISSING_KEYS_1a("c1 = 1 AND c2 = 10 OR c1 = 2", TABLE_C1_C2),
+ MISSING_KEYS_1b("c1 = 1 OR c1 = 2 AND c2 = 10 ", TABLE_C1_C2),
+
+ MISSING_KEYS_2a("c1 = 1 AND c2 = 10 OR c1 = 2", TABLE_C1_C2),
+ MISSING_KEYS_2b("c1 = 1 OR c1 = 2 AND c2 = 10 ", TABLE_C1_C2),
+
+ // AND restrict (additional condition on non colocated column still allows to extract column values)
+
+ AND_RESTRICT_1a("c1 = 42 AND c2 > 10", TABLE_C1, "[c1=42]"),
+ AND_RESTRICT_1b("c1 = 42 AND c2 >= 10", TABLE_C1, "[c1=42]"),
+ AND_RESTRICT_1c("c1 = 42 AND c2 < 10", TABLE_C1, "[c1=42]"),
+ AND_RESTRICT_1d("c1 = 42 AND c2 <= 10", TABLE_C1, "[c1=42]"),
+ AND_RESTRICT_1e("c1 = 42 AND c2 != 10", TABLE_C1, "[c1=42]"),
+ AND_RESTRICT_1f("c1 = 42 AND c2 <> 10", TABLE_C1, "[c1=42]"),
+
+ // both c2 AND c1 are not nullable.
+ AND_BOOL_EXPR_1a("c1 = 42 AND c2 IS NULL", TABLE_C1),
+ AND_BOOL_EXPR_1b("c1 = 42 AND c2 IS NOT NULL", TABLE_C1, "[c1=42]"),
+ AND_BOOL_EXPR_1c("c1 = 42 AND c2 IS TRUE", TABLE_C1_BOOLS, "[c1=42]"),
+ AND_BOOL_EXPR_1d("c1 = 42 AND c2 IS NOT TRUE", TABLE_C1_BOOLS, "[c1=42]"),
+ AND_BOOL_EXPR_1e("c1 = 42 AND c2 IS FALSE", TABLE_C1_BOOLS, "[c1=42]"),
+ AND_BOOL_EXPR_1f("c1 = 42 AND c2 IS NOT FALSE", TABLE_C1_BOOLS, "[c1=42]"),
+ AND_BOOL_EXPR_1g("c1 = 42 AND true IS NOT NULL", TABLE_C1, "[c1=42]"),
+ AND_BOOL_EXPR_1h("c1 = 42 AND true", TABLE_C1, "[c1=42]"),
+ AND_BOOL_EXPR_1i("c1 = 42 AND false", TABLE_C1),
+
+ AND_OTHER_1a("c1 = 42 AND 5 > 1", TABLE_C1, "[c1=42]"),
+
+ // DISTINCT FROM
+
+ DISTINCT_FROM_1a("c1 IS NOT DISTINCT FROM 42", TABLE_C1, "[c1=42]"),
+ DISTINCT_FROM_1b("c1 IS NOT DISTINCT FROM 42 AND c2 = 10", TABLE_C1, "[c1=42]"),
+ DISTINCT_FROM_1c("c1 IS NOT DISTINCT FROM 42 OR c2 = 10", TABLE_C1),
+
+ DISTINCT_FROM_2a("c1 IS DISTINCT FROM 42", TABLE_C1),
+ DISTINCT_FROM_2b("c1 IS NOT DISTINCT FROM null", TABLE_C1),
+ DISTINCT_FROM_2c("c1 IS DISTINCT FROM null", TABLE_C1),
+
+ // IN
+
+ IN_1a("c1 IN (1, 2, 3)", TABLE_C1, "[c1=1]", "[c1=2]", "[c1=3]"),
+ IN_1b("c1 NOT IN (1, 2, 3)", TABLE_C1),
+ IN_1c("c1 IN (1, 2, c2)", TABLE_C1),
+
+ // BETWEEN
+
+ BETWEEN_1a("c1 BETWEEN 1 AND 10", TABLE_C1),
+ // Rewritten to Sarg
+ BETWEEN_1b("c1 BETWEEN 1 AND 1", TABLE_C1, "[c1=1]"),
+
+ // IS NULL
+
+ IS_NULL_1a("c1 IS NULL", TABLE_C1),
+ // should be converted to false by the optimizer
+ IS_NULL_1b("c1 = 42 AND c1 IS NULL", TABLE_C1),
+ IS_NULL_1c("c1 = 42 OR c1 IS NOT NULL", TABLE_C1),
+
+ IS_NOT_NULL_1a("c1 IS NOT NULL", TABLE_C1),
+ IS_NOT_NULL_1b("c1 = 42 AND c1 IS NOT NULL", TABLE_C1, "[c1=42]"),
+ IS_NOT_NULL_1c("c1 = 42 OR c1 IS NOT NULL", TABLE_C1),
+
+ // c2 is nullable.
+ IS_NOT_NULL_2a("c1 = 42 AND c2 IS NULL", TABLE_C1_NULLABLE_C2, "[c1=42]"),
+
+ // Negation
+
+ NEGATE_1a("c1 != 42", TABLE_C1),
+ // filters=[<>($t0, 42)]
+ NEGATE_1b("NOT(c1 = 42)", TABLE_C1),
+ // filters=[=($t0, 42)]
+ NEGATE_1c("NOT(c1 != 42)", TABLE_C1, "[c1=42]"),
+
+ // filters=[OR(=($t0, 42), =($t1, 10))]
+ NEGATE_2a("NOT(c1 != 42 AND c2 != 10)", TABLE_C1_C2),
+ // filters=[AND(=($t0, 42), =($t1, 10))]
+ NEGATE_2b("NOT(c1 != 42 OR c2 != 10)", TABLE_C1_C2, "[c1=42, c2=10]"),
+ // filters=[OR(=($t0, 42), <>($t1, 10))]
+ NEGATE_2c("NOT(c1 != 42 AND c2 = 10)", TABLE_C1_C2),
+ // filters=[AND(=($t0, 42), =($t1, 10), <>($t2, 1))]
+ NEGATE_2d("NOT(c1 != 42 OR c2 != 10 OR c3=1)", TABLE_C1_C2, "[c1=42, c2=10]"),
+
+ NO_META_1a("c1 = c2", TABLE_C1),
+ NO_META_1b("c1 = c1", TABLE_C1),
+ NO_META_1c("c2 = 42", TABLE_C1),
+ NO_META_1d("true", TABLE_C1),
+ NO_META_1e("false", TABLE_C1),
+ NO_META_1f("c1 = ABS(c1)", TABLE_C1),
+ NO_META_1g("c1 = ABS(c2)", TABLE_C1),
+
+ // No metadata, complex expressions:
+ NO_META_2a("c1 = c1 + 0", TABLE_C1),
+ NO_META_2b("c1 = 0 + c1", TABLE_C1),
+ NO_META_2c("c1 = 1 * c1", TABLE_C1),
+ NO_META_2d("c1 = 2 * c1", TABLE_C1),
+ NO_META_2e("c1 = (10 + c1)", TABLE_C1),
+ NO_META_2f("c1 = (10 + c2)", TABLE_C1),
+
+ NO_META_4a("c1 > 10", TABLE_C1),
+ NO_META_4b("c1 < 10", TABLE_C1),
+ NO_META_4c("c1 <= 10", TABLE_C1),
+ NO_META_4d("c1 >= 10", TABLE_C1),
+
+ // We do not know the result of an function expression, so we can not do extract column values
+ NO_META_5a("c1 = 42 OR LENGTH('abc') = 3", TABLE_C1),
+ NO_META_5b("c1 = 42 OR SUBSTRING(c2::VARCHAR, 2) = 'a'", TABLE_C1),
+ NO_META_5c("c1 = 42 OR SUBSTRING(c1::VARCHAR, 2) = '2'", TABLE_C1),
+
+ NO_META_5d("c1 = 42 AND LENGTH(SUBSTRING('aaaaaaa', 2)) > 0", TABLE_C1_C2),
+ NO_META_5e("c1 = 42 AND c1 = SUBSTRING(c3::VARCHAR, 2)::INTEGER", TABLE_C1_C2),
+ NO_META_5f("c1 = 42 AND c2 = SUBSTRING(c3::VARCHAR, 2)::INTEGER", TABLE_C1_C2),
+
+ // No constant folding
+ CONST_FOLDING_1a("c1 = 10 + 4", TABLE_C1),
+ CONST_FOLDING_1b("c1 = ? + 4", TABLE_C1),
+
+ // 0s removed by Calcite.
+ CONST_FOLDING_1c("c1 = 0 + 42", TABLE_C1, "[c1=42]"),
+ CONST_FOLDING_1d("c1 = 42 + 0", TABLE_C1, "[c1=42]"),
+ ;
+
+ private final TestCase data;
+
+ TestCaseBasic(String condition, IgniteSchema schema, String... expected) {
+ this.data = new TestCase(condition, schema, expected);
+ }
+
+ @Override
+ public String toString() {
+ return data.toString();
+ }
+ }
+
+ /** Test cases for bool columns. */
+ @ParameterizedTest
+ @EnumSource(TestCaseBool.class)
+ public void testBool(TestCaseBool testCaseBool) {
+ checkPruningMetadata(testCaseBool.data);
+ }
+
+ enum TestCaseBool {
+
+ BOOL_1a("c1 IS TRUE", TABLE_BOOL_C1, "[c1=true]"),
+ BOOL_1b("c1 IS FALSE", TABLE_BOOL_C1, "[c1=false]"),
+ BOOL_1c("NOT c1", TABLE_BOOL_C1, "[c1=false]"),
+ BOOL_1d("NOT NOT c1", TABLE_BOOL_C1, "[c1=true]"),
+
+ // converted to filters=[true]
+ BOOL_1e("c1 OR true", TABLE_BOOL_C1),
+ BOOL_1f("c1 OR false", TABLE_BOOL_C1, "[c1=true]"),
+
+ BOOL_2a("c1 AND c2", TABLE_BOOL_C1_C3, "[c1=true, c2=true]"),
+ BOOL_2b("c1 AND c2 IS FALSE", TABLE_BOOL_C1_C3, "[c1=true, c2=false]"),
+ BOOL_2c("NOT c1 AND c2", TABLE_BOOL_C1_C3, "[c1=false, c2=true]"),
+ BOOL_2d("NOT c1 AND NOT c2", TABLE_BOOL_C1_C3, "[c1=false, c2=false]"),
+
+ BOOL_3a("c1 OR c2", TABLE_BOOL_C1_C3),
+ BOOL_3b("c1 IS TRUE OR c2", TABLE_BOOL_C1_C3),
+ BOOL_3c("c1 OR c2 IS TRUE", TABLE_BOOL_C1_C3),
+
+ BOOL_3d("c1 AND c2 IS TRUE", TABLE_ALL_BOOLS_C1, "[c1=true]"),
+ BOOL_3e("c1 AND c2 IS FALSE", TABLE_ALL_BOOLS_C1, "[c1=true]"),
+
+ // filters=[IS NOT NULL($t0)]
+ BOOL_4e("c1 = c1", TABLE_BOOL_C1),
+ ;
+
+ private final TestCase data;
+
+ TestCaseBool(String condition, IgniteSchema schema, String... expected) {
+ this.data = new TestCase(condition, schema, expected);
+ }
+
+ @Override
+ public String toString() {
+ return data.toString();
+ }
+ }
+
+ /** Test cases for CASE expression. */
+ @ParameterizedTest
+ @EnumSource(TestCaseCaseExpr.class)
+ public void testCaseExpr(TestCaseCaseExpr testCaseBool) {
+ checkPruningMetadata(testCaseBool.data);
+ }
+
+ enum TestCaseCaseExpr {
+ // filters=[=($t0, 42)]
+ CASE_1a("CASE c1 WHEN 42 THEN true ELSE false END", TABLE_C1, "[c1=42]"),
+ // filters=[IS NOT TRUE(=($t0, 2))]
+ CASE_1b("CASE c1 WHEN 42 THEN false ELSE true END", TABLE_C1),
+ // filters=[=($t0, 42)]
+ CASE_1c("CASE c1 WHEN 42 THEN true WHEN 3 THEN false ELSE false END", TABLE_C1, "[c1=42]"),
+ // filters=[OR(=($t0, 42), IS NOT TRUE(=($t0, 3)))]
+ CASE_1d("CASE c1 WHEN 42 THEN true WHEN 3 THEN false ELSE true END", TABLE_C1),
+
+ // filters=[=($t0, $t1)]
+ CASE_2a("CASE c1 WHEN c2 THEN true ELSE false END", TABLE_C1),
+ // filters=[IS NOT TRUE(=($t0, $t1))]
+ CASE_2b("CASE c1 WHEN c2 THEN false ELSE true END", TABLE_C1),
+ // filters=[=($t0, 42)]
+ CASE_2c("CASE c1 WHEN 42 THEN true WHEN c2 THEN false ELSE false END", TABLE_C1, "[c1=42]"),
+ // filters=[OR(=($t0, 42), IS NOT TRUE(=($t0, $t1)))]
+ CASE_2d("CASE c1 WHEN 42 THEN true WHEN c2 THEN false ELSE true END", TABLE_C1),
+
+ // [OR(=($t0, 10), =($t0, 42))]
+ CASE_3a("CASE c1 WHEN 42 THEN true WHEN 10 THEN true ELSE false END", TABLE_C1, "[c1=10]", "[c1=42]"),
+ ;
+ final TestCase data;
+
+ TestCaseCaseExpr(String condition, IgniteSchema schema, String... expected) {
+ this.data = new TestCase(condition, schema, expected);
+ }
+
+ @Override
+ public String toString() {
+ return data.toString();
+ }
+ }
+
+ private static class TestCase {
+
+ private final String condition;
+
+ private final IgniteSchema schema;
+
+ private final String[] expected;
+
+ TestCase(String condition, IgniteSchema schema, String... expected) {
+ this.condition = condition;
+ this.schema = schema;
+ this.expected = expected;
+ }
+
+ List<Integer> colocationKeys() {
+ IgniteTable table = (IgniteTable) schema.getTable("T");
+ return table.distribution().getKeys();
+ }
+
+ List<String> columnNames() {
+ IgniteTable table = (IgniteTable) schema.getTable("T");
+ List<String> names = new ArrayList<>();
+ TableDescriptor tableDescriptor = table.descriptor();
+
+ for (int i = 0; i < tableDescriptor.columnsCount(); i++) {
+ ColumnDescriptor col = tableDescriptor.columnDescriptor(i);
+ names.add(col.name());
+ }
+
+ return names;
+ }
+
+ @Override
+ public String toString() {
+ return condition + " -> " + Arrays.toString(expected);
+ }
+ }
+
+ private void checkPruningMetadata(TestCase testCase) {
+ String statement = "SELECT * FROM t WHERE " + testCase.condition;
+
+ List<String> expectedMetadata = Arrays.asList(testCase.expected);
+ List<Integer> colocationKeys = testCase.colocationKeys();
+ List<String> columnNames = testCase.columnNames();
+
+ log.info("Statement: {}", statement);
+ log.info("Keys: {}", colocationKeys.stream().map(columnNames::get).collect(Collectors.toList()));
+ log.info("Expected: {}", expectedMetadata);
+
+ IgniteRel igniteRel;
+ try {
+ igniteRel = physicalPlan(statement, testCase.schema);
+ } catch (Exception e) {
+ throw new IllegalStateException("Unable to build a physical plan", e);
+ }
+ log.info("Plan: {}", RelOptUtil.dumpPlan("\n", igniteRel, SqlExplainFormat.TEXT, SqlExplainLevel.NON_COST_ATTRIBUTES));
+
+ extractMetadataAndCheck(igniteRel, columnNames, expectedMetadata);
+ }
+
+ private void extractMetadataAndCheck(IgniteRel rel, List<String> columnNames, List<String> expectedMetadata) {
+ PartitionPruningMetadataExtractor extractor = new PartitionPruningMetadataExtractor();
+ PartitionPruningMetadata actual = extractor.go(rel);
+
+ List<String> actualMetadata;
+
+ if (actual.data().isEmpty()) {
+ actualMetadata = Collections.emptyList();
+ } else {
+ PartitionPruningColumns columns = actual.data().long2ObjectEntrySet().iterator().next().getValue();
+
+ // replace column indices with column names in lower case
+ actualMetadata = PartitionPruningColumns.canonicalForm(columns).stream().map(cols -> cols.stream().map(col -> {
+ String columnName = columnNames.get(col.getKey()).toLowerCase(Locale.US);
+ return Map.entry(columnName, col.getValue());
+ }).collect(Collectors.toList()).toString()).collect(Collectors.toList());
+ }
+
+ log.info("Actual metadata: {}", actualMetadata);
+
+ assertEquals(expectedMetadata, actualMetadata, "Pruning metadata");
+ }
+}
diff --git a/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/planner/PartitionPruningTest.java b/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/planner/PartitionPruningTest.java
new file mode 100644
index 0000000000..7cec257e61
--- /dev/null
+++ b/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/planner/PartitionPruningTest.java
@@ -0,0 +1,137 @@
+/*
+ * 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.ignite.internal.sql.engine.planner;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+import java.util.List;
+import org.apache.ignite.internal.sql.engine.framework.TestBuilders;
+import org.apache.ignite.internal.sql.engine.prepare.IgniteRelShuttle;
+import org.apache.ignite.internal.sql.engine.prepare.pruning.PartitionPruningColumns;
+import org.apache.ignite.internal.sql.engine.prepare.pruning.PartitionPruningMetadata;
+import org.apache.ignite.internal.sql.engine.prepare.pruning.PartitionPruningMetadataExtractor;
+import org.apache.ignite.internal.sql.engine.rel.IgniteRel;
+import org.apache.ignite.internal.sql.engine.rel.SourceAwareIgniteRel;
+import org.apache.ignite.internal.sql.engine.schema.IgniteIndex.Collation;
+import org.apache.ignite.internal.sql.engine.schema.IgniteSchema;
+import org.apache.ignite.internal.sql.engine.schema.IgniteTable;
+import org.apache.ignite.internal.sql.engine.trait.IgniteDistributions;
+import org.apache.ignite.internal.type.NativeTypes;
+import org.junit.jupiter.api.Test;
+
+/**
+ * Tests for {@link PartitionPruningMetadataExtractor} against different physical operators.
+ */
+public class PartitionPruningTest extends AbstractPlannerTest {
+
+ @Test
+ public void testTableScan() throws Exception {
+ IgniteTable table = TestBuilders.table()
+ .name("T")
+ .addKeyColumn("C1", NativeTypes.INT32)
+ .addColumn("C2", NativeTypes.INT32)
+ .distribution(IgniteDistributions.affinity(List.of(0), 1, 2))
+ .build();
+
+ PartitionPruningMetadata actual = extractMetadata(
+ "SELECT * FROM t WHERE c1 = 42",
+ table
+ );
+
+ PartitionPruningColumns cols = actual.get(1);
+ assertNotNull(cols, "No metadata for source=1");
+ assertEquals("[[0=42]]", PartitionPruningColumns.canonicalForm(cols).toString());
+ }
+
+ @Test
+ public void testIndexScan() throws Exception {
+ IgniteTable table = TestBuilders.table()
+ .name("T")
+ .addKeyColumn("C1", NativeTypes.INT32)
+ .addColumn("C2", NativeTypes.INT32)
+ .distribution(IgniteDistributions.affinity(List.of(0), 1, 2))
+ .sortedIndex().name("C1_SORTED")
+ .addColumn("C1", Collation.ASC_NULLS_FIRST)
+ .end()
+ .build();
+
+ PartitionPruningMetadata actual = extractMetadata(
+ "SELECT * FROM t WHERE c1 = 42",
+ table
+ );
+
+ PartitionPruningColumns cols = actual.get(1);
+ assertNotNull(cols, "No metadata for source=1");
+ assertEquals("[[0=42]]", PartitionPruningColumns.canonicalForm(cols).toString());
+ }
+
+ @Test
+ public void testMultipleScans() throws Exception {
+ IgniteTable table1 = TestBuilders.table()
+ .name("T1")
+ .addKeyColumn("C1", NativeTypes.INT32)
+ .addColumn("C2", NativeTypes.INT32)
+ .distribution(IgniteDistributions.affinity(List.of(0), 1, 2))
+ .build();
+
+ IgniteTable table2 = TestBuilders.table()
+ .name("T2")
+ .addKeyColumn("C1", NativeTypes.INT32)
+ .addColumn("C2", NativeTypes.INT32)
+ .distribution(IgniteDistributions.affinity(List.of(0), 1, 2))
+ .build();
+
+ PartitionPruningMetadata actual = extractMetadata(
+ "SELECT * FROM t1 WHERE c1 = 42 UNION SELECT * FROM t2 WHERE c1 = 99",
+ table1, table2
+ );
+
+ PartitionPruningColumns cols1 = actual.get(1);
+ assertNotNull(cols1, "No metadata for source=1");
+ assertEquals("[[0=42]]", PartitionPruningColumns.canonicalForm(cols1).toString());
+
+ PartitionPruningColumns cols2 = actual.get(2);
+ assertNotNull(cols2, "No metadata for source=2");
+ assertEquals("[[0=99]]", PartitionPruningColumns.canonicalForm(cols2).toString());
+ }
+
+ private PartitionPruningMetadata extractMetadata(String query, IgniteTable... table) throws Exception {
+ IgniteSchema schema = createSchema(table);
+
+ IgniteRel rel = physicalPlan(query, schema);
+
+ PartitionPruningMetadataExtractor extractor = new PartitionPruningMetadataExtractor();
+ return extractor.go(rel.accept(new AssignSourceIds()));
+ }
+
+ private static class AssignSourceIds extends IgniteRelShuttle {
+
+ private long sourceId = 1;
+
+ @Override
+ public IgniteRel visit(IgniteRel rel) {
+ if (rel instanceof SourceAwareIgniteRel) {
+ SourceAwareIgniteRel s = (SourceAwareIgniteRel) rel;
+ return s.clone(sourceId++);
+ } else {
+ return super.visit(rel);
+ }
+ }
+ }
+}