@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 testNonPartitionedTablePartitionStats() throws Exception {
    final String testTable = "db.test_table_partition_stats_non_partitioned";
    try (SparkSession spark = getSparkSession()) {
      spark.sql("USE openhouse");
      spark.sql(String.format("CREATE TABLE %s (id INT, data STRING)", testTable));
      spark.sql(String.format("INSERT INTO %s VALUES (0, '0')", testTable));
      spark.sql(String.format("INSERT INTO %s VALUES (1, '1')", testTable));
      TablePartitionStats tablePartitionStats =
          TablePartitionStats.builder().spark(spark).tableName(testTable).build();
      List<PartitionStat> stats = tablePartitionStats.get().collectAsList();
      Assertions.assertEquals(1, stats.size());
      Assertions.assertTrue(stats.get(0).getValues().isEmpty());
      Assertions.assertEquals(2, stats.get(0).getFileCount());
    }
  }