Fix stream ordering in ORC writer

presto-orc/src/main/java/com/facebook/presto/orc/metadata/Stream.java

@@ -13,14 +13,16 @@
  */
 package com.facebook.presto.orc.metadata;
 
+import java.util.Objects;
 import java.util.Optional;
 
 import static com.google.common.base.MoreObjects.toStringHelper;
 import static java.util.Objects.requireNonNull;
 
 public class Stream
 {
-    public enum StreamArea {
+    public enum StreamArea
+    {
         INDEX,
         DATA,
     }
@@ -130,4 +132,28 @@ public Stream withOffset(long offset)
                 this.sequence,
                 Optional.of(offset));
     }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        Stream stream = (Stream) o;
+        return column == stream.column
+                && length == stream.length
+                && useVInts == stream.useVInts
+                && sequence == stream.sequence
+                && streamKind == stream.streamKind
+                && Objects.equals(offset, stream.offset);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(column, streamKind, length, useVInts, sequence, offset);
+    }
 }

presto-orc/src/main/java/com/facebook/presto/orc/stream/StreamDataOutput.java

@@ -20,6 +20,7 @@
 
 import java.util.function.ToLongFunction;
 
+import static com.google.common.base.MoreObjects.toStringHelper;
 import static com.google.common.base.Verify.verify;
 import static java.util.Objects.requireNonNull;
 
@@ -68,4 +69,12 @@ public void writeData(SliceOutput sliceOutput)
         long size = writer.applyAsLong(sliceOutput);
         verify(stream.getLength() == size, "Data stream did not write expected size");
     }
+
+    @Override
+    public String toString()
+    {
+        return toStringHelper(this)
+                .add("stream", stream)
+                .toString();
+    }
 }

presto-orc/src/main/java/com/facebook/presto/orc/writer/ColumnSequenceKey.java

@@ -0,0 +1,47 @@
+/*
+ * 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 com.facebook.presto.orc.writer;
+
+import java.util.Objects;
+
+class ColumnSequenceKey
+{
+    final int column;
+    final int sequence;
+
+    public ColumnSequenceKey(int column, int sequence)
+    {
+        this.column = column;
+        this.sequence = sequence;
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        ColumnSequenceKey that = (ColumnSequenceKey) o;
+        return column == that.column && sequence == that.sequence;
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(column, sequence);
+    }
+}

presto-orc/src/main/java/com/facebook/presto/orc/writer/ColumnSizeLayout.java

@@ -0,0 +1,164 @@
+/*
+ * 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 com.facebook.presto.orc.writer;
+
+import com.facebook.presto.orc.metadata.ColumnEncoding;
+import com.facebook.presto.orc.metadata.Stream;
+import com.facebook.presto.orc.stream.StreamDataOutput;
+
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import static com.facebook.presto.orc.metadata.ColumnEncoding.ColumnEncodingKind.DWRF_MAP_FLAT;
+import static com.google.common.base.MoreObjects.toStringHelper;
+import static java.util.Objects.requireNonNull;
+
+/**
+ * Order streams by total column size in the desc order preserving the type-tree
+ * iteration order for complex types.
+ * <p>
+ * For flatmaps consider streams for the same flatmap key as a pseudo-column: do
+ * additional grouping by sequence to keep streams for the same key together,
+ * and then sort streams belonging to the same flatmap key by the total size in
+ * desc order.
+ */
+public class ColumnSizeLayout
+        implements StreamLayout
+{
+    @Override
+    public void reorder(List<StreamDataOutput> dataStreams, Map<Integer, Integer> nodeToColumn, Map<Integer, ColumnEncoding> nodeIdToColumnEncodings)
+    {
+        requireNonNull(dataStreams, "dataStreams is null");
+        requireNonNull(nodeToColumn, "nodeToColumn is null");
+        requireNonNull(nodeIdToColumnEncodings, "nodeIdToColumnEncodings is null");
+
+        if (dataStreams.isEmpty()) {
+            return;
+        }
+
+        Set<Integer> flatMapColumns = getFlatMapColumns(nodeToColumn, nodeIdToColumnEncodings);
+
+        // gather column sizes on the column and column+seq levels
+        Map<Integer, long[]> columnSize = new HashMap<>();
+        Map<ColumnSequenceKey, long[]> flatMapColumnSize = new HashMap<>();
+
+        for (StreamDataOutput dataStream : dataStreams) {
+            Stream stream = dataStream.getStream();
+            int node = stream.getColumn();
+            Integer column = nodeToColumn.get(node);
+
+            long[] storedColumnSize = columnSize.computeIfAbsent(column, (n) -> new long[] {0});
+            storedColumnSize[0] += dataStream.size();
+
+            if (flatMapColumns.contains(column)) {
+                ColumnSequenceKey key = new ColumnSequenceKey(column, stream.getSequence());
+                long[] storedFlatMapColumnSize = flatMapColumnSize.computeIfAbsent(key, (n) -> new long[] {0});
+                storedFlatMapColumnSize[0] += dataStream.size();
+            }
+        }
+
+        // do the ordering
+        dataStreams.sort((streamDataA, streamDataB) -> {
+            Stream streamA = streamDataA.getStream();
+            Stream streamB = streamDataB.getStream();
+
+            int nodeA = streamA.getColumn();
+            int nodeB = streamB.getColumn();
+
+            int columnA = nodeToColumn.get(nodeA);
+            int columnB = nodeToColumn.get(nodeB);
+
+            boolean isFlatMapA = flatMapColumns.contains(columnA);
+            boolean isFlatMapB = flatMapColumns.contains(columnB);
+
+            // split non-flatmap and flatmap columns into separate groups
+            if (isFlatMapA != isFlatMapB) {
+                return Boolean.compare(isFlatMapA, isFlatMapB);
+            }
+
+            long columnSizeA = columnSize.get(columnA)[0];
+            long columnSizeB = columnSize.get(columnB)[0];
+
+            // order columns by total column size in desc order
+            if (columnSizeA != columnSizeB) {
+                return Long.compare(columnSizeB, columnSizeA);
+            }
+
+            // group streams by the column
+            if (columnA != columnB) {
+                return Integer.compare(columnA, columnB);
+            }
+
+            if (isFlatMapA) {
+                int sequenceA = streamA.getSequence();
+                int sequenceB = streamB.getSequence();
+
+                // special handling for seq 0 before sorting by the col+seq size
+                // to keep it on top of the group
+                if (sequenceA != sequenceB) {
+                    if (sequenceA == 0) {
+                        return -1;
+                    }
+                    if (sequenceB == 0) {
+                        return 1;
+                    }
+                }
+
+                long columnSeqSizeA = flatMapColumnSize.get(new ColumnSequenceKey(columnA, sequenceA))[0];
+                long columnSeqSizeB = flatMapColumnSize.get(new ColumnSequenceKey(columnB, sequenceB))[0];
+
+                // order sequences by total column+seq size in desc order
+                if (columnSeqSizeA != columnSeqSizeB) {
+                    return Long.compare(columnSeqSizeB, columnSeqSizeA);
+                }
+
+                // group by the sequence
+                if (sequenceA != sequenceB) {
+                    return Integer.compare(sequenceA, sequenceB);
+                }
+            }
+
+            // order by the node in asc order
+            if (nodeA != nodeB) {
+                return Integer.compare(nodeA, nodeB);
+            }
+
+            // sort by the stream kind, we don't really need it, but it makes testing easier
+            return Integer.compare(streamA.getStreamKind().ordinal(), streamB.getStreamKind().ordinal());
+        });
+    }
+
+    @Override
+    public String toString()
+    {
+        return toStringHelper(this)
+                .toString();
+    }
+
+    private static Set<Integer> getFlatMapColumns(Map<Integer, Integer> nodeIdToColumn, Map<Integer, ColumnEncoding> nodeIdToColumnEncodings)
+    {
+        Set<Integer> flatMapColumns = new HashSet<>();
+        for (Map.Entry<Integer, ColumnEncoding> e : nodeIdToColumnEncodings.entrySet()) {
+            Integer node = e.getKey();
+            ColumnEncoding encoding = e.getValue();
+            if (encoding.getColumnEncodingKind() == DWRF_MAP_FLAT) {
+                flatMapColumns.add(nodeIdToColumn.get(node));
+            }
+        }
+        return flatMapColumns;
+    }
+}

presto-orc/src/main/java/com/facebook/presto/orc/writer/StreamLayout.java

@@ -14,14 +14,12 @@
 package com.facebook.presto.orc.writer;
 
 import com.facebook.presto.orc.metadata.ColumnEncoding;
-import com.facebook.presto.orc.metadata.Stream;
 import com.facebook.presto.orc.stream.StreamDataOutput;
 
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
-import static com.google.common.collect.ImmutableMap.toImmutableMap;
 import static java.util.Objects.requireNonNull;
 
 /**
@@ -63,62 +61,4 @@ public String toString()
             return "ByStreamSize{}";
         }
     }
-
-    /**
-     * Streams are ordered by Column Size. If two columns have same size then
-     * columnId, stream size and stream kind are used for ordering. It orders them
-     * by ascending order of column size to read multiple small columns in one IO.
-     * All streams for a column are stored together to read one column in one IO.
-     */
-    class ByColumnSize
-            implements StreamLayout
-    {
-        public void reorder(List<StreamDataOutput> dataStreams)
-        {
-            requireNonNull(dataStreams, "dataStreams is null");
-            if (dataStreams.isEmpty()) {
-                return;
-            }
-
-            Map<Integer, Long> columnSizes = dataStreams.stream()
-                    .collect(toImmutableMap(
-                            s -> s.getStream().getColumn(),
-                            s -> (long) s.getStream().getLength(),
-                            Long::sum));
-
-            dataStreams.sort((left, right) -> {
-                Stream leftStream = left.getStream();
-                Stream rightStream = right.getStream();
-
-                long sizeDelta = columnSizes.get(leftStream.getColumn()) - columnSizes.get(rightStream.getColumn());
-                if (sizeDelta != 0) {
-                    return sizeDelta < 0 ? -1 : 1;
-                }
-
-                int columnDelta = leftStream.getColumn() - rightStream.getColumn();
-                if (columnDelta != 0) {
-                    return columnDelta;
-                }
-
-                sizeDelta = leftStream.getLength() - rightStream.getLength();
-                if (sizeDelta != 0) {
-                    return sizeDelta < 0 ? -1 : 1;
-                }
-
-                return leftStream.getStreamKind().compareTo(rightStream.getStreamKind());
-            });
-        }
-
-        @Override
-        public void reorder(List<StreamDataOutput> dataStreams, Map<Integer, Integer> nodeIdToColumn, Map<Integer, ColumnEncoding> nodeIdToColumnEncodings)
-        {
-            reorder(dataStreams);
-        }
-
-        @Override
-        public String toString()
-        {
-            return "ByColumnSize{}";
-        }
-    }
 }

presto-orc/src/main/java/com/facebook/presto/orc/writer/StreamLayoutFactory.java

@@ -23,7 +23,7 @@ class ColumnSizeLayoutFactory
         @Override
         public StreamLayout create()
         {
-            return new StreamLayout.ByColumnSize();
+            return new ColumnSizeLayout();
         }
 
         @Override