Skip to content

Commit 8680ba1

Browse files
opensearch-trigger-bot[bot]github-actions[bot]
opensearch-trigger-bot[bot]
and
github-actions[bot]
authored
Add internal MAP_REMOVE function for Calcite PPL (opensearch-project#4511) (opensearch-project#4585)
(cherry picked from commit e3ab9d0) Signed-off-by: Tomoyuki Morita <moritato@amazon.com> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent e5bc2cd commit 8680ba1

6 files changed

Lines changed: 495 additions & 3 deletions

File tree

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,8 +62,9 @@ public enum BuiltinFunctionName {
6262
/** Collection functions */
6363
ARRAY(FunctionName.of("array")),
6464
ARRAY_LENGTH(FunctionName.of("array_length")),
65-
MAP_CONCAT(FunctionName.of("map_concat"), true),
6665
MAP_APPEND(FunctionName.of("map_append"), true),
66+
MAP_CONCAT(FunctionName.of("map_concat"), true),
67+
MAP_REMOVE(FunctionName.of("map_remove"), true),
6768
MVAPPEND(FunctionName.of("mvappend")),
6869
MVJOIN(FunctionName.of("mvjoin")),
6970
FORALL(FunctionName.of("forall")),

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

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.expression.function.CollectionUDF;
7+
8+
import java.util.HashMap;
9+
import java.util.List;
10+
import java.util.Map;
11+
import org.apache.calcite.adapter.enumerable.NotNullImplementor;
12+
import org.apache.calcite.adapter.enumerable.NullPolicy;
13+
import org.apache.calcite.adapter.enumerable.RexToLixTranslator;
14+
import org.apache.calcite.linq4j.tree.Expression;
15+
import org.apache.calcite.linq4j.tree.Expressions;
16+
import org.apache.calcite.linq4j.tree.Types;
17+
import org.apache.calcite.rel.type.RelDataType;
18+
import org.apache.calcite.rex.RexCall;
19+
import org.apache.calcite.sql.type.SqlReturnTypeInference;
20+
import org.opensearch.sql.expression.function.ImplementorUDF;
21+
import org.opensearch.sql.expression.function.UDFOperandMetadata;
22+
23+
/**
24+
* Internal MAP_REMOVE function that removes specified keys from a map. Function signature:
25+
* map_remove(map, array_of_keys) -> map Used internally for dynamic fields implementation to dedupe
26+
* field names in _MAP.
27+
*/
28+
public class MapRemoveFunctionImpl extends ImplementorUDF {
29+
30+
public MapRemoveFunctionImpl() {
31+
super(new MapRemoveImplementor(), NullPolicy.ARG0);
32+
}
33+
34+
@Override
35+
public SqlReturnTypeInference getReturnTypeInference() {
36+
return sqlOperatorBinding -> {
37+
// Return type is the same as the first argument (the map)
38+
RelDataType mapType = sqlOperatorBinding.getOperandType(0);
39+
return sqlOperatorBinding.getTypeFactory().createTypeWithNullability(mapType, true);
40+
};
41+
}
42+
43+
@Override
44+
public UDFOperandMetadata getOperandMetadata() {
45+
return null;
46+
}
47+
48+
public static class MapRemoveImplementor implements NotNullImplementor {
49+
@Override
50+
public Expression implement(
51+
RexToLixTranslator translator, RexCall call, List<Expression> translatedOperands) {
52+
return Expressions.call(
53+
Types.lookupMethod(MapRemoveFunctionImpl.class, "mapRemove", Object.class, Object.class),
54+
translatedOperands.get(0),
55+
translatedOperands.get(1));
56+
}
57+
}
58+
59+
/**
60+
* Removes specified keys from a map.
61+
*
62+
* @param mapArg the input map
63+
* @param keysArg the array/list of keys to remove
64+
* @return a new map with the specified keys removed, or null if input map is null
65+
*/
66+
@SuppressWarnings("unchecked")
67+
public static Object mapRemove(Object mapArg, Object keysArg) {
68+
if (mapArg == null || keysArg == null) {
69+
return mapArg;
70+
}
71+
72+
verifyArgTypes(mapArg, keysArg);
73+
74+
return mapRemove((Map<String, Object>) mapArg, (List<Object>) keysArg);
75+
}
76+
77+
private static void verifyArgTypes(Object mapArg, Object keysArg) {
78+
if (!(mapArg instanceof Map)) {
79+
throw new IllegalArgumentException("First argument must be a map, got: " + mapArg.getClass());
80+
}
81+
82+
if (!(keysArg instanceof List)) {
83+
throw new IllegalArgumentException(
84+
"Second argument must be an array/list, got: " + keysArg.getClass());
85+
}
86+
}
87+
88+
private static Map<String, Object> mapRemove(
89+
Map<String, Object> originalMap, List<Object> keysToRemove) {
90+
Map<String, Object> resultMap = new HashMap<>(originalMap);
91+
92+
for (Object keyObj : keysToRemove) {
93+
if (keyObj != null) {
94+
String key = keyObj.toString();
95+
resultMap.remove(key);
96+
}
97+
}
98+
99+
return resultMap;
100+
}
101+
}

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@
4848
import org.opensearch.sql.expression.function.CollectionUDF.ForallFunctionImpl;
4949
import org.opensearch.sql.expression.function.CollectionUDF.MVAppendFunctionImpl;
5050
import org.opensearch.sql.expression.function.CollectionUDF.MapAppendFunctionImpl;
51+
import org.opensearch.sql.expression.function.CollectionUDF.MapRemoveFunctionImpl;
5152
import org.opensearch.sql.expression.function.CollectionUDF.ReduceFunctionImpl;
5253
import org.opensearch.sql.expression.function.CollectionUDF.TransformFunctionImpl;
5354
import org.opensearch.sql.expression.function.jsonUDF.JsonAppendFunctionImpl;
@@ -388,8 +389,9 @@ public class PPLBuiltinOperators extends ReflectiveSqlOperatorTable {
388389
public static final SqlOperator FORALL = new ForallFunctionImpl().toUDF("forall");
389390
public static final SqlOperator EXISTS = new ExistsFunctionImpl().toUDF("exists");
390391
public static final SqlOperator ARRAY = new ArrayFunctionImpl().toUDF("array");
391-
public static final SqlOperator MVAPPEND = new MVAppendFunctionImpl().toUDF("mvappend");
392392
public static final SqlOperator MAP_APPEND = new MapAppendFunctionImpl().toUDF("map_append");
393+
public static final SqlOperator MAP_REMOVE = new MapRemoveFunctionImpl().toUDF("MAP_REMOVE");
394+
public static final SqlOperator MVAPPEND = new MVAppendFunctionImpl().toUDF("mvappend");
393395
public static final SqlOperator FILTER = new FilterFunctionImpl().toUDF("filter");
394396
public static final SqlOperator TRANSFORM = new TransformFunctionImpl().toUDF("transform");
395397
public static final SqlOperator REDUCE = new ReduceFunctionImpl().toUDF("reduce");

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,7 @@
126126
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAKETIME;
127127
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAP_APPEND;
128128
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAP_CONCAT;
129+
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MAP_REMOVE;
129130
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH;
130131
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH_BOOL_PREFIX;
131132
import static org.opensearch.sql.expression.function.BuiltinFunctionName.MATCH_PHRASE;
@@ -908,8 +909,9 @@ void populate() {
908909
registerOperator(ARRAY, PPLBuiltinOperators.ARRAY);
909910
registerOperator(MVAPPEND, PPLBuiltinOperators.MVAPPEND);
910911
registerOperator(MAP_APPEND, PPLBuiltinOperators.MAP_APPEND);
911-
registerOperator(ARRAY_LENGTH, SqlLibraryOperators.ARRAY_LENGTH);
912912
registerOperator(MAP_CONCAT, SqlLibraryOperators.MAP_CONCAT);
913+
registerOperator(MAP_REMOVE, PPLBuiltinOperators.MAP_REMOVE);
914+
registerOperator(ARRAY_LENGTH, SqlLibraryOperators.ARRAY_LENGTH);
913915
registerOperator(FORALL, PPLBuiltinOperators.FORALL);
914916
registerOperator(EXISTS, PPLBuiltinOperators.EXISTS);
915917
registerOperator(FILTER, PPLBuiltinOperators.FILTER);

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

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.sql.expression.function.CollectionUDF;
7+
8+
import static org.junit.jupiter.api.Assertions.assertEquals;
9+
import static org.junit.jupiter.api.Assertions.assertNull;
10+
import static org.junit.jupiter.api.Assertions.assertThrows;
11+
12+
import java.util.Arrays;
13+
import java.util.HashMap;
14+
import java.util.List;
15+
import java.util.Map;
16+
import org.junit.jupiter.api.Test;
17+
18+
public class MapRemoveFunctionImplTest {
19+
20+
@Test
21+
public void testMapRemoveWithNullMap() {
22+
Object result = MapRemoveFunctionImpl.mapRemove(null, Arrays.asList("key1", "key2"));
23+
assertNull(result);
24+
}
25+
26+
@Test
27+
public void testMapRemoveWithNullKeys() {
28+
Map<String, Object> map = getBaseMap();
29+
30+
Object result = MapRemoveFunctionImpl.mapRemove(map, null);
31+
assertEquals(map, result);
32+
}
33+
34+
@Test
35+
public void testMapRemoveWithInvalidMapArgument() {
36+
String notAMap = "not a map";
37+
List<String> keysToRemove = Arrays.asList("key1");
38+
39+
IllegalArgumentException exception =
40+
assertThrows(
41+
IllegalArgumentException.class,
42+
() -> MapRemoveFunctionImpl.mapRemove(notAMap, keysToRemove));
43+
44+
assertEquals(
45+
"First argument must be a map, got: class java.lang.String", exception.getMessage());
46+
}
47+
48+
@Test
49+
public void testMapRemoveWithInvalidKeysArgument() {
50+
Map<String, Object> map = getBaseMap();
51+
String notAList = "not a list";
52+
53+
IllegalArgumentException exception =
54+
assertThrows(
55+
IllegalArgumentException.class, () -> MapRemoveFunctionImpl.mapRemove(map, notAList));
56+
57+
assertEquals(
58+
"Second argument must be an array/list, got: class java.lang.String",
59+
exception.getMessage());
60+
}
61+
62+
@Test
63+
public void testMapRemoveExistingKeys() {
64+
Map<String, Object> map = getBaseMap();
65+
List<String> keysToRemove = Arrays.asList("key1", "key3");
66+
67+
@SuppressWarnings("unchecked")
68+
Map<String, Object> result =
69+
(Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
70+
71+
assertEquals(1, result.size());
72+
assertEquals("value2", result.get("key2"));
73+
assertNull(result.get("key1"));
74+
assertNull(result.get("key3"));
75+
76+
// Verify original map is not modified
77+
assertEqualToBaseMap(map);
78+
}
79+
80+
@Test
81+
public void testMapRemoveNonExistingKeys() {
82+
Map<String, Object> map = getBaseMap();
83+
List<String> keysToRemove = Arrays.asList("key4", "key5");
84+
85+
@SuppressWarnings("unchecked")
86+
Map<String, Object> result =
87+
(Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
88+
89+
assertEqualToBaseMap(result);
90+
}
91+
92+
@Test
93+
public void testMapRemoveEmptyKeysList() {
94+
Map<String, Object> map = getBaseMap();
95+
List<String> keysToRemove = Arrays.asList();
96+
97+
@SuppressWarnings("unchecked")
98+
Map<String, Object> result =
99+
(Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
100+
101+
assertEqualToBaseMap(result);
102+
}
103+
104+
@Test
105+
public void testMapRemoveMixedExistingAndNonExistingKeys() {
106+
Map<String, Object> map = getBaseMap();
107+
List<String> keysToRemove = Arrays.asList("key1", "key4", "key2");
108+
109+
@SuppressWarnings("unchecked")
110+
Map<String, Object> result =
111+
(Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
112+
113+
assertEquals(1, result.size());
114+
assertEquals("value3", result.get("key3"));
115+
}
116+
117+
@Test
118+
public void testMapRemoveWithNullKeysInList() {
119+
Map<String, Object> map = getBaseMap();
120+
List<Object> keysToRemove = Arrays.asList("key1", null, "key3");
121+
122+
@SuppressWarnings("unchecked")
123+
Map<String, Object> result =
124+
(Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
125+
126+
assertEquals(1, result.size());
127+
assertEquals("value2", result.get("key2"));
128+
}
129+
130+
@Test
131+
public void testMapRemoveWithDifferentValueTypes() {
132+
Map<String, Object> map = new HashMap<>();
133+
map.put("string", "value");
134+
map.put("number", 42);
135+
map.put("boolean", true);
136+
map.put("list", Arrays.asList(1, 2, 3));
137+
List<String> keysToRemove = Arrays.asList("number", "boolean");
138+
139+
@SuppressWarnings("unchecked")
140+
Map<String, Object> result =
141+
(Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
142+
143+
assertEquals(2, result.size());
144+
assertEquals("value", result.get("string"));
145+
assertEquals(Arrays.asList(1, 2, 3), result.get("list"));
146+
assertNull(result.get("number"));
147+
assertNull(result.get("boolean"));
148+
}
149+
150+
@Test
151+
public void testMapRemoveAllKeys() {
152+
Map<String, Object> map = getBaseMap();
153+
List<String> keysToRemove = Arrays.asList("key1", "key2", "key3");
154+
155+
@SuppressWarnings("unchecked")
156+
Map<String, Object> result =
157+
(Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
158+
159+
assertEquals(0, result.size());
160+
}
161+
162+
@Test
163+
public void testMapRemoveWithEmptyMap() {
164+
Map<String, Object> map = new HashMap<>();
165+
List<String> keysToRemove = Arrays.asList("key1", "key2");
166+
167+
@SuppressWarnings("unchecked")
168+
Map<String, Object> result =
169+
(Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
170+
171+
assertEquals(0, result.size());
172+
}
173+
174+
@Test
175+
public void testMapRemoveWithNonStringKeys() {
176+
Map<String, Object> map = new HashMap<>();
177+
map.put("key1", "value1");
178+
map.put("key2", "value2");
179+
map.put("123", "numeric_key_value");
180+
181+
List<Object> keysToRemove = Arrays.asList("key1", 123); // 123 will be converted to string "123"
182+
183+
@SuppressWarnings("unchecked")
184+
Map<String, Object> result =
185+
(Map<String, Object>) MapRemoveFunctionImpl.mapRemove(map, keysToRemove);
186+
187+
assertEquals(1, result.size());
188+
assertEquals("value2", result.get("key2"));
189+
}
190+
191+
private Map<String, Object> getBaseMap() {
192+
Map<String, Object> map = new HashMap<>();
193+
map.put("key1", "value1");
194+
map.put("key2", "value2");
195+
map.put("key3", "value3");
196+
return map;
197+
}
198+
199+
private void assertEqualToBaseMap(Map<String, Object> map) {
200+
assertEquals(3, map.size());
201+
assertEquals("value1", map.get("key1"));
202+
assertEquals("value2", map.get("key2"));
203+
assertEquals("value3", map.get("key3"));
204+
}
205+
}

0 commit comments

Comments
 (0)