Fix updating column statistics when inserting with overwrite mode

anton-kutuzov · raunaqmorarka · commit 571fd670d290 · 2025-11-07T10:08:39.000+05:30
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
@@ -181,6 +181,7 @@
 import static io.trino.metastore.PrincipalPrivileges.NO_PRIVILEGES;
 import static io.trino.metastore.PrincipalPrivileges.fromHivePrivilegeInfos;
 import static io.trino.metastore.StatisticsUpdateMode.MERGE_INCREMENTAL;
+import static io.trino.metastore.StatisticsUpdateMode.OVERWRITE_ALL;
 import static io.trino.metastore.StorageFormat.VIEW_STORAGE_FORMAT;
 import static io.trino.metastore.type.Category.PRIMITIVE;
 import static io.trino.parquet.writer.ParquetWriter.SUPPORTED_BLOOM_FILTER_TYPES;
@@ -2360,7 +2361,8 @@ else if (partitionUpdate.getUpdateMode() == APPEND) {
                                 partitionValues,
                                 partitionUpdate.getWritePath(),
                                 partitionUpdate.getFileNames(),
-                                partitionStatistics));
+                                partitionStatistics,
+                                MERGE_INCREMENTAL));
             }
             else if (partitionUpdate.getUpdateMode() == NEW || partitionUpdate.getUpdateMode() == OVERWRITE) {
                 // insert into new partition or overwrite existing partition
@@ -2385,6 +2387,13 @@ else if (partitionUpdate.getUpdateMode() == NEW || partitionUpdate.getUpdateMode
                             TrinoFileSystem fileSystem = fileSystemFactory.create(session);
                             cleanExtraOutputFiles(fileSystem, session.getQueryId(), partitionUpdate.getTargetPath(), ImmutableSet.copyOf(partitionUpdate.getFileNames()));
                         }
+                        partitionUpdateInfosBuilder.add(
+                                new PartitionUpdateInfo(
+                                        partitionValues,
+                                        partitionUpdate.getWritePath(),
+                                        partitionUpdate.getFileNames(),
+                                        partitionStatistics,
+                                        OVERWRITE_ALL));
                     }
                     else {
                         metastore.dropPartition(session, handle.getSchemaName(), handle.getTableName(), partition.getValues(), true);
@@ -2690,7 +2699,8 @@ private void finishOptimize(ConnectorSession session, ConnectorTableExecuteHandl
                                 partitionValues,
                                 partitionUpdate.getWritePath(),
                                 partitionUpdate.getFileNames(),
-                                PartitionStatistics.empty()));
+                                PartitionStatistics.empty(),
+                                MERGE_INCREMENTAL));
             }
         }
 
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/SemiTransactionalHiveMetastore.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/metastore/SemiTransactionalHiveMetastore.java
@@ -171,6 +171,8 @@ public class SemiTransactionalHiveMetastore
             AcidOperation.INSERT, ActionType.INSERT_EXISTING,
             AcidOperation.MERGE, ActionType.MERGE);
 
+    private static final boolean SHOULD_MERGE_STATISTICS = true;
+
     private final HiveMetastore delegate;
     private final TypeManager typeManager;
     private final boolean partitionProjectionEnabled;
@@ -911,7 +913,7 @@ public synchronized void addPartition(
         if (oldPartitionAction == null) {
             partitionActionsOfTable.put(
                     partition.getValues(),
-                    new Action<>(ActionType.ADD, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, cleanExtraOutputFilesOnCommit), session.getIdentity(), session.getQueryId()));
+                    new Action<>(ActionType.ADD, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, SHOULD_MERGE_STATISTICS, cleanExtraOutputFilesOnCommit), session.getIdentity(), session.getQueryId()));
             return;
         }
         switch (oldPartitionAction.type()) {
@@ -921,7 +923,7 @@ public synchronized void addPartition(
                 }
                 partitionActionsOfTable.put(
                         partition.getValues(),
-                        new Action<>(ActionType.ALTER, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, cleanExtraOutputFilesOnCommit), session.getIdentity(), session.getQueryId()));
+                        new Action<>(ActionType.ALTER, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, SHOULD_MERGE_STATISTICS, cleanExtraOutputFilesOnCommit), session.getIdentity(), session.getQueryId()));
             }
             case ADD, ALTER, INSERT_EXISTING, MERGE ->
                     throw new TrinoException(ALREADY_EXISTS, format("Partition already exists for table '%s.%s': %s", databaseName, tableName, partition.getValues()));
@@ -1003,8 +1005,9 @@ public synchronized void finishInsertIntoExistingPartitions(
                                         partition,
                                         partitionInfo.currentLocation(),
                                         Optional.of(partitionInfo.fileNames()),
-                                        MERGE_INCREMENTAL.updatePartitionStatistics(currentStatistics, partitionInfo.statisticsUpdate()),
+                                        partitionInfo.statisticsUpdateMode().updatePartitionStatistics(currentStatistics, partitionInfo.statisticsUpdate()),
                                         partitionInfo.statisticsUpdate(),
+                                        partitionInfo.statisticsUpdateMode() != OVERWRITE_ALL,
                                         cleanExtraOutputFilesOnCommit),
                                 session.getIdentity(),
                                 session.getQueryId()));
@@ -1989,7 +1992,7 @@ private void prepareInsertExistingPartition(ConnectorIdentity identity, String q
                     partition.getSchemaTableName(),
                     Optional.of(getPartitionName(partition.getDatabaseName(), partition.getTableName(), partition.getValues())),
                     partitionAndMore.statisticsUpdate(),
-                    true));
+                    partitionAndMore.mergeStatistic()));
         }
 
         private void executeCleanupTasksForAbort(Collection<DeclaredIntentionToWrite> declaredIntentionsToWrite)
@@ -2840,6 +2843,7 @@ private record PartitionAndMore(
             Optional<List<String>> fileNames,
             PartitionStatistics statistics,
             PartitionStatistics statisticsUpdate,
+            boolean mergeStatistic,
             boolean cleanExtraOutputFilesOnCommit)
     {
         private PartitionAndMore
@@ -3362,14 +3366,15 @@ public static void cleanExtraOutputFiles(TrinoFileSystem fileSystem, String quer
         }
     }
 
-    public record PartitionUpdateInfo(List<String> partitionValues, Location currentLocation, List<String> fileNames, PartitionStatistics statisticsUpdate)
+    public record PartitionUpdateInfo(List<String> partitionValues, Location currentLocation, List<String> fileNames, PartitionStatistics statisticsUpdate, StatisticsUpdateMode statisticsUpdateMode)
     {
         public PartitionUpdateInfo
         {
             requireNonNull(partitionValues, "partitionValues is null");
             requireNonNull(currentLocation, "currentLocation is null");
             requireNonNull(fileNames, "fileNames is null");
             requireNonNull(statisticsUpdate, "statisticsUpdate is null");
+            requireNonNull(statisticsUpdateMode, "statisticsUpdateMode is null");
         }
     }
 }
diff --git a/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseTestHiveOnDataLake.java b/plugin/trino-hive/src/test/java/io/trino/plugin/hive/BaseTestHiveOnDataLake.java
@@ -50,6 +50,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.OptionalLong;
 import java.util.TimeZone;
 import java.util.stream.Collectors;
 
@@ -281,6 +282,80 @@ public void testSyncPartitionOnBucketRoot()
         assertUpdate("DROP TABLE " + fullyQualifiedTestTableName);
     }
 
+    @Test
+    public void testUpdateStatisticInsertOverwritePartitionedTable()
+    {
+        String partitionValue = "0";
+        Session session = Session.builder(getQueryRunner().getDefaultSession())
+                .setCatalogSessionProperty("hive", "insert_existing_partitions_behavior", "OVERWRITE")
+                .setCatalogSessionProperty("hive", "collect_column_statistics_on_write", "true")
+                .build();
+        String tableName = "test_statistic" + randomNameSuffix();
+        String testTable = getFullyQualifiedTestTableName(tableName);
+        computeActual(getCreateTableStatement(
+                testTable,
+                "partitioned_by=ARRAY['regionkey']"));
+        copyTpchNationToTable(session, testTable);
+        Table hiveTable = metastoreClient.getTable(HIVE_TEST_SCHEMA, tableName).orElseThrow();
+        Partition partition = metastoreClient.getPartition(hiveTable, List.of(partitionValue)).orElseThrow();
+        Map<String, Map<String, HiveColumnStatistics>> partitionStatistics = metastoreClient.getPartitionColumnStatistics(
+                HIVE_TEST_SCHEMA,
+                tableName,
+                ImmutableSet.of("regionkey=0"),
+                partition.getColumns().stream().map(Column::getName).collect(toSet()));
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().isPresent()).isTrue();
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().get().getMin()).isEqualTo(OptionalLong.of(0));
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().get().getMax()).isEqualTo(OptionalLong.of(16));
+
+        assertUpdate(session, "INSERT INTO " + testTable + "(name, comment, nationkey, regionkey) values ('name1', 'comment1', 20, 0)", 1);
+
+        partitionStatistics = metastoreClient.getPartitionColumnStatistics(
+                HIVE_TEST_SCHEMA,
+                tableName,
+                ImmutableSet.of("regionkey=0"),
+                partition.getColumns().stream().map(Column::getName).collect(toSet()));
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().isPresent()).isTrue();
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().get().getMin()).isEqualTo(OptionalLong.of(20));
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().get().getMax()).isEqualTo(OptionalLong.of(20));
+    }
+
+    @Test
+    public void testUpdateStatisticInsertAppendPartitionedTable()
+    {
+        String partitionValue = "0";
+        Session session = Session.builder(getQueryRunner().getDefaultSession())
+                .setCatalogSessionProperty("hive", "insert_existing_partitions_behavior", "APPEND")
+                .setCatalogSessionProperty("hive", "collect_column_statistics_on_write", "true")
+                .build();
+        String tableName = "test_statistic" + randomNameSuffix();
+        String testTable = getFullyQualifiedTestTableName(tableName);
+        computeActual(session, getCreateTableStatement(
+                testTable,
+                "partitioned_by=ARRAY['regionkey']"));
+        copyTpchNationToTable(session, testTable);
+        Table hiveTable = metastoreClient.getTable(HIVE_TEST_SCHEMA, tableName).orElseThrow();
+        Partition partition = metastoreClient.getPartition(hiveTable, List.of(partitionValue)).orElseThrow();
+        Map<String, Map<String, HiveColumnStatistics>> partitionStatistics = metastoreClient.getPartitionColumnStatistics(
+                HIVE_TEST_SCHEMA,
+                tableName,
+                ImmutableSet.of("regionkey=0"),
+                partition.getColumns().stream().map(Column::getName).collect(toSet()));
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().isPresent()).isTrue();
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().get().getMin()).isEqualTo(OptionalLong.of(0));
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().get().getMax()).isEqualTo(OptionalLong.of(16));
+
+        computeActual(session, "INSERT INTO " + testTable + "(name, comment, nationkey, regionkey) values ('name1', 'comment1', 20, 0)");
+
+        partitionStatistics = metastoreClient.getPartitionColumnStatistics(
+                HIVE_TEST_SCHEMA,
+                tableName,
+                ImmutableSet.of("regionkey=0"),
+                partition.getColumns().stream().map(Column::getName).collect(toSet()));
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().isPresent()).isTrue();
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().get().getMin()).isEqualTo(OptionalLong.of(0));
+        assertThat(partitionStatistics.get("regionkey=0").get("nationkey").getIntegerStatistics().get().getMax()).isEqualTo(OptionalLong.of(20));
+    }
+
     @Test
     public void testSyncPartitionCaseSensitivePathVariation()
     {
@@ -2470,9 +2545,13 @@ protected String getCreateTableStatement(String tableName, List<String> properti
                 tableName);
     }
 
-    protected void copyTpchNationToTable(String testTable)
+    protected void copyTpchNationToTable(String testTable) {
+        copyTpchNationToTable(getSession(), testTable);
+    }
+
+    protected void copyTpchNationToTable(Session session, String testTable)
     {
-        computeActual(format("INSERT INTO " + testTable + " SELECT name, comment, nationkey, regionkey FROM tpch.tiny.nation"));
+        computeActual(session, format("INSERT INTO " + testTable + " SELECT name, comment, nationkey, regionkey FROM tpch.tiny.nation"));
     }
 
     private void testWriteWithFileSize(String testTable, int scaleFactorInThousands, long fileSizeRangeStart, long fileSizeRangeEnd)

Original file line number	Diff line number	Diff line change
`@@ -171,6 +171,8 @@ public class SemiTransactionalHiveMetastore`
`171`	`171`	`AcidOperation.INSERT, ActionType.INSERT_EXISTING,`
`172`	`172`	`AcidOperation.MERGE, ActionType.MERGE);`
`173`	`173`
	`174`	`+ private static final boolean SHOULD_MERGE_STATISTICS = true;`
	`175`	`+`
`174`	`176`	`private final HiveMetastore delegate;`
`175`	`177`	`private final TypeManager typeManager;`
`176`	`178`	`private final boolean partitionProjectionEnabled;`
`@@ -911,7 +913,7 @@ public synchronized void addPartition(`
`911`	`913`	`if (oldPartitionAction == null) {`
`912`	`914`	`partitionActionsOfTable.put(`
`913`	`915`	`partition.getValues(),`
`914`		`- new Action<>(ActionType.ADD, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, cleanExtraOutputFilesOnCommit), session.getIdentity(), session.getQueryId()));`
	`916`	`+ new Action<>(ActionType.ADD, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, SHOULD_MERGE_STATISTICS, cleanExtraOutputFilesOnCommit), session.getIdentity(), session.getQueryId()));`
`915`	`917`	`return;`
`916`	`918`	`}`
`917`	`919`	`switch (oldPartitionAction.type()) {`
`@@ -921,7 +923,7 @@ public synchronized void addPartition(`
`921`	`923`	`}`
`922`	`924`	`partitionActionsOfTable.put(`
`923`	`925`	`partition.getValues(),`
`924`		`- new Action<>(ActionType.ALTER, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, cleanExtraOutputFilesOnCommit), session.getIdentity(), session.getQueryId()));`
	`926`	`+ new Action<>(ActionType.ALTER, new PartitionAndMore(partition, currentLocation, files, statistics, statistics, SHOULD_MERGE_STATISTICS, cleanExtraOutputFilesOnCommit), session.getIdentity(), session.getQueryId()));`
`925`	`927`	`}`
`926`	`928`	`case ADD, ALTER, INSERT_EXISTING, MERGE ->`
`927`	`929`	`throw new TrinoException(ALREADY_EXISTS, format("Partition already exists for table '%s.%s': %s", databaseName, tableName, partition.getValues()));`
`@@ -1003,8 +1005,9 @@ public synchronized void finishInsertIntoExistingPartitions(`
`1003`	`1005`	`partition,`
`1004`	`1006`	`partitionInfo.currentLocation(),`
`1005`	`1007`	`Optional.of(partitionInfo.fileNames()),`
`1006`		`- MERGE_INCREMENTAL.updatePartitionStatistics(currentStatistics, partitionInfo.statisticsUpdate()),`
	`1008`	`+ partitionInfo.statisticsUpdateMode().updatePartitionStatistics(currentStatistics, partitionInfo.statisticsUpdate()),`
`1007`	`1009`	`partitionInfo.statisticsUpdate(),`
	`1010`	`+ partitionInfo.statisticsUpdateMode() != OVERWRITE_ALL,`
`1008`	`1011`	`cleanExtraOutputFilesOnCommit),`
`1009`	`1012`	`session.getIdentity(),`
`1010`	`1013`	`session.getQueryId()));`
`@@ -1989,7 +1992,7 @@ private void prepareInsertExistingPartition(ConnectorIdentity identity, String q`
`1989`	`1992`	`partition.getSchemaTableName(),`
`1990`	`1993`	`Optional.of(getPartitionName(partition.getDatabaseName(), partition.getTableName(), partition.getValues())),`
`1991`	`1994`	`partitionAndMore.statisticsUpdate(),`
`1992`		`- true));`
	`1995`	`+ partitionAndMore.mergeStatistic()));`
`1993`	`1996`	`}`
`1994`	`1997`
`1995`	`1998`	`private void executeCleanupTasksForAbort(Collection<DeclaredIntentionToWrite> declaredIntentionsToWrite)`
`@@ -2840,6 +2843,7 @@ private record PartitionAndMore(`
`2840`	`2843`	`Optional<List<String>> fileNames,`
`2841`	`2844`	`PartitionStatistics statistics,`
`2842`	`2845`	`PartitionStatistics statisticsUpdate,`
	`2846`	`+ boolean mergeStatistic,`
`2843`	`2847`	`boolean cleanExtraOutputFilesOnCommit)`
`2844`	`2848`	`{`
`2845`	`2849`	`private PartitionAndMore`
`@@ -3362,14 +3366,15 @@ public static void cleanExtraOutputFiles(TrinoFileSystem fileSystem, String quer`
`3362`	`3366`	`}`
`3363`	`3367`	`}`
`3364`	`3368`
`3365`		`- public record PartitionUpdateInfo(List<String> partitionValues, Location currentLocation, List<String> fileNames, PartitionStatistics statisticsUpdate)`
	`3369`	`+ public record PartitionUpdateInfo(List<String> partitionValues, Location currentLocation, List<String> fileNames, PartitionStatistics statisticsUpdate, StatisticsUpdateMode statisticsUpdateMode)`
`3366`	`3370`	`{`
`3367`	`3371`	`public PartitionUpdateInfo`
`3368`	`3372`	`{`
`3369`	`3373`	`requireNonNull(partitionValues, "partitionValues is null");`
`3370`	`3374`	`requireNonNull(currentLocation, "currentLocation is null");`
`3371`	`3375`	`requireNonNull(fileNames, "fileNames is null");`
`3372`	`3376`	`requireNonNull(statisticsUpdate, "statisticsUpdate is null");`
	`3377`	`+ requireNonNull(statisticsUpdateMode, "statisticsUpdateMode is null");`
`3373`	`3378`	`}`
`3374`	`3379`	`}`
`3375`	`3380`	`}`