Fix for SystemTables in allowSplittingReadIntoMultipleSubQueries

trinodb · Sep 23, 2024 · 82b55e5 · 82b55e5
1 parent 57ec0d2
commit 82b55e5
Show file tree

Hide file tree

Showing 6 changed files with 37 additions and 8 deletions.
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
@@ -4058,8 +4058,8 @@ private Optional<SystemTable> getRawSystemTable(ConnectorSession session, Schema
     @Override
     public boolean allowSplittingReadIntoMultipleSubQueries(ConnectorSession session, ConnectorTableHandle tableHandle)
     {
-        // delta lake supports only a columnar (parquet) storage format
-        return true;
+        // dont split to subqueries if tableHandle is systemTableHandle, delta lake supports only a columnar (parquet) storage format
+        return tableHandle instanceof DeltaLakeTableHandle;
     }
 
     @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
@@ -3984,7 +3984,12 @@ private static Optional<CatalogSchemaTableName> redirectTableToHudi(Optional<Str
     @Override
     public boolean allowSplittingReadIntoMultipleSubQueries(ConnectorSession session, ConnectorTableHandle tableHandle)
     {
-        SchemaTableName tableName = ((HiveTableHandle) tableHandle).getSchemaTableName();
+        // dont split to subqueries if tableHandle is systemTableHandle
+        if (!(tableHandle instanceof HiveTableHandle hiveTableHandle)) {
+            return false;
+        }
+
+        SchemaTableName tableName = hiveTableHandle.getSchemaTableName();
 
         Table table = metastore.getTable(tableName.getSchemaName(), tableName.getTableName())
                 .orElseThrow(() -> new TableNotFoundException(tableName));

diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveDistributedAggregations.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/TestHiveDistributedAggregations.java
@@ -15,6 +15,10 @@
 
 import io.trino.testing.AbstractTestAggregations;
 import io.trino.testing.QueryRunner;
+import org.intellij.lang.annotations.Language;
+import org.junit.jupiter.api.Test;
+
+import static io.trino.testing.TestingNames.randomNameSuffix;
 
 public class TestHiveDistributedAggregations
         extends AbstractTestAggregations
@@ -27,4 +31,20 @@ protected QueryRunner createQueryRunner()
                 .setInitialTables(REQUIRED_TPCH_TABLES)
                 .build();
     }
+
+    @Test
+    public void testDistinctAggregationWithSystemTable()
+    {
+        String tableName = "test_dist_aggr_" + randomNameSuffix();
+        @Language("SQL") String createTable = """
+                CREATE TABLE %s
+                WITH (
+                partitioned_by = ARRAY[ 'regionkey', 'nationkey' ]
+                ) AS (SELECT name, comment, regionkey, nationkey FROM nation)
+                """.formatted(tableName);
+
+        assertUpdate(getSession(), createTable, 25);
+
+        assertQuerySucceeds("SELECT count(distinct regionkey), count(distinct nationkey) FROM \"%s$partitions\"".formatted(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
@@ -290,8 +290,8 @@ public void validateScan(ConnectorSession session, ConnectorTableHandle handle)
     @Override
     public boolean allowSplittingReadIntoMultipleSubQueries(ConnectorSession session, ConnectorTableHandle tableHandle)
     {
-        // hudi supports only a columnar (parquet) storage format
-        return true;
+        // dont split to subqueries if tableHandle is systemTableHandle, hudi supports only a columnar (parquet) storage format
+        return tableHandle instanceof HudiTableHandle;
     }
 
     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
@@ -3312,8 +3312,11 @@ public Optional<CatalogSchemaTableName> redirectTable(ConnectorSession session,
     @Override
     public boolean allowSplittingReadIntoMultipleSubQueries(ConnectorSession session, ConnectorTableHandle connectorTableHandle)
     {
-        IcebergTableHandle tableHandle = (IcebergTableHandle) connectorTableHandle;
-        IcebergFileFormat storageFormat = getFileFormat(tableHandle.getStorageProperties());
+        // dont split to subqueries if tableHandle is systemTableHandle
+        if (!(connectorTableHandle instanceof IcebergTableHandle icebergTableHandle)) {
+            return false;
+        }
+        IcebergFileFormat storageFormat = getFileFormat(icebergTableHandle.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
@@ -640,7 +640,8 @@ public Optional<SampleApplicationResult<ConnectorTableHandle>> applySample(Conne
     @Override
     public boolean allowSplittingReadIntoMultipleSubQueries(ConnectorSession session, ConnectorTableHandle tableHandle)
     {
-        return true;
+        // dont split to subqueries if tableHandle is systemTableHandle
+        return tableHandle instanceof MemoryTableHandle;
     }
 
     @Override