Add ConnectorNodePartitioningProvider#getBucketCount

adityaramkumar · Jun 15, 2020 · 21f0b04 · 21f0b04
1 parent cbe1cb1
commit 21f0b04
Show file tree

Hide file tree

Showing 8 changed files with 54 additions and 0 deletions.
diff --git a/...src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleNodePartitioningProvider.java b/...src/main/java/com/facebook/presto/plugin/blackhole/BlackHoleNodePartitioningProvider.java
@@ -80,4 +80,11 @@ public BucketFunction getBucketFunction(
             return (int) (hash % bucketCount);
         };
     }
+
+    @Override
+    public int getBucketCount(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle)
+    {
+        // create one bucket per node
+        return nodeManager.getRequiredWorkerNodes().size();
+    }
 }
diff --git a/presto-hive/src/main/java/com/facebook/presto/hive/HiveNodePartitioningProvider.java b/presto-hive/src/main/java/com/facebook/presto/hive/HiveNodePartitioningProvider.java
@@ -80,6 +80,13 @@ public ToIntFunction<ConnectorSplit> getSplitBucketFunction(
                 .orElseThrow(() -> new IllegalArgumentException("Bucket number not set in split"));
     }
 
+    @Override
+    public int getBucketCount(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle)
+    {
+        HivePartitioningHandle handle = (HivePartitioningHandle) partitioningHandle;
+        return handle.getBucketCount();
+    }
+
     @Override
     public List<ConnectorPartitionHandle> listPartitionHandles(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle)
     {

diff --git a/...-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotNodePartitioningProvider.java b/...-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotNodePartitioningProvider.java
@@ -58,4 +58,10 @@ public BucketFunction getBucketFunction(
     {
         return null;
     }
+
+    @Override
+    public int getBucketCount(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle)
+    {
+        return 1;
+    }
 }
diff --git a/presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorNodePartitioningProvider.java b/presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorNodePartitioningProvider.java
@@ -77,4 +77,11 @@ public BucketFunction getBucketFunction(ConnectorTransactionHandle transaction,
     {
         return new RaptorBucketFunction(bucketCount, partitionChannelTypes);
     }
+
+    @Override
+    public int getBucketCount(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle)
+    {
+        RaptorPartitioningHandle handle = (RaptorPartitioningHandle) partitioningHandle;
+        return handle.getBucketToNode().size();
+    }
 }
diff --git a/...pi/src/main/java/com/facebook/presto/spi/connector/ConnectorNodePartitioningProvider.java b/...pi/src/main/java/com/facebook/presto/spi/connector/ConnectorNodePartitioningProvider.java
@@ -51,4 +51,9 @@ BucketFunction getBucketFunction(
             ConnectorPartitioningHandle partitioningHandle,
             List<Type> partitionChannelTypes,
             int bucketCount);
+
+    int getBucketCount(
+            ConnectorTransactionHandle transactionHandle,
+            ConnectorSession session,
+            ConnectorPartitioningHandle partitioningHandle);
 }
diff --git a/...om/facebook/presto/spi/connector/classloader/ClassLoaderSafeNodePartitioningProvider.java b/...om/facebook/presto/spi/connector/classloader/ClassLoaderSafeNodePartitioningProvider.java
@@ -78,4 +78,12 @@ public ToIntFunction<ConnectorSplit> getSplitBucketFunction(ConnectorTransaction
             return delegate.getSplitBucketFunction(transactionHandle, session, partitioningHandle);
         }
     }
+
+    @Override
+    public int getBucketCount(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle)
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
+            return delegate.getBucketCount(transactionHandle, session, partitioningHandle);
+        }
+    }
 }
diff --git a/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsNodePartitioningProvider.java b/presto-tpcds/src/main/java/com/facebook/presto/tpcds/TpcdsNodePartitioningProvider.java
@@ -70,4 +70,12 @@ public BucketFunction getBucketFunction(ConnectorTransactionHandle transactionHa
     {
         throw new UnsupportedOperationException();
     }
+
+    @Override
+    public int getBucketCount(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle)
+    {
+        Set<Node> nodes = nodeManager.getRequiredWorkerNodes();
+        checkState(!nodes.isEmpty(), "No TPCDS nodes available");
+        return nodes.size() * splitsPerNode;
+    }
 }
diff --git a/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchNodePartitioningProvider.java b/presto-tpch/src/main/java/com/facebook/presto/tpch/TpchNodePartitioningProvider.java
@@ -70,4 +70,10 @@ public BucketFunction getBucketFunction(ConnectorTransactionHandle transactionHa
         checkArgument(partitionChannelTypes.equals(ImmutableList.of(BIGINT)), "Expected one BIGINT parameter");
         return new TpchBucketFunction(bucketCount, rowsPerBucket);
     }
+
+    @Override
+    public int getBucketCount(ConnectorTransactionHandle transactionHandle, ConnectorSession session, ConnectorPartitioningHandle partitioningHandle)
+    {
+        return nodeManager.getRequiredWorkerNodes().size() * splitsPerNode;
+    }
 }