Skip to content

Commit

Permalink
[Refactor](UDF) Refactor the java udf code to reduce the unless code
Browse files Browse the repository at this point in the history
  • Loading branch information
HappenLee committed Feb 5, 2025
1 parent e042e9f commit 4f144a5
Show file tree
Hide file tree
Showing 3 changed files with 56 additions and 69 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,8 @@
import java.math.BigInteger;
import java.net.InetAddress;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

// Data types that are supported as return or argument types in Java UDFs.
Expand Down Expand Up @@ -63,32 +64,36 @@ public class JavaUdfDataType {
public static final JavaUdfDataType MAP_TYPE = new JavaUdfDataType("MAP_TYPE", TPrimitiveType.MAP, 0);
public static final JavaUdfDataType STRUCT_TYPE = new JavaUdfDataType("STRUCT_TYPE", TPrimitiveType.STRUCT, 0);

private static Set<JavaUdfDataType> JavaUdfDataTypeSet = new HashSet<>();
private static final Map<TPrimitiveType, JavaUdfDataType> javaUdfDataTypeMap = new HashMap<>();

public static void addJavaUdfDataType(JavaUdfDataType dataType) {
javaUdfDataTypeMap.put(dataType.getPrimitiveType(), dataType);
}

static {
JavaUdfDataTypeSet.add(INVALID_TYPE);
JavaUdfDataTypeSet.add(BOOLEAN);
JavaUdfDataTypeSet.add(TINYINT);
JavaUdfDataTypeSet.add(SMALLINT);
JavaUdfDataTypeSet.add(INT);
JavaUdfDataTypeSet.add(BIGINT);
JavaUdfDataTypeSet.add(FLOAT);
JavaUdfDataTypeSet.add(DOUBLE);
JavaUdfDataTypeSet.add(STRING);
JavaUdfDataTypeSet.add(DATE);
JavaUdfDataTypeSet.add(DATETIME);
JavaUdfDataTypeSet.add(LARGEINT);
JavaUdfDataTypeSet.add(DECIMALV2);
JavaUdfDataTypeSet.add(DATEV2);
JavaUdfDataTypeSet.add(DATETIMEV2);
JavaUdfDataTypeSet.add(DECIMAL32);
JavaUdfDataTypeSet.add(DECIMAL64);
JavaUdfDataTypeSet.add(DECIMAL128);
JavaUdfDataTypeSet.add(ARRAY_TYPE);
JavaUdfDataTypeSet.add(MAP_TYPE);
JavaUdfDataTypeSet.add(STRUCT_TYPE);
JavaUdfDataTypeSet.add(IPV4);
JavaUdfDataTypeSet.add(IPV6);
addJavaUdfDataType(INVALID_TYPE);
addJavaUdfDataType(BOOLEAN);
addJavaUdfDataType(TINYINT);
addJavaUdfDataType(SMALLINT);
addJavaUdfDataType(INT);
addJavaUdfDataType(BIGINT);
addJavaUdfDataType(FLOAT);
addJavaUdfDataType(DOUBLE);
addJavaUdfDataType(STRING);
addJavaUdfDataType(DATE);
addJavaUdfDataType(DATETIME);
addJavaUdfDataType(LARGEINT);
addJavaUdfDataType(DECIMALV2);
addJavaUdfDataType(DATEV2);
addJavaUdfDataType(DATETIMEV2);
addJavaUdfDataType(DECIMAL32);
addJavaUdfDataType(DECIMAL64);
addJavaUdfDataType(DECIMAL128);
addJavaUdfDataType(ARRAY_TYPE);
addJavaUdfDataType(MAP_TYPE);
addJavaUdfDataType(STRUCT_TYPE);
addJavaUdfDataType(IPV4);
addJavaUdfDataType(IPV6);
}

private final String description;
Expand Down Expand Up @@ -117,17 +122,33 @@ public JavaUdfDataType(JavaUdfDataType other) {

@Override
public String toString() {
return description;
}
StringBuilder res = new StringBuilder();
res.append(description);
// TODO: the item/key/value type should be dispose in child class
if (getItemType() != null) {
res.append(" item: ").append(getItemType().toString()).append(" sql: ")
.append(getItemType().toSql());
}
if (getKeyType() != null) {
res.append(" key: ").append(getKeyType().toString()).append(" sql: ")
.append(getKeyType().toSql());
}
if (getValueType() != null) {
res.append(" value: ").append(getValueType().toString()).append(" sql: ")
.append(getValueType().toSql());
}

public TPrimitiveType getPrimitiveType() {
return thriftType;
return res.toString();
}

public int getLen() {
return len;
}

public TPrimitiveType getPrimitiveType() {
return thriftType;
}

public static Set<JavaUdfDataType> getCandidateTypes(Class<?> c) {
if (c == boolean.class || c == Boolean.class) {
return Sets.newHashSet(JavaUdfDataType.BOOLEAN);
Expand Down Expand Up @@ -169,19 +190,14 @@ public static Set<JavaUdfDataType> getCandidateTypes(Class<?> c) {
}

public static boolean isSupported(Type t) {
for (JavaUdfDataType javaType : JavaUdfDataTypeSet) {
if (javaType == JavaUdfDataType.INVALID_TYPE) {
continue;
}
if (javaType.getPrimitiveType() == t.getPrimitiveType().toThrift()) {
return true;
}
}
if (t.getPrimitiveType().toThrift() == TPrimitiveType.VARCHAR
|| t.getPrimitiveType().toThrift() == TPrimitiveType.CHAR) {
TPrimitiveType thriftType = t.getPrimitiveType().toThrift();
// varchar and char are supported in java udf, type is String
if (thriftType == TPrimitiveType.VARCHAR
|| thriftType == TPrimitiveType.CHAR) {
return true;
}
return false;
return !thriftType.equals(TPrimitiveType.INVALID_TYPE)
&& javaUdfDataTypeMap.containsKey(thriftType);
}

public int getPrecision() {
Expand Down Expand Up @@ -209,7 +225,6 @@ public void setItemType(Type type) throws InternalException {
this.itemType = type;
} else {
if (!this.itemType.matchesType(type)) {
LOG.info("set error");
throw new InternalException("udf type not matches origin type :" + this.itemType.toSql()
+ " set type :" + type.toSql());
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -102,32 +102,8 @@ public String debugString() {
StringBuilder res = new StringBuilder();
for (JavaUdfDataType type : argTypes) {
res.append(type.toString());
if (type.getItemType() != null) {
res.append(" item: ").append(type.getItemType().toString()).append(" sql: ")
.append(type.getItemType().toSql());
}
if (type.getKeyType() != null) {
res.append(" key: ").append(type.getKeyType().toString()).append(" sql: ")
.append(type.getKeyType().toSql());
}
if (type.getValueType() != null) {
res.append(" key: ").append(type.getValueType().toString()).append(" sql: ")
.append(type.getValueType().toSql());
}
}
res.append(" return type: ").append(retType.toString());
if (retType.getItemType() != null) {
res.append(" item: ").append(retType.getItemType().toString()).append(" sql: ")
.append(retType.getItemType().toSql());
}
if (retType.getKeyType() != null) {
res.append(" key: ").append(retType.getKeyType().toString()).append(" sql: ")
.append(retType.getKeyType().toSql());
}
if (retType.getValueType() != null) {
res.append(" key: ").append(retType.getValueType().toString()).append(" sql: ")
.append(retType.getValueType().toSql());
}
res.append(" methodAccess: ").append(methodAccess.toString());
res.append(" fn.toString(): ").append(fn.toString());
return res.toString();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -216,16 +216,12 @@ private void checkAndCacheUdfClass(String className, UdfClassCache cache, Type f
} else {
cache.retType = returnType.second;
}
Type keyType = cache.retType.getKeyType();
Type valueType = cache.retType.getValueType();
Pair<Boolean, JavaUdfDataType[]> inputType = UdfUtils.setArgTypes(parameterTypes, cache.argClass, false);
if (!inputType.first) {
continue;
} else {
cache.argTypes = inputType.second;
}
cache.retType.setKeyType(keyType);
cache.retType.setValueType(valueType);
return;
}
StringBuilder sb = new StringBuilder();
Expand Down

0 comments on commit 4f144a5

Please sign in to comment.