feat(java): row encoder supports custom types and collections

Author: stevenschlansker

java/fury-core/src/main/java/org/apache/fury/type/CustomTypeRegistry.java

@@ -0,0 +1,43 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.fury.type;
+
+import org.apache.fury.annotation.Internal;
+
+@Internal
+public interface CustomTypeRegistry {
+  CustomTypeRegistry EMPTY =
+      new CustomTypeRegistry() {
+        @Override
+        public boolean hasCodec(final Class<?> beanType, final Class<?> fieldType) {
+          return false;
+        }
+
+        @Override
+        public boolean canConstructCollection(
+            final Class<?> collectionType, final Class<?> elementType) {
+          return false;
+        }
+      };
+
+  boolean hasCodec(Class<?> beanType, Class<?> fieldType);
+
+  boolean canConstructCollection(Class<?> collectionType, Class<?> elementType);
+}

java/fury-core/src/main/java/org/apache/fury/type/TypeUtils.java

@@ -568,6 +568,10 @@ public static boolean isBean(Class<?> clz) {
     return isBean(TypeRef.of(clz));
   }
 
+  public static boolean isBean(Type type, CustomTypeRegistry customTypes) {
+    return isBean(TypeRef.of(type), new LinkedHashSet<>(), customTypes);
+  }
+
   /**
    * Returns true if class is not array/iterable/map, and all fields is {@link
    * TypeUtils#isSupported(TypeRef)}. Bean class can't be a non-static inner class. Public static
@@ -577,7 +581,14 @@ public static boolean isBean(TypeRef<?> typeRef) {
     return isBean(typeRef, new LinkedHashSet<>());
   }
 
-  private static boolean isBean(TypeRef<?> typeRef, LinkedHashSet<TypeRef> walkedTypePath) {
+  private static boolean isBean(TypeRef<?> typeRef, LinkedHashSet<TypeRef<?>> walkedTypePath) {
+    return isBean(typeRef, walkedTypePath, CustomTypeRegistry.EMPTY);
+  }
+
+  private static boolean isBean(
+      TypeRef<?> typeRef,
+      LinkedHashSet<TypeRef<?>> walkedTypePath,
+      CustomTypeRegistry customTypes) {
     Class<?> cls = getRawType(typeRef);
     if (Modifier.isAbstract(cls.getModifiers()) || Modifier.isInterface(cls.getModifiers())) {
       return false;
@@ -589,7 +600,7 @@ private static boolean isBean(TypeRef<?> typeRef, LinkedHashSet<TypeRef> walkedT
       if (cls.getEnclosingClass() != null && !Modifier.isStatic(cls.getModifiers())) {
         return false;
       }
-      LinkedHashSet<TypeRef> newTypePath = new LinkedHashSet<>(walkedTypePath);
+      LinkedHashSet<TypeRef<?>> newTypePath = new LinkedHashSet<>(walkedTypePath);
       newTypePath.add(typeRef);
       if (cls == Object.class) {
         // return false for typeToken that point to un-specialized generic type.
@@ -602,14 +613,17 @@ private static boolean isBean(TypeRef<?> typeRef, LinkedHashSet<TypeRef> walkedT
               && !ITERABLE_TYPE.isSupertypeOf(typeRef)
               && !MAP_TYPE.isSupertypeOf(typeRef);
       if (maybe) {
+        Class<?> enclosingType = enclosingType(newTypePath);
         return Descriptor.getDescriptors(cls).stream()
             .allMatch(
                 d -> {
                   TypeRef<?> t = d.getTypeRef();
                   // do field modifiers and getter/setter validation here, not in getDescriptors.
                   // If Modifier.isFinal(d.getModifiers()), use reflection
                   // private field that doesn't have getter/setter will be handled by reflection.
-                  return isSupported(t, newTypePath) || isBean(t, newTypePath);
+                  return customTypes.hasCodec(enclosingType, t.getRawType())
+                      || isSupported(t, newTypePath, customTypes)
+                      || isBean(t, newTypePath, customTypes);
                 });
       } else {
         return false;
@@ -621,10 +635,13 @@ private static boolean isBean(TypeRef<?> typeRef, LinkedHashSet<TypeRef> walkedT
 
   /** Check if <code>typeToken</code> is supported by row-format. */
   public static boolean isSupported(TypeRef<?> typeRef) {
-    return isSupported(typeRef, new LinkedHashSet<>());
+    return isSupported(typeRef, new LinkedHashSet<>(), CustomTypeRegistry.EMPTY);
   }
 
-  private static boolean isSupported(TypeRef<?> typeRef, LinkedHashSet<TypeRef> walkedTypePath) {
+  private static boolean isSupported(
+      TypeRef<?> typeRef,
+      LinkedHashSet<TypeRef<?>> walkedTypePath,
+      CustomTypeRegistry customTypes) {
     Class<?> cls = getRawType(typeRef);
     if (!Modifier.isPublic(cls.getModifiers())) {
       return false;
@@ -639,6 +656,10 @@ private static boolean isSupported(TypeRef<?> typeRef, LinkedHashSet<TypeRef> wa
     } else if (typeRef.isArray()) {
       return isSupported(Objects.requireNonNull(typeRef.getComponentType()));
     } else if (ITERABLE_TYPE.isSupertypeOf(typeRef)) {
+      TypeRef<?> elementType = getElementType(typeRef);
+      if (customTypes.canConstructCollection(typeRef.getRawType(), elementType.getRawType())) {
+        return true;
+      }
       boolean isSuperOfArrayList = cls.isAssignableFrom(ArrayList.class);
       boolean isSuperOfHashSet = cls.isAssignableFrom(HashSet.class);
       if ((!isSuperOfArrayList && !isSuperOfHashSet)
@@ -658,7 +679,7 @@ private static boolean isSupported(TypeRef<?> typeRef, LinkedHashSet<TypeRef> wa
         throw new UnsupportedOperationException(
             "cyclic type is not supported. walkedTypePath: " + walkedTypePath);
       } else {
-        LinkedHashSet<TypeRef> newTypePath = new LinkedHashSet<>(walkedTypePath);
+        LinkedHashSet<TypeRef<?>> newTypePath = new LinkedHashSet<>(walkedTypePath);
         newTypePath.add(typeRef);
         return isBean(typeRef, newTypePath);
       }
@@ -673,13 +694,24 @@ private static boolean isSupported(TypeRef<?> typeRef, LinkedHashSet<TypeRef> wa
    *     parameters recursively
    */
   public static LinkedHashSet<Class<?>> listBeansRecursiveInclusive(Class<?> beanClass) {
-    return listBeansRecursiveInclusive(beanClass, new LinkedHashSet<>());
+    return listBeansRecursiveInclusive(beanClass, CustomTypeRegistry.EMPTY);
+  }
+
+  public static LinkedHashSet<Class<?>> listBeansRecursiveInclusive(
+      Class<?> beanClass, CustomTypeRegistry customTypes) {
+    return listBeansRecursiveInclusive(beanClass, new LinkedHashSet<>(), customTypes);
   }
 
   private static LinkedHashSet<Class<?>> listBeansRecursiveInclusive(
-      Class<?> beanClass, LinkedHashSet<TypeRef<?>> walkedTypePath) {
+      Class<?> beanClass,
+      LinkedHashSet<TypeRef<?>> walkedTypePath,
+      CustomTypeRegistry customTypes) {
     LinkedHashSet<Class<?>> beans = new LinkedHashSet<>();
-    if (isBean(beanClass)) {
+    Class<?> enclosingType = enclosingType(walkedTypePath);
+    if (customTypes.hasCodec(enclosingType, beanClass)) {
+      return beans;
+    }
+    if (isBean(beanClass, customTypes)) {
       beans.add(beanClass);
     }
     LinkedHashSet<TypeRef<?>> typeRefs = new LinkedHashSet<>();
@@ -692,33 +724,34 @@ private static LinkedHashSet<Class<?>> listBeansRecursiveInclusive(
 
     for (TypeRef<?> typeToken : typeRefs) {
       Class<?> type = getRawType(typeToken);
-      if (isBean(type)) {
-        beans.add(type);
+      if (isBean(type, customTypes)) {
         if (walkedTypePath.contains(typeToken)) {
           throw new UnsupportedOperationException(
               "cyclic type is not supported. walkedTypePath: " + walkedTypePath);
         } else {
           LinkedHashSet<TypeRef<?>> newPath = new LinkedHashSet<>(walkedTypePath);
           newPath.add(typeToken);
-          beans.addAll(listBeansRecursiveInclusive(type, newPath));
+          beans.addAll(listBeansRecursiveInclusive(type, newPath, customTypes));
         }
       } else if (isCollection(type)) {
         TypeRef<?> elementType = getElementType(typeToken);
         LinkedHashSet<TypeRef<?>> newPath = new LinkedHashSet<>(walkedTypePath);
         newPath.add(elementType);
-        beans.addAll(listBeansRecursiveInclusive(elementType.getClass(), newPath));
+        beans.addAll(listBeansRecursiveInclusive(elementType.getClass(), newPath, customTypes));
       } else if (isMap(type)) {
         Tuple2<TypeRef<?>, TypeRef<?>> mapKeyValueType = getMapKeyValueType(typeToken);
         LinkedHashSet<TypeRef<?>> newPath = new LinkedHashSet<>(walkedTypePath);
         newPath.add(mapKeyValueType.f0);
         newPath.add(mapKeyValueType.f1);
-        beans.addAll(listBeansRecursiveInclusive(mapKeyValueType.f0.getRawType(), newPath));
-        beans.addAll(listBeansRecursiveInclusive(mapKeyValueType.f1.getRawType(), newPath));
+        beans.addAll(
+            listBeansRecursiveInclusive(mapKeyValueType.f0.getRawType(), newPath, customTypes));
+        beans.addAll(
+            listBeansRecursiveInclusive(mapKeyValueType.f1.getRawType(), newPath, customTypes));
       } else if (type.isArray()) {
         Class<?> arrayComponent = getArrayComponent(type);
         LinkedHashSet<TypeRef<?>> newPath = new LinkedHashSet<>(walkedTypePath);
         newPath.add(TypeRef.of(arrayComponent));
-        beans.addAll(listBeansRecursiveInclusive(arrayComponent, newPath));
+        beans.addAll(listBeansRecursiveInclusive(arrayComponent, newPath, customTypes));
       }
     }
     return beans;
@@ -737,7 +770,7 @@ public static int computeStringHash(String str) {
   }
 
   /** Returns generic type arguments of <code>typeToken</code>. */
-  public static List<TypeRef<?>> getTypeArguments(TypeRef typeRef) {
+  public static List<TypeRef<?>> getTypeArguments(TypeRef<?> typeRef) {
     if (typeRef.getType() instanceof ParameterizedType) {
       ParameterizedType parameterizedType = (ParameterizedType) typeRef.getType();
       return Arrays.stream(parameterizedType.getActualTypeArguments())
@@ -752,7 +785,7 @@ public static List<TypeRef<?>> getTypeArguments(TypeRef typeRef) {
    * Returns generic type arguments of <code>typeToken</code>, includes generic type arguments of
    * generic type arguments recursively.
    */
-  public static List<TypeRef<?>> getAllTypeArguments(TypeRef typeRef) {
+  public static List<TypeRef<?>> getAllTypeArguments(TypeRef<?> typeRef) {
     List<TypeRef<?>> types = getTypeArguments(typeRef);
     LinkedHashSet<TypeRef<?>> allTypeArguments = new LinkedHashSet<>(types);
     for (TypeRef<?> type : types) {
@@ -776,4 +809,12 @@ public static String qualifiedName(String pkg, String className) {
       return pkg + "." + className;
     }
   }
+
+  private static Class<?> enclosingType(LinkedHashSet<TypeRef<?>> newTypePath) {
+    Class<?> result = Object.class;
+    for (TypeRef<?> type : newTypePath) {
+      result = type.getRawType();
+    }
+    return result;
+  }
 }

java/fury-format/README.md

@@ -11,3 +11,45 @@ Fury row format is heavily inspired by spark tungsten row format, but with chang
 - Support adding fields without breaking compatibility.
 
 The initial fury java row data structure implementation is modified from spark unsafe row/writer.
+
+It is possible to register custom type handling and collection factories for the row format -
+see Encoders.registerCustomCodec and Encoders.registerCustomCollectionFactory.
+
+A short example:
+
+```
+@Data
+public static class UuidType {
+  public UUID f1;
+  public UUID[] f2;
+  public SortedSet<UUID> f3;
+}
+
+static class UuidEncoder implements CustomCodec.MemoryBufferCodec<UUID> {
+  @Override
+  public MemoryBuffer encode(final UUID value) {
+    final MemoryBuffer result = MemoryBuffer.newHeapBuffer(16);
+    result.putInt64(0, value.getMostSignificantBits());
+    result.putInt64(8, value.getLeastSignificantBits());
+    return result;
+  }
+
+  @Override
+  public UUID decode(final MemoryBuffer value) {
+    return new UUID(value.readInt64(), value.readInt64());
+  }
+}
+
+static class SortedSetOfUuidDecoder implements CustomCollectionFactory<UUID, SortedSet<UUID>> {
+  @Override
+  public SortedSet<UUID> newCollection(final int size) {
+    return new TreeSet<>(UnsignedUuidComparator.INSTANCE);
+  }
+}
+
+Encoders.registerCustomCodec(UUID.class, new UuidEncoder());
+Encoders.registerCustomCollectionFactory(
+    SortedSet.class, UUID.class, new SortedSetOfUuidDecoder());
+
+RowEncoder<UuidType> encoder = Encoders.bean(UuidType.class);
+```

java/fury-format/src/main/java/org/apache/fury/format/encoder/ArrayDataForEach.java

@@ -30,6 +30,7 @@
 import org.apache.fury.codegen.Expression.AbstractExpression;
 import org.apache.fury.format.row.binary.BinaryArray;
 import org.apache.fury.format.row.binary.BinaryUtils;
+import org.apache.fury.format.type.CustomTypeEncoderRegistry;
 import org.apache.fury.reflect.TypeRef;
 import org.apache.fury.type.TypeUtils;
 import org.apache.fury.util.Preconditions;
@@ -75,8 +76,17 @@ public ArrayDataForEach(
     super(inputArrayData);
     Preconditions.checkArgument(getRawType(inputArrayData.type()) == BinaryArray.class);
     this.inputArrayData = inputArrayData;
-    this.accessMethod = BinaryUtils.getElemAccessMethodName(elemType);
-    this.elemType = BinaryUtils.getElemReturnType(elemType);
+    TypeRef<?> accessType;
+    CustomCodec<?, ?> customEncoder =
+        CustomTypeEncoderRegistry.customTypeHandler()
+            .findCodec(BinaryArray.class, elemType.getRawType());
+    if (customEncoder == null) {
+      accessType = elemType;
+    } else {
+      accessType = TypeRef.of(customEncoder.encodedType());
+    }
+    this.accessMethod = BinaryUtils.getElemAccessMethodName(accessType);
+    this.elemType = BinaryUtils.getElemReturnType(accessType);
     this.notNullAction = notNullAction;
     this.nullAction = nullAction;
   }