This is an automated email from the ASF dual-hosted git repository.
korlov 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 87eb584f765 IGNITE-25559 Sql. Implement collocation metadata builder.
(#6030)
87eb584f765 is described below
commit 87eb584f765e38c309fce7a16fc0e077eed68cf5
Author: Max Zhuravkov <[email protected]>
AuthorDate: Wed Jun 18 10:01:00 2025 +0300
IGNITE-25559 Sql. Implement collocation metadata builder. (#6030)
---
.../internal/sql/engine/prepare/DdlPlan.java | 8 +
.../internal/sql/engine/prepare/ExplainPlan.java | 8 +
.../sql/engine/prepare/KeyValueGetPlan.java | 20 +-
.../sql/engine/prepare/KeyValueModifyPlan.java | 14 +-
.../internal/sql/engine/prepare/KillPlan.java | 8 +
.../internal/sql/engine/prepare/MultiStepPlan.java | 7 +
.../sql/engine/prepare/PrepareServiceImpl.java | 19 +-
.../internal/sql/engine/prepare/QueryPlan.java | 7 +
.../sql/engine/prepare/SelectCountPlan.java | 7 +
.../PartitionAwarenessMetadata.java | 87 ++++++++
.../PartitionAwarenessMetadataExtractor.java | 121 +++++++++++
.../PartitionAwarenessMetadataTest.java | 235 +++++++++++++++++++++
12 files changed, 535 insertions(+), 6 deletions(-)
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/DdlPlan.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/DdlPlan.java
index 10f886bc8c3..2ca3da83560 100644
---
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/DdlPlan.java
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/DdlPlan.java
@@ -23,9 +23,11 @@ import org.apache.ignite.internal.catalog.CatalogCommand;
import org.apache.ignite.internal.sql.ColumnMetadataImpl;
import org.apache.ignite.internal.sql.ResultSetMetadataImpl;
import org.apache.ignite.internal.sql.engine.SqlQueryType;
+import
org.apache.ignite.internal.sql.engine.prepare.partitionawareness.PartitionAwarenessMetadata;
import org.apache.ignite.sql.ColumnMetadata;
import org.apache.ignite.sql.ColumnType;
import org.apache.ignite.sql.ResultSetMetadata;
+import org.jetbrains.annotations.Nullable;
/**
* DdlPlan.
@@ -75,6 +77,12 @@ public class DdlPlan implements QueryPlan {
return EMPTY_PARAMETERS;
}
+ /** {@inheritDoc} */
+ @Override
+ public @Nullable PartitionAwarenessMetadata partitionAwarenessMetadata() {
+ return null;
+ }
+
/** {@inheritDoc} */
@Override
public String toString() {
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/ExplainPlan.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/ExplainPlan.java
index b1e8d2bf206..d9cc4da9681 100644
---
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/ExplainPlan.java
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/ExplainPlan.java
@@ -21,10 +21,12 @@ import java.util.List;
import org.apache.ignite.internal.sql.ColumnMetadataImpl;
import org.apache.ignite.internal.sql.ResultSetMetadataImpl;
import org.apache.ignite.internal.sql.engine.SqlQueryType;
+import
org.apache.ignite.internal.sql.engine.prepare.partitionawareness.PartitionAwarenessMetadata;
import org.apache.ignite.internal.sql.engine.sql.IgniteSqlExplainMode;
import org.apache.ignite.sql.ColumnMetadata;
import org.apache.ignite.sql.ColumnType;
import org.apache.ignite.sql.ResultSetMetadata;
+import org.jetbrains.annotations.Nullable;
/**
* Query explain plan.
@@ -73,6 +75,12 @@ public class ExplainPlan implements QueryPlan {
return plan.parameterMetadata();
}
+ /** {@inheritDoc} */
+ @Override
+ public @Nullable PartitionAwarenessMetadata partitionAwarenessMetadata() {
+ return null;
+ }
+
public ExplainablePlan plan() {
return plan;
}
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KeyValueGetPlan.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KeyValueGetPlan.java
index 7425235dc20..0078584356d 100644
---
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KeyValueGetPlan.java
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KeyValueGetPlan.java
@@ -47,6 +47,7 @@ import
org.apache.ignite.internal.sql.engine.exec.exp.SqlPredicate;
import org.apache.ignite.internal.sql.engine.exec.exp.SqlProjection;
import org.apache.ignite.internal.sql.engine.exec.exp.SqlRowProvider;
import org.apache.ignite.internal.sql.engine.exec.row.RowSchema;
+import
org.apache.ignite.internal.sql.engine.prepare.partitionawareness.PartitionAwarenessMetadata;
import org.apache.ignite.internal.sql.engine.rel.IgniteKeyValueGet;
import org.apache.ignite.internal.sql.engine.rel.IgniteRel;
import org.apache.ignite.internal.sql.engine.rel.explain.ExplainUtils;
@@ -70,16 +71,25 @@ public class KeyValueGetPlan implements ExplainablePlan,
ExecutablePlan {
private final IgniteKeyValueGet lookupNode;
private final ResultSetMetadata meta;
private final ParameterMetadata parameterMetadata;
+ @Nullable
+ private final PartitionAwarenessMetadata partitionAwarenessMetadata;
private volatile Performable<?> operation;
- KeyValueGetPlan(PlanId id, int catalogVersion, IgniteKeyValueGet
lookupNode, ResultSetMetadata meta,
- ParameterMetadata parameterMetadata) {
+ KeyValueGetPlan(
+ PlanId id,
+ int catalogVersion,
+ IgniteKeyValueGet lookupNode,
+ ResultSetMetadata meta,
+ ParameterMetadata parameterMetadata,
+ @Nullable PartitionAwarenessMetadata partitionAwarenessMetadata
+ ) {
this.id = id;
this.catalogVersion = catalogVersion;
this.lookupNode = lookupNode;
this.meta = meta;
this.parameterMetadata = parameterMetadata;
+ this.partitionAwarenessMetadata = partitionAwarenessMetadata;
}
/** {@inheritDoc} */
@@ -106,6 +116,12 @@ public class KeyValueGetPlan implements ExplainablePlan,
ExecutablePlan {
return parameterMetadata;
}
+ /** {@inheritDoc} */
+ @Override
+ public @Nullable PartitionAwarenessMetadata partitionAwarenessMetadata() {
+ return partitionAwarenessMetadata;
+ }
+
/** Returns a table in question. */
private IgniteTable table() {
IgniteTable table = lookupNode.getTable().unwrap(IgniteTable.class);
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KeyValueModifyPlan.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KeyValueModifyPlan.java
index eae69f78d0a..0b3bee355dc 100644
---
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KeyValueModifyPlan.java
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KeyValueModifyPlan.java
@@ -33,6 +33,7 @@ import
org.apache.ignite.internal.sql.engine.exec.ExecutableTableRegistry;
import org.apache.ignite.internal.sql.engine.exec.ExecutionContext;
import org.apache.ignite.internal.sql.engine.exec.UpdatableTable;
import org.apache.ignite.internal.sql.engine.exec.exp.SqlRowProvider;
+import
org.apache.ignite.internal.sql.engine.prepare.partitionawareness.PartitionAwarenessMetadata;
import org.apache.ignite.internal.sql.engine.rel.IgniteKeyValueModify;
import org.apache.ignite.internal.sql.engine.rel.IgniteRel;
import org.apache.ignite.internal.sql.engine.rel.explain.ExplainUtils;
@@ -42,6 +43,7 @@ import org.apache.ignite.internal.sql.engine.util.Commons;
import org.apache.ignite.internal.sql.engine.util.IteratorToDataCursorAdapter;
import org.apache.ignite.internal.tx.InternalTransaction;
import org.apache.ignite.sql.ResultSetMetadata;
+import org.jetbrains.annotations.Nullable;
/**
* Plan representing simple modify operation that can be executed by Key-Value
API.
@@ -52,6 +54,8 @@ public class KeyValueModifyPlan implements ExplainablePlan,
ExecutablePlan {
private final IgniteKeyValueModify modifyNode;
private final ResultSetMetadata meta;
private final ParameterMetadata parameterMetadata;
+ @Nullable
+ private final PartitionAwarenessMetadata partitionAwarenessMetadata;
private volatile InsertExecution<?> operation;
@@ -60,13 +64,15 @@ public class KeyValueModifyPlan implements ExplainablePlan,
ExecutablePlan {
int catalogVersion,
IgniteKeyValueModify modifyNode,
ResultSetMetadata meta,
- ParameterMetadata parameterMetadata
+ ParameterMetadata parameterMetadata,
+ @Nullable PartitionAwarenessMetadata partitionAwarenessMetadata
) {
this.id = id;
this.catalogVersion = catalogVersion;
this.modifyNode = modifyNode;
this.meta = meta;
this.parameterMetadata = parameterMetadata;
+ this.partitionAwarenessMetadata = partitionAwarenessMetadata;
}
/** {@inheritDoc} */
@@ -93,6 +99,12 @@ public class KeyValueModifyPlan implements ExplainablePlan,
ExecutablePlan {
return parameterMetadata;
}
+ /** {@inheritDoc} */
+ @Override
+ public @Nullable PartitionAwarenessMetadata partitionAwarenessMetadata() {
+ return partitionAwarenessMetadata;
+ }
+
/** Returns a table in question. */
private IgniteTable table() {
IgniteTable table = modifyNode.getTable().unwrap(IgniteTable.class);
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KillPlan.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KillPlan.java
index 5a8511f3996..96fce97d26d 100644
---
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KillPlan.java
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/KillPlan.java
@@ -23,9 +23,11 @@ import org.apache.ignite.internal.sql.ColumnMetadataImpl;
import org.apache.ignite.internal.sql.ResultSetMetadataImpl;
import org.apache.ignite.internal.sql.engine.SqlQueryType;
import org.apache.ignite.internal.sql.engine.exec.kill.KillCommand;
+import
org.apache.ignite.internal.sql.engine.prepare.partitionawareness.PartitionAwarenessMetadata;
import org.apache.ignite.sql.ColumnMetadata;
import org.apache.ignite.sql.ColumnType;
import org.apache.ignite.sql.ResultSetMetadata;
+import org.jetbrains.annotations.Nullable;
/**
* KILL command execution plan.
@@ -73,6 +75,12 @@ public class KillPlan implements QueryPlan {
return EMPTY_PARAMETERS;
}
+ /** {@inheritDoc} */
+ @Override
+ public @Nullable PartitionAwarenessMetadata partitionAwarenessMetadata() {
+ return null;
+ }
+
/** {@inheritDoc} */
@Override
public String toString() {
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/MultiStepPlan.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/MultiStepPlan.java
index 2adf84080e6..20ace2cce4d 100644
---
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/MultiStepPlan.java
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/MultiStepPlan.java
@@ -18,6 +18,7 @@
package org.apache.ignite.internal.sql.engine.prepare;
import org.apache.ignite.internal.sql.engine.SqlQueryType;
+import
org.apache.ignite.internal.sql.engine.prepare.partitionawareness.PartitionAwarenessMetadata;
import org.apache.ignite.internal.sql.engine.rel.IgniteRel;
import org.apache.ignite.internal.sql.engine.rel.explain.ExplainUtils;
import org.apache.ignite.internal.sql.engine.util.Cloner;
@@ -81,6 +82,12 @@ public class MultiStepPlan implements ExplainablePlan {
return parameterMetadata;
}
+ /** {@inheritDoc} */
+ @Override
+ public @Nullable PartitionAwarenessMetadata partitionAwarenessMetadata() {
+ return null;
+ }
+
/** {@inheritDoc} */
@Override
public SqlQueryType type() {
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/PrepareServiceImpl.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/PrepareServiceImpl.java
index 4a435be21ad..98651b549be 100644
---
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/PrepareServiceImpl.java
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/PrepareServiceImpl.java
@@ -60,6 +60,8 @@ import
org.apache.ignite.internal.sql.engine.SqlOperationContext;
import org.apache.ignite.internal.sql.engine.SqlQueryType;
import org.apache.ignite.internal.sql.engine.exec.kill.KillCommand;
import
org.apache.ignite.internal.sql.engine.prepare.ddl.DdlSqlToCommandConverter;
+import
org.apache.ignite.internal.sql.engine.prepare.partitionawareness.PartitionAwarenessMetadata;
+import
org.apache.ignite.internal.sql.engine.prepare.partitionawareness.PartitionAwarenessMetadataExtractor;
import org.apache.ignite.internal.sql.engine.rel.IgniteKeyValueGet;
import org.apache.ignite.internal.sql.engine.rel.IgniteKeyValueModify;
import org.apache.ignite.internal.sql.engine.rel.IgniteRel;
@@ -434,8 +436,13 @@ public class PrepareServiceImpl implements PrepareService {
int catalogVersion = ctx.catalogVersion();
if (optimizedRel instanceof IgniteKeyValueGet) {
+ IgniteKeyValueGet kvGet = (IgniteKeyValueGet) optimizedRel;
+ PartitionAwarenessMetadata partitionAwarenessMetadata =
+
PartitionAwarenessMetadataExtractor.getMetadata(kvGet);
+
return new KeyValueGetPlan(
- nextPlanId(), catalogVersion, (IgniteKeyValueGet)
optimizedRel, resultSetMetadata, parameterMetadata
+ nextPlanId(), catalogVersion, kvGet,
resultSetMetadata,
+ parameterMetadata, partitionAwarenessMetadata
);
}
@@ -502,7 +509,8 @@ public class PrepareServiceImpl implements PrepareService {
ExplainablePlan plan;
if (optimizedRel instanceof IgniteKeyValueModify) {
plan = new KeyValueModifyPlan(
- nextPlanId(), ctx.catalogVersion(), (IgniteKeyValueModify)
optimizedRel, DML_METADATA, parameterMetadata
+ nextPlanId(), ctx.catalogVersion(), (IgniteKeyValueModify)
optimizedRel, DML_METADATA,
+ parameterMetadata, null
);
} else {
plan = new MultiStepPlan(
@@ -557,8 +565,13 @@ public class PrepareServiceImpl implements PrepareService {
ExplainablePlan plan;
if (optimizedRel instanceof IgniteKeyValueModify) {
+ IgniteKeyValueModify kvModify = (IgniteKeyValueModify)
optimizedRel;
+ PartitionAwarenessMetadata partitionAwarenessMetadata =
+
PartitionAwarenessMetadataExtractor.getMetadata(kvModify);
+
plan = new KeyValueModifyPlan(
- nextPlanId(), catalogVersion,
(IgniteKeyValueModify) optimizedRel, DML_METADATA, parameterMetadata
+ nextPlanId(), catalogVersion,
(IgniteKeyValueModify) optimizedRel, DML_METADATA,
+ parameterMetadata, partitionAwarenessMetadata
);
} else {
plan = new MultiStepPlan(
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/QueryPlan.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/QueryPlan.java
index 4fb88e9c5a6..68138c112c4 100644
---
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/QueryPlan.java
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/QueryPlan.java
@@ -17,7 +17,9 @@
package org.apache.ignite.internal.sql.engine.prepare;
+import javax.annotation.Nullable;
import org.apache.ignite.internal.sql.engine.SqlQueryType;
+import
org.apache.ignite.internal.sql.engine.prepare.partitionawareness.PartitionAwarenessMetadata;
import org.apache.ignite.sql.ResultSetMetadata;
/**
@@ -44,4 +46,9 @@ public interface QueryPlan {
* Returns parameters metadata.
*/
ParameterMetadata parameterMetadata();
+
+ /**
+ * Returns partition-awareness metadata or {@code null} if it not present.
+ */
+ @Nullable PartitionAwarenessMetadata partitionAwarenessMetadata();
}
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/SelectCountPlan.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/SelectCountPlan.java
index 8a19b1055f1..77e549c13b5 100644
---
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/SelectCountPlan.java
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/SelectCountPlan.java
@@ -41,6 +41,7 @@ import
org.apache.ignite.internal.sql.engine.exec.ExecutionContext;
import org.apache.ignite.internal.sql.engine.exec.RowHandler;
import org.apache.ignite.internal.sql.engine.exec.exp.SqlProjection;
import org.apache.ignite.internal.sql.engine.exec.row.RowSchema;
+import
org.apache.ignite.internal.sql.engine.prepare.partitionawareness.PartitionAwarenessMetadata;
import org.apache.ignite.internal.sql.engine.rel.IgniteRel;
import org.apache.ignite.internal.sql.engine.rel.IgniteSelectCount;
import org.apache.ignite.internal.sql.engine.rel.explain.ExplainUtils;
@@ -52,6 +53,7 @@ import org.apache.ignite.internal.sql.engine.util.TypeUtils;
import org.apache.ignite.internal.tx.InternalTransaction;
import org.apache.ignite.internal.type.NativeTypes;
import org.apache.ignite.sql.ResultSetMetadata;
+import org.jetbrains.annotations.Nullable;
/**
* Plan representing a COUNT(*) query.
@@ -143,6 +145,11 @@ public class SelectCountPlan implements ExplainablePlan,
ExecutablePlan {
return parameterMetadata;
}
+ @Override
+ public @Nullable PartitionAwarenessMetadata partitionAwarenessMetadata() {
+ return null;
+ }
+
private <RowT> Function<Long, Iterator<InternalSqlRow>>
createResultProjection(ExecutionContext<RowT> ctx) {
RelDataType getCountType = new
RelDataTypeFactory.Builder(ctx.getTypeFactory())
.add("ROWCOUNT", SqlTypeName.BIGINT)
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/partitionawareness/PartitionAwarenessMetadata.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/partitionawareness/PartitionAwarenessMetadata.java
new file mode 100644
index 00000000000..ccc3fa81ad4
--- /dev/null
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/partitionawareness/PartitionAwarenessMetadata.java
@@ -0,0 +1,87 @@
+/*
+ * 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.partitionawareness;
+
+import java.util.Arrays;
+import org.apache.ignite.internal.tostring.S;
+
+/**
+ * Partition awareness metadata. The index array stores a mapping between
colocation key indexes and dynamic parameters and hash array.
+ * The mapping uses negated 1-based indexes for hashes and 0-based non-negated
indexes for dynamic parameters
+ *
+ * @see PartitionAwarenessMetadataExtractor
+ */
+public final class PartitionAwarenessMetadata {
+
+ private final int tableId;
+
+ private final int[] indexes;
+
+ private final int[] hash;
+
+ /**
+ * Constructor.
+ *
+ * @param tableId Table Id.
+ * @param indexes Mapping between positions in colocation key and dynamic
parameters.
+ * @param hash Array of computed hashes.
+ */
+ public PartitionAwarenessMetadata(int tableId, int[] indexes, int[] hash) {
+ this.tableId = tableId;
+ this.indexes = indexes;
+ this.hash = hash;
+ }
+
+ /** Return table id. */
+ public int tableId() {
+ return tableId;
+ }
+
+ /** Returns the number of colocation key columns. */
+ public int size() {
+ return indexes.length;
+ }
+
+ /**
+ * Returns a mapping between positions in colocation key columns and
dynamic parameter indices.
+ * If a colocation key column has a value, returns a negated 1-based index
into the hash array.
+ *
+ * @return Mapping.
+ */
+ public int[] indexes() {
+ return indexes;
+ }
+
+ /**
+ * Returns an array of precomputed hashes.
+ *
+ * @return Hash array.
+ */
+ public int[] hash() {
+ return hash;
+ }
+
+ /** {@inheritDoc} */
+ @Override
+ public String toString() {
+ return S.toString(PartitionAwarenessMetadata.class, this,
+ "indexes", Arrays.toString(indexes),
+ "hash", Arrays.toString(hash)
+ );
+ }
+}
diff --git
a/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/partitionawareness/PartitionAwarenessMetadataExtractor.java
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/partitionawareness/PartitionAwarenessMetadataExtractor.java
new file mode 100644
index 00000000000..49513a89633
--- /dev/null
+++
b/modules/sql-engine/src/main/java/org/apache/ignite/internal/sql/engine/prepare/partitionawareness/PartitionAwarenessMetadataExtractor.java
@@ -0,0 +1,121 @@
+/*
+ * 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.partitionawareness;
+
+import java.util.List;
+import org.apache.calcite.plan.RelOptTable;
+import org.apache.calcite.rex.RexDynamicParam;
+import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.util.ImmutableIntList;
+import org.apache.ignite.internal.sql.engine.rel.IgniteKeyValueGet;
+import org.apache.ignite.internal.sql.engine.rel.IgniteKeyValueModify;
+import org.apache.ignite.internal.sql.engine.schema.IgniteTable;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * Extracts partition awareness metadata from physical plans. Examples:
+ *
+ * <pre>
+ * SELECT * FROM t WHERE pk=?
+ * colocation key: [pk]
+ * =>
+ * indexes: [0], hash: []
+ *
+ * SELECT * FROM t WHERE pk1=? and pk2=?
+ * colocation key: [pk1, pk2]
+ * =>
+ * indexes: [0, 1], hash: []
+ *
+ * SELECT * FROM t WHERE pk1=? and pk2=V1 and pk3=?
+ * colocation key: [pk1, pk2, pk3]
+ * =>
+ * indexes: [0, -1, 1], hash: [hash(V1)]
+ * </pre>
+ *
+ * @see PartitionAwarenessMetadata
+ */
+public class PartitionAwarenessMetadataExtractor {
+
+ /**
+ * Extracts partition awareness metadata from the given IgniteKeyValueGet
plan.
+ *
+ * @param kv IgniteKeyValueGet Plan.
+ * @return Metadata.
+ */
+ @Nullable
+ public static PartitionAwarenessMetadata getMetadata(IgniteKeyValueGet kv)
{
+ RelOptTable optTable = kv.getTable();
+ assert optTable != null;
+
+ List<RexNode> expressions = kv.keyExpressions();
+
+ return buildMetadata(optTable, false, expressions);
+ }
+
+ /**
+ * Extracts partition awareness metadata from the given
IgniteKeyValueModify plan.
+ *
+ * @param kv IgniteKeyValueModify Plan.
+ * @return Metadata.
+ */
+ @Nullable
+ public static PartitionAwarenessMetadata getMetadata(IgniteKeyValueModify
kv) {
+ RelOptTable optTable = kv.getTable();
+ assert optTable != null;
+
+ List<RexNode> expressions = kv.expressions();
+
+ return buildMetadata(optTable, true, expressions);
+ }
+
+ private static @Nullable PartitionAwarenessMetadata buildMetadata(
+ RelOptTable optTable,
+ boolean fullRow,
+ List<RexNode> expressions
+ ) {
+ IgniteTable igniteTable = optTable.unwrap(IgniteTable.class);
+ assert igniteTable != null;
+
+ ImmutableIntList colocationKeys = igniteTable.distribution().getKeys();
+
+ // colocation key index to dynamic param index
+ int[] indexes = new int[colocationKeys.size()];
+ int[] hash = new int[0];
+
+ for (int i = 0; i < colocationKeys.size(); i++) {
+ int colIdx = colocationKeys.get(i);
+ RexNode expr;
+
+ if (fullRow) {
+ expr = expressions.get(colIdx);
+ } else {
+ int keyIdx = igniteTable.keyColumns().indexOf(colIdx);
+ expr = expressions.get(keyIdx);
+ }
+
+ if (expr instanceof RexDynamicParam) {
+ RexDynamicParam dynamicParam = (RexDynamicParam) expr;
+ indexes[i] = dynamicParam.getIndex();
+ } else {
+ return null;
+ }
+ }
+
+ return new PartitionAwarenessMetadata(igniteTable.id(), indexes, hash);
+ }
+}
diff --git
a/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/prepare/partitionawareness/PartitionAwarenessMetadataTest.java
b/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/prepare/partitionawareness/PartitionAwarenessMetadataTest.java
new file mode 100644
index 00000000000..14da955a818
--- /dev/null
+++
b/modules/sql-engine/src/test/java/org/apache/ignite/internal/sql/engine/prepare/partitionawareness/PartitionAwarenessMetadataTest.java
@@ -0,0 +1,235 @@
+/*
+ * 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.partitionawareness;
+
+import static org.apache.ignite.internal.testframework.IgniteTestUtils.await;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
+import org.apache.ignite.internal.catalog.CatalogCommand;
+import org.apache.ignite.internal.catalog.CatalogManager;
+import org.apache.ignite.internal.catalog.commands.DropTableCommand;
+import org.apache.ignite.internal.catalog.descriptors.CatalogTableDescriptor;
+import org.apache.ignite.internal.sql.SqlCommon;
+import org.apache.ignite.internal.sql.engine.framework.TestBuilders;
+import org.apache.ignite.internal.sql.engine.framework.TestCluster;
+import org.apache.ignite.internal.sql.engine.framework.TestNode;
+import org.apache.ignite.internal.sql.engine.prepare.QueryPlan;
+import org.apache.ignite.internal.sql.engine.util.Commons;
+import org.apache.ignite.internal.testframework.BaseIgniteAbstractTest;
+import org.jetbrains.annotations.Nullable;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
+
+/**
+ * Tests forr {@link PartitionAwarenessMetadata} in query plans.
+ */
+public class PartitionAwarenessMetadataTest extends BaseIgniteAbstractTest {
+
+ private static final String NODE_NAME = "N1";
+
+ private static final TestCluster CLUSTER = TestBuilders.cluster()
+ .nodes(NODE_NAME)
+ .build();
+
+ private final TestNode node = CLUSTER.node(NODE_NAME);
+
+ @BeforeAll
+ static void start() {
+ CLUSTER.start();
+ }
+
+ @AfterAll
+ static void stop() throws Exception {
+ CLUSTER.stop();
+ }
+
+ @AfterEach
+ void clearCatalog() {
+ Commons.resetFastQueryOptimizationFlag();
+
+ int version = CLUSTER.catalogManager().latestCatalogVersion();
+
+ List<CatalogCommand> commands = new ArrayList<>();
+ for (CatalogTableDescriptor table :
CLUSTER.catalogManager().catalog(version).tables()) {
+ commands.add(
+ DropTableCommand.builder()
+ .schemaName(SqlCommon.DEFAULT_SCHEMA_NAME)
+ .tableName(table.name())
+ .build()
+ );
+ }
+
+ await(CLUSTER.catalogManager().execute(commands));
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {
+ "SELECT 1",
+ "KILL QUERY '1'",
+ "KILL TRANSACTION '1'",
+ "CREATE TABLE x (a INT, b INT, PRIMARY KEY (a))",
+ "SELECT * FROM table(system_range(1, 10)) WHERE x = 1",
+ "SELECT count(*) FROM t WHERE c1=?",
+ "UPDATE t SET c2=1 WHERE c1=?",
+ "DELETE FROM t WHERE c1=?",
+ })
+ public void noMetadata(String query) {
+ node.initSchema("CREATE TABLE t (c1 INT PRIMARY KEY, c2 INT)");
+
+ QueryPlan plan = node.prepare(query);
+ PartitionAwarenessMetadata metadata =
plan.partitionAwarenessMetadata();
+ assertNull(metadata);
+ }
+
+ @ParameterizedTest
+ @MethodSource("simpleKeyMetadata")
+ public void simpleKey(String query, PartitionAwarenessMetadata expected) {
+ node.initSchema("CREATE TABLE t (c1 INT PRIMARY KEY, c2 INT)");
+
+ QueryPlan plan = node.prepare(query);
+ PartitionAwarenessMetadata metadata =
plan.partitionAwarenessMetadata();
+
+ expectMedata(expected, metadata);
+ }
+
+ private static Stream<Arguments> simpleKeyMetadata() {
+ return Stream.of(
+ // KV GET
+ Arguments.of("SELECT * FROM t WHERE c1=?", dynamicParams(0)),
+ Arguments.of("SELECT * FROM t WHERE c2=? and c1=?",
dynamicParams(1)),
+ Arguments.of("SELECT * FROM t WHERE c1=1", null),
+ Arguments.of("SELECT * FROM t WHERE c1=1+1", null),
+ // the first condition goes into key lookup other into a
post-lookup filter.
+ Arguments.of("SELECT * FROM t WHERE c1=? and c1=?",
dynamicParams(0)),
+ Arguments.of("SELECT * FROM t WHERE c2=? and c1=? and c1=?",
dynamicParams(1)),
+
+ // KV PUT
+ Arguments.of("INSERT INTO t VALUES(?, ?)", dynamicParams(0)),
+ Arguments.of("INSERT INTO t VALUES(1, ?)", null),
+ Arguments.of("INSERT INTO t VALUES(1+1, ?)", null),
+ Arguments.of("INSERT INTO t(c2, c1) VALUES(?, ?)",
dynamicParams(1)),
+ Arguments.of("INSERT INTO t(c2, c1) VALUES(1, ?)",
dynamicParams(0)),
+ Arguments.of("INSERT INTO t(c2, c1) VALUES(?, 1)", null)
+ );
+ }
+
+ @ParameterizedTest
+ @MethodSource("shortKeyMetadata")
+ public void shortKey(String query, PartitionAwarenessMetadata expected) {
+ node.initSchema("CREATE TABLE t (c1 INT, c2 INT, c3 INT, PRIMARY KEY
(c3, c2)) COLOCATE BY (c3)");
+
+ QueryPlan plan = node.prepare(query);
+ PartitionAwarenessMetadata metadata =
plan.partitionAwarenessMetadata();
+
+ expectMedata(expected, metadata);
+ }
+
+ private static Stream<Arguments> shortKeyMetadata() {
+ return Stream.of(
+ // KV GET
+ Arguments.of("SELECT * FROM t WHERE c3=? and c2=?",
dynamicParams(0)),
+ Arguments.of("SELECT * FROM t WHERE c2=? and c3=?",
dynamicParams(1)),
+ Arguments.of("SELECT * FROM t WHERE c3=? and c2=?",
dynamicParams(0)),
+ Arguments.of("SELECT * FROM t WHERE c1=? and c2=? and c3=?",
dynamicParams(2)),
+ Arguments.of("SELECT * FROM t WHERE c3=? and c1=? and c2=?",
dynamicParams(0)),
+
+ Arguments.of("SELECT * FROM t WHERE c3=3", null),
+ Arguments.of("SELECT * FROM t WHERE c2=? and c3=3", null),
+ Arguments.of("SELECT * FROM t WHERE c1=? and c2=? and c3=3",
null),
+
+ // KV PUT
+ Arguments.of("INSERT INTO t VALUES (?, ?, ?)",
dynamicParams(2)),
+ Arguments.of("INSERT INTO t (c1, c2, c3) VALUES (?, ?, ?)",
dynamicParams(2)),
+ Arguments.of("INSERT INTO t (c3, c1, c2) VALUES (?, ?, ?)",
dynamicParams(0)),
+
+ Arguments.of("INSERT INTO t (c1, c2, c3) VALUES (?, ?, 3)",
null),
+ Arguments.of("INSERT INTO t (c1, c3, c2) VALUES (?, 3, ?)",
null),
+ Arguments.of("INSERT INTO t (c3, c1, c2) VALUES (3, ?, ?)",
null)
+ );
+ }
+
+ @ParameterizedTest
+ @MethodSource("compoundKeyMetadata")
+ public void compoundKey(String query, PartitionAwarenessMetadata expected)
{
+ node.initSchema("CREATE TABLE t (c1 INT, c2 INT, c3 INT, c4 INT,
PRIMARY KEY(c1, c2, c3)) COLOCATE BY (c3, c1, c2)");
+
+ QueryPlan plan = node.prepare(query);
+ PartitionAwarenessMetadata metadata =
plan.partitionAwarenessMetadata();
+
+ expectMedata(expected, metadata);
+ }
+
+ private static Stream<Arguments> compoundKeyMetadata() {
+ return Stream.of(
+ // KV GET
+ Arguments.of("SELECT * FROM t WHERE c1=? and c2=? and c3=?",
dynamicParams(2, 0, 1)),
+ Arguments.of("SELECT * FROM t WHERE c3=? and c1=? and c2=?",
dynamicParams(0, 1, 2)),
+ Arguments.of("SELECT * FROM t WHERE c3=? and c2=? and c1=?",
dynamicParams(0, 2, 1)),
+ Arguments.of("SELECT * FROM t WHERE c4=? and c1=? and c2=? and
1=? and c3=?", dynamicParams(4, 1, 2)),
+ // duplicate condition goes to a post lookup filter.
+ Arguments.of("SELECT * FROM t WHERE c1=? and c2=? and c3=? and
c2=?", dynamicParams(2, 0, 1)),
+
+ Arguments.of("SELECT * FROM t WHERE c1=1 and c2=? and c3=?",
null),
+ Arguments.of("SELECT * FROM t WHERE c1=? and c2=2 and c3=?",
null),
+ Arguments.of("SELECT * FROM t WHERE c1=? and c2=? and c3=3",
null),
+ Arguments.of("SELECT * FROM t WHERE c1=1 and c2=2 and c3=3",
null),
+
+ // KV PUT
+ Arguments.of("INSERT INTO t VALUES (?, ?, ?, ?)",
dynamicParams(2, 0, 1)),
+ Arguments.of("INSERT INTO t (c3, c2, c4, c1) VALUES (?, ?, ?,
?)", dynamicParams(0, 3, 1)),
+ Arguments.of("INSERT INTO t (c3, c2, c4, c1) VALUES (?, ?, 1,
?)", dynamicParams(0, 2, 1))
+ );
+ }
+
+ private static PartitionAwarenessMetadata dynamicParams(int...
dynamicParams) {
+ return new PartitionAwarenessMetadata(1, dynamicParams, new int[0]);
+ }
+
+ private static void expectMedata(PartitionAwarenessMetadata expected,
@Nullable PartitionAwarenessMetadata actual) {
+ if (expected == null) {
+ assertNull(actual, "Metadata should not be present");
+ } else {
+ assertNotNull(actual, "Metadata not found");
+
+ CatalogManager catalogManager = CLUSTER.catalogManager();
+ int v = catalogManager.latestCatalogVersion();
+
+ CatalogTableDescriptor table =
catalogManager.catalog(v).table("PUBLIC", "T");
+ assertNotNull(table, "table");
+
+ assertEquals(table.id(), actual.tableId(), "metadata tableId");
+ assertEquals(
+
Arrays.stream(expected.indexes()).boxed().collect(Collectors.toList()),
+
Arrays.stream(actual.indexes()).boxed().collect(Collectors.toList()),
+ "indexes"
+ );
+ }
+ }
+}
