feat(ingest/bigquery): improve profiler to support multiple partition columns and support external table profiling

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py

@@ -41,7 +41,7 @@
     BigQueryIdentifierBuilder,
 )
 from datahub.ingestion.source.bigquery_v2.lineage import BigqueryLineageExtractor
-from datahub.ingestion.source.bigquery_v2.profiler import BigqueryProfiler
+from datahub.ingestion.source.bigquery_v2.profiling.profiler import BigqueryProfiler
 from datahub.ingestion.source.bigquery_v2.queries_extractor import (
     BigQueryQueriesExtractor,
     BigQueryQueriesExtractorConfig,

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_config.py

@@ -28,6 +28,7 @@
     BigQueryConnectionConfig,
 )
 from datahub.ingestion.source.data_lake_common.path_spec import PathSpec
+from datahub.ingestion.source.ge_profiling_config import GEProfilingConfig
 from datahub.ingestion.source.sql.sql_config import SQLCommonConfig, SQLFilterConfig
 from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulLineageConfigMixin,
@@ -50,6 +51,49 @@
 )
 
 
+class BigQueryProfilingConfig(GEProfilingConfig):
+    fallback_partition_values: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="Fallback values for partition columns when timeout occurs. Keys are column names, "
+        "values are the fallback values to use. For non-date columns, the values are used directly. "
+        "Example: {'batch': 'default', 'region': 'us-east-1'}",
+    )
+
+    partition_fetch_timeout: int = Field(
+        default=30,
+        description="Timeout in seconds for partition value fetch operations. If exceeded, fallback "
+        "partition values will be used.",
+    )
+
+    profiling_row_limit: int = Field(
+        default=1000000,
+        description="The number of rows to sample for profiling. This is a low level config property which "
+        "should be touched with care. This restriction is needed because excessively wide tables can "
+        "result in failure to ingest the schema.",
+    )
+
+    skip_stale_tables: bool = Field(
+        default=True,
+        description="Skip profiling for tables that haven't been modified in over a year. "
+        "Uses last_altered timestamp (which contains BigQuery's last_modified_time) for both regular and external tables. "
+        "This helps avoid profiling abandoned or archived tables.",
+    )
+
+    staleness_threshold_days: int = Field(
+        default=365,
+        description="Number of days after which a table is considered stale and profiling will be skipped "
+        "if skip_stale_tables is enabled.",
+    )
+
+    partition_datetime_window_days: Optional[int] = Field(
+        default=30,
+        description="Limit profiling to partitions within this many days from the selected partition date. "
+        "For example, if set to 30 and the selected partition is '2025-08-15', only partitions from "
+        "'2025-07-16' to '2025-08-15' will be included in profiling. Set to None to disable date windowing. "
+        "This helps focus profiling on recent data patterns and improves performance.",
+    )
+
+
 class BigQueryBaseConfig(ConfigModel):
     rate_limit: bool = Field(
         default=False, description="Should we rate limit requests made to API."
@@ -482,6 +526,11 @@ def have_table_data_read_permission(self) -> bool:
         "See [this](https://cloud.google.com/bigquery/docs/information-schema-jobs#scope_and_syntax) for details.",
     )
 
+    profiling: BigQueryProfilingConfig = Field(
+        default=BigQueryProfilingConfig(),
+        description="Profiling related configs",
+    )
+
     pushdown_deny_usernames: List[str] = Field(
         default=[],
         description="List of user email patterns using SQL LIKE syntax (e.g., 'bot_%', '%@%.iam.gserviceaccount.com') "

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_schema.py

@@ -13,6 +13,7 @@
     List,
     Optional,
     Set,
+    Tuple,
 )
 
 from google.api_core import retry
@@ -70,22 +71,29 @@ class BigqueryTableConstraint:
 RANGE_PARTITION_NAME: str = "RANGE"
 
 
-@dataclass
+@dataclass(frozen=True)
 class PartitionInfo:
-    field: str
-    # Data type is optional as we not have it when we set it from TimePartitioning
-    column: Optional[BigqueryColumn] = None
+    fields: Tuple[str, ...]
+    columns: Optional[Tuple[BigqueryColumn, ...]] = None
     type: str = TimePartitioningType.DAY
     expiration_ms: Optional[int] = None
     require_partition_filter: bool = False
 
-    # TimePartitioning field doesn't provide data_type so we have to add it afterwards
+    def __post_init__(self) -> None:
+        if not self.fields:
+            raise ValueError("PartitionInfo must have at least one field")
+        if self.columns is not None and len(self.fields) != len(self.columns):
+            raise ValueError(
+                f"fields/columns length mismatch: {len(self.fields)} fields vs {len(self.columns)} columns"
+            )
+
     @classmethod
     def from_time_partitioning(
         cls, time_partitioning: TimePartitioning
     ) -> "PartitionInfo":
+        """Convert BigQuery time partitioning to PartitionInfo."""
         return cls(
-            field=time_partitioning.field or "_PARTITIONTIME",
+            fields=(time_partitioning.field or "_PARTITIONTIME",),
             type=time_partitioning.type_,
             expiration_ms=time_partitioning.expiration_ms,
             require_partition_filter=time_partitioning.require_partition_filter,
@@ -100,7 +108,7 @@ def from_range_partitioning(
             return None
 
         return cls(
-            field=field,
+            fields=(field,),
             type=RANGE_PARTITION_NAME,
         )
 

metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery_schema_gen.py

@@ -2,6 +2,7 @@
 import re
 from base64 import b32decode
 from collections import defaultdict
+from dataclasses import replace
 from typing import Dict, Iterable, List, Optional, Set, Tuple, Type, Union, cast
 
 from google.cloud.bigquery.table import TableListItem
@@ -55,7 +56,7 @@
     BigQueryFilter,
     BigQueryIdentifierBuilder,
 )
-from datahub.ingestion.source.bigquery_v2.profiler import BigqueryProfiler
+from datahub.ingestion.source.bigquery_v2.profiling.profiler import BigqueryProfiler
 from datahub.ingestion.source.common.subtypes import (
     DatasetContainerSubTypes,
     DatasetSubTypes,
@@ -765,15 +766,16 @@ def _process_table(
             )
 
         # If table has time partitioning, set the data type of the partitioning field
-        if table.partition_info:
-            table.partition_info.column = next(
-                (
-                    column
-                    for column in columns
-                    if column.name == table.partition_info.field
-                ),
-                None,
+        if table.partition_info and table.partition_info.fields:
+            matching_columns = tuple(
+                column
+                for column in columns
+                if column.name in table.partition_info.fields
             )
+            if matching_columns:
+                table.partition_info = replace(
+                    table.partition_info, columns=matching_columns
+                )
         yield from self.gen_table_dataset_workunits(
             table, columns, project_id, dataset_name
         )