use schema in template and add table verification

Author: Linchin

google/cloud/bigquery_storage_v1/writer.py

@@ -72,11 +72,13 @@ def _process_request_template(
     # The protobuf payload will be decoded as proto2 on the server side. The
     # schema is also specified as proto2. Hence we must clear proto3-only
     # features. This works since proto2 and proto3 are binary-compatible.
-    proto_descriptor = template_copy.proto_rows.writer_schema.proto_descriptor
-    for field in proto_descriptor.field:
-        field.ClearField("oneof_index")
-        field.ClearField("proto3_optional")
-    proto_descriptor.ClearField("oneof_decl")
+    oneof_field = template_copy._pb.WhichOneof("rows")
+    if oneof_field == "proto_rows":
+        proto_descriptor = template_copy.proto_rows.writer_schema.proto_descriptor
+        for field in proto_descriptor.field:
+            field.ClearField("oneof_index")
+            field.ClearField("proto3_optional")
+        proto_descriptor.ClearField("oneof_decl")
 
     return template_copy
 

samples/pyarrow/append_rows_with_arrow.py

@@ -24,43 +24,76 @@
 from google.cloud.bigquery_storage_v1.writer import AppendRowsStream
 
 
+TABLE_LENGTH = 1000
+
+BQ_SCHEMA = [
+    bigquery.SchemaField("bool_col", enums.SqlTypeNames.BOOLEAN),
+    bigquery.SchemaField("int64_col", enums.SqlTypeNames.INT64),
+    bigquery.SchemaField("float64_col", enums.SqlTypeNames.FLOAT64),
+    bigquery.SchemaField("numeric_col", enums.SqlTypeNames.NUMERIC),
+    bigquery.SchemaField("bignumeric_col", enums.SqlTypeNames.BIGNUMERIC),
+    bigquery.SchemaField("string_col", enums.SqlTypeNames.STRING),
+    bigquery.SchemaField("bytes_col", enums.SqlTypeNames.BYTES),
+    bigquery.SchemaField("date_col", enums.SqlTypeNames.DATE),
+    bigquery.SchemaField("datetime_col", enums.SqlTypeNames.DATETIME),
+    bigquery.SchemaField("time_col", enums.SqlTypeNames.TIME),
+    bigquery.SchemaField("timestamp_col", enums.SqlTypeNames.TIMESTAMP),
+    bigquery.SchemaField("geography_col", enums.SqlTypeNames.GEOGRAPHY),
+    bigquery.SchemaField(
+        "range_date_col", enums.SqlTypeNames.RANGE, range_element_type="DATE"
+    ),
+    bigquery.SchemaField(
+        "range_datetime_col",
+        enums.SqlTypeNames.RANGE,
+        range_element_type="DATETIME",
+    ),
+    bigquery.SchemaField(
+        "range_timestamp_col",
+        enums.SqlTypeNames.RANGE,
+        range_element_type="TIMESTAMP",
+    ),
+]
+
+PYARROW_SCHEMA = pa.schema(
+    [
+        pa.field("bool_col", pa.bool_()),
+        pa.field("int64_col", pa.int64()),
+        pa.field("float64_col", pa.float64()),
+        pa.field("numeric_col", pa.decimal128(38, scale=9)),
+        pa.field("bignumeric_col", pa.decimal256(76, scale=38)),
+        pa.field("string_col", pa.string()),
+        pa.field("bytes_col", pa.binary()),
+        pa.field("date_col", pa.date32()),
+        pa.field("datetime_col", pa.timestamp("us")),
+        pa.field("time_col", pa.time64("us")),
+        pa.field("timestamp_col", pa.timestamp("us")),
+        pa.field("geography_col", pa.string()),
+        pa.field(
+            "range_date_col",
+            pa.struct([("start", pa.date32()), ("end", pa.date32())]),
+        ),
+        pa.field(
+            "range_datetime_col",
+            pa.struct([("start", pa.timestamp("us")), ("end", pa.timestamp("us"))]),
+        ),
+        pa.field(
+            "range_timestamp_col",
+            pa.struct([("start", pa.timestamp("us")), ("end", pa.timestamp("us"))]),
+        ),
+    ]
+)
+
+
 def bqstorage_write_client():
     from google.cloud import bigquery_storage_v1
 
     return bigquery_storage_v1.BigQueryWriteClient()
 
 
 def make_table(project_id, dataset_id, bq_client):
-    schema = [
-        bigquery.SchemaField("bool_col", enums.SqlTypeNames.BOOLEAN),
-        bigquery.SchemaField("int64_col", enums.SqlTypeNames.INT64),
-        bigquery.SchemaField("float64_col", enums.SqlTypeNames.FLOAT64),
-        bigquery.SchemaField("numeric_col", enums.SqlTypeNames.NUMERIC),
-        bigquery.SchemaField("bignumeric_col", enums.SqlTypeNames.BIGNUMERIC),
-        bigquery.SchemaField("string_col", enums.SqlTypeNames.STRING),
-        bigquery.SchemaField("bytes_col", enums.SqlTypeNames.BYTES),
-        bigquery.SchemaField("date_col", enums.SqlTypeNames.DATE),
-        bigquery.SchemaField("datetime_col", enums.SqlTypeNames.DATETIME),
-        bigquery.SchemaField("time_col", enums.SqlTypeNames.TIME),
-        bigquery.SchemaField("timestamp_col", enums.SqlTypeNames.TIMESTAMP),
-        bigquery.SchemaField("geography_col", enums.SqlTypeNames.GEOGRAPHY),
-        bigquery.SchemaField(
-            "range_date_col", enums.SqlTypeNames.RANGE, range_element_type="DATE"
-        ),
-        bigquery.SchemaField(
-            "range_datetime_col",
-            enums.SqlTypeNames.RANGE,
-            range_element_type="DATETIME",
-        ),
-        bigquery.SchemaField(
-            "range_timestamp_col",
-            enums.SqlTypeNames.RANGE,
-            range_element_type="TIMESTAMP",
-        ),
-    ]
     table_id = "append_rows_w_arrow_test"
     table_id_full = f"{project_id}.{dataset_id}.{table_id}"
-    bq_table = bigquery.Table(table_id_full, schema=schema)
+    bq_table = bigquery.Table(table_id_full, schema=BQ_SCHEMA)
     created_table = bq_client.create_table(bq_table)
 
     return created_table
@@ -71,14 +104,19 @@ def create_stream(bqstorage_write_client, table):
     request_template = gapic_types.AppendRowsRequest()
     request_template.write_stream = stream_name
 
+    # Add schema to the template.
+    arrow_data = gapic_types.AppendRowsRequest.ArrowData()
+    arrow_data.writer_schema.serialized_schema = PYARROW_SCHEMA.serialize().to_pybytes()
+    request_template.arrow_rows = arrow_data
+
     append_rows_stream = AppendRowsStream(
         bqstorage_write_client,
         request_template,
     )
     return append_rows_stream
 
 
-def generate_write_request_with_pyarrow(row_num=10):
+def generate_write_request_with_pyarrow(num_rows=TABLE_LENGTH):
     date_1 = datetime.date(2020, 10, 1)
     date_2 = datetime.date(2021, 10, 1)
 
@@ -94,7 +132,7 @@ def generate_write_request_with_pyarrow(row_num=10):
 
     # Pandas Dataframe.
     rows = []
-    for i in range(row_num):
+    for i in range(num_rows):
         row = {
             "bool_col": True,
             "int64_col": i,
@@ -116,57 +154,40 @@ def generate_write_request_with_pyarrow(row_num=10):
     df = pd.DataFrame(rows)
 
     # Dataframe to PyArrow Table.
-    schema = pa.schema(
-        [
-            pa.field("bool_col", pa.bool_()),
-            pa.field("int64_col", pa.int64()),
-            pa.field("float64_col", pa.float64()),
-            pa.field("numeric_col", pa.decimal128(38, scale=9)),
-            pa.field("bignumeric_col", pa.decimal256(76, scale=38)),
-            pa.field("string_col", pa.string()),
-            pa.field("bytes_col", pa.binary()),
-            pa.field("date_col", pa.date32()),
-            pa.field("datetime_col", pa.timestamp("us")),
-            pa.field("time_col", pa.time64("us")),
-            pa.field("timestamp_col", pa.timestamp("us")),
-            pa.field("geography_col", pa.string()),
-            pa.field(
-                "range_date_col",
-                pa.struct([("start", pa.date32()), ("end", pa.date32())]),
-            ),
-            pa.field(
-                "range_datetime_col",
-                pa.struct([("start", pa.timestamp("us")), ("end", pa.timestamp("us"))]),
-            ),
-            pa.field(
-                "range_timestamp_col",
-                pa.struct([("start", pa.timestamp("us")), ("end", pa.timestamp("us"))]),
-            ),
-        ]
-    )
-    table = pa.Table.from_pandas(df, schema=schema)
+    table = pa.Table.from_pandas(df, schema=PYARROW_SCHEMA)
 
     # Construct request.
     request = gapic_types.AppendRowsRequest()
-    arrow_data = gapic_types.AppendRowsRequest.ArrowData()
-    arrow_data.writer_schema.serialized_schema = table.schema.serialize().to_pybytes()
-    arrow_data.rows.serialized_record_batch = (
+    request.arrow_rows.rows.serialized_record_batch = (
         table.to_batches()[0].serialize().to_pybytes()
     )
-    request.arrow_rows = arrow_data
     return request
 
 
 def append_rows(bqstorage_write_client, table):
     append_rows_stream = create_stream(bqstorage_write_client, table)
-    request = generate_write_request_with_pyarrow(row_num=20)
+    request = generate_write_request_with_pyarrow(num_rows=TABLE_LENGTH)
 
     response_future = append_rows_stream.send(request)
     print(response_future.result())
 
 
+def verify_table(client, table):
+    bq_table = client.get_table(table)
+
+    # Verify table schema.
+    assert bq_table.schema == BQ_SCHEMA
+
+    # Verify table size.
+    query = client.query(f"SELECT COUNT(1) FROM `{bq_table}`;")
+    query_result = query.result().to_dataframe()
+    assert query_result.iloc[0, 0] == TABLE_LENGTH
+
+
 def main(project_id, dataset):
     write_client = bqstorage_write_client()
     bq_client = bigquery.Client()
     table = make_table(project_id, dataset.dataset_id, bq_client)
+
     append_rows(write_client, table)
+    verify_table(bq_client, table)