Skip to content

Commit 1908100

Browse files
dukewy123坤羽
dukewy123
and
坤羽
authored
[DSL] Support UDF json get object (#536)
* dsl support json get * PullRequest: 927 add udf get json object, refine code style Merge branch dev_ky_opensource_json_get of git@code.alipay.com:AntGraph/GeaFlow.git into dev_opensource https://code.alipay.com/AntGraph/GeaFlow/pull_requests/927?tab=diff Reviewed-by: 知尘 <zhichen.zq@antgroup.com> * add udf get json object, refine code style * change import test --------- Co-authored-by: 坤羽 <ky.wy@antgroup.com>
1 parent 6351ac7 commit 1908100

7 files changed

Lines changed: 405 additions & 0 deletions

File tree

‎geaflow/geaflow-dsl/geaflow-dsl-plan/src/main/java/com/antgroup/geaflow/dsl/schema/function/BuildInSqlFunctionTable.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,7 @@
7575
import com.antgroup.geaflow.dsl.udf.table.string.Base64Encode;
7676
import com.antgroup.geaflow.dsl.udf.table.string.Concat;
7777
import com.antgroup.geaflow.dsl.udf.table.string.ConcatWS;
78+
import com.antgroup.geaflow.dsl.udf.table.string.GetJsonObject;
7879
import com.antgroup.geaflow.dsl.udf.table.string.Hash;
7980
import com.antgroup.geaflow.dsl.udf.table.string.IndexOf;
8081
import com.antgroup.geaflow.dsl.udf.table.string.Instr;
@@ -176,6 +177,7 @@ public class BuildInSqlFunctionTable extends ListSqlOperatorTable {
176177
.add(GeaFlowFunction.of(Substr.class))
177178
.add(GeaFlowFunction.of(UrlDecode.class))
178179
.add(GeaFlowFunction.of(UrlEncode.class))
180+
.add(GeaFlowFunction.of(GetJsonObject.class))
179181

180182
// udf.table.other
181183
.add(GeaFlowFunction.of(If.class))

‎geaflow/geaflow-dsl/geaflow-dsl-plan/src/main/java/com/antgroup/geaflow/dsl/udf/table/string/GetJsonObject.java‎

Lines changed: 266 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,266 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package com.antgroup.geaflow.dsl.udf.table.string;
21+
22+
import com.antgroup.geaflow.dsl.common.function.Description;
23+
import com.antgroup.geaflow.dsl.common.function.UDF;
24+
import com.google.common.collect.Iterators;
25+
import java.util.ArrayList;
26+
import java.util.Iterator;
27+
import java.util.LinkedHashMap;
28+
import java.util.List;
29+
import java.util.Map;
30+
import java.util.regex.Matcher;
31+
import java.util.regex.Pattern;
32+
import org.codehaus.jackson.JsonFactory;
33+
import org.codehaus.jackson.JsonParser;
34+
import org.codehaus.jackson.map.ObjectMapper;
35+
import org.codehaus.jackson.map.type.TypeFactory;
36+
import org.codehaus.jackson.type.JavaType;
37+
38+
@Description(name = "get_json_object", description = "parse string from json string.")
39+
public class GetJsonObject extends UDF {
40+
41+
private static final JsonFactory JSON_FACTORY = new JsonFactory();
42+
private static final JavaType MAP_TYPE = TypeFactory.fromClass(Map.class);
43+
private static final JavaType LIST_TYPE = TypeFactory.fromClass(List.class);
44+
45+
static {
46+
JSON_FACTORY.enable(JsonParser.Feature.ALLOW_UNQUOTED_CONTROL_CHARS);
47+
}
48+
49+
private final Pattern patternKey = Pattern.compile("^([a-zA-Z0-9_\\-\\:\\s]+).*");
50+
private final Pattern patternIndex = Pattern.compile("\\[([0-9]+|\\*)\\]");
51+
private final ObjectMapper mapper = new ObjectMapper(JSON_FACTORY);
52+
53+
private Map<String, Object> extractObjectCache = new HashCache<String, Object>();
54+
private Map<String, String[]> pathExprCache = new HashCache<String, String[]>();
55+
private Map<String, ArrayList<String>> indexListCache = new HashCache<String, ArrayList<String>>();
56+
private Map<String, String> mKeyGroup1Cache = new HashCache<String, String>();
57+
private Map<String, Boolean> mKeyMatchesCache = new HashCache<String, Boolean>();
58+
59+
private transient AddingList jsonList = new AddingList();
60+
61+
static class HashCache<K, V> extends LinkedHashMap<K, V> {
62+
63+
private static final int CACHE_SIZE = 16;
64+
private static final int INIT_SIZE = 32;
65+
private static final float LOAD_FACTOR = 0.6f;
66+
67+
HashCache() {
68+
super(INIT_SIZE, LOAD_FACTOR);
69+
}
70+
71+
private static final long serialVersionUID = 1;
72+
73+
@Override
74+
protected boolean removeEldestEntry(Map.Entry<K, V> eldest) {
75+
return size() > CACHE_SIZE;
76+
}
77+
}
78+
79+
public String eval(String jsonString, String pathString) {
80+
81+
if (pathString != null && !pathString.startsWith("$.")) {
82+
pathString = "$." + pathString;
83+
}
84+
if (jsonString == null || jsonString.isEmpty() || pathString == null
85+
|| pathString.isEmpty() || pathString.charAt(0) != '$') {
86+
return null;
87+
}
88+
89+
int pathExprStart = 1;
90+
boolean isRootArray = false;
91+
92+
if (pathString.length() > 1) {
93+
if (pathString.charAt(1) == '[') {
94+
pathExprStart = 0;
95+
isRootArray = true;
96+
} else if (pathString.charAt(1) == '.') {
97+
isRootArray = pathString.length() > 2 && pathString.charAt(2) == '[';
98+
} else {
99+
return null;
100+
}
101+
}
102+
103+
String[] pathExpr = pathExprCache.get(pathString);
104+
if (pathExpr == null) {
105+
pathExpr = pathString.split("\\.", -1);
106+
pathExprCache.put(pathString, pathExpr);
107+
}
108+
109+
Object extractObject = extractObjectCache.get(jsonString);
110+
if (extractObject == null) {
111+
JavaType javaType = isRootArray ? LIST_TYPE : MAP_TYPE;
112+
try {
113+
extractObject = mapper.readValue(jsonString, javaType);
114+
} catch (Exception e) {
115+
return null;
116+
}
117+
extractObjectCache.put(jsonString, extractObject);
118+
}
119+
for (int i = pathExprStart; i < pathExpr.length; i++) {
120+
if (extractObject == null) {
121+
return null;
122+
}
123+
extractObject = extract(extractObject, pathExpr[i], i == pathExprStart && isRootArray);
124+
}
125+
String result = null;
126+
if (extractObject instanceof Map || extractObject instanceof List) {
127+
try {
128+
result = mapper.writeValueAsString(extractObject);
129+
} catch (Exception e) {
130+
return null;
131+
}
132+
} else if (extractObject != null) {
133+
result = extractObject.toString();
134+
} else {
135+
return null;
136+
}
137+
return result;
138+
}
139+
140+
private Object extract(Object json, String path, boolean skipMapProc) {
141+
if (!skipMapProc) {
142+
Matcher mKey = null;
143+
Boolean mKeyMatches = mKeyMatchesCache.get(path);
144+
if (mKeyMatches == null) {
145+
mKey = patternKey.matcher(path);
146+
mKeyMatches = mKey.matches() ? Boolean.TRUE : Boolean.FALSE;
147+
mKeyMatchesCache.put(path, mKeyMatches);
148+
}
149+
if (!mKeyMatches.booleanValue()) {
150+
return null;
151+
}
152+
153+
String mKeyGroup1 = mKeyGroup1Cache.get(path);
154+
if (mKeyGroup1 == null) {
155+
if (mKey == null) {
156+
mKey = patternKey.matcher(path);
157+
mKeyMatches = mKey.matches() ? Boolean.TRUE : Boolean.FALSE;
158+
mKeyMatchesCache.put(path, mKeyMatches);
159+
if (!mKeyMatches.booleanValue()) {
160+
return null;
161+
}
162+
}
163+
mKeyGroup1 = mKey.group(1);
164+
mKeyGroup1Cache.put(path, mKeyGroup1);
165+
}
166+
json = extract_json_withkey(json, mKeyGroup1);
167+
}
168+
// Cache indexList
169+
ArrayList<String> indexList = indexListCache.get(path);
170+
if (indexList == null) {
171+
Matcher mIndex = patternIndex.matcher(path);
172+
indexList = new ArrayList<String>();
173+
while (mIndex.find()) {
174+
indexList.add(mIndex.group(1));
175+
}
176+
indexListCache.put(path, indexList);
177+
}
178+
179+
if (indexList.size() > 0) {
180+
json = extract_json_withindex(json, indexList);
181+
}
182+
183+
return json;
184+
}
185+
186+
private static class AddingList extends ArrayList<Object> {
187+
188+
@Override
189+
public Iterator<Object> iterator() {
190+
return Iterators.forArray(toArray());
191+
}
192+
193+
@Override
194+
public void removeRange(int fromIndex, int toIndex) {
195+
super.removeRange(fromIndex, toIndex);
196+
}
197+
}
198+
199+
@SuppressWarnings("unchecked")
200+
private Object extract_json_withindex(Object json, ArrayList<String> indexList) {
201+
202+
jsonList.clear();
203+
jsonList.add(json);
204+
for (String index : indexList) {
205+
int targets = jsonList.size();
206+
if (index.equalsIgnoreCase("*")) {
207+
for (Object array : jsonList) {
208+
if (array instanceof List) {
209+
for (int j = 0; j < ((List<Object>) array).size(); j++) {
210+
jsonList.add(((List<Object>) array).get(j));
211+
}
212+
}
213+
}
214+
} else {
215+
for (Object array : jsonList) {
216+
int indexValue = Integer.parseInt(index);
217+
if (!(array instanceof List)) {
218+
continue;
219+
}
220+
List<Object> list = (List<Object>) array;
221+
if (indexValue >= list.size()) {
222+
continue;
223+
}
224+
jsonList.add(list.get(indexValue));
225+
}
226+
}
227+
if (jsonList.size() == targets) {
228+
return null;
229+
}
230+
jsonList.removeRange(0, targets);
231+
}
232+
if (jsonList.isEmpty()) {
233+
return null;
234+
}
235+
return (jsonList.size() > 1) ? new ArrayList<Object>(jsonList) : jsonList.get(0);
236+
}
237+
238+
@SuppressWarnings("unchecked")
239+
private Object extract_json_withkey(Object json, String path) {
240+
if (json instanceof List) {
241+
List<Object> jsonList = new ArrayList<Object>();
242+
for (int i = 0; i < ((List<Object>) json).size(); i++) {
243+
Object jsonElem = ((List<Object>) json).get(i);
244+
Object jsonObj = null;
245+
if (jsonElem instanceof Map) {
246+
jsonObj = ((Map<String, Object>) jsonElem).get(path);
247+
} else {
248+
continue;
249+
}
250+
if (jsonObj instanceof List) {
251+
for (int j = 0; j < ((List<Object>) jsonObj).size(); j++) {
252+
jsonList.add(((List<Object>) jsonObj).get(j));
253+
}
254+
} else if (jsonObj != null) {
255+
jsonList.add(jsonObj);
256+
}
257+
}
258+
return (jsonList.isEmpty()) ? null : jsonList;
259+
} else if (json instanceof Map) {
260+
return ((Map<String, Object>) json).get(path);
261+
} else {
262+
return null;
263+
}
264+
}
265+
266+
}

‎geaflow/geaflow-dsl/geaflow-dsl-plan/src/test/java/com/antgroup/geaflow/dsl/udf/string/UDFGetJsonObjectTest.java‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package com.antgroup.geaflow.dsl.udf.string;
21+
22+
import static org.testng.Assert.assertEquals;
23+
24+
import com.antgroup.geaflow.dsl.udf.table.string.GetJsonObject;
25+
import org.testng.annotations.Test;
26+
27+
public class UDFGetJsonObjectTest {
28+
29+
@Test
30+
public void test() {
31+
32+
GetJsonObject udf = new GetJsonObject();
33+
34+
String jsonStr = "{\"name\": \"Bob\", \"age\": 30, \"address\": " +
35+
"{\"city\": \"Los Angeles\", \"zip\": \"90001\"},\"items\": [\"item1\", \"item2\", \"item3\"]}";
36+
37+
assertEquals("Bob", udf.eval(jsonStr, "name"));
38+
assertEquals("30", udf.eval(jsonStr, "$.age"));
39+
assertEquals("Los Angeles", udf.eval(jsonStr, "$.address.city"));
40+
assertEquals("item3", udf.eval(jsonStr, "$.items[2]"));
41+
assertEquals(null, udf.eval(jsonStr, "gender"));
42+
assertEquals(null, udf.eval(jsonStr, "$.items[3]"));
43+
}
44+
}

‎geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/java/com/antgroup/geaflow/dsl/runtime/query/udf/JsonParserTest.java‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package com.antgroup.geaflow.dsl.runtime.query.udf;
21+
22+
import com.antgroup.geaflow.dsl.runtime.query.QueryTester;
23+
import org.testng.annotations.Test;
24+
25+
public class JsonParserTest {
26+
27+
@Test
28+
public void testJsonPathGet() throws Exception {
29+
QueryTester
30+
.build()
31+
.withQueryPath("/query/json_path_get_001.sql")
32+
.execute()
33+
.checkSinkResult();
34+
}
35+
36+
}

‎geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/resources/data/users_json_object.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
1|jim|15|{"name": "Bob", "age": 30, "address": {"city": "Los Angeles", "zip": "90001"},"items": ["word", "size", "metrics"]}
2+
2|kate|16|{"name": "Alen", "age": 20, "address": {"city": "Zhe Jiang", "zip": "310005"},"items": ["red", "green", "yellow"]}
3+
3|same|20|{"name": "John", "age": 25,"items": ["round", "triangle"]}
4+
4|lucy|21|{"name": "Bush", "age": 11, "address": {"city": "Zhe Jiang", "zip": "310006"}}
5+
5|brown|22|{"name": "Ray"}

‎geaflow/geaflow-dsl/geaflow-dsl-runtime/src/test/resources/expect/json_path_get_001.txt‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
Bob|30|Los Angeles|metrics
2+
Alen|20|Zhe Jiang|yellow
3+
John|25|null|null
4+
Bush|11|Zhe Jiang|null
5+
Ray|null|null|null

0 commit comments

Comments
 (0)