fix

JingsongLi · JingsongLi · commit 1a2daf6c1fea · 2025-04-08T21:10:26.000+08:00
diff --git a/paimon-common/src/main/java/org/apache/paimon/format/SimpleStatsExtractor.java b/paimon-common/src/main/java/org/apache/paimon/format/SimpleStatsExtractor.java
@@ -32,6 +32,8 @@ public interface SimpleStatsExtractor {
     Pair<SimpleColStats[], FileInfo> extractWithFileInfo(FileIO fileIO, Path path, long length)
             throws IOException;
 
+    boolean isStatsDisabled();
+
     /** File info fetched from physical file. */
     class FileInfo {
 
diff --git a/paimon-core/src/main/java/org/apache/paimon/io/KeyValueDataFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/KeyValueDataFileWriter.java
@@ -103,7 +103,7 @@ public KeyValueDataFileWriter(
                 writeRowType,
                 simpleStatsExtractor,
                 compression,
-                StatsCollectorFactories.createStatsFactories(
+                StatsCollectorFactories.createStatsFactoriesForAvro(
                         statsMode, options, writeRowType.getFieldNames()),
                 options.asyncFileWrite());
 
diff --git a/paimon-core/src/main/java/org/apache/paimon/io/KeyValueFileWriterFactory.java b/paimon-core/src/main/java/org/apache/paimon/io/KeyValueFileWriterFactory.java
@@ -41,6 +41,7 @@
 import javax.annotation.Nullable;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
@@ -291,7 +292,10 @@ private WriteFormatContext(
                                     StatsCollectorFactories.createStatsFactories(
                                             level2Stats.apply(level),
                                             options,
-                                            writeRowType.getFieldNames());
+                                            writeRowType.getFieldNames(),
+                                            thinModeEnabled
+                                                    ? keyType.getFieldNames()
+                                                    : Collections.emptyList());
                             return fileFormat
                                     .createStatsExtractor(writeRowType, statsFactories)
                                     .orElse(null);
diff --git a/paimon-core/src/main/java/org/apache/paimon/io/StatsCollectingSingleFileWriter.java b/paimon-core/src/main/java/org/apache/paimon/io/StatsCollectingSingleFileWriter.java
@@ -25,15 +25,13 @@
 import org.apache.paimon.format.SimpleStatsExtractor;
 import org.apache.paimon.fs.FileIO;
 import org.apache.paimon.fs.Path;
-import org.apache.paimon.statistics.NoneSimpleColStatsCollector;
 import org.apache.paimon.statistics.SimpleColStatsCollector;
 import org.apache.paimon.types.RowType;
 import org.apache.paimon.utils.Preconditions;
 
 import javax.annotation.Nullable;
 
 import java.io.IOException;
-import java.util.Arrays;
 import java.util.function.Function;
 import java.util.stream.IntStream;
 
@@ -62,15 +60,15 @@ public StatsCollectingSingleFileWriter(
             boolean asyncWrite) {
         super(fileIO, factory, path, converter, compression, asyncWrite);
         this.simpleStatsExtractor = simpleStatsExtractor;
-        if (this.simpleStatsExtractor == null) {
+        if (this.simpleStatsExtractor != null) {
+            this.isStatsDisabled = simpleStatsExtractor.isStatsDisabled();
+        } else {
             this.simpleStatsCollector = new SimpleStatsCollector(writeSchema, statsCollectors);
+            this.isStatsDisabled = simpleStatsCollector.isDisabled();
         }
         Preconditions.checkArgument(
                 statsCollectors.length == writeSchema.getFieldCount(),
                 "The stats collector is not aligned to write schema.");
-        this.isStatsDisabled =
-                Arrays.stream(SimpleColStatsCollector.create(statsCollectors))
-                        .allMatch(p -> p instanceof NoneSimpleColStatsCollector);
         if (isStatsDisabled) {
             this.noneStats =
                     IntStream.range(0, statsCollectors.length)
diff --git a/paimon-core/src/main/java/org/apache/paimon/utils/StatsCollectorFactories.java b/paimon-core/src/main/java/org/apache/paimon/utils/StatsCollectorFactories.java
@@ -21,7 +21,10 @@
 import org.apache.paimon.CoreOptions;
 import org.apache.paimon.options.Options;
 import org.apache.paimon.statistics.SimpleColStatsCollector;
+import org.apache.paimon.statistics.TruncateSimpleColStatsCollector;
+import org.apache.paimon.table.SpecialFields;
 
+import java.util.Collections;
 import java.util.List;
 
 import static org.apache.paimon.CoreOptions.FIELDS_PREFIX;
@@ -33,21 +36,56 @@ public class StatsCollectorFactories {
 
     public static SimpleColStatsCollector.Factory[] createStatsFactories(
             String statsMode, CoreOptions options, List<String> fields) {
-        Options cfg = options.toConfiguration();
+        return createStatsFactories(statsMode, options, fields, Collections.emptyList());
+    }
+
+    public static SimpleColStatsCollector.Factory[] createStatsFactories(
+            String statsMode, CoreOptions coreOptions, List<String> fields, List<String> keyNames) {
+        Options options = coreOptions.toConfiguration();
         SimpleColStatsCollector.Factory[] modes =
                 new SimpleColStatsCollector.Factory[fields.size()];
         for (int i = 0; i < fields.size(); i++) {
             String field = fields.get(i);
-            String fieldMode =
-                    cfg.get(
-                            key(String.format("%s.%s.%s", FIELDS_PREFIX, field, STATS_MODE_SUFFIX))
-                                    .stringType()
-                                    .noDefaultValue());
+            String fieldMode = fieldMode(options, field);
+            if (fieldMode != null) {
+                modes[i] = SimpleColStatsCollector.from(fieldMode);
+            } else if (SpecialFields.isSystemField(field)
+                    ||
+                    // If we config DATA_FILE_THIN_MODE to true, we need to maintain the
+                    // stats for key fields.
+                    keyNames.contains(SpecialFields.KEY_FIELD_PREFIX + field)) {
+                modes[i] = () -> new TruncateSimpleColStatsCollector(128);
+            } else {
+                modes[i] = SimpleColStatsCollector.from(statsMode);
+            }
+        }
+        return modes;
+    }
+
+    /**
+     * If all are None, return all None to Avro Writer, which can greatly accelerate the writing
+     * speed.
+     */
+    public static SimpleColStatsCollector.Factory[] createStatsFactoriesForAvro(
+            String statsMode, CoreOptions coreOptions, List<String> fields) {
+        Options options = coreOptions.toConfiguration();
+        SimpleColStatsCollector.Factory[] modes =
+                new SimpleColStatsCollector.Factory[fields.size()];
+        for (int i = 0; i < fields.size(); i++) {
+            String field = fields.get(i);
+            String fieldMode = fieldMode(options, field);
             modes[i] =
                     fieldMode != null
                             ? SimpleColStatsCollector.from(fieldMode)
                             : SimpleColStatsCollector.from(statsMode);
         }
         return modes;
     }
+
+    private static String fieldMode(Options options, String field) {
+        return options.get(
+                key(String.format("%s.%s.%s", FIELDS_PREFIX, field, STATS_MODE_SUFFIX))
+                        .stringType()
+                        .noDefaultValue());
+    }
 }
diff --git a/paimon-core/src/test/java/org/apache/paimon/stats/TestSimpleStatsExtractor.java b/paimon-core/src/test/java/org/apache/paimon/stats/TestSimpleStatsExtractor.java
@@ -26,6 +26,7 @@
 import org.apache.paimon.format.SimpleStatsExtractor;
 import org.apache.paimon.fs.FileIO;
 import org.apache.paimon.fs.Path;
+import org.apache.paimon.statistics.NoneSimpleColStatsCollector;
 import org.apache.paimon.statistics.SimpleColStatsCollector;
 import org.apache.paimon.types.RowType;
 import org.apache.paimon.utils.ObjectSerializer;
@@ -34,6 +35,7 @@
 
 import java.io.IOException;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 
 import static org.apache.paimon.utils.FileUtils.createFormatReader;
@@ -77,6 +79,12 @@ public Pair<SimpleColStats[], FileInfo> extractWithFileInfo(
         return Pair.of(statsCollector.extract(), new FileInfo(records.size()));
     }
 
+    @Override
+    public boolean isStatsDisabled() {
+        return Arrays.stream(SimpleColStatsCollector.create(stats))
+                .allMatch(p -> p instanceof NoneSimpleColStatsCollector);
+    }
+
     private static <T> List<T> readListFromFile(
             FileIO fileIO,
             Path path,
diff --git a/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroSimpleStatsExtractor.java b/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroSimpleStatsExtractor.java
@@ -72,6 +72,11 @@ public Pair<SimpleColStats[], FileInfo> extractWithFileInfo(
                 new FileInfo(rowCount));
     }
 
+    @Override
+    public boolean isStatsDisabled() {
+        return true;
+    }
+
     private long getRowCount(InputStream inStream) throws java.io.IOException {
         // an Avro file's layout looks like this:
         //   header|block|block|...
diff --git a/paimon-format/src/main/java/org/apache/paimon/format/orc/filter/OrcSimpleStatsExtractor.java b/paimon-format/src/main/java/org/apache/paimon/format/orc/filter/OrcSimpleStatsExtractor.java
@@ -26,6 +26,7 @@
 import org.apache.paimon.format.orc.OrcReaderFactory;
 import org.apache.paimon.fs.FileIO;
 import org.apache.paimon.fs.Path;
+import org.apache.paimon.statistics.NoneSimpleColStatsCollector;
 import org.apache.paimon.statistics.SimpleColStatsCollector;
 import org.apache.paimon.types.DataField;
 import org.apache.paimon.types.DecimalType;
@@ -50,6 +51,7 @@
 
 import java.io.IOException;
 import java.sql.Date;
+import java.util.Arrays;
 import java.util.List;
 import java.util.stream.IntStream;
 
@@ -116,6 +118,12 @@ public Pair<SimpleColStats[], FileInfo> extractWithFileInfo(
         }
     }
 
+    @Override
+    public boolean isStatsDisabled() {
+        return Arrays.stream(SimpleColStatsCollector.create(statsCollectors))
+                .allMatch(p -> p instanceof NoneSimpleColStatsCollector);
+    }
+
     private SimpleColStats toFieldStats(
             DataField field,
             ColumnStatistics stats,
diff --git a/paimon-format/src/main/java/org/apache/paimon/format/parquet/ParquetSimpleStatsExtractor.java b/paimon-format/src/main/java/org/apache/paimon/format/parquet/ParquetSimpleStatsExtractor.java
@@ -25,6 +25,7 @@
 import org.apache.paimon.format.SimpleStatsExtractor;
 import org.apache.paimon.fs.FileIO;
 import org.apache.paimon.fs.Path;
+import org.apache.paimon.statistics.NoneSimpleColStatsCollector;
 import org.apache.paimon.statistics.SimpleColStatsCollector;
 import org.apache.paimon.types.DataField;
 import org.apache.paimon.types.DecimalType;
@@ -46,6 +47,7 @@
 import java.io.IOException;
 import java.math.BigDecimal;
 import java.math.BigInteger;
+import java.util.Arrays;
 import java.util.Map;
 import java.util.stream.IntStream;
 
@@ -91,6 +93,12 @@ public Pair<SimpleColStats[], FileInfo> extractWithFileInfo(
                 statsPair.getRight());
     }
 
+    @Override
+    public boolean isStatsDisabled() {
+        return Arrays.stream(SimpleColStatsCollector.create(statsCollectors))
+                .allMatch(p -> p instanceof NoneSimpleColStatsCollector);
+    }
+
     private SimpleColStats toFieldStats(
             DataField field, Statistics<?> stats, SimpleColStatsCollector collector) {
         if (stats == null) {
