Update spark/src/main/scala/ai/chronon/spark/TableUtils.scala

david-zlai · nikhil-zlai · web-flow · commit 9159d3d2d110 · 2025-01-17T14:45:41.000-05:00
Co-authored-by: Nikhil Simha &lt;nikhil@zipline.ai&gt;
diff --git a/spark/src/main/scala/ai/chronon/spark/TableUtils.scala b/spark/src/main/scala/ai/chronon/spark/TableUtils.scala
@@ -786,8 +786,9 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable
       val whereStr = allWheres.map(w => s"($w)").mkString(" AND ")
       df = df.where(whereStr)
     }
-    val partitionCount = sparkSession.sparkContext.getConf.getInt("spark.default.parallelism", 1000)
-    df.coalesce(partitionCount)
+    val parallelism = sparkSession.sparkContext.getConf.getInt("spark.default.parallelism", 1000)
+    val coalesceFactor = sparkSession.sparkContext.getConf.getInt("spark.chronon.coalesce.factor", 4)
+    df.coalesce(coalesceFactor * parallelism)
   }
 
   def whereClauses(partitionRange: PartitionRange, partitionColumn: String = partitionColumn): Seq[String] = {

Original file line number	Diff line number	Diff line change
`@@ -786,8 +786,9 @@ class TableUtils(@transient val sparkSession: SparkSession) extends Serializable`
`786`	`786`	`val whereStr = allWheres.map(w => s"($w)").mkString(" AND ")`
`787`	`787`	`df = df.where(whereStr)`
`788`	`788`	`}`
`789`		`- val partitionCount = sparkSession.sparkContext.getConf.getInt("spark.default.parallelism", 1000)`
`790`		`- df.coalesce(partitionCount)`
	`789`	`+ val parallelism = sparkSession.sparkContext.getConf.getInt("spark.default.parallelism", 1000)`
	`790`	`+ val coalesceFactor = sparkSession.sparkContext.getConf.getInt("spark.chronon.coalesce.factor", 4)`
	`791`	`+ df.coalesce(coalesceFactor * parallelism)`
`791`	`792`	`}`
`792`	`793`
`793`	`794`	`def whereClauses(partitionRange: PartitionRange, partitionColumn: String = partitionColumn): Seq[String] = {`