diff --git a/integration_tests/src/main/python/delta_lake_liquid_clustering_test.py b/integration_tests/src/main/python/delta_lake_liquid_clustering_test.py index 8d71a64aeef..d1c44fd477e 100644 --- a/integration_tests/src/main/python/delta_lake_liquid_clustering_test.py +++ b/integration_tests/src/main/python/delta_lake_liquid_clustering_test.py @@ -53,6 +53,18 @@ def assert_liquid_clustering_delta_logs_equivalent(data_path): with_cpu_session(lambda spark: assert_gpu_and_cpu_delta_logs_equivalent(spark, data_path)) +def delta_path_identifier(path): + return f"delta.`{path}`" + + +def optimize_liquid_clustered_table(spark, table_identifier): + if is_databricks173_or_later(): + with_cpu_session(lambda cpu_spark: cpu_spark.sql( + f"OPTIMIZE {table_identifier}").collect()) + else: + spark.sql(f"OPTIMIZE {table_identifier}").collect() + + @allow_non_gpu(*delta_meta_allow) @allow_non_gpu_conditional(is_databricks173_or_later(), f"{delta_write_fallback_allow},AtomicCreateTableAsSelectExec," @@ -77,6 +89,7 @@ def write_func(spark, path): CLUSTER BY (a) AS SELECT * FROM {view_name} """) + optimize_liquid_clustered_table(spark, delta_path_identifier(path)) data_path = spark_tmp_path + "/DELTA_LIQUID_CLUSTER" @@ -122,6 +135,7 @@ def setup_clustered_table_sql(spark, path, table_name, view_name): INSERT INTO {table_name} SELECT * FROM {view_name} """) + optimize_liquid_clustered_table(spark, table_name) @allow_non_gpu(*delta_meta_allow, "CreateTableExec") @@ -147,6 +161,7 @@ def write_func(spark, path): CLUSTER BY (a) AS SELECT * FROM {view_name} """) + optimize_liquid_clustered_table(spark, table_name) data_path = spark_tmp_path + "/DELTA_LIQUID_CLUSTER" @@ -176,6 +191,7 @@ def write_func(spark, path): INSERT INTO {table_name} SELECT * FROM {view_name} """) + optimize_liquid_clustered_table(spark, table_name) data_path = spark_tmp_path + "/DELTA_LIQUID_CLUSTER" @@ -206,6 +222,7 @@ def write_func(spark, path): INSERT OVERWRITE TABLE {table_name} SELECT * FROM {view_name} """) + optimize_liquid_clustered_table(spark, table_name) data_path = spark_tmp_path + "/DELTA_LIQUID_CLUSTER" @@ -240,6 +257,7 @@ def write_func(spark, path): INSERT OVERWRITE TABLE {table_name} SELECT * FROM {view_name} """) + optimize_liquid_clustered_table(spark, table_name) data_path = spark_tmp_path + "/DELTA_LIQUID_CLUSTER" @@ -279,6 +297,7 @@ def write_func(spark, path): REPLACE WHERE b = 'x' SELECT * FROM {view_name} """) + optimize_liquid_clustered_table(spark, table_name) data_path = spark_tmp_path + "/DELTA_LIQUID_CLUSTER" @@ -311,6 +330,7 @@ def do_test_delta_dml_sql_liquid_clustering(spark_tmp_path, def modify_table(spark, path): table_name = cpu_table_name if path == cpu_data_path else gpu_table_name spark.sql(sql_func(table_name)) + optimize_liquid_clustered_table(spark, table_name) assert_gpu_and_cpu_writes_are_equal_collect( modify_table, @@ -402,6 +422,7 @@ def merge_table(spark, path): INSERT * """ spark.sql(sql).collect() + optimize_liquid_clustered_table(spark, target_name) read_table = lambda spark, path: spark.read.format("delta").load(path) if is_databricks173_or_later(): @@ -465,6 +486,7 @@ def write_func(spark, path): create_clustered_delta_table_df(table_name, path) # Create a temp view to select from for the CTAS operation write_to_delta_table_df(spark, path, "append") + optimize_liquid_clustered_table(spark, table_name) data_path = spark_tmp_path + "/DELTA_LIQUID_CLUSTER" @@ -496,6 +518,7 @@ def write_func(spark, path): create_clustered_delta_table_df(table_name, path) write_to_delta_table_df(spark, path, "overwrite", opts = {'partitionOverwriteMode': overwrite_mode}) + optimize_liquid_clustered_table(spark, table_name) data_path = spark_tmp_path + "/DELTA_LIQUID_CLUSTER" @@ -527,6 +550,7 @@ def write_func(spark, path): create_clustered_delta_table_df(table_name, path) write_to_delta_table_df(spark, path, "overwrite", opts = {'replaceWhere': " b='x' "}) + optimize_liquid_clustered_table(spark, table_name) data_path = spark_tmp_path + "/DELTA_LIQUID_CLUSTER"