Add support for Delta lake tables in TableUtils #869

piyush-zlai · 2024-10-21T13:59:06Z

Summary

Add delta lake support in Chronon. This PR refactors TableUtils a bit to allow us to support Delta Lake (and other table formats in the future like Hudi). The crux of the PR is the refactor of TableUtils - I've pulled out some of the format specific stuff into per-format objects (e.g. Hive / Iceberg / Delta) to help us cover aspects like creating tables, listing partitions.

As the versions of the various libraries we build against in Chronon is on the older side (e.g. delta lake - 2.0.x for Spark 3.2), we've structured the code to allow users to override their format providers. This allows folks say on Delta 3.x (Spark 3.5) to write their own Delta3xFormatProvider that is built against the newer version and configure it via the spark.chronon.table.format_provider Spark config setting. For users on the same versions / rails as OSS, they can skip configuring this and we'll default to the DefaultFormatProvider (which maintains the existing Iceberg / Hive behavior along with support for delta).

Notes on testing:

TableUtilsTest is un-changed. This test covers a bunch of Hive related tests (like creating tables, inserting partitions) and some format agnostic stuff (like schema evolution type of checks).
Added a TableUtilsFormat test - for now this covers testing Delta. Can add Iceberg in a follow up (as Iceberg format support is currently untested in the project).
Unfortunately the way we create Spark sessions in not ideal in the Chronon project. We kick off all our tests in parallel at the same time in CI (same JVM) and we end up triggering code like: SchemaEvolutionTest where we are creating the SparkSession at the class object construction time. This results in many tests vying to create sessions at the same time - most with the same spark settings but some like the new delta format test with different settings (using the delta catalog, delta spark session extension). Thanks to this we're not able to run the delta tests in the test JVM as the others AND we need to ensure that the other test spark sessions don't end up getting created with the wrong settings as that will cause the delta test to fail.
The potentially nice way to handle this would be to refactor all our existing spark tests to postpone session creation till the test body and handle session cleanup etc but as this is a fairly involved / risky change, I've chosen to instead use an env var to have the Spark session builder choose the right delta / non-delta options. It's not the cleanest but it does minimize the blast radius of these changes considerably.

Why / Goal

Allow folks keen on experimenting with Chronon that use delta lake to do so. Set the stage for supporting additional table formats in the future.

Test Plan

Added Unit Tests
Covered by existing CI
Integration tested

@mickjermsurawong-openai was also able to pull these changes and test things out using delta lake tables. We were able to sort out issues with Delta 3.x (as they're on Spark 3.5) using the provider interface.

Checklist

Documentation update

Reviewers

mickjermsurawong-openai · 2024-10-21T16:35:59Z

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

+    sparkSession.sqlContext
+      .sql(s"SHOW PARTITIONS $tableName")
+      .collect()
+      .map(row => parseHivePartition(row.getString(0)))


Curious if the partition keys here are enforced by Chronon? i.e I'm wondering if users can specified partition key and we have increasing cardinality of this and we end up returning large entries of this

the partition keys chosen / used here are based on the partitionColumn that you configure in the table util setup for your job(s). You are right that you could have a very large number of partitions here loaded up in mem.

piyush-zlai · 2024-10-28T13:53:11Z

@nikhilsimha / @mears-stripe / @pengyu-hou - would you folks have cycles to take a look?
@mickjermsurawong-openai - feel free to give this a spin as well.

.circleci/config.yml

pengyu-hou · 2024-11-08T04:52:15Z

spark/src/main/scala/ai/chronon/spark/SparkSessionBuilder.scala

+    // allow us to override the format by specifying env vars. This allows us to not have to worry about interference
+    // between Spark sessions created in existing chronon tests that need the hive format and some specific tests
+    // that require a format override like delta lake.
+    val (formatConfigs, kryoRegistrator) = sys.env.get(FormatTestEnvVar) match {


looks like this is only used for local testing. Could we move it within the if (local) block?

We could - I actually thought the code is cleaner with the match up front - this allows us to skip the use of vars and a lot of if/else type of checks in a couple of places.
Let me know if you feel strongly and I can add there.

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

mears-stripe

Excited for this change!

add thrift 18 using user format foo more log add log lines morel ong lines only use delta format need to commit try add duck wip wip add log line provider revert thrift

pengyu-hou · 2024-11-15T20:16:21Z

spark/src/test/scala/ai/chronon/spark/test/TableUtilsFormatTest.scala

+
+import scala.util.Try
+
+class TableUtilsFormatTest {


@piyush-zlai If this test specific to delta lake?

Is it reasonable to add other formats like Hive or Iceberg to the CI?

This is largely agnostic to the format. As part of the regular "sbt test" run as we don't specify the format, this suite gets run for the default (hive).
I'll circle back with Iceberg tests in a follow up - For that I might need to wire up the iceberg spark extensions etc.

pengyu-hou

Thanks for the changes.
Looking forward to more changes from zipline-ai!

## Summary Port of our OSS delta lake PR - airbnb/chronon#869. Largely the same aside from delta lake versions. We don't need this immediately atm but we'll need this if we have other users come along that need delta lake (or we need to add support for formats like hudi) ## Checklist - [X] Added Unit Tests - [X] Covered by existing CI - [ ] Integration tested - [ ] Documentation update  ## Summary by CodeRabbit - **New Features** - Added support for Delta Lake operations with new dependencies and configurations. - Introduced new traits and case objects for handling different table formats, enhancing data management capabilities. - Added a new job in the CI workflow for testing Delta Lake format functionality. - **Bug Fixes** - Improved error handling in class registration processes. - **Tests** - Implemented a suite of unit tests for the `TableUtils` class to validate partitioned data insertions with schema modifications.

## Summary Port of our OSS delta lake PR - airbnb/chronon#869. Largely the same aside from delta lake versions. We don't need this immediately atm but we'll need this if we have other users come along that need delta lake (or we need to add support for formats like hudi) ## Cheour clientslist - [X] Added Unit Tests - [X] Covered by existing CI - [ ] Integration tested - [ ] Documentation update  ## Summary by CodeRabbit - **New Features** - Added support for Delta Lake operations with new dependencies and configurations. - Introduced new traits and case objects for handling different table formats, enhancing data management capabilities. - Added a new job in the CI workflow for testing Delta Lake format functionality. - **Bug Fixes** - Improved error handling in class registration processes. - **Tests** - Implemented a suite of unit tests for the `TableUtils` class to validate partitioned data insertions with schema modifications.

mickjermsurawong-openai reviewed Oct 21, 2024

View reviewed changes

piyush-zlai force-pushed the piyush--deltalake branch 2 times, most recently from 620a789 to df4f9cc Compare October 27, 2024 21:14

piyush-zlai changed the title ~~[wip] Add support for Delta lake tables in TableUtils~~ Add support for Delta lake tables in TableUtils Oct 28, 2024

mickjermsurawong-openai force-pushed the piyush--deltalake branch from c1df43b to 87e0986 Compare November 5, 2024 11:24

piyush-zlai force-pushed the piyush--deltalake branch from 1fffb45 to 7225a69 Compare November 7, 2024 16:22

pengyu-hou reviewed Nov 8, 2024

View reviewed changes

mears-stripe requested changes Nov 12, 2024

View reviewed changes

spark/src/main/scala/ai/chronon/spark/TableUtils.scala Outdated Show resolved Hide resolved

spark/src/main/scala/ai/chronon/spark/TableUtils.scala Outdated Show resolved Hide resolved

mears-stripe reviewed Nov 12, 2024

View reviewed changes

piyush-zlai added 20 commits November 12, 2024 11:37

First stab at adding support for Delta lake tables in TableUtils

c3aae9b

Scala fmt

1a64baa

Update versions and code to work across older delta versions

5e134ed

Add working delta tests

cd801ca

Add logging

51ae067

Unset confs

6a326a5

Scala fmt

c72ef70

Tweak logs

02dda37

Try unsetting spark context

fe4cc83

Swap things around

d22525d

Tweak how we infer write format

81dd918

Play with test configs

2ff3ed5

Revert circle ci sbt and ignore tests for now

f3a26ba

Try commenting kryo code to see if that makes tests happy

b7b084a

Trigger Build

ad14e67

Switch to env var driven approach

f15cc5b

Shuffle around kryo registrator

7a9396e

Add missing circle ci workflows

f93d83e

Drop the sealed and enable tests

48a7604

Trigger Build

837c943

piyush-zlai and others added 12 commits November 12, 2024 11:37

Switch table utils format test to junit tests

2dd5a33

Enable one test

e24a9d6

Enable more tests

1d36d63

support format provider

5c8ed17

add thrift 18 using user format foo more log add log lines morel ong lines only use delta format need to commit try add duck wip wip add log line provider revert thrift

Scala fmt

73577af

Add default format provider

46748b6

Add read and write format providers

b3cb5a5

Unify table format trait

aa76d05

Combine format provider config setting

3066719

Add serializable

2b381c2

Address comments

6affb63

Refactor on comments

e3d6632

piyush-zlai force-pushed the piyush--deltalake branch from 087f81b to e3d6632 Compare November 12, 2024 22:10

scala fmt fixes

48cccf1

piyush-zlai force-pushed the piyush--deltalake branch from ae010c1 to 268f765 Compare November 14, 2024 23:39

Change db name

5673ac4

piyush-zlai force-pushed the piyush--deltalake branch from 268f765 to 5673ac4 Compare November 14, 2024 23:50

Trigger Build

1af4545

pengyu-hou reviewed Nov 15, 2024

View reviewed changes

mears-stripe approved these changes Nov 18, 2024

View reviewed changes

pengyu-hou approved these changes Nov 19, 2024

View reviewed changes

caiocamatta-stripe merged commit b14f886 into airbnb:main Nov 19, 2024
8 checks passed

piyush-zlai mentioned this pull request Nov 26, 2024

Add delta lake integration zipline-ai/chronon#51

Merged

4 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add support for Delta lake tables in TableUtils #869

Add support for Delta lake tables in TableUtils #869

Uh oh!

piyush-zlai commented Oct 21, 2024 •

edited

Loading

Uh oh!

mickjermsurawong-openai Oct 21, 2024

Uh oh!

piyush-zlai Oct 21, 2024

Uh oh!

piyush-zlai commented Oct 28, 2024

Uh oh!

Uh oh!

pengyu-hou Nov 8, 2024

Uh oh!

piyush-zlai Nov 8, 2024

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mears-stripe left a comment

Uh oh!

pengyu-hou Nov 15, 2024

Uh oh!

piyush-zlai Nov 15, 2024

Uh oh!

pengyu-hou left a comment

Uh oh!

Uh oh!

Uh oh!

Add support for Delta lake tables in TableUtils #869

Add support for Delta lake tables in TableUtils #869

Uh oh!

Conversation

piyush-zlai commented Oct 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Why / Goal

Test Plan

Checklist

Reviewers

Uh oh!

mickjermsurawong-openai Oct 21, 2024

Choose a reason for hiding this comment

Uh oh!

piyush-zlai Oct 21, 2024

Choose a reason for hiding this comment

Uh oh!

piyush-zlai commented Oct 28, 2024

Uh oh!

Uh oh!

pengyu-hou Nov 8, 2024

Choose a reason for hiding this comment

Uh oh!

piyush-zlai Nov 8, 2024

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mears-stripe left a comment

Choose a reason for hiding this comment

Uh oh!

pengyu-hou Nov 15, 2024

Choose a reason for hiding this comment

Uh oh!

piyush-zlai Nov 15, 2024

Choose a reason for hiding this comment

Uh oh!

pengyu-hou left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

piyush-zlai commented Oct 21, 2024 •

edited

Loading