Add more logging and error handling to partition pruning

airbnb · Feb 1, 2024 · 5cb59f3 · 5cb59f3
1 parent e663aef
commit 5cb59f3
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/spark/src/main/scala/ai/chronon/spark/Analyzer.scala b/spark/src/main/scala/ai/chronon/spark/Analyzer.scala
@@ -295,7 +295,7 @@ class Analyzer(tableUtils: TableUtils,
       val gbStartPartition = part.groupBy.sources.toScala
         .map(_.query.startPartition)
         .filter(_ != null)
-      if (!gbStartPartition.isEmpty)
+      if (gbStartPartition.nonEmpty)
         gbStartPartitions += (part.groupBy.metaData.name -> gbStartPartition)
     }
     val noAccessTables = runTablePermissionValidation((gbTables.toList ++ List(joinConf.left.table)).toSet)

diff --git a/spark/src/main/scala/ai/chronon/spark/Extensions.scala b/spark/src/main/scala/ai/chronon/spark/Extensions.scala
@@ -64,9 +64,13 @@ object Extensions {
     val count: Long = partitionCounts.values.sum
 
     def prunePartitions(range: PartitionRange): Option[DfWithStats] = {
+      println(
+        s"Pruning down to new range $range, original range: $partitionRange." +
+          s"\nOriginal partition counts: $partitionCounts")
       val intersected = partitionRange.intersect(range)
       if (!intersected.wellDefined) return None
       val intersectedCounts = partitionCounts.filter(intersected.partitions contains _._1)
+      if (intersectedCounts.isEmpty) return None
       Some(DfWithStats(df.prunePartition(range), intersectedCounts))
     }
     def stats: DfStats = DfStats(count, partitionRange)