diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java index 1b828a499986..338f94859de6 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java +++ b/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 { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java index 432ea445a517..d6f270d6f567 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java +++ b/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); diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java index 01ba0c1d8dc7..4458d4ed23aa 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java +++ b/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); } diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/parser/CalciteRexExpressionParserTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/parser/CalciteRexExpressionParserTest.java new file mode 100644 index 000000000000..c604a8dce598 --- /dev/null +++ b/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)); + } +} diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RexExpressionUtilsTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RexExpressionUtilsTest.java index 44c6227508c7..a0883b349a2a 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RexExpressionUtilsTest.java +++ b/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)