Skip to content

Commit df5e252

Browse files
authored
Merge index schema meta (#596)
* merge MetaData Signed-off-by: xinyual <xinyual@amazon.com> * remove logs Signed-off-by: xinyual <xinyual@amazon.com> * apply spotless Signed-off-by: xinyual <xinyual@amazon.com> * change equal logic Signed-off-by: xinyual <xinyual@amazon.com> * apply spotless Signed-off-by: xinyual <xinyual@amazon.com> * skip NPE Signed-off-by: xinyual <xinyual@amazon.com> --------- Signed-off-by: xinyual <xinyual@amazon.com>
1 parent df667b7 commit df5e252

6 files changed

Lines changed: 334 additions & 5 deletions

File tree

src/main/java/org/opensearch/agent/tools/PPLTool.java

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import org.opensearch.action.admin.indices.mapping.get.GetMappingsRequest;
3838
import org.opensearch.action.search.SearchRequest;
3939
import org.opensearch.agent.tools.utils.ToolHelper;
40+
import org.opensearch.agent.tools.utils.mergeMetaData.MergeRuleHelper;
4041
import org.opensearch.cluster.metadata.MappingMetadata;
4142
import org.opensearch.core.action.ActionListener;
4243
import org.opensearch.core.action.ActionResponse;
@@ -296,6 +297,7 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
296297

297298
return;
298299
}
300+
299301
GetMappingsRequest getMappingsRequest = buildGetMappingRequest(indexName);
300302
client.admin().indices().getMappings(getMappingsRequest, ActionListener.wrap(getMappingsResponse -> {
301303
Map<String, MappingMetadata> mappings = getMappingsResponse.getMappings();
@@ -512,16 +514,19 @@ private String constructTableInfoByPPLResultForSpark(Map<String, Object> schema,
512514
}
513515

514516
private String constructTableInfo(SearchHit[] searchHits, Map<String, MappingMetadata> mappings) throws PrivilegedActionException {
515-
String firstIndexName = (String) mappings.keySet().toArray()[0];
516-
MappingMetadata mappingMetadata = mappings.get(firstIndexName);
517-
Map<String, Object> mappingSource = (Map<String, Object>) mappingMetadata.getSourceAsMap().get("properties");
518-
if (Objects.isNull(mappingSource)) {
517+
if (mappings.keySet().size() == 0) {
519518
throw new IllegalArgumentException(
520519
"The querying index doesn't have mapping metadata, please add data to it or using another index."
521520
);
522521
}
522+
Map<String, Object> allFields = new HashMap<>();
523+
for (MappingMetadata mappingMetadata : mappings.values()) {
524+
Map<String, Object> mappingSource = (Map<String, Object>) mappingMetadata.getSourceAsMap().get("properties");
525+
MergeRuleHelper.merge(mappingSource, allFields);
526+
}
523527
Map<String, String> fieldsToType = new HashMap<>();
524-
ToolHelper.extractFieldNamesTypes(mappingSource, fieldsToType, "", false);
528+
ToolHelper.extractFieldNamesTypes(allFields, fieldsToType, "", false);
529+
525530
StringJoiner tableInfoJoiner = new StringJoiner("\n");
526531
List<String> sortedKeys = new ArrayList<>(fieldsToType.keySet());
527532
Collections.sort(sortedKeys);
Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.agent.tools.utils.mergeMetaData;
7+
8+
import java.util.Map;
9+
10+
/** This rule will merge two array/struct object and merge their properties */
11+
public class DeepMergeRule implements MergeRule {
12+
13+
@Override
14+
public boolean isMatch(Map<String, Object> source, Map<String, Object> target) {
15+
return source != null
16+
&& target != null
17+
&& source.get("properties") != null
18+
&& target.get("properties") != null
19+
&& source.getOrDefault("type", "object").equals(target.getOrDefault("type", "object"));
20+
}
21+
22+
@Override
23+
public void mergeInto(String key, Map<String, Object> source, Map<String, Object> target) {
24+
Map<String, Object> existing = (Map<String, Object>) target.get(key);
25+
MergeRuleHelper.merge((Map<String, Object>) source.get("properties"), (Map<String, Object>) existing.get("properties"));
26+
target.put(key, existing);
27+
}
28+
}
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.agent.tools.utils.mergeMetaData;
7+
8+
import java.util.Map;
9+
10+
/** The rule always keep the latest one. */
11+
public class LatestRule implements MergeRule {
12+
13+
@Override
14+
public boolean isMatch(Map<String, Object> source, Map<String, Object> target) {
15+
return true;
16+
}
17+
18+
@Override
19+
public void mergeInto(String key, Map<String, Object> source, Map<String, Object> target) {
20+
target.put(key, source);
21+
}
22+
}
Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.agent.tools.utils.mergeMetaData;
7+
8+
import java.util.Map;
9+
10+
/**
11+
* The Interface to merge index schemas. Need to implement isMatch: Whether match this rule,
12+
* mergeInto, how to merge the source type to target map.
13+
*/
14+
public interface MergeRule {
15+
boolean isMatch(Map<String, Object> source, Map<String, Object> target);
16+
17+
void mergeInto(String key, Map<String, Object> source, Map<String, Object> target);
18+
}
Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.agent.tools.utils.mergeMetaData;
7+
8+
import java.util.List;
9+
import java.util.Map;
10+
11+
public class MergeRuleHelper {
12+
private static final List<MergeRule> RULES = List
13+
.of(
14+
new DeepMergeRule(),
15+
new LatestRule() // must come last
16+
);
17+
18+
public static MergeRule selectRule(Map<String, Object> source, Map<String, Object> target) {
19+
MergeRule resultRule = RULES.stream().filter(rule -> rule.isMatch(source, target)).findFirst().orElseThrow(); // logically
20+
// unreachable if
21+
// fallback exists
22+
return resultRule;
23+
}
24+
25+
public static void merge(Map<String, Object> source, Map<String, Object> target) {
26+
for (Map.Entry<String, Object> entry : source.entrySet()) {
27+
String key = entry.getKey();
28+
Map<String, Object> sourceValue = (Map<String, Object>) entry.getValue();
29+
Map<String, Object> targetValue = (Map<String, Object>) target.get(key);
30+
MergeRuleHelper.selectRule(sourceValue, targetValue).mergeInto(key, sourceValue, target);
31+
}
32+
}
33+
}

src/test/java/org/opensearch/agent/tools/ToolHelperTests.java

Lines changed: 223 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212

1313
import org.junit.Test;
1414
import org.opensearch.agent.tools.utils.ToolHelper;
15+
import org.opensearch.agent.tools.utils.mergeMetaData.MergeRuleHelper;
16+
17+
import com.google.gson.Gson;
1518

1619
import lombok.extern.log4j.Log4j2;
1720

@@ -87,4 +90,224 @@ private void assertMapEquals(Map<String, String> expected, Map<String, String> a
8790
assertEquals(entry.getValue(), actual.get(entry.getKey()));
8891
}
8992
}
93+
94+
private Gson gson = new Gson();
95+
96+
private Map<String, Object> prepareMap1() {
97+
String mapBlock = """
98+
{
99+
"event": {
100+
"properties": {
101+
"field1": {
102+
"type": "string"
103+
}
104+
}
105+
}
106+
}
107+
108+
""";
109+
Map<String, Object> tmpMap = gson.fromJson(mapBlock, Map.class);
110+
return tmpMap;
111+
}
112+
113+
private Map<String, Object> prepareMap2() {
114+
String mapBlock = """
115+
{
116+
"event": {
117+
"properties": {
118+
"field2": {
119+
"type": "string"
120+
}
121+
}
122+
}
123+
}
124+
125+
""";
126+
Map<String, Object> tmpMap = gson.fromJson(mapBlock, Map.class);
127+
return tmpMap;
128+
}
129+
130+
private Map<String, Object> prepareNormalMap1() {
131+
String mapBlock = """
132+
{
133+
"event1": {
134+
"properties": {
135+
"field1": {
136+
"type": "string"
137+
}
138+
}
139+
},
140+
"replace" : {
141+
"type":"string"
142+
}
143+
144+
}
145+
146+
""";
147+
Map<String, Object> tmpMap = gson.fromJson(mapBlock, Map.class);
148+
return tmpMap;
149+
}
150+
151+
private Map<String, Object> prepareNormalMap2() {
152+
String mapBlock = """
153+
{
154+
"event2": {
155+
"properties": {
156+
"field2": {
157+
"type": "string"
158+
}
159+
}
160+
},
161+
"replace" : {
162+
"type":"keyword"
163+
}
164+
}
165+
166+
""";
167+
Map<String, Object> tmpMap = gson.fromJson(mapBlock, Map.class);
168+
return tmpMap;
169+
}
170+
171+
@Test
172+
public void testMergeTwoObjectMaps() {
173+
String mapBlock = """
174+
{
175+
"event": {
176+
"properties": {
177+
"field1": {
178+
"type": "string"
179+
},
180+
"field2": {
181+
"type": "string"
182+
}
183+
}
184+
}
185+
}
186+
187+
""";
188+
Map<String, Object> allFields = new HashMap<>();
189+
Map<String, Object> map1 = prepareMap1();
190+
Map<String, Object> map2 = prepareMap2();
191+
MergeRuleHelper.merge(map1, allFields);
192+
MergeRuleHelper.merge(map2, allFields);
193+
assertEquals(allFields, gson.fromJson(mapBlock, Map.class));
194+
}
195+
196+
@Test
197+
public void testMergeTwoNormalMaps() {
198+
String mapBlock = """
199+
{
200+
"event1": {
201+
"properties": {
202+
"field1": {
203+
"type": "string"
204+
}
205+
}
206+
},
207+
"event2": {
208+
"properties": {
209+
"field2": {
210+
"type": "string"
211+
}
212+
}
213+
},
214+
"replace" : {
215+
"type":"keyword"
216+
}
217+
}
218+
219+
""";
220+
Map<String, Object> allFields = new HashMap<>();
221+
Map<String, Object> map1 = prepareNormalMap1();
222+
Map<String, Object> map2 = prepareNormalMap2();
223+
MergeRuleHelper.merge(map1, allFields);
224+
MergeRuleHelper.merge(map2, allFields);
225+
assertEquals(allFields, gson.fromJson(mapBlock, Map.class));
226+
}
227+
228+
@Test
229+
public void testMergeTwoDeepMaps() {
230+
String mapBlock = """
231+
{
232+
"event": {
233+
"properties": {
234+
"field1": {
235+
"type": "string"
236+
},
237+
"field2": {
238+
"type": "string"
239+
},
240+
"deep": {
241+
"properties": {
242+
"field1": {
243+
"type": "string"
244+
},
245+
"field2": {
246+
"type": "string"
247+
}
248+
}
249+
}
250+
}
251+
}
252+
253+
}
254+
255+
""";
256+
Map<String, Object> allFields = new HashMap<>();
257+
Map<String, Object> map1 = prepareDeepMap1();
258+
Map<String, Object> map2 = prepareDeepMap2();
259+
MergeRuleHelper.merge(map1, allFields);
260+
MergeRuleHelper.merge(map2, allFields);
261+
assertEquals(allFields, gson.fromJson(mapBlock, Map.class));
262+
}
263+
264+
private Map<String, Object> prepareDeepMap1() {
265+
String mapBlock = """
266+
{
267+
"event": {
268+
"properties": {
269+
"field1": {
270+
"type": "string"
271+
},
272+
"deep": {
273+
"properties": {
274+
"field1": {
275+
"type": "string"
276+
}
277+
}
278+
}
279+
}
280+
}
281+
282+
}
283+
284+
""";
285+
Map<String, Object> tmpMap = gson.fromJson(mapBlock, Map.class);
286+
return tmpMap;
287+
}
288+
289+
private Map<String, Object> prepareDeepMap2() {
290+
String mapBlock = """
291+
{
292+
"event": {
293+
"properties": {
294+
"field2": {
295+
"type": "string"
296+
},
297+
"deep": {
298+
"properties": {
299+
"field2": {
300+
"type": "string"
301+
}
302+
}
303+
}
304+
}
305+
}
306+
}
307+
308+
""";
309+
Map<String, Object> tmpMap = gson.fromJson(mapBlock, Map.class);
310+
return tmpMap;
311+
}
312+
90313
}

0 commit comments

Comments
 (0)