Reuse the already-fetched index mappings, and stop the merge mutating them

Partial-result partitioning needs per-index mappings, which the merged field
types cached on OpenSearchIndex discard, so it was fetching them a second time.

Retain the per-index mappings on the describe request that already fetches them
and cache them alongside the merged types, so partitioning reuses that result.

The first attempt at this was reverted because MergeRuleHelper rewrites the
accumulated type's nested properties in place, mutating the very mappings being
retained: a nested text/keyword conflict then read back as no conflict, produced
no partitioning plan, and fell through to the PIT-exhausting scan. Merge deep
copies instead, via a new OpenSearchDataType.cloneDeep() that carries the nested
properties subtree (cloneEmpty drops it).

Covered by a regression test that fails without the copy. Stress-verified:
CalcitePartialResultOnMappingConflictIT passes 8/8 (it failed 4/5 before).

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

Author: ahkcs

opensearch/src/main/java/org/opensearch/sql/opensearch/data/type/OpenSearchDataType.java

@@ -264,6 +264,23 @@ protected OpenSearchDataType cloneEmpty() {
         : new OpenSearchDataType(this.mappingType);
   }
 
+  /**
+   * Clone this type including its nested {@link #properties} subtree, so the copy shares no mutable
+   * state with the original. Needed by callers that must keep a mapping intact across an in-place
+   * merge (see {@code MergeRuleHelper}), which rewrites the target's {@code properties}.
+   *
+   * @return A deep copy of this type.
+   */
+  public OpenSearchDataType cloneDeep() {
+    OpenSearchDataType copy = cloneEmpty();
+    if (!properties.isEmpty()) {
+      Map<String, OpenSearchDataType> copiedProperties = new LinkedHashMap<>();
+      properties.forEach((field, type) -> copiedProperties.put(field, type.cloneDeep()));
+      copy.properties = copiedProperties;
+    }
+    return copy;
+  }
+
   /**
    * Flattens mapping tree into a single layer list of objects (pairs of name-types actually), which
    * don't have nested types. See {@link OpenSearchDataTypeTest#traverseAndFlatten() test} for

opensearch/src/main/java/org/opensearch/sql/opensearch/request/system/OpenSearchDescribeIndexRequest.java

@@ -15,6 +15,7 @@
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import lombok.Getter;
 import lombok.extern.log4j.Log4j2;
 import org.opensearch.sql.data.model.ExprTupleValue;
 import org.opensearch.sql.data.model.ExprValue;
@@ -92,6 +93,14 @@ public List<ExprValue> search() {
     return results;
   }
 
+  /**
+   * The per-index mappings behind the last {@link #getFieldTypes()} call, keyed by concrete index
+   * name. Retained because merging discards which index mapped a field which way, and callers that
+   * need that detail (e.g. partitioning a wildcard by whether a field is aggregatable) would
+   * otherwise have to fetch the mappings a second time.
+   */
+  @Getter private Map<String, IndexMapping> lastIndexMappings = Map.of();
+
   /**
    * Get the mapping of field and type.
    *
@@ -102,18 +111,30 @@ public Map<String, OpenSearchDataType> getFieldTypes() {
     Map<String, OpenSearchDataType> fieldTypes = new HashMap<>();
     Map<String, IndexMapping> indexMappings =
         client.getIndexMappings(getLocalIndexNames(indexName.getIndexNames()));
+    this.lastIndexMappings = indexMappings;
     if (indexMappings.size() <= 1) {
       for (IndexMapping indexMapping : indexMappings.values()) {
         fieldTypes.putAll(indexMapping.getFieldMappings());
       }
     } else {
+      // Merge deep copies: MergeRuleHelper rewrites the accumulated type's nested `properties` in
+      // place, which would otherwise mutate the per-index mappings retained above (they are reused
+      // by partial-result partitioning, which needs to see each index's original mapping).
       for (IndexMapping indexMapping : indexMappings.values()) {
-        MergeRuleHelper.merge(fieldTypes, indexMapping.getFieldMappings());
+        MergeRuleHelper.merge(fieldTypes, deepCopy(indexMapping.getFieldMappings()));
       }
     }
     return fieldTypes;
   }
 
+  /** Copy a field-mapping map so an in-place merge cannot mutate the source types. */
+  private static Map<String, OpenSearchDataType> deepCopy(
+      Map<String, OpenSearchDataType> mappings) {
+    Map<String, OpenSearchDataType> copy = new LinkedHashMap<>();
+    mappings.forEach((field, type) -> copy.put(field, type.cloneDeep()));
+    return copy;
+  }
+
   /**
    * Get the minimum of the max result windows of the indices.
    *

opensearch/src/main/java/org/opensearch/sql/opensearch/storage/OpenSearchIndex.java

@@ -29,6 +29,7 @@
 import org.opensearch.sql.opensearch.client.OpenSearchClient;
 import org.opensearch.sql.opensearch.data.type.OpenSearchDataType;
 import org.opensearch.sql.opensearch.data.value.OpenSearchExprValueFactory;
+import org.opensearch.sql.opensearch.mapping.IndexMapping;
 import org.opensearch.sql.opensearch.monitor.OpenSearchMemoryHealthy;
 import org.opensearch.sql.opensearch.monitor.OpenSearchResourceMonitor;
 import org.opensearch.sql.opensearch.planner.physical.ADOperator;
@@ -91,6 +92,13 @@ public class OpenSearchIndex extends AbstractOpenSearchTable {
   /** The cached mapping of alias type field to its original path. */
   private Map<String, String> aliasMapping = null;
 
+  /**
+   * The cached per-index field mappings, keyed by concrete index name. Populated as a by-product of
+   * resolving {@link #cachedFieldOpenSearchTypes}, since merging those types discards which index
+   * mapped a field which way.
+   */
+  private Map<String, IndexMapping> cachedIndexMappings = null;
+
   /** The cached max result window setting of index. */
   private Integer cachedMaxResultWindow = null;
 
@@ -137,10 +145,7 @@ public void create(Map<String, ExprType> schema) {
    */
   @Override
   public Map<String, ExprType> getFieldTypes() {
-    if (cachedFieldOpenSearchTypes == null) {
-      cachedFieldOpenSearchTypes =
-          new OpenSearchDescribeIndexRequest(client, indexName).getFieldTypes();
-    }
+    resolveFieldOpenSearchTypes();
     if (cachedFieldTypes == null) {
       cachedFieldTypes =
           OpenSearchDataType.traverseAndFlatten(cachedFieldOpenSearchTypes).entrySet().stream()
@@ -163,10 +168,7 @@ public Map<String, ExprType> getAllFieldTypes() {
   }
 
   public Map<String, String> getAliasMapping() {
-    if (cachedFieldOpenSearchTypes == null) {
-      cachedFieldOpenSearchTypes =
-          new OpenSearchDescribeIndexRequest(client, indexName).getFieldTypes();
-    }
+    resolveFieldOpenSearchTypes();
     if (aliasMapping == null) {
       aliasMapping =
           OpenSearchDataType.traverseAndFlatten(cachedFieldOpenSearchTypes).entrySet().stream()
@@ -184,11 +186,29 @@ public Map<String, String> getAliasMapping() {
    * @return A complete map between field names and their types.
    */
   public Map<String, OpenSearchDataType> getFieldOpenSearchTypes() {
+    resolveFieldOpenSearchTypes();
+    return cachedFieldOpenSearchTypes;
+  }
+
+  /**
+   * The per-index field mappings behind this index's merged types, keyed by concrete index name
+   * (the wildcard, if any, is already resolved). Needed by callers that must know which index
+   * mapped a field which way -- the merged view in {@link #getFieldOpenSearchTypes()} discards
+   * that. Shares the mapping fetch with the merged types, so this costs no extra round trip.
+   */
+  public Map<String, IndexMapping> getIndexMappings() {
+    resolveFieldOpenSearchTypes();
+    return cachedIndexMappings;
+  }
+
+  /** Fetch and cache the merged field types, retaining the per-index mappings behind them. */
+  private void resolveFieldOpenSearchTypes() {
     if (cachedFieldOpenSearchTypes == null) {
-      cachedFieldOpenSearchTypes =
-          new OpenSearchDescribeIndexRequest(client, indexName).getFieldTypes();
+      OpenSearchDescribeIndexRequest request =
+          new OpenSearchDescribeIndexRequest(client, indexName);
+      cachedFieldOpenSearchTypes = request.getFieldTypes();
+      cachedIndexMappings = request.getLastIndexMappings();
     }
-    return cachedFieldOpenSearchTypes;
   }
 
   /** Get the max result window setting of the table. */

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

@@ -471,11 +471,9 @@ private AbstractRelNode tryPartialResultAggregate(
     try {
       List<String> outputFields = aggregate.getRowType().getFieldNames();
       List<String> bucketNames = outputFields.subList(0, aggregate.getGroupSet().cardinality());
-      // getIndexNames() returns the raw pattern (e.g. ["logs-*"]) for a wildcard, not the resolved
-      // concrete indices; getIndexMappings expands it server-side to a map keyed by concrete index
-      // name, which is what the partitioning classifies.
-      Map<String, IndexMapping> mappings =
-          osIndex.getClient().getIndexMappings(osIndex.getIndexName().getIndexNames());
+      // Per-index mappings keyed by concrete index name (the wildcard is already resolved), reused
+      // from the fetch that resolved this index's field types rather than re-requesting them.
+      Map<String, IndexMapping> mappings = osIndex.getIndexMappings();
       PartialResultAggregatePushdown.Plan plan =
           PartialResultAggregatePushdown.plan(bucketNames, mappings);
       if (plan == null) {

opensearch/src/test/java/org/opensearch/sql/opensearch/request/system/OpenSearchDescribeIndexRequestTest.java

@@ -35,6 +35,48 @@ class OpenSearchDescribeIndexRequestTest {
 
   @Mock private IndexMapping mapping2;
 
+  /**
+   * Merging must not mutate the per-index mappings it reads. {@code MergeRuleHelper} rewrites the
+   * accumulated type's nested {@code properties} in place, so without copying, the first index's
+   * nested field would be merged into the second's -- and callers of {@link
+   * OpenSearchDescribeIndexRequest#getLastIndexMappings()} (partial-result partitioning) would see
+   * a text/keyword conflict as no conflict at all.
+   */
+  @Test
+  void getFieldTypesLeavesRetainedPerIndexMappingsIntact() {
+    Map<String, OpenSearchDataType> keywordSide =
+        Map.of("attrs", nestedObject("env", OpenSearchDataType.MappingType.Keyword));
+    Map<String, OpenSearchDataType> textSide =
+        Map.of("attrs", nestedObject("env", OpenSearchDataType.MappingType.Text));
+    when(mapping.getFieldMappings()).thenReturn(keywordSide);
+    when(mapping2.getFieldMappings()).thenReturn(textSide);
+    when(client.getIndexMappings("idx-*"))
+        .thenReturn(ImmutableMap.of("idx-keyword", mapping, "idx-text", mapping2));
+
+    OpenSearchDescribeIndexRequest request = new OpenSearchDescribeIndexRequest(client, "idx-*");
+    request.getFieldTypes();
+
+    // Each index must still report the type it actually declared.
+    assertEquals(
+        OpenSearchDataType.MappingType.Keyword,
+        nestedFieldType(request.getLastIndexMappings().get("idx-keyword"), "attrs", "env"));
+    assertEquals(
+        OpenSearchDataType.MappingType.Text,
+        nestedFieldType(request.getLastIndexMappings().get("idx-text"), "attrs", "env"));
+  }
+
+  private static OpenSearchDataType nestedObject(
+      String innerField, OpenSearchDataType.MappingType innerType) {
+    return OpenSearchDataType.of(
+        OpenSearchDataType.MappingType.Object,
+        Map.of("properties", Map.of(innerField, Map.of("type", innerType.toString()))));
+  }
+
+  private static OpenSearchDataType.MappingType nestedFieldType(
+      IndexMapping indexMapping, String parent, String child) {
+    return indexMapping.getFieldMappings().get(parent).getProperties().get(child).getMappingType();
+  }
+
   @Test
   void testSearch() {
     when(mapping.getFieldMappings())