Use boolean return from TypedSet.add instead of TypedSet.contains

trinodb · Nov 22, 2021 · 0f78fa9 · 0f78fa9
1 parent 4fa5116
commit 0f78fa9
Show file tree

Hide file tree

Showing 9 changed files with 12 additions and 21 deletions.
diff --git a/.../src/main/java/io/trino/operator/aggregation/multimapagg/MultimapAggregationFunction.java b/.../src/main/java/io/trino/operator/aggregation/multimapagg/MultimapAggregationFunction.java
@@ -182,8 +182,7 @@ public static void output(Type keyType, BlockPositionEqual keyEqual, BlockPositi
 
             state.forEach((key, value, keyValueIndex) -> {
                 // Merge values of the same key into an array
-                if (!keySet.contains(key, keyValueIndex)) {
-                    keySet.add(key, keyValueIndex);
+                if (keySet.add(key, keyValueIndex)) {
                     keyType.appendTo(key, keyValueIndex, distinctKeyBlockBuilder);
                     BlockBuilder valueArrayBuilder = valueType.createBlockBuilder(null, 10, expectedValueSize(valueType, EXPECTED_ENTRY_SIZE));
                     valueArrayBlockBuilders.set(keySet.positionOf(key, keyValueIndex), valueArrayBuilder);

diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/ArrayDistinctFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/ArrayDistinctFunction.java
@@ -84,8 +84,7 @@ public Block distinct(
 
         BlockBuilder distinctElementBlockBuilder = pageBuilder.getBlockBuilder(0);
         for (int i = 0; i < array.getPositionCount(); i++) {
-            if (!typedSet.contains(array, i)) {
-                typedSet.add(array, i);
+            if (typedSet.add(array, i)) {
                 distinctCount++;
                 type.appendTo(array, i, distinctElementBlockBuilder);
             }

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
@@ -66,8 +66,7 @@ public static Block except(
             typedSet.add(rightArray, i);
         }
         for (int i = 0; i < leftPositionCount; i++) {
-            if (!typedSet.contains(leftArray, i)) {
-                typedSet.add(leftArray, i);
+            if (typedSet.add(leftArray, i)) {
                 type.appendTo(leftArray, i, distinctElementBlockBuilder);
             }
         }

diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/ArrayUnionFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/ArrayUnionFunction.java
@@ -72,8 +72,7 @@ public static Block union(
     private static void appendTypedArray(Block array, Type type, TypedSet typedSet, BlockBuilder blockBuilder)
     {
         for (int i = 0; i < array.getPositionCount(); i++) {
-            if (!typedSet.contains(array, i)) {
-                typedSet.add(array, i);
+            if (typedSet.add(array, i)) {
                 type.appendTo(array, i, blockBuilder);
             }
         }

diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapConcatFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapConcatFunction.java
@@ -160,8 +160,7 @@ public static Block mapConcat(MapType mapType, BlockPositionEqual keyEqual, Bloc
         for (int idx = lastMapIndex - 1; idx > firstMapIndex; idx--) {
             map = maps[idx];
             for (int i = 0; i < map.getPositionCount(); i += 2) {
-                if (!typedSet.contains(map, i)) {
-                    typedSet.add(map, i);
+                if (typedSet.add(map, i)) {
                     keyType.appendTo(map, i, blockBuilder);
                     valueType.appendTo(map, i + 1, blockBuilder);
                 }

diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapFromEntriesFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapFromEntriesFunction.java
@@ -100,12 +100,11 @@ public Block mapFromEntries(
                 throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null");
             }
 
-            if (uniqueKeys.contains(rowBlock, 0)) {
+            if (!uniqueKeys.add(rowBlock, 0)) {
                 mapBlockBuilder.closeEntry();
                 pageBuilder.declarePosition();
                 throw new TrinoException(INVALID_FUNCTION_ARGUMENT, format("Duplicate keys (%s) are not allowed", keyType.getObjectValue(session, rowBlock, 0)));
             }
-            uniqueKeys.add(rowBlock, 0);
 
             keyType.appendTo(rowBlock, 0, resultBuilder);
             valueType.appendTo(rowBlock, 1, resultBuilder);

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
@@ -258,8 +258,7 @@ public static Block mapCast(
         BlockBuilder mapBlockBuilder = toMapType.createBlockBuilder(null, 1);
         BlockBuilder blockBuilder = mapBlockBuilder.beginBlockEntry();
         for (int i = 0; i < fromMap.getPositionCount(); i += 2) {
-            if (!typedSet.contains(keyBlock, i / 2)) {
-                typedSet.add(keyBlock, i / 2);
+            if (typedSet.add(keyBlock, i / 2)) {
                 toKeyType.appendTo(keyBlock, i / 2, blockBuilder);
                 if (fromMap.isNull(i + 1)) {
                     blockBuilder.appendNull();

diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MapTransformKeysFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MapTransformKeysFunction.java
@@ -276,9 +276,8 @@ private MethodHandle generateTransformKey(Type keyType, Type transformedKeyType,
                         .append(loadValueElement)
                         .append(writeKeyElement)
                         .append(new IfStatement()
-                                .condition(typedSet.invoke("contains", boolean.class, blockBuilder.cast(Block.class), position))
-                                .ifTrue(throwDuplicatedKeyException)
-                                .ifFalse(typedSet.invoke("add", boolean.class, blockBuilder.cast(Block.class), position)))));
+                                .condition(typedSet.invoke("add", boolean.class, blockBuilder.cast(Block.class), position))
+                                .ifFalse(throwDuplicatedKeyException))));
 
         body.append(mapBlockBuilder
                 .invoke("closeEntry", BlockBuilder.class)

diff --git a/core/trino-main/src/main/java/io/trino/operator/scalar/MultimapFromEntriesFunction.java b/core/trino-main/src/main/java/io/trino/operator/scalar/MultimapFromEntriesFunction.java
@@ -104,12 +104,11 @@ public Block multimapFromEntries(
                 throw new TrinoException(INVALID_FUNCTION_ARGUMENT, "map key cannot be null");
             }
 
-            if (keySet.contains(rowBlock, 0)) {
-                entryIndicesList[keySet.positionOf(rowBlock, 0)].add(i);
+            if (keySet.add(rowBlock, 0)) {
+                entryIndicesList[keySet.size() - 1].add(i);
             }
             else {
-                keySet.add(rowBlock, 0);
-                entryIndicesList[keySet.size() - 1].add(i);
+                entryIndicesList[keySet.positionOf(rowBlock, 0)].add(i);
             }
         }