Merge branch 'main' into vz/add_test_case_for_different_partition_formats

Author: nikhil-zlai

api/python/ai/chronon/airflow_helpers.py

@@ -12,6 +12,7 @@ def create_airflow_dependency(table, partition_column, additional_partitions=Non
     Args:
         table: The table name (with namespace)
         partition_column: The partition column to use (defaults to 'ds')
+        additional_partitions: Additional partitions to include in the dependency
 
     Returns:
         A dictionary with name and spec for the Airflow dependency
@@ -47,6 +48,12 @@ def _get_partition_col_from_query(query):
         return query.partitionColumn
     return None
 
+def _get_additional_subPartitionsToWaitFor_from_query(query):
+    """Gets additional subPartitionsToWaitFor from query if available"""
+    if query:
+        return query.subPartitionsToWaitFor
+    return None
+
 
 def _get_airflow_deps_from_source(source, partition_column=None):
     """
@@ -60,20 +67,22 @@ def _get_airflow_deps_from_source(source, partition_column=None):
         A list of Airflow dependency objects
     """
     tables = []
+    additional_partitions = None
     # Assumes source has already been normalized
     if source.events:
         tables = [source.events.table]
         # Use partition column from query if available, otherwise use the provided one
-        source_partition_column = (
-            _get_partition_col_from_query(source.events.query) or partition_column
+        source_partition_column, additional_partitions = (
+            _get_partition_col_from_query(source.events.query) or partition_column, _get_additional_subPartitionsToWaitFor_from_query(source.events.query)
         )
+
     elif source.entities:
         # Given the setup of Query, we currently mandate the same partition column for snapshot and mutations tables
         tables = [source.entities.snapshotTable]
         if source.entities.mutationTable:
             tables.append(source.entities.mutationTable)
-        source_partition_column = (
-            _get_partition_col_from_query(source.entities.query) or partition_column
+        source_partition_column, additional_partitions = (
+            _get_partition_col_from_query(source.entities.query) or partition_column, _get_additional_subPartitionsToWaitFor_from_query(source.entities.query)
         )
     elif source.joinSource:
         # TODO: Handle joinSource -- it doesn't work right now because the metadata isn't set on joinSource at this point
@@ -83,7 +92,7 @@ def _get_airflow_deps_from_source(source, partition_column=None):
         return []
 
     return [
-        create_airflow_dependency(table, source_partition_column) for table in tables
+        create_airflow_dependency(table, source_partition_column, additional_partitions) for table in tables
     ]
 
 

api/python/ai/chronon/group_by.py

@@ -648,6 +648,7 @@ def _normalize_source(source):
                 for output_col in get_output_col_names(agg):
                     column_tags[output_col] = agg.tags
 
+
     metadata = ttypes.MetaData(
         online=online,
         production=production,

api/python/ai/chronon/query.py

@@ -29,6 +29,7 @@ def Query(
     reversal_column: str = None,
     partition_column: str = None,
     partition_format: str = None,
+    sub_partitions_to_wait_for: List[str] = None,
 ) -> api.Query:
     """
     Create a query object that is used to scan data from various data sources.
@@ -76,6 +77,9 @@ def Query(
     :param partition_column:
         Specify this to override spark.chronon.partition.column set in teams.py for this particular query.
     :type partition_column: str, optional
+    :param sub_partitions_to_wait_for:
+        Additional partitions to be used in sensing that the source data has landed. Should be a full partition string, such as `hr=23:00'
+    :type partition_column: List[str], optional
     :param partition_format:
         Date format string to expect the partition values to be in.
     :type partition_format: str, optional
@@ -91,6 +95,7 @@ def Query(
         mutationTimeColumn=mutation_time_column,
         reversalColumn=reversal_column,
         partitionColumn=partition_column,
+        subPartitionsToWaitFor=sub_partitions_to_wait_for,
         partitionFormat=partition_format
     )
 

api/thrift/api.thrift

@@ -40,6 +40,12 @@ struct Query {
     **/
     23: optional common.Window partitionLag
 
+    /**
+    * Additional partitions to be used in sensing that the source data has landed.
+    * Should be a full partition string, such as `hr=23:00'
+    **/
+    24: optional list<string> subPartitionsToWaitFor
+
 }
 
 /**