Make array_except use 'is distinct' semantics

trinodb · Apr 7, 2022 · 06be474 · 06be474
1 parent a92c22c
commit 06be474
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/ArrayExceptFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/ArrayExceptFunction.java
@@ -23,15 +23,14 @@
 import io.trino.spi.function.SqlType;
 import io.trino.spi.function.TypeParameter;
 import io.trino.spi.type.Type;
-import io.trino.type.BlockTypeOperators.BlockPositionEqual;
 import io.trino.type.BlockTypeOperators.BlockPositionHashCode;
+import io.trino.type.BlockTypeOperators.BlockPositionIsDistinctFrom;
 
-import static io.trino.operator.aggregation.TypedSet.createEqualityTypedSet;
+import static io.trino.operator.aggregation.TypedSet.createDistinctTypedSet;
 import static io.trino.spi.function.InvocationConvention.InvocationArgumentConvention.BLOCK_POSITION;
 import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.FAIL_ON_NULL;
-import static io.trino.spi.function.InvocationConvention.InvocationReturnConvention.NULLABLE_RETURN;
-import static io.trino.spi.function.OperatorType.EQUAL;
 import static io.trino.spi.function.OperatorType.HASH_CODE;
+import static io.trino.spi.function.OperatorType.IS_DISTINCT_FROM;
 
 @ScalarFunction("array_except")
 @Description("Returns an array of elements that are in the first array but not the second, without duplicates.")
@@ -44,9 +43,9 @@ private ArrayExceptFunction() {}
     public static Block except(
             @TypeParameter("E") Type type,
             @OperatorDependency(
-                    operator = EQUAL,
+                    operator = IS_DISTINCT_FROM,
                     argumentTypes = {"E", "E"},
-                    convention = @Convention(arguments = {BLOCK_POSITION, BLOCK_POSITION}, result = NULLABLE_RETURN)) BlockPositionEqual elementEqual,
+                    convention = @Convention(arguments = {BLOCK_POSITION, BLOCK_POSITION}, result = FAIL_ON_NULL)) BlockPositionIsDistinctFrom isDistinctOperator,
             @OperatorDependency(
                     operator = HASH_CODE,
                     argumentTypes = "E",
@@ -61,7 +60,7 @@ public static Block except(
             return leftArray;
         }
 
-        TypedSet typedSet = createEqualityTypedSet(type, elementEqual, elementHashCode, leftPositionCount, "array_except");
+        TypedSet typedSet = createDistinctTypedSet(type, isDistinctOperator, elementHashCode, leftPositionCount, "array_except");
         BlockBuilder distinctElementBlockBuilder = type.createBlockBuilder(null, leftPositionCount);
         for (int i = 0; i < rightPositionCount; i++) {
             typedSet.add(rightArray, i);

diff --git a/core/trino-main/src/test/java/io/trino/operator/scalar/TestArrayExceptFunction.java b/core/trino-main/src/test/java/io/trino/operator/scalar/TestArrayExceptFunction.java
@@ -66,4 +66,11 @@ public void testDuplicates()
         assertFunction("array_except(ARRAY[VARCHAR 'x', 'x', 'y', 'z'], ARRAY['x', 'y', 'x'])", new ArrayType(VARCHAR), ImmutableList.of("z"));
         assertFunction("array_except(ARRAY[true, false, null, true, false, null], ARRAY[true, true, true])", new ArrayType(BOOLEAN), asList(false, null));
     }
+
+    @Test
+    public void testNonDistinctNonEqualValues()
+    {
+        assertFunction("array_except(ARRAY[NaN()], ARRAY[NaN()])", new ArrayType(DOUBLE), ImmutableList.of());
+        assertFunction("array_except(ARRAY[1, NaN(), 3], ARRAY[NaN(), 3])", new ArrayType(DOUBLE), ImmutableList.of(1.0));
+    }
 }