[UUID 6/8] UUID multi-stage engine (planner + runtime)

pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java

@@ -138,6 +138,16 @@ public void testDistinctOnUuidColumn(boolean useMultiStageQueryEngine)
     }
   }
 
+  @Test(dataProvider = "useV2QueryEngine")
+  public void testMultiStageUuidLiteralPredicate(boolean useMultiStageQueryEngine)
+      throws Exception {
+    setUseMultiStageQueryEngine(useMultiStageQueryEngine);
+    JsonNode rows = query(String.format(
+        "SELECT COUNT(*) FROM %1$s WHERE %2$s = CAST('%3$s' AS UUID)",
+        getTableName(), UUID_RAW_SV_COLUMN, UUID_0));
+    assertCounts(rows.get(0), 2L);
+  }
+
   @Test
   public void testDistinctCountOnUuidColumns()
       throws Exception {

pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java

@@ -144,7 +144,7 @@ public static Literal toLiteral(RexExpression.Literal literal) {
     ColumnDataType dataType = literal.getDataType();
     if (dataType == ColumnDataType.BOOLEAN) {
       value = BooleanUtils.isTrueInternalValue(value);
-    } else if (dataType == ColumnDataType.BYTES) {
+    } else if (dataType == ColumnDataType.BYTES || dataType == ColumnDataType.UUID) {
       value = ((ByteArray) value).getBytes();
     }
     return RequestUtils.getLiteral(value);

pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java

@@ -26,6 +26,7 @@
 import java.util.Calendar;
 import java.util.List;
 import java.util.Set;
+import java.util.UUID;
 import javax.annotation.Nullable;
 import org.apache.calcite.avatica.util.ByteString;
 import org.apache.calcite.plan.RelOptCluster;
@@ -53,6 +54,7 @@
 import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
 import org.apache.pinot.spi.utils.BooleanUtils;
 import org.apache.pinot.spi.utils.ByteArray;
+import org.apache.pinot.spi.utils.UuidUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -149,6 +151,9 @@ public static RexLiteral toRexLiteral(RelBuilder builder, RexExpression.Literal
         ByteString byteString = new ByteString(bytes);
         return rexBuilder.makeBinaryLiteral(byteString);
       }
+      case UUID:
+        assert value != null;
+        return rexBuilder.makeUuidLiteral(UuidUtils.toUUID((ByteArray) value));
       default:
         throw new IllegalStateException("Unsupported ColumnDataType: " + literal.getDataType());
     }
@@ -264,6 +269,9 @@ private static RexExpression.Literal fromRexLiteralValue(ColumnDataType dataType
       case BYTES:
         value = new ByteArray(((ByteString) value).getBytes());
         break;
+      case UUID:
+        value = new ByteArray(UuidUtils.toBytes((UUID) value));
+        break;
       default:
         throw new IllegalStateException("Unsupported ColumnDataType: " + dataType);
     }

pinot-query-planner/src/test/java/org/apache/pinot/query/parser/CalciteRexExpressionParserTest.java

@@ -0,0 +1,46 @@
+/**
+ * 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.query.parser;
+
+import org.apache.pinot.common.request.Literal;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
+import org.apache.pinot.query.planner.logical.RexExpression;
+import org.apache.pinot.spi.utils.ByteArray;
+import org.apache.pinot.spi.utils.UuidUtils;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+/// Tests UUID literal conversion from multi-stage Rex expressions to single-stage request expressions.
+public class CalciteRexExpressionParserTest {
+  private static final String UUID_VALUE = "550e8400-e29b-41d4-a716-446655440000";
+
+  @Test
+  public void testUuidLiteralUsesBinaryValue() {
+    RexExpression.Literal uuidLiteral =
+        new RexExpression.Literal(ColumnDataType.UUID, new ByteArray(UuidUtils.toBytes(UUID_VALUE)));
+
+    Literal literal = CalciteRexExpressionParser.toLiteral(uuidLiteral);
+
+    assertTrue(literal.isSetBinaryValue());
+    assertEquals(literal.getBinaryValue(), UuidUtils.toBytes(UUID_VALUE));
+  }
+}

pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RexExpressionUtilsTest.java

@@ -21,6 +21,7 @@
 import com.google.common.collect.ImmutableRangeSet;
 import com.google.common.collect.Range;
 import java.math.BigDecimal;
+import java.util.UUID;
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCall;
@@ -31,9 +32,14 @@
 import org.apache.calcite.sql.SqlKind;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.type.SqlTypeName;
+import org.apache.calcite.tools.Frameworks;
+import org.apache.calcite.tools.RelBuilder;
 import org.apache.calcite.util.NlsString;
 import org.apache.calcite.util.Sarg;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
 import org.apache.pinot.query.type.TypeFactory;
+import org.apache.pinot.spi.utils.ByteArray;
+import org.apache.pinot.spi.utils.UuidUtils;
 import org.testng.Assert;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
@@ -50,6 +56,20 @@ public void setup() {
     _rexBuilder = new RexBuilder(_typeFactory);
   }
 
+  @Test
+  public void testUuidLiteralRoundTrip() {
+    UUID uuid = UUID.fromString("550e8400-e29b-41d4-a716-446655440000");
+    RelBuilder relBuilder = RelBuilder.create(Frameworks.newConfigBuilder().build());
+
+    RexExpression.Literal literal = RexExpressionUtils.fromRexLiteral(_rexBuilder.makeUuidLiteral(uuid));
+    Assert.assertEquals(literal.getDataType(), ColumnDataType.UUID);
+    Assert.assertEquals(literal.getValue(), new ByteArray(UuidUtils.toBytes(uuid)));
+
+    RexLiteral roundTrip = RexExpressionUtils.toRexLiteral(relBuilder, literal);
+    Assert.assertEquals(roundTrip.getTypeName(), SqlTypeName.UUID);
+    Assert.assertEquals(roundTrip.getValue(), uuid);
+  }
+
   @Test
   public void testHandleSearchNullLiteralInWithNullAsUnknown() {
     // Test: NULL IN (1, 2, 3) (when nullAs = UNKNOWN)