Allow setting partition column name in sources #381
Conversation
WalkthroughThis pull request introduces a new optional Changes
Possibly related PRs
Suggested reviewers
Poem
Warning Review ran into problems🔥 ProblemsGitHub Actions and Pipeline Checks: Resource not accessible by integration - https://docs.github.com/rest/actions/workflow-runs#list-workflow-runs-for-a-repository. Please grant the required permissions to the CodeRabbit GitHub App under the organization or repository settings. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
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)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
spark/src/main/scala/ai/chronon/spark/LabelJoin.scala (1)
195-197: LGTM! Good replacement of null with explicit Query.Improves code clarity by removing null and being explicit about partition column.
Consider adding a comment explaining the significance of
Constants.LabelPartitionColumn.spark/src/main/scala/ai/chronon/spark/Extensions.scala (1)
394-397: Consider simplifying the Option nesting.The nested Options could be flattened for better readability.
- def effectivePartitionColumn(implicit tableUtils: TableUtils): String = - Option(query).flatMap(q => Option(q.partitionColumn)).getOrElse(tableUtils.partitionColumn) + def effectivePartitionColumn(implicit tableUtils: TableUtils): String = + Option(query.partitionColumn).getOrElse(tableUtils.partitionColumn)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro (Legacy)
📒 Files selected for processing (10)
api/src/main/scala/ai/chronon/api/Builders.scala(2 hunks)api/thrift/api.thrift(2 hunks)spark/src/main/scala/ai/chronon/spark/Analyzer.scala(3 hunks)spark/src/main/scala/ai/chronon/spark/Extensions.scala(2 hunks)spark/src/main/scala/ai/chronon/spark/GroupBy.scala(1 hunks)spark/src/main/scala/ai/chronon/spark/JoinBase.scala(4 hunks)spark/src/main/scala/ai/chronon/spark/LabelJoin.scala(2 hunks)spark/src/main/scala/ai/chronon/spark/TableUtils.scala(4 hunks)spark/src/test/scala/ai/chronon/spark/test/DataFrameGen.scala(1 hunks)spark/src/test/scala/ai/chronon/spark/test/JoinTest.scala(28 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (4)
- GitHub Check: non_spark_tests
- GitHub Check: spark_tests
- GitHub Check: scala_compile_fmt_fix
- GitHub Check: enforce_triggered_workflows
🔇 Additional comments (20)
spark/src/test/scala/ai/chronon/spark/test/DataFrameGen.scala (3)
41-44: LGTM! Clean addition of partition column parameter.The optional parameter with default value maintains backward compatibility while adding flexibility.
52-61: LGTM! Good extraction of partition column string.Early extraction of
partitionColumnStringimproves code clarity.
65-71: LGTM! Consistent handling of partition column.Properly propagates partition column to both local usage and downstream call.
api/src/main/scala/ai/chronon/api/Builders.scala (1)
49-63: LGTM! Clean addition to Query builder.Maintains builder pattern while adding partition column support.
spark/src/main/scala/ai/chronon/spark/Extensions.scala (1)
387-392: LGTM! Clean implementation of Source extension.Good use of Option for null safety in partition column handling.
api/thrift/api.thrift (2)
19-19: LGTM! Added optional partition column support to Query struct.The new field enhances partition handling while maintaining backward compatibility.
53-56: LGTM! Added optional partition column support to StagingQuery struct with clear documentation.The field's purpose for
max_datetemplate is well documented.spark/src/main/scala/ai/chronon/spark/JoinBase.scala (3)
54-54: LGTM! Added implicit tableUtils for better code organization.Reduces parameter passing and improves readability.
173-174: LGTM! Added partition column handling to unfilledRanges call.Correctly propagates the effective partition column from join configuration.
488-489: LGTM! Added partition column handling to unfilledRanges call.Correctly propagates the effective partition column from join configuration.
spark/src/main/scala/ai/chronon/spark/Analyzer.scala (3)
35-35: LGTM! Added QuerySparkOps import.Provides access to query-related extensions.
94-94: LGTM! Added implicit tableUtils for better code organization.Reduces parameter passing and improves readability.
316-319: LGTM! Added partition column handling to unfilledRanges call.Correctly propagates the effective partition column from join configuration.
spark/src/main/scala/ai/chronon/spark/GroupBy.scala (1)
629-629: LGTM! Updated metaColumns to use source query's partition column.Correctly uses the actual partition column from the source query instead of null.
spark/src/main/scala/ai/chronon/spark/TableUtils.scala (3)
202-204: LGTM! Good addition of the partition column parameter.The optional parameter with a default value maintains backward compatibility.
597-603: LGTM! Good propagation of partition column parameter.The optional parameter enables flexible partition column handling in unfilled ranges computation.
836-850: LGTM! Good handling of partition column renaming.The code correctly handles partition column mapping between query and table.
spark/src/test/scala/ai/chronon/spark/test/JoinTest.scala (3)
129-130: Good test data optimization.Reduced entity count from 300/50 to improve test execution time.
434-556: LGTM! Good test coverage for partition column feature.Comprehensive test case validating custom partition column functionality.
189-191: LGTM! Good query configuration update.Added partition column to join configuration.
varant-zlai
force-pushed
the
vz--partition_column_in_source_v2
branch
from
0207e63 to
2a05b64
Compare
| @@ -265,16 +260,25 @@ class LabelJoin(joinConf: api.Join, tableUtils: TableUtils, labelDS: String) { | |||
| leftDf | |||
| } | |||
|
|
|||
| val updatedRight = | |||
Co-authored-by: Nikhil Simha <[email protected]>
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
spark/src/test/scala/ai/chronon/spark/test/bootstrap/TableBootstrapTest.scala (1)
116-116: Remove commented out code.Clean up unused code to improve maintainability.
-//val bootstrapDf2 = bootstrapDf2Partition.withColumnRenamed("date", "ds")
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro (Legacy)
📒 Files selected for processing (4)
spark/src/main/scala/ai/chronon/spark/BootstrapInfo.scala(3 hunks)spark/src/main/scala/ai/chronon/spark/Join.scala(1 hunks)spark/src/test/scala/ai/chronon/spark/test/JoinTest.scala(27 hunks)spark/src/test/scala/ai/chronon/spark/test/bootstrap/TableBootstrapTest.scala(3 hunks)
⏰ 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 (9)
spark/src/main/scala/ai/chronon/spark/BootstrapInfo.scala (2)
89-89: LGTM!Clean implicit parameter addition for dependency injection.
187-191: LGTM!Correctly uses
effectivePartitionColumnand improves error message with schema context.spark/src/main/scala/ai/chronon/spark/Join.scala (1)
518-518: LGTM!Correctly uses
effectivePartitionColumnfor partition column customization.spark/src/test/scala/ai/chronon/spark/test/bootstrap/TableBootstrapTest.scala (2)
47-48: LGTM! Well-structured partition column handling.The changes correctly implement partition column name customization while maintaining backward compatibility.
Also applies to: 59-59, 67-69, 75-76, 81-81
113-114: LGTM! Good test coverage.Test case effectively validates custom partition column functionality.
spark/src/test/scala/ai/chronon/spark/test/JoinTest.scala (4)
128-129: LGTM! Improved test performance and partition handling.Reduced test data size and added proper partition column support.
Also applies to: 176-177
229-236: LGTM! SQL query properly updated.Query correctly handles custom partition column.
433-555: LGTM! Comprehensive test for partition columns.Test thoroughly validates the new partition column functionality.
565-565: LGTM! Consistent test optimization.Test data size reduced across all test cases for better performance.
Also applies to: 863-863, 864-864, 879-880, 887-887, 1073-1074, 1139-1140, 1160-1161, 1269-1270, 1290-1291, 1326-1327, 1339-1340, 712-713
## Summary - #381 introduced the ability to configure a partition column at the node-level. This PR simply fixes a missed spot on the plumbing of the new StagingQuery attribute. ## Checklist - [ ] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enhanced the query builder to support specifying a partition column, providing greater customization for query formation and partitioning. - **Improvements** - Improved handling of partition columns by introducing a fallback mechanism to ensure valid values are used when necessary. <!-- end of auto-generated comment: release notes by coderabbit.ai --> <!-- av pr metadata This information is embedded by the av CLI when creating PRs to track the status of stacks when using Aviator. Please do not delete or edit this section of the PR. ``` {"parent":"main","parentHead":"","trunk":"main"} ``` --> --------- Co-authored-by: Thomas Chow <[email protected]>
## Summary Allow setting partition column name in sources. Maps it to the default partition name upon read and partition checking. ## Checklist - [x] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enabled configurable partition columns in query, join, and data generation operations for improved data partitioning. - **Refactor** - Streamlined partition handling and consolidated import structures to enhance workflow efficiency. - **Tests** - Added test cases for verifying partition column functionality and adjusted data generation volumes for better validation. - Introduced new tests specifically for different partition columns to ensure accurate handling of partitioned data. These enhancements provide increased flexibility and accuracy in managing partitioned datasets during data processing and join operations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: ezvz <[email protected]> Co-authored-by: Nikhil Simha <[email protected]>
## Summary - #381 introduced the ability to configure a partition column at the node-level. This PR simply fixes a missed spot on the plumbing of the new StagingQuery attribute. ## Checklist - [ ] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enhanced the query builder to support specifying a partition column, providing greater customization for query formation and partitioning. - **Improvements** - Improved handling of partition columns by introducing a fallback mechanism to ensure valid values are used when necessary. <!-- end of auto-generated comment: release notes by coderabbit.ai --> <!-- av pr metadata This information is embedded by the av CLI when creating PRs to track the status of stacks when using Aviator. Please do not delete or edit this section of the PR. ``` {"parent":"main","parentHead":"","trunk":"main"} ``` --> --------- Co-authored-by: Thomas Chow <[email protected]>
## Summary Allow setting partition column name in sources. Maps it to the default partition name upon read and partition checking. ## Checklist - [x] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enabled configurable partition columns in query, join, and data generation operations for improved data partitioning. - **Refactor** - Streamlined partition handling and consolidated import structures to enhance workflow efficiency. - **Tests** - Added test cases for verifying partition column functionality and adjusted data generation volumes for better validation. - Introduced new tests specifically for different partition columns to ensure accurate handling of partitioned data. These enhancements provide increased flexibility and accuracy in managing partitioned datasets during data processing and join operations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: ezvz <[email protected]> Co-authored-by: Nikhil Simha <[email protected]>
## Summary - #381 introduced the ability to configure a partition column at the node-level. This PR simply fixes a missed spot on the plumbing of the new StagingQuery attribute. ## Checklist - [ ] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enhanced the query builder to support specifying a partition column, providing greater customization for query formation and partitioning. - **Improvements** - Improved handling of partition columns by introducing a fallback mechanism to ensure valid values are used when necessary. <!-- end of auto-generated comment: release notes by coderabbit.ai --> <!-- av pr metadata This information is embedded by the av CLI when creating PRs to track the status of stacks when using Aviator. Please do not delete or edit this section of the PR. ``` {"parent":"main","parentHead":"","trunk":"main"} ``` --> --------- Co-authored-by: Thomas Chow <[email protected]>
## Summary Allow setting partition column name in sources. Maps it to the default partition name upon read and partition checking. ## Checklist - [x] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enabled configurable partition columns in query, join, and data generation operations for improved data partitioning. - **Refactor** - Streamlined partition handling and consolidated import structures to enhance workflow efficiency. - **Tests** - Added test cases for verifying partition column functionality and adjusted data generation volumes for better validation. - Introduced new tests specifically for different partition columns to ensure accurate handling of partitioned data. These enhancements provide increased flexibility and accuracy in managing partitioned datasets during data processing and join operations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: ezvz <[email protected]> Co-authored-by: Nikhil Simha <[email protected]>
## Summary - #381 introduced the ability to configure a partition column at the node-level. This PR simply fixes a missed spot on the plumbing of the new StagingQuery attribute. ## Checklist - [ ] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enhanced the query builder to support specifying a partition column, providing greater customization for query formation and partitioning. - **Improvements** - Improved handling of partition columns by introducing a fallback mechanism to ensure valid values are used when necessary. <!-- end of auto-generated comment: release notes by coderabbit.ai --> <!-- av pr metadata This information is embedded by the av CLI when creating PRs to track the status of stacks when using Aviator. Please do not delete or edit this section of the PR. ``` {"parent":"main","parentHead":"","trunk":"main"} ``` --> --------- Co-authored-by: Thomas Chow <[email protected]>
## Summary Allow setting partition column name in sources. Maps it to the default partition name upon read and partition checking. ## Checklist - [x] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enabled configurable partition columns in query, join, and data generation operations for improved data partitioning. - **Refactor** - Streamlined partition handling and consolidated import structures to enhance workflow efficiency. - **Tests** - Added test cases for verifying partition column functionality and adjusted data generation volumes for better validation. - Introduced new tests specifically for different partition columns to ensure accurate handling of partitioned data. These enhancements provide increased flexibility and accuracy in managing partitioned datasets during data processing and join operations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: ezvz <[email protected]> Co-authored-by: Nikhil Simha <[email protected]>
## Summary - #381 introduced the ability to configure a partition column at the node-level. This PR simply fixes a missed spot on the plumbing of the new StagingQuery attribute. ## Checklist - [ ] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enhanced the query builder to support specifying a partition column, providing greater customization for query formation and partitioning. - **Improvements** - Improved handling of partition columns by introducing a fallback mechanism to ensure valid values are used when necessary. <!-- end of auto-generated comment: release notes by coderabbit.ai --> <!-- av pr metadata This information is embedded by the av CLI when creating PRs to track the status of stacks when using Aviator. Please do not delete or edit this section of the PR. ``` {"parent":"main","parentHead":"","trunk":"main"} ``` --> --------- Co-authored-by: Thomas Chow <[email protected]>
## Summary Allow setting partition column name in sources. Maps it to the default partition name upon read and partition cheour clientsing. ## Cheour clientslist - [x] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enabled configurable partition columns in query, join, and data generation operations for improved data partitioning. - **Refactor** - Streamlined partition handling and consolidated import structures to enhance workflow efficiency. - **Tests** - Added test cases for verifying partition column functionality and adjusted data generation volumes for better validation. - Introduced new tests specifically for different partition columns to ensure accurate handling of partitioned data. These enhancements provide increased flexibility and accuracy in managing partitioned datasets during data processing and join operations. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: ezvz <[email protected]> Co-authored-by: Nikhil Simha <[email protected]>
## Summary - #381 introduced the ability to configure a partition column at the node-level. This PR simply fixes a missed spot on the plumbing of the new StagingQuery attribute. ## Cheour clientslist - [ ] Added Unit Tests - [ ] Covered by existing CI - [ ] Integration tested - [ ] Documentation update <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Enhanced the query builder to support specifying a partition column, providing greater customization for query formation and partitioning. - **Improvements** - Improved handling of partition columns by introducing a fallbaour clients mechanism to ensure valid values are used when necessary. <!-- end of auto-generated comment: release notes by coderabbit.ai --> <!-- av pr metadata This information is embedded by the av CLI when creating PRs to traour clients the status of staour clientss when using Aviator. Please do not delete or edit this section of the PR. ``` {"parent":"main","parentHead":"","trunk":"main"} ``` --> --------- Co-authored-by: Thomas Chow <[email protected]>
Summary
Allow setting partition column name in sources. Maps it to the default partition name upon read and partition checking.
Checklist
Summary by CodeRabbit
These enhancements provide increased flexibility and accuracy in managing partitioned datasets during data processing and join operations.