chore: refactor some existing methods

[tchow-zlai]

Summary

Don't rely on TableUtils directly as much as possible, use the dataframe extensions.

Checklist

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

Summary by CodeRabbit

• Refactor
  • Updated DataFrame saving methods across multiple Spark job classes
  • Replaced tableUtils.insertPartitions and tableUtils.insertUnPartitioned with direct DataFrame save and saveUnPartitioned methods
  • Simplified data persistence mechanism in LogFlattenerJob, StagingQuery, CompareJob, and ConsistencyJob

[coderabbitai]

Walkthrough

The pull request introduces a systematic refactoring of DataFrame saving methods across multiple Spark-related classes. The changes replace tableUtils method calls with direct DataFrame save methods, streamlining the data persistence mechanism. This modification appears to be part of a broader effort to simplify and standardize DataFrame handling in the Chronon Spark library.

Changes

File	Change Summary
spark/src/main/scala/ai/chronon/spark/LogFlattenerJob.scala	Replaced tableUtils.insertPartitions with flattenedDf.save for DataFrame persistence
spark/src/main/scala/ai/chronon/spark/StagingQuery.scala	Substituted tableUtils.insertPartitions with df.save method
spark/src/main/scala/ai/chronon/spark/stats/CompareJob.scala	Replaced tableUtils.insertUnPartitioned with direct saveUnPartitioned calls on DataFrames
spark/src/main/scala/ai/chronon/spark/stats/ConsistencyJob.scala	Changed insertPartitions to save method for DataFrame persistence

Possibly related PRs

• Summary upload #50: Similar DataFrame saving method updates
• chore: separate column predicates from partition filters #149: Column predicate and partition filter separation
• chore: small cleanup of TableUtils #153: TableUtils class cleanup
• chore: simplify some Extensions #165: DataFrame extension simplification
• chore: modify tool-versions to align with requirements #166: Tool version updates

Suggested Reviewers

• nikhil-zlai
• piyush-zlai
• varant-zlai

Poem

🚀 DataFrame dance, a new refrain,
Saving methods shed their old chain,
Spark's magic, now more sleek and bright,
Code flows with elegant might!
Chronon's library takes flight! 🌟

---

📜 Recent 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 0f79d45 and 8709e08.

📒 Files selected for processing (4)

• spark/src/main/scala/ai/chronon/spark/LogFlattenerJob.scala (2 hunks)
• spark/src/main/scala/ai/chronon/spark/StagingQuery.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/stats/CompareJob.scala (2 hunks)
• spark/src/main/scala/ai/chronon/spark/stats/ConsistencyJob.scala (1 hunks)

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

• spark/src/main/scala/ai/chronon/spark/stats/ConsistencyJob.scala
• spark/src/main/scala/ai/chronon/spark/StagingQuery.scala
• spark/src/main/scala/ai/chronon/spark/stats/CompareJob.scala
• spark/src/main/scala/ai/chronon/spark/LogFlattenerJob.scala

---

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR. (Beta)
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 1

📜 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 1e87f5c and 8256ef4.

📒 Files selected for processing (5)

• spark/src/main/scala/ai/chronon/spark/LogFlattenerJob.scala (2 hunks)
• spark/src/main/scala/ai/chronon/spark/StagingQuery.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/stats/CompareJob.scala (2 hunks)
• spark/src/main/scala/ai/chronon/spark/stats/ConsistencyJob.scala (1 hunks)
• spark/src/test/scala/ai/chronon/spark/test/TableUtilsTest.scala (0 hunks)

💤 Files with no reviewable changes (1)

• spark/src/test/scala/ai/chronon/spark/test/TableUtilsTest.scala

🔇 Additional comments (6)

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

83-83: Check partition expansions
Use a try/catch or confirm that the table exists.

spark/src/main/scala/ai/chronon/spark/stats/ConsistencyJob.scala (1)

134-134: Validate table write
Ensure correct partition data after saving.

spark/src/main/scala/ai/chronon/spark/stats/CompareJob.scala (2)

28-28: New import usage
Extensions needed for saveUnPartitioned. Looks fine.

---

91-91: Metrics table save
Make sure metrics are consistent with compare logic.

✅ Verification successful

Let me gather information about the metrics and compare logic to verify their consistency.

This will help us understand:

1. The context around metrics table save
2. The compare logic implementation
3. How metrics are calculated and used

---

Let me search for the compare logic implementation to verify its consistency with metrics.

---

Let me check how the metrics are computed and what they contain.

---

Let me check the metric transformations and compare logic implementation.

---

Let me check the final implementation of metrics computation and comparison.

---

Based on the code analysis, I can now provide a final response:

Metrics are correctly aligned with compare logic
The metrics computation is properly integrated with the comparison process:

• Metrics are generated from value fields and mapped columns
• Comparison results feed into metrics through CompareBaseJob.compare
• Test cases verify metrics consistency with actual data mismatches

🏁 Scripts executed

The following scripts were executed for the analysis:

Error: No shell scripts found

---

Script:

#!/bin/bash
# Find the metrics table save implementation and surrounding context
rg -A 10 -B 10 "saveUnPartitioned.*metricsTableName" 

# Find compare logic implementation
ast-grep --pattern 'def compare($$$) {
  $$$
}'

# Search for metric calculation logic
rg -A 5 "metrics" --type scala

Length of output: 76428

---

Script:

#!/bin/bash
# Find compare logic implementation
ast-grep --pattern 'def compare($$$) {
  $$$
  CompareBaseJob.compare($$$)
  $$$
}'

# Find metrics computation in CompareBaseJob
rg -A 10 "CompareMetrics.compute" 

Length of output: 778

---

Script:

#!/bin/bash
# Find metrics computation implementation
rg -A 15 "def compute.*fields.*compareDf.*keys.*name.*mapping" --type scala

# Find metric transformations
ast-grep --pattern 'val metrics = $$$'

Length of output: 2642

---

Script:

#!/bin/bash
# Find metric transformations for numeric and any columns
rg -A 10 "def numericTransforms|def anyTransforms" --type scala

# Find CompareMetrics compute implementation
rg -A 20 "object CompareMetrics" --type scala

Length of output: 4228

---

Script:

#!/bin/bash
# Find CompareMetrics compute implementation
rg -B 5 -A 20 "def compute.*valueSchema.*compareDf.*keys.*name.*mapping" --type scala

# Find how metrics are used in comparison
rg -A 10 "CompareBaseJob.compare.*leftDf.*rightDf.*keys" --type scala

Length of output: 1468

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

25-25: New import
Required for the new save method. OK.

---

229-231: Auto-expand caution
Check columns added at runtime won't break downstream.

✅ Verification successful

Let me gather information about the codebase to understand the auto-expand feature and its potential impact.

---

Based on the search results, I can now generate a final response:

Auto-expand is safe here
The autoExpand parameter is used consistently across the codebase for schema evolution support, with proper validation in TableUtils. Tests show that incompatible schema changes are caught by IncompatibleSchemaException, protecting downstream dependencies.

🏁 Scripts executed

The following scripts were executed for the analysis:

Error: No shell scripts found

---

Script:

#!/bin/bash
# Search for references to autoExpand parameter and its usage
rg -A 5 "autoExpand\s*=" 

# Look for save method implementation
ast-grep --pattern 'def save($$$) { $$$ }'

# Search for any references to schema evolution or column addition
rg -i "schema.*evolution|add.*column"

Length of output: 18441

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

build.sbt (1)

218-218: Check version alignment

Consider matching the version with gcsio (3.0.3) for consistency.

-    libraryDependencies += "com.google.cloud.bigdataoss" % "util-hadoop" % "3.0.0"
+    libraryDependencies += "com.google.cloud.bigdataoss" % "util-hadoop" % "3.0.3"

📜 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 603c857 and 0f79d45.

📒 Files selected for processing (1)

• build.sbt (1 hunks)