Create the list of type-specific unnesters in UnnestOperator only once

trinodb · Mar 22, 2017 · 46cedd4 · 46cedd4
1 parent cda4143
commit 46cedd4
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 18 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/operator/ArrayUnnester.java b/presto-main/src/main/java/com/facebook/presto/operator/ArrayUnnester.java
@@ -27,11 +27,11 @@ public class ArrayUnnester
         implements Unnester
 {
     private final Type elementType;
-    private final Block arrayBlock;
+    private Block arrayBlock;
     private final int channelCount;
 
     private int position;
-    private final int positionCount;
+    private int positionCount;
 
     public ArrayUnnester(ArrayType arrayType, @Nullable Block arrayBlock)
     {
@@ -66,4 +66,12 @@ public final void appendNext(PageBuilder pageBuilder, int outputChannelOffset)
     {
         appendTo(pageBuilder, outputChannelOffset);
     }
+
+    @Override
+    public void setBlock(@Nullable Block arrayBlock)
+    {
+        this.arrayBlock = arrayBlock;
+        this.position = 0;
+        this.positionCount = arrayBlock == null ? 0 : arrayBlock.getPositionCount();
+    }
 }
diff --git a/presto-main/src/main/java/com/facebook/presto/operator/MapUnnester.java b/presto-main/src/main/java/com/facebook/presto/operator/MapUnnester.java
@@ -28,11 +28,11 @@ public class MapUnnester
 {
     private final Type keyType;
     private final Type valueType;
-    private final Block block;
+    private Block block;
     private final int channelCount;
 
     private int position;
-    private final int positionCount;
+    private int positionCount;
 
     public MapUnnester(MapType mapType, @Nullable Block mapBlock)
     {
@@ -70,4 +70,12 @@ public final void appendNext(PageBuilder pageBuilder, int outputChannelOffset)
     {
         appendTo(pageBuilder, outputChannelOffset);
     }
+
+    @Override
+    public void setBlock(@Nullable Block mapBlock)
+    {
+        this.block = mapBlock;
+        this.position = 0;
+        this.positionCount = mapBlock == null ? 0 : mapBlock.getPositionCount();
+    }
 }
diff --git a/presto-main/src/main/java/com/facebook/presto/operator/UnnestOperator.java b/presto-main/src/main/java/com/facebook/presto/operator/UnnestOperator.java
@@ -125,7 +125,18 @@ public UnnestOperator(OperatorContext operatorContext, List<Integer> replicateCh
         }
         this.outputTypes = outputTypesBuilder.build();
         this.pageBuilder = new PageBuilder(outputTypes);
-        this.unnesters = new ArrayList<>();
+        this.unnesters = new ArrayList<>(unnestTypes.size());
+        for (Type type : unnestTypes) {
+            if (type instanceof ArrayType) {
+                unnesters.add(new ArrayUnnester((ArrayType) type, null));
+            }
+            else if (type instanceof MapType) {
+                unnesters.add(new MapUnnester((MapType) type, null));
+            }
+            else {
+                throw new IllegalArgumentException("Cannot unnest type: " + type);
+            }
+        }
     }
 
     private static List<Type> getUnnestedTypes(List<Type> types)
@@ -178,28 +189,19 @@ public void addInput(Page page)
 
         currentPage = page;
         currentPosition = 0;
-        initializeUnnesters();
+        fillUnnesters();
     }
 
-    private void initializeUnnesters()
+    private void fillUnnesters()
     {
-        unnesters.clear();
         for (int i = 0; i < unnestTypes.size(); i++) {
             Type type = unnestTypes.get(i);
             int channel = unnestChannels.get(i);
             Block block = null;
             if (!currentPage.getBlock(channel).isNull(currentPosition)) {
                 block = (Block) type.getObject(currentPage.getBlock(channel), currentPosition);
             }
-            if (type instanceof ArrayType) {
-                unnesters.add(new ArrayUnnester((ArrayType) type, block));
-            }
-            else if (type instanceof MapType) {
-                unnesters.add(new MapUnnester((MapType) type, block));
-            }
-            else {
-                throw new IllegalArgumentException("Cannot unnest type: " + type);
-            }
+            unnesters.get(i).setBlock(block);
         }
         ordinalityCount = 0;
     }
@@ -226,7 +228,7 @@ public Page getOutput()
                     currentPosition = 0;
                     break;
                 }
-                initializeUnnesters();
+                fillUnnesters();
             }
             while (!pageBuilder.isFull() && anyUnnesterHasData()) {
                 // Copy all the channels marked for replication

diff --git a/presto-main/src/main/java/com/facebook/presto/operator/Unnester.java b/presto-main/src/main/java/com/facebook/presto/operator/Unnester.java
@@ -14,6 +14,7 @@
 package com.facebook.presto.operator;
 
 import com.facebook.presto.spi.PageBuilder;
+import com.facebook.presto.spi.block.Block;
 
 public interface Unnester
 {
@@ -22,4 +23,6 @@ public interface Unnester
     void appendNext(PageBuilder pageBuilder, int outputChannelOffset);
 
     boolean hasNext();
+
+    void setBlock(Block block);
 }