Refine index-based DISTINCT operators (JSON / inverted)

pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java

@@ -192,6 +192,13 @@ public static Double getInvertedIndexDistinctCostRatio(Map<String, String> query
         queryOptions.get(QueryOptionKey.INVERTED_INDEX_DISTINCT_COST_RATIO));
   }
 
+  /// When true, [org.apache.pinot.core.operator.query.JsonIndexDistinctOperator] skips its missing-path handling —
+  /// does not add a 4-arg default, does not add null, and does not throw `Illegal Json Path`. The distinct set is
+  /// purely the values returned by the JSON-index lookup.
+  public static boolean isJsonIndexDistinctSkipMissingPath(Map<String, String> queryOptions) {
+    return Boolean.parseBoolean(queryOptions.get(QueryOptionKey.JSON_INDEX_DISTINCT_SKIP_MISSING_PATH));
+  }
+
   public static boolean isSkipScanFilterReorder(Map<String, String> queryOptions) {
     return "false".equalsIgnoreCase(queryOptions.get(QueryOptionKey.USE_SCAN_REORDER_OPTIMIZATION));
   }

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

@@ -24,10 +24,12 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import javax.annotation.Nullable;
 import org.apache.pinot.common.function.JsonPathCache;
 import org.apache.pinot.core.operator.ColumnContext;
 import org.apache.pinot.core.operator.blocks.ValueBlock;
 import org.apache.pinot.core.operator.transform.TransformResultMetadata;
+import org.apache.pinot.segment.spi.datasource.DataSource;
 import org.apache.pinot.segment.spi.index.IndexService;
 import org.apache.pinot.segment.spi.index.IndexType;
 import org.apache.pinot.segment.spi.index.reader.JsonIndexReader;
@@ -45,14 +47,15 @@
 public class JsonExtractIndexTransformFunction extends BaseTransformFunction {
   public static final String FUNCTION_NAME = "jsonExtractIndex";
 
-  private TransformFunction _jsonFieldTransformFunction;
-  private String _jsonPathString;
-  private TransformResultMetadata _resultMetadata;
   private JsonIndexReader _jsonIndexReader;
+  private String _jsonPathString;
+  private boolean _isSingleValue;
+  @Nullable
   private Object _defaultValue;
+  @Nullable
+  private String _filterJsonExpression;
+  private TransformResultMetadata _resultMetadata;
   private Map<String, RoaringBitmap> _valueToMatchingDocsMap;
-  private boolean _isSingleValue;
-  private String _filterJsonPath;
 
   @Override
   public String getName() {
@@ -62,23 +65,25 @@ public String getName() {
   @Override
   public void init(List<TransformFunction> arguments, Map<String, ColumnContext> columnContextMap) {
     super.init(arguments, columnContextMap);
+
+    int numArguments = arguments.size();
     // Check that there are exactly 3 or 4 or 5 arguments
-    if (arguments.size() < 3 || arguments.size() > 5) {
+    if (numArguments < 3 || numArguments > 5) {
       throw new IllegalArgumentException(
           "Expected 3/4/5 arguments for transform function: jsonExtractIndex(jsonFieldName, 'jsonPath', 'resultsType',"
               + " ['defaultValue'], ['jsonFilterExpression'])");
     }
 
     TransformFunction firstArgument = arguments.get(0);
     if (firstArgument instanceof IdentifierTransformFunction) {
-      String columnName = ((IdentifierTransformFunction) firstArgument).getColumnName();
-      _jsonIndexReader = columnContextMap.get(columnName).getDataSource().getJsonIndex();
-      if (_jsonIndexReader == null) { //TODO: rework
-        Optional<IndexType<?, ?, ?>> compositeIndex =
-            IndexService.getInstance().getOptional("composite_json_index");
+      DataSource dataSource =
+          columnContextMap.get(((IdentifierTransformFunction) firstArgument).getColumnName()).getDataSource();
+      _jsonIndexReader = dataSource.getJsonIndex();
+      // TODO: rework
+      if (_jsonIndexReader == null) {
+        Optional<IndexType<?, ?, ?>> compositeIndex = IndexService.getInstance().getOptional("composite_json_index");
         if (compositeIndex.isPresent()) {
-          _jsonIndexReader = (JsonIndexReader) columnContextMap.get(columnName)
-              .getDataSource().getIndex(compositeIndex.get());
+          _jsonIndexReader = (JsonIndexReader) dataSource.getIndex(compositeIndex.get());
         }
       }
       if (_jsonIndexReader == null) {
@@ -87,7 +92,6 @@ public void init(List<TransformFunction> arguments, Map<String, ColumnContext> c
     } else {
       throw new IllegalArgumentException("jsonExtractIndex can only be applied to a raw column");
     }
-    _jsonFieldTransformFunction = firstArgument;
 
     TransformFunction secondArgument = arguments.get(1);
     if (!(secondArgument instanceof LiteralTransformFunction)) {
@@ -113,12 +117,11 @@ public void init(List<TransformFunction> arguments, Map<String, ColumnContext> c
     DataType dataType = _isSingleValue ? DataType.valueOf(resultsType)
         : DataType.valueOf(resultsType.substring(0, resultsType.length() - 6));
 
-    if (arguments.size() >= 4) {
+    if (numArguments >= 4) {
       TransformFunction fourthArgument = arguments.get(3);
       if (!(fourthArgument instanceof LiteralTransformFunction)) {
         throw new IllegalArgumentException("Default value must be a literal");
       }
-
       if (_isSingleValue) {
         _defaultValue = dataType.convert(((LiteralTransformFunction) fourthArgument).getStringLiteral());
       } else {
@@ -138,12 +141,12 @@ public void init(List<TransformFunction> arguments, Map<String, ColumnContext> c
       }
     }
 
-    if (arguments.size() == 5) {
+    if (numArguments == 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();
+      _filterJsonExpression = ((LiteralTransformFunction) fifthArgument).getStringLiteral();
     }
 
     _resultMetadata = new TransformResultMetadata(dataType, _isSingleValue, false);
@@ -425,7 +428,7 @@ public String[][] transformToStringValuesMV(ValueBlock valueBlock) {
    */
   private Map<String, RoaringBitmap> getValueToMatchingDocsMap() {
     if (_valueToMatchingDocsMap == null) {
-      _valueToMatchingDocsMap = _jsonIndexReader.getMatchingFlattenedDocsMap(_jsonPathString, _filterJsonPath);
+      _valueToMatchingDocsMap = _jsonIndexReader.getMatchingFlattenedDocsMap(_jsonPathString, _filterJsonExpression);
       if (_isSingleValue) {
         // For single value result type, it's more efficient to use original docIDs map
         _jsonIndexReader.convertFlattenedDocIdsToDocIds(_valueToMatchingDocsMap);

pinot-core/src/main/java/org/apache/pinot/core/plan/DistinctPlanNode.java

@@ -84,9 +84,9 @@ public Operator<DistinctResultsBlock> run() {
       ExpressionContext expr = expressions.get(0);
 
       // JSON index path
-      if (JsonIndexDistinctOperator.canUseJsonIndexDistinct(_indexSegment, expr)) {
+      if (JsonIndexDistinctOperator.canUseJsonIndexDistinct(expr)) {
         BaseFilterOperator filterOperator = new FilterPlanNode(_segmentContext, _queryContext).run();
-        return new JsonIndexDistinctOperator(_indexSegment, _segmentContext, _queryContext, filterOperator);
+        return new JsonIndexDistinctOperator(_indexSegment, _queryContext, filterOperator);
       }
 
       // Inverted/sorted index path. For unsorted dictionaries the operator still avoids the scan/projection path,