Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 1 addition & 1 deletion .github/trigger_files/beam_PostCommit_Python.json
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run.",
"pr": "38701",
"modification": 57
"modification": 58
}
2 changes: 1 addition & 1 deletion .github/trigger_files/beam_PostCommit_Python_Arm.json
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 1
"modification": 2
}
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 4
"modification": 5
}
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run.",
"pr": "37360",
"modification": 2
"modification": 3
}
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 5
"modification": 6
}
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 5,
"modification": 6,
"https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test"
}
Original file line number Diff line number Diff line change
Expand Up @@ -2,5 +2,6 @@
"https://github.com/apache/beam/pull/32648": "testing addition of Flink 1.19 support",
"https://github.com/apache/beam/pull/34830": "testing",
"trigger-2026-04-04": "portable_runner expand_sdf opt-in",
"https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test"
"https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test",
"modification": 1
}
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,6 @@
"https://github.com/apache/beam/issues/35429": "testing",
"trigger-2026-04-04": "portable_runner expand_sdf opt-in",
"https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test",
"modification": 1,
"modification": 2,
"https://github.com/apache/beam/issues/19468": "SDF self-checkpointing and bundle finalization"
}
2 changes: 1 addition & 1 deletion .github/trigger_files/beam_PostCommit_Python_Versions.json
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"revision": 5
"revision": 6
}
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 16
"modification": 17
}
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 2
"modification": 3
}
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 3
"modification": 4
}
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run",
"modification": 2
"modification": 3
}
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run.",
"modification": 1
"modification": 2
}
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
{
"comment": "Modify this file in a trivial way to cause this test suite to run.",
"modification": 1
"modification": 2
}
85 changes: 68 additions & 17 deletions sdks/python/apache_beam/io/gcp/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -389,8 +389,14 @@ def chain_after(result):
from apache_beam.io.gcp.bigquery_read_internal import _PassThroughThenCleanup
from apache_beam.io.gcp.bigquery_read_internal import _PassThroughThenCleanupTempDatasets
from apache_beam.io.gcp.bigquery_read_internal import bigquery_export_destination_uri
from apache_beam.io.gcp.bigquery_tools import DatasetReference
from apache_beam.io.gcp.bigquery_tools import JobReference
from apache_beam.io.gcp.bigquery_tools import RetryStrategy
from apache_beam.io.gcp.internal.clients import bigquery
from apache_beam.io.gcp.bigquery_tools import TableCell
from apache_beam.io.gcp.bigquery_tools import TableFieldSchema
from apache_beam.io.gcp.bigquery_tools import TableReference
from apache_beam.io.gcp.bigquery_tools import TableRow
from apache_beam.io.gcp.bigquery_tools import TableSchema
from apache_beam.io.iobase import BoundedSource
from apache_beam.io.iobase import RangeTracker
from apache_beam.io.iobase import SDFBoundedSourceReader
Expand Down Expand Up @@ -421,13 +427,12 @@ def chain_after(result):
from apache_beam.utils.annotations import deprecated

try:
from apache_beam.io.gcp.internal.clients.bigquery import DatasetReference
from apache_beam.io.gcp.internal.clients.bigquery import JobReference
from apache_beam.io.gcp.internal.clients.bigquery import TableReference
from apache_beam.io.gcp.internal.clients import bigquery
except ImportError:
DatasetReference = None
TableReference = None
JobReference = None
bigquery = None

if bigquery is None or not hasattr(bigquery, 'TableReference'):
import apache_beam.io.gcp.bigquery_tools as bigquery

_LOGGER = logging.getLogger(__name__)

Expand Down Expand Up @@ -524,8 +529,12 @@ def __init__(self, table_schema=None):
self.table_schema = table_schema
# Precompute field names since we need them for row encoding.
if self.table_schema:
self.field_names = tuple(fs.name for fs in self.table_schema.fields)
self.field_types = tuple(fs.type for fs in self.table_schema.fields)
fields = (
self.table_schema.fields
if hasattr(self.table_schema, 'fields') else self.table_schema)
self.field_names = tuple(fs.name for fs in fields)
self.field_types = tuple(
getattr(fs, 'type', getattr(fs, 'field_type', None)) for fs in fields)

def encode(self, table_row):
if self.table_schema is None:
Expand Down Expand Up @@ -767,14 +776,42 @@ def estimate_size(self):
# no access to the query that we're running.
return None

def _get_temp_dataset_id(self):
if self.temp_dataset is None:
return None
elif hasattr(self.temp_dataset, 'datasetId'):
return self.temp_dataset.datasetId
elif hasattr(self.temp_dataset, 'dataset_id'):
return self.temp_dataset.dataset_id
elif isinstance(self.temp_dataset, str):
if ':' in self.temp_dataset:
return self.temp_dataset.split(':', 1)[1]
elif '.' in self.temp_dataset:
return self.temp_dataset.split('.', 1)[1]
return self.temp_dataset
return None

def _get_temp_dataset_project(self):
if hasattr(self.temp_dataset, 'projectId') and self.temp_dataset.projectId:
return self.temp_dataset.projectId
elif hasattr(self.temp_dataset, 'project') and self.temp_dataset.project:
return self.temp_dataset.project
elif isinstance(self.temp_dataset, str):
if ':' in self.temp_dataset:
return self.temp_dataset.split(':', 1)[0]
elif '.' in self.temp_dataset:
return self.temp_dataset.split('.', 1)[0]
return None

def _get_project(self):
"""Returns the project that queries and exports will be billed to."""
temp_project = self._get_temp_dataset_project()
if temp_project:
return temp_project

project = self.options.view_as(GoogleCloudOptions).project
if isinstance(project, vp.ValueProvider):
project = project.get()
if self.temp_dataset:
return self.temp_dataset.projectId
if not project:
project = self.project
return project
Expand All @@ -794,8 +831,7 @@ def _create_source(self, path, coder):
def split(self, desired_bundle_size, start_position=None, stop_position=None):
if self.export_result is None:
bq = bigquery_tools.BigQueryWrapper(
temp_dataset_id=(
self.temp_dataset.datasetId if self.temp_dataset else None),
temp_dataset_id=self._get_temp_dataset_id(),
client=bigquery_tools.BigQueryWrapper._bigquery_client(self.options))

if self.query is not None:
Expand Down Expand Up @@ -3088,10 +3124,25 @@ def file_path_to_remove(unused_elm):
def _expand_direct_read(self, pcoll):
project_id = None
temp_table_ref = None
if 'temp_dataset' in self._kwargs:
temp_table_ref = bigquery.TableReference(
projectId=self._kwargs['temp_dataset'].projectId,
datasetId=self._kwargs['temp_dataset'].datasetId,
temp_dataset = self._kwargs.get('temp_dataset')
if temp_dataset is not None:
if isinstance(temp_dataset, str):
if ':' in temp_dataset:
project_id, dataset_id = temp_dataset.split(':', 1)
elif '.' in temp_dataset:
project_id, dataset_id = temp_dataset.split('.', 1)
else:
project_id = pcoll.pipeline.options.view_as(
GoogleCloudOptions).project
dataset_id = temp_dataset
else:
project_id = getattr(temp_dataset, 'projectId', None) or getattr(
temp_dataset, 'project', None)
dataset_id = getattr(temp_dataset, 'datasetId', None) or getattr(
temp_dataset, 'dataset_id', None)
temp_table_ref = TableReference(
projectId=project_id,
datasetId=dataset_id,
tableId='beam_temp_table_' + uuid.uuid4().hex)
else:
project_id = pcoll.pipeline.options.view_as(GoogleCloudOptions).project
Expand Down
11 changes: 7 additions & 4 deletions sdks/python/apache_beam/io/gcp/bigquery_avro_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,8 @@
"precision": 38,
"scale": 9,
},
"BIGNUMERIC": "string",
"JSON": "string",
"GEOGRAPHY": "string",
}

Expand Down Expand Up @@ -108,11 +110,12 @@ def table_field_to_avro_field(table_field: dict[str, Any],
"""
assert "type" in table_field, \
"Unable to get type for table field {}".format(table_field)
assert table_field["type"] in BIG_QUERY_TO_AVRO_TYPES, \
field_type_str = table_field["type"].upper()
assert field_type_str in BIG_QUERY_TO_AVRO_TYPES, \
"Unable to map BigQuery field type {} to avro type".format(
table_field["type"])

avro_type = BIG_QUERY_TO_AVRO_TYPES[table_field["type"]]
avro_type = BIG_QUERY_TO_AVRO_TYPES[field_type_str]

if avro_type == "record":
element_type = get_record_schema_from_dict_table_schema(
Expand All @@ -122,9 +125,9 @@ def table_field_to_avro_field(table_field: dict[str, Any],
else:
element_type = avro_type

field_mode = table_field.get("mode", "NULLABLE")
field_mode = (table_field.get("mode") or "NULLABLE").upper()

if field_mode in (None, "NULLABLE"):
if field_mode in (None, "NULLABLE", ""):
field_type = ["null", element_type]
elif field_mode == "REQUIRED":
field_type = element_type
Expand Down
Loading
Loading