fix(ingest): map bigquery nested types properly (datahub-project#10409)

darnaut · May 2, 2024 · 8bebf63 · 8bebf63
1 parent f3c42f6
commit 8bebf63
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 2 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py b/metadata-ingestion/src/datahub/ingestion/source/bigquery_v2/bigquery.py
@@ -178,6 +178,8 @@ def cleanup(config: BigQueryV2Config) -> None:
 )
 class BigqueryV2Source(StatefulIngestionSourceBase, TestableSource):
  # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types
+ # Note: We use the hive schema parser to parse nested BigQuery types. We also have
+ # some extra type mappings in that file.
  BIGQUERY_FIELD_TYPE_MAPPINGS: Dict[
  str,
  Type[
@@ -1264,7 +1266,6 @@ def gen_schema_fields(self, columns: List[BigqueryColumn]) -> List[SchemaField]:
  type=SchemaFieldDataType(
  self.BIGQUERY_FIELD_TYPE_MAPPINGS.get(col.data_type, NullType)()
  ),
- # NOTE: nativeDataType will not be in sync with older connector
  nativeDataType=col.data_type,
  description=col.comment,
  nullable=col.is_nullable,

diff --git a/metadata-ingestion/src/datahub/utilities/hive_schema_to_avro.py b/metadata-ingestion/src/datahub/utilities/hive_schema_to_avro.py
@@ -28,6 +28,20 @@ class HiveColumnToAvroConverter:
  "bigint": "long",
  "varchar": "string",
  "char": "string",
+ "bytes": "bytes",
+ }
+ _EXTRA_BIGQUERY_TYPE_TO_AVRO_TYPE = {
+ # A few extra types, purely to map BigQuery things correctly.
+ "bool": "boolean",
+ "decimal": "double",
+ "numeric": "int",
+ "bignumeric": "long",
+ "bigdecimal": "double",
+ "float64": "double",
+ "int64": "long",
+ "smallint": "int",
+ "tinyint": "int",
+ "byteint": "int",
  }
 
  _COMPLEX_TYPE = re.compile("^(struct|map|array|uniontype)")
@@ -180,13 +194,19 @@ def _parse_basic_datatype_string(s: str) -> Dict[str, object]:
  "native_data_type": s,
  "_nullable": True,
  }
- elif s == "timestamp":
+ elif s in {"timestamp", "datetime"}:
  return {
  "type": "int",
  "logicalType": "timestamp-millis",
  "native_data_type": s,
  "_nullable": True,
  }
+ elif s in HiveColumnToAvroConverter._EXTRA_BIGQUERY_TYPE_TO_AVRO_TYPE:
+ return {
+ "type": HiveColumnToAvroConverter._EXTRA_BIGQUERY_TYPE_TO_AVRO_TYPE[s],
+ "native_data_type": s,
+ "_nullable": True,
+ }
  else:
  return {"type": "null", "native_data_type": s, "_nullable": True}