Complete UUID v1 coverage and fix CI regressions

Author: xiangfu0

pinot-common/src/main/java/org/apache/pinot/common/request/context/RequestContextUtils.java

@@ -41,12 +41,16 @@
 import org.apache.pinot.common.utils.RegexpPatternConverterUtils;
 import org.apache.pinot.common.utils.request.RequestUtils;
 import org.apache.pinot.segment.spi.AggregationFunctionType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
 import org.apache.pinot.spi.exception.BadQueryRequestException;
 import org.apache.pinot.sql.FilterKind;
 import org.apache.pinot.sql.parsers.CalciteSqlParser;
 
 
 public class RequestContextUtils {
+  private static final String UNSUPPORTED_RHS_MESSAGE =
+      "Pinot does not support column or function on the right-hand side of the predicate";
+
   private RequestContextUtils() {
   }
 
@@ -217,23 +221,23 @@ private static FilterContext getFilterInner(Function thriftFunction) {
       case GREATER_THAN:
         return FilterContext.forPredicate(
             new RangePredicate(getExpression(operands.get(0)), false, getStringValue(operands.get(1)), false,
-                RangePredicate.UNBOUNDED, new LiteralContext(operands.get(1).getLiteral()).getType()));
+                RangePredicate.UNBOUNDED, getLiteralType(operands.get(1))));
       case GREATER_THAN_OR_EQUAL:
         return FilterContext.forPredicate(
             new RangePredicate(getExpression(operands.get(0)), true, getStringValue(operands.get(1)), false,
-                RangePredicate.UNBOUNDED, new LiteralContext(operands.get(1).getLiteral()).getType()));
+                RangePredicate.UNBOUNDED, getLiteralType(operands.get(1))));
       case LESS_THAN:
         return FilterContext.forPredicate(
             new RangePredicate(getExpression(operands.get(0)), false, RangePredicate.UNBOUNDED, false,
-                getStringValue(operands.get(1)), new LiteralContext(operands.get(1).getLiteral()).getType()));
+                getStringValue(operands.get(1)), getLiteralType(operands.get(1))));
       case LESS_THAN_OR_EQUAL:
         return FilterContext.forPredicate(
             new RangePredicate(getExpression(operands.get(0)), false, RangePredicate.UNBOUNDED, true,
-                getStringValue(operands.get(1)), new LiteralContext(operands.get(1).getLiteral()).getType()));
+                getStringValue(operands.get(1)), getLiteralType(operands.get(1))));
       case BETWEEN:
         return FilterContext.forPredicate(
             new RangePredicate(getExpression(operands.get(0)), true, getStringValue(operands.get(1)), true,
-                getStringValue(operands.get(2)), new LiteralContext(operands.get(1).getLiteral()).getType()));
+                getStringValue(operands.get(2)), getLiteralType(operands.get(1))));
       case RANGE:
         return FilterContext.forPredicate(
             new RangePredicate(getExpression(operands.get(0)), getStringValue(operands.get(1))));
@@ -278,12 +282,7 @@ private static FilterContext getFilterInner(Function thriftFunction) {
   }
 
   public static String getStringValue(Expression thriftExpression) {
-    Literal literal = thriftExpression.getLiteral();
-    if (literal == null) {
-      throw new BadQueryRequestException(
-          "Pinot does not support column or function on the right-hand side of the predicate");
-    }
-    return RequestUtils.getLiteralString(literal);
+    return evaluateLiteralValue(thriftExpression).getStringValue();
   }
 
   /**
@@ -402,23 +401,23 @@ private static FilterContext getFilterInner(FunctionContext filterFunction) {
       case GREATER_THAN:
         return FilterContext.forPredicate(
             new RangePredicate(operands.get(0), false, getStringValue(operands.get(1)), false, RangePredicate.UNBOUNDED,
-                operands.get(1).getLiteral().getType()));
+                getLiteralType(operands.get(1))));
       case GREATER_THAN_OR_EQUAL:
         return FilterContext.forPredicate(
             new RangePredicate(operands.get(0), true, getStringValue(operands.get(1)), false, RangePredicate.UNBOUNDED,
-                operands.get(1).getLiteral().getType()));
+                getLiteralType(operands.get(1))));
       case LESS_THAN:
         return FilterContext.forPredicate(
             new RangePredicate(operands.get(0), false, RangePredicate.UNBOUNDED, false, getStringValue(operands.get(1)),
-                operands.get(1).getLiteral().getType()));
+                getLiteralType(operands.get(1))));
       case LESS_THAN_OR_EQUAL:
         return FilterContext.forPredicate(
             new RangePredicate(operands.get(0), false, RangePredicate.UNBOUNDED, true, getStringValue(operands.get(1)),
-                operands.get(1).getLiteral().getType()));
+                getLiteralType(operands.get(1))));
       case BETWEEN:
         return FilterContext.forPredicate(
             new RangePredicate(operands.get(0), true, getStringValue(operands.get(1)), true,
-                getStringValue(operands.get(2)), operands.get(1).getLiteral().getType()));
+                getStringValue(operands.get(2)), getLiteralType(operands.get(1))));
       case RANGE:
         return FilterContext.forPredicate(new RangePredicate(operands.get(0), getStringValue(operands.get(1))));
       case REGEXP_LIKE:
@@ -462,11 +461,96 @@ private static FilterContext getFilterInner(FunctionContext filterFunction) {
   //       literal context doesn't support float, and we cannot differentiate explicit string literal and literal
   //       without explicit type, so we always convert the literal into string.
   private static String getStringValue(ExpressionContext expressionContext) {
-    if (expressionContext.getType() != ExpressionContext.Type.LITERAL) {
-      throw new BadQueryRequestException(
-          "Pinot does not support column or function on the right-hand side of the predicate");
+    return evaluateLiteralValue(expressionContext).getStringValue();
+  }
+
+  private static DataType getLiteralType(Expression thriftExpression) {
+    return evaluateLiteralValue(thriftExpression).getType();
+  }
+
+  private static DataType getLiteralType(ExpressionContext expressionContext) {
+    return evaluateLiteralValue(expressionContext).getType();
+  }
+
+  private static EvaluatedLiteralValue evaluateLiteralValue(Expression thriftExpression) {
+    Literal literal = thriftExpression.getLiteral();
+    if (literal != null) {
+      return fromLiteralContext(new LiteralContext(literal));
     }
-    return expressionContext.getLiteral().getStringValue();
+    Function function = thriftExpression.getFunctionCall();
+    if (function != null) {
+      return evaluateFunctionLiteral(function.getOperator(), function.getOperands(),
+          RequestContextUtils::evaluateLiteralValue);
+    }
+    throw new BadQueryRequestException(UNSUPPORTED_RHS_MESSAGE);
+  }
+
+  private static EvaluatedLiteralValue evaluateLiteralValue(ExpressionContext expressionContext) {
+    if (expressionContext.getType() == ExpressionContext.Type.LITERAL) {
+      return fromLiteralContext(expressionContext.getLiteral());
+    }
+    if (expressionContext.getType() == ExpressionContext.Type.FUNCTION) {
+      FunctionContext function = expressionContext.getFunction();
+      return evaluateFunctionLiteral(function.getFunctionName(), function.getArguments(),
+          RequestContextUtils::evaluateLiteralValue);
+    }
+    throw new BadQueryRequestException(UNSUPPORTED_RHS_MESSAGE);
+  }
+
+  private static <T> EvaluatedLiteralValue evaluateFunctionLiteral(String functionName, List<T> operands,
+      java.util.function.Function<T, EvaluatedLiteralValue> evaluator) {
+    if (!functionName.equalsIgnoreCase("cast")) {
+      throw new BadQueryRequestException(UNSUPPORTED_RHS_MESSAGE);
+    }
+    Preconditions.checkState(operands.size() == 2, "CAST function must have exactly 2 operands");
+    EvaluatedLiteralValue source = evaluator.apply(operands.get(0));
+    EvaluatedLiteralValue target = evaluator.apply(operands.get(1));
+    DataType targetType = getCastTargetType(target.getStringValue());
+    if (source.getType() == DataType.UNKNOWN) {
+      return new EvaluatedLiteralValue(source.getStringValue(), targetType);
+    }
+    Object converted = targetType.convert(source.getStringValue());
+    return new EvaluatedLiteralValue(targetType.toString(converted), targetType);
+  }
+
+  private static DataType getCastTargetType(String targetTypeString) {
+    switch (targetTypeString.toUpperCase()) {
+      case "INT":
+      case "INTEGER":
+        return DataType.INT;
+      case "LONG":
+      case "BIGINT":
+        return DataType.LONG;
+      case "FLOAT":
+        return DataType.FLOAT;
+      case "DOUBLE":
+        return DataType.DOUBLE;
+      case "DECIMAL":
+      case "BIGDECIMAL":
+      case "BIG_DECIMAL":
+        return DataType.BIG_DECIMAL;
+      case "BOOL":
+      case "BOOLEAN":
+        return DataType.BOOLEAN;
+      case "TIMESTAMP":
+        return DataType.TIMESTAMP;
+      case "STRING":
+      case "VARCHAR":
+        return DataType.STRING;
+      case "JSON":
+        return DataType.JSON;
+      case "BYTES":
+      case "VARBINARY":
+        return DataType.BYTES;
+      case "UUID":
+        return DataType.UUID;
+      default:
+        throw new BadQueryRequestException("Unable to cast expression to type - " + targetTypeString);
+    }
+  }
+
+  private static EvaluatedLiteralValue fromLiteralContext(LiteralContext literalContext) {
+    return new EvaluatedLiteralValue(literalContext.getStringValue(), literalContext.getType());
   }
 
   private static float[] getVectorValue(ExpressionContext expressionContext) {
@@ -556,4 +640,22 @@ private static float getFloatValue(Expression thriftExpression) {
         throw new IllegalStateException("Unsupported literal type: " + type);
     }
   }
+
+  private static final class EvaluatedLiteralValue {
+    private final String _stringValue;
+    private final DataType _type;
+
+    private EvaluatedLiteralValue(String stringValue, DataType type) {
+      _stringValue = stringValue;
+      _type = type;
+    }
+
+    private String getStringValue() {
+      return _stringValue;
+    }
+
+    private DataType getType() {
+      return _type;
+    }
+  }
 }

pinot-common/src/test/java/org/apache/pinot/common/request/context/RequestContextUtilsTest.java

@@ -0,0 +1,57 @@
+/**
+ * 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.common.request.context;
+
+import java.util.List;
+import org.apache.pinot.common.request.context.predicate.EqPredicate;
+import org.apache.pinot.common.request.context.predicate.InPredicate;
+import org.apache.pinot.sql.parsers.CalciteSqlParser;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+
+/**
+ * Tests filter conversion for literal-only expressions on the right-hand side.
+ */
+public class RequestContextUtilsTest {
+  private static final String UUID_1 = "550e8400-e29b-41d4-a716-446655440000";
+  private static final String UUID_2 = "550e8400-e29b-41d4-a716-446655440001";
+
+  @Test
+  public void testGetFilterWithUuidCastLiteralRhs() {
+    FilterContext filter =
+        RequestContextUtils.getFilter(CalciteSqlParser.compileToExpression("uuidCol = CAST('" + UUID_1 + "' AS UUID)"));
+
+    Assert.assertEquals(filter.getType(), FilterContext.Type.PREDICATE);
+    EqPredicate predicate = (EqPredicate) filter.getPredicate();
+    Assert.assertEquals(predicate.getLhs().getIdentifier(), "uuidCol");
+    Assert.assertEquals(predicate.getValue(), UUID_1);
+  }
+
+  @Test
+  public void testGetFilterExpressionContextWithUuidCastLiteralIn() {
+    FilterContext filter = RequestContextUtils.getFilter(
+        RequestContextUtils.getExpression("uuidCol IN (CAST('" + UUID_1 + "' AS UUID), CAST('" + UUID_2 + "' AS UUID))"));
+
+    Assert.assertEquals(filter.getType(), FilterContext.Type.PREDICATE);
+    InPredicate predicate = (InPredicate) filter.getPredicate();
+    Assert.assertEquals(predicate.getLhs().getIdentifier(), "uuidCol");
+    Assert.assertEquals(predicate.getValues(), List.of(UUID_1, UUID_2));
+  }
+}

pinot-common/src/test/java/org/apache/pinot/common/utils/DataSchemaTest.java

@@ -21,8 +21,8 @@
 import java.math.BigDecimal;
 import java.nio.ByteBuffer;
 import java.sql.Timestamp;
-import org.apache.pinot.spi.utils.ByteArray;
 import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.utils.ByteArray;
 import org.apache.pinot.spi.utils.BytesUtils;
 import org.apache.pinot.spi.utils.UuidUtils;
 import org.testng.Assert;

pinot-core/src/main/java/org/apache/pinot/core/operator/filter/RawValueInvertedIndexFilterOperator.java

@@ -36,6 +36,8 @@
 import org.apache.pinot.segment.local.segment.index.readers.RawValueBitmapInvertedIndexReader;
 import org.apache.pinot.segment.spi.datasource.DataSource;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.utils.BytesUtils;
+import org.apache.pinot.spi.utils.UuidUtils;
 import org.roaringbitmap.buffer.ImmutableRoaringBitmap;
 import org.roaringbitmap.buffer.MutableRoaringBitmap;
 
@@ -138,6 +140,12 @@ private void addMatchingValueBitmap(MutableRoaringBitmap bitmap, String value) {
       case STRING:
         valueBitmap = _invertedIndexReader.getDocIdsForString(value);
         break;
+      case BYTES:
+        valueBitmap = _invertedIndexReader.getDocIdsForBytes(BytesUtils.toBytes(value));
+        break;
+      case UUID:
+        valueBitmap = _invertedIndexReader.getDocIdsForBytes(UuidUtils.toBytes(value));
+        break;
       default:
         throw new IllegalStateException("Unsupported data type for raw inverted index: " + _dataType);
     }

pinot-core/src/main/java/org/apache/pinot/core/query/pruner/ValueBasedSegmentPruner.java

@@ -236,7 +236,7 @@ public void ensureDataType(DataType dt) {
       public boolean mightBeContained(BloomFilterReader bloomFilter) {
         if (!_hashed) {
           GuavaBloomFilterReaderUtils.Hash128AsLongs hash128AsLongs =
-              GuavaBloomFilterReaderUtils.hashAsLongs(_comparableValue.toString());
+              GuavaBloomFilterReaderUtils.hashAsLongs(_dt.toString(_comparableValue));
           _hash1 = hash128AsLongs.getHash1();
           _hash2 = hash128AsLongs.getHash2();
           _hashed = true;

pinot-core/src/test/java/org/apache/pinot/core/query/pruner/BloomFilterSegmentPrunerTest.java

@@ -96,6 +96,22 @@ public void testBloomFilterPruning()
     assertTrue(runPruner(indexSegment, "SELECT COUNT(*) FROM testTable WHERE column = 21.0 AND column = 30.0"));
   }
 
+  @Test
+  public void testUuidBloomFilterPruning()
+      throws IOException {
+    IndexSegment indexSegment = mockIndexSegment(new String[]{
+        "550e8400-e29b-41d4-a716-446655440000",
+        "550e8400-e29b-41d4-a716-446655440001"
+    }, DataType.UUID);
+
+    assertFalse(runPruner(indexSegment,
+        "SELECT COUNT(*) FROM testTable WHERE column = '550e8400-e29b-41d4-a716-446655440000'"));
+    assertFalse(runPruner(indexSegment,
+        "SELECT COUNT(*) FROM testTable WHERE column IN ('550e8400-e29b-41d4-a716-446655440001')"));
+    assertTrue(runPruner(indexSegment,
+        "SELECT COUNT(*) FROM testTable WHERE column = '550e8400-e29b-41d4-a716-44665544ffff'"));
+  }
+
   @Test(expectedExceptions = RuntimeException.class)
   public void testQueryTimeoutOnPruning()
       throws IOException {
@@ -162,6 +178,11 @@ public void testIsApplicableTo() {
 
   private IndexSegment mockIndexSegment(String[] values)
       throws IOException {
+    return mockIndexSegment(values, DataType.DOUBLE);
+  }
+
+  private IndexSegment mockIndexSegment(String[] values, DataType dataType)
+      throws IOException {
     IndexSegment indexSegment = mock(IndexSegment.class);
     when(indexSegment.getColumnNames()).thenReturn(ImmutableSet.of("column"));
     SegmentMetadata segmentMetadata = mock(SegmentMetadata.class);
@@ -176,7 +197,7 @@ private IndexSegment mockIndexSegment(String[] values)
     for (String v : values) {
       builder.put(v);
     }
-    when(dataSourceMetadata.getDataType()).thenReturn(DataType.DOUBLE);
+    when(dataSourceMetadata.getDataType()).thenReturn(dataType);
     when(dataSource.getDataSourceMetadata()).thenReturn(dataSourceMetadata);
     when(dataSource.getBloomFilter()).thenReturn(builder.build());
 

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

@@ -81,6 +81,21 @@ protected int getRealtimeSegmentFlushSize() {
     return 2;
   }
 
+  @Override
+  protected List<String> getInvertedIndexColumns() {
+    return List.of(UUID_FROM_BYTES_COLUMN);
+  }
+
+  @Override
+  protected List<String> getNoDictionaryColumns() {
+    return List.of(UUID_FROM_BYTES_COLUMN);
+  }
+
+  @Override
+  protected List<String> getBloomFilterColumns() {
+    return List.of(UUID_FROM_STRING_COLUMN, UUID_FROM_BYTES_COLUMN);
+  }
+
   @Override
   public Schema createSchema() {
     return new Schema.SchemaBuilder().setSchemaName(getTableName())

pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroIngestionSchemaValidator.java

@@ -139,7 +139,7 @@ private void validateSchemas() {
         if (fieldSpec.getDataType() != dataTypeForSVColumn) {
           _dataTypeMismatch.addMismatchReason(String
               .format("The Pinot column: (%s: %s) doesn't match with the column (%s: %s) in input %s schema.",
-                  columnName, fieldSpec.getDataType().name(), avroColumnName, avroColumnType.name(),
+                  columnName, fieldSpec.getDataType().name(), avroColumnName, dataTypeForSVColumn.name(),
                   getInputSchemaType()));
         }
       } else {