Fix NPE for IN clause on constant STRING dictionary (apache#11930)

vvivekiyer · Nov 2, 2023 · fee11d6 · fee11d6
1 parent 403790b
commit fee11d6
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/...tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java b/...tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java
@@ -1767,21 +1767,41 @@ public void testBrokerResponseMetadata(boolean useMultiStageQueryEngine)
   public void testInBuiltVirtualColumns(boolean useMultiStageQueryEngine)
       throws Exception {
     setUseMultiStageQueryEngine(useMultiStageQueryEngine);
-    String query = "SELECT $docId, $hostName, $segmentName FROM mytable";
+
+    String query = "SELECT $docId, $hostName, $segmentName FROM mytable LIMIT 10";
     JsonNode response = postQuery(query);
     JsonNode resultTable = response.get("resultTable");
     JsonNode dataSchema = resultTable.get("dataSchema");
     assertEquals(dataSchema.get("columnNames").toString(), "[\"$docId\",\"$hostName\",\"$segmentName\"]");
     assertEquals(dataSchema.get("columnDataTypes").toString(), "[\"INT\",\"STRING\",\"STRING\"]");
     JsonNode rows = resultTable.get("rows");
+    assertEquals(rows.size(), 10);
     String expectedHostName = NetUtils.getHostnameOrAddress();
     String expectedSegmentNamePrefix = "mytable_";
     for (int i = 0; i < 10; i++) {
       JsonNode row = rows.get(i);
       assertEquals(row.get(0).asInt(), i);
       assertEquals(row.get(1).asText(), expectedHostName);
-      assertTrue(row.get(2).asText().startsWith(expectedSegmentNamePrefix));
+      String segmentName = row.get(2).asText();
+      assertTrue(segmentName.startsWith(expectedSegmentNamePrefix));
     }
+
+    // Collect all segment names
+    query = "SELECT DISTINCT $segmentName FROM mytable LIMIT 10000";
+    response = postQuery(query);
+    rows = response.get("resultTable").get("rows");
+    int numSegments = rows.size();
+    List<String> segmentNames = new ArrayList<>(numSegments);
+    for (int i = 0; i < numSegments; i++) {
+      segmentNames.add(rows.get(i).get(0).asText());
+    }
+    // Test IN clause on $segmentName
+    Collections.shuffle(segmentNames);
+    int numSegmentsToQuery = RANDOM.nextInt(numSegments) + 1;
+    query = "SELECT COUNT(*) FROM mytable WHERE $segmentName IN ('" + String.join("','",
+        segmentNames.subList(0, numSegmentsToQuery)) + "')";
+    response = postQuery(query);
+    assertEquals(response.get("numSegmentsMatched").asInt(), numSegmentsToQuery);
   }
 
   @Test(dataProvider = "useBothQueryEngines")

diff --git a/...a/org/apache/pinot/segment/local/segment/index/readers/ConstantValueStringDictionary.java b/...a/org/apache/pinot/segment/local/segment/index/readers/ConstantValueStringDictionary.java
@@ -18,7 +18,10 @@
  */
 package org.apache.pinot.segment.local.segment.index.readers;
 
+import it.unimi.dsi.fastutil.ints.IntSet;
 import java.math.BigDecimal;
+import java.util.Collections;
+import java.util.List;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -113,4 +116,18 @@ public String getStringValue(int dictId) {
   public byte[] getBytesValue(int dictId) {
     return _bytes;
   }
+
+  @Override
+  public void getDictIds(List<String> values, IntSet dictIds) {
+    if (values.contains(_value)) {
+      dictIds.add(0);
+    }
+  }
+
+  @Override
+  public void getDictIds(List<String> sortedValues, IntSet dictIds, SortedBatchLookupAlgorithm algorithm) {
+    if (Collections.binarySearch(sortedValues, _value) >= 0) {
+      dictIds.add(0);
+    }
+  }
 }