Move MethodHandleUtil.nativeValueWriter() to MapToMapCast

trinodb · Aug 9, 2022 · 76ba96b · 76ba96b
1 parent e4e0d19
commit 76ba96b
Showing 1 changed file with 29 additions and 1 deletion.
diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapToMapCast.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapToMapCast.java
@@ -44,7 +44,6 @@
 import static io.trino.operator.aggregation.TypedSet.createDistinctTypedSet;
 import static io.trino.spi.StandardErrorCode.INVALID_CAST_ARGUMENT;
 import static io.trino.spi.block.MethodHandleUtil.compose;
-import static io.trino.spi.block.MethodHandleUtil.nativeValueWriter;
 import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION;
 import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.NEVER_NULL;
 import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL;
@@ -77,6 +76,11 @@ public final class MapToMapCast
     private static final MethodHandle CHECK_SLICE_IS_NOT_NULL = methodHandle(MapToMapCast.class, "checkSliceIsNotNull", Slice.class);
     private static final MethodHandle CHECK_BLOCK_IS_NOT_NULL = methodHandle(MapToMapCast.class, "checkBlockIsNotNull", Block.class);
 
+    private static final MethodHandle WRITE_LONG = methodHandle(Type.class, "writeLong", BlockBuilder.class, long.class);
+    private static final MethodHandle WRITE_DOUBLE = methodHandle(Type.class, "writeDouble", BlockBuilder.class, double.class);
+    private static final MethodHandle WRITE_BOOLEAN = methodHandle(Type.class, "writeBoolean", BlockBuilder.class, boolean.class);
+    private static final MethodHandle WRITE_OBJECT = methodHandle(Type.class, "writeObject", BlockBuilder.class, Object.class);
+
     private final BlockTypeOperators blockTypeOperators;
 
     public MapToMapCast(BlockTypeOperators blockTypeOperators)
@@ -287,4 +291,28 @@ public static Block mapCast(
         mapBlockBuilder.closeEntry();
         return (Block) targetType.getObject(mapBlockBuilder, mapBlockBuilder.getPositionCount() - 1);
     }
+
+    public static MethodHandle nativeValueWriter(Type type)
+    {
+        Class<?> javaType = type.getJavaType();
+
+        MethodHandle methodHandle;
+        if (javaType == long.class) {
+            methodHandle = WRITE_LONG;
+        }
+        else if (javaType == double.class) {
+            methodHandle = WRITE_DOUBLE;
+        }
+        else if (javaType == boolean.class) {
+            methodHandle = WRITE_BOOLEAN;
+        }
+        else if (!javaType.isPrimitive()) {
+            methodHandle = WRITE_OBJECT;
+        }
+        else {
+            throw new IllegalArgumentException("Unknown java type " + javaType + " from type " + type);
+        }
+
+        return methodHandle.bindTo(type);
+    }
 }