Skip to content

Commit

Permalink
[CALCITE-6776] Multiple expanded IS NOT DISTINCT FROM cannot be col…
Browse files Browse the repository at this point in the history
…lapsed back
  • Loading branch information
linorosa authored and mihaibudiu committed Jan 14, 2025
1 parent 5f767dd commit 2f096ff
Show file tree
Hide file tree
Showing 2 changed files with 275 additions and 17 deletions.
37 changes: 20 additions & 17 deletions core/src/main/java/org/apache/calcite/plan/RelOptUtil.java
Original file line number Diff line number Diff line change
Expand Up @@ -1748,32 +1748,35 @@ private static void splitJoinCondition(
* and
* {@link #splitJoinCondition(List, List, RexNode, List, List, List, List)}.
*
* <p>If the given expr <code>call</code> is an expanded version of
* <p>If the given expr <code>rexCall</code> contains an expanded version of
* {@code IS NOT DISTINCT FROM} function call, collapses it and return a
* {@code IS NOT DISTINCT FROM} function call.
*
* <p>For example: {@code t1.key IS NOT DISTINCT FROM t2.key}
* can rewritten in expanded form as
* can be rewritten in expanded form as
* {@code t1.key = t2.key OR (t1.key IS NULL AND t2.key IS NULL)}.
*
* @param call Function expression to try collapsing
* @param rexCall Function expression to try collapsing
* @param rexBuilder {@link RexBuilder} instance to create new {@link RexCall} instances.
* @return If the given function is an expanded IS NOT DISTINCT FROM function call,
* return a IS NOT DISTINCT FROM function call. Otherwise return the input
* function call as it is.
* @return A function where all IS NOT DISTINCT FROM are collapsed.
*/
public static RexCall collapseExpandedIsNotDistinctFromExpr(final RexCall call,
public static RexCall collapseExpandedIsNotDistinctFromExpr(final RexCall rexCall,
final RexBuilder rexBuilder) {
switch (call.getKind()) {
case OR:
return doCollapseExpandedIsNotDistinctFromOrExpr(call, rexBuilder);

case CASE:
return doCollapseExpandedIsNotDistinctFromCaseExpr(call, rexBuilder);

default:
return call;
}
final RexShuttle shuttle = new RexShuttle() {
@Override public RexNode visitCall(RexCall call) {
RexCall recursivelyExpanded = (RexCall) super.visitCall(call);

switch (recursivelyExpanded.getKind()) {
case OR:
return doCollapseExpandedIsNotDistinctFromOrExpr(recursivelyExpanded, rexBuilder);
case CASE:
return doCollapseExpandedIsNotDistinctFromCaseExpr(recursivelyExpanded, rexBuilder);
default:
return recursivelyExpanded;
}
}
};
return (RexCall) rexCall.accept(shuttle);
}

private static RexCall doCollapseExpandedIsNotDistinctFromOrExpr(final RexCall call,
Expand Down
255 changes: 255 additions & 0 deletions core/src/test/java/org/apache/calcite/plan/RelOptUtilTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import org.apache.calcite.rel.type.RelDataTypeField;
import org.apache.calcite.rel.type.RelDataTypeSystem;
import org.apache.calcite.rex.RexBuilder;
import org.apache.calcite.rex.RexCall;
import org.apache.calcite.rex.RexInputRef;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.schema.SchemaPlus;
Expand Down Expand Up @@ -420,6 +421,260 @@ private void splitJoinConditionHelper(RexNode joinCond, List<Integer> expLeftKey
assertThat(actRightKeys, is(expRightKeys));
}

/**
* Test that {@link RelOptUtil#collapseExpandedIsNotDistinctFromExpr(RexCall, RexBuilder)}
* collapses an expanded version of IS NOT DISTINCT using OR.
*/
@Test void testCollapseExpandedIsNotDistinctFromUsingOr() {
final RexBuilder rexBuilder = relBuilder.getRexBuilder();

final RexNode leftEmpNo =
RexInputRef.of(empScan.getRowType().getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);
final RexNode rightEmpNo =
RexInputRef.of(empRow.getFieldCount() + deptRow.getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);

// OR(AND(IS NULL($0), IS NULL($7)), IS TRUE(=($0, $7)))
RexNode expanded = relBuilder.isNotDistinctFrom(leftEmpNo, rightEmpNo);

// IS NOT DISTINCT FROM($0, $7)
RexNode collapsed =
RelOptUtil.collapseExpandedIsNotDistinctFromExpr((RexCall) expanded, rexBuilder);

RexNode expected =
rexBuilder.makeCall(SqlStdOperatorTable.IS_NOT_DISTINCT_FROM, leftEmpNo, rightEmpNo);

assertThat(collapsed, is(expected));
}

/**
* Test that {@link RelOptUtil#collapseExpandedIsNotDistinctFromExpr(RexCall, RexBuilder)}
* collapses an expanded version of IS NOT DISTINCT using CASE.
*/
@Test void testCollapseExpandedIsNotDistinctFromUsingCase() {
final RexBuilder rexBuilder = relBuilder.getRexBuilder();

final RexNode leftEmpNo =
RexInputRef.of(empScan.getRowType().getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);
final RexNode rightEmpNo =
RexInputRef.of(empRow.getFieldCount() + deptRow.getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);

// CASE(IS NULL($0), IS NULL($7), IS NULL($7), IS NULL($0), =($0, $7))
RexNode expanded =
relBuilder.call(SqlStdOperatorTable.CASE, relBuilder.isNull(leftEmpNo),
relBuilder.isNull(rightEmpNo),
relBuilder.isNull(rightEmpNo),
relBuilder.isNull(leftEmpNo),
relBuilder.equals(leftEmpNo, rightEmpNo));

// IS NOT DISTINCT FROM($0, $7)
RexNode collapsed =
RelOptUtil.collapseExpandedIsNotDistinctFromExpr((RexCall) expanded, rexBuilder);

RexNode expected =
rexBuilder.makeCall(SqlStdOperatorTable.IS_NOT_DISTINCT_FROM, leftEmpNo, rightEmpNo);

assertThat(collapsed, is(expected));
}

/**
* Test that {@link RelOptUtil#collapseExpandedIsNotDistinctFromExpr(RexCall, RexBuilder)}
* collapses an expression with expanded versions of IS NOT DISTINCT using OR and CASE.
*/
@Test void testCollapseExpandedIsNotDistinctFromUsingOrAndCase() {
final RexBuilder rexBuilder = relBuilder.getRexBuilder();

final RexNode leftEmpNo =
RexInputRef.of(empScan.getRowType().getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);
final RexNode rightEmpNo =
RexInputRef.of(empRow.getFieldCount() + deptRow.getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);

final RexNode leftDeptNo =
RexInputRef.of(empScan.getRowType().getFieldNames().indexOf("DEPTNO"),
empDeptJoinRelFields);
final RexNode rightDeptNo =
RexInputRef.of(empRow.getFieldCount() + deptRow.getFieldNames().indexOf("DEPTNO"),
empDeptJoinRelFields);

// An IS NOT DISTINCT FROM expanded in "CASE" shape
// CASE(IS NULL($0), IS NULL($7), IS NULL($7), IS NULL($0), =($0, $7))
RexNode expandedCase =
relBuilder.call(SqlStdOperatorTable.CASE, relBuilder.isNull(leftEmpNo),
relBuilder.isNull(rightEmpNo),
relBuilder.isNull(rightEmpNo),
relBuilder.isNull(leftEmpNo),
relBuilder.equals(leftEmpNo, rightEmpNo));

// An IS NOT DISTINCT FROM expanded in "OR" shape
// OR(AND(IS NULL($7), IS NULL($8)), =($7, $8))
RexNode expandedOr =
relBuilder.call(
SqlStdOperatorTable.OR, relBuilder.call(SqlStdOperatorTable.AND,
relBuilder.isNull(leftDeptNo),
relBuilder.isNull(rightDeptNo)),
relBuilder.call(SqlStdOperatorTable.EQUALS, leftDeptNo, rightDeptNo));

// AND(
// OR(AND(IS NULL($7), IS NULL($8)), =($7, $8)),
// CASE(IS NULL($0), IS NULL($7), IS NULL($7), IS NULL($0), =($0, $7))
// )
RexNode expanded = relBuilder.and(expandedOr, expandedCase);

// AND(IS NOT DISTINCT FROM($7, $8), IS NOT DISTINCT FROM($0, $7))
RexNode collapsed =
RelOptUtil.collapseExpandedIsNotDistinctFromExpr((RexCall) expanded, rexBuilder);

RexNode expected =
rexBuilder.makeCall(
// Expected is nullable because `expandedCase` is nullable
relBuilder.getTypeFactory().createTypeWithNullability(expanded.getType(), true),
SqlStdOperatorTable.AND,
ImmutableList.of(
rexBuilder.makeCall(
SqlStdOperatorTable.IS_NOT_DISTINCT_FROM,
leftEmpNo,
rightEmpNo),
rexBuilder.makeCall(
SqlStdOperatorTable.IS_NOT_DISTINCT_FROM,
leftDeptNo,
rightDeptNo)));

assertThat(collapsed, is(expected));
}

/**
* Test that {@link RelOptUtil#collapseExpandedIsNotDistinctFromExpr(RexCall, RexBuilder)}
* recursively collapses expanded versions of IS NOT DISTINCT.
*/
@Test void testCollapseExpandedIsNotDistinctFromRecursively() {
final RexBuilder rexBuilder = relBuilder.getRexBuilder();

final RexNode leftEmpNo =
RexInputRef.of(empScan.getRowType().getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);
final RexNode rightEmpNo =
RexInputRef.of(empRow.getFieldCount() + deptRow.getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);

// OR(
// AND(
// IS NULL(OR(AND(IS NULL($0), IS NULL($7)), IS TRUE(=($0, $7)))),
// IS NULL(OR(AND(IS NULL($0), IS NULL($7)), IS TRUE(=($0, $7))))),
// IS TRUE(=(
// OR(AND(IS NULL($0), IS NULL($7)), IS TRUE(=($0, $7))),
// OR(AND(IS NULL($0), IS NULL($7)), IS TRUE(=($0, $7)))
// )
// )
// )
RexNode expanded =
relBuilder.isNotDistinctFrom(
relBuilder.isNotDistinctFrom(leftEmpNo, rightEmpNo),
relBuilder.isNotDistinctFrom(leftEmpNo, rightEmpNo));

// IS NOT DISTINCT FROM(IS NOT DISTINCT FROM($0, $7), IS NOT DISTINCT FROM($0, $7))
RexNode collapsed =
RelOptUtil.collapseExpandedIsNotDistinctFromExpr((RexCall) expanded, rexBuilder);

RexNode expected =
rexBuilder.makeCall(SqlStdOperatorTable.IS_NOT_DISTINCT_FROM,
rexBuilder.makeCall(SqlStdOperatorTable.IS_NOT_DISTINCT_FROM,
leftEmpNo,
rightEmpNo),

rexBuilder.makeCall(SqlStdOperatorTable.IS_NOT_DISTINCT_FROM,
leftEmpNo,
rightEmpNo));

assertThat(collapsed, is(expected));
}

/**
* Test that {@link RelOptUtil#collapseExpandedIsNotDistinctFromExpr(RexCall, RexBuilder)}
* will collapse IS NOT DISTINCT FROM nested within RexNodes.
*/
@Test void testCollapseExpandedIsNotDistinctFromInsideRexNode() {
final RexBuilder rexBuilder = relBuilder.getRexBuilder();

final RexNode leftEmpNo =
RexInputRef.of(empScan.getRowType().getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);
final RexNode rightEmpNo =
RexInputRef.of(empRow.getFieldCount() + deptRow.getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);

RexNode expandedIsNotDistinctFrom = relBuilder.isNotDistinctFrom(leftEmpNo, rightEmpNo);
// NULLIF(
// NOT(OR(AND(IS NULL($0), IS NULL($7)), IS TRUE(=($0, $7)))),
// IS NOT NULL(OR(AND(IS NULL($0), IS NULL($7)), IS TRUE(=($0, $7))))
// )
RexNode expanded =
relBuilder.call(SqlStdOperatorTable.NULLIF,
relBuilder.not(expandedIsNotDistinctFrom),
relBuilder.isNotNull(expandedIsNotDistinctFrom));

// NULLIF(NOT(IS NOT DISTINCT FROM($0, $7)), IS NOT NULL(IS NOT DISTINCT FROM($0, $7)))
RexNode collapsed =
RelOptUtil.collapseExpandedIsNotDistinctFromExpr((RexCall) expanded, rexBuilder);

RexNode collapsedIsNotDistinctFrom =
rexBuilder.makeCall(SqlStdOperatorTable.IS_NOT_DISTINCT_FROM, leftEmpNo, rightEmpNo);
RexNode expected =
rexBuilder.makeCall(
SqlStdOperatorTable.NULLIF,
relBuilder.not(collapsedIsNotDistinctFrom),
relBuilder.isNotNull(collapsedIsNotDistinctFrom));

assertThat(collapsed, is(expected));
}

/**
* Test that {@link RelOptUtil#collapseExpandedIsNotDistinctFromExpr(RexCall, RexBuilder)}
* can handle collapsing IS NOT DISTINCT FROM composed of other RexNodes.
*/
@Test void testCollapseExpandedIsNotDistinctFromOnContainingRexNodes() {
final RexBuilder rexBuilder = relBuilder.getRexBuilder();

final RexNode leftEmpNo =
RexInputRef.of(empScan.getRowType().getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);
final RexNode rightEmpNo =
RexInputRef.of(empRow.getFieldCount() + deptRow.getFieldNames().indexOf("EMPNO"),
empDeptJoinRelFields);

final RexNode leftDeptNo =
RexInputRef.of(empScan.getRowType().getFieldNames().indexOf("DEPTNO"),
empDeptJoinRelFields);
final RexNode rightDeptNo =
RexInputRef.of(empRow.getFieldCount() + deptRow.getFieldNames().indexOf("DEPTNO"),
empDeptJoinRelFields);

// OR(
// AND(IS NULL(NULLIF($0, $7)), IS NULL(NULLIF($7, $8))),
// IS TRUE(=(NULLIF($0, $7), NULLIF($7, $8)))
// )
RexNode expanded =
relBuilder.isNotDistinctFrom(
relBuilder.call(SqlStdOperatorTable.NULLIF, leftEmpNo, rightEmpNo),
relBuilder.call(SqlStdOperatorTable.NULLIF, leftDeptNo, rightDeptNo));

// IS NOT DISTINCT FROM(NULLIF($0, $7), NULLIF($7, $8))
RexNode collapsed =
RelOptUtil.collapseExpandedIsNotDistinctFromExpr((RexCall) expanded, rexBuilder);

RexNode expected =
rexBuilder.makeCall(
SqlStdOperatorTable.IS_NOT_DISTINCT_FROM,
rexBuilder.makeCall(SqlStdOperatorTable.NULLIF, leftEmpNo, rightEmpNo),
rexBuilder.makeCall(SqlStdOperatorTable.NULLIF, leftDeptNo, rightDeptNo));

assertThat(collapsed, is(expected));
}

/**
* Tests {@link RelOptUtil#pushDownJoinConditions(org.apache.calcite.rel.core.Join, RelBuilder)}
* where the join condition contains a complex expression.
Expand Down

0 comments on commit 2f096ff

Please sign in to comment.