Fix MAP and ARRAY flat write

core/trino-main/src/test/java/io/trino/type/TestBigintVarcharMapType.java

@@ -41,6 +41,7 @@ public static Block createTestBlock(Type mapType)
         BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 2);
         mapType.writeObject(blockBuilder, mapBlockOf(BIGINT, VARCHAR, ImmutableMap.of(1, "hi")));
         mapType.writeObject(blockBuilder, mapBlockOf(BIGINT, VARCHAR, ImmutableMap.of(1, "2", 2, "hello")));
+        mapType.writeObject(blockBuilder, mapBlockOf(BIGINT, VARCHAR, ImmutableMap.of(1, "123456789012345", 2, "hello-world-hello-world-hello-world")));
         return blockBuilder.build();
     }
 

core/trino-main/src/test/java/io/trino/type/TestIntegerVarcharMapType.java

@@ -41,6 +41,7 @@ public static Block createTestBlock(Type mapType)
         BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 2);
         mapType.writeObject(blockBuilder, mapBlockOf(INTEGER, VARCHAR, ImmutableMap.of(1, "hi")));
         mapType.writeObject(blockBuilder, mapBlockOf(INTEGER, VARCHAR, ImmutableMap.of(1, "2", 2, "hello")));
+        mapType.writeObject(blockBuilder, mapBlockOf(INTEGER, VARCHAR, ImmutableMap.of(1, "123456789012345", 2, "hello-world-hello-world-hello-world")));
         return blockBuilder.build();
     }
 

core/trino-main/src/test/java/io/trino/type/TestSmallintVarcharMapType.java

@@ -41,6 +41,7 @@ public static Block createTestBlock(Type mapType)
         BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 2);
         mapType.writeObject(blockBuilder, mapBlockOf(SMALLINT, VARCHAR, ImmutableMap.of(1, "hi")));
         mapType.writeObject(blockBuilder, mapBlockOf(SMALLINT, VARCHAR, ImmutableMap.of(1, "2", 2, "hello")));
+        mapType.writeObject(blockBuilder, mapBlockOf(SMALLINT, VARCHAR, ImmutableMap.of(1, "123456789012345", 2, "hello-world-hello-world-hello-world")));
         return blockBuilder.build();
     }
 

core/trino-main/src/test/java/io/trino/type/TestTinyintVarcharMapType.java

@@ -41,6 +41,7 @@ public static Block createTestBlock(Type mapType)
         BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 2);
         mapType.writeObject(blockBuilder, mapBlockOf(TINYINT, VARCHAR, ImmutableMap.of(1, "hi")));
         mapType.writeObject(blockBuilder, mapBlockOf(TINYINT, VARCHAR, ImmutableMap.of(1, "2", 2, "hello")));
+        mapType.writeObject(blockBuilder, mapBlockOf(TINYINT, VARCHAR, ImmutableMap.of(1, "123456789012345", 2, "hello-world-hello-world-hello-world")));
         return blockBuilder.build();
     }
 

core/trino-main/src/test/java/io/trino/type/TestVarcharArrayType.java

@@ -0,0 +1,80 @@
+/*
+ * 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 io.trino.type;
+
+import io.trino.spi.block.Block;
+import io.trino.spi.block.BlockBuilder;
+import io.trino.spi.type.Type;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+
+import static io.airlift.slice.Slices.utf8Slice;
+import static io.trino.spi.type.TypeSignature.arrayType;
+import static io.trino.spi.type.VarcharType.VARCHAR;
+import static io.trino.type.InternalTypeManager.TESTING_TYPE_MANAGER;
+import static io.trino.util.StructuralTestUtil.arrayBlockOf;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class TestVarcharArrayType
+        extends AbstractTestType
+{
+    public TestVarcharArrayType()
+    {
+        super(TESTING_TYPE_MANAGER.getType(arrayType(VARCHAR.getTypeSignature())), List.class, createTestBlock(TESTING_TYPE_MANAGER.getType(arrayType(VARCHAR.getTypeSignature()))));
+    }
+
+    public static Block createTestBlock(Type arrayType)
+    {
+        BlockBuilder blockBuilder = arrayType.createBlockBuilder(null, 4);
+        arrayType.writeObject(blockBuilder, arrayBlockOf(VARCHAR, "1", "2"));
+        arrayType.writeObject(blockBuilder, arrayBlockOf(VARCHAR, "the", "quick", "brown", "fox"));
+        arrayType.writeObject(blockBuilder, arrayBlockOf(VARCHAR, "one-two-three-four-five", "123456789012345", "the quick brown fox", "hello-world-hello-world-hello-world"));
+        return blockBuilder.build();
+    }
+
+    @Override
+    protected Object getGreaterValue(Object value)
+    {
+        Block block = (Block) value;
+        BlockBuilder blockBuilder = VARCHAR.createBlockBuilder(null, block.getPositionCount() + 1);
+        for (int i = 0; i < block.getPositionCount(); i++) {
+            VARCHAR.appendTo(block, i, blockBuilder);
+        }
+        VARCHAR.writeSlice(blockBuilder, utf8Slice("_"));
+
+        return blockBuilder.build();
+    }
+
+    @Test
+    public void testRange()
+    {
+        assertThat(type.getRange())
+                .isEmpty();
+    }
+
+    @Test
+    public void testPreviousValue()
+    {
+        assertThat(type.getPreviousValue(getSampleValue()))
+                .isEmpty();
+    }
+
+    @Test
+    public void testNextValue()
+    {
+        assertThat(type.getNextValue(getSampleValue()))
+                .isEmpty();
+    }
+}

core/trino-main/src/test/java/io/trino/type/TestVarcharVarcharMapType.java

@@ -0,0 +1,87 @@
+/*
+ * 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 io.trino.type;
+
+import com.google.common.collect.ImmutableMap;
+import io.trino.spi.block.Block;
+import io.trino.spi.block.BlockBuilder;
+import io.trino.spi.type.Type;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+
+import static io.trino.spi.type.VarcharType.VARCHAR;
+import static io.trino.util.StructuralTestUtil.mapBlockOf;
+import static io.trino.util.StructuralTestUtil.mapType;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+public class TestVarcharVarcharMapType
+        extends AbstractTestType
+{
+    public TestVarcharVarcharMapType()
+    {
+        super(mapType(VARCHAR, VARCHAR), Map.class, createTestBlock(mapType(VARCHAR, VARCHAR)));
+    }
+
+    public static Block createTestBlock(Type mapType)
+    {
+        BlockBuilder blockBuilder = mapType.createBlockBuilder(null, 2);
+        mapType.writeObject(blockBuilder, mapBlockOf(VARCHAR, VARCHAR, ImmutableMap.of("hi", "there")));
+        mapType.writeObject(blockBuilder, mapBlockOf(VARCHAR, VARCHAR, ImmutableMap.of("one", "1", "hello", "world")));
+        mapType.writeObject(blockBuilder, mapBlockOf(VARCHAR, VARCHAR, ImmutableMap.of("one-two-three-four-five", "123456789012345", "the quick brown fox", "hello-world-hello-world-hello-world")));
+        return blockBuilder.build();
+    }
+
+    @Override
+    protected Object getGreaterValue(Object value)
+    {
+        throw new UnsupportedOperationException();
+    }
+
+    @Test
+    public void testRange()
+    {
+        assertThat(type.getRange())
+                .isEmpty();
+    }
+
+    @Test
+    public void testPreviousValue()
+    {
+        Object sampleValue = getSampleValue();
+        if (!type.isOrderable()) {
+            assertThatThrownBy(() -> type.getPreviousValue(sampleValue))
+                    .isInstanceOf(IllegalStateException.class)
+                    .hasMessage("Type is not orderable: " + type);
+            return;
+        }
+        assertThat(type.getPreviousValue(sampleValue))
+                .isEmpty();
+    }
+
+    @Test
+    public void testNextValue()
+    {
+        Object sampleValue = getSampleValue();
+        if (!type.isOrderable()) {
+            assertThatThrownBy(() -> type.getNextValue(sampleValue))
+                    .isInstanceOf(IllegalStateException.class)
+                    .hasMessage("Type is not orderable: " + type);
+            return;
+        }
+        assertThat(type.getNextValue(sampleValue))
+                .isEmpty();
+    }
+}

core/trino-spi/src/main/java/io/trino/spi/type/ArrayType.java

@@ -271,6 +271,28 @@ public void writeObject(BlockBuilder blockBuilder, Object value)
         });
     }
 
+    // FLAT MEMORY LAYOUT
+    //
+    // All data of the array is stored in the variable width section. Within the variable width section,
+    // fixed data for all elements is stored first, followed by variable length data for all elements
+    // This simplifies the read implementation as we can simply step through the fixed section without
+    // knowing the variable length of each element, since each element stores the offset to its variable
+    // length data inside its fixed length data.
+    //
+    // In the current implementation, the element and null flag are stored in an interleaved flat record.
+    // This layout is not required by the format, and could be changed to a columnar if it is determined
+    // to be more efficient.
+    //
+    // Fixed:
+    //   int positionCount, int variableSizeOffset
+    // Variable:
+    //   byte element1Null, elementFixedSize element1FixedData
+    //   byte element2Null, elementFixedSize element2FixedData
+    //   ...
+    //   element1VariableSize element1VariableData
+    //   element2VariableSize element2VariableData
+    //   ...
+
     @Override
     public int getFlatFixedSize()
     {
@@ -320,21 +342,24 @@ public int relocateFlatVariableWidthOffsets(byte[] fixedSizeSlice, int fixedSize
 
     private int relocateVariableWidthData(int positionCount, int elementFixedSize, byte[] slice, int offset)
     {
-        int writeVariableWidthOffset = positionCount * (1 + elementFixedSize);
+        int writeFixedOffset = offset;
+        // variable width data starts after fixed width data
+        // there is one extra byte per position for the null flag
+        int writeVariableWidthOffset = offset + positionCount * (1 + elementFixedSize);
         for (int index = 0; index < positionCount; index++) {
-            if (slice[offset] != 0) {
-                offset++;
+            if (slice[writeFixedOffset] != 0) {
+                writeFixedOffset++;
             }
             else {
                 // skip null byte
-                offset++;
+                writeFixedOffset++;
 
-                int elementVariableSize = elementType.relocateFlatVariableWidthOffsets(slice, offset, slice, offset + writeVariableWidthOffset);
+                int elementVariableSize = elementType.relocateFlatVariableWidthOffsets(slice, writeFixedOffset, slice, writeVariableWidthOffset);
                 writeVariableWidthOffset += elementVariableSize;
             }
-            offset += elementFixedSize;
+            writeFixedOffset += elementFixedSize;
         }
-        return writeVariableWidthOffset;
+        return writeVariableWidthOffset - offset;
     }
 
     @Override
@@ -433,6 +458,8 @@ private static void writeFlatElements(Type elementType, MethodHandle elementWrit
             throws Throwable
     {
         int positionCount = array.getPositionCount();
+        // variable width data starts after fixed width data
+        // there is one extra byte per position for the null flag
         int writeVariableWidthOffset = offset + positionCount * (1 + elementFixedSize);
         for (int index = 0; index < positionCount; index++) {
             if (array.isNull(index)) {

core/trino-spi/src/main/java/io/trino/spi/type/MapType.java

@@ -334,6 +334,29 @@ public void writeObject(BlockBuilder blockBuilder, Object value)
         });
     }
 
+    // FLAT MEMORY LAYOUT
+    //
+    // All data of the map is stored in the variable width section. Within the variable width section,
+    // fixed data for all keys and values are stored first, followed by variable length data for all keys
+    // and values. This simplifies the read implementation as we can simply step through the fixed
+    // section without knowing the variable length of each value, since each value stores the offset
+    // to its variable length data inside its fixed length data.
+    //
+    // In the current implementation, the keys and values are stored in an interleaved flat record along
+    // with null flags. This layout is not required by the format, and could be changed to a columnar
+    // if it is determined to be more efficient. Additionally, this layout allows for a null key, since
+    // non-null keys is not always enforced, and null keys may be allowed in the future.
+    //
+    // Fixed:
+    //   int positionCount, int variableSizeOffset
+    // Variable:
+    //   byte key1Null, keyFixedSize key1FixedData, byte value1Null, valueFixedSize value1FixedData
+    //   byte key2Null, keyFixedSize key2FixedData, byte value2Null, valueFixedSize value2FixedData
+    //   ...
+    //   key1VariableSize key1VariableData, value1VariableSize value1VariableData
+    //   key2VariableSize key2VariableData, value2VariableSize value2VariableData
+    //   ...
+
     @Override
     public int getFlatFixedSize()
     {
@@ -387,33 +410,36 @@ public int relocateFlatVariableWidthOffsets(byte[] fixedSizeSlice, int fixedSize
 
     private int relocateVariableWidthData(int positionCount, int keyFixedSize, int valueFixedSize, byte[] slice, int offset)
     {
-        int writeVariableWidthOffset = positionCount / 2 * (2 + keyFixedSize + valueFixedSize);
+        int writeFixedOffset = offset;
+        // variable width data starts after fixed width data for the keys and values
+        // there is one extra byte per key and value for a null flag
+        int writeVariableWidthOffset = offset + (positionCount / 2 * (2 + keyFixedSize + valueFixedSize));
         for (int index = 0; index < positionCount; index += 2) {
-            if (!keyType.isFlatVariableWidth() || slice[offset] != 0) {
-                offset++;
+            if (!keyType.isFlatVariableWidth() || slice[writeFixedOffset] != 0) {
+                writeFixedOffset++;
             }
             else {
                 // skip null byte
-                offset++;
+                writeFixedOffset++;
 
-                int keyVariableSize = keyType.relocateFlatVariableWidthOffsets(slice, offset, slice, offset + writeVariableWidthOffset);
+                int keyVariableSize = keyType.relocateFlatVariableWidthOffsets(slice, writeFixedOffset, slice, writeVariableWidthOffset);
                 writeVariableWidthOffset += keyVariableSize;
             }
-            offset += keyFixedSize;
+            writeFixedOffset += keyFixedSize;
 
-            if (!valueType.isFlatVariableWidth() || slice[offset] != 0) {
-                offset++;
+            if (!valueType.isFlatVariableWidth() || slice[writeFixedOffset] != 0) {
+                writeFixedOffset++;
             }
             else {
                 // skip null byte
-                offset++;
+                writeFixedOffset++;
 
-                int valueVariableSize = valueType.relocateFlatVariableWidthOffsets(slice, offset, slice, offset + writeVariableWidthOffset);
+                int valueVariableSize = valueType.relocateFlatVariableWidthOffsets(slice, writeFixedOffset, slice, writeVariableWidthOffset);
                 writeVariableWidthOffset += valueVariableSize;
             }
-            offset += valueFixedSize;
+            writeFixedOffset += valueFixedSize;
         }
-        return writeVariableWidthOffset;
+        return writeVariableWidthOffset - offset;
     }
 
     @Override
@@ -636,7 +662,9 @@ private static void writeFlatEntries(
             int offset)
             throws Throwable
     {
-        int writeVariableWidthOffset = offset + map.getPositionCount() / 2 * (2 + keyFixedSize + valueFixedSize);
+        // variable width data starts after fixed width data for the keys and values
+        // there is one extra byte per key and value for a null flag
+        int writeVariableWidthOffset = offset + (map.getPositionCount() / 2 * (2 + keyFixedSize + valueFixedSize));
         for (int index = 0; index < map.getPositionCount(); index += 2) {
             if (map.isNull(index)) {
                 slice[offset] = 1;
@@ -656,7 +684,7 @@ private static void writeFlatEntries(
                         slice,
                         offset,
                         slice,
-                        offset + writeVariableWidthOffset);
+                        writeVariableWidthOffset);
                 writeVariableWidthOffset += keyVariableSize;
             }
             offset += keyFixedSize;
@@ -679,7 +707,7 @@ private static void writeFlatEntries(
                         slice,
                         offset,
                         slice,
-                        offset + writeVariableWidthOffset);
+                        writeVariableWidthOffset);
                 writeVariableWidthOffset += valueVariableSize;
             }
             offset += valueFixedSize;