Do not allocate memory for all queries scanning system.* tables

Extend the logic to specially treat metadata queries and not allocate any memory for those, so they are scheduled immediatelly, to cover all schematas from system catalog (previously only jdbc schama was covered).
trinodb · Oct 19, 2023 · a098d04 · a098d04
1 parent 74012e3
commit a098d04
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 7 deletions.
diff --git a/...ava/io/trino/execution/scheduler/faulttolerant/NoMemoryAwarePartitionMemoryEstimator.java b/...ava/io/trino/execution/scheduler/faulttolerant/NoMemoryAwarePartitionMemoryEstimator.java
@@ -95,9 +95,7 @@ private boolean isNoMemoryFragment(PlanFragment fragment, Function<PlanFragmentI
         private static boolean isMetadataTableScan(TableScanNode tableScanNode)
         {
             return (tableScanNode.getTable().getConnectorHandle() instanceof InformationSchemaTableHandle) ||
-                    (tableScanNode.getTable().getCatalogHandle().getCatalogName().equals(GlobalSystemConnector.NAME) &&
-                            (tableScanNode.getTable().getConnectorHandle() instanceof SystemTableHandle systemHandle) &&
-                            systemHandle.getSchemaName().equals("jdbc"));
+                    (tableScanNode.getTable().getCatalogHandle().getCatalogName().equals(GlobalSystemConnector.NAME) && (tableScanNode.getTable().getConnectorHandle() instanceof SystemTableHandle));
         }
     }
 }
diff --git a/...io/trino/execution/scheduler/faulttolerant/TestNoMemoryAwarePartitionMemoryEstimator.java b/...io/trino/execution/scheduler/faulttolerant/TestNoMemoryAwarePartitionMemoryEstimator.java
@@ -21,6 +21,8 @@
 import io.trino.Session;
 import io.trino.connector.informationschema.InformationSchemaTable;
 import io.trino.connector.informationschema.InformationSchemaTableHandle;
+import io.trino.connector.system.GlobalSystemConnector;
+import io.trino.connector.system.SystemTableHandle;
 import io.trino.cost.StatsAndCosts;
 import io.trino.metadata.TableHandle;
 import io.trino.operator.RetryPolicy;
@@ -130,6 +132,34 @@ public void testRemoteFromInformationSchemaAndTpchTableScans()
         assertThat(estimator).isInstanceOf(MockDelegatePatitionMemoryEstimator.class);
     }
 
+    @Test
+    public void testSystemJdbcTableScan()
+    {
+        PartitionMemoryEstimator estimator = createEstimator(tableScanPlanFragment(
+                "ts",
+                new TableHandle(
+                        GlobalSystemConnector.CATALOG_HANDLE,
+                        new SystemTableHandle("jdbc", "tables", TupleDomain.all()),
+                        TestingTransactionHandle.create())));
+        assertThat(estimator).isInstanceOf(NoMemoryPartitionMemoryEstimator.class);
+        PartitionMemoryEstimator.MemoryRequirements noMemoryRequirements = new PartitionMemoryEstimator.MemoryRequirements(DataSize.ofBytes(0));
+        assertThat(estimator.getInitialMemoryRequirements()).isEqualTo(noMemoryRequirements);
+    }
+
+    @Test
+    public void testSystemMetadataTableScan()
+    {
+        PartitionMemoryEstimator estimator = createEstimator(tableScanPlanFragment(
+                "ts",
+                new TableHandle(
+                        GlobalSystemConnector.CATALOG_HANDLE,
+                        new SystemTableHandle("metadata", "blah", TupleDomain.all()),
+                        TestingTransactionHandle.create())));
+        assertThat(estimator).isInstanceOf(NoMemoryPartitionMemoryEstimator.class);
+        PartitionMemoryEstimator.MemoryRequirements noMemoryRequirements = new PartitionMemoryEstimator.MemoryRequirements(DataSize.ofBytes(0));
+        assertThat(estimator.getInitialMemoryRequirements()).isEqualTo(noMemoryRequirements);
+    }
+
     private static PlanFragment getParentFragment(PlanFragment... childFragments)
     {
         ImmutableList<PlanFragmentId> childFragmentIds = Stream.of(childFragments)
@@ -161,13 +191,18 @@ private PartitionMemoryEstimator createEstimator(PlanFragment planFragment, Plan
     }
 
     private static PlanFragment tableScanPlanFragment(String fragmentId, ConnectorTableHandle tableHandle)
+    {
+        return tableScanPlanFragment(fragmentId, new TableHandle(
+                TEST_CATALOG_HANDLE,
+                tableHandle,
+                TestingTransactionHandle.create()));
+    }
+
+    private static PlanFragment tableScanPlanFragment(String fragmentId, TableHandle tableHandle)
     {
         TableScanNode informationSchemaViewsTableScan = new TableScanNode(
                 new PlanNodeId("tableScan"),
-                new TableHandle(
-                        TEST_CATALOG_HANDLE,
-                        tableHandle,
-                        TestingTransactionHandle.create()),
+                tableHandle,
                 ImmutableList.of(),
                 ImmutableMap.of(),
                 TupleDomain.all(),