Add support for LazyBlock in columnar decompose helpers

trinodb · Oct 9, 2019 · 774fab9 · 774fab9
1 parent adc9c78
commit 774fab9
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 5 deletions.
diff --git a/presto-main/src/main/java/io/prestosql/operator/PageUtils.java b/presto-main/src/main/java/io/prestosql/operator/PageUtils.java
@@ -37,7 +37,6 @@ public static Page recordMaterializedBytes(Page page, LongConsumer sizeInBytesCo
                 blocks[i] = block;
             }
             else {
-                // TODO: block might be partially loaded
                 blocks[i] = new LazyBlock(page.getPositionCount(), lazyBlock -> {
                     Block loadedBlock = block.getLoadedBlock();
                     sizeInBytesConsumer.accept(loadedBlock.getSizeInBytes());

diff --git a/presto-spi/src/main/java/io/prestosql/spi/block/Block.java b/presto-spi/src/main/java/io/prestosql/spi/block/Block.java
@@ -285,16 +285,18 @@ default boolean mayHaveNull()
     boolean isNull(int position);
 
     /**
-     * Returns true if block data is loaded into memory.
+     * Returns true if block data is fully loaded into memory.
      */
     default boolean isLoaded()
     {
         return true;
     }
 
     /**
-     * Returns a block that assures all data is in memory.
-     * May return the same block if all block data is already in memory.
+     * Returns a fully loaded block that assures all data is in memory.
+     * Neither the returned block nor any nested block will be a {@link LazyBlock}.
+     * The same block will be returned if neither the current block nor any
+     * nested blocks are {@link LazyBlock},
      * <p>
      * This allows streaming data sources to skip sections that are not
      * accessed in a query.

diff --git a/presto-spi/src/main/java/io/prestosql/spi/block/ColumnarArray.java b/presto-spi/src/main/java/io/prestosql/spi/block/ColumnarArray.java
@@ -26,6 +26,9 @@ public static ColumnarArray toColumnarArray(Block block)
     {
         requireNonNull(block, "block is null");
 
+        if (block instanceof LazyBlock) {
+            block = ((LazyBlock) block).getBlock();
+        }
         if (block instanceof DictionaryBlock) {
             return toColumnarArray((DictionaryBlock) block);
         }

diff --git a/presto-spi/src/main/java/io/prestosql/spi/block/ColumnarMap.java b/presto-spi/src/main/java/io/prestosql/spi/block/ColumnarMap.java
@@ -27,6 +27,9 @@ public static ColumnarMap toColumnarMap(Block block)
     {
         requireNonNull(block, "block is null");
 
+        if (block instanceof LazyBlock) {
+            block = ((LazyBlock) block).getBlock();
+        }
         if (block instanceof DictionaryBlock) {
             return toColumnarMap((DictionaryBlock) block);
         }

diff --git a/presto-spi/src/main/java/io/prestosql/spi/block/ColumnarRow.java b/presto-spi/src/main/java/io/prestosql/spi/block/ColumnarRow.java
@@ -24,6 +24,9 @@ public static ColumnarRow toColumnarRow(Block block)
     {
         requireNonNull(block, "block is null");
 
+        if (block instanceof LazyBlock) {
+            block = ((LazyBlock) block).getBlock();
+        }
         if (block instanceof DictionaryBlock) {
             return toColumnarRow((DictionaryBlock) block);
         }

diff --git a/presto-spi/src/main/java/io/prestosql/spi/block/LazyBlock.java b/presto-spi/src/main/java/io/prestosql/spi/block/LazyBlock.java
@@ -248,6 +248,12 @@ public boolean isNull(int position)
         return block.isNull(position);
     }
 
+    public Block getBlock()
+    {
+        assureLoaded();
+        return block;
+    }
+
     public void setBlock(Block block)
     {
         if (this.block != null) {
@@ -259,13 +265,14 @@ public void setBlock(Block block)
     @Override
     public boolean isLoaded()
     {
-        return block != null;
+        return block != null && block.isLoaded();
     }
 
     @Override
     public Block getLoadedBlock()
     {
         assureLoaded();
+        block = block.getLoadedBlock();
         return block;
     }
 
@@ -276,6 +283,10 @@ private void assureLoaded()
         }
         loader.load(this);
 
+        while (block instanceof LazyBlock) {
+            block = ((LazyBlock) block).getBlock();
+        }
+
         if (block == null) {
             throw new IllegalArgumentException("Lazy block loader did not load this block");
         }