WIP - feat: multiple partition format support

[nikhil-zlai]

Summary

We recently added support to specify partition format and partition column to sources.

So when we use these sources in groupBy-s and join left side - we need to first convert the query range to the source specs and convert back the dataframes to contain partitionSpecs in the output spec format (this is set once globally and available via tableUtils.partitionSpec).

the translate methods in the PR are added to date strings, partition ranges and dataframes and re-used across code.

Checklist

• Added Unit Tests
• Covered by existing CI -- varant to add a breaking test and fix it.
• Integration tested
• Documentation update

Summary by CodeRabbit

• Bug Fixes

  • Improved error handling for partition range validation to provide clearer input validation errors.
  • Enhanced partition range computation for joins, ensuring only relevant partitions are processed.

• New Features

  • Added support for specifying partition formats when generating test data, allowing for more flexible and accurate test setups.
  • Introduced partition specification translation and interval window support for improved date and partition handling.
  • Added functionality to translate and reformat partition columns in data frames for consistent partitioning.

• Tests

  • Updated tests to use explicit partition formats and improved configuration for join scenarios.
  • Ensured consistency in partitioning semantics across test cases.

• Refactor

  • Renamed internal methods for clarity and consistency in partition range computation logic.
  • Refactored range and partition condition computations to leverage implicit context and source partition specifications for improved consistency.

[coderabbitai]

Walkthrough

This update refactors the partition range utility method from getRangesToFill to getRangeToFill across core Spark join logic and related tests. Partition format handling is improved in data generation, and error checking is made more idiomatic. Test cases are updated for consistency with these changes.

Changes

File(s)	Change Summary
spark/src/main/scala/ai/chronon/spark/Analyzer.scala spark/src/main/scala/ai/chronon/spark/JoinBase.scala spark/src/main/scala/ai/chronon/spark/JoinUtils.scala	Replaced all usage of JoinUtils.getRangesToFill with JoinUtils.getRangeToFill. Updated method name in utility object. Added logic in JoinBase to filter fillable partitions and a require check for input partitions. Modified JoinUtils.leftDf to translate partition ranges and added implicit TableUtils. Minor formatting fixes. No signature changes except utility method rename.
spark/src/main/scala/ai/chronon/spark/catalog/TableUtils.scala	Changed an assert to a require for input validation when determining partition range start. No other logic changed.
spark/src/test/scala/ai/chronon/spark/test/DataFrameGen.scala	Added optional partitionFormat parameter to gen, events, and entities methods. Calls now propagate partition format for consistent partitioning. Adjusted internal logic to use format in partition spec and transformations.
spark/src/test/scala/ai/chronon/spark/test/LocalTableExporterTest.scala	Updated test to use partition column and partition format from TableUtils when generating DataFrames. Removed redundant TableUtils instantiation.
spark/src/test/scala/ai/chronon/spark/test/join/JoinTest.scala	Modified test to use explicit partition format "yyyyMMdd" for weight table and query. Changed join configuration to use a deep copy with a shifted date range.
spark/src/test/scala/ai/chronon/spark/test/join/JoinUtilsTest.scala	Updated tests to call JoinUtils.getRangeToFill instead of the old method. No other logic changed.
api/src/main/scala/ai/chronon/api/DataRange.scala	Added translate method to PartitionRange to convert partition range strings between partition specs.
api/src/main/scala/ai/chronon/api/PartitionSpec.scala	Added intervalWindow method to get a Window for day/hour spans, throwing on unsupported intervals. Added translate method to convert date strings between partition specs.
spark/src/main/scala/ai/chronon/spark/Extensions.scala	Added translatePartitionSpec method to DataframeOps to rename and reformat partition columns. Modified SourceSparkOps to take implicit TableUtils and added accessors for partition column, format, interval, and spec. Updated imports to support new functionality.
spark/src/main/scala/ai/chronon/spark/GroupBy.scala	Introduced implicit TableUtils and PartitionSpec in getIntersectedRange. Translated query range to source partition spec. Simplified partition conditions construction. Added partition spec translation on source DataFrame. Reformatted logging in computeBackfill. Added import for SourceSparkOps. No signature changes.

Sequence Diagram(s)

sequenceDiagram
    participant Test as Test Suite
    participant DataGen as DataFrameGen
    participant Join as Join Logic
    participant Utils as JoinUtils
    participant Table as TableUtils

    Test->>DataGen: gen(..., partitionFormat)
    DataGen->>DataGen: Use partitionFormat in partition spec
    Test->>Join: computeJoinOpt(...)
    Join->>Utils: getRangeToFill(...)
    Utils->>Table: get partition range info
    Join->>Join: Filter fillable partitions
    Join-->>Test: Return join results

Loading

Possibly related PRs

• Moving PartitionRange out of spark into online #15: Related usage of PartitionRange and partition specs in join logic.
• Refactor spark to create modular jobs for orchestrator #402: Refactors join computations and updates JoinUtils.getRangeToFill usage.
• Remove redundant partition listing #328: Related changes in GroupBy.scala on partition range handling and conditions.

Suggested reviewers

• varant-zlai
• piyush-zlai

Poem

🐇✨
Partition specs align with grace,
Joins now fill their rightful space.
Formats flow through data streams,
Tests confirm our coding dreams.
Require checks guard the gate,
Cleaner code, we celebrate!
🎉🚂

Tip

⚡️ Faster reviews with caching

• CodeRabbit now supports caching for code and dependencies, helping speed up reviews. This means quicker feedback, reduced wait times, and a smoother review experience overall. Cached data is encrypted and stored securely. This feature will be automatically enabled for all accounts on May 16th. To opt out, configure Review - Disable Cache at either the organization or repository level. If you prefer to disable all data retention across your organization, simply turn off the Data Retention setting under your Organization Settings.

Enjoy the performance boost—your workflow just got faster.

---

📜 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 a804788 and c33cbd3.

📒 Files selected for processing (2)

• spark/src/main/scala/ai/chronon/spark/GroupBy.scala (6 hunks)
• spark/src/main/scala/ai/chronon/spark/JoinUtils.scala (3 hunks)

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

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

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

• GitHub Check: streaming_tests
• GitHub Check: groupby_tests
• GitHub Check: streaming_tests
• GitHub Check: join_tests
• GitHub Check: service_tests
• GitHub Check: analyzer_tests
• GitHub Check: fetcher_tests
• GitHub Check: online_tests
• GitHub Check: join_tests
• GitHub Check: batch_tests
• GitHub Check: cloud_gcp_tests
• GitHub Check: groupby_tests
• GitHub Check: analyzer_tests
• GitHub Check: cloud_aws_tests
• GitHub Check: fetcher_tests
• GitHub Check: spark_tests
• GitHub Check: batch_tests
• GitHub Check: api_tests
• GitHub Check: service_tests
• GitHub Check: spark_tests
• GitHub Check: flink_tests
• GitHub Check: service_commons_tests
• GitHub Check: aggregator_tests
• GitHub Check: cloud_gcp_tests
• GitHub Check: cloud_aws_tests
• GitHub Check: api_tests
• GitHub Check: online_tests
• GitHub Check: flink_tests
• GitHub Check: aggregator_tests
• GitHub Check: scala_compile_fmt_fix
• GitHub Check: enforce_triggered_workflows

🔇 Additional comments (6)

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

42-42: LGTM: New import for partition spec translation support.

---

587-590: Good use of implicits for cleaner partition handling.

Making TableUtils and SourcePartitionSpec implicit simplifies downstream code and adds clarity.

---

592-596: Effective query range usage aligns with source partitioning.

Using translated query range ensures proper partition alignment with source.

---

641-642: Consistent use of implicit TableUtils.

Consistent with earlier pattern in getIntersectedRange.

---

671-671: Simplified partition condition generation.

Direct use of whereClauses improves code conciseness.

---

697-707:

✅ Verification successful

Partition spec translation ensures consistent formats.

Final translatePartitionSpec call ensures DataFrame uses correct partition format.

---

🏁 Script executed:

#!/bin/bash
# Check for DataFrame translation uses across the codebase
rg -A 2 "translatePartitionSpec" --type scala

Length of output: 830

---

Partition spec translation consistent across codebase

• This call mirrors usages in JoinUtils.scala and Extensions.scala, ensuring the DataFrame always has the correct partition format.
• No further changes needed.

---

🪧 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.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

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 sequence diagram to generate a sequence diagram of the changes in this PR.
• @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: 0

🧹 Nitpick comments (1)

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

290-299: Added partition existence validation.

Good addition to filter partitions that actually exist in the source table.

Consider uncommenting the require statement once the implementation is fully tested, or remove it if not needed.

-//    require(
-//      fillableRanges.nonEmpty,
-//      s"""No relevant input partitions present in ${joinConfCloned.left.table}
-//         |on join.left for the requested range ${rangeToFill.start} - ${rangeToFill.end} """.stripMargin
-//    )
+    if (fillableRanges.isEmpty) {
+      logger.warn(s"""No relevant input partitions present in ${joinConfCloned.left.table}
+        |on join.left for the requested range ${rangeToFill.start} - ${rangeToFill.end} """.stripMargin)
+    }

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 3779e99 and 4ce3bcb.

📒 Files selected for processing (8)

• spark/src/main/scala/ai/chronon/spark/Analyzer.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/JoinBase.scala (2 hunks)
• spark/src/main/scala/ai/chronon/spark/JoinUtils.scala (1 hunks)
• spark/src/main/scala/ai/chronon/spark/catalog/TableUtils.scala (1 hunks)
• spark/src/test/scala/ai/chronon/spark/test/DataFrameGen.scala (3 hunks)
• spark/src/test/scala/ai/chronon/spark/test/LocalTableExporterTest.scala (2 hunks)
• spark/src/test/scala/ai/chronon/spark/test/join/JoinTest.scala (2 hunks)
• spark/src/test/scala/ai/chronon/spark/test/join/JoinUtilsTest.scala (2 hunks)

🧰 Additional context used

🧬 Code Graph Analysis (2)

spark/src/test/scala/ai/chronon/spark/test/LocalTableExporterTest.scala (2)

spark/src/main/scala/ai/chronon/spark/catalog/TableUtils.scala (2)

• TableUtils (45-584)
• TableUtils (586-588)

spark/src/test/scala/ai/chronon/spark/test/DataFrameGen.scala (2)

• DataFrameGen (39-177)
• gen (41-57)

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

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

• JoinUtils (39-532)
• getRangeToFill (137-166)

spark/src/main/scala/ai/chronon/spark/catalog/TableUtils.scala (1)

• partitions (131-156)

api/src/main/scala/ai/chronon/api/DataRange.scala (1)

• partitions (90-95)

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

• GitHub Check: streaming_tests
• GitHub Check: join_tests
• GitHub Check: streaming_tests
• GitHub Check: groupby_tests
• GitHub Check: join_tests
• GitHub Check: fetcher_tests
• GitHub Check: groupby_tests
• GitHub Check: analyzer_tests
• GitHub Check: fetcher_tests
• GitHub Check: spark_tests
• GitHub Check: analyzer_tests
• GitHub Check: batch_tests
• GitHub Check: spark_tests
• GitHub Check: batch_tests
• GitHub Check: scala_compile_fmt_fix

🔇 Additional comments (15)

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

267-267: Updated method name to align with codebase standardization.

Method name changed from getRangesToFill to getRangeToFill for consistency across the codebase.

spark/src/test/scala/ai/chronon/spark/test/join/JoinUtilsTest.scala (2)

288-288: Method name updated to match implementation.

Updated test to use renamed method getRangeToFill instead of getRangesToFill.

---

305-305: Method name updated to match implementation.

Updated test with override parameter to use renamed method getRangeToFill.

spark/src/main/scala/ai/chronon/spark/catalog/TableUtils.scala (1)

336-342: Improved error handling using require instead of assert.

Changed from assertion to requirement validation for better runtime behavior. This is more idiomatic Scala and ensures validation happens even in production environments where assertions might be disabled.

spark/src/test/scala/ai/chronon/spark/test/LocalTableExporterTest.scala (2)

92-93: Added explicit partition formatting for test data.

Moved TableUtils initialization earlier and explicitly passed partition information to the test data generator for consistency.

---

105-106: Consistent partition formatting for weight test data.

Using the same TableUtils instance to ensure consistent partition formatting across test data generation.

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

137-137: Method renaming from getRangesToFill to getRangeToFill.

The rename provides clearer indication that the method returns a single range.

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

134-138: Method call updated to match renamed method.

The method call has been correctly updated to match the renamed method in JoinUtils.

spark/src/test/scala/ai/chronon/spark/test/join/JoinTest.scala (3)

316-319: Added partition format support to DataFrameGen call.

Good addition of explicit partition format for consistent test behavior.

---

321-324: Added partition format to query.

Ensures partition format consistency between data generation and query execution.

---

366-370: Uses deep copy for join configuration.

Improved test to properly clone configuration and use future date for testing.

spark/src/test/scala/ai/chronon/spark/test/DataFrameGen.scala (4)

41-49: Enhanced gen method with partition format support.

Good implementation for flexible partition format handling.

---

64-71: Updated events method to support partition format.

Properly handles partition format in timestamp conversion.

---

79-86: Updated entities method with partition format support.

Consistent implementation across generator methods.

---

108-112: Updated mutations method to pass partition format.

Ensures format consistency in mutation test data generation.