Memory Connector: optimize split/logicalpart allocation and remove memory.splits-per-node usage

Author: peiwangdb

hetu-hive-functions/src/test/java/io/hetu/core/hive/dynamicfunctions/TestDynamicHiveScalarFunction.java

@@ -94,8 +94,7 @@ public void setUpClass()
             metastoreConfig.put("hetu.metastore.hetufilesystem.path", folder.newFolder("metastore").getAbsolutePath());
             queryRunner.getCoordinator().loadMetastore(metastoreConfig);
             queryRunner.createCatalog("memory", "memory",
-                    ImmutableMap.of("memory.spill-path", folder.newFolder("memory-connector").getAbsolutePath(),
-                            "memory.splits-per-node", "1"));
+                    ImmutableMap.of("memory.spill-path", folder.newFolder("memory-connector").getAbsolutePath()));
         }
         catch (Exception e) {
             closeAllSuppress(e, queryRunner);

presto-geospatial/src/test/java/io/prestosql/plugin/geospatial/BenchmarkGeometryAggregations.java

@@ -86,8 +86,7 @@ public void setUp()
             metastoreConfig.put("hetu.metastore.cache.type", "local");
             queryRunner.loadMetastore(metastoreConfig);
             queryRunner.createCatalog("memory", new MemoryConnectorFactory(),
-                    ImmutableMap.of("memory.spill-path", folder.newFolder("memory-connector").getAbsolutePath(),
-                            "memory.splits-per-node", "1"));
+                    ImmutableMap.of("memory.spill-path", folder.newFolder("memory-connector").getAbsolutePath()));
 
             Path path = Paths.get(BenchmarkGeometryAggregations.class.getClassLoader().getResource("us-states.tsv").getPath());
             String polygonValues = Files.lines(path)

presto-geospatial/src/test/java/io/prestosql/plugin/geospatial/BenchmarkSpatialJoin.java

@@ -98,8 +98,7 @@ public void setUp()
             metastoreConfig.put("hetu.metastore.cache.type", "local");
             queryRunner.loadMetastore(metastoreConfig);
             queryRunner.createCatalog("memory", new MemoryConnectorFactory(),
-                    ImmutableMap.of("memory.spill-path", folder.newFolder("memory-connector").getAbsolutePath(),
-                            "memory.splits-per-node", "1"));
+                    ImmutableMap.of("memory.spill-path", folder.newFolder("memory-connector").getAbsolutePath()));
 
             Path path = Paths.get(BenchmarkSpatialJoin.class.getClassLoader().getResource("us-states.tsv").getPath());
             String polygonValues = Files.lines(path)

presto-geospatial/src/test/java/io/prestosql/plugin/geospatial/TestSpatialJoinPlanning.java

@@ -90,8 +90,7 @@ private static LocalQueryRunner createQueryRunner()
         metastoreConfig.put("hetu.metastore.cache.type", "local");
         queryRunner.loadMetastore(metastoreConfig);
         queryRunner.createCatalog("memory", new MemoryConnectorFactory(),
-                ImmutableMap.of("memory.spill-path", folder.newFolder("memory-connector").getAbsolutePath(),
-                        "memory.splits-per-node", "1"));
+                ImmutableMap.of("memory.spill-path", folder.newFolder("memory-connector").getAbsolutePath()));
 
         queryRunner.execute(format("CREATE TABLE kdb_tree AS SELECT '%s' AS v", KDB_TREE_JSON));
         queryRunner.execute("CREATE TABLE points (lng, lat, name) AS (VALUES (2.1e0, 2.1e0, 'x'))");

presto-jdbc/src/test/java/io/hetu/core/jdbc/TestHetuConnection.java

@@ -54,8 +54,7 @@ public void setupServer()
         server.installPlugin(new MemoryPlugin());
         server.loadMetastore(metastoreConfig);
         server.createCatalog("memory", "memory",
-                ImmutableMap.of("memory.spill-path", folder.newFolder("memory-connector").getAbsolutePath(),
-                        "memory.splits-per-node", "1"));
+                ImmutableMap.of("memory.spill-path", folder.newFolder("memory-connector").getAbsolutePath()));
 
         try (Connection connection = createConnection();
                 Statement statement = connection.createStatement()) {

presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryConfig.java

@@ -37,29 +37,13 @@
 
 public class MemoryConfig
 {
-    private int splitsPerNode = Math.max(Runtime.getRuntime().availableProcessors(), 1);
     private DataSize maxDataPerNode = new DataSize(256, DataSize.Unit.MEGABYTE);
     private DataSize maxLogicalPartSize = new DataSize(256, MEGABYTE);
     private DataSize maxPageSize = new DataSize(512, DataSize.Unit.KILOBYTE);
     private Duration processingDelay = new Duration(5, TimeUnit.SECONDS);
     private Path spillRoot;
     private int threadPoolSize = Math.max((Runtime.getRuntime().availableProcessors() / 2), 1);
 
-    public int getSplitsPerNode()
-    {
-        return splitsPerNode;
-    }
-
-    @Mandatory(name = "memory.splits-per-node",
-            description = "Number of splits to create per node. Default value is number of available processors on the coordinator." +
-                    " Value is ignored on the workers.")
-    @Config("memory.splits-per-node")
-    public MemoryConfig setSplitsPerNode(int splitsPerNode)
-    {
-        this.splitsPerNode = splitsPerNode;
-        return this;
-    }
-
     @NotNull
     public Path getSpillRoot()
     {

presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryDataFragment.java

@@ -30,15 +30,18 @@ public class MemoryDataFragment
 
     private final HostAddress hostAddress;
     private final long rows;
+    private final int logicalPartCount;
 
     @JsonCreator
     public MemoryDataFragment(
             @JsonProperty("hostAddress") HostAddress hostAddress,
-            @JsonProperty("rows") long rows)
+            @JsonProperty("rows") long rows,
+            @JsonProperty("logicalPartCount") int logicalPartCount)
     {
         this.hostAddress = requireNonNull(hostAddress, "hostAddress is null");
         checkArgument(rows >= 0, "Rows number can not be negative");
         this.rows = rows;
+        this.logicalPartCount = logicalPartCount;
     }
 
     @JsonProperty
@@ -53,6 +56,12 @@ public long getRows()
         return rows;
     }
 
+    @JsonProperty
+    public int getLogicalPartCount()
+    {
+        return logicalPartCount;
+    }
+
     public Slice toSlice()
     {
         return Slices.wrappedBuffer(MEMORY_DATA_FRAGMENT_CODEC.toJsonBytes(this));
@@ -66,6 +75,6 @@ public static MemoryDataFragment fromSlice(Slice fragment)
     public static MemoryDataFragment merge(MemoryDataFragment a, MemoryDataFragment b)
     {
         checkArgument(a.getHostAddress().equals(b.getHostAddress()), "Can not merge fragments from different hosts");
-        return new MemoryDataFragment(a.getHostAddress(), a.getRows() + b.getRows());
+        return new MemoryDataFragment(a.getHostAddress(), a.getRows() + b.getRows(), Math.max(a.getLogicalPartCount(), b.getLogicalPartCount()));
     }
 }

presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryMetadata.java

@@ -283,8 +283,7 @@ public void renameTable(ConnectorSession session, ConnectorTableHandle tableHand
                 oldTableHandle.getActiveTableIds(),
                 oldTableHandle.getColumns(),
                 oldTableHandle.getSortedBy(),
-                oldTableHandle.getIndexColumns(),
-                oldTableHandle.getSplitsPerNode());
+                oldTableHandle.getIndexColumns());
 
         metastore.alterTable(MEM_KEY, oldInfo.getSchemaName(), oldInfo.getTableName(),
                 TableEntity.builder()
@@ -398,8 +397,7 @@ public synchronized MemoryWriteTableHandle beginCreateTable(ConnectorSession ses
                 getTableIdSet(nextId),
                 columnHandles,
                 sortedBy,
-                indexColumns,
-                config.getSplitsPerNode());
+                indexColumns);
     }
 
     private void checkSchemaExists(String schemaName, boolean expectExist)

presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryPageSinkProvider.java

@@ -73,7 +73,6 @@ public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHa
         pagesStore.refreshTables(memoryOutputTableHandle.getActiveTableIds());
         pagesStore.initialize(tableId,
                 memoryOutputTableHandle.isCompressionEnabled(),
-                memoryOutputTableHandle.getSplitsPerNode(),
                 memoryOutputTableHandle.getColumns(),
                 memoryOutputTableHandle.getSortedBy(),
                 memoryOutputTableHandle.getIndexColumns());
@@ -114,7 +113,6 @@ public ConnectorPageSink createPageSink(ConnectorTransactionHandle transactionHa
             //
             pagesStore.initialize(tableId,
                     memoryOutputTableHandle.isCompressionEnabled(),
-                    memoryOutputTableHandle.getSplitsPerNode(),
                     memoryOutputTableHandle.getColumns(),
                     memoryOutputTableHandle.getSortedBy(),
                     memoryOutputTableHandle.getIndexColumns());
@@ -160,7 +158,8 @@ public CompletableFuture<?> appendPage(Page page)
         public CompletableFuture<Collection<Slice>> finish()
         {
             tablesManager.finishUpdatingTable(tableId);
-            return completedFuture(ImmutableList.of(new MemoryDataFragment(currentHostAddress, addedRows).toSlice()));
+            int lpCount = tablesManager.getTableLpCount(tableId);
+            return completedFuture(ImmutableList.of(new MemoryDataFragment(currentHostAddress, addedRows, lpCount).toSlice()));
         }
 
         @Override

presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryPageSourceProvider.java

@@ -74,8 +74,7 @@ public ConnectorPageSource createPageSource(
     {
         MemorySplit memorySplit = (MemorySplit) split;
         long tableId = memorySplit.getTable();
-        int partNumber = memorySplit.getPartNumber();
-        int totalParts = memorySplit.getTotalPartsPerWorker();
+        int lpNum = memorySplit.getLogicalPartNumber();
         long expectedRows = memorySplit.getExpectedRows();
         MemoryTableHandle memoryTable = (MemoryTableHandle) table;
         OptionalDouble sampleRatio = memoryTable.getSampleRatio();
@@ -88,8 +87,7 @@ public ConnectorPageSource createPageSource(
                 .map(MemoryColumnHandle::getColumnIndex).collect(toList());
         List<Page> pages = pagesStore.getPages(
                 tableId,
-                partNumber,
-                totalParts,
+                lpNum,
                 columnIndexes,
                 expectedRows,
                 memorySplit.getLimit(),

presto-memory/src/main/java/io/prestosql/plugin/memory/MemorySplit.java

@@ -30,28 +30,23 @@ public class MemorySplit
         implements ConnectorSplit
 {
     private final long table;
-    private final int totalPartsPerWorker; // how many concurrent reads there will be from one worker
-    private final int partNumber; // part of the pages on one worker that this splits is responsible
+    private final int lpNum; // the order of the logicalPart (equivalent to a split) in each node, starting from 0 (1 split, 0 is the available index)
     private final HostAddress address;
     private final long expectedRows;
     private final OptionalLong limit;
 
     @JsonCreator
     public MemorySplit(
             @JsonProperty("table") long table,
-            @JsonProperty("partNumber") int partNumber,
-            @JsonProperty("totalPartsPerWorker") int totalPartsPerWorker,
+            @JsonProperty("lpNum") int lpNum,
             @JsonProperty("address") HostAddress address,
             @JsonProperty("expectedRows") long expectedRows,
             @JsonProperty("limit") OptionalLong limit)
     {
-        checkState(partNumber >= 0, "partNumber must be >= 0");
-        checkState(totalPartsPerWorker >= 1, "totalPartsPerWorker must be >= 1");
-        checkState(totalPartsPerWorker > partNumber, "totalPartsPerWorker must be > partNumber");
+        checkState(lpNum >= 0, "lpCount must be >= 0");
 
         this.table = requireNonNull(table, "table is null");
-        this.partNumber = partNumber;
-        this.totalPartsPerWorker = totalPartsPerWorker;
+        this.lpNum = lpNum;
         this.address = requireNonNull(address, "address is null");
         this.expectedRows = expectedRows;
         this.limit = limit;
@@ -64,15 +59,9 @@ public long getTable()
     }
 
     @JsonProperty
-    public int getTotalPartsPerWorker()
+    public int getLogicalPartNumber()
     {
-        return totalPartsPerWorker;
-    }
-
-    @JsonProperty
-    public int getPartNumber()
-    {
-        return partNumber;
+        return lpNum;
     }
 
     @Override
@@ -122,23 +111,21 @@ public boolean equals(Object obj)
         }
         MemorySplit other = (MemorySplit) obj;
         return Objects.equals(this.table, other.table) &&
-                Objects.equals(this.totalPartsPerWorker, other.totalPartsPerWorker) &&
-                Objects.equals(this.partNumber, other.partNumber);
+                Objects.equals(this.lpNum, other.lpNum);
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(table, totalPartsPerWorker, partNumber);
+        return Objects.hash(table, lpNum);
     }
 
     @Override
     public String toString()
     {
         return "MemorySplit{" +
                 "table=" + table +
-                ", totalPartsPerWorker=" + totalPartsPerWorker +
-                ", partNumber=" + partNumber +
+                ", logicalPart number=" + lpNum +
                 '}';
     }
 }

presto-memory/src/main/java/io/prestosql/plugin/memory/MemorySplitManager.java

@@ -30,13 +30,11 @@
 public final class MemorySplitManager
         implements ConnectorSplitManager
 {
-    private final int splitsPerNode;
     private final MemoryMetadata metadata;
 
     @Inject
     public MemorySplitManager(MemoryConfig config, MemoryMetadata metadata)
     {
-        this.splitsPerNode = config.getSplitsPerNode();
         this.metadata = metadata;
     }
 
@@ -47,16 +45,15 @@ public ConnectorSplitSource getSplits(ConnectorTransactionHandle transactionHand
 
         List<MemoryDataFragment> dataFragments = metadata.getDataFragments(table.getId());
 
-        int totalRows = 0;
-
+        long totalRows = 0;
         ImmutableList.Builder<ConnectorSplit> splits = ImmutableList.builder();
 
         for (MemoryDataFragment dataFragment : dataFragments) {
             long rows = dataFragment.getRows();
+            int lpCount = dataFragment.getLogicalPartCount();
             totalRows += rows;
-
-            for (int i = 0; i < splitsPerNode; i++) {
-                splits.add(new MemorySplit(table.getId(), i, splitsPerNode, dataFragment.getHostAddress(), rows, OptionalLong.empty()));
+            for (int i = 0; i <= lpCount; i++) {
+                splits.add(new MemorySplit(table.getId(), i, dataFragment.getHostAddress(), rows, OptionalLong.empty()));
             }
         }
         return new FixedSplitSource(splits.build());

presto-memory/src/main/java/io/prestosql/plugin/memory/MemoryWriteTableHandle.java

@@ -38,7 +38,6 @@ public final class MemoryWriteTableHandle
     private final List<String> indexColumns;
     private final String schemaName;
     private final String tableName;
-    private final int splitsPerNode;
 
     @JsonCreator
     public MemoryWriteTableHandle(
@@ -49,14 +48,12 @@ public MemoryWriteTableHandle(
             @JsonProperty("activeTableIds") Set<Long> activeTableIds,
             @JsonProperty("columns") List<MemoryColumnHandle> columns,
             @JsonProperty("sortedBy") List<SortingColumn> sortedBy,
-            @JsonProperty("indexColumns") List<String> indexColumns,
-            @JsonProperty("splitsPerNode") int splitsPerNode)
+            @JsonProperty("indexColumns") List<String> indexColumns)
     {
         this.table = table;
         this.schemaName = schemaName;
         this.tableName = tableName;
         this.compressionEnabled = compressionEnabled;
-        this.splitsPerNode = splitsPerNode;
         this.activeTableIds = requireNonNull(activeTableIds, "activeTableIds is null");
         this.columns = requireNonNull(columns, "columns is null");
         this.sortedBy = requireNonNull(sortedBy, "sortedBy is null");
@@ -66,7 +63,7 @@ public MemoryWriteTableHandle(
     @VisibleForTesting
     MemoryWriteTableHandle(long table, Set<Long> activeTableIds)
     {
-        this(table, "", "", SPILL_COMPRESSION_DEFAULT_VALUE, activeTableIds, Collections.emptyList(), Collections.emptyList(), Collections.emptyList(), 1);
+        this(table, "", "", SPILL_COMPRESSION_DEFAULT_VALUE, activeTableIds, Collections.emptyList(), Collections.emptyList(), Collections.emptyList());
     }
 
     @JsonProperty
@@ -117,12 +114,6 @@ public String getTableName()
         return tableName;
     }
 
-    @JsonProperty
-    public int getSplitsPerNode()
-    {
-        return splitsPerNode;
-    }
-
     @Override
     public String toString()
     {

presto-memory/src/main/java/io/prestosql/plugin/memory/data/LogicalPart.java

@@ -89,8 +89,7 @@ enum LogicalPartState
     private final Set<Integer> indexChannels;
     private final long maxLogicalPartBytes;
     private final int maxPageSizeBytes;
-    private final int splitNum;
-    private final int logicalPartNum;
+    private final int lpNum;
     private final boolean compressionEnabled;
 
     // indexes
@@ -144,13 +143,11 @@ public LogicalPart(
             int maxPageSizeBytes,
             TypeManager typeManager,
             PagesSerde pagesSerde,
-            int splitNum,
-            int logicalPartNum,
+            int lpNum,
             boolean compressionEnabled)
     {
         this.tableDataRoot = tableDataRoot;
-        this.splitNum = splitNum;
-        this.logicalPartNum = logicalPartNum;
+        this.lpNum = lpNum;
         this.pages = new ArrayList<>();
         this.maxLogicalPartBytes = maxLogicalPartBytes;
         this.maxPageSizeBytes = maxPageSizeBytes;
@@ -701,7 +698,7 @@ void process()
 
     private String getPageFileName()
     {
-        return "split" + splitNum + "lp" + logicalPartNum;
+        return "logicalPart" + lpNum;
     }
 
     private synchronized void readPages()