anuragrai16 commented on code in PR #17317: URL: https://github.com/apache/pinot/pull/17317#discussion_r2617988889
########## pinot-query-planner/src/main/java/org/apache/pinot/calcite/sql2rel/RowComparisonConvertlet.java: ########## @@ -0,0 +1,244 @@ +/** + * 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.pinot.calcite.sql2rel; + +import java.util.ArrayList; +import java.util.List; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexUtil; +import org.apache.calcite.sql.SqlCall; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlNode; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql2rel.SqlRexContext; +import org.apache.calcite.sql2rel.SqlRexConvertlet; + + +/** + * Convertlet that rewrites ROW comparisons into scalar comparisons. + * + * Transforms: + * (A, B) = (X, Y) → (A = X) AND (B = Y) + * (A, B) > (X, Y) → (A > X) OR ((A = X) AND (B > Y)) + * (A, B) >= (X, Y) → (A > X) OR ((A = X) AND (B > Y)) OR ((A = X) AND (B = Y)) + * (A, B) < (X, Y) → (A < X) OR ((A = X) AND (B < Y)) + * (A, B) <= (X, Y) → (A < X) OR ((A = X) AND (B < Y)) OR ((A = X) AND (B = Y)) + * (A, B) <> (X, Y) → NOT((A = X) AND (B = Y)) + */ +public class RowComparisonConvertlet implements SqlRexConvertlet { + + public static final RowComparisonConvertlet INSTANCE = + new RowComparisonConvertlet(); + + @Override + public RexNode convertCall(SqlRexContext cx, SqlCall call) { + if (call.getOperandList().size() != 2) { + throw new IllegalArgumentException( + "ROW comparisons must have exactly 2 operands, got: " + call.getOperandList().size()); + } + + SqlNode leftNode = call.operand(0); + SqlNode rightNode = call.operand(1); + + // Validate both operands are ROW expressions + if (!isRowExpression(leftNode)) { + throw new IllegalArgumentException( + "Left operand of ROW comparison must be a ROW expression. Got: " + getNodeKind(leftNode) + + ". Both operands must be ROW expressions for comparison."); + } + if (!isRowExpression(rightNode)) { + throw new IllegalArgumentException( + "Right operand of ROW comparison must be a ROW expression. Got: " + getNodeKind(rightNode) + + ". Both operands must be ROW expressions for comparison."); + } + + // Unwrap CAST to get the actual ROW expressions + SqlCall leftRow = unwrapCastToRow(leftNode); + SqlCall rightRow = unwrapCastToRow(rightNode); + + // Validate both ROW expressions have the same number of fields + int leftSize = leftRow.getOperandList().size(); + int rightSize = rightRow.getOperandList().size(); + if (leftSize != rightSize) { + throw new IllegalArgumentException( + String.format("ROW comparison operands must have the same number of fields. Left has %d fields, " + + "right has %d fields.", + leftSize, rightSize)); + } + if (leftSize == 0) { + throw new IllegalArgumentException("ROW expressions cannot be empty"); + } + + RexBuilder rexBuilder = cx.getRexBuilder(); + // Convert ROW fields to RexNodes + List<RexNode> leftFields = new ArrayList<>(leftSize); + for (SqlNode field : leftRow.getOperandList()) { + leftFields.add(cx.convertExpression(field)); + } + List<RexNode> rightFields = new ArrayList<>(rightSize); + for (SqlNode field : rightRow.getOperandList()) { + rightFields.add(cx.convertExpression(field)); + } + + SqlKind kind = call.getKind(); + switch (kind) { + case EQUALS: + return rewriteEquals(rexBuilder, leftFields, rightFields); + case NOT_EQUALS: + return rewriteNotEquals(rexBuilder, leftFields, rightFields); + case GREATER_THAN: + return rewriteGreaterThan(rexBuilder, leftFields, rightFields); + case GREATER_THAN_OR_EQUAL: + return rewriteGreaterThanOrEqual(rexBuilder, leftFields, rightFields); + case LESS_THAN: + return rewriteLessThan(rexBuilder, leftFields, rightFields); + case LESS_THAN_OR_EQUAL: + return rewriteLessThanOrEqual(rexBuilder, leftFields, rightFields); + default: + throw new IllegalArgumentException("Unsupported ROW comparison operator: " + kind); + } + } + + private static boolean isRowExpression(SqlNode node) { + if (!(node instanceof SqlCall)) { + return false; + } + + SqlCall call = (SqlCall) node; + if (call.getKind() == SqlKind.ROW) { + return true; + } + // ROW wrapped in CAST + if (call.getKind() == SqlKind.CAST && call.getOperandList().size() > 0) { + SqlNode castOperand = call.getOperandList().get(0); + if (castOperand instanceof SqlCall && ((SqlCall) castOperand).getKind() == SqlKind.ROW) { + return true; + } + } + return false; + } + + private static SqlCall unwrapCastToRow(SqlNode node) { + if (!(node instanceof SqlCall)) { + throw new IllegalArgumentException("Expected SqlCall, got: " + node.getClass()); + } + + SqlCall call = (SqlCall) node; + if (call.getKind() == SqlKind.ROW) { + return call; + } + // CAST(ROW(...) AS ...) - extract the ROW + if (call.getKind() == SqlKind.CAST && call.getOperandList().size() > 0) { + SqlNode operand = call.getOperandList().get(0); + if (operand instanceof SqlCall && ((SqlCall) operand).getKind() == SqlKind.ROW) { + return (SqlCall) operand; + } + } Review Comment: Hey, yes, that is expected in some cases. I noticed that when you do (colA, colB) <> (2, 3), the right hand is rewritten by Calcite to `CAST(ROW(2, 3) AS ROW(`EXPR$0` INTEGER, `EXPR$1` REAL))`. This needs explicit unwrapping and verification, so the code is written to handle both cases. ie, ROW() occuring independently and ROW() occuring as part of the cast. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
