[Prototype & Benchmark] Support directly writing to output block builder for scalar functions

presto-main/src/main/java/com/facebook/presto/metadata/FunctionRegistry.java

@@ -214,6 +214,8 @@
 import static com.facebook.presto.operator.scalar.ArrayConcatFunction.ARRAY_CONCAT_FUNCTION;
 import static com.facebook.presto.operator.scalar.ArrayConstructor.ARRAY_CONSTRUCTOR;
 import static com.facebook.presto.operator.scalar.ArrayFlattenFunction.ARRAY_FLATTEN_FUNCTION;
+import static com.facebook.presto.operator.scalar.ArrayIdentityDirectFunction.ARRAY_IDENTITY_DIRECT_FUNCTION;
+import static com.facebook.presto.operator.scalar.ArrayIdentityFunction.ARRAY_IDENTITY_FUNCTION;
 import static com.facebook.presto.operator.scalar.ArrayJoin.ARRAY_JOIN;
 import static com.facebook.presto.operator.scalar.ArrayJoin.ARRAY_JOIN_WITH_NULL_REPLACEMENT;
 import static com.facebook.presto.operator.scalar.ArrayReduceFunction.ARRAY_REDUCE_FUNCTION;
@@ -546,6 +548,7 @@ public WindowFunctionSupplier load(SpecializedFunctionKey key)
                 .function(MAP_ELEMENT_AT)
                 .function(MAP_CONCAT_FUNCTION)
                 .function(ARRAY_FLATTEN_FUNCTION)
+                .functions(ARRAY_IDENTITY_DIRECT_FUNCTION, ARRAY_IDENTITY_FUNCTION)
                 .function(ARRAY_CONCAT_FUNCTION)
                 .functions(ARRAY_CONSTRUCTOR, ARRAY_SUBSCRIPT, ARRAY_TO_JSON, JSON_TO_ARRAY, JSON_STRING_TO_ARRAY)
                 .functions(new MapSubscriptOperator(featuresConfig.isLegacyMapSubscript()))

presto-main/src/main/java/com/facebook/presto/operator/scalar/ArrayConcatFunction.java

@@ -101,6 +101,7 @@ public ScalarFunctionImplementation specialize(BoundVariables boundVariables, in
                 nCopies(arity, Optional.empty()),
                 methodHandleAndConstructor.getMethodHandle(),
                 Optional.of(methodHandleAndConstructor.getConstructor()),
+                false,
                 isDeterministic());
     }
 

presto-main/src/main/java/com/facebook/presto/operator/scalar/ArrayIdentityDirectFunction.java

@@ -0,0 +1,102 @@
+/*
+ * Licensed 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 com.facebook.presto.operator.scalar;
+
+import com.facebook.presto.metadata.BoundVariables;
+import com.facebook.presto.metadata.FunctionKind;
+import com.facebook.presto.metadata.FunctionRegistry;
+import com.facebook.presto.metadata.Signature;
+import com.facebook.presto.metadata.SqlScalarFunction;
+import com.facebook.presto.spi.block.Block;
+import com.facebook.presto.spi.block.BlockBuilder;
+import com.facebook.presto.spi.type.StandardTypes;
+import com.facebook.presto.spi.type.Type;
+import com.facebook.presto.spi.type.TypeManager;
+import com.facebook.presto.spi.type.TypeSignatureParameter;
+import com.google.common.collect.ImmutableList;
+
+import java.lang.invoke.MethodHandle;
+import java.util.Optional;
+
+import static com.facebook.presto.metadata.Signature.typeVariable;
+import static com.facebook.presto.spi.type.TypeSignature.parseTypeSignature;
+import static com.facebook.presto.util.Reflection.methodHandle;
+
+public class ArrayIdentityDirectFunction
+        extends SqlScalarFunction
+{
+    public static final ArrayIdentityDirectFunction ARRAY_IDENTITY_DIRECT_FUNCTION = new ArrayIdentityDirectFunction();
+    private static final String FUNCTION_NAME = "array_identity_direct";
+    private static final MethodHandle METHOD_HANDLE = methodHandle(ArrayIdentityDirectFunction.class, "identity", Type.class, Type.class, BlockBuilder.class, Block.class);
+
+    private ArrayIdentityDirectFunction()
+    {
+        super(new Signature(FUNCTION_NAME,
+                FunctionKind.SCALAR,
+                ImmutableList.of(typeVariable("E")),
+                ImmutableList.of(),
+                parseTypeSignature("array(E)"),
+                ImmutableList.of(parseTypeSignature("array(E)")),
+                false));
+    }
+
+    @Override
+    public boolean isHidden()
+    {
+        return false;
+    }
+
+    @Override
+    public boolean isDeterministic()
+    {
+        return true;
+    }
+
+    @Override
+    public String getDescription()
+    {
+        return "array identity";
+    }
+
+    @Override
+    public ScalarFunctionImplementation specialize(BoundVariables boundVariables, int arity, TypeManager typeManager, FunctionRegistry functionRegistry)
+    {
+        Type elementType = boundVariables.getTypeVariable("E");
+        Type arrayType = typeManager.getParameterizedType(StandardTypes.ARRAY, ImmutableList.of(TypeSignatureParameter.of(elementType.getTypeSignature())));
+        MethodHandle methodHandle = METHOD_HANDLE.bindTo(elementType).bindTo(arrayType);
+        return new ScalarFunctionImplementation(
+                false,
+                ImmutableList.of(false),
+                ImmutableList.of(false),
+                ImmutableList.of(Optional.empty()),
+                methodHandle,
+                Optional.empty(),
+                true,
+                isDeterministic());
+    }
+
+    public static void identity(Type type, Type arrayType, BlockBuilder outputBlock, Block array)
+    {
+        BlockBuilder entryBlockBuilder = outputBlock.beginBlockEntry();
+        for (int i = 0; i < array.getPositionCount(); i++) {
+            if (array.isNull(i)) {
+                entryBlockBuilder.appendNull();
+            }
+            else {
+                type.appendTo(array, i, entryBlockBuilder);
+            }
+        }
+        outputBlock.closeEntry();
+    }
+}

presto-main/src/main/java/com/facebook/presto/operator/scalar/ArrayIdentityFunction.java

@@ -0,0 +1,122 @@
+/*
+ * Licensed 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 com.facebook.presto.operator.scalar;
+
+import com.facebook.presto.annotation.UsedByGeneratedCode;
+import com.facebook.presto.metadata.BoundVariables;
+import com.facebook.presto.metadata.FunctionKind;
+import com.facebook.presto.metadata.FunctionRegistry;
+import com.facebook.presto.metadata.Signature;
+import com.facebook.presto.metadata.SqlScalarFunction;
+import com.facebook.presto.spi.PageBuilder;
+import com.facebook.presto.spi.block.Block;
+import com.facebook.presto.spi.block.BlockBuilder;
+import com.facebook.presto.spi.type.ArrayType;
+import com.facebook.presto.spi.type.StandardTypes;
+import com.facebook.presto.spi.type.Type;
+import com.facebook.presto.spi.type.TypeManager;
+import com.facebook.presto.spi.type.TypeSignatureParameter;
+import com.google.common.collect.ImmutableList;
+
+import java.lang.invoke.MethodHandle;
+import java.util.Optional;
+
+import static com.facebook.presto.metadata.Signature.typeVariable;
+import static com.facebook.presto.spi.type.TypeSignature.parseTypeSignature;
+import static com.facebook.presto.util.Reflection.methodHandle;
+
+public class ArrayIdentityFunction
+        extends SqlScalarFunction
+{
+    public static final ArrayIdentityFunction ARRAY_IDENTITY_FUNCTION = new ArrayIdentityFunction();
+    private static final String FUNCTION_NAME = "array_identity";
+
+    private static final MethodHandle METHOD_HANDLE = methodHandle(ArrayIdentityFunction.class, "identity", Type.class, Type.class, Object.class, Block.class);
+    private static final MethodHandle STATE_FACTORY = methodHandle(ArrayIdentityFunction.class, "createState", ArrayType.class);
+
+    private ArrayIdentityFunction()
+    {
+        super(new Signature(FUNCTION_NAME,
+                FunctionKind.SCALAR,
+                ImmutableList.of(typeVariable("E")),
+                ImmutableList.of(),
+                parseTypeSignature("array(E)"),
+                ImmutableList.of(parseTypeSignature("array(E)")),
+                false));
+    }
+
+    @Override
+    public boolean isHidden()
+    {
+        return false;
+    }
+
+    @Override
+    public boolean isDeterministic()
+    {
+        return true;
+    }
+
+    @Override
+    public String getDescription()
+    {
+        return "array identity";
+    }
+
+    @Override
+    public ScalarFunctionImplementation specialize(BoundVariables boundVariables, int arity, TypeManager typeManager, FunctionRegistry functionRegistry)
+    {
+        Type elementType = boundVariables.getTypeVariable("E");
+        Type arrayType = typeManager.getParameterizedType(StandardTypes.ARRAY, ImmutableList.of(TypeSignatureParameter.of(elementType.getTypeSignature())));
+        MethodHandle methodHandle = METHOD_HANDLE.bindTo(elementType).bindTo(arrayType);
+        return new ScalarFunctionImplementation(
+                false,
+                ImmutableList.of(false),
+                ImmutableList.of(false),
+                ImmutableList.of(Optional.empty()),
+                methodHandle,
+                Optional.of(STATE_FACTORY.bindTo(arrayType)),
+                false,
+                isDeterministic());
+    }
+
+    @UsedByGeneratedCode
+    public static Object createState(ArrayType arrayType)
+    {
+        return new PageBuilder(ImmutableList.of(arrayType));
+    }
+
+    public static Block identity(Type type, Type arrayType, Object state, Block array)
+    {
+        PageBuilder pageBuilder = (PageBuilder) state;
+        if (pageBuilder.isFull()) {
+            pageBuilder.reset();
+        }
+        BlockBuilder arrayBlockBuilder = pageBuilder.getBlockBuilder(0);
+        BlockBuilder blockBuilder = arrayBlockBuilder.beginBlockEntry();
+
+        for (int i = 0; i < array.getPositionCount(); i++) {
+            if (array.isNull(i)) {
+                blockBuilder.appendNull();
+            }
+            else {
+                type.appendTo(array, i, blockBuilder);
+            }
+        }
+
+        arrayBlockBuilder.closeEntry();
+        pageBuilder.declarePosition();
+        return (Block) arrayType.getObject(arrayBlockBuilder, arrayBlockBuilder.getPositionCount() - 1);
+    }
+}

presto-main/src/main/java/com/facebook/presto/operator/scalar/ArrayTransformFunction.java

@@ -111,6 +111,7 @@ public ScalarFunctionImplementation specialize(BoundVariables boundVariables, in
                 ImmutableList.of(Optional.empty(), Optional.of(UnaryFunctionInterface.class)),
                 methodHandle(generatedClass, "transform", PageBuilder.class, Block.class, UnaryFunctionInterface.class),
                 Optional.of(methodHandle(generatedClass, "createPageBuilder")),
+                false,
                 isDeterministic());
     }
 

presto-main/src/main/java/com/facebook/presto/operator/scalar/MapConcatFunction.java

@@ -110,6 +110,7 @@ public ScalarFunctionImplementation specialize(BoundVariables boundVariables, in
                 nCopies(arity, Optional.empty()),
                 methodHandleAndConstructor.getMethodHandle(),
                 Optional.of(methodHandleAndConstructor.getConstructor()),
+                false,
                 isDeterministic());
     }
 

presto-main/src/main/java/com/facebook/presto/operator/scalar/MapConstructor.java

@@ -98,6 +98,7 @@ public ScalarFunctionImplementation specialize(BoundVariables boundVariables, in
                 ImmutableList.of(Optional.empty(), Optional.empty()),
                 METHOD_HANDLE.bindTo(mapType).bindTo(keyEqual).bindTo(keyHashCode),
                 Optional.of(instanceFactory),
+                false,
                 isDeterministic());
     }
 

presto-main/src/main/java/com/facebook/presto/operator/scalar/MapFilterFunction.java

@@ -119,6 +119,7 @@ public ScalarFunctionImplementation specialize(BoundVariables boundVariables, in
                 ImmutableList.of(Optional.empty(), Optional.of(BinaryFunctionInterface.class)),
                 generateFilter(mapType),
                 Optional.of(STATE_FACTORY.bindTo(mapType)),
+                false,
                 isDeterministic());
     }
 

presto-main/src/main/java/com/facebook/presto/operator/scalar/MapTransformKeyFunction.java

@@ -130,6 +130,7 @@ public ScalarFunctionImplementation specialize(BoundVariables boundVariables, in
                 ImmutableList.of(Optional.empty(), Optional.of(BinaryFunctionInterface.class)),
                 generateTransformKey(keyType, transformedKeyType, valueType, resultMapType),
                 Optional.of(STATE_FACTORY.bindTo(resultMapType)),
+                false,
                 isDeterministic());
     }
 

presto-main/src/main/java/com/facebook/presto/operator/scalar/MapTransformValueFunction.java

@@ -125,6 +125,7 @@ public ScalarFunctionImplementation specialize(BoundVariables boundVariables, in
                 ImmutableList.of(Optional.empty(), Optional.of(BinaryFunctionInterface.class)),
                 generateTransform(keyType, valueType, transformedValueType, resultMapType),
                 Optional.of(STATE_FACTORY.bindTo(resultMapType)),
+                false,
                 isDeterministic());
     }
 

presto-main/src/main/java/com/facebook/presto/operator/scalar/ParametricScalar.java

@@ -89,6 +89,7 @@ public ScalarFunctionImplementation specialize(BoundVariables boundVariables, in
                     implementation.getLambdaInterface(),
                     methodHandleAndConstructor.get().getMethodHandle(),
                     methodHandleAndConstructor.get().getConstructor(),
+                    false,
                     isDeterministic());
         }
 
@@ -104,6 +105,7 @@ public ScalarFunctionImplementation specialize(BoundVariables boundVariables, in
                         implementation.getLambdaInterface(),
                         methodHandle.get().getMethodHandle(),
                         methodHandle.get().getConstructor(),
+                        false,
                         isDeterministic());
             }
         }
@@ -122,6 +124,7 @@ public ScalarFunctionImplementation specialize(BoundVariables boundVariables, in
                         implementation.getLambdaInterface(),
                         methodHandle.get().getMethodHandle(),
                         methodHandle.get().getConstructor(),
+                        false,
                         isDeterministic());
             }
         }

presto-main/src/main/java/com/facebook/presto/operator/scalar/ScalarFunctionImplementation.java

@@ -33,6 +33,7 @@ public final class ScalarFunctionImplementation
     private final List<Optional<Class>> lambdaInterface;
     private final MethodHandle methodHandle;
     private final Optional<MethodHandle> instanceFactory;
+    private final boolean writeToBlockBuilderParamater;
     private final boolean deterministic;
 
     public ScalarFunctionImplementation(boolean nullable, List<Boolean> nullableArguments, MethodHandle methodHandle, boolean deterministic)
@@ -44,6 +45,7 @@ public ScalarFunctionImplementation(boolean nullable, List<Boolean> nullableArgu
                 nCopies(nullableArguments.size(), Optional.empty()),
                 methodHandle,
                 Optional.empty(),
+                false,
                 deterministic);
     }
 
@@ -56,6 +58,7 @@ public ScalarFunctionImplementation(boolean nullable, List<Boolean> nullableArgu
                 nCopies(nullableArguments.size(), Optional.empty()),
                 methodHandle,
                 Optional.empty(),
+                false,
                 deterministic);
     }
 
@@ -74,6 +77,7 @@ public ScalarFunctionImplementation(
                 lambdaInterface,
                 methodHandle,
                 Optional.empty(),
+                false,
                 deterministic);
     }
 
@@ -84,6 +88,7 @@ public ScalarFunctionImplementation(
             List<Optional<Class>> lambdaInterface,
             MethodHandle methodHandle,
             Optional<MethodHandle> instanceFactory,
+            boolean writeToBlockBuilderParamater,
             boolean deterministic)
     {
         this.nullable = nullable;
@@ -92,6 +97,7 @@ public ScalarFunctionImplementation(
         this.lambdaInterface = ImmutableList.copyOf(requireNonNull(lambdaInterface, "lambdaInterface is null"));
         this.methodHandle = requireNonNull(methodHandle, "methodHandle is null");
         this.instanceFactory = requireNonNull(instanceFactory, "instanceFactory is null");
+        this.writeToBlockBuilderParamater = writeToBlockBuilderParamater;
         this.deterministic = deterministic;
 
         if (instanceFactory.isPresent()) {
@@ -146,6 +152,11 @@ public Optional<MethodHandle> getInstanceFactory()
         return instanceFactory;
     }
 
+    public boolean isWriteToBlockBuilderParamater()
+    {
+        return writeToBlockBuilderParamater;
+    }
+
     public boolean isDeterministic()
     {
         return deterministic;