tarun11Mavani commented on code in PR #19040:
URL: https://github.com/apache/pinot/pull/19040#discussion_r3664097757


##########
pinot-core/src/main/java/org/apache/pinot/core/plan/AggregationPlanNode.java:
##########
@@ -204,6 +221,50 @@ private boolean isFitForNonScanBasedPlan() {
     return true;
   }
 
+  @Nullable
+  private DataSource resolveDataSource(ExpressionContext expression) {
+    return resolveDataSource(expression, _indexSegment, 
_queryContext.getSchema());
+  }
+
+  @Nullable
+  static DataSource resolveDataSource(ExpressionContext expression, 
IndexSegment segment,
+      @Nullable org.apache.pinot.spi.data.Schema schema) {
+    if (expression.getType() == ExpressionContext.Type.IDENTIFIER) {
+      return segment.getDataSource(expression.getIdentifier(), schema);
+    }
+    if (expression.getType() == ExpressionContext.Type.FUNCTION) {
+      return tryResolveKeyedDataSource(expression, segment, schema);
+    }
+    return null;
+  }
+
+  @Nullable
+  static DataSource tryResolveKeyedDataSource(ExpressionContext expression, 
IndexSegment segment,
+      @Nullable org.apache.pinot.spi.data.Schema schema) {
+    FunctionContext function = expression.getFunction();
+    if (function == null
+        || 
!ItemTransformFunction.FUNCTION_NAME.equals(function.getFunctionName())) {
+      return null;
+    }
+    List<ExpressionContext> args = function.getArguments();
+    if (args.size() != 2
+        || args.get(0).getType() != ExpressionContext.Type.IDENTIFIER
+        || args.get(1).getType() != ExpressionContext.Type.LITERAL) {
+      return null;
+    }
+    String columnName = args.get(0).getIdentifier();
+    String key = args.get(1).getLiteral().getStringValue();
+    DataSource columnDs = segment.getDataSource(columnName, schema);
+    if (columnDs instanceof MapDataSource) {
+      return ((MapDataSource) columnDs).getDataSource(key);

Review Comment:
   Checked both concerns. Neither can bite today, but I've added the guard.
    
   (1) **COUNT** never reaches the per-key DataSource — 
`isFitForNonScanBasedPlan` skips COUNT before resolving it, and the operator 
counts from segment metadata. So `NullDataSource`'s `numDocs=0` is never read.
    
   (2) **MIN/MAX** is the real hole, and it's min/max rather than the null 
vector: `NullDataSourceMetadata` returns `Integer.MIN_VALUE` for both 
(`NullDataSource:181-189`), so an absent MAP key would pass the metadata check 
and `MIN(mapCol['absentKey'])` would return `-2147483648` instead of NULL.
    
   That path is unreachable on master, though. `BaseMapDataSource` only returns 
a `NullDataSource` when `getMapIndexReader().getIndexes(key)` is null, and the 
only `MapIndexReader` implementation returns a non-null map unconditionally. 
Absent keys land on `ImmutableDataSource` instead, whose min/max are null, so 
they fail the metadata check and fall back to a scan.
    
   Added the guard regardless — `return keyDs instanceof NullDataSource ? null 
: keyDs`, mirroring the OPEN_STRUCT branch. `ItemTransformFunction` already 
null-checks the same call, so the nullable contract is clearly intended; this 
just makes the aggregation path honour it too. 
    



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to