Skip to content

Commit

Permalink
IGNITE-14964 SQL Calcite: Optimize IN statement with literals
Browse files Browse the repository at this point in the history
  • Loading branch information
alex-plekhanov committed Feb 13, 2025
1 parent a26aa8e commit b4280a8
Show file tree
Hide file tree
Showing 9 changed files with 175 additions and 7 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,8 @@ private PartitionNode processCondition(
return new PartitionParameterNode(cacheId, (RexDynamicParam)right, fldType);
}
case SEARCH:
// Intentionally use of RexUtil.expandSearch (not RexUtils.expandSearchNullable), since here we
// expand operator not for bytecode generation and expect output with OR/AND operators.
RexNode condition0 = RexUtil.expandSearch(Commons.emptyCluster().getRexBuilder(), null, condition);

return processCondition(condition0, types, keys, requiredCols, cacheId);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,6 @@
import org.apache.calcite.rex.RexRangeRef;
import org.apache.calcite.rex.RexSubQuery;
import org.apache.calcite.rex.RexTableInputRef;
import org.apache.calcite.rex.RexUtil;
import org.apache.calcite.rex.RexVisitor;
import org.apache.calcite.sql.SqlIntervalQualifier;
import org.apache.calcite.sql.SqlOperator;
Expand All @@ -72,6 +71,7 @@
import org.apache.calcite.util.Pair;
import org.apache.ignite.internal.processors.query.calcite.type.IgniteTypeFactory;
import org.apache.ignite.internal.processors.query.calcite.util.IgniteMethod;
import org.apache.ignite.internal.processors.query.calcite.util.RexUtils;

import static org.apache.calcite.sql.fun.SqlStdOperatorTable.CASE;
import static org.apache.calcite.sql.fun.SqlStdOperatorTable.SEARCH;
Expand Down Expand Up @@ -1034,7 +1034,7 @@ private ConstantExpression getTypedNullLiteral(RexLiteral literal) {
return implementCaseWhen(call);

if (operator == SEARCH)
return RexUtil.expandSearch(builder, program, call).accept(this);
return RexUtils.expandSearchNullable(builder, program, call).accept(this);

final RexImpTable.RexCallImplementor implementor =
RexImpTable.INSTANCE.get(operator);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,6 @@
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.rex.RexOver;
import org.apache.calcite.rex.RexSlot;
import org.apache.calcite.rex.RexUtil;
import org.apache.calcite.rex.RexVariable;
import org.apache.calcite.rex.RexWindow;
import org.apache.calcite.rex.RexWindowBound;
Expand Down Expand Up @@ -113,6 +112,7 @@
import org.apache.ignite.internal.processors.query.calcite.type.IgniteCustomType;
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.processors.query.calcite.util.RexUtils;
import org.apache.ignite.internal.util.typedef.F;
import org.apache.ignite.internal.util.typedef.internal.U;

Expand Down Expand Up @@ -808,8 +808,8 @@ private Object toJson(CorrelationId node) {

/** */
private Object toJson(RexNode node) {
// removes calls to SEARCH and the included Sarg and converts them to comparisons
node = RexUtil.expandSearch(Commons.emptyCluster().getRexBuilder(), null, node);
// Removes calls to SEARCH and the included Sarg and converts them to comparisons.
node = RexUtils.expandSearchNullableRecursive(Commons.emptyCluster().getRexBuilder(), null, node);

Map<String, Object> map;
switch (node.getKind()) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,8 @@ else if (predKind == SqlKind.NOT) {
continue;
}
else if (predKind == SqlKind.SEARCH) {
// Intentionally use of RexUtil.expandSearch (not RexUtils.expandSearchNullable), since here we
// expand operator not for bytecode generation and expect output with OR/AND operators.
sel *= getTablePredicateBasedSelectivity(rel, mq,
RexUtil.expandSearch(rel.getCluster().getRexBuilder(), null, pred));

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ public ImmutableBitSet requiredColumns() {
protected RelWriter explainTerms0(RelWriter pw) {
if (condition != null) {
pw.item("filters", pw.nest() ? condition :
RexUtil.expandSearch(getCluster().getRexBuilder(), null, condition));
RexUtils.expandSearchNullableRecursive(getCluster().getRexBuilder(), null, condition));
}

return pw
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@
import org.apache.calcite.rex.RexLiteral;
import org.apache.calcite.rex.RexLocalRef;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.rex.RexProgram;
import org.apache.calcite.rex.RexShuttle;
import org.apache.calcite.rex.RexSimplify;
import org.apache.calcite.rex.RexSlot;
Expand Down Expand Up @@ -866,6 +867,44 @@ public static Set<CorrelationId> extractCorrelationIds(List<RexNode> nodes) {
return cors;
}

/**
* Expand SEARCH/SARG operator with preceding NULLs check for nullable input operator.
* This allows to eliminate NULL checks for every argument of SARG and reduces generated bytecode.
*/
public static RexNode expandSearchNullable(RexBuilder rexBuilder, @Nullable RexProgram program, RexCall call) {
assert call.getOperator() == SqlStdOperatorTable.SEARCH : "Unexpected operator: " + call.getOperator();

RexNode op0 = call.getOperands().get(0);
RexNode op1 = call.getOperands().get(1);

if (!op0.getType().isNullable())
return RexUtil.expandSearch(rexBuilder, program, call);

Sarg<?> arg = ((RexLiteral)op1).getValueAs(Sarg.class);

RexNode nullAs = arg.nullAs == RexUnknownAs.TRUE || arg.nullAs == RexUnknownAs.FALSE
? rexBuilder.makeLiteral(arg.nullAs.toBoolean())
: rexBuilder.makeNullLiteral(rexBuilder.getTypeFactory().createSqlType(SqlTypeName.BOOLEAN));

RexNode expandedSearch = RexUtil.expandSearch(rexBuilder, program,
rexBuilder.makeCall(call.getOperator(), rexBuilder.makeNotNull(op0), op1));

return rexBuilder.makeCall(SqlStdOperatorTable.CASE,
rexBuilder.makeCall(SqlStdOperatorTable.IS_NULL, op0),
nullAs,
expandedSearch
);
}

/**
* Traverse {@code node} and expand all SEARCH/SARG operators (with preceding NULLs check).
*
* @see #expandSearchNullable
*/
public static RexNode expandSearchNullableRecursive(RexBuilder rexBuilder, @Nullable RexProgram program, RexNode node) {
return node.accept(new SearchExpandingShuttle(rexBuilder, program));
}

/** */
private static Boolean containsRef(RexNode node) {
RexVisitor<Void> v = new RexVisitorImpl<Void>(true) {
Expand Down Expand Up @@ -927,4 +966,27 @@ private static class InputRefReplacer extends RexShuttle {
return new RexLocalRef(inputRef.getIndex(), inputRef.getType());
}
}

/** */
private static class SearchExpandingShuttle extends RexShuttle {
/** */
private final RexBuilder rexBuilder;

/** */
private final @Nullable RexProgram program;

/** */
public SearchExpandingShuttle(RexBuilder rexBuilder, @Nullable RexProgram program) {
this.rexBuilder = rexBuilder;
this.program = program;
}

/** {@inheritDoc} */
@Override public RexNode visitCall(RexCall call) {
if (call.getOperator() == SqlStdOperatorTable.SEARCH)
return expandSearchNullable(rexBuilder, program, call);

return super.visitCall(call);
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
/*
* 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.integration;

import java.util.stream.Collectors;
import java.util.stream.IntStream;
import org.junit.Test;

/**
* Test scalar IN operator.
*/
public class ScalarInIntegrationTest extends AbstractBasicIntegrationTest {
/** {@inheritDoc} */
@Override protected void beforeTest() throws Exception {
super.beforeTest();

sql("CREATE TABLE t(i INT, s VARCHAR)");
sql("INSERT INTO t(i, s) VALUES (1, '1'), (3, '3'), (null, null)");
}

/** */
@Test
public void testInWithNull() {
assertQuery("SELECT i FROM t WHERE i IN (0, 1, 2) OR i IS NULL")
.returns(1).returns(NULL_RESULT).check();
assertQuery("SELECT i FROM t WHERE i IN (0, 1, 2, NULL)")
.returns(1).check();
assertQuery("SELECT i FROM t WHERE i IN (0, 1, 2)")
.returns(1).check();

assertQuery("SELECT i, i IN (0, 1, 2) OR i IS NULL FROM t")
.returns(1, true).returns(3, false).returns(null, true).check();
assertQuery("SELECT i, i IN (0, 1, 2, NULL) FROM t")
.returns(1, true).returns(3, null).returns(null, null).check();
assertQuery("SELECT i, i IN (0, 1, 2) FROM t")
.returns(1, true).returns(3, false).returns(null, null).check();

assertQuery("SELECT s, s IN ('0', '1', '2') OR s IS NULL FROM t")
.returns("1", true).returns("3", false).returns(null, true).check();
assertQuery("SELECT s, s IN ('0', '1', '2', NULL) FROM t")
.returns("1", true).returns("3", null).returns(null, null).check();
assertQuery("SELECT s, s IN ('0', '1', '2') FROM t")
.returns("1", true).returns("3", false).returns(null, null).check();

assertQuery("SELECT i FROM t WHERE i NOT IN (0, 1, 2) OR i IS NULL")
.returns(3).returns(NULL_RESULT).check();
assertQuery("SELECT i FROM t WHERE i NOT IN (0, 1, 2, NULL)")
.resultSize(0).check();
assertQuery("SELECT i FROM t WHERE i NOT IN (0, 1, 2)")
.returns(3).check();

assertQuery("SELECT i, i NOT IN (0, 1, 2) OR i IS NULL FROM t")
.returns(1, false).returns(3, true).returns(null, true).check();
assertQuery("SELECT i, i NOT IN (0, 1, 2, NULL) FROM t")
.returns(1, false).returns(3, null).returns(null, null).check();
assertQuery("SELECT i, i NOT IN (0, 1, 2) FROM t")
.returns(1, false).returns(3, true).returns(null, null).check();

assertQuery("SELECT s, s NOT IN ('0', '1', '2') OR s IS NULL FROM t")
.returns("1", false).returns("3", true).returns(null, true).check();
assertQuery("SELECT s, s NOT IN ('0', '1', '2', NULL) FROM t")
.returns("1", false).returns("3", null).returns(null, null).check();
assertQuery("SELECT s, s NOT IN ('0', '1', '2') FROM t")
.returns("1", false).returns("3", true).returns(null, null).check();
}

/** */
@Test
public void testLargeIn() {
String in = IntStream.range(2, 1000).mapToObj(Integer::toString).collect(Collectors.joining(", "));

assertQuery("SELECT i, i IN (" + in + ") FROM t")
.returns(1, false).returns(3, true).returns(null, null).check();

assertQuery("SELECT i, i NOT IN (" + in + ") FROM t")
.returns(1, true).returns(3, false).returns(null, null).check();

in = IntStream.range(2, 1000).mapToObj(i -> "'" + i + "'").collect(Collectors.joining(", "));

assertQuery("SELECT s, s IN (" + in + ") FROM t")
.returns("1", false).returns("3", true).returns(null, null).check();

assertQuery("SELECT s, s NOT IN (" + in + ") FROM t")
.returns("1", true).returns("3", false).returns(null, null).check();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@
import org.apache.ignite.internal.processors.query.calcite.integration.QueryMetadataIntegrationTest;
import org.apache.ignite.internal.processors.query.calcite.integration.QueryWithPartitionsIntegrationTest;
import org.apache.ignite.internal.processors.query.calcite.integration.RunningQueriesIntegrationTest;
import org.apache.ignite.internal.processors.query.calcite.integration.ScalarInIntegrationTest;
import org.apache.ignite.internal.processors.query.calcite.integration.ServerStatisticsIntegrationTest;
import org.apache.ignite.internal.processors.query.calcite.integration.SetOpIntegrationTest;
import org.apache.ignite.internal.processors.query.calcite.integration.SortAggregateIntegrationTest;
Expand Down Expand Up @@ -153,6 +154,7 @@
SessionContextSqlFunctionTest.class,
SqlPlanHistoryIntegrationTest.class,
QueryBlockingTaskExecutorIntegrationTest.class,
ScalarInIntegrationTest.class,
})
public class IntegrationTestSuite {
}
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
# name: test/sql/function/generic/test_large_in.test_slow
# description: Test large IN statement with varchar
# group: [generic]
# Ignored: https://issues.apache.org/jira/browse/IGNITE-14964

statement ok
CREATE TABLE strings(s VARCHAR)
Expand Down

0 comments on commit b4280a8

Please sign in to comment.