claudevdm commented on code in PR #39236:
URL: https://github.com/apache/beam/pull/39236#discussion_r4219754055
##########
sdks/python/apache_beam/io/gcp/bigquery.py:
##########
@@ -2883,25 +3089,175 @@ def __enter__(self):
def __exit__(self, *args):
pass
+ class _ConvertDynamicRowDoFn(DoFn):
+ def __init__(
+ self,
+ schema,
+ union_field_names,
+ create_disposition=BigQueryDisposition.CREATE_IF_NEEDED,
+ write_disposition=BigQueryDisposition.WRITE_APPEND,
+ additional_bq_parameters=None,
+ test_client=None):
+ self.schema = schema
+ self.union_field_names = union_field_names
+ self.create_disposition = create_disposition
+ self.write_disposition = write_disposition
+ self.additional_bq_parameters = additional_bq_parameters
+ self.test_client = test_client
+ self.bigquery_wrapper = None
+
+ def start_bundle(self):
+ if self.create_disposition == BigQueryDisposition.CREATE_IF_NEEDED:
+ if self.test_client:
+ self.bigquery_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self.test_client)
+ else:
+ try:
+ self.bigquery_wrapper = bigquery_tools.BigQueryWrapper()
+ except Exception:
+ self.bigquery_wrapper = None
+
+ def _create_table_if_needed(self, dest, record_schema):
+ if self.create_disposition != BigQueryDisposition.CREATE_IF_NEEDED:
+ return
+ try:
+ table_ref = bigquery_tools.parse_table_reference(dest)
+ except ValueError:
+ return
+ if not table_ref.datasetId or not table_ref.tableId:
+ return
+
+ if self.bigquery_wrapper is None:
+ if self.test_client:
+ self.bigquery_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self.test_client)
+ else:
+ try:
+ self.bigquery_wrapper = bigquery_tools.BigQueryWrapper()
+ except Exception:
+ return
+
+ project_id = (
+ table_ref.projectId or self.bigquery_wrapper._get_project_id())
+ str_table_ref = '%s:%s.%s' % (
+ project_id, table_ref.datasetId, table_ref.tableId)
+ if str_table_ref in _KNOWN_TABLES or dest in _KNOWN_TABLES:
+ return
+
+ table_schema = bigquery_tools.get_bq_tableschema(record_schema)
+ self.bigquery_wrapper.get_or_create_table(
+ project_id=project_id,
+ dataset_id=table_ref.datasetId,
+ table_id=table_ref.tableId,
+ schema=table_schema,
+ create_disposition=self.create_disposition,
+ write_disposition=self.write_disposition,
Review Comment:
This passes the user's write disposition to `get_or_create_table`. For
WRITE_TRUNCATE that helper deletes the table, creates it again and sleeps 150
seconds, and here it runs inside `process`.
`_KNOWN_TABLES` is per process, so every worker process does this the first
time it sees a row for a table.
Can we always pass WRITE_APPEND here?
##########
sdks/python/apache_beam/io/gcp/bigquery.py:
##########
@@ -2883,25 +3089,175 @@ def __enter__(self):
def __exit__(self, *args):
pass
+ class _ConvertDynamicRowDoFn(DoFn):
+ def __init__(
+ self,
+ schema,
+ union_field_names,
+ create_disposition=BigQueryDisposition.CREATE_IF_NEEDED,
+ write_disposition=BigQueryDisposition.WRITE_APPEND,
+ additional_bq_parameters=None,
+ test_client=None):
+ self.schema = schema
+ self.union_field_names = union_field_names
+ self.create_disposition = create_disposition
+ self.write_disposition = write_disposition
+ self.additional_bq_parameters = additional_bq_parameters
+ self.test_client = test_client
+ self.bigquery_wrapper = None
+
+ def start_bundle(self):
+ if self.create_disposition == BigQueryDisposition.CREATE_IF_NEEDED:
+ if self.test_client:
+ self.bigquery_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self.test_client)
+ else:
+ try:
+ self.bigquery_wrapper = bigquery_tools.BigQueryWrapper()
+ except Exception:
+ self.bigquery_wrapper = None
+
+ def _create_table_if_needed(self, dest, record_schema):
+ if self.create_disposition != BigQueryDisposition.CREATE_IF_NEEDED:
+ return
+ try:
+ table_ref = bigquery_tools.parse_table_reference(dest)
+ except ValueError:
+ return
+ if not table_ref.datasetId or not table_ref.tableId:
+ return
+
+ if self.bigquery_wrapper is None:
+ if self.test_client:
+ self.bigquery_wrapper = bigquery_tools.BigQueryWrapper(
+ client=self.test_client)
+ else:
+ try:
+ self.bigquery_wrapper = bigquery_tools.BigQueryWrapper()
+ except Exception:
+ return
+
+ project_id = (
+ table_ref.projectId or self.bigquery_wrapper._get_project_id())
+ str_table_ref = '%s:%s.%s' % (
+ project_id, table_ref.datasetId, table_ref.tableId)
+ if str_table_ref in _KNOWN_TABLES or dest in _KNOWN_TABLES:
+ return
+
+ table_schema = bigquery_tools.get_bq_tableschema(record_schema)
+ self.bigquery_wrapper.get_or_create_table(
+ project_id=project_id,
+ dataset_id=table_ref.datasetId,
+ table_id=table_ref.tableId,
+ schema=table_schema,
+ create_disposition=self.create_disposition,
+ write_disposition=self.write_disposition,
+ additional_create_parameters=self.additional_bq_parameters)
+ _KNOWN_TABLES.add(str_table_ref)
+ _KNOWN_TABLES.add(dest)
+
+ def process(self, row, *schema_side_inputs):
+ dest, dict_row = row[0], row[1]
+ record_schema = self.schema(dest, *schema_side_inputs)
+ self._create_table_if_needed(dest, record_schema)
+
+ if record_schema is None or record_schema == SCHEMA_AUTODETECT:
+ schema_json = ''
+ else:
+ schema_dict = bigquery_tools.get_dict_table_schema(record_schema)
+ schema_json = json.dumps(schema_dict)
+
+ record_row = bigquery_tools.beam_row_from_dict(dict_row, record_schema)
+ if self.union_field_names:
+ record_dict = record_row._asdict()
+ record_row = beam.Row(
+ **{
+ name: record_dict.get(name, None)
+ for name in self.union_field_names
+ })
+ yield beam.Row(
+ **{
+ StorageWriteToBigQuery.DESTINATION: dest,
+ StorageWriteToBigQuery.SCHEMA: schema_json,
Review Comment:
A full json schema seems like a lot of byte overhead per row. Should we
consider proto encoding or something to make this more efficient for shuffle
(streaming exactly once)?
##########
sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/providers/PortableBigQueryDestinations.java:
##########
@@ -23,29 +23,47 @@
import com.google.api.services.bigquery.model.Clustering;
import com.google.api.services.bigquery.model.TableConstraints;
+import com.google.api.services.bigquery.model.TableFieldSchema;
import com.google.api.services.bigquery.model.TableRow;
import com.google.api.services.bigquery.model.TableSchema;
+import java.util.ArrayList;
import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
import org.apache.avro.generic.GenericRecord;
import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.coders.NullableCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.extensions.avro.schemas.utils.AvroUtils;
import org.apache.beam.sdk.io.gcp.bigquery.AvroWriteRequest;
+import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers;
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils;
import org.apache.beam.sdk.io.gcp.bigquery.DynamicDestinations;
import org.apache.beam.sdk.io.gcp.bigquery.TableDestination;
import org.apache.beam.sdk.schemas.Schema;
import org.apache.beam.sdk.transforms.SerializableFunction;
import org.apache.beam.sdk.util.RowFilter;
import org.apache.beam.sdk.util.RowStringInterpolator;
+import org.apache.beam.sdk.values.KV;
import org.apache.beam.sdk.values.Row;
import org.apache.beam.sdk.values.ValueInSingleWindow;
+import
org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
import org.checkerframework.checker.nullness.qual.Nullable;
@Internal
-public class PortableBigQueryDestinations extends DynamicDestinations<Row,
String> {
+public class PortableBigQueryDestinations
+ extends DynamicDestinations<Row, KV<String, @Nullable String>> {
Review Comment:
Is this a breaking change for streaming compatibility?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
