[Backport 2.19-dev] Add MAP_REMOVE internal function for Calcite PPL

core/src/main/java/org/opensearch/sql/expression/function/BuiltinFunctionName.java

@@ -62,8 +62,9 @@ public enum BuiltinFunctionName {
   /** Collection functions */
   ARRAY(FunctionName.of("array")),
   ARRAY_LENGTH(FunctionName.of("array_length")),
-  MAP_CONCAT(FunctionName.of("map_concat"), true),
   MAP_APPEND(FunctionName.of("map_append"), true),
+  MAP_CONCAT(FunctionName.of("map_concat"), true),
+  MAP_REMOVE(FunctionName.of("map_remove"), true),
   MVAPPEND(FunctionName.of("mvappend")),
   MVJOIN(FunctionName.of("mvjoin")),
   FORALL(FunctionName.of("forall")),

core/src/main/java/org/opensearch/sql/expression/function/CollectionUDF/MapRemoveFunctionImpl.java

@@ -0,0 +1,101 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.function.CollectionUDF;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.calcite.adapter.enumerable.NotNullImplementor;
+import org.apache.calcite.adapter.enumerable.NullPolicy;
+import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
+import org.apache.calcite.linq4j.tree.Expression;
+import org.apache.calcite.linq4j.tree.Expressions;
+import org.apache.calcite.linq4j.tree.Types;
+import org.apache.calcite.rel.type.RelDataType;
+import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.sql.type.SqlReturnTypeInference;
+import org.opensearch.sql.expression.function.ImplementorUDF;
+import org.opensearch.sql.expression.function.UDFOperandMetadata;
+
+/**
+ * Internal MAP_REMOVE function that removes specified keys from a map. Function signature:
+ * map_remove(map, array_of_keys) -> map Used internally for dynamic fields implementation to dedupe
+ * field names in _MAP.
+ */
+public class MapRemoveFunctionImpl extends ImplementorUDF {
+
+  public MapRemoveFunctionImpl() {
+    super(new MapRemoveImplementor(), NullPolicy.ARG0);
+  }
+
+  @Override
+  public SqlReturnTypeInference getReturnTypeInference() {
+    return sqlOperatorBinding -> {
+      // Return type is the same as the first argument (the map)
+      RelDataType mapType = sqlOperatorBinding.getOperandType(0);
+      return sqlOperatorBinding.getTypeFactory().createTypeWithNullability(mapType, true);
+    };
+  }
+
+  @Override
+  public UDFOperandMetadata getOperandMetadata() {
+    return null;
+  }
+
+  public static class MapRemoveImplementor implements NotNullImplementor {
+    @Override
+    public Expression implement(
+        RexToLixTranslator translator, RexCall call, List<Expression> translatedOperands) {
+      return Expressions.call(
+          Types.lookupMethod(MapRemoveFunctionImpl.class, "mapRemove", Object.class, Object.class),
+          translatedOperands.get(0),
+          translatedOperands.get(1));
+    }
+  }
+
+  /**
+   * Removes specified keys from a map.
+   *
+   * @param mapArg the input map
+   * @param keysArg the array/list of keys to remove
+   * @return a new map with the specified keys removed, or null if input map is null
+   */
+  @SuppressWarnings("unchecked")
+  public static Object mapRemove(Object mapArg, Object keysArg) {
+    if (mapArg == null || keysArg == null) {
+      return mapArg;
+    }
+
+    verifyArgTypes(mapArg, keysArg);
+
+    return mapRemove((Map<String, Object>) mapArg, (List<Object>) keysArg);
+  }
+
+  private static void verifyArgTypes(Object mapArg, Object keysArg) {
+    if (!(mapArg instanceof Map)) {
+      throw new IllegalArgumentException("First argument must be a map, got: " + mapArg.getClass());
+    }
+
+    if (!(keysArg instanceof List)) {
+      throw new IllegalArgumentException(
+          "Second argument must be an array/list, got: " + keysArg.getClass());
+    }
+  }
+
+  private static Map<String, Object> mapRemove(
+      Map<String, Object> originalMap, List<Object> keysToRemove) {
+    Map<String, Object> resultMap = new HashMap<>(originalMap);
+
+    for (Object keyObj : keysToRemove) {
+      if (keyObj != null) {
+        String key = keyObj.toString();
+        resultMap.remove(key);
+      }
+    }
+
+    return resultMap;
+  }
+}

core/src/main/java/org/opensearch/sql/expression/function/PPLBuiltinOperators.java

@@ -48,6 +48,7 @@
 import org.opensearch.sql.expression.function.CollectionUDF.ForallFunctionImpl;
 import org.opensearch.sql.expression.function.CollectionUDF.MVAppendFunctionImpl;
 import org.opensearch.sql.expression.function.CollectionUDF.MapAppendFunctionImpl;
+import org.opensearch.sql.expression.function.CollectionUDF.MapRemoveFunctionImpl;
 import org.opensearch.sql.expression.function.CollectionUDF.ReduceFunctionImpl;
 import org.opensearch.sql.expression.function.CollectionUDF.TransformFunctionImpl;
 import org.opensearch.sql.expression.function.jsonUDF.JsonAppendFunctionImpl;
@@ -388,8 +389,9 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
   public static final SqlOperator FORALL = new ForallFunctionImpl().toUDF("forall");
   public static final SqlOperator EXISTS = new ExistsFunctionImpl().toUDF("exists");
   public static final SqlOperator ARRAY = new ArrayFunctionImpl().toUDF("array");
-  public static final SqlOperator MVAPPEND = new MVAppendFunctionImpl().toUDF("mvappend");
   public static final SqlOperator MAP_APPEND = new MapAppendFunctionImpl().toUDF("map_append");
+  public static final SqlOperator MAP_REMOVE = new MapRemoveFunctionImpl().toUDF("MAP_REMOVE");
+  public static final SqlOperator MVAPPEND = new MVAppendFunctionImpl().toUDF("mvappend");
   public static final SqlOperator FILTER = new FilterFunctionImpl().toUDF("filter");
   public static final SqlOperator TRANSFORM = new TransformFunctionImpl().toUDF("transform");
   public static final SqlOperator REDUCE = new ReduceFunctionImpl().toUDF("reduce");

core/src/main/java/org/opensearch/sql/expression/function/PPLFuncImpTable.java

@@ -126,6 +126,7 @@
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAKETIME;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAP_APPEND;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAP_CONCAT;
+import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAP_REMOVE;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH_BOOL_PREFIX;
 import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH_PHRASE;
@@ -863,8 +864,9 @@ void populate() {
       registerOperator(ARRAY, PPLBuiltinOperators.ARRAY);
       registerOperator(MVAPPEND, PPLBuiltinOperators.MVAPPEND);
       registerOperator(MAP_APPEND, PPLBuiltinOperators.MAP_APPEND);
-      registerOperator(ARRAY_LENGTH, SqlLibraryOperators.ARRAY_LENGTH);
       registerOperator(MAP_CONCAT, SqlLibraryOperators.MAP_CONCAT);
+      registerOperator(MAP_REMOVE, PPLBuiltinOperators.MAP_REMOVE);
+      registerOperator(ARRAY_LENGTH, SqlLibraryOperators.ARRAY_LENGTH);
       registerOperator(FORALL, PPLBuiltinOperators.FORALL);
       registerOperator(EXISTS, PPLBuiltinOperators.EXISTS);
       registerOperator(FILTER, PPLBuiltinOperators.FILTER);

core/src/test/java/org/opensearch/sql/expression/function/CollectionUDF/MapRemoveFunctionImplTest.java

@@ -0,0 +1,205 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.expression.function.CollectionUDF;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.junit.jupiter.api.Test;
+
+public class MapRemoveFunctionImplTest {
+
+  @Test
+  public void testMapRemoveWithNullMap() {
+    Object result = MapRemoveFunctionImpl.mapRemove(null, Arrays.asList("key1", "key2"));
+    assertNull(result);
+  }
+
+  @Test
+  public void testMapRemoveWithNullKeys() {
+    Map<String, Object> map = getBaseMap();
+
+    Object result = MapRemoveFunctionImpl.mapRemove(map, null);
+    assertEquals(map, result);
+  }
+
+  @Test
+  public void testMapRemoveWithInvalidMapArgument() {
+    String notAMap = "not a map";
+    List<String> keysToRemove = Arrays.asList("key1");
+
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class,
+            () -> MapRemoveFunctionImpl.mapRemove(notAMap, keysToRemove));
+
+    assertEquals(
+        "First argument must be a map, got: class java.lang.String", exception.getMessage());
+  }
+
+  @Test
+  public void testMapRemoveWithInvalidKeysArgument() {
+    Map<String, Object> map = getBaseMap();
+    String notAList = "not a list";
+
+    IllegalArgumentException exception =
+        assertThrows(
+            IllegalArgumentException.class, () -> MapRemoveFunctionImpl.mapRemove(map, notAList));
+
+    assertEquals(
+        "Second argument must be an array/list, got: class java.lang.String",
+        exception.getMessage());
+  }
+
+  @Test
+  public void testMapRemoveExistingKeys() {
+    Map<String, Object> map = getBaseMap();
+    List<String> keysToRemove = Arrays.asList("key1", "key3");
+
+    @SuppressWarnings("unchecked")
+    Map<String, Object> result =
+        (Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
+
+    assertEquals(1, result.size());
+    assertEquals("value2", result.get("key2"));
+    assertNull(result.get("key1"));
+    assertNull(result.get("key3"));
+
+    // Verify original map is not modified
+    assertEqualToBaseMap(map);
+  }
+
+  @Test
+  public void testMapRemoveNonExistingKeys() {
+    Map<String, Object> map = getBaseMap();
+    List<String> keysToRemove = Arrays.asList("key4", "key5");
+
+    @SuppressWarnings("unchecked")
+    Map<String, Object> result =
+        (Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
+
+    assertEqualToBaseMap(result);
+  }
+
+  @Test
+  public void testMapRemoveEmptyKeysList() {
+    Map<String, Object> map = getBaseMap();
+    List<String> keysToRemove = Arrays.asList();
+
+    @SuppressWarnings("unchecked")
+    Map<String, Object> result =
+        (Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
+
+    assertEqualToBaseMap(result);
+  }
+
+  @Test
+  public void testMapRemoveMixedExistingAndNonExistingKeys() {
+    Map<String, Object> map = getBaseMap();
+    List<String> keysToRemove = Arrays.asList("key1", "key4", "key2");
+
+    @SuppressWarnings("unchecked")
+    Map<String, Object> result =
+        (Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
+
+    assertEquals(1, result.size());
+    assertEquals("value3", result.get("key3"));
+  }
+
+  @Test
+  public void testMapRemoveWithNullKeysInList() {
+    Map<String, Object> map = getBaseMap();
+    List<Object> keysToRemove = Arrays.asList("key1", null, "key3");
+
+    @SuppressWarnings("unchecked")
+    Map<String, Object> result =
+        (Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
+
+    assertEquals(1, result.size());
+    assertEquals("value2", result.get("key2"));
+  }
+
+  @Test
+  public void testMapRemoveWithDifferentValueTypes() {
+    Map<String, Object> map = new HashMap<>();
+    map.put("string", "value");
+    map.put("number", 42);
+    map.put("boolean", true);
+    map.put("list", Arrays.asList(1, 2, 3));
+    List<String> keysToRemove = Arrays.asList("number", "boolean");
+
+    @SuppressWarnings("unchecked")
+    Map<String, Object> result =
+        (Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
+
+    assertEquals(2, result.size());
+    assertEquals("value", result.get("string"));
+    assertEquals(Arrays.asList(1, 2, 3), result.get("list"));
+    assertNull(result.get("number"));
+    assertNull(result.get("boolean"));
+  }
+
+  @Test
+  public void testMapRemoveAllKeys() {
+    Map<String, Object> map = getBaseMap();
+    List<String> keysToRemove = Arrays.asList("key1", "key2", "key3");
+
+    @SuppressWarnings("unchecked")
+    Map<String, Object> result =
+        (Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
+
+    assertEquals(0, result.size());
+  }
+
+  @Test
+  public void testMapRemoveWithEmptyMap() {
+    Map<String, Object> map = new HashMap<>();
+    List<String> keysToRemove = Arrays.asList("key1", "key2");
+
+    @SuppressWarnings("unchecked")
+    Map<String, Object> result =
+        (Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
+
+    assertEquals(0, result.size());
+  }
+
+  @Test
+  public void testMapRemoveWithNonStringKeys() {
+    Map<String, Object> map = new HashMap<>();
+    map.put("key1", "value1");
+    map.put("key2", "value2");
+    map.put("123", "numeric_key_value");
+
+    List<Object> keysToRemove = Arrays.asList("key1", 123); // 123 will be converted to string "123"
+
+    @SuppressWarnings("unchecked")
+    Map<String, Object> result =
+        (Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
+
+    assertEquals(1, result.size());
+    assertEquals("value2", result.get("key2"));
+  }
+
+  private Map<String, Object> getBaseMap() {
+    Map<String, Object> map = new HashMap<>();
+    map.put("key1", "value1");
+    map.put("key2", "value2");
+    map.put("key3", "value3");
+    return map;
+  }
+
+  private void assertEqualToBaseMap(Map<String, Object> map) {
+    assertEquals(3, map.size());
+    assertEquals("value1", map.get("key1"));
+    assertEquals("value2", map.get("key2"));
+    assertEquals("value3", map.get("key3"));
+  }
+}