WIP -- Not ready to merge -- Spark refactor

[varant-zlai]

Summary

Checklist

• Added Unit Tests
• Covered by existing CI
• Integration tested
• Documentation update

Summary by CodeRabbit

• New Features

  • Introduced new classes for managing join operations, including Bootstrap, BootstrapJob, JoinPartJob, DerivationJob, and SourceJob, enhancing the data processing capabilities.
  • Added functionalities for generating bootstrap tables, handling derivation jobs, and managing skew filtering, improving the efficiency of join operations.
  • Added a new structure SourceWithFilter to enhance data modeling capabilities within the orchestration API.

• Bug Fixes

  • Improved error handling and logging during join processing to ensure robustness.

• Refactor

  • Streamlined existing join logic by removing outdated methods and consolidating functionalities, enhancing maintainability and clarity.

[coderabbitai]

Actionable comments posted: 5

♻️ Duplicate comments (2)

spark/src/main/scala/ai/chronon/spark/DerivationJob.scala (2)

27-32: ⚠️ Potential issue

Add parameter validation and replace table name placeholders.

Validate input parameters and define actual table names.

   def fromJoin(join: api.Join, dateRange: PartitionRange): DerivationJob = {
+    require(join != null, "join cannot be null")
+    require(dateRange != null, "dateRange cannot be null")
+    require(join.derivations != null, "derivations cannot be null")
     val baseOutputTable = "TODO" // Output of the base Join pre-derivation
     val finalOutputTable = "TODO" // The actual output table
     val derivations = join.derivations.asScala

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

  def fromJoin(join: api.Join, dateRange: PartitionRange): DerivationJob = {
+    require(join != null, "join cannot be null")
+    require(dateRange != null, "dateRange cannot be null")
+    require(join.derivations != null, "derivations cannot be null")
    val baseOutputTable = "TODO" // Output of the base Join pre-derivation
    val finalOutputTable = "TODO" // The actual output table
    val derivations = join.derivations.asScala
    new DerivationJob(baseOutputTable, finalOutputTable, derivations, dateRange)
  }

---

20-25: ⚠️ Potential issue

Add parameter validation and replace table name placeholders.

Validate input parameters and define actual table names.

   def fromGroupBy(groupBy: api.GroupBy, dateRange: PartitionRange): DerivationJob = {
+    require(groupBy != null, "groupBy cannot be null")
+    require(dateRange != null, "dateRange cannot be null")
+    require(groupBy.derivations != null, "derivations cannot be null")
     val baseOutputTable = "TODO" // Output of the base GroupBy pre-derivation
     val finalOutputTable = "TODO" // The actual output table
     val derivations = groupBy.derivations.asScala

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

  def fromGroupBy(groupBy: api.GroupBy, dateRange: PartitionRange): DerivationJob = {
    require(groupBy != null, "groupBy cannot be null")
    require(dateRange != null, "dateRange cannot be null")
    require(groupBy.derivations != null, "derivations cannot be null")
    val baseOutputTable = "TODO" // Output of the base GroupBy pre-derivation
    val finalOutputTable = "TODO" // The actual output table
    val derivations = groupBy.derivations.asScala
    new DerivationJob(baseOutputTable, finalOutputTable, derivations, dateRange)
  }

🧹 Nitpick comments (5)

spark/src/main/scala/ai/chronon/spark/SourceJob.scala (1)

69-74: Use constants for null values.

Replace hardcoded nulls with constants.

-    val nulls = Seq("null", "Null", "NULL")
+    private val NULL_VALUES = Seq("null", "Null", "NULL")
+    def generateSkewFilterSql(key: String, values: Seq[String]): String = {
+      val nonNullFilters = Some(s"$key NOT IN (${values.filterNot(NULL_VALUES.contains).mkString(", ")})")
+      val nullFilters = if (values.exists(NULL_VALUES.contains)) Some(s"$key IS NOT NULL") else None

spark/src/main/scala/ai/chronon/spark/BootstrapJob.scala (1)

72-119: Consider optimizing join operations.

The fold operation with joins could be optimized using broadcast joins for small tables.

spark/src/main/scala/ai/chronon/spark/Join.scala (3)

301-328: Remove commented code.

Good refactoring moving join part logic to a dedicated class. However, remove the commented-out code before finalizing the PR.

-                // val df =
-                // computeRightTable(unfilledLeftDf, joinPart, leftRange, leftTimeRangeOpt, bloomFilterOpt, runSmallMode)
-                //  .map(df => joinPart -> df)
-

---

301-304: Remove commented out code.

Delete the commented out code as it's no longer needed.

-                // val df =
-                // computeRightTable(unfilledLeftDf, joinPart, leftRange, leftTimeRangeOpt, bloomFilterOpt, runSmallMode)
-                //  .map(df => joinPart -> df)
-

---

301-328: Remove commented out code.

Since the code has been refactored to use JoinPartJob, the commented out code can be safely removed.

Apply this diff to remove the commented out code:

-                // val df =
-                // computeRightTable(unfilledLeftDf, joinPart, leftRange, leftTimeRangeOpt, bloomFilterOpt, runSmallMode)
-                //  .map(df => joinPart -> df)
-

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro (Legacy)

📥 Commits

Reviewing files that changed from the base of the PR and between 531477f and 0ee77c9.

📒 Files selected for processing (9)

• api/thrift/orchestration.thrift (1 hunks)
• orchestration/src/main/scala/ai/chronon/orchestration/Bootstrap.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/BootstrapJob.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/DerivationJob.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/Join.scala (2 hunks)
• spark/src/main/scala/ai/chronon/spark/JoinBase.scala (2 hunks)
• spark/src/main/scala/ai/chronon/spark/JoinPartJob.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/JoinUtils.scala (5 hunks)
• spark/src/main/scala/ai/chronon/spark/SourceJob.scala (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (3)

• orchestration/src/main/scala/ai/chronon/orchestration/Bootstrap.scala
• spark/src/main/scala/ai/chronon/spark/JoinPartJob.scala
• spark/src/main/scala/ai/chronon/spark/JoinUtils.scala

⏰ Context from checks skipped due to timeout of 90000ms (4)

• GitHub Check: spark_tests
• GitHub Check: non_spark_tests
• GitHub Check: scala_compile_fmt_fix
• GitHub Check: enforce_triggered_workflows

🔇 Additional comments (11)

spark/src/main/scala/ai/chronon/spark/JoinBase.scala (6)

174-175: LGTM! Good refactoring.

Moving bootstrap logic to a dedicated class improves modularity.

---

296-296: LGTM! Good refactoring.

Moving small mode determination to JoinUtils improves reusability.

---

174-175: LGTM! Good refactoring.

Moving bootstrap logic to a dedicated class improves modularity.

---

296-296: LGTM! Good refactoring.

Centralizing small mode determination logic in JoinUtils improves reusability.

---

174-175: LGTM! Good refactoring.

Moving bootstrap logic to a dedicated class improves code organization.

---

296-296: LGTM! Good refactoring.

Moving small mode determination to JoinUtils improves code reusability.

spark/src/main/scala/ai/chronon/spark/Join.scala (5)

222-223: LGTM! Consistent refactoring.

Changes align with the bootstrap logic refactoring in JoinBase.scala.

---

222-223: LGTM! Consistent refactoring.

Moving bootstrap logic to a dedicated class is consistent with changes in JoinBase.scala.

---

305-328: LGTM! Good refactoring.

Moving join part logic to a dedicated class with proper context improves modularity.

---

1-1: Add missing items before merge.

As indicated in the PR description, please ensure the following are completed before merge:

1. Add unit tests
2. Ensure CI coverage
3. Perform integration testing
4. Update documentation

---

222-223: LGTM! Good refactoring.

Moving bootstrap logic to a dedicated class improves code organization.

[coderabbitai]

⚠️ Potential issue

Implement the run method.

Empty implementation needs to be completed.

[coderabbitai]

⚠️ Potential issue

Add null check for source.

Prevent potential NPE.

-    val timeProjection = if (source.dataModel == Events) {
+    val timeProjection = if (source != null && source.dataModel == Events) {

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	val timeProjection = if (source.dataModel == Events) {
	Seq(Constants.TimeColumn -> Option(source.query).map(_.timeColumn).orNull)
	} else {
	Seq()
	}
	val timeProjection = if (source != null && source.dataModel == Events) {
	Seq(Constants.TimeColumn -> Option(source.query).map(_.timeColumn).orNull)
	} else {
	Seq()
	}

[coderabbitai]

⚠️ Potential issue

Implement the run method.

Empty implementation needs to be completed.

[coderabbitai]

💡 Verification agent

🧩 Analysis chain

Address TODO about deduplication logic customization.

Implement customizable deduplication strategy.

---

🏁 Script executed:

#!/bin/bash
# Check if there are any existing customization patterns in the codebase
rg -A 5 "deduplication.*logic" 

Length of output: 826

---

Custom deduplication logic is not implemented.

• The TODO remains, and the code still hard-codes .dropDuplicates(...) without a customization hook.

[coderabbitai]

🛠️ Refactor suggestion

Add documentation for SourceWithFilter struct.

Add Thrift documentation comments explaining the purpose and usage of this struct.

+/**
+ * Represents a source with optional key exclusion filters.
+ * @param source The source to filter
+ * @param excludeKeys Map of column names to values that should be excluded
+ */
 struct SourceWithFilter {
 }

Committable suggestion skipped: line range outside the PR's diff.

[coderabbitai]

Actionable comments posted: 2

♻️ Duplicate comments (1)

spark/src/main/scala/ai/chronon/spark/SourceJob.scala (1)

24-28: ⚠️ Potential issue

Add null check for source.

Prevent potential NPE by checking if source is null before accessing dataModel.

-    val timeProjection = if (source.dataModel == Events) {
+    val timeProjection = if (source != null && source.dataModel == Events) {

🧹 Nitpick comments (5)

spark/src/main/scala/ai/chronon/spark/JoinUtils.scala (2)

446-470: Add input validation for skewKeys.

Consider validating that skewKeys values are non-empty when present.

 def skewFilter(keys: Option[Seq[String]] = None,
                skewKeys: Option[Map[String, Seq[String]]],
                leftKeyCols: Seq[String],
                joiner: String = " OR "): Option[String] = {
+    require(skewKeys.forall(_.values.forall(_.nonEmpty)), "skewKeys values must not be empty")
     skewKeys.map { keysMap =>

---

521-535: Extract magic numbers into constants.

The threshold count appears in multiple log messages. Consider extracting it to avoid repetition.

+  private val SmallModeLogMsg = "Counted %d rows, running join in small mode."
+  private val NormalModeLogMsg = "Counted greater than %d rows, proceeding with normal computation."
   def runSmallMode(tableUtils: TableUtils, leftDf: DataFrame): Boolean = {
     if (tableUtils.smallModelEnabled) {
       val thresholdCount = leftDf.limit(Some(tableUtils.smallModeNumRowsCutoff + 1).get).count()
       val result = thresholdCount <= tableUtils.smallModeNumRowsCutoff
       if (result) {
-        logger.info(s"Counted $thresholdCount rows, running join in small mode.")
+        logger.info(SmallModeLogMsg.format(thresholdCount))
       } else {
-        logger.info(
-          s"Counted greater than ${tableUtils.smallModeNumRowsCutoff} rows, proceeding with normal computation.")
+        logger.info(NormalModeLogMsg.format(tableUtils.smallModeNumRowsCutoff))
       }

spark/src/main/scala/ai/chronon/spark/SourceJob.scala (1)

47-49: Enhance error message for empty DataFrame.

Add more context to help debug empty result sets.

-      throw new RuntimeException(s"Query produced 0 rows in range $range.")
+      throw new RuntimeException(s"Query for source ${source.table} produced 0 rows in range $range. Please verify the source data and filters.")

spark/src/main/scala/ai/chronon/spark/BootstrapJob.scala (1)

191-195: Enhance error handling with specific error types.

Current error handling is too generic.

-    } catch {
-      case e: Exception =>
-        logger.error(s"Error while processing groupBy: ${joinPart.groupBy.getMetaData.getName}")
-        throw e
+    } catch {
+      case e: IllegalArgumentException =>
+        logger.error(s"Invalid arguments while processing groupBy: ${joinPart.groupBy.getMetaData.getName}", e)
+        throw e
+      case e: Exception =>
+        logger.error(s"Unexpected error while processing groupBy: ${joinPart.groupBy.getMetaData.getName}", e)
+        throw new RuntimeException(s"Bootstrap computation failed: ${e.getMessage}", e)

spark/src/main/scala/ai/chronon/spark/Join.scala (1)

343-345: Remove commented code.

Clean up commented code that's no longer needed.

-                // val df =
-                // computeRightTable(unfilledLeftDf, joinPart, leftRange, leftTimeRangeOpt, bloomFilterOpt, runSmallMode)
-                //  .map(df => joinPart -> df)

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro (Legacy)

📥 Commits

Reviewing files that changed from the base of the PR and between bd4e828 and b4bd285.

📒 Files selected for processing (9)

• api/thrift/orchestration.thrift (1 hunks)
• orchestration/src/main/scala/ai/chronon/orchestration/Bootstrap.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/BootstrapJob.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/DerivationJob.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/Join.scala (4 hunks)
• spark/src/main/scala/ai/chronon/spark/JoinBase.scala (2 hunks)
• spark/src/main/scala/ai/chronon/spark/JoinPartJob.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/JoinUtils.scala (5 hunks)
• spark/src/main/scala/ai/chronon/spark/SourceJob.scala (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (3)

• api/thrift/orchestration.thrift
• orchestration/src/main/scala/ai/chronon/orchestration/Bootstrap.scala
• spark/src/main/scala/ai/chronon/spark/DerivationJob.scala

⏰ Context from checks skipped due to timeout of 90000ms (4)

• GitHub Check: spark_tests
• GitHub Check: scala_compile_fmt_fix
• GitHub Check: non_spark_tests
• GitHub Check: enforce_triggered_workflows

🔇 Additional comments (11)

spark/src/main/scala/ai/chronon/spark/JoinUtils.scala (4)

308-313: LGTM! Good refactoring.

The parameter change from joinConf to leftDataModel improves method focus and reduces coupling.

---

495-500: LGTM! Robust null handling.

The method correctly handles both null and non-null cases in the SQL filter generation.

---

502-519: LGTM! Good early returns.

The method has clear early return conditions and proper null checks.

---

537-542: LGTM! Clean conversion utility.

The method is focused and handles the conversion well.

spark/src/main/scala/ai/chronon/spark/SourceJob.scala (1)

69-74: LGTM! Robust SQL filter generation.

The SQL generation handles both null and non-null cases correctly.

spark/src/main/scala/ai/chronon/spark/BootstrapJob.scala (2)

117-118: Implement customizable deduplication logic.

The TODO indicates a need for flexible deduplication strategies.

Would you like me to propose a design for customizable deduplication logic?

---

148-177: LGTM! Robust external field padding.

The implementation correctly handles both contextual and non-contextual fields.

spark/src/main/scala/ai/chronon/spark/JoinPartJob.scala (1)

289-302: LGTM! Comprehensive join logic.

The pattern matching covers all combinations of data models and accuracies correctly.

spark/src/main/scala/ai/chronon/spark/JoinBase.scala (2)

174-175: LGTM! Clean bootstrap job integration.

Good modularization of bootstrap logic into a separate job class.

---

295-295: LGTM! Centralized small mode determination.

Good refactoring to use JoinUtils for consistent small mode logic.

spark/src/main/scala/ai/chronon/spark/Join.scala (1)

347-370: LGTM! Good refactoring to improve modularity.

The introduction of JoinPartJobContext and JoinPartJob improves code organization and maintainability.

[coderabbitai]

🛠️ Refactor suggestion

Replace println with logger.

Use proper logging instead of println for debugging.

-    println("leftSkewFilter: " + leftSkewFilter)
+    logger.debug(s"leftSkewFilter: $leftSkewFilter")

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	println("leftSkewFilter: " + leftSkewFilter)
	logger.debug(s"leftSkewFilter: $leftSkewFilter")

[coderabbitai]

⚠️ Potential issue

Remove duplicated code block.

This block is an exact duplicate of the code above at lines 318-324.

-                // if the join part contains ChrononRunDs macro, then we need to make sure the join is for a single day
-                val selects = Option(joinPart.groupBy.sources.toScala.map(_.query.selects).map(_.toScala))
-                if (
-                  selects.isDefined && selects.get.nonEmpty && selects.get.exists(selectsMap =>
-                    Option(selectsMap).isDefined && selectsMap.values.exists(_.contains(Constants.ChrononRunDs)))
-                ) {
-                  assert(
-                    leftRange.isSingleDay,
-                    s"Macro ${Constants.ChrononRunDs} is only supported for single day join, current range is $leftRange")
-                }

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	// if the join part contains ChrononRunDs macro, then we need to make sure the join is for a single day
	val selects = Option(joinPart.groupBy.sources.toScala.map(_.query.selects).map(_.toScala))
	if (
	selects.isDefined && selects.get.nonEmpty && selects.get.exists(selectsMap =>
	Option(selectsMap).isDefined && selectsMap.values.exists(_.contains(Constants.ChrononRunDs)))
	) {
	assert(
	leftRange.isSingleDay,
	s"Macro ${Constants.ChrononRunDs} is only supported for single day join, current range is $leftRange")
	}

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

spark/src/main/scala/ai/chronon/spark/SourceJob.scala (2)

14-17: Enhance class documentation.

Add parameter descriptions and example usage.

 /*
 Runs and materializes a `Source` for a given `dateRange`. Used in the Join computation flow to first compute the Source,
 then each join may have a further Bootstrap computation to produce the left side to effectively use.
+ 
+ @param sourceWithFilter Source configuration with skew filtering options
+ @param tableUtils Utility for table operations
+ @param range Partition range to process
+ 
+ Example:
+ {{{
+ val job = new SourceJob(source, utils, range)
+ job.run()
+ }}}
  */

---

51-53: Resolve design decision and add error handling.

The comment suggests uncertainty about parameter design. Also, add error handling for digest generation.

-      f"${source.table}_${ThriftJsonCodec.md5Digest(sourceWithFilter)}" // Or should we pass this in as an arg?
+      f"${source.table}_${Try(ThriftJsonCodec.md5Digest(sourceWithFilter))
+        .getOrElse(throw new RuntimeException("Failed to generate output name digest"))}"

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro (Legacy)

📥 Commits

Reviewing files that changed from the base of the PR and between b4bd285 and 2c40253.

📒 Files selected for processing (3)

• spark/src/main/scala/ai/chronon/spark/BootstrapJob.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/Join.scala (4 hunks)
• spark/src/main/scala/ai/chronon/spark/SourceJob.scala (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• spark/src/main/scala/ai/chronon/spark/BootstrapJob.scala

⏰ Context from checks skipped due to timeout of 90000ms (4)

• GitHub Check: non_spark_tests
• GitHub Check: scala_compile_fmt_fix
• GitHub Check: spark_tests
• GitHub Check: enforce_triggered_workflows

🔇 Additional comments (4)

spark/src/main/scala/ai/chronon/spark/Join.scala (3)

325-334: Remove duplicated code block.

This block is identical to lines 318-324.

---

347-369: LGTM! Good modularization with JoinPartJob.

The refactoring improves code organization by encapsulating join part processing logic.

---

354-365: LGTM! Clear configuration handling.

Good separation of concerns with skew handling and bootstrap table configurations.

spark/src/main/scala/ai/chronon/spark/SourceJob.scala (1)

24-28: Add null check for source.

Prevent potential NPE.

-    val timeProjection = if (source.dataModel == Events) {
+    val timeProjection = if (source != null && source.dataModel == Events) {

[coderabbitai]

⚠️ Potential issue

Prevent SQL injection in filter generation.

Sanitize input values before using in SQL.

   def generateSkewFilterSql(key: String, values: Seq[String]): String = {
     val nulls = Seq("null", "Null", "NULL")
-    val nonNullFilters = Some(s"$key NOT IN (${values.filterNot(nulls.contains).mkString(", ")})")
+    val sanitizedValues = values.filterNot(nulls.contains).map(_.replaceAll("[^a-zA-Z0-9_]", ""))
+    val nonNullFilters = Some(s"$key NOT IN (${sanitizedValues.mkString(", ")})")
     val nullFilters = if (values.exists(nulls.contains)) Some(s"$key IS NOT NULL") else None
     (nonNullFilters ++ nullFilters).mkString(" AND ")
   }

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	def generateSkewFilterSql(key: String, values: Seq[String]): String = {
	val nulls = Seq("null", "Null", "NULL")
	val nonNullFilters = Some(s"$key NOT IN (${values.filterNot(nulls.contains).mkString(", ")})")
	val nullFilters = if (values.exists(nulls.contains)) Some(s"$key IS NOT NULL") else None
	(nonNullFilters ++ nullFilters).mkString(" AND ")
	}
	def generateSkewFilterSql(key: String, values: Seq[String]): String = {
	val nulls = Seq("null", "Null", "NULL")
	val sanitizedValues = values.filterNot(nulls.contains).map(_.replaceAll("[^a-zA-Z0-9_]", ""))
	val nonNullFilters = Some(s"$key NOT IN (${sanitizedValues.mkString(", ")})")
	val nullFilters = if (values.exists(nulls.contains)) Some(s"$key IS NOT NULL") else None
	(nonNullFilters ++ nullFilters).mkString(" AND ")
	}

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (4)

spark/src/main/scala/ai/chronon/spark/JoinUtils.scala (4)

447-470: Add input validation for empty sequences.

Consider adding validation for empty sequences in keys and leftKeyCols parameters.

 def skewFilter(keys: Option[Seq[String]] = None,
                skewKeys: Option[Map[String, Seq[String]]],
                leftKeyCols: Seq[String],
                joiner: String = " OR "): Option[String] = {
+  require(leftKeyCols.nonEmpty, "leftKeyCols cannot be empty")
+  keys.foreach(k => require(k.nonEmpty, "keys cannot be empty when provided"))
   skewKeys.map { keysMap =>

---

494-499: Extract null values as constants.

Move magic strings for null values to constants.

+  private val NullValues = Set("null", "Null", "NULL")
   private def generateSkewFilterSql(key: String, values: Seq[String]): String = {
-    val nulls = Seq("null", "Null", "NULL")
     val nonNullFilters = Some(s"$key NOT IN (${values.filterNot(nulls.contains).mkString(", ")})")
-    val nullFilters = if (values.exists(nulls.contains)) Some(s"$key IS NOT NULL") else None
+    val nullFilters = if (values.exists(NullValues.contains)) Some(s"$key IS NOT NULL") else None

---

501-518: Add method documentation.

Add ScalaDoc to explain the purpose and parameters of findUnfilledRecords.

+  /**
+   * Identifies unfilled records in a bootstrap DataFrame based on covering sets.
+   *
+   * @param bootstrapDfWithStats The bootstrap DataFrame with statistics
+   * @param coveringSets The sequence of covering sets to check against
+   * @return Option[DfWithStats] containing unfilled records if any exist
+   */
   def findUnfilledRecords(bootstrapDfWithStats: DfWithStats, coveringSets: Seq[CoveringSet]): Option[DfWithStats] = {

---

520-534: Extract magic numbers in logging.

Move the cutoff value to a constant to avoid repetition.

   def runSmallMode(tableUtils: TableUtils, leftDf: DataFrame): Boolean = {
     if (tableUtils.smallModelEnabled) {
-      val thresholdCount = leftDf.limit(Some(tableUtils.smallModeNumRowsCutoff + 1).get).count()
+      val cutoff = tableUtils.smallModeNumRowsCutoff
+      val thresholdCount = leftDf.limit(Some(cutoff + 1).get).count()
       val result = thresholdCount <= tableUtils.smallModeNumRowsCutoff

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro (Legacy)

📥 Commits

Reviewing files that changed from the base of the PR and between 2c40253 and 130b318.

📒 Files selected for processing (1)

• spark/src/main/scala/ai/chronon/spark/JoinUtils.scala (4 hunks)

⏰ Context from checks skipped due to timeout of 90000ms (4)

• GitHub Check: scala_compile_fmt_fix
• GitHub Check: non_spark_tests
• GitHub Check: spark_tests
• GitHub Check: enforce_triggered_workflows

🔇 Additional comments (2)

spark/src/main/scala/ai/chronon/spark/JoinUtils.scala (2)

309-346: LGTM! Good refactoring.

Simplified method signature by requiring only the necessary data model information.

---

536-540: LGTM! Clean implementation.

Simple and effective conversion from Java to Scala collections.