Json extract index filter support

pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java

@@ -117,7 +117,8 @@ public enum TransformFunctionType {
       ReturnTypes.cascade(opBinding -> positionalReturnTypeInferenceFromStringLiteral(opBinding, 2,
           SqlTypeName.VARCHAR), SqlTypeTransforms.FORCE_NULLABLE),
       OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER,
-          SqlTypeFamily.CHARACTER), ordinal -> ordinal > 2), "json_extract_index"),
+          SqlTypeFamily.CHARACTER, SqlTypeFamily.CHARACTER), ordinal -> ordinal > 2), "json_extract_index"),
+
   JSON_EXTRACT_KEY("jsonExtractKey", ReturnTypes.TO_ARRAY,
       OperandTypes.family(ImmutableList.of(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER)), "json_extract_key"),
 

pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/JsonExtractIndexTransformFunction.java

@@ -53,11 +53,11 @@ public String getName() {
 
   @Override
   public void init(List<TransformFunction> arguments, Map<String, ColumnContext> columnContextMap) {
-    // Check that there are exactly 3 or 4 arguments
-    if (arguments.size() < 3 || arguments.size() > 4) {
+    // Check that there are exactly 3 or 4 or 5 arguments
+    if (arguments.size() < 3 || arguments.size() > 5) {
       throw new IllegalArgumentException(
-          "Expected 3/4 arguments for transform function: jsonExtractIndex(jsonFieldName, 'jsonPath', 'resultsType',"
-              + " ['defaultValue'])");
+          "Expected 3/4/5 arguments for transform function: jsonExtractIndex(jsonFieldName, 'jsonPath', 'resultsType',"
+              + " ['defaultValue'], ['jsonFilterExpression'])");
     }
 
     TransformFunction firstArgument = arguments.get(0);
@@ -76,37 +76,45 @@ public void init(List<TransformFunction> arguments, Map<String, ColumnContext> c
     if (!(secondArgument instanceof LiteralTransformFunction)) {
       throw new IllegalArgumentException("JSON path argument must be a literal");
     }
-    String inputJsonPath = ((LiteralTransformFunction) secondArgument).getStringLiteral();
+    _jsonPathString = ((LiteralTransformFunction) secondArgument).getStringLiteral();
     try {
-      JsonPathCache.INSTANCE.getOrCompute(inputJsonPath);
+      JsonPathCache.INSTANCE.getOrCompute(_jsonPathString);
     } catch (Exception e) {
       throw new IllegalArgumentException("JSON path argument is not a valid JSON path");
     }
-    _jsonPathString = inputJsonPath.substring(1); // remove $ prefix
 
     TransformFunction thirdArgument = arguments.get(2);
     if (!(thirdArgument instanceof LiteralTransformFunction)) {
       throw new IllegalArgumentException("Result type argument must be a literal");
     }
     String resultsType = ((LiteralTransformFunction) thirdArgument).getStringLiteral().toUpperCase();
     boolean isSingleValue = !resultsType.endsWith("_ARRAY");
-    if (isSingleValue && inputJsonPath.contains("[*]")) {
+    if (isSingleValue && _jsonPathString.contains("[*]")) {
       throw new IllegalArgumentException(
           "[*] syntax in json path is unsupported for singleValue field json_extract_index");
     }
     DataType dataType = isSingleValue ? DataType.valueOf(resultsType)
         : DataType.valueOf(resultsType.substring(0, resultsType.length() - 6));
 
-    if (arguments.size() == 4) {
+    if (arguments.size() >= 4) {
       TransformFunction fourthArgument = arguments.get(3);
       if (!(fourthArgument instanceof LiteralTransformFunction)) {
         throw new IllegalArgumentException("Default value must be a literal");
       }
       _defaultValue = dataType.convert(((LiteralTransformFunction) fourthArgument).getStringLiteral());
     }
 
+    String filterJsonPath = null;
+    if (arguments.size() == 5) {
+      TransformFunction fifthArgument = arguments.get(4);
+      if (!(fifthArgument instanceof LiteralTransformFunction)) {
+        throw new IllegalArgumentException("JSON path filter argument must be a literal");
+      }
+      filterJsonPath = ((LiteralTransformFunction) fifthArgument).getStringLiteral();
+    }
+
     _resultMetadata = new TransformResultMetadata(dataType, isSingleValue, false);
-    _valueToMatchingDocsMap = _jsonIndexReader.getMatchingFlattenedDocsMap(_jsonPathString);
+    _valueToMatchingDocsMap = _jsonIndexReader.getMatchingFlattenedDocsMap(_jsonPathString, filterJsonPath);
     if (isSingleValue) {
       // For single value result type, it's more efficient to use original docIDs map
       _jsonIndexReader.convertFlattenedDocIdsToDocIds(_valueToMatchingDocsMap);

pinot-core/src/test/java/org/apache/pinot/core/operator/transform/function/JsonExtractIndexTransformFunctionTest.java

@@ -247,6 +247,23 @@ private void addMvTests(List<Object[]> testArguments) {
         String.format("jsonExtractIndex(%s,'%s','STRING_ARRAY')", JSON_STRING_SV_COLUMN,
             "$.arrayField[*].arrStringField"), "$.arrayField[*].arrStringField", DataType.STRING, false
     });
+
+    // MV with filters
+    testArguments.add(new Object[]{
+        String.format(
+            "jsonExtractIndex(%s,'%s','INT_ARRAY', '0', 'REGEXP_LIKE(\"$.arrayField[*].arrStringField\", ''.*y.*'')')",
+            JSON_STRING_SV_COLUMN,
+            "$.arrayField[*].arrIntField"), "$.arrayField[?(@.arrStringField =~ /.*y.*/)].arrIntField", DataType.INT,
+        false
+    });
+
+    testArguments.add(new Object[]{
+        String.format(
+            "jsonExtractIndex(%s,'%s','STRING_ARRAY', '0', '\"$.arrayField[*].arrIntField\" > 2')",
+            JSON_STRING_SV_COLUMN,
+            "$.arrayField[*].arrStringField"), "$.arrayField[?(@.arrIntField > 2)].arrStringField", DataType.STRING,
+        false
+    });
   }
 
   @Test(dataProvider = "testJsonExtractIndexDefaultValue")

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/json/MutableJsonIndexImpl.java

@@ -33,6 +33,7 @@
 import java.util.TreeMap;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
 import java.util.regex.Pattern;
+import javax.annotation.Nullable;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pinot.common.request.context.ExpressionContext;
 import org.apache.pinot.common.request.context.FilterContext;
@@ -367,10 +368,32 @@ public void convertFlattenedDocIdsToDocIds(Map<String, RoaringBitmap> valueToFla
   }
 
   @Override
-  public Map<String, RoaringBitmap> getMatchingFlattenedDocsMap(String jsonPathKey) {
+  public Map<String, RoaringBitmap> getMatchingFlattenedDocsMap(String jsonPathKey, @Nullable String filterString) {
     Map<String, RoaringBitmap> valueToMatchingFlattenedDocIdsMap = new HashMap<>();
     _readLock.lock();
     try {
+      RoaringBitmap filteredFlattenedDocIds = null;
+      FilterContext filter;
+      if (filterString != null) {
+        filter = RequestContextUtils.getFilter(CalciteSqlParser.compileToExpression(filterString));
+        Preconditions.checkArgument(!filter.isConstant(), "Invalid json match filter: " + filterString);
+        if (filter.getType() == FilterContext.Type.PREDICATE && isExclusive(filter.getPredicate().getType())) {
+          // Handle exclusive predicate separately because the flip can only be applied to the
+          // unflattened doc ids in order to get the correct result, and it cannot be nested
+          filteredFlattenedDocIds = getMatchingFlattenedDocIds(filter.getPredicate());
+          filteredFlattenedDocIds.flip(0, (long) _nextFlattenedDocId);
+        } else {
+          filteredFlattenedDocIds = getMatchingFlattenedDocIds(filter);
+        }
+      }
+      // Support 2 formats:
+      // - JSONPath format (e.g. "$.a[1].b"='abc', "$[0]"=1, "$"='abc')
+      // - Legacy format (e.g. "a[1].b"='abc')
+      if (jsonPathKey.startsWith("$")) {
+        jsonPathKey = jsonPathKey.substring(1);
+      } else {
+        jsonPathKey = JsonUtils.KEY_SEPARATOR + jsonPathKey;
+      }
       Pair<String, RoaringBitmap> result = getKeyAndFlattenedDocIds(jsonPathKey);
       jsonPathKey = result.getLeft();
       RoaringBitmap arrayIndexFlattenDocIds = result.getRight();
@@ -380,6 +403,9 @@ public Map<String, RoaringBitmap> getMatchingFlattenedDocsMap(String jsonPathKey
       Map<String, RoaringBitmap> subMap = getMatchingKeysMap(jsonPathKey);
       for (Map.Entry<String, RoaringBitmap> entry : subMap.entrySet()) {
         RoaringBitmap flattenedDocIds = entry.getValue().clone();
+        if (filteredFlattenedDocIds != null) {
+          flattenedDocIds.and(filteredFlattenedDocIds);
+        }
         if (arrayIndexFlattenDocIds != null) {
           flattenedDocIds.and(arrayIndexFlattenDocIds);
         }

pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/json/ImmutableJsonIndexReader.java

@@ -28,6 +28,7 @@
 import java.util.Map;
 import java.util.PriorityQueue;
 import java.util.regex.Pattern;
+import javax.annotation.Nullable;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pinot.common.request.context.ExpressionContext;
 import org.apache.pinot.common.request.context.FilterContext;
@@ -134,7 +135,8 @@ private MutableRoaringBitmap getMatchingFlattenedDocIds(FilterContext filter) {
       case AND: {
         List<FilterContext> children = filter.getChildren();
         int numChildren = children.size();
-        MutableRoaringBitmap matchingDocIds = getMatchingFlattenedDocIds(children.get(0));
+        MutableRoaringBitmap matchingDocIds =
+            getMatchingFlattenedDocIds(children.get(0));
         for (int i = 1; i < numChildren; i++) {
           matchingDocIds.and(getMatchingFlattenedDocIds(children.get(i)));
         }
@@ -143,7 +145,8 @@ private MutableRoaringBitmap getMatchingFlattenedDocIds(FilterContext filter) {
       case OR: {
         List<FilterContext> children = filter.getChildren();
         int numChildren = children.size();
-        MutableRoaringBitmap matchingDocIds = getMatchingFlattenedDocIds(children.get(0));
+        MutableRoaringBitmap matchingDocIds =
+            getMatchingFlattenedDocIds(children.get(0));
         for (int i = 1; i < numChildren; i++) {
           matchingDocIds.or(getMatchingFlattenedDocIds(children.get(i)));
         }
@@ -331,7 +334,40 @@ public void convertFlattenedDocIdsToDocIds(Map<String, RoaringBitmap> valueToFla
   }
 
   @Override
-  public Map<String, RoaringBitmap> getMatchingFlattenedDocsMap(String jsonPathKey) {
+  public Map<String, RoaringBitmap> getMatchingFlattenedDocsMap(String jsonPathKey, @Nullable String filterString) {
+    RoaringBitmap filteredFlattenedDocIds = null;
+    if (filterString != null) {
+      FilterContext filter;
+      try {
+        filter = RequestContextUtils.getFilter(CalciteSqlParser.compileToExpression(filterString));
+        Preconditions.checkArgument(!filter.isConstant());
+      } catch (Exception e) {
+        throw new BadQueryRequestException("Invalid json match filter: " + filterString);
+      }
+      if (filter.getType() == FilterContext.Type.PREDICATE && isExclusive(filter.getPredicate().getType())) {
+        // Handle exclusive predicate separately because the flip can only be applied to the
+        // unflattened doc ids in order to get the correct result, and it cannot be nested
+        filteredFlattenedDocIds = getMatchingFlattenedDocIds(filter.getPredicate()).toRoaringBitmap();
+        filteredFlattenedDocIds.flip(0, _numFlattenedDocs);
+      } else {
+        filteredFlattenedDocIds = getMatchingFlattenedDocIds(filter).toRoaringBitmap();
+      }
+    }
+    // Support 2 formats:
+    // - JSONPath format (e.g. "$.a[1].b"='abc', "$[0]"=1, "$"='abc')
+    // - Legacy format (e.g. "a[1].b"='abc')
+    if (_version == BaseJsonIndexCreator.VERSION_2) {
+      if (jsonPathKey.startsWith("$")) {
+        jsonPathKey = jsonPathKey.substring(1);
+      } else {
+        jsonPathKey = JsonUtils.KEY_SEPARATOR + jsonPathKey;
+      }
+    } else {
+      // For V1 backward-compatibility
+      if (jsonPathKey.startsWith("$.")) {
+        jsonPathKey = jsonPathKey.substring(2);
+      }
+    }
     Map<String, RoaringBitmap> result = new HashMap<>();
     Pair<String, MutableRoaringBitmap> pathKey = getKeyAndFlattenedDocIds(jsonPathKey);
     if (pathKey.getRight() != null && pathKey.getRight().isEmpty()) {
@@ -347,6 +383,10 @@ public Map<String, RoaringBitmap> getMatchingFlattenedDocsMap(String jsonPathKey
     for (int dictId = dictIds[0]; dictId < dictIds[1]; dictId++) {
       String key = _dictionary.getStringValue(dictId);
       RoaringBitmap docIds = _invertedIndex.getDocIds(dictId).toRoaringBitmap();
+      if (filteredFlattenedDocIds != null) {
+        docIds.and(filteredFlattenedDocIds);
+      }
+
       if (arrayIndexFlattenDocIds != null) {
         docIds.and(arrayIndexFlattenDocIds);
       }