Add REMOVE_NULLS function for arrays

presto-docs/src/main/sphinx/functions/array.rst

@@ -259,6 +259,10 @@ Array Functions
                       (s, x) -> CAST(ROW(x + s.sum, s.count + 1) AS ROW(sum DOUBLE, count INTEGER)),
                       s -> IF(s.count = 0, NULL, s.sum / s.count));
 
+.. function:: remove_nulls(array(T)) -> array
+
+    Remove all null elements in the array.
+
 .. function:: repeat(element, count) -> array
 
     Repeat ``element`` for ``count`` times.

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

@@ -121,6 +121,7 @@
 import com.facebook.presto.operator.scalar.ArrayPositionFunction;
 import com.facebook.presto.operator.scalar.ArrayPositionWithIndexFunction;
 import com.facebook.presto.operator.scalar.ArrayRemoveFunction;
+import com.facebook.presto.operator.scalar.ArrayRemoveNullsFunction;
 import com.facebook.presto.operator.scalar.ArrayReverseFunction;
 import com.facebook.presto.operator.scalar.ArrayShuffleFunction;
 import com.facebook.presto.operator.scalar.ArraySliceFunction;
@@ -781,6 +782,7 @@ private List<? extends SqlFunction> getBuildInFunctions(FeaturesConfig featuresC
                 .scalar(ArrayLessThanOperator.class)
                 .scalar(ArrayLessThanOrEqualOperator.class)
                 .scalar(ArrayRemoveFunction.class)
+                .scalar(ArrayRemoveNullsFunction.class)
                 .scalar(ArrayGreaterThanOperator.class)
                 .scalar(ArrayGreaterThanOrEqualOperator.class)
                 .scalar(ArrayElementAtFunction.class)

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

@@ -0,0 +1,70 @@
+/*
+ * 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.common.block.Block;
+import com.facebook.presto.common.block.BlockBuilder;
+import com.facebook.presto.common.type.Type;
+import com.facebook.presto.spi.function.Description;
+import com.facebook.presto.spi.function.ScalarFunction;
+import com.facebook.presto.spi.function.SqlType;
+import com.facebook.presto.spi.function.TypeParameter;
+
+@ScalarFunction("remove_nulls")
+@Description("Removes null values from an array")
+public final class ArrayRemoveNullsFunction
+{
+    @TypeParameter("E")
+    public ArrayRemoveNullsFunction(@TypeParameter("E") Type elementType) {}
+
+    @TypeParameter("E")
+    @SqlType("array(E)")
+    public Block remove(
+            @TypeParameter("E") Type type,
+            @SqlType("array(E)") Block array)
+    {
+        if (!array.mayHaveNull()) {
+            return array;
+        }
+
+        int found = -1;
+        for (int i = 0; i < array.getPositionCount(); i++) {
+            if (array.isNull(i)) {
+                found = i;
+                break;
+            }
+        }
+
+        if (found == -1) {
+            // all elements are non-null
+            return array;
+        }
+
+        BlockBuilder blockBuilder = type.createBlockBuilder(null, array.getPositionCount() - 1);
+
+        // copy all elements up to found-1
+        for (int i = 0; i < found; i++) {
+            type.appendTo(array, i, blockBuilder);
+        }
+
+        // and then copy non-null elements from found+1 till the end of the array
+        for (int i = found + 1; i < array.getPositionCount(); i++) {
+            if (!array.isNull(i)) {
+                type.appendTo(array, i, blockBuilder);
+            }
+        }
+
+        return blockBuilder.build();
+    }
+}

presto-main/src/test/java/com/facebook/presto/type/TestArrayOperators.java

@@ -1591,6 +1591,19 @@ public void testArrayRemove()
         assertInvalidFunction("ARRAY_REMOVE(ARRAY [ARRAY[CAST(1 AS BIGINT)]], ARRAY[CAST(null AS BIGINT)])", NOT_SUPPORTED);
     }
 
+    @Test
+    public void testRemoveNulls()
+    {
+        assertFunction("REMOVE_NULLS(ARRAY ['foo', 'bar'])", new ArrayType(createVarcharType(3)), asList("foo", "bar"));
+        assertFunction("REMOVE_NULLS(ARRAY ['foo', NULL, 'bar'])", new ArrayType(createVarcharType(3)), asList("foo", "bar"));
+        assertFunction("REMOVE_NULLS(ARRAY [1, NULL, NULL, 3])", new ArrayType(INTEGER), asList(1, 3));
+        assertFunction("REMOVE_NULLS(ARRAY [ARRAY ['foo'], NULL, ARRAY['bar']])", new ArrayType(new ArrayType(createVarcharType(3))), asList(singletonList("foo"), singletonList("bar")));
+        assertFunction("REMOVE_NULLS(ARRAY [TRUE, FALSE, TRUE])", new ArrayType(BOOLEAN), asList(true, false, true));
+        assertFunction("REMOVE_NULLS(ARRAY [TRUE, FALSE, NULL])", new ArrayType(BOOLEAN), asList(true, false));
+        assertFunction("REMOVE_NULLS(ARRAY [ARRAY[NULL]])", new ArrayType(new ArrayType(UNKNOWN)), ImmutableList.of(singletonList(null)));
+        assertFunction("REMOVE_NULLS(ARRAY [ARRAY[NULL], NULL])", new ArrayType(new ArrayType(UNKNOWN)), ImmutableList.of(singletonList(null)));
+    }
+
     @Test
     public void testRepeat()
     {