[native] Add followup e2e test for table format

unix280 · Nov 5, 2024 · b036356 · b036356
1 parent 0e489de
commit b036356
Showing 1 changed file with 54 additions and 0 deletions.
diff --git a/...o-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestWriter.java b/...o-native-execution/src/test/java/com/facebook/presto/nativeworker/AbstractTestWriter.java
@@ -74,6 +74,40 @@ public void testCreateTableWithUnsupportedFormats()
         }
     }
 
+    @Test
+    public void testTableFormats()
+    {
+        // With different storage_format than table format, verify the data can be correctly read back
+        Session[] sessions = buildSessionsForTableWriteFormat();
+        for (String tableFormat : TABLE_FORMATS) {
+            for (Session session : sessions) {
+                // Unpartitioned table
+                String tmpTableName = generateRandomTableName();
+                try {
+                    getQueryRunner().execute(session, String.format("CREATE TABLE %s WITH (format = '" + tableFormat + "') AS SELECT * FROM nation", tmpTableName));
+                    assertQuery(String.format("SELECT * FROM %s", tmpTableName), "SELECT * FROM nation");
+                    // TODO add support for presto to query each partition's format then verify written format is correct
+                }
+                finally {
+                    dropTableIfExists(tmpTableName);
+                }
+
+                // Partitioned table
+                String partitionedOrdersTableName = generateRandomTableName();
+                try {
+                    getQueryRunner().execute(session, String.format(
+                            "CREATE TABLE %s WITH (format = '" + tableFormat + "', " +
+                                    "partitioned_by = ARRAY[ 'orderstatus' ]) " +
+                                    "AS SELECT custkey, comment, orderstatus FROM orders", partitionedOrdersTableName));
+                    assertQuery(String.format("SELECT * FROM %s", partitionedOrdersTableName), "SELECT custkey, comment, orderstatus FROM orders");
+                }
+                finally {
+                    dropTableIfExists(partitionedOrdersTableName);
+                }
+            }
+        }
+    }
+
     @Test
     public void testCreateUnpartitionedTableAsSelect()
     {
@@ -411,4 +445,24 @@ private Session buildSessionForTableWrite()
                 .setCatalogSessionProperty("hive", "orc_compression_codec", "ZSTD")
                 .build();
     }
+
+    private Session[] buildSessionsForTableWriteFormat()
+    {
+        Session.SessionBuilder sessionBuilder = Session.builder(getSession())
+                .setSystemProperty("scale_writers", "true")
+                .setSystemProperty("table_writer_merge_operator_enabled", "true")
+                .setSystemProperty("task_writer_count", "1")
+                .setSystemProperty("task_partitioned_writer_count", "2")
+                .setCatalogSessionProperty("hive", "collect_column_statistics_on_write", "true")
+                .setCatalogSessionProperty("hive", "orc_compression_codec", "ZSTD");
+        Session s1 = sessionBuilder
+                .setCatalogSessionProperty("hive", "hive_storage_format", "PARQUET")
+                .setCatalogSessionProperty("hive", "respect_table_format", "true").build();
+        Session s2 = sessionBuilder
+                .setCatalogSessionProperty("hive", "hive_storage_format", "PARQUET")
+                .setCatalogSessionProperty("hive", "respect_table_format", "false").build();
+        Session[] sessions = {s1, s2};
+
+        return sessions;
+    }
 }