Add methods to get ArrayBlock elements and MapBlock keys and values

nineinchnick · Oct 8, 2024 · 596b31f · 596b31f
1 parent 3999f62
commit 596b31f
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/ArrayBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/ArrayBlock.java
@@ -165,6 +165,13 @@ public void retainedBytesForEachPart(ObjLongConsumer<Object> consumer)
         consumer.accept(this, INSTANCE_SIZE);
     }
 
+    public Block getElementsBlock()
+    {
+        int start = offsets[arrayOffset];
+        int end = offsets[arrayOffset + positionCount];
+        return values.getRegion(start, end - start);
+    }
+
     Block getRawElementBlock()
     {
         return values;

diff --git a/core/trino-spi/src/main/java/io/trino/spi/block/MapBlock.java b/core/trino-spi/src/main/java/io/trino/spi/block/MapBlock.java
@@ -186,6 +186,20 @@ private MapBlock(
         this.retainedSizeInBytes = INSTANCE_SIZE + sizeOf(offsets) + sizeOf(mapIsNull);
     }
 
+    public Block getKeyBlock()
+    {
+        int start = offsets[startOffset];
+        int end = offsets[startOffset + positionCount];
+        return keyBlock.getRegion(start, end - start);
+    }
+
+    public Block getValueBlock()
+    {
+        int start = offsets[startOffset];
+        int end = offsets[startOffset + positionCount];
+        return valueBlock.getRegion(start, end - start);
+    }
+
     Block getRawKeyBlock()
     {
         return keyBlock;