From 9bf99987b555a787b100a5d9431542c4c47dd772 Mon Sep 17 00:00:00 2001
From: Shubham Tagra <stagra@qubole.com>
Date: Wed, 16 Jul 2014 09:57:57 +0530
Subject: [PATCH] Fixes over Dain's changes:
 1. Fix for files created by hive0.11
 2. Changes to pass down the predicate to OrcReader
 3. Fix to avoid using hadoop's FSInputStream.readFully method which causes unwanted seeks

---
 .../java/com/facebook/presto/hive/HiveClient.java  |   10 +++++++-
 .../presto/hive/HiveSplitSourceProvider.java       |   11 +++++++-
 .../facebook/presto/hive/orc/OrcRecordReader.java  |    9 +++++++-
 .../com/facebook/presto/hive/orc/StripeReader.java |   23 +++++++++++++++++--
 .../InformationSchemaSplitManager.java             |    6 +++++
 .../presto/connector/jmx/JmxSplitManager.java      |    6 +++++
 .../connector/system/SystemSplitManager.java       |    6 +++++
 .../com/facebook/presto/split/SplitManager.java    |   17 +++++++++++++-
 .../sql/planner/DistributedExecutionPlanner.java   |   11 ++++++++-
 .../facebook/presto/raptor/RaptorSplitManager.java |    6 +++++
 .../facebook/presto/spi/ConnectorSplitManager.java |    2 +
 .../ClassLoaderSafeConnectorSplitManager.java      |    8 ++++++-
 .../com/facebook/presto/tpch/TpchSplitManager.java |    6 +++++
 13 files changed, 111 insertions(+), 10 deletions(-)

diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveClient.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveClient.java
index 9a9faad..cb2deb4 100644
--- a/presto-hive/src/main/java/com/facebook/presto/hive/HiveClient.java
+++ b/presto-hive/src/main/java/com/facebook/presto/hive/HiveClient.java
@@ -847,7 +847,14 @@ public class HiveClient
     @Override
     public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle tableHandle, List<ConnectorPartition> connectorPartitions)
     {
+        return getPartitionSplits(tableHandle, connectorPartitions, TupleDomain.<ConnectorColumnHandle>all());
+    }
+
+    @Override
+    public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle tableHandle, List<ConnectorPartition> connectorPartitions, TupleDomain<ConnectorColumnHandle> tupleDomain)
+    {
         HiveTableHandle hiveTableHandle = checkType(tableHandle, HiveTableHandle.class, "tableHandle");
+        TupleDomain<HiveColumnHandle> td = (TupleDomain<HiveColumnHandle>) (Object) tupleDomain;
 
         checkNotNull(connectorPartitions, "connectorPartitions is null");
         List<HivePartition> partitions = Lists.transform(connectorPartitions, new Function<ConnectorPartition, HivePartition>() {
@@ -893,7 +900,8 @@ public class HiveClient
                 hiveTableHandle.getSession(),
                 maxInitialSplitSize,
                 maxInitialSplits,
-                recursiveDfsWalkerEnabled).get();
+                recursiveDfsWalkerEnabled,
+                td).get();
     }
 
     private Iterable<HivePartitionMetadata> getPartitionMetadata(final Table table, final SchemaTableName tableName, List<HivePartition> partitions)
diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveSplitSourceProvider.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveSplitSourceProvider.java
index ebc5fdb..4face71 100644
--- a/presto-hive/src/main/java/com/facebook/presto/hive/HiveSplitSourceProvider.java
+++ b/presto-hive/src/main/java/com/facebook/presto/hive/HiveSplitSourceProvider.java
@@ -31,7 +31,9 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+
 import io.airlift.units.DataSize;
+
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.BlockLocation;
 import org.apache.hadoop.fs.FileStatus;
@@ -120,6 +122,8 @@ class HiveSplitSourceProvider
     private final ConnectorSession session;
     private final boolean recursiveDirWalkerEnabled;
 
+    private final TupleDomain<HiveColumnHandle> tupleDomain;
+
     HiveSplitSourceProvider(String connectorId,
             Table table,
             Iterable<HivePartitionMetadata> partitions,
@@ -135,7 +139,8 @@ class HiveSplitSourceProvider
             ConnectorSession session,
             DataSize maxInitialSplitSize,
             int maxInitialSplits,
-            boolean recursiveDirWalkerEnabled)
+            boolean recursiveDirWalkerEnabled,
+            TupleDomain<HiveColumnHandle> tupleDomain)
     {
         this.connectorId = connectorId;
         this.table = table;
@@ -154,6 +159,7 @@ class HiveSplitSourceProvider
         this.maxInitialSplitSize = maxInitialSplitSize;
         this.remainingInitialSplits = maxInitialSplits;
         this.recursiveDirWalkerEnabled = recursiveDirWalkerEnabled;
+        this.tupleDomain = tupleDomain;
     }
 
     public ConnectorSplitSource get()
@@ -189,7 +195,7 @@ class HiveSplitSourceProvider
                 final String partitionName = partition.getHivePartition().getPartitionId();
                 final Properties schema = getPartitionSchema(table, partition.getPartition());
                 final List<HivePartitionKey> partitionKeys = getPartitionKeys(table, partition.getPartition());
-                final TupleDomain<HiveColumnHandle> tupleDomain = (TupleDomain<HiveColumnHandle>) (Object) partition.getHivePartition().getTupleDomain();
+                final TupleDomain<HiveColumnHandle> tupleDomain = this.tupleDomain;
 
                 Path path = new Path(getPartitionLocation(table, partition.getPartition()));
                 final Configuration configuration = hdfsEnvironment.getConfiguration(path);
@@ -253,6 +259,7 @@ class HiveSplitSourceProvider
                     {
                         try {
                             boolean splittable = isSplittable(inputFormat, file.getPath().getFileSystem(configuration), file.getPath());
+                            splittable = true;
 
                             hiveSplitSource.addToQueue(createHiveSplits(
                                     partitionName,
diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcRecordReader.java b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcRecordReader.java
index 125bd4e..a45a34f 100644
--- a/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcRecordReader.java
+++ b/presto-hive/src/main/java/com/facebook/presto/hive/orc/OrcRecordReader.java
@@ -91,7 +91,14 @@ public class OrcRecordReader
         for (int i = 0; i < fileStripes.size(); i++) {
             StripeInformation stripe = fileStripes.get(i);
             if (splitOffset <= stripe.getOffset() && stripe.getOffset() < splitEndOffset) {
-                TupleDomain<HiveColumnHandle> stripeDomain = extractDomain(columnHandleStreamIndex, stripe.getNumberOfRows(), stripeStats.get(i).getColStatsList());
+                // ORC files created by Hive0.11 don't have Stripe Stats, consider full stripe in such cases
+                TupleDomain<HiveColumnHandle> stripeDomain;
+                if (stripeStats.size() == 0) {
+                    stripeDomain = TupleDomain.<HiveColumnHandle>all();
+                }
+                else {
+                    stripeDomain = extractDomain(columnHandleStreamIndex, stripe.getNumberOfRows(), stripeStats.get(i).getColStatsList());
+                }
                 if (tupleDomain.overlaps(stripeDomain)) {
                     stripes.add(stripe);
                     totalRowCount += stripe.getNumberOfRows();
diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/orc/StripeReader.java b/presto-hive/src/main/java/com/facebook/presto/hive/orc/StripeReader.java
index acc4944..b450490 100644
--- a/presto-hive/src/main/java/com/facebook/presto/hive/orc/StripeReader.java
+++ b/presto-hive/src/main/java/com/facebook/presto/hive/orc/StripeReader.java
@@ -21,7 +21,9 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.primitives.Booleans;
 import com.google.common.primitives.Ints;
+
 import io.airlift.slice.Slices;
+
 import org.apache.hadoop.fs.FSDataInputStream;
 import org.apache.hadoop.hive.ql.io.orc.OrcProto.ColumnEncoding;
 import org.apache.hadoop.hive.ql.io.orc.OrcProto.CompressionKind;
@@ -154,7 +156,7 @@ public class StripeReader
 
         // read the footer
         byte[] tailBuf = new byte[tailLength];
-        file.readFully(offset, tailBuf);
+        readFully(offset, tailBuf);
         InputStream inputStream = new OrcInputStream(Slices.wrappedBuffer(tailBuf).getInput(), compressionKind, bufferSize);
         return StripeFooter.parseFrom(inputStream);
     }
@@ -168,7 +170,7 @@ public class StripeReader
         for (Stream stream : stripeFooter.getStreamsList()) {
             if (includedStreams[stream.getColumn()] && stream.getKind() == ROW_INDEX) {
                 byte[] buffer = new byte[(int) stream.getLength()];
-                file.readFully(offset, buffer);
+                readFully(offset, buffer);
                 indexes[stream.getColumn()] = RowIndex.parseFrom(new OrcInputStream(Slices.wrappedBuffer(buffer).getInput(), compressionKind, bufferSize));
             }
             offset += stream.getLength();
@@ -340,12 +342,27 @@ public class StripeReader
 
             byte[] buffer = new byte[length];
             long readPosition = base + offset;
-            file.readFully(readPosition, buffer);
+            readFully(readPosition, buffer);
             result.add(new StripeSlice(Slices.wrappedBuffer(buffer), offset));
         }
         return result;
     }
 
+    private void readFully(long position, byte[] buffer) throws IOException
+    {
+        file.seek(position);
+        int nread = 0;
+        int offset = 0;
+        int length = buffer.length;
+        while (nread < length) {
+            int nbytes = file.read(buffer, offset + nread, length - nread);
+            if (nbytes < 0) {
+              throw new RuntimeException("End of file reached before reading fully.");
+            }
+            nread += nbytes;
+        }
+    }
+
     private List<Long> getOffsetPositions(
             StreamId streamId,
             ColumnEncoding.Kind columnEncoding,
diff --git a/presto-main/src/main/java/com/facebook/presto/connector/informationSchema/InformationSchemaSplitManager.java b/presto-main/src/main/java/com/facebook/presto/connector/informationSchema/InformationSchemaSplitManager.java
index 2973684..4af3a85 100644
--- a/presto-main/src/main/java/com/facebook/presto/connector/informationSchema/InformationSchemaSplitManager.java
+++ b/presto-main/src/main/java/com/facebook/presto/connector/informationSchema/InformationSchemaSplitManager.java
@@ -65,6 +65,12 @@ public class InformationSchemaSplitManager
     }
 
     @Override
+    public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions, TupleDomain<ConnectorColumnHandle> tupleDomain)
+    {
+        return getPartitionSplits(table, partitions);
+    }
+
+    @Override
     public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions)
     {
         checkNotNull(partitions, "partitions is null");
diff --git a/presto-main/src/main/java/com/facebook/presto/connector/jmx/JmxSplitManager.java b/presto-main/src/main/java/com/facebook/presto/connector/jmx/JmxSplitManager.java
index 06972d4..75e9e79 100644
--- a/presto-main/src/main/java/com/facebook/presto/connector/jmx/JmxSplitManager.java
+++ b/presto-main/src/main/java/com/facebook/presto/connector/jmx/JmxSplitManager.java
@@ -59,6 +59,12 @@ public class JmxSplitManager
     }
 
     @Override
+    public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions, TupleDomain<ConnectorColumnHandle> td)
+    {
+        return getPartitionSplits(table, partitions);
+    }
+
+    @Override
     public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions)
     {
         checkNotNull(partitions, "partitions is null");
diff --git a/presto-main/src/main/java/com/facebook/presto/connector/system/SystemSplitManager.java b/presto-main/src/main/java/com/facebook/presto/connector/system/SystemSplitManager.java
index f6a8dbe..1e1aecf 100644
--- a/presto-main/src/main/java/com/facebook/presto/connector/system/SystemSplitManager.java
+++ b/presto-main/src/main/java/com/facebook/presto/connector/system/SystemSplitManager.java
@@ -72,6 +72,12 @@ public class SystemSplitManager
     }
 
     @Override
+    public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions, TupleDomain<ConnectorColumnHandle> tupleDomain)
+    {
+        return getPartitionSplits(table, partitions);
+    }
+
+    @Override
     public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions)
     {
         checkNotNull(partitions, "partitions is null");
diff --git a/presto-main/src/main/java/com/facebook/presto/split/SplitManager.java b/presto-main/src/main/java/com/facebook/presto/split/SplitManager.java
index f17efb1..ad5ce9c 100644
--- a/presto-main/src/main/java/com/facebook/presto/split/SplitManager.java
+++ b/presto-main/src/main/java/com/facebook/presto/split/SplitManager.java
@@ -19,6 +19,7 @@ import com.facebook.presto.metadata.ColumnHandle;
 import com.facebook.presto.metadata.Partition;
 import com.facebook.presto.metadata.PartitionResult;
 import com.facebook.presto.metadata.TableHandle;
+import com.facebook.presto.spi.ConnectorColumnHandle;
 import com.facebook.presto.spi.ConnectorPartitionResult;
 import com.facebook.presto.spi.ConnectorSplitManager;
 import com.facebook.presto.spi.ConnectorSplitSource;
@@ -53,8 +54,22 @@ public class SplitManager
 
     public SplitSource getPartitionSplits(TableHandle handle, List<Partition> partitions)
     {
+        return getPartitionSplits(handle, partitions, TupleDomain.<ColumnHandle>all());
+    }
+
+    public SplitSource getPartitionSplits(TableHandle handle, List<Partition> partitions, TupleDomain<ColumnHandle> tupleDomain)
+    {
         ConnectorTableHandle table = handle.getConnectorHandle();
-        ConnectorSplitSource source = getConnectorSplitManager(handle).getPartitionSplits(table, Lists.transform(partitions, connectorPartitionGetter()));
+        TupleDomain<ConnectorColumnHandle> td = tupleDomain.transform(new TupleDomain.Function<ColumnHandle, ConnectorColumnHandle>()
+                {
+                    @Override
+                    public ConnectorColumnHandle apply(ColumnHandle handle)
+                    {
+                        return handle.getConnectorHandle();
+                    }
+                });
+        ConnectorSplitSource source = getConnectorSplitManager(handle).getPartitionSplits(table, Lists.transform(partitions, connectorPartitionGetter()), td);
+
         return new ConnectorAwareSplitSource(handle.getConnectorId(), source);
     }
 
diff --git a/presto-main/src/main/java/com/facebook/presto/sql/planner/DistributedExecutionPlanner.java b/presto-main/src/main/java/com/facebook/presto/sql/planner/DistributedExecutionPlanner.java
index 7a47f9c..d19a2b2 100644
--- a/presto-main/src/main/java/com/facebook/presto/sql/planner/DistributedExecutionPlanner.java
+++ b/presto-main/src/main/java/com/facebook/presto/sql/planner/DistributedExecutionPlanner.java
@@ -89,11 +89,20 @@ public class DistributedExecutionPlanner
         public Optional<SplitSource> visitTableScan(TableScanNode node, Void context)
         {
             // get dataSource for table
-            SplitSource splitSource = splitManager.getPartitionSplits(node.getTable(), getPartitions(node));
+            SplitSource splitSource = splitManager.getPartitionSplits(node.getTable(), getPartitions(node), getGeneratedTupleDomain(node));
 
             return Optional.of(splitSource);
         }
 
+        private TupleDomain<ColumnHandle> getGeneratedTupleDomain(TableScanNode node)
+        {
+            if (!node.getGeneratedPartitions().isPresent()) {
+                return TupleDomain.<ColumnHandle>all();
+            }
+
+            return node.getGeneratedPartitions().get().getTupleDomainInput();
+        }
+
         private List<Partition> getPartitions(TableScanNode node)
         {
             if (node.getGeneratedPartitions().isPresent()) {
diff --git a/presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorSplitManager.java b/presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorSplitManager.java
index f71e20a..2978f24 100644
--- a/presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorSplitManager.java
+++ b/presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorSplitManager.java
@@ -108,6 +108,12 @@ public class RaptorSplitManager
     }
 
     @Override
+    public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle tableHandle, List<ConnectorPartition> partitions, TupleDomain<ConnectorColumnHandle> td)
+    {
+        return getPartitionSplits(tableHandle, partitions);
+    }
+
+    @Override
     public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle tableHandle, List<ConnectorPartition> partitions)
     {
         Stopwatch splitTimer = Stopwatch.createStarted();
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorSplitManager.java b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorSplitManager.java
index 5e77714..5b42c6e 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorSplitManager.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/ConnectorSplitManager.java
@@ -30,4 +30,6 @@ public interface ConnectorSplitManager
      * Gets the Splits for the specified Partitions in the indicated table.
      */
     ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions);
+
+    ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions, TupleDomain<ConnectorColumnHandle> tupleDomain);
 }
diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/classloader/ClassLoaderSafeConnectorSplitManager.java b/presto-spi/src/main/java/com/facebook/presto/spi/classloader/ClassLoaderSafeConnectorSplitManager.java
index 3480b85..f19335c 100644
--- a/presto-spi/src/main/java/com/facebook/presto/spi/classloader/ClassLoaderSafeConnectorSplitManager.java
+++ b/presto-spi/src/main/java/com/facebook/presto/spi/classloader/ClassLoaderSafeConnectorSplitManager.java
@@ -48,8 +48,14 @@ public final class ClassLoaderSafeConnectorSplitManager
     @Override
     public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions)
     {
+        return getPartitionSplits(table, partitions, TupleDomain.<ConnectorColumnHandle>all());
+    }
+
+    @Override
+    public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions, TupleDomain<ConnectorColumnHandle> tupleDomain)
+    {
         try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
-            return delegate.getPartitionSplits(table, partitions);
+            return delegate.getPartitionSplits(table, partitions, tupleDomain);
         }
     }
 
diff --git a/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchSplitManager.java b/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchSplitManager.java
index fb0faab..d46183e 100644
--- a/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchSplitManager.java
+++ b/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchSplitManager.java
@@ -59,6 +59,12 @@ public class TpchSplitManager
     }
 
     @Override
+    public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions, TupleDomain<ConnectorColumnHandle> tupleDomain)
+    {
+        return getPartitionSplits(table, partitions);
+    }
+
+    @Override
     public ConnectorSplitSource getPartitionSplits(ConnectorTableHandle table, List<ConnectorPartition> partitions)
     {
         checkNotNull(partitions, "partitions is null");
-- 
1.7.1

