Partition partial results on the group key's source fields

Resolve each aggregation group key through the eval Project to the scan
fields it reads, so an expression key (e.g. eval g = lower(city) | stats
count() by g) gets partial results over the keyword subset just like a
bare 'by city'. Previously only a bare group field matched the per-index
mapping; a derived key looked up its output alias, found nothing, and
bailed to the complete (script) path. A constant group key resolves to no
field and cleanly bails.

Signed-off-by: Kai Huang <ahkcs@amazon.com>

Author: ahkcs

integ-test/src/test/java/org/opensearch/sql/calcite/remote/CalcitePartialResultOnMappingConflictIT.java

@@ -248,6 +248,30 @@ public void partialResultOnHandlesNestedDottedField() throws IOException {
         warning.getString("detail").contains(NESTED_TEXT_INDEX));
   }
 
+  @Test
+  public void partialResultOnHandlesEvalDerivedGroupKey() throws IOException {
+    setPartialResult(true);
+    setPitContextLimit("1");
+    // The group key is an expression over the conflicting field (upper(env)), not the bare field.
+    // Partitioning traces it back to env, so the keyword index is kept and the text index excluded
+    // just as for a bare group key. Only the keyword index contributes: PROD=2, DEV=1.
+    JSONObject result =
+        executeQuery(
+            String.format(
+                "source=%s | eval g = upper(env) | stats count() by g | sort g", PATTERN));
+    verifyDataRows(result, rows(1, "DEV"), rows(2, "PROD"));
+
+    assertTrue("response should carry a warnings array", result.has("warnings"));
+    JSONObject warning = result.getJSONArray("warnings").getJSONObject(0);
+    assertEquals("PARTIAL_RESULT", warning.getString("type"));
+    assertTrue(
+        "warning should name the underlying field the expression reads",
+        warning.getString("detail").contains("env"));
+    assertTrue(
+        "warning should name the excluded text index",
+        warning.getString("detail").contains(TEXT_INDEX));
+  }
+
   @Test
   public void partialResultKeepsKeywordGroupEvenWhenOutnumbered() throws IOException {
     setPartialResult(true);

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/CalciteLogicalIndexScan.java

@@ -7,9 +7,11 @@
 
 import com.google.common.collect.ImmutableList;
 import java.util.ArrayList;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
 import lombok.Getter;
@@ -34,7 +36,9 @@
 import org.apache.calcite.rel.type.RelDataTypeFactory;
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexBuilder;
+import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.rex.RexVisitorImpl;
 import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.commons.lang3.tuple.Pair;
@@ -414,9 +418,12 @@ private AbstractRelNode pushDownAggregate(
       // Try partial mode before analyze: since #5646 a text/keyword conflict pushes down as a slow
       // _source script instead of failing, so a post-failure fallback would never fire.
       if (allowPartialFallback) {
-        AbstractRelNode partial = tryPartialResultAggregate(aggregate, project, bucketNames);
-        if (partial != null) {
-          return partial;
+        List<String> partitionFields = resolvePartitionFields(aggregate, project);
+        if (partitionFields != null) {
+          AbstractRelNode partial = tryPartialResultAggregate(aggregate, project, partitionFields);
+          if (partial != null) {
+            return partial;
+          }
         }
       }
       int queryBucketSize = osIndex.getQueryBucketSize();
@@ -451,15 +458,58 @@ private AbstractRelNode pushDownAggregate(
     return null;
   }
 
+  /**
+   * Resolve the aggregation's group keys to the underlying scan fields to partition indices on. A
+   * key may be a bare field ({@code ... by city}) or an expression over fields ({@code eval g =
+   * lower(city) | ... by g}); in the latter case we partition on every field the expression reads,
+   * since a kept index must map all of them aggregatably. The {@code project} (when present) sits
+   * directly on the scan, so its input refs index into this scan's row type. Returns {@code null}
+   * if any key is a pure constant with no field to key on.
+   */
+  @Nullable
+  private List<String> resolvePartitionFields(Aggregate aggregate, @Nullable Project project) {
+    List<String> scanFields = getRowType().getFieldNames();
+    List<String> fields = new ArrayList<>();
+    for (int group : aggregate.getGroupSet()) {
+      Set<Integer> refs = new LinkedHashSet<>();
+      if (project == null) {
+        refs.add(group); // group key indexes directly into the scan
+      } else {
+        project
+            .getProjects()
+            .get(group)
+            .accept(
+                new RexVisitorImpl<Void>(true) {
+                  @Override
+                  public Void visitInputRef(RexInputRef ref) {
+                    refs.add(ref.getIndex());
+                    return null;
+                  }
+                });
+      }
+      if (refs.isEmpty()) {
+        return null; // constant group key -> nothing to partition on
+      }
+      for (int ref : refs) {
+        String name = scanFields.get(ref);
+        if (!fields.contains(name)) {
+          fields.add(name);
+        }
+      }
+    }
+    return fields;
+  }
+
   /**
    * On a text/keyword mapping conflict, narrow the scan to the index subset where the group field
    * is aggregatable, push the aggregation over just that subset, and record a warning naming the
    * excluded indices. Only runs behind the opt-in setting and only when the response format can
    * carry the warning ({@link QueryContext#isWarningsSupported}); returns {@code null} otherwise.
-   * Partitioning lives in {@link PartialResultAggregatePushdown}.
+   * {@code partitionFields} are the scan fields the group keys resolve to (see {@link
+   * #resolvePartitionFields}). Partitioning lives in {@link PartialResultAggregatePushdown}.
    */
   private AbstractRelNode tryPartialResultAggregate(
-      Aggregate aggregate, @Nullable Project project, List<String> bucketNames) {
+      Aggregate aggregate, @Nullable Project project, List<String> partitionFields) {
     if (!QueryContext.isPartialResultEnabled(osIndex.getSettings())) {
       return null;
     }
@@ -470,7 +520,7 @@ private AbstractRelNode tryPartialResultAggregate(
     try {
       Map<String, IndexMapping> mappings = osIndex.getIndexMappings();
       PartialResultAggregatePushdown.Plan plan =
-          PartialResultAggregatePushdown.plan(bucketNames, mappings);
+          PartialResultAggregatePushdown.plan(partitionFields, mappings);
       if (plan == null) {
         return null;
       }

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/scan/PartialResultAggregatePushdown.java

@@ -51,7 +51,8 @@ record Plan(List<String> keptIndices, List<String> excludedIndices, Warning warn
   /**
    * Decide the partial-result plan for a group key over a set of per-index mappings.
    *
-   * @param bucketNames the aggregation's group-by field names (dotted paths)
+   * @param bucketNames the storage fields the group keys resolve to (dotted paths); an expression
+   *     key like {@code lower(city)} resolves to the field(s) it reads, e.g. {@code city}
    * @param mappings per-index field mappings, keyed by concrete index name (from {@code
    *     getIndexMappings}); the wildcard has already been resolved to concrete indices
    * @return a plan naming the kept and excluded indices plus the warning, or {@code null} when