[UUID 4b] Bloom filter support for the logical UUID type

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

@@ -31,29 +31,50 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.Executors;
+import org.apache.commons.io.FileUtils;
 import org.apache.pinot.core.query.request.context.QueryContext;
 import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
+import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig;
 import org.apache.pinot.segment.local.segment.index.readers.bloom.OnHeapGuavaBloomFilterReader;
+import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader;
+import org.apache.pinot.segment.spi.ImmutableSegment;
 import org.apache.pinot.segment.spi.IndexSegment;
 import org.apache.pinot.segment.spi.SegmentMetadata;
+import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig;
 import org.apache.pinot.segment.spi.datasource.DataSource;
 import org.apache.pinot.segment.spi.datasource.DataSourceMetadata;
 import org.apache.pinot.segment.spi.index.reader.BloomFilterReader;
 import org.apache.pinot.segment.spi.memory.PinotDataBuffer;
+import org.apache.pinot.spi.config.table.BloomFilterConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
 import org.apache.pinot.spi.env.PinotConfiguration;
+import org.apache.pinot.spi.utils.ReadMode;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
 import org.testng.annotations.BeforeClass;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotNull;
 import static org.testng.Assert.assertTrue;
 
 
 public class BloomFilterSegmentPrunerTest {
   private static final BloomFilterSegmentPruner PRUNER = new BloomFilterSegmentPruner();
+  private static final String UUID_COLUMN = "uuidColumn";
+  private static final String UUID_0 = "550e8400-e29b-41d4-a716-446655440000";
+  private static final String ABSENT_UUID = "550e8400-e29b-41d4-a716-446655440001";
+  private static final String UUID_2 = "550e8400-e29b-41d4-a716-446655440002";
+  private static final String UUID_3 = "550e8400-e29b-41d4-a716-446655440003";
 
   @BeforeClass
   public void setUp() {
@@ -95,6 +116,104 @@ public void testBloomFilterPruning()
     assertTrue(runPruner(indexSegment, "SELECT COUNT(*) FROM testTable WHERE column = 21.0 AND column = 30.0"));
   }
 
+  @Test(dataProvider = "uuidBloomFilterCreationModes")
+  public void testUuidBloomFilterPruningEndToEnd(boolean noDictionary, boolean multiValue, boolean createOnLoad)
+      throws Exception {
+    File indexDir = Files.createTempDirectory("uuidBloomFilterSegment").toFile();
+    ImmutableSegment segment = null;
+    try {
+      TableConfig tableConfig = createTableConfig(noDictionary, !createOnLoad);
+      Schema schema = createSchema(multiValue);
+
+      List<GenericRow> rows = new ArrayList<>();
+      rows.add(row(UUID_0, multiValue));
+      rows.add(row(UUID_2, multiValue));
+
+      String segmentName = (noDictionary ? "raw" : "dictionary") + (multiValue ? "Mv" : "Sv") + "UuidSegment";
+      SegmentGeneratorConfig generatorConfig = new SegmentGeneratorConfig(tableConfig, schema);
+      generatorConfig.setSegmentName(segmentName);
+      generatorConfig.setOutDir(indexDir.getAbsolutePath());
+      SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
+      driver.init(generatorConfig, new GenericRowRecordReader(rows));
+      driver.build();
+
+      File segmentDir = new File(indexDir, segmentName);
+      if (createOnLoad) {
+        segment = ImmutableSegmentLoader.load(segmentDir,
+            new IndexLoadingConfig(createTableConfig(noDictionary, true), schema));
+      } else {
+        segment = ImmutableSegmentLoader.load(segmentDir, ReadMode.mmap);
+      }
+      assertEquals(segment.getSegmentMetadata().getColumnMetadataFor(UUID_COLUMN).hasDictionary(), !noDictionary);
+      assertEquals(segment.getSegmentMetadata().getColumnMetadataFor(UUID_COLUMN).isSingleValue(), !multiValue);
+      assertEquals(segment.getDataSource(UUID_COLUMN).getDataSourceMetadata().getDataType(), DataType.UUID);
+      assertNotNull(segment.getDataSource(UUID_COLUMN).getBloomFilter());
+
+      assertFalse(runPruner(segment,
+          "SELECT COUNT(*) FROM testTable WHERE uuidColumn = '" + UUID_0 + "'"));
+      assertFalse(runPruner(segment,
+          "SELECT COUNT(*) FROM testTable WHERE uuidColumn = '550E8400-E29B-41D4-A716-446655440000'"));
+      assertFalse(runPruner(segment,
+          "SELECT COUNT(*) FROM testTable WHERE uuidColumn = '550e8400e29b41d4a716446655440000'"));
+      if (multiValue) {
+        assertFalse(runPruner(segment,
+            "SELECT COUNT(*) FROM testTable WHERE uuidColumn = '" + UUID_3 + "'"));
+      }
+      assertFalse(runPruner(segment,
+          "SELECT COUNT(*) FROM testTable WHERE uuidColumn IN ('" + ABSENT_UUID + "', '" + UUID_2 + "')"));
+      assertTrue(runPruner(segment,
+          "SELECT COUNT(*) FROM testTable WHERE uuidColumn = '" + ABSENT_UUID + "'"));
+    } finally {
+      if (segment != null) {
+        segment.destroy();
+      }
+      FileUtils.deleteDirectory(indexDir);
+    }
+  }
+
+  @DataProvider(name = "uuidBloomFilterCreationModes")
+  public Object[][] uuidBloomFilterCreationModes() {
+    return new Object[][]{
+        {false, false, false},
+        {true, false, false},
+        {false, true, false},
+        {true, true, false},
+        {false, false, true},
+        {true, false, true},
+        {false, true, true},
+        {true, true, true}
+    };
+  }
+
+  private static TableConfig createTableConfig(boolean noDictionary, boolean enableBloomFilter) {
+    TableConfigBuilder builder = new TableConfigBuilder(TableType.OFFLINE).setTableName("testTable");
+    if (noDictionary) {
+      builder.setNoDictionaryColumns(List.of(UUID_COLUMN));
+    }
+    TableConfig tableConfig = builder.build();
+    if (enableBloomFilter) {
+      tableConfig.getIndexingConfig().setBloomFilterConfigs(
+          Map.of(UUID_COLUMN, new BloomFilterConfig(1e-9, 0, false)));
+    }
+    return tableConfig;
+  }
+
+  private static Schema createSchema(boolean multiValue) {
+    Schema.SchemaBuilder builder = new Schema.SchemaBuilder();
+    if (multiValue) {
+      builder.addMultiValueDimension(UUID_COLUMN, DataType.UUID);
+    } else {
+      builder.addSingleValueDimension(UUID_COLUMN, DataType.UUID);
+    }
+    return builder.build();
+  }
+
+  private static GenericRow row(String uuid, boolean multiValue) {
+    GenericRow row = new GenericRow();
+    row.putValue(UUID_COLUMN, multiValue ? new String[]{uuid, UUID_3} : uuid);
+    return row;
+  }
+
   @Test(expectedExceptions = RuntimeException.class)
   public void testQueryTimeoutOnPruning()
       throws IOException {

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

@@ -0,0 +1,124 @@
+/**
+ * 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.integration.tests.custom;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import java.io.File;
+import java.util.List;
+import java.util.Map;
+import org.apache.avro.file.DataFileWriter;
+import org.apache.avro.generic.GenericData;
+import org.apache.pinot.spi.config.table.BloomFilterConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+
+
+/// End-to-end coverage for querying a UUID column backed by a Bloom filter. A single segment contains UUIDs ending in
+/// `0000` and `0002`; the absent `0001` value is inside the segment min/max range, and the pruning metrics verify that
+/// the Bloom value pruner eliminated it.
+@Test(suiteName = "CustomClusterIntegrationTest")
+public class UuidBloomFilterTest extends CustomDataQueryClusterIntegrationTest {
+  private static final String TABLE_NAME = "UuidBloomFilterTest";
+  private static final String UUID_COLUMN = "uuidColumn";
+  private static final String UUID_0 = "550e8400-e29b-41d4-a716-446655440000";
+  private static final String UUID_0_HEX = "550e8400e29b41d4a716446655440000";
+  private static final String UUID_1_HEX = "550e8400e29b41d4a716446655440001";
+  private static final String UUID_2 = "550e8400-e29b-41d4-a716-446655440002";
+
+  @Override
+  public String getTableName() {
+    return TABLE_NAME;
+  }
+
+  @Override
+  protected long getCountStarResult() {
+    return 2;
+  }
+
+  @Override
+  public int getNumAvroFiles() {
+    return 1;
+  }
+
+  @Override
+  public TableConfig createOfflineTableConfig() {
+    TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName()).build();
+    tableConfig.getIndexingConfig().setBloomFilterConfigs(
+        Map.of(UUID_COLUMN, new BloomFilterConfig(1e-9, 0, false)));
+    return tableConfig;
+  }
+
+  @Override
+  public Schema createSchema() {
+    return new Schema.SchemaBuilder().setSchemaName(getTableName())
+        .addSingleValueDimension(UUID_COLUMN, DataType.UUID)
+        .build();
+  }
+
+  @Override
+  public List<File> createAvroFiles()
+      throws Exception {
+    org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("uuidRecord", null, null, false);
+    avroSchema.setFields(List.of(new org.apache.avro.Schema.Field(UUID_COLUMN,
+        org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING), null, null)));
+
+    try (AvroFilesAndWriters avroFilesAndWriters = createAvroFilesAndWriters(avroSchema)) {
+      DataFileWriter<GenericData.Record> writer = avroFilesAndWriters.getWriters().get(0);
+      for (String uuid : List.of(UUID_0, UUID_2)) {
+        GenericData.Record record = new GenericData.Record(avroSchema);
+        record.put(UUID_COLUMN, uuid);
+        writer.append(record);
+      }
+      return avroFilesAndWriters.getAvroFiles();
+    }
+  }
+
+  @Test
+  public void testUuidBloomFilterQueries()
+      throws Exception {
+    setUseMultiStageQueryEngine(false);
+
+    assertCountAndPrunedSegments(
+        String.format("SELECT COUNT(*) FROM %s WHERE %s = '%s'", getTableName(), UUID_COLUMN, UUID_0_HEX), 1, 1, 0);
+    assertCountAndPrunedSegments(String.format(
+        "SELECT COUNT(*) FROM %s WHERE %s = CAST('%s' AS UUID)", getTableName(), UUID_COLUMN, UUID_2), 1, 1, 0);
+    assertCountAndPrunedSegments(
+        String.format("SELECT COUNT(*) FROM %s WHERE %s = '%s'", getTableName(), UUID_COLUMN, UUID_1_HEX), 0, 0, 1);
+  }
+
+  private void assertCountAndPrunedSegments(String query, long expectedCount, int expectedProcessedSegments,
+      int expectedPrunedSegments)
+      throws Exception {
+    JsonNode response = postQuery(query);
+    assertTrue(response.path("exceptions").isEmpty(), response.toPrettyString());
+    assertEquals(response.path("resultTable").path("rows").path(0).path(0).asLong(), expectedCount,
+        response.toPrettyString());
+    assertEquals(response.path("numSegmentsQueried").asInt(), 1, response.toPrettyString());
+    assertEquals(response.path("numSegmentsProcessed").asInt(), expectedProcessedSegments, response.toPrettyString());
+    assertEquals(response.path("numSegmentsPrunedByValue").asInt(), expectedPrunedSegments,
+        response.toPrettyString());
+  }
+}

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/bloomfilter/BloomFilterHandler.java

@@ -130,6 +130,7 @@ private void createAndSealBloomFilterForNonDictionaryColumn(File indexDir, Colum
       BloomFilterConfig bloomFilterConfig, SegmentDirectory.Writer segmentWriter)
       throws Exception {
     int numDocs = columnMetadata.getTotalDocs();
+    DataType storedType = columnMetadata.getDataType().getStoredType();
     IndexCreationContext context = new IndexCreationContext.Builder(indexDir, _tableConfig, columnMetadata).build();
     IndexReaderFactory<ForwardIndexReader> readerFactory = StandardIndexes.forward().getReaderFactory();
     try (BloomFilterCreator bloomFilterCreator = StandardIndexes.bloomFilter()
@@ -139,7 +140,7 @@ private void createAndSealBloomFilterForNonDictionaryColumn(File indexDir, Colum
         ForwardIndexReaderContext readerContext = forwardIndexReader.createContext()) {
       if (columnMetadata.isSingleValue()) {
         // SV
-        switch (columnMetadata.getDataType()) {
+        switch (storedType) {
           case INT:
             for (int i = 0; i < numDocs; i++) {
               bloomFilterCreator.add(Integer.toString(forwardIndexReader.getInt(i, readerContext)));
@@ -177,7 +178,7 @@ private void createAndSealBloomFilterForNonDictionaryColumn(File indexDir, Colum
         bloomFilterCreator.seal();
       } else {
         // MV
-        switch (columnMetadata.getDataType()) {
+        switch (storedType) {
           case INT:
             for (int i = 0; i < numDocs; i++) {
               int[] buffer = new int[columnMetadata.getMaxNumberOfMultiValues()];
@@ -286,9 +287,9 @@ private void createBloomFilterForColumn(SegmentDirectory.Writer segmentWriter, C
 
   private Dictionary getDictionaryReader(ColumnMetadata columnMetadata, SegmentDirectory.Writer segmentWriter)
       throws IOException {
-    DataType dataType = columnMetadata.getDataType();
+    DataType storedType = columnMetadata.getDataType().getStoredType();
 
-    switch (dataType) {
+    switch (storedType) {
       case INT:
       case LONG:
       case FLOAT:
@@ -299,7 +300,7 @@ private Dictionary getDictionaryReader(ColumnMetadata columnMetadata, SegmentDir
         return DictionaryIndexType.read(buf, columnMetadata, DictionaryIndexConfig.DEFAULT);
       default:
         throw new IllegalStateException(
-            "Unsupported data type: " + dataType + " for column: " + columnMetadata.getColumnName());
+            "Unsupported data type: " + storedType + " for column: " + columnMetadata.getColumnName());
     }
   }
 }

pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/BloomFilterCreatorTest.java

@@ -83,6 +83,8 @@ public void testUuidBloomFilterCreatorWithBytesValues()
     String columnName = "uuidColumn";
     String uuid0 = "550e8400-e29b-41d4-a716-446655440000";
     String uuid1 = "550e8400-e29b-41d4-a716-446655440001";
+    String uuid0Hex = "550e8400e29b41d4a716446655440000";
+    String uuid1Hex = "550e8400e29b41d4a716446655440001";
     try (BloomFilterCreator bloomFilterCreator = new OnHeapGuavaBloomFilterCreator(TEMP_DIR, columnName, cardinality,
         new BloomFilterConfig(BloomFilterConfig.DEFAULT_FPP, 0, false), FieldSpec.DataType.UUID)) {
       bloomFilterCreator.add(UuidUtils.toBytes(uuid0), -1);
@@ -94,12 +96,12 @@ public void testUuidBloomFilterCreatorWithBytesValues()
     try (PinotDataBuffer dataBuffer = PinotDataBuffer.mapReadOnlyBigEndianFile(bloomFilterFile);
         BloomFilterReader onHeapBloomFilter = BloomFilterReaderFactory.getBloomFilterReader(dataBuffer, true);
         BloomFilterReader offHeapBloomFilter = BloomFilterReaderFactory.getBloomFilterReader(dataBuffer, false)) {
-      Assert.assertTrue(onHeapBloomFilter.mightContain(uuid0));
-      Assert.assertTrue(onHeapBloomFilter.mightContain(uuid1));
-      Assert.assertFalse(onHeapBloomFilter.mightContain("550e8400-e29b-41d4-a716-4466554400ff"));
-      Assert.assertTrue(offHeapBloomFilter.mightContain(uuid0));
-      Assert.assertTrue(offHeapBloomFilter.mightContain(uuid1));
-      Assert.assertFalse(offHeapBloomFilter.mightContain("550e8400-e29b-41d4-a716-4466554400ff"));
+      Assert.assertTrue(onHeapBloomFilter.mightContain(uuid0Hex));
+      Assert.assertTrue(onHeapBloomFilter.mightContain(uuid1Hex));
+      Assert.assertFalse(onHeapBloomFilter.mightContain("550e8400e29b41d4a7164466554400ff"));
+      Assert.assertTrue(offHeapBloomFilter.mightContain(uuid0Hex));
+      Assert.assertTrue(offHeapBloomFilter.mightContain(uuid1Hex));
+      Assert.assertFalse(offHeapBloomFilter.mightContain("550e8400e29b41d4a7164466554400ff"));
     }
   }