Skip to content

Add support for parquetOptions in GCSToBigQueryOperator #60876

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shahar1 merged 5 commits into from
Jan 31, 2026
Merged

Add support for parquetOptions in GCSToBigQueryOperator #60876

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
5afbcb6
Add tests for gcs_to_bigquery parquet format
mlauter Jan 16, 2026
a7a21ea
Add support for `parquetOptions` in GCSToBigQueryOperator
mlauter Jan 20, 2026
e5666b9
Refactor _validate_src_fmt_configs
mlauter Jan 20, 2026
603a232
Add newsfragment
mlauter Jan 21, 2026
9d8f5b2
Delete newsfragment
mlauter Jan 25, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 49 additions & 7 deletions providers/google/src/airflow/providers/google/cloud/transfers/gcs_to_bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -544,7 +544,11 @@ def _create_external_table(self):
"allowJaggedRows": self.allow_jagged_rows,
"encoding": self.encoding,
}
src_fmt_to_param_mapping = {"CSV": "csvOptions", "GOOGLE_SHEETS": "googleSheetsOptions"}
src_fmt_to_param_mapping = {
"CSV": "csvOptions",
"GOOGLE_SHEETS": "googleSheetsOptions",
"PARQUET": "parquetOptions",
}
src_fmt_to_configs_mapping = {
"csvOptions": [
"allowJaggedRows",
Expand All @@ -557,6 +561,7 @@ def _create_external_table(self):
"columnNameCharacterMap",
],
"googleSheetsOptions": ["skipLeadingRows"],
"parquetOptions": ["enumAsString", "enableListInference", "mapTargetType"],
}
if self.source_format in src_fmt_to_param_mapping:
valid_configs = src_fmt_to_configs_mapping[src_fmt_to_param_mapping[self.source_format]]
Expand Down Expand Up @@ -687,8 +692,18 @@ def _use_existing_table(self):
"ORC": ["autodetect"],
}

# Some source formats have nested configuration options which are not available
# at the top level of the load configuration.
src_fmt_to_param_mapping = {"PARQUET": "parquetOptions"}
src_fmt_to_nested_configs_mapping = {
"parquetOptions": ["enumAsString", "enableListInference", "mapTargetType"],
}

valid_configs = src_fmt_to_configs_mapping[self.source_format]

src_fmt_param = src_fmt_to_param_mapping.get(self.source_format)
valid_nested_configs = src_fmt_to_nested_configs_mapping[src_fmt_param] if src_fmt_param else None

# if following fields are not specified in src_fmt_configs,
# honor the top-level params for backward-compatibility
backward_compatibility_configs = {
Expand All @@ -701,7 +716,12 @@ def _use_existing_table(self):
}

self.src_fmt_configs = self._validate_src_fmt_configs(
self.source_format, self.src_fmt_configs, valid_configs, backward_compatibility_configs
self.source_format,
self.src_fmt_configs,
valid_configs,
backward_compatibility_configs,
src_fmt_param,
valid_nested_configs,
)

self.configuration["load"].update(self.src_fmt_configs)
Expand All @@ -716,29 +736,51 @@ def _validate_src_fmt_configs(
src_fmt_configs: dict,
valid_configs: list[str],
backward_compatibility_configs: dict | None = None,
src_fmt_param: str | None = None,
valid_nested_configs: list[str] | None = None,
) -> dict:
"""
Validate the given src_fmt_configs against a valid configuration for the source format.
Validate and format the given src_fmt_configs against a valid configuration for the source format.

Adds the backward compatibility config to the src_fmt_configs.

Adds nested source format configurations if valid_nested_configs is provided.

:param source_format: File format to export.
:param src_fmt_configs: Configure optional fields specific to the source format.
:param valid_configs: Valid configuration specific to the source format
:param backward_compatibility_configs: The top-level params for backward-compatibility
:param src_fmt_param: The source format parameter for nested configurations.
Required when valid_nested_configs is provided.
:param valid_nested_configs: Valid nested configuration specific to the source format.
"""
valid_src_fmt_configs = {}

if backward_compatibility_configs is None:
backward_compatibility_configs = {}

for k, v in backward_compatibility_configs.items():
if k not in src_fmt_configs and k in valid_configs:
src_fmt_configs[k] = v
valid_src_fmt_configs[k] = v

if valid_nested_configs is None:
valid_nested_configs = []

if valid_nested_configs:
if src_fmt_param is None:
raise ValueError("src_fmt_param is required when valid_nested_configs is provided.")

for k in src_fmt_configs:
if k not in valid_configs:
valid_src_fmt_configs[src_fmt_param] = {}

for k, v in src_fmt_configs.items():
if k in valid_configs:
valid_src_fmt_configs[k] = v
elif k in valid_nested_configs:
valid_src_fmt_configs[src_fmt_param][k] = v
else:
Comment thread
shahar1 marked this conversation as resolved.
raise ValueError(f"{k} is not a valid src_fmt_configs for type {source_format}.")

return src_fmt_configs
return valid_src_fmt_configs

def _cleanse_time_partitioning(
self, destination_dataset_table: str | None, time_partitioning_in: dict | None
Expand Down
206 changes: 206 additions & 0 deletions providers/google/tests/unit/google/cloud/transfers/test_gcs_to_bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1682,6 +1682,212 @@ def test_external_table_should_accept_orc_source_format(self, hook):
},
)

@mock.patch(GCS_TO_BQ_PATH.format("BigQueryHook"))
def test_external_table_should_accept_parquet_format(self, hook):
hook.return_value.insert_job.side_effect = [
MagicMock(job_id=REAL_JOB_ID, error_result=False),
REAL_JOB_ID,
]
hook.return_value.generate_job_id.return_value = REAL_JOB_ID
hook.return_value.split_tablename.return_value = (PROJECT_ID, DATASET, TABLE)

operator = GCSToBigQueryOperator(
task_id=TASK_ID,
bucket=TEST_BUCKET,
source_objects=TEST_SOURCE_OBJECTS,
destination_project_dataset_table=TEST_EXPLICIT_DEST,
schema_fields=SCHEMA_FIELDS,
write_disposition=WRITE_DISPOSITION,
external_table=True,
project_id=JOB_PROJECT_ID,
source_format="PARQUET",
)

operator.execute(context=MagicMock())

hook.return_value.create_table.assert_called_once_with(
exists_ok=True,
location=None,
project_id=JOB_PROJECT_ID,
dataset_id=DATASET,
table_id=TABLE,
table_resource={
"tableReference": {
"projectId": PROJECT_ID,
"datasetId": DATASET,
"tableId": TABLE,
},
"externalDataConfiguration": {
"autodetect": True,
"sourceFormat": "PARQUET",
"sourceUris": [f"gs://{TEST_BUCKET}/{TEST_SOURCE_OBJECTS}"],
"compression": "NONE",
"ignoreUnknownValues": False,
"schema": {"fields": SCHEMA_FIELDS},
"parquetOptions": {},
},
"labels": {},
},
)

@mock.patch(GCS_TO_BQ_PATH.format("BigQueryHook"))
def test_external_table_should_accept_parquet_format_and_options(self, hook):
hook.return_value.insert_job.side_effect = [
MagicMock(job_id=REAL_JOB_ID, error_result=False),
REAL_JOB_ID,
]
hook.return_value.generate_job_id.return_value = REAL_JOB_ID
hook.return_value.split_tablename.return_value = (PROJECT_ID, DATASET, TABLE)

operator = GCSToBigQueryOperator(
task_id=TASK_ID,
bucket=TEST_BUCKET,
source_objects=TEST_SOURCE_OBJECTS,
destination_project_dataset_table=TEST_EXPLICIT_DEST,
schema_fields=SCHEMA_FIELDS,
write_disposition=WRITE_DISPOSITION,
external_table=True,
project_id=JOB_PROJECT_ID,
source_format="PARQUET",
src_fmt_configs={
"enableListInference": True,
},
)

operator.execute(context=MagicMock())

hook.return_value.create_table.assert_called_once_with(
exists_ok=True,
location=None,
project_id=JOB_PROJECT_ID,
dataset_id=DATASET,
table_id=TABLE,
table_resource={
"tableReference": {
"projectId": PROJECT_ID,
"datasetId": DATASET,
"tableId": TABLE,
},
"externalDataConfiguration": {
"autodetect": True,
"sourceFormat": "PARQUET",
"sourceUris": [f"gs://{TEST_BUCKET}/{TEST_SOURCE_OBJECTS}"],
"compression": "NONE",
"ignoreUnknownValues": False,
"schema": {"fields": SCHEMA_FIELDS},
"parquetOptions": {
"enableListInference": True,
},
},
"labels": {},
},
)

@mock.patch(GCS_TO_BQ_PATH.format("BigQueryHook"))
def test_without_external_table_should_accept_parquet_format(self, hook):
hook.return_value.insert_job.side_effect = [
MagicMock(job_id=REAL_JOB_ID, error_result=False),
REAL_JOB_ID,
]
hook.return_value.generate_job_id.return_value = REAL_JOB_ID
hook.return_value.split_tablename.return_value = (PROJECT_ID, DATASET, TABLE)
operator = GCSToBigQueryOperator(
task_id=TASK_ID,
bucket=TEST_BUCKET,
source_objects=TEST_SOURCE_OBJECTS,
write_disposition=WRITE_DISPOSITION,
destination_project_dataset_table=TEST_EXPLICIT_DEST,
external_table=False,
project_id=JOB_PROJECT_ID,
source_format="PARQUET",
)

operator.execute(context=MagicMock())

calls = [
call(
configuration={
"load": {
"autodetect": True,
"createDisposition": "CREATE_IF_NEEDED",
"destinationTable": {
"projectId": "test-project",
"datasetId": "dataset",
"tableId": "table",
},
"sourceFormat": "PARQUET",
"sourceUris": ["gs://test-bucket/test/objects/test.csv"],
"writeDisposition": "WRITE_TRUNCATE",
"ignoreUnknownValues": False,
"parquetOptions": {},
}
},
job_id=REAL_JOB_ID,
location=None,
nowait=True,
project_id=JOB_PROJECT_ID,
retry=DEFAULT_RETRY,
timeout=None,
)
]

hook.return_value.insert_job.assert_has_calls(calls)

@mock.patch(GCS_TO_BQ_PATH.format("BigQueryHook"))
def test_without_external_table_should_accept_parquet_format_and_options(self, hook):
hook.return_value.insert_job.side_effect = [
MagicMock(job_id=REAL_JOB_ID, error_result=False),
REAL_JOB_ID,
]
hook.return_value.generate_job_id.return_value = REAL_JOB_ID
hook.return_value.split_tablename.return_value = (PROJECT_ID, DATASET, TABLE)
operator = GCSToBigQueryOperator(
task_id=TASK_ID,
bucket=TEST_BUCKET,
source_objects=TEST_SOURCE_OBJECTS,
write_disposition=WRITE_DISPOSITION,
destination_project_dataset_table=TEST_EXPLICIT_DEST,
external_table=False,
project_id=JOB_PROJECT_ID,
source_format="PARQUET",
src_fmt_configs={
"enableListInference": True,
},
)

operator.execute(context=MagicMock())

calls = [
call(
configuration={
"load": {
"autodetect": True,
"createDisposition": "CREATE_IF_NEEDED",
"destinationTable": {
"projectId": "test-project",
"datasetId": "dataset",
"tableId": "table",
},
"sourceFormat": "PARQUET",
"sourceUris": ["gs://test-bucket/test/objects/test.csv"],
"writeDisposition": "WRITE_TRUNCATE",
"ignoreUnknownValues": False,
"parquetOptions": {
"enableListInference": True,
},
}
},
job_id=REAL_JOB_ID,
location=None,
nowait=True,
project_id=JOB_PROJECT_ID,
retry=DEFAULT_RETRY,
timeout=None,
)
]

hook.return_value.insert_job.assert_has_calls(calls)


@pytest.fixture
def create_task_instance(create_task_instance_of_operator, session):
Expand Down