@Override
  public Dataset<PartitionStat> get() {
    StructType partitionSchema =
        spark.sql(String.format("SELECT * FROM %s.partitions", tableName)).schema();
    try {
      partitionSchema.apply("partition");
      return spark
          .sql(String.format("SELECT partition, file_count FROM %s.partitions", tableName))
          .map(new TablePartitionStats.PartitionStatMapper(), Encoders.bean(PartitionStat.class));
    } catch (IllegalArgumentException e) {
      return spark
          .sql(String.format("SELECT null, file_count FROM %s.partitions", tableName))
          .map(new TablePartitionStats.PartitionStatMapper(), Encoders.bean(PartitionStat.class));
    }
  }
@Test
  public void testPartitionedTablePartitionStats() throws Exception {
    final String testTable = "db.test_table_partition_stats_partitioned";
    try (SparkSession spark = getSparkSession()) {
      spark.sql("USE openhouse");
      spark.sql(
          String.format(
              "CREATE TABLE %s (id INT, data STRING, dt STRING) PARTITIONED BY (dt, id)",
              testTable));
      spark.sql(String.format("INSERT INTO %s VALUES (0, '0', '2024-01-01')", testTable));
      spark.sql(String.format("INSERT INTO %s VALUES (1, '1', '2024-01-02')", testTable));
      spark.sql(String.format("INSERT INTO %s VALUES (1, '2', '2024-01-02')", testTable));
      TablePartitionStats tablePartitionStats =
          TablePartitionStats.builder().spark(spark).tableName(testTable).build();
      List<PartitionStat> stats = tablePartitionStats.get().collectAsList();
      Assertions.assertEquals(2, stats.size());
      stats.sort(Comparator.comparing(a -> a.getValues().get(0)));
      Assertions.assertEquals(Arrays.asList("2024-01-01", "0"), stats.get(0).getValues());
      Assertions.assertEquals(1, stats.get(0).getFileCount());
      Assertions.assertEquals(Arrays.asList("2024-01-02", "1"), stats.get(1).getValues());
      Assertions.assertEquals(2, stats.get(1).getFileCount());
    }
  }