Skip to content

Commit 0ab4a98

Browse files
luchunliangluchunliang
committed
[INLONG-12125][SDK] Enhance Transform SDK protobuf processing and SQL alias parsing
1 parent 4307016 commit 0ab4a98

6 files changed

Lines changed: 116 additions & 30 deletions

File tree

‎inlong-sdk/transform-sdk/src/main/java/org/apache/inlong/sdk/transform/decode/PbSourceData.java‎

Lines changed: 89 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import com.google.protobuf.Descriptors.FieldDescriptor.JavaType;
2828
import com.google.protobuf.DynamicMessage;
2929
import org.apache.commons.lang3.StringUtils;
30+
import org.apache.commons.lang3.math.NumberUtils;
3031
import org.apache.flink.table.data.GenericArrayData;
3132
import org.apache.flink.table.data.GenericMapData;
3233
import org.apache.flink.table.data.GenericRowData;
@@ -35,8 +36,10 @@
3536
import org.slf4j.LoggerFactory;
3637

3738
import java.nio.charset.Charset;
39+
import java.nio.charset.StandardCharsets;
3840
import java.util.ArrayList;
3941
import java.util.HashMap;
42+
import java.util.IdentityHashMap;
4043
import java.util.List;
4144
import java.util.Map;
4245
import java.util.Objects;
@@ -68,8 +71,8 @@ public class PbSourceData extends AbstractSourceData {
6871

6972
protected Charset srcCharset;
7073

71-
private Map<DynamicMessage, Map<String, Object>> nodeValueCache = new HashMap<>();
72-
private Map<DynamicMessage, Map<String, Map<Object, Object>>> mapNodeCache = new HashMap<>();
74+
private Map<DynamicMessage, Map<String, Object>> nodeValueCache = new IdentityHashMap<>();
75+
private Map<DynamicMessage, Map<String, Map<Object, Object>>> mapNodeCache = new IdentityHashMap<>();
7376

7477
/**
7578
* Constructor
@@ -158,7 +161,10 @@ public Object getField(int rowNum, String fieldName) {
158161
List<Object> valueList = (List) fieldValue;
159162
List<Object> result = new ArrayList<>(valueList.size());
160163
for (Object value : valueList) {
161-
result.add(this.buildFieldValue(lastNode.getFieldDesc(), value));
164+
Object itemValue = this.buildFieldValue(lastNode.getFieldDesc(), value);
165+
if (itemValue != null) {
166+
result.add(itemValue);
167+
}
162168
}
163169
return new GenericArrayData(result.toArray());
164170
}
@@ -184,28 +190,56 @@ public Object buildFieldValue(FieldDescriptor fieldDesc, Object nodeValue) {
184190
}
185191
switch (fieldDesc.getJavaType()) {
186192
case STRING:
193+
if (nodeValue instanceof BinaryStringData) {
194+
return nodeValue;
195+
}
196+
if (nodeValue instanceof String) {
197+
return new BinaryStringData((String) nodeValue);
198+
}
199+
return new BinaryStringData(String.valueOf(nodeValue));
187200
case INT:
201+
if (nodeValue instanceof Integer) {
202+
return nodeValue;
203+
}
204+
return NumberUtils.toInt(String.valueOf(nodeValue), 0);
188205
case LONG:
206+
if (nodeValue instanceof Long) {
207+
return nodeValue;
208+
}
209+
return NumberUtils.toLong(String.valueOf(nodeValue), 0);
189210
case FLOAT:
211+
if (nodeValue instanceof Float) {
212+
return nodeValue;
213+
}
214+
return NumberUtils.toFloat(String.valueOf(nodeValue), 0);
190215
case DOUBLE:
216+
if (nodeValue instanceof Double) {
217+
return nodeValue;
218+
}
219+
return NumberUtils.toDouble(String.valueOf(nodeValue), 0);
191220
case BOOLEAN:
192-
return nodeValue;
221+
if (nodeValue instanceof Boolean) {
222+
return nodeValue;
223+
}
224+
return Boolean.TRUE.toString().equals(String.valueOf(nodeValue));
193225
case ENUM:
194226
if (nodeValue instanceof EnumValueDescriptor) {
195227
EnumValueDescriptor enumDesc = (EnumValueDescriptor) nodeValue;
196228
return enumDesc.getIndex();
197229
}
198230
return null;
199231
case BYTE_STRING:
200-
if (nodeValue instanceof ByteString) {
232+
if (nodeValue instanceof byte[]) {
233+
return nodeValue;
234+
} else if (nodeValue instanceof ByteString) {
201235
return ((ByteString) nodeValue).toByteArray();
202236
} else {
203-
return nodeValue;
237+
return String.valueOf(nodeValue).getBytes(StandardCharsets.ISO_8859_1);
204238
}
205239
case MESSAGE:
206240
return this.buildStructData(fieldDesc.getMessageType(), nodeValue);
207241
default:
208-
return String.valueOf(nodeValue);
242+
return nodeValue;
209243
}
210244
}
211245

@@ -269,7 +303,9 @@ protected Object buildMapData(Descriptors.Descriptor messageType, Object nodeVal
269303
DynamicMessage subnodeValue = (DynamicMessage) value;
270304
Object keyValue = buildFieldValue(keyField, subnodeValue.getField(keyField));
271305
Object valueValue = buildFieldValue(valueField, subnodeValue.getField(valueField));
272-
result.put(keyValue, valueValue);
306+
if (keyValue != null && valueValue != null) {
307+
result.put(keyValue, valueValue);
308+
}
273309
}
274310
return new GenericMapData(result);
275311
}
@@ -307,7 +343,7 @@ public List<DynamicMessage> getChildRoot() {
307343
}
308344

309345
public Object findFieldNode(int rowNum, String fieldName) {
310-
Object fieldValue = "";
346+
Object fieldValue = null;
311347
try {
312348
if (StringUtils.startsWith(fieldName, ROOT_KEY)) {
313349
fieldValue = this.findRootField(fieldName);
@@ -326,7 +362,7 @@ public Object findFieldNode(int rowNum, String fieldName) {
326362
}
327363
// error config
328364
if (childNodes.size() == 0) {
329-
return "";
365+
return null;
330366
}
331367
// parse other node
332368
fieldValue = this.findNodeValueByCache(childNodes, root);
@@ -377,7 +413,7 @@ private Object findRootField(String srcFieldName) {
377413

378414
private Object findChildField(int rowNum, String srcFieldName) {
379415
if (this.childRoot == null || this.childDesc == null) {
380-
return "";
416+
return null;
381417
}
382418
List<PbNode> childNodes = this.columnNodeMap.get(srcFieldName);
383419
if (childNodes == null) {
@@ -390,7 +426,7 @@ private Object findChildField(int rowNum, String srcFieldName) {
390426
}
391427
// error config
392428
if (childNodes.size() == 0) {
393-
return "";
429+
return null;
394430
}
395431
// parse other node
396432
DynamicMessage child = childRoot.get(rowNum);
@@ -423,7 +459,47 @@ public Object findNodeValueByCache(List<PbNode> childNodes, DynamicMessage root)
423459
subNodeValue = subNodeValueCache.get(node.getCurrentPath());
424460
if (subNodeValue != null) {
425461
if (i == childNodes.size() - 1) {
426-
return subNodeValue;
462+
// primitive
463+
if (node.isPrimitiveType()) {
464+
return subNodeValue;
465+
}
466+
// struct
467+
if (node.isStructType()) {
468+
return subNodeValue;
469+
}
470+
// array
471+
if (node.isArrayType()) {
472+
if (!node.isHasArrayIndex()) {
473+
return subNodeValue;
474+
}
475+
if (!(subNodeValue instanceof List)) {
476+
return null;
477+
}
478+
List<?> nodeValueList = (List<?>) subNodeValue;
479+
if (node.getArrayIndex() >= nodeValueList.size()) {
480+
return null;
481+
}
482+
Object newNode = nodeValueList.get(node.getArrayIndex());
483+
if (!(newNode instanceof DynamicMessage)) {
484+
return null;
485+
}
486+
return newNode;
487+
}
488+
// map
489+
if (node.isMapType()) {
490+
if (!node.isHasMapKey()) {
491+
return subNodeValue;
492+
}
493+
final Object mapNodeValue = subNodeValue;
494+
Map<Object, Object> mapCache = subMapNodeCache.computeIfAbsent(node.getCurrentPath(),
495+
k -> parseMapNode(mapNodeValue, node));
496+
Object fieldValue = mapCache.get(node.getMapKey());
497+
if (fieldValue == null || !(fieldValue instanceof DynamicMessage)) {
498+
return null;
499+
}
500+
return fieldValue;
501+
}
502+
return null;
427503
} else {
428504
// primitive
429505
if (node.isPrimitiveType()) {

‎inlong-sdk/transform-sdk/src/main/java/org/apache/inlong/sdk/transform/encode/DefaultSinkData.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ public void addField(String fieldName, Object fieldValue) {
5151
*/
5252
@Override
5353
public Object getField(String fieldName) {
54-
return this.currentRow.getOrDefault(fieldName, "");
54+
return this.currentRow.get(fieldName);
5555
}
5656

5757
/**

‎inlong-sdk/transform-sdk/src/main/java/org/apache/inlong/sdk/transform/process/TransformProcessor.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,11 @@ private void initTransformSql(String sql) throws JSQLParserException {
116116
fieldName = exprItem.toString();
117117
} else {
118118
fieldName = exprItem.getAlias().getName();
119+
// Strip surrounding backticks if present
120+
if (fieldName != null && fieldName.length() >= 2
121+
&& fieldName.startsWith("`") && fieldName.endsWith("`")) {
122+
fieldName = fieldName.substring(1, fieldName.length() - 1);
123+
}
119124
}
120125
if (!this.checkSelectField(fieldName)) {
121126
throw new JSQLParserException(

‎inlong-sdk/transform-sdk/src/main/java/org/apache/inlong/sdk/transform/utils/FieldToRowDataUtils.java‎

Lines changed: 13 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,8 @@
2323
import org.apache.flink.table.data.GenericArrayData;
2424
import org.apache.flink.table.data.GenericMapData;
2525
import org.apache.flink.table.data.GenericRowData;
26-
import org.apache.flink.table.data.StringData;
2726
import org.apache.flink.table.data.TimestampData;
27+
import org.apache.flink.table.data.binary.BinaryStringData;
2828
import org.apache.flink.table.types.logical.ArrayType;
2929
import org.apache.flink.table.types.logical.DecimalType;
3030
import org.apache.flink.table.types.logical.LogicalType;
@@ -139,7 +139,7 @@ private static FieldToRowDataConverter createFieldRowConverter(LogicalType field
139139
FieldToRowDataConverter keyConverter = createFieldRowConverter(((MapType) fieldType).getKeyType());
140140
FieldToRowDataConverter valueConverter = createFieldRowConverter(
141141
((MapType) fieldType).getValueType());
142-
Map map = (Map) obj;
142+
Map<?, ?> map = (Map<?, ?>) obj;
143143
Map<Object, Object> internalMap = new HashMap<>();
144144
for (Object k : map.keySet()) {
145145
internalMap.put(keyConverter.convert(k),
@@ -279,10 +279,16 @@ private static Object parseVarchar(Object obj) {
279279
if (obj == null) {
280280
return null;
281281
}
282+
if (obj instanceof BinaryStringData) {
283+
return obj;
284+
}
285+
if (obj instanceof String) {
286+
return new BinaryStringData((String) obj);
287+
}
282288
if (obj instanceof byte[]) {
283-
return StringData.fromString(new String((byte[]) obj));
289+
return new BinaryStringData(new String((byte[]) obj));
284290
}
285-
return StringData.fromString(String.valueOf(obj));
291+
return new BinaryStringData(String.valueOf(obj));
286292
} catch (RuntimeException e) {
287293
if (isIgnoreError()) {
288294
return null;
@@ -428,7 +434,7 @@ private static Object parseArray(Object obj) {
428434
if (obj instanceof List<?>) {
429435
return new GenericArrayData(((List<?>) obj).toArray());
430436
}
431-
return new GenericArrayData(new Object[]{obj});
437+
return null;
432438
} catch (RuntimeException e) {
433439
if (isIgnoreError()) {
434440
return null;
@@ -448,9 +454,7 @@ private static Object parseMap(Object obj) {
448454
if (obj instanceof Map<?, ?>) {
449455
return new GenericMapData((Map<?, ?>) obj);
450456
}
451-
Map<Object, Object> mapObj = new HashMap<>();
452-
mapObj.put(obj, obj);
453-
return new GenericMapData(mapObj);
457+
return null;
454458
} catch (RuntimeException e) {
455459
if (isIgnoreError()) {
456460
return null;
@@ -467,9 +471,7 @@ private static Object parseRow(Object obj) {
467471
if (obj instanceof GenericRowData) {
468472
return obj;
469473
}
470-
GenericRowData result = new GenericRowData(1);
471-
result.setField(0, obj);
472-
return result;
474+
return null;
473475
} catch (RuntimeException e) {
474476
if (isIgnoreError()) {
475477
return null;

‎inlong-sdk/transform-sdk/src/test/java/org/apache/inlong/sdk/transform/process/processor/TestCsv2CsvForErrorOrderProcessor.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -74,7 +74,7 @@ public void testKv2KvForErrorOrder() throws Exception {
7474
List<String> output1 = processor1.transform("key1=string11&key2=string12&key3=number11&key4=number12",
7575
new HashMap<>());
7676
Assert.assertEquals(1, output1.size());
77-
Assert.assertEquals("field1=string11&field2=&field3=number12", output1.get(0));
77+
Assert.assertEquals("field1=string11&field2=null&field3=number12", output1.get(0));
7878
}
7979

8080
@Test

‎inlong-sdk/transform-sdk/src/test/java/org/apache/inlong/sdk/transform/process/processor/TestPb2RowDataProcessor.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.apache.flink.table.data.GenericMapData;
3737
import org.apache.flink.table.data.GenericRowData;
3838
import org.apache.flink.table.data.RowData;
39+
import org.apache.flink.table.data.binary.BinaryStringData;
3940
import org.junit.Assert;
4041
import org.junit.Test;
4142

@@ -98,8 +99,9 @@ public void testPb2RowData() throws Exception {
9899
Assert.assertEquals(output.get(0).getString(1).toString(), "1");
99100
Assert.assertEquals(output.get(0).getString(2).toString(), "1713243918000");
100101
Assert.assertEquals(new String(output.get(0).getBinary(3)), "msgValue4");
101-
Assert.assertEquals(((GenericMapData) output.get(0).getMap(4)).get("key"), "value");
102-
Assert.assertEquals(((GenericMapData) output.get(0).getMap(4)).get("value"), null);
102+
Assert.assertEquals(((GenericMapData) output.get(0).getMap(4)).get(new BinaryStringData("key")).toString(),
103+
"value");
104+
Assert.assertEquals(((GenericMapData) output.get(0).getMap(4)).get(new BinaryStringData("value")), null);
103105
Assert.assertEquals(new String(((GenericRowData) output.get(0).getRow(5, 3)).getBinary(0)), "msgValue42");
104106
Assert.assertEquals(((GenericRowData) output.get(0).getRow(5, 3)).getLong(1), 1713243918002L);
105107
Assert.assertEquals(((GenericRowData) output.get(0).getRow(5, 3)).getMap(2).size(), 1);
@@ -109,8 +111,9 @@ public void testPb2RowData() throws Exception {
109111
Assert.assertEquals(output.get(1).getString(1).toString(), "1");
110112
Assert.assertEquals(output.get(1).getString(2).toString(), "1713243918002");
111113
Assert.assertEquals(new String(output.get(1).getBinary(3)), "msgValue42");
112-
Assert.assertEquals(((GenericMapData) output.get(1).getMap(4)).get("key2"), "value2");
113-
Assert.assertEquals(((GenericMapData) output.get(1).getMap(4)).get("value"), null);
114+
Assert.assertEquals(((GenericMapData) output.get(1).getMap(4)).get(new BinaryStringData("key2")).toString(),
115+
"value2");
116+
Assert.assertEquals(((GenericMapData) output.get(1).getMap(4)).get(new BinaryStringData("value")), null);
114117
Assert.assertEquals(new String(((GenericRowData) output.get(1).getRow(5, 3)).getBinary(0)), "msgValue42");
115118
Assert.assertEquals(((GenericRowData) output.get(1).getRow(5, 3)).getLong(1), 1713243918002L);
116119
Assert.assertEquals(((GenericRowData) output.get(1).getRow(5, 3)).getMap(2).size(), 1);

0 commit comments

Comments
 (0)