Allow using key columns as input to derivations

[xiaohui-sun]

Summary

Add support for using key columns as input to derivations.

Why / Goal

Test Plan

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

Checklist

• Documentation update

Reviewers

@hzding621
@yuli-han

[yuli-han]

Could you try to test the fetcher behavior on a gateway? Usually the step is:

1. update an existing testing join derivation with derivation on key in the zipline_test folder
2. run metadata-upload to upload the new join schema to k-v store metadata store
3. upload the local jar to a remote gateway with the new change.
4. run fetcher cli to validate the change
   This is the procedure I am running when changing the fetcher, cc @hzding621 @pengyu-hou could you pls share if you know any simpler way to test or if there are existing doc for the testing procedure? Also pls lemme know if you need the command to run the steps.

[hzding621]

+1 to fetchercli testing as to what @yuli-han said

[xiaohui-sun]

Tested with the following example:

Added a rename only derivation to group by.
https://git.musta.ch/airbnb/ml_models/compare/0228_test?expand=1#diff-861ca6d9862bdc850d3de0cb56d324078afd9b0ef4ce23ff54097991d0a480fbR18-R20

===== Before Fix =======

===== Query =====
SPARK_VERSION=3.2.1 python3 ~/.local/bin/run.py --mode=fetch --name zipline_test/iceberg_realtime.v0 -k '{"username":"svc_di_trust"}' -t join --chronon-jar ~/test/chronon-embedded.jar | tee ~/test/fetch.log

===== Query Result =====
2025-02-28 21:42:21 ERROR Fetcher:192 - Exception Message: java.lang.RuntimeException: Couldn't fetch joinName = zipline_test/iceberg_realtime.v0 or build join codec due to java.lang.IllegalArgumentException: Failed to run expression username for username_renamed. Please ensure the derivation is correct.
at ai.chronon.online.OnlineDerivationUtil$.$anonfun$buildDerivedFields$2(OnlineDerivationUtil.scala:110)
at scala.collection.immutable.List.map(List.scala:246)
at ai.chronon.online.OnlineDerivationUtil$.buildDerivedFields(OnlineDerivationUtil.scala:105)
at ai.chronon.online.JoinCodec.valueSchema$lzycompute(JoinCodec.scala:38)
at ai.chronon.online.JoinCodec.valueSchema(JoinCodec.scala:34)
at ai.chronon.online.JoinCodec.(JoinCodec.scala:78)
at ai.chronon.online.Fetcher.buildJoinCodec(Fetcher.scala:170)
at ai.chronon.online.Fetcher.$anonfun$getJoinCodecs$3(Fetcher.scala:180)

===== After Fix =======

===== Query =====
SPARK_VERSION=3.2.1 python3 ~/.local/bin/run.py --mode=fetch --name zipline_test/iceberg_realtime.v0 -k '{"username":"svc_di_trust"}' -t join --chronon-jar ~/test/chronon-embedded.jar | tee ~/test/fetch.log

===== Query Result =====

2025-02-28 21:37:15 INFO Driver$FetcherCli$:750 - --- [FETCHED RESULT] ---
{
"username_renamed" : "svc_di_trust"
}
2025-02-28 21:37:15 INFO Driver$FetcherCli$:752 - Fetched in: 773.667865 ms