Fix memory size calculation when memory table is truncated

Author: takezoe

Diff for: plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryPagesStore.java

@@ -136,7 +136,8 @@ public synchronized boolean contains(Long tableId)
 
     public synchronized void purge(long tableId)
     {
-        tables.remove(tableId);
+        TableData tableData = tables.remove(tableId);
+        currentBytes = currentBytes - tableData.getPages().stream().mapToLong(Page::getRetainedSizeInBytes).sum();
     }
 
     public synchronized void cleanUp(Set<Long> activeTableIds)

Diff for: plugin/trino-memory/src/test/java/io/trino/plugin/memory/TestMemoryPagesStore.java

@@ -133,6 +133,15 @@ public void testMemoryLimitExceeded()
                 .hasMessageMatching("Memory limit.*");
     }
 
+    @Test
+    public void testTruncate()
+    {
+        createTable(0L, 0L);
+        insertToTable(0L, createOneMegaBytePage(), 0L);
+        truncateTable(0L, 0L);
+        insertToTable(0L, createOneMegaBytePage(), 0L);
+    }
+
     private void insertToTable(long tableId, Long... activeTableIds)
     {
         insertToTable(tableId, createPage(), activeTableIds);
@@ -159,6 +168,16 @@ private void createTable(long tableId, Long... activeTableIds)
         pageSink.finish();
     }
 
+    private void truncateTable(long tableId, Long... activeTableIds)
+    {
+        ConnectorPageSink pageSink = pageSinkProvider.createPageSink(
+                MemoryTransactionHandle.INSTANCE,
+                SESSION,
+                createOverwriteMemoryInsertTableHandle(tableId, activeTableIds),
+                TESTING_PAGE_SINK_ID);
+        pageSink.finish();
+    }
+
     private static ConnectorOutputTableHandle createMemoryOutputTableHandle(long tableId, Long... activeTableIds)
     {
         return new MemoryOutputTableHandle(tableId, ImmutableSet.copyOf(activeTableIds));
@@ -169,6 +188,11 @@ private static ConnectorInsertTableHandle createMemoryInsertTableHandle(long tab
         return new MemoryInsertTableHandle(tableId, InsertMode.APPEND, ImmutableSet.copyOf(activeTableIds));
     }
 
+    private static ConnectorInsertTableHandle createOverwriteMemoryInsertTableHandle(long tableId, Long[] activeTableIds)
+    {
+        return new MemoryInsertTableHandle(tableId, InsertMode.OVERWRITE, ImmutableSet.copyOf(activeTableIds));
+    }
+
     private static Page createPage()
     {
         BlockBuilder blockBuilder = BIGINT.createFixedSizeBlockBuilder(POSITIONS_PER_PAGE);