Do not store node assignments for bucketed shards

presto-raptor/src/main/java/com/facebook/presto/raptor/RaptorSplitManager.java

@@ -103,7 +103,8 @@ public ConnectorSplitSource getSplits(ConnectorTransactionHandle transaction, Co
         boolean merged = bucketed && !table.isDelete() && (table.getBucketCount().getAsInt() >= getOneSplitPerBucketThreshold(session));
         OptionalLong transactionId = table.getTransactionId();
         Optional<Map<Integer, String>> bucketToNode = handle.getPartitioning().map(RaptorPartitioningHandle::getBucketToNode);
-        return new RaptorSplitSource(tableId, bucketed, merged, effectivePredicate, transactionId, bucketToNode);
+        verify(bucketed == bucketToNode.isPresent(), "mismatched bucketCount and bucketToNode presence");
+        return new RaptorSplitSource(tableId, merged, effectivePredicate, transactionId, bucketToNode);
     }
 
     private static List<HostAddress> getAddressesForNodes(Map<String, Node> nodeMap, Iterable<String> nodeIdentifiers)
@@ -145,7 +146,6 @@ private class RaptorSplitSource
 
         public RaptorSplitSource(
                 long tableId,
-                boolean bucketed,
                 boolean merged,
                 TupleDomain<RaptorColumnHandle> effectivePredicate,
                 OptionalLong transactionId,
@@ -155,7 +155,15 @@ public RaptorSplitSource(
             this.effectivePredicate = requireNonNull(effectivePredicate, "effectivePredicate is null");
             this.transactionId = requireNonNull(transactionId, "transactionId is null");
             this.bucketToNode = requireNonNull(bucketToNode, "bucketToNode is null");
-            this.iterator = new SynchronizedResultIterator<>(shardManager.getShardNodes(tableId, bucketed, merged, effectivePredicate));
+
+            ResultIterator<BucketShards> iterator;
+            if (bucketToNode.isPresent()) {
+                iterator = shardManager.getShardNodesBucketed(tableId, merged, bucketToNode.get(), effectivePredicate);
+            }
+            else {
+                iterator = shardManager.getShardNodes(tableId, effectivePredicate);
+            }
+            this.iterator = new SynchronizedResultIterator<>(iterator);
         }
 
         @Override
@@ -238,26 +246,15 @@ private ConnectorSplit createSplit(BucketShards bucketShards)
 
         private ConnectorSplit createBucketSplit(int bucketNumber, Set<ShardNodes> shards)
         {
+            // Bucket splits contain all the shards for the bucket
+            // and run on the node assigned to the bucket.
+
             String nodeId = bucketToNode.get().get(bucketNumber);
             Node node = nodesById.get(nodeId);
             if (node == null) {
                 throw new PrestoException(NO_NODES_AVAILABLE, "Node for bucket is offline: " + nodeId);
             }
 
-            // Bucket splits contain all the shards for the bucket and run on
-            // the node assigned to the bucket. Reassign any of the bucket
-            // shards not currently assigned to the node for the bucket.
-            for (ShardNodes shard : shards) {
-                if (!shard.getNodeIdentifiers().contains(nodeId)) {
-                    shardManager.assignShard(tableId, shard.getShardUuid(), nodeId, false);
-                }
-                for (String shardNodeId : shard.getNodeIdentifiers()) {
-                    if (!shardNodeId.equals(nodeId)) {
-                        shardManager.unassignShard(tableId, shard.getShardUuid(), shardNodeId);
-                    }
-                }
-            }
-
             Set<UUID> shardUuids = shards.stream()
                     .map(ShardNodes::getShardUuid)
                     .collect(toSet());

presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/DatabaseShardManager.java

@@ -183,7 +183,6 @@ public void createTable(long tableId, List<ColumnInfo> columns, boolean bucketed
                     "  shard_id BIGINT NOT NULL,\n" +
                     "  shard_uuid BINARY(16) NOT NULL,\n" +
                     "  bucket_number INT NOT NULL\n," +
-                    "  node_ids VARBINARY(128) NOT NULL,\n" +
                     tableColumns +
                     "  PRIMARY KEY (bucket_number, shard_uuid),\n" +
                     "  UNIQUE (shard_id),\n" +
@@ -402,11 +401,19 @@ private static void bindLongs(PreparedStatement statement, Iterable<Long> values
     private static void insertShardsAndIndex(long tableId, List<ColumnInfo> columns, Collection<ShardInfo> shards, Map<String, Integer> nodeIds, Handle handle)
             throws SQLException
     {
+        if (shards.isEmpty()) {
+            return;
+        }
+        boolean bucketed = shards.iterator().next().getBucketNumber().isPresent();
+
         Connection connection = handle.getConnection();
         try (IndexInserter indexInserter = new IndexInserter(connection, tableId, columns)) {
             for (List<ShardInfo> batch : partition(shards, batchSize(connection))) {
                 List<Long> shardIds = insertShards(connection, tableId, batch);
-                insertShardNodes(connection, nodeIds, shardIds, batch);
+
+                if (!bucketed) {
+                    insertShardNodes(connection, nodeIds, shardIds, batch);
+                }
 
                 for (int i = 0; i < batch.size(); i++) {
                     ShardInfo shard = batch.get(i);
@@ -448,9 +455,15 @@ public Set<ShardMetadata> getNodeShards(String nodeIdentifier)
     }
 
     @Override
-    public ResultIterator<BucketShards> getShardNodes(long tableId, boolean bucketed, boolean merged, TupleDomain<RaptorColumnHandle> effectivePredicate)
+    public ResultIterator<BucketShards> getShardNodes(long tableId, TupleDomain<RaptorColumnHandle> effectivePredicate)
+    {
+        return new ShardIterator(tableId, false, Optional.empty(), effectivePredicate, dbi);
+    }
+
+    @Override
+    public ResultIterator<BucketShards> getShardNodesBucketed(long tableId, boolean merged, Map<Integer, String> bucketToNode, TupleDomain<RaptorColumnHandle> effectivePredicate)
     {
-        return new ShardIterator(tableId, bucketed, merged, effectivePredicate, dbi);
+        return new ShardIterator(tableId, merged, Optional.of(bucketToNode), effectivePredicate, dbi);
     }
 
     @Override

presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/IndexInserter.java

@@ -65,14 +65,15 @@ public IndexInserter(Connection connection, long tableId, List<ColumnInfo> colum
         StringJoiner valueJoiner = new StringJoiner(", ");
         int index = 1;
 
-        nameJoiner.add("shard_id").add("shard_uuid").add("node_ids");
+        nameJoiner.add("shard_id").add("shard_uuid");
         valueJoiner.add("?").add("?").add("?");
         index += 3;
 
         if (bucketed) {
             nameJoiner.add("bucket_number");
-            valueJoiner.add("?");
-            index++;
+        }
+        else {
+            nameJoiner.add("node_ids");
         }
 
         for (ColumnInfo column : columns) {
@@ -121,14 +122,14 @@ public void insert(long shardId, UUID shardUuid, OptionalInt bucketNumber, Set<I
     {
         statement.setLong(1, shardId);
         statement.setBytes(2, uuidToBytes(shardUuid));
-        statement.setBytes(3, intArrayToBytes(nodeIds));
 
         if (bucketed) {
             checkArgument(bucketNumber.isPresent(), "shard bucket missing for bucketed table");
-            statement.setInt(4, bucketNumber.getAsInt());
+            statement.setInt(3, bucketNumber.getAsInt());
         }
         else {
             checkArgument(!bucketNumber.isPresent(), "shard bucket present for non-bucketed table");
+            statement.setBytes(3, intArrayToBytes(nodeIds));
         }
 
         for (ColumnInfo column : columns) {

presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardIterator.java

@@ -30,6 +30,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.OptionalInt;
 import java.util.Set;
 import java.util.UUID;
@@ -51,23 +52,28 @@ final class ShardIterator
     private static final Logger log = Logger.get(ShardIterator.class);
     private final Map<Integer, String> nodeMap = new HashMap<>();
 
-    private final boolean bucketed;
     private final boolean merged;
+    private final Map<Integer, String> bucketToNode;
     private final ShardDao dao;
     private final Connection connection;
     private final PreparedStatement statement;
     private final ResultSet resultSet;
     private boolean first = true;
 
-    public ShardIterator(long tableId, boolean bucketed, boolean merged, TupleDomain<RaptorColumnHandle> effectivePredicate, IDBI dbi)
+    public ShardIterator(
+            long tableId,
+            boolean merged,
+            Optional<Map<Integer, String>> bucketToNode,
+            TupleDomain<RaptorColumnHandle> effectivePredicate,
+            IDBI dbi)
     {
-        this.bucketed = bucketed;
         this.merged = merged;
-        ShardPredicate predicate = ShardPredicate.create(effectivePredicate, bucketed);
+        this.bucketToNode = bucketToNode.orElse(null);
+        ShardPredicate predicate = ShardPredicate.create(effectivePredicate, bucketToNode.isPresent());
 
         String sql;
-        if (bucketed) {
-            sql = "SELECT shard_uuid, node_ids, bucket_number FROM %s WHERE %s ORDER BY bucket_number";
+        if (bucketToNode.isPresent()) {
+            sql = "SELECT shard_uuid, bucket_number FROM %s WHERE %s ORDER BY bucket_number";
         }
         else {
             sql = "SELECT shard_uuid, node_ids FROM %s WHERE %s";
@@ -127,10 +133,20 @@ private BucketShards compute()
         }
 
         UUID shardUuid = uuidFromBytes(resultSet.getBytes("shard_uuid"));
-        List<Integer> nodeIds = intArrayFromBytes(resultSet.getBytes("node_ids"));
-        OptionalInt bucketNumber = bucketed ? OptionalInt.of(resultSet.getInt("bucket_number")) : OptionalInt.empty();
+        Set<String> nodeIdentifiers;
+        OptionalInt bucketNumber = OptionalInt.empty();
 
-        ShardNodes shard = new ShardNodes(shardUuid, getNodeIdentifiers(nodeIds, shardUuid));
+        if (bucketToNode != null) {
+            int bucket = resultSet.getInt("bucket_number");
+            bucketNumber = OptionalInt.of(bucket);
+            nodeIdentifiers = ImmutableSet.of(getBucketNode(bucket));
+        }
+        else {
+            List<Integer> nodeIds = intArrayFromBytes(resultSet.getBytes("node_ids"));
+            nodeIdentifiers = getNodeIdentifiers(nodeIds, shardUuid);
+        }
+
+        ShardNodes shard = new ShardNodes(shardUuid, nodeIdentifiers);
         return new BucketShards(bucketNumber, ImmutableSet.of(shard));
     }
 
@@ -155,8 +171,8 @@ private BucketShards computeMerged()
 
         do {
             UUID shardUuid = uuidFromBytes(resultSet.getBytes("shard_uuid"));
-            List<Integer> nodeIds = intArrayFromBytes(resultSet.getBytes("node_ids"));
-            Set<String> nodeIdentifiers = getNodeIdentifiers(nodeIds, shardUuid);
+            int bucket = resultSet.getInt("bucket_number");
+            Set<String> nodeIdentifiers = ImmutableSet.of(getBucketNode(bucket));
 
             shards.add(new ShardNodes(shardUuid, nodeIdentifiers));
         }
@@ -165,6 +181,15 @@ private BucketShards computeMerged()
         return new BucketShards(OptionalInt.of(bucketNumber), shards.build());
     }
 
+    private String getBucketNode(int bucket)
+    {
+        String node = bucketToNode.get(bucket);
+        if (node == null) {
+            throw new PrestoException(RAPTOR_ERROR, "No node mapping for bucket: " + bucket);
+        }
+        return node;
+    }
+
     private Set<String> getNodeIdentifiers(List<Integer> nodeIds, UUID shardUuid)
     {
         Function<Integer, String> fetchNode = id -> fetchNode(id, shardUuid);

presto-raptor/src/main/java/com/facebook/presto/raptor/metadata/ShardManager.java

@@ -57,9 +57,14 @@ public interface ShardManager
     Set<ShardMetadata> getNodeShards(String nodeIdentifier);
 
     /**
-     * Return the shard nodes a given table.
+     * Return the shard nodes for a non-bucketed table.
      */
-    ResultIterator<BucketShards> getShardNodes(long tableId, boolean bucketed, boolean merged, TupleDomain<RaptorColumnHandle> effectivePredicate);
+    ResultIterator<BucketShards> getShardNodes(long tableId, TupleDomain<RaptorColumnHandle> effectivePredicate);
+
+    /**
+     * Return the shard nodes for a bucketed table.
+     */
+    ResultIterator<BucketShards> getShardNodesBucketed(long tableId, boolean merged, Map<Integer, String> bucketToNode, TupleDomain<RaptorColumnHandle> effectivePredicate);
 
     /**
      * Assign a shard to a node.

presto-raptor/src/test/java/com/facebook/presto/raptor/metadata/TestDatabaseShardManager.java

@@ -299,7 +299,7 @@ public void testReplaceShardUuids()
         expectedAllUuids.addAll(expectedUuids);
 
         // check that shards are replaced in index table as well
-        Set<BucketShards> shardNodes = ImmutableSet.copyOf(shardManager.getShardNodes(tableId, false, false, TupleDomain.all()));
+        Set<BucketShards> shardNodes = ImmutableSet.copyOf(shardManager.getShardNodes(tableId, TupleDomain.all()));
         Set<UUID> actualAllUuids = shardNodes.stream()
                 .map(BucketShards::getShards)
                 .flatMap(Collection::stream)
@@ -391,7 +391,7 @@ public void testEmptyTable()
         List<ColumnInfo> columns = ImmutableList.of(new ColumnInfo(1, BIGINT));
         shardManager.createTable(tableId, columns, false);
 
-        try (ResultIterator<BucketShards> iterator = shardManager.getShardNodes(tableId, false, false, TupleDomain.all())) {
+        try (ResultIterator<BucketShards> iterator = shardManager.getShardNodes(tableId, TupleDomain.all())) {
             assertFalse(iterator.hasNext());
         }
     }
@@ -403,7 +403,7 @@ public void testEmptyTableBucketed()
         List<ColumnInfo> columns = ImmutableList.of(new ColumnInfo(1, BIGINT));
         shardManager.createTable(tableId, columns, true);
 
-        try (ResultIterator<BucketShards> iterator = shardManager.getShardNodes(tableId, true, true, TupleDomain.all())) {
+        try (ResultIterator<BucketShards> iterator = shardManager.getShardNodesBucketed(tableId, true, ImmutableMap.of(), TupleDomain.all())) {
             assertFalse(iterator.hasNext());
         }
     }
@@ -601,7 +601,7 @@ public void testShardPruningNoStats()
 
     private Set<ShardNodes> getShardNodes(long tableId, TupleDomain<RaptorColumnHandle> predicate)
     {
-        try (ResultIterator<BucketShards> iterator = shardManager.getShardNodes(tableId, false, false, predicate)) {
+        try (ResultIterator<BucketShards> iterator = shardManager.getShardNodes(tableId, predicate)) {
             return ImmutableSet.copyOf(concat(transform(iterator, i -> i.getShards().iterator())));
         }
     }

presto-raptor/src/test/java/com/facebook/presto/raptor/storage/TestShardEjector.java

@@ -136,7 +136,7 @@ public void testEjector()
 
         ejector.process();
 
-        shardManager.getShardNodes(tableId, false, false, TupleDomain.all());
+        shardManager.getShardNodes(tableId, TupleDomain.all());
 
         Set<UUID> ejectedShards = shards.subList(0, 4).stream()
                 .map(ShardInfo::getShardUuid)