paimon: expanded tests for stats extractor

Author: mao-liu

xtable-core/src/main/java/org/apache/xtable/paimon/PaimonDataFileExtractor.java

@@ -27,6 +27,7 @@
 import org.apache.paimon.Snapshot;
 import org.apache.paimon.data.BinaryArray;
 import org.apache.paimon.data.BinaryRow;
+import org.apache.paimon.data.Timestamp;
 import org.apache.paimon.io.DataFileMeta;
 import org.apache.paimon.manifest.ManifestEntry;
 import org.apache.paimon.stats.SimpleStats;
@@ -62,13 +63,11 @@ public List<InternalDataFile> toInternalDataFiles(
     while (manifestEntryIterator.hasNext()) {
       result.add(toInternalDataFile(table, manifestEntryIterator.next(), internalSchema));
     }
-    log.info("PaimonPartitionExtractor: Returning " + result.size() + " data files for " + table.name());
     return result;
   }
 
   private InternalDataFile toInternalDataFile(
       FileStoreTable table, ManifestEntry entry, InternalSchema internalSchema) {
-//    log.info("Adding manifest entry {}", entry.fileName());
     return InternalDataFile.builder()
         .physicalPath(toFullPhysicalPath(table, entry))
         .fileSizeBytes(entry.file().fileSize())
@@ -102,9 +101,7 @@ private List<ColumnStat> toColumnStats(DataFileMeta file, InternalSchema interna
     // stats for all columns are present in valueStats, we can safely ignore file.keyStats() - TODO: validate this assumption
     SimpleStats valueStats = file.valueStats();
     if (valueStats != null) {
-      //      log.info("Processing valueStats: {}", valueStats.toRow());
       List<String> colNames = file.valueStatsCols();
-      //      log.info("valueStatsCols: {}", colNames);
       if (colNames == null || colNames.isEmpty()) {
         // if column names are not present, we assume all columns in the schema are present in the same order as the schema - TODO: validate this assumption
         colNames =
@@ -137,11 +134,6 @@ private void extractStats(
     BinaryRow maxValues = stats.maxValues();
     BinaryArray nullCounts = stats.nullCounts();
 
-    //    log.info("Extracting stats for columns: {}", colNames);
-    //    log.info("minValues: arity={}, {}", minValues.getFieldCount(), minValues);
-    //    log.info("maxValues: arity={}, {}", maxValues.getFieldCount(), maxValues);
-    //    log.info("fieldMap: {}", fieldMap.toString());
-
     for (int i = 0; i < colNames.size(); i++) {
       String colName = colNames.get(i);
       InternalField field = fieldMap.get(colName);
@@ -161,14 +153,6 @@ private void extractStats(
       Object max = getValue(maxValues, i, type, field.getSchema());
       Long nullCount = (nullCounts != null && i < nullCounts.size()) ? nullCounts.getLong(i) : 0L;
 
-      //      log.info(
-      //          "Column: {}, Index: {}, Min: {}, Max: {}, NullCount: {}",
-      //          colName,
-      //          i,
-      //          min,
-      //          max,
-      //          nullCount);
-
       columnStats.add(
           ColumnStat.builder()
               .field(field)
@@ -210,16 +194,14 @@ private Object getValue(BinaryRow row, int index, InternalType type, InternalSch
               fieldSchema.getName());
           tsPrecision = TimestampType.DEFAULT_PRECISION;
         }
-        // TODO: BinaryRow.getTimestamp().toInstant() is deprecated (use LocalZoneTimestamp), but BinaryRow does not have a method to get LocalZoneTimestamp?
-        Instant timestamp = row.getTimestamp(index, tsPrecision).toInstant();
-        long tsMillis = timestamp.toEpochMilli();
+        Timestamp ts = row.getTimestamp(index, tsPrecision);
 
         // according to docs for org.apache.xtable.model.stat.Range, timestamp is stored as millis
         // or micros - even if precision is higher than micros, return micros
         if (tsPrecisionEnum == InternalSchema.MetadataValue.MILLIS) {
-          return tsMillis;
+          return ts.getMillisecond();
         } else {
-          return tsMillis * 1000 + timestamp.getNano() / 1000L;
+          return ts.toMicros();
         }
       case FLOAT:
         return row.getFloat(index);

xtable-core/src/test/java/org/apache/xtable/TestPaimonTable.java

@@ -70,9 +70,23 @@ public static GenericTable<GenericRow, String> createTable(
       Path tempDir,
       Configuration hadoopConf,
       boolean additionalColumns) {
+
+    Schema schema = buildGenericSchema(partitionField, additionalColumns);
+    return createTable(
+        tableName, partitionField, tempDir, hadoopConf, additionalColumns, schema);
+  }
+
+  public static GenericTable<GenericRow, String> createTable(
+      String tableName,
+      String partitionField,
+      Path tempDir,
+      Configuration hadoopConf,
+      boolean additionalColumns,
+      Schema schema) {
     String basePath = initBasePath(tempDir, tableName);
     Catalog catalog = createFilesystemCatalog(basePath, hadoopConf);
-    FileStoreTable paimonTable = createTable(catalog, partitionField, additionalColumns);
+    FileStoreTable paimonTable =
+        createTable(catalog, tableName, schema);
 
     System.out.println(
         "Initialized Paimon test table at base path: "
@@ -91,19 +105,20 @@ public static Catalog createFilesystemCatalog(String basePath, Configuration had
   }
 
   public static FileStoreTable createTable(
-      Catalog catalog, String partitionField, boolean additionalColumns) {
+      Catalog catalog,
+      String tableName,
+      Schema schema) {
     try {
       catalog.createDatabase("test_db", true);
-      Identifier identifier = Identifier.create("test_db", "test_table");
-      Schema schema = buildSchema(partitionField, additionalColumns);
+      Identifier identifier = Identifier.create("test_db", tableName);
       catalog.createTable(identifier, schema, true);
       return (FileStoreTable) catalog.getTable(identifier);
     } catch (Exception e) {
       throw new RuntimeException(e);
     }
   }
 
-  private static Schema buildSchema(String partitionField, boolean additionalColumns) {
+  private static Schema buildGenericSchema(String partitionField, boolean additionalColumns) {
     Schema.Builder builder =
         Schema.newBuilder()
             .primaryKey("id")
@@ -179,20 +194,12 @@ public List<GenericRow> insertRecordsForSpecialPartition(int numRows) {
   }
 
   private List<GenericRow> insertRecordsToPartition(int numRows, String partitionValue) {
-    BatchWriteBuilder batchWriteBuilder = paimonTable.newBatchWriteBuilder();
-    try (BatchTableWrite writer = batchWriteBuilder.newWrite()) {
-      List<GenericRow> rows = new ArrayList<>(numRows);
-      for (int i = 0; i < numRows; i++) {
-        GenericRow row = buildGenericRow(i, paimonTable.schema(), partitionValue);
-        writer.write(row);
-        rows.add(row);
-      }
-      commitWrites(batchWriteBuilder, writer);
-      compactTable();
-      return rows;
-    } catch (Exception e) {
-      throw new RuntimeException("Failed to insert rows into Paimon table", e);
+    List<GenericRow> rows = new ArrayList<>(numRows);
+    for (int i = 0; i < numRows; i++) {
+      rows.add(buildGenericRow(i, paimonTable.schema(), partitionValue));
     }
+    writeRows(paimonTable, rows);
+    return rows;
   }
 
   @Override
@@ -225,8 +232,12 @@ public void deleteRows(List<GenericRow> rows) {
   }
 
   private void compactTable() {
-    BatchWriteBuilder batchWriteBuilder = paimonTable.newBatchWriteBuilder();
-    SnapshotReader snapshotReader = paimonTable.newSnapshotReader();
+    compactTable(paimonTable);
+  }
+
+  public static void compactTable(FileStoreTable table) {
+    BatchWriteBuilder batchWriteBuilder = table.newBatchWriteBuilder();
+    SnapshotReader snapshotReader = table.newSnapshotReader();
     try (BatchTableWrite writer = batchWriteBuilder.newWrite()) {
       for (BucketEntry bucketEntry : snapshotReader.bucketEntries()) {
         writer.compact(bucketEntry.partition(), bucketEntry.bucket(), true);
@@ -237,6 +248,19 @@ private void compactTable() {
     }
   }
 
+  public static void writeRows(FileStoreTable table, List<GenericRow> rows) {
+    BatchWriteBuilder batchWriteBuilder = table.newBatchWriteBuilder();
+    try (BatchTableWrite writer = batchWriteBuilder.newWrite()) {
+      for (GenericRow row : rows) {
+        writer.write(row);
+      }
+      commitWrites(batchWriteBuilder, writer);
+      compactTable(table);
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to write rows into Paimon table", e);
+    }
+  }
+
   private static void commitWrites(BatchWriteBuilder batchWriteBuilder, BatchTableWrite writer)
       throws Exception {
     BatchTableCommit commit = batchWriteBuilder.newCommit();

xtable-core/src/test/java/org/apache/xtable/paimon/TestPaimonConversionSource.java

@@ -98,7 +98,7 @@ void testGetTableWithUnpartitionedTable() {
     InternalTable result = unpartitionedSource.getTable(snapshot);
 
     assertNotNull(result);
-    assertEquals("test_table", result.getName());
+    assertEquals("unpartitioned_table", result.getName());
     assertEquals(TableFormat.PAIMON, result.getTableFormat());
     assertNotNull(result.getReadSchema());
     assertEquals(DataLayoutStrategy.HIVE_STYLE_PARTITION, result.getLayoutStrategy());