IGNITE-28331 Fix SQL search by _key for a composite pk in Calcite engine

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/IndexWrappedKeyScan.java

@@ -0,0 +1,105 @@
+/*
+ * 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.processors.query.calcite.exec;
+
+import java.util.Map;
+import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.calcite.util.ImmutableIntList;
+import org.apache.ignite.IgniteException;
+import org.apache.ignite.binary.BinaryObject;
+import org.apache.ignite.internal.cache.query.index.sorted.IndexKeyDefinition;
+import org.apache.ignite.internal.cache.query.index.sorted.IndexKeyType;
+import org.apache.ignite.internal.cache.query.index.sorted.IndexPlainRowImpl;
+import org.apache.ignite.internal.cache.query.index.sorted.IndexRow;
+import org.apache.ignite.internal.cache.query.index.sorted.InlineIndexRowHandler;
+import org.apache.ignite.internal.cache.query.index.sorted.inline.InlineIndex;
+import org.apache.ignite.internal.cache.query.index.sorted.keys.IndexKey;
+import org.apache.ignite.internal.cache.query.index.sorted.keys.IndexKeyFactory;
+import org.apache.ignite.internal.processors.query.QueryUtils;
+import org.apache.ignite.internal.processors.query.calcite.exec.exp.RangeIterable;
+import org.apache.ignite.internal.processors.query.calcite.schema.CacheTableDescriptor;
+import org.apache.ignite.internal.processors.query.calcite.schema.ColumnDescriptor;
+import org.apache.ignite.internal.processors.query.calcite.util.Commons;
+import org.apache.ignite.internal.processors.query.calcite.util.TypeUtils;
+import org.jetbrains.annotations.Nullable;
+
+/** Extension for column {@value QueryUtils#KEY_FIELD_NAME} in case of composite primary key. */
+public class IndexWrappedKeyScan<Row> extends IndexScan<Row> {
+    /** */
+    public IndexWrappedKeyScan(
+        ExecutionContext<Row> ectx,
+        CacheTableDescriptor desc,
+        InlineIndex idx,
+        ImmutableIntList idxFieldMapping,
+        int[] parts,
+        RangeIterable<Row> ranges,
+        @Nullable ImmutableBitSet requiredColumns
+    ) {
+        super(ectx, desc, idx, idxFieldMapping, parts, ranges, requiredColumns);
+    }
+
+    /** */
+    @Override protected IndexRow row2indexRow(Row bound) {
+        if (bound == null)
+            return null;
+
+        RowHandler<Row> rowHnd = ectx.rowHandler();
+
+        Object key = rowHnd.get(QueryUtils.KEY_COL, bound);
+        assert key != null : String.format("idxName=%s, bound=%s", idx.name(), Commons.toString(rowHnd, bound));
+
+        if (key instanceof BinaryObject)
+            return binaryObject2indexRow((BinaryObject)key);
+
+        throw new IgniteException(String.format(
+            "Unsupported type for index boundary: [expected=%s, current=%s]",
+            BinaryObject.class.getName(), key.getClass().getName()
+        ));
+    }
+
+    /** */
+    private IndexRow binaryObject2indexRow(BinaryObject o) {
+        assert o.type().typeName().equals(idx.indexDefinition().typeDescriptor().keyTypeName()) : String.format(
+            "idx=%s, o=%s, oType=%s, idxKeyType=%s",
+            idx.name(), o, o.type().typeName(), idx.indexDefinition().typeDescriptor().keyTypeName()
+        );
+
+        InlineIndexRowHandler idxRowHnd = idx.segment(0).rowHandler();
+        IndexKey[] keys = new IndexKey[idx.indexDefinition().indexKeyDefinitions().size()];
+
+        int i = 0;
+        for (Map.Entry<String, IndexKeyDefinition> e : idx.indexDefinition().indexKeyDefinitions().entrySet()) {
+            String keyName = e.getKey();
+
+            ColumnDescriptor fieldDesc = desc.columnDescriptor(keyName);
+            assert fieldDesc != null : String.format("idx=%s, o=%s, keyName=%s", idx.name(), o, keyName);
+
+            Object field = o.field(keyName);
+            Object key = TypeUtils.fromInternal(ectx, field, fieldDesc.storageType());
+
+            keys[i++] = wrapIndexKey(key, e.getValue().indexKeyType());
+        }
+
+        return new IndexPlainRowImpl(keys, idxRowHnd);
+    }
+
+    /** */
+    private IndexKey wrapIndexKey(Object key, IndexKeyType keyType) {
+        return IndexKeyFactory.wrap(key, keyType, cctx.cacheObjectContext(), idx.indexDefinition().keyTypeSettings());
+    }
+}

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/exp/ExpressionFactoryImpl.java

@@ -33,7 +33,7 @@
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
-import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Primitives;
 import org.apache.calcite.DataContext;
@@ -46,6 +46,7 @@
 import org.apache.calcite.linq4j.tree.ParameterExpression;
 import org.apache.calcite.plan.RelOptUtil;
 import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
 import org.apache.calcite.rel.RelFieldCollation;
 import org.apache.calcite.rel.core.AggregateCall;
 import org.apache.calcite.rel.type.RelDataType;
@@ -81,6 +82,8 @@
 import org.apache.ignite.internal.util.GridBoundedConcurrentLinkedHashMap;
 import org.apache.ignite.internal.util.typedef.F;
 
+import static java.util.stream.Collectors.toList;
+
 /**
  * Implements rex expression into a function object. Uses JaninoRexCompiler under the hood.
  * Each expression compiles into a class and a wrapper over it is returned.
@@ -338,6 +341,14 @@ else if (o2 == null)
 
         List<RangeConditionImpl> ranges = new ArrayList<>();
 
+        if (collation.getKeys().isEmpty()) {
+            collation = RelCollations.of(IntStream.range(0, searchBounds.size())
+                .filter(i -> searchBounds.get(i) != null)
+                .mapToObj(RelFieldCollation::new)
+                .collect(toList())
+            );
+        }
+
         Comparator<Row> rowComparator = comparator(collation);
 
         expandBounds(
@@ -1003,7 +1014,7 @@ public RangeIterableImpl(List<RangeConditionImpl> ranges) {
             // should not affect ordering.
             if (!sorted) {
                 ranges = ranges.stream().filter(r -> !r.skip()).sorted(RangeConditionImpl::compareTo)
-                    .collect(Collectors.toList());
+                    .collect(toList());
 
                 List<RangeConditionImpl> ranges0 = new ArrayList<>(ranges.size());
 

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/exp/IgniteSqlFunctions.java

@@ -27,6 +27,7 @@
 import org.apache.calcite.linq4j.Linq4j;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
+import org.apache.calcite.runtime.SqlFunctions;
 import org.apache.calcite.schema.ScannableTable;
 import org.apache.calcite.schema.Schema;
 import org.apache.calcite.schema.Statistic;
@@ -293,4 +294,52 @@ private long convertToLongArg(Object val, String name) {
             return true;
         }
     }
+
+    /** SQL >=. */
+    public static boolean geAny(Object a, Object b) {
+        if (Commons.isBinaryComparable(a, b))
+            return Commons.compareBinary(a, b) >= 0;
+
+        return SqlFunctions.geAny(a, b);
+    }
+
+    /** SQL >. */
+    public static boolean gtAny(Object a, Object b) {
+        if (Commons.isBinaryComparable(a, b))
+            return Commons.compareBinary(a, b) > 0;
+
+        return SqlFunctions.gtAny(a, b);
+    }
+
+    /** SQL <=. */
+    public static boolean leAny(Object a, Object b) {
+        if (Commons.isBinaryComparable(a, b))
+            return Commons.compareBinary(a, b) <= 0;
+
+        return SqlFunctions.leAny(a, b);
+    }
+
+    /** SQL <. */
+    public static boolean ltAny(Object a, Object b) {
+        if (Commons.isBinaryComparable(a, b))
+            return Commons.compareBinary(a, b) < 0;
+
+        return SqlFunctions.ltAny(a, b);
+    }
+
+    /** SQL =. */
+    public static boolean eqAny(Object a, Object b) {
+        if (Commons.isBinaryComparable(a, b))
+            return Commons.compareBinary(a, b) == 0;
+
+        return SqlFunctions.eqAny(a, b);
+    }
+
+    /** SQL <>. */
+    public static boolean neAny(Object a, Object b) {
+        if (Commons.isBinaryComparable(a, b))
+            return Commons.compareBinary(a, b) != 0;
+
+        return SqlFunctions.neAny(a, b);
+    }
 }

modules/calcite/src/main/java/org/apache/ignite/internal/processors/query/calcite/exec/exp/RexImpTable.java

@@ -1298,7 +1298,7 @@ private Expression callBackupMethodAnyType(RexToLixTranslator translator,
             // one or both of parameter(s) is(are) ANY type
             final Expression expression0 = maybeBox(expressions.get(0));
             final Expression expression1 = maybeBox(expressions.get(1));
-            return Expressions.call(SqlFunctions.class, backupMethodNameForAnyType,
+            return Expressions.call(IgniteSqlFunctions.class, backupMethodNameForAnyType,
                 expression0, expression1);
         }
 

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

@@ -38,6 +38,7 @@
 import org.apache.ignite.internal.processors.query.calcite.exec.ExecutionContext;
 import org.apache.ignite.internal.processors.query.calcite.exec.RowHandler;
 import org.apache.ignite.internal.processors.query.calcite.type.IgniteTypeFactory;
+import org.apache.ignite.internal.processors.query.calcite.util.Commons;
 import org.apache.ignite.internal.util.typedef.F;
 
 import static org.apache.calcite.sql.type.SqlTypeName.ANY;
@@ -234,7 +235,11 @@
                 return () -> new ComparableMinMax<Row, UUID>(call, hnd, true,
                     tf -> tf.createTypeWithNullability(tf.createSqlType(SqlTypeName.UUID), true));
             case ANY:
-                throw new UnsupportedOperationException("MIN() is not supported for type '" + call.type + "'.");
+                return () -> new ComparableMinMax<>(call, hnd, true,
+                    tf -> tf.createTypeWithNullability(tf.createSqlType(ANY), true));
+            case OTHER:
+                return () -> new ComparableMinMax<>(call, hnd, true,
+                    tf -> tf.createTypeWithNullability(tf.createSqlType(SqlTypeName.OTHER), true));
             case BIGINT:
             default:
                 return () -> new LongMinMax<>(call, hnd, true);
@@ -263,7 +268,11 @@
                 return () -> new ComparableMinMax<Row, UUID>(call, hnd, false,
                     tf -> tf.createTypeWithNullability(tf.createSqlType(SqlTypeName.UUID), true));
             case ANY:
-                throw new UnsupportedOperationException("MAX() is not supported for type '" + call.type + "'.");
+                return () -> new ComparableMinMax<>(call, hnd, false,
+                    tf -> tf.createTypeWithNullability(tf.createSqlType(ANY), true));
+            case OTHER:
+                return () -> new ComparableMinMax<>(call, hnd, false,
+                    tf -> tf.createTypeWithNullability(tf.createSqlType(SqlTypeName.OTHER), true));
             case BIGINT:
             default:
                 return () -> new LongMinMax<>(call, hnd, false);
@@ -1116,7 +1125,7 @@
     }
 
     /** */
-    private static class ComparableMinMax<Row, T extends Comparable<T>> extends AbstractAccumulator<Row> {
+    private static class ComparableMinMax<Row, T> extends AbstractAccumulator<Row> {
         /** */
         private final boolean min;
 
@@ -1149,8 +1158,8 @@
                 return;
 
             val = empty ? in : min ?
-                (val.compareTo(in) < 0 ? val : in) :
-                (val.compareTo(in) < 0 ? in : val);
+                (compare(val, in) < 0 ? val : in) :
+                (compare(val, in) < 0 ? in : val);
 
             empty = false;
         }
@@ -1163,8 +1172,8 @@
                 return;
 
             val = empty ? other0.val : min ?
-                (val.compareTo(other0.val) < 0 ? val : other0.val) :
-                (val.compareTo(other0.val) < 0 ? other0.val : val);
+                (compare(val, other0.val) < 0 ? val : other0.val) :
+                (compare(val, other0.val) < 0 ? other0.val : val);
 
             empty = false;
         }
@@ -1183,6 +1192,22 @@
         @Override public RelDataType returnType(IgniteTypeFactory typeFactory) {
             return typeSupplier.apply(typeFactory);
         }
+
+        /** */
+        @SuppressWarnings({"rawtypes", "unchecked"})
+        private int compare(Object a, Object b) {
+            if (Commons.isBinaryComparable(a, b))
+                return Commons.compareBinary(a, b);
+
+            if (a.getClass() != b.getClass()) {
+                throw new UnsupportedOperationException(String.format(
+                    "%s() is not supported for different value types: [type0=%s, type1=%s]",
+                    min ? "MIN" : "MAX", a.getClass().getName(), b.getClass().getName()
+                ));
+            }
+
+            return ((Comparable)a).compareTo(b);
+        }
     }
 
     /** */