From bad6578fdd0c47983b47ed4a0f6fdd453c1d8a29 Mon Sep 17 00:00:00 2001
From: Kamil Endruszkiewicz <kamil.endruszkiewicz@starburstdata.com>
Date: Tue, 24 Sep 2024 11:05:54 +0200
Subject: [PATCH] Alter allowSplittingReadIntoMultipleSubQueries check for
 systemTables

---
 .../src/main/java/io/trino/metadata/MetadataManager.java   | 3 +++
 .../java/io/trino/plugin/deltalake/DeltaLakeMetadata.java  | 4 ++--
 .../src/main/java/io/trino/plugin/hive/HiveMetadata.java   | 7 +------
 .../src/main/java/io/trino/plugin/hudi/HudiMetadata.java   | 4 ++--
 .../main/java/io/trino/plugin/iceberg/IcebergMetadata.java | 7 ++-----
 .../main/java/io/trino/plugin/memory/MemoryMetadata.java   | 3 +--
 6 files changed, 11 insertions(+), 17 deletions(-)

diff --git a/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java b/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java
index 1e21e88fa3e0..e52e2bd25683 100644
--- a/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java
+++ b/core/trino-main/src/main/java/io/trino/metadata/MetadataManager.java
@@ -2803,6 +2803,9 @@ public OptionalInt getMaxWriterTasks(Session session, String catalogName)
     public boolean allowSplittingReadIntoMultipleSubQueries(Session session, TableHandle tableHandle)
     {
         CatalogHandle catalogHandle = tableHandle.catalogHandle();
+        if (catalogHandle.getType().isInternal()) {
+            return false;
+        }
         CatalogMetadata catalogMetadata = getCatalogMetadata(session, catalogHandle);
         ConnectorSession connectorSession = session.toConnectorSession(catalogHandle);
         return catalogMetadata.getMetadata(session).allowSplittingReadIntoMultipleSubQueries(connectorSession, tableHandle.connectorHandle());
diff --git a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java
index 09c669cbf402..fca1db0fd290 100644
--- a/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java
+++ b/plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeMetadata.java
@@ -4061,8 +4061,8 @@ private Optional<SystemTable> getRawSystemTable(ConnectorSession session, Schema
     @Override
     public boolean allowSplittingReadIntoMultipleSubQueries(ConnectorSession session, ConnectorTableHandle tableHandle)
     {
-        // dont split to subqueries if tableHandle is systemTableHandle, delta lake supports only a columnar (parquet) storage format
-        return tableHandle instanceof DeltaLakeTableHandle;
+        // delta lake supports only a columnar (parquet) storage format
+        return true;
     }
 
     @Override
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java
index f9860fbf5eb2..6c547ebfddc4 100644
--- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java
+++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveMetadata.java
@@ -3990,12 +3990,7 @@ private static Optional<CatalogSchemaTableName> redirectTableToHudi(Optional<Str
     @Override
     public boolean allowSplittingReadIntoMultipleSubQueries(ConnectorSession session, ConnectorTableHandle tableHandle)
     {
-        // dont split to subqueries if tableHandle is systemTableHandle
-        if (!(tableHandle instanceof HiveTableHandle hiveTableHandle)) {
-            return false;
-        }
-
-        SchemaTableName tableName = hiveTableHandle.getSchemaTableName();
+        SchemaTableName tableName = ((HiveTableHandle) tableHandle).getSchemaTableName();
 
         Table table = metastore.getTable(tableName.getSchemaName(), tableName.getTableName())
                 .orElseThrow(() -> new TableNotFoundException(tableName));
diff --git a/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiMetadata.java b/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiMetadata.java
index e5de78673b7c..f86182de78ad 100644
--- a/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiMetadata.java
+++ b/plugin/trino-hudi/src/main/java/io/trino/plugin/hudi/HudiMetadata.java
@@ -290,8 +290,8 @@ public void validateScan(ConnectorSession session, ConnectorTableHandle handle)
     @Override
     public boolean allowSplittingReadIntoMultipleSubQueries(ConnectorSession session, ConnectorTableHandle tableHandle)
     {
-        // dont split to subqueries if tableHandle is systemTableHandle, hudi supports only a columnar (parquet) storage format
-        return tableHandle instanceof HudiTableHandle;
+        // hudi supports only a columnar (parquet) storage format
+        return true;
     }
 
     HiveMetastore getMetastore()
diff --git a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java
index c77c4abe4954..31bce54dcb11 100644
--- a/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java
+++ b/plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/IcebergMetadata.java
@@ -3314,11 +3314,8 @@ public Optional<CatalogSchemaTableName> redirectTable(ConnectorSession session,
     @Override
     public boolean allowSplittingReadIntoMultipleSubQueries(ConnectorSession session, ConnectorTableHandle connectorTableHandle)
     {
-        // dont split to subqueries if tableHandle is systemTableHandle
-        if (!(connectorTableHandle instanceof IcebergTableHandle icebergTableHandle)) {
-            return false;
-        }
-        IcebergFileFormat storageFormat = getFileFormat(icebergTableHandle.getStorageProperties());
+        IcebergTableHandle tableHandle = (IcebergTableHandle) connectorTableHandle;
+        IcebergFileFormat storageFormat = getFileFormat(tableHandle.getStorageProperties());
 
         return storageFormat == IcebergFileFormat.ORC || storageFormat == IcebergFileFormat.PARQUET;
     }
diff --git a/plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java b/plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java
index 13ef811ddb84..a9da77ac5127 100644
--- a/plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java
+++ b/plugin/trino-memory/src/main/java/io/trino/plugin/memory/MemoryMetadata.java
@@ -640,8 +640,7 @@ public Optional<SampleApplicationResult<ConnectorTableHandle>> applySample(Conne
     @Override
     public boolean allowSplittingReadIntoMultipleSubQueries(ConnectorSession session, ConnectorTableHandle tableHandle)
     {
-        // dont split to subqueries if tableHandle is systemTableHandle
-        return tableHandle instanceof MemoryTableHandle;
+        return true;
     }
 
     @Override