diff --git a/presto-docs/src/main/sphinx/functions/array.rst b/presto-docs/src/main/sphinx/functions/array.rst index 3b91ef71d5572..a3e85481c05f4 100644 --- a/presto-docs/src/main/sphinx/functions/array.rst +++ b/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. diff --git a/presto-main/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java b/presto-main/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java index 1bb27a6b19acb..aaa0fc260961e 100644 --- a/presto-main/src/main/java/com/facebook/presto/metadata/BuiltInTypeAndFunctionNamespaceManager.java +++ b/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 getBuildInFunctions(FeaturesConfig featuresC .scalar(ArrayLessThanOperator.class) .scalar(ArrayLessThanOrEqualOperator.class) .scalar(ArrayRemoveFunction.class) + .scalar(ArrayRemoveNullsFunction.class) .scalar(ArrayGreaterThanOperator.class) .scalar(ArrayGreaterThanOrEqualOperator.class) .scalar(ArrayElementAtFunction.class) diff --git a/presto-main/src/main/java/com/facebook/presto/operator/scalar/ArrayRemoveNullsFunction.java b/presto-main/src/main/java/com/facebook/presto/operator/scalar/ArrayRemoveNullsFunction.java new file mode 100644 index 0000000000000..617455b98cb8c --- /dev/null +++ b/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(); + } +} diff --git a/presto-main/src/test/java/com/facebook/presto/type/TestArrayOperators.java b/presto-main/src/test/java/com/facebook/presto/type/TestArrayOperators.java index 6afee6293d8be..05d05c64314e4 100644 --- a/presto-main/src/test/java/com/facebook/presto/type/TestArrayOperators.java +++ b/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() {