[Spark][Optimization] Modify MergeJob to implement storage partition join using row_id

api/python/ai/chronon/join.py

@@ -356,7 +356,7 @@ def Join(
     left: api.Source,
     right_parts: List[api.JoinPart],
     version: int,
-    row_ids: Union[str, List[str]],
+    row_ids: Union[str, List[str]] = None,
     online_external_parts: List[api.ExternalPart] = None,
     bootstrap_parts: List[api.BootstrapPart] = None,
     bootstrap_from_log: bool = False,
@@ -478,7 +478,10 @@ def Join(
     # create a deep copy for case: multiple LeftOuterJoin use the same left,
     # validation will fail after the first iteration
     updated_left = copy.deepcopy(left)
-    if left.events and left.events.query.selects:
+
+    selects = None
+    if left.events:
+        selects = left.events.query.selects
         assert "ts" not in left.events.query.selects.keys(), (
             "'ts' is a reserved key word for Chronon,"
             " please specify the expression in timeColumn"
@@ -487,6 +490,13 @@ def Join(
         updated_left.events.query.selects.update(
             {"ts": updated_left.events.query.timeColumn}
         )
+    elif left.entities:
+        selects = left.entities.query.selects
+
+    if selects:
+        # For JoinSource, we can rely on the validation at the base join level
+        # TODO add more docs about row ID and link here
+        assert "row_id" in selects, "Left side of the join must contain `row_id` as a column."
 
     if label_part:
         label_metadata = api.MetaData(

api/python/ai/chronon/query.py

@@ -19,7 +19,7 @@
 
 
 def Query(
-    selects: Dict[str, str] = None,
+    selects: Dict[str, str],
     wheres: List[str] = None,
     start_partition: str = None,
     end_partition: str = None,

api/python/test/canary/compiled/joins/gcp/item_event_join.canary_batch_v1__0

@@ -17,6 +17,7 @@
       "user_id_purchase_price_last10": "5609e04d61a47a8cafc67970785a3a59",
       "listing_id": "f2f6c814b8ae1521176b22a1ae7f2d0d",
       "user_id": "493d3df28f80664abd11b19fcd33b6e6",
+      "row_id": "6c80b6474532731fdb20d0a43bcb287f",
       "ts": "ad9fd4c611e20ad833819a4ce9d752bf"
     },
     "online": 1,
@@ -144,6 +145,7 @@
         "selects": {
           "listing_id": "EXPLODE(TRANSFORM(SPLIT(COALESCE(attributes['sold_listing_ids'], attributes['listing_id']), ','), e -> CAST(e AS LONG)))",
           "user_id": "attributes['user_id']",
+          "row_id": "request_UUID",
           "ts": "timestamp"
         },
         "timeColumn": "timestamp"

api/python/test/canary/compiled/joins/gcp/item_event_join.canary_combined_v1__0

@@ -21,6 +21,7 @@
       "user_id_purchase_price_last10": "5609e04d61a47a8cafc67970785a3a59",
       "listing_id": "f2f6c814b8ae1521176b22a1ae7f2d0d",
       "user_id": "493d3df28f80664abd11b19fcd33b6e6",
+      "row_id": "6c80b6474532731fdb20d0a43bcb287f",
       "ts": "ad9fd4c611e20ad833819a4ce9d752bf"
     },
     "online": 1,
@@ -148,6 +149,7 @@
         "selects": {
           "listing_id": "EXPLODE(TRANSFORM(SPLIT(COALESCE(attributes['sold_listing_ids'], attributes['listing_id']), ','), e -> CAST(e AS LONG)))",
           "user_id": "attributes['user_id']",
+          "row_id": "request_UUID",
           "ts": "timestamp"
         },
         "timeColumn": "timestamp"

api/python/test/canary/compiled/joins/gcp/item_event_join.canary_streaming_v1__0

@@ -11,6 +11,7 @@
       "listing_id_favorite_sum_1d": "bb07e39a4b7ea81aca5baf027fd8f555",
       "listing_id": "f2f6c814b8ae1521176b22a1ae7f2d0d",
       "user_id": "493d3df28f80664abd11b19fcd33b6e6",
+      "row_id": "6c80b6474532731fdb20d0a43bcb287f",
       "ts": "ad9fd4c611e20ad833819a4ce9d752bf"
     },
     "online": 1,
@@ -138,6 +139,7 @@
         "selects": {
           "listing_id": "EXPLODE(TRANSFORM(SPLIT(COALESCE(attributes['sold_listing_ids'], attributes['listing_id']), ','), e -> CAST(e AS LONG)))",
           "user_id": "attributes['user_id']",
+          "row_id": "request_UUID",
           "ts": "timestamp"
         },
         "timeColumn": "timestamp"
@@ -251,8 +253,5 @@
       "useLongNames": 0
     }
   ],
-  "rowIds": [
-    "user_id"
-  ],
   "useLongNames": 0
 }

api/python/test/canary/compiled/joins/gcp/training_set.v1_dev__0

@@ -16,6 +16,7 @@
       "user_id_purchase_price_average_7d": "e3351cc3265d17cfd1b99f8da7fb083d",
       "user_id_purchase_price_last10": "e3351cc3265d17cfd1b99f8da7fb083d",
       "user_id": "a673c80ffc3260ee5ca173869ba07f7c",
+      "row_id": "cd335c9bb6f12c6e71b482b33109e07c",
       "ts": "a173b2adccfe3c386f262c4cd7ffe9bd"
     },
     "online": 0,
@@ -142,6 +143,7 @@
       "query": {
         "selects": {
           "user_id": "user_id",
+          "row_id": "request_UUID",
           "ts": "ts"
         },
         "startPartition": "2023-11-01",
@@ -265,8 +267,5 @@
       "useLongNames": 0
     }
   ],
-  "rowIds": [
-    "user_id"
-  ],
   "useLongNames": 0
 }

api/python/test/canary/compiled/joins/gcp/training_set.v1_dev_notds__0

@@ -16,6 +16,7 @@
       "user_id_purchase_price_average_7d": "1d85a493771d48ac9fb6581c18f186e2",
       "user_id_purchase_price_last10": "1d85a493771d48ac9fb6581c18f186e2",
       "user_id": "fbc25deb4a5dd5d9c1d6e960c3922d0e",
+      "row_id": "70230580430bc568b5ab415e20b802dc",
       "ts": "c4364195d7639fb8b71e2f87ee017d4d"
     },
     "online": 0,
@@ -142,6 +143,7 @@
       "query": {
         "selects": {
           "user_id": "user_id",
+          "row_id": "request_UUID",
           "ts": "ts"
         },
         "timeColumn": "ts",
@@ -266,8 +268,5 @@
       "useLongNames": 0
     }
   ],
-  "rowIds": [
-    "user_id"
-  ],
   "useLongNames": 0
 }

api/python/test/canary/compiled/joins/gcp/training_set.v1_test__0

@@ -16,6 +16,7 @@
       "user_id_purchase_price_average_7d": "e3351cc3265d17cfd1b99f8da7fb083d",
       "user_id_purchase_price_last10": "e3351cc3265d17cfd1b99f8da7fb083d",
       "user_id": "a673c80ffc3260ee5ca173869ba07f7c",
+      "row_id": "cd335c9bb6f12c6e71b482b33109e07c",
       "ts": "a173b2adccfe3c386f262c4cd7ffe9bd"
     },
     "online": 0,
@@ -142,6 +143,7 @@
       "query": {
         "selects": {
           "user_id": "user_id",
+          "row_id": "request_UUID",
           "ts": "ts"
         },
         "startPartition": "2023-11-01",
@@ -265,8 +267,5 @@
       "useLongNames": 0
     }
   ],
-  "rowIds": [
-    "user_id"
-  ],
   "useLongNames": 0
 }

api/python/test/canary/compiled/joins/gcp/training_set.v1_test_notds__0

@@ -16,6 +16,7 @@
       "user_id_purchase_price_average_7d": "1d85a493771d48ac9fb6581c18f186e2",
       "user_id_purchase_price_last10": "1d85a493771d48ac9fb6581c18f186e2",
       "user_id": "fbc25deb4a5dd5d9c1d6e960c3922d0e",
+      "row_id": "70230580430bc568b5ab415e20b802dc",
       "ts": "c4364195d7639fb8b71e2f87ee017d4d"
     },
     "online": 0,
@@ -142,6 +143,7 @@
       "query": {
         "selects": {
           "user_id": "user_id",
+          "row_id": "request_UUID",
           "ts": "ts"
         },
         "timeColumn": "ts",
@@ -266,8 +268,5 @@
       "useLongNames": 0
     }
   ],
-  "rowIds": [
-    "user_id"
-  ],
   "useLongNames": 0
 }

api/python/test/canary/joins/gcp/item_event_join.py

@@ -11,6 +11,7 @@
             selects=selects(
               listing_id="EXPLODE(TRANSFORM(SPLIT(COALESCE(attributes['sold_listing_ids'], attributes['listing_id']), ','), e -> CAST(e AS LONG)))",
               user_id="attributes['user_id']",
+              row_id="request_UUID",
             ),
             time_column="timestamp",
         ),
@@ -20,7 +21,6 @@
 # Join with just a streaming GB
 canary_streaming_v1 = Join(
     left=source,
-    row_ids="user_id",
     right_parts=[
         JoinPart(group_by=item_event_canary.actions_pubsub_v2)
     ],

api/python/test/canary/joins/gcp/training_set.py

@@ -13,7 +13,8 @@
         table="data.checkouts",
         query=Query(
             selects=selects(
-                "user_id"
+                user_id="user_id",
+                row_id="request_UUID",
             ),  # The primary key used to join various GroupBys together
             start_partition="2023-11-01",
             time_column="ts",
@@ -23,7 +24,6 @@
 
 v1_test = Join(
     left=source,
-    row_ids="user_id",
     right_parts=[
         JoinPart(group_by=purchases.v1_test)
     ],
@@ -41,7 +41,6 @@
 
 v1_dev = Join(
     left=source,
-    row_ids="user_id",
     right_parts=[
         JoinPart(group_by=purchases.v1_dev)
     ],
@@ -53,7 +52,8 @@
         table="data.checkouts_notds",
         query=Query(
             selects=selects(
-                "user_id"
+                user_id = "user_id",
+                row_id="request_UUID",
             ),  # The primary key used to join various GroupBys together
             time_column="ts",
             partition_column="notds"
@@ -63,7 +63,6 @@
 
 v1_test_notds = Join(
     left=source_notds,
-    row_ids=["user_id"],
     right_parts=[
         JoinPart(group_by=purchases.v1_test_notds)
     ],
@@ -72,7 +71,6 @@
 
 v1_dev_notds = Join(
     left=source_notds,
-    row_ids=["user_id"],
     right_parts=[
         JoinPart(group_by=purchases.v1_dev_notds)
     ],

api/python/test/sample/joins/kaggle/outbrain.py

@@ -21,8 +21,7 @@
 
 training_set = Join(  # left equi join
     left=outbrain_left_events(
-        "uuid", "display_id", "ad_id", "document_id", "clicked", "geo_location", "platform"),
-    row_ids="uuid",
+        "uuid", "display_id", "ad_id", "document_id", "clicked", "geo_location", "platform", "row_id"),
     right_parts=[JoinPart(group_by=group_by) for group_by in [ad_doc, ad_uuid, ad_streaming, ad_platform]],
     use_long_names = True,
     version=0,

api/python/test/sample/joins/quickstart/training_set.py

@@ -29,7 +29,8 @@
         table="data.checkouts",
         query=Query(
             selects=selects(
-                "user_id"
+                user_id="user_id",
+                row_id="request_UUID",
             ),  # The primary key used to join various GroupBys together
             time_column="ts",
         ),  # The event time used to compute feature values as-of
@@ -38,7 +39,6 @@
 
 v1 = Join(
     left=source,
-    row_ids="user_id",
     right_parts=[
         JoinPart(group_by=group_by) for group_by in [purchases_v1, returns_v1, users]
     ],  # Include the three GroupBys
@@ -47,7 +47,6 @@
 
 v2 = Join(
     left=source,
-    row_ids=["user_id"],
     right_parts=[
         JoinPart(group_by=group_by) for group_by in [purchases_v1, returns_v1]
     ],  # Include the two online GroupBys

api/python/test/sample/joins/risk/user_transactions.py

@@ -8,13 +8,12 @@
 
 source_users = Source(
     events=EventSource(
-        table="data.users", query=Query(selects=selects("user_id"), time_column="ts")
+        table="data.users", query=Query(selects=selects("user_id", "row_id"), time_column="ts")
     )
 )
 
 txn_join = Join(
     left=source_users,
-    row_ids="user_id",
     right_parts=[
         JoinPart(group_by=txn_group_by_user, prefix="user"),
         JoinPart(group_by=txn_group_by_merchant, prefix="merchant"),

api/python/test/sample/joins/sample_team/sample_backfill_mutation_join.py

@@ -23,7 +23,6 @@
 
 v0 = Join(
     left=test_sources.event_source,
-    row_ids="subject",
     right_parts=[JoinPart(group_by=mutation_sample_group_by.v0)],
     online=False,
     version=0,

api/python/test/sample/joins/sample_team/sample_chaining_join.py

@@ -23,7 +23,6 @@
 
 v1 = Join(
     left=test_sources.event_source,
-    row_ids=["subject", "event"],
     right_parts=[
         JoinPart(
             group_by=chaining_group_by_v1,

api/python/test/sample/joins/sample_team/sample_chaining_join_parent.py

@@ -8,7 +8,6 @@
 
 parent_join = Join(
     left=test_sources.event_source,
-    row_ids="group_by_subject",
     right_parts=[
         JoinPart(
             group_by=event_sample_group_by.v1,

api/python/test/sample/joins/sample_team/sample_join.py

@@ -25,7 +25,6 @@
 v1 = Join(
     left=test_sources.staging_entities,
     right_parts=[JoinPart(group_by=sample_group_by.v1)],
-    row_ids="place_id",
     table_properties={"config_json": """{"sample_key": "sample_value"}"""},
     output_namespace="sample_namespace",
     env_vars=EnvironmentVariables(
@@ -41,7 +40,6 @@
 never = Join(
     left=test_sources.staging_entities,
     right_parts=[JoinPart(group_by=sample_group_by.v1)],
-    row_ids=["s2CellId", "place_id"],
     output_namespace="sample_namespace",
     offline_schedule="@never",
     version=0,
@@ -50,15 +48,13 @@
 group_by_of_group_by = Join(
     left=test_sources.staging_entities,
     right_parts=[JoinPart(group_by=sample_group_by_group_by.v1)],
-    row_ids="s2CellId",
     output_namespace="sample_namespace",
     version=0,
 )
 
 consistency_check = Join(
     left=test_sources.staging_entities,
     right_parts=[JoinPart(group_by=sample_group_by.v1)],
-    row_ids="place_id",
     output_namespace="sample_namespace",
     check_consistency=True,
     version=0,
@@ -67,7 +63,6 @@
 no_log_flattener = Join(
     left=test_sources.staging_entities,
     right_parts=[JoinPart(group_by=sample_group_by.v1)],
-    row_ids=["place_id"],
     output_namespace="sample_namespace",
     sample_percent=0.0,
     version=0,