From b127a4351589db0581ff7921912f08bb32633c72 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Knut=20Olav=20L=C3=B8ite?= Date: Thu, 10 Sep 2026 12:26:03 +0200 Subject: [PATCH] perf(spanner): optimize query result decoding and type resolution Improve query result decoding performance in StreamedResultSet and simplify type decoder resolution in _helpers. * Optimize row construction in StreamedResultSet._merge_values: - Process complete rows in batch using direct column indexing rather than iterating cell-by-cell with repeated append() calls and boundary checks. - For lazy decoding, slice raw values directly into rows. - Add an explicit guard for width == 0 to prevent unbounded buffer growth on empty schemas. - Separate partial row boundary handling from complete row decoding by extracting _append_to_current_row, _decode_eager_rows, and _decode_lazy_rows, removing duplicated cell-appending logic. * Optimize type resolution in _helpers: - Replace the if/elif cascade in _get_type_decoder with an integer-keyed _SCALAR_DECODERS lookup table. - Pre-allocate static decoder callables at module scope to avoid lambda recreation when resolving column metadata. --- .../cloud/spanner_v1/_async/streamed.py | 116 ++++-- .../google/cloud/spanner_v1/_helpers.py | 64 ++-- .../google/cloud/spanner_v1/streamed.py | 117 +++++-- .../tests/unit/_async/test_streamed.py | 329 +++++++++++++++++- .../tests/unit/test__helpers.py | 101 ++++++ .../tests/unit/test_streamed.py | 315 ++++++++++++++++- 6 files changed, 939 insertions(+), 103 deletions(-) diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py index c47cc0ef0a17..c0cfb6405b84 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py @@ -120,48 +120,102 @@ def _merge_chunk(self, value): self._pending_chunk = None return merged + def _append_to_current_row(self, values): + """Append cells to the in-progress partial row.""" + if self._lazy_decode: + self._current_row.extend(values) + else: + decoders = self._decoders + start_column = len(self._current_row) + for column_offset, value in enumerate(values): + if value.HasField("null_value"): + self._current_row.append(None) + else: + self._current_row.append( + decoders[start_column + column_offset](value) + ) + + def _decode_lazy_rows(self, values, values_offset, batch_end, width): + """Slice raw protobuf values into rows for lazy decoding.""" + if width == 1: + self._rows.extend([[value] for value in values[values_offset:batch_end]]) + else: + self._rows.extend( + [ + values[row_start : row_start + width] + for row_start in range(values_offset, batch_end, width) + ] + ) + + def _decode_eager_rows(self, values, values_offset, batch_end, width): + """Decode complete row batches into typed Python values.""" + if width == 1: + decoder = self._decoders[0] + self._rows.extend( + [ + [None if value.HasField("null_value") else decoder(value)] + for value in values[values_offset:batch_end] + ] + ) + else: + decoders = self._decoders + rows_append = self._rows.append + column_indices = list(range(width)) + for row_start in range(values_offset, batch_end, width): + rows_append( + [ + None + if values[row_start + column_index].HasField("null_value") + else decoders[column_index](values[row_start + column_index]) + for column_index in column_indices + ] + ) + def _merge_values(self, values): """Merge values into rows. :type values: list of :class:`~google.protobuf.struct_pb2.Value` :param values: non-chunked values from partial result set. """ - decoders = self._decoders + if not values: + return + width = len(self.fields) - index = len(self._current_row) - current_row = self._current_row - rows = self._rows + if width == 0: + return + + values_offset = 0 + total_values = len(values) + + # 1. Complete pending partial row from previous chunk (if any) + if self._current_row: + needed = width - len(self._current_row) + fill_count = min(needed, total_values) + self._append_to_current_row(values[:fill_count]) + values_offset = fill_count + if len(self._current_row) == width: + self._rows.append(self._current_row) + self._current_row = [] + else: + return + + remaining_values = total_values - values_offset + if remaining_values == 0: + return - current_row_append = current_row.append - rows_append = rows.append + row_count = remaining_values // width + full_values_count = row_count * width + batch_end = values_offset + full_values_count + # 2. Batch-decode complete rows if self._lazy_decode: - for value in values: - current_row_append(value) - index += 1 - if index == width: - rows_append(current_row) - current_row = [] - current_row_append = current_row.append - index = 0 + self._decode_lazy_rows(values, values_offset, batch_end, width) else: - for value in values: - # Note: We manually check value.HasField("null_value") here instead of - # wrapping every decoder in _parse_nullable to avoid the overhead of - # an extra Python function call layer for every cell value decoded in this loop. - # If the nullable check logic is updated in _parse_nullable, update this check. - if value.HasField("null_value"): - current_row_append(None) - else: - current_row_append(decoders[index](value)) - index += 1 - if index == width: - rows_append(current_row) - current_row = [] - current_row_append = current_row.append - index = 0 - - self._current_row = current_row + self._decode_eager_rows(values, values_offset, batch_end, width) + + # 3. Buffer trailing partial row remainder for the next chunk (if any) + if remaining_values > full_values_count: + self._append_to_current_row(values[batch_end:]) @CrossSync.convert async def _consume_next(self): diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py index 06b137db1e28..d68028705a49 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py @@ -505,47 +505,28 @@ def _get_type_decoder(field_type, field_name, column_info=None): """ type_code = field_type.code - # Note: STRING and BOOL use operator.attrgetter because direct attribute extraction - # is faster in Python. Other types require type transformation, so they use lambdas. - if type_code == TypeCode.STRING: - return operator.attrgetter("string_value") - elif type_code == TypeCode.BYTES: - return lambda value_pb: value_pb.string_value.encode("utf8") - elif type_code == TypeCode.BOOL: - return operator.attrgetter("bool_value") - elif type_code == TypeCode.INT64: - return lambda value_pb: int(value_pb.string_value) - elif type_code == TypeCode.FLOAT64: - return _parse_float - elif type_code == TypeCode.FLOAT32: - return _parse_float - elif type_code == TypeCode.DATE: - return lambda value_pb: _date_fromisoformat(value_pb.string_value) - elif type_code == TypeCode.TIMESTAMP: - return _parse_timestamp - elif type_code == TypeCode.NUMERIC: - return lambda value_pb: _Decimal(value_pb.string_value) - elif type_code == TypeCode.JSON: - return lambda value_pb: _json_from_str(value_pb.string_value) - elif type_code == TypeCode.UUID: - return lambda value_pb: _uuid_UUID(value_pb.string_value) - elif type_code == TypeCode.PROTO: + try: + type_code_integer = int(type_code) + except (TypeError, ValueError): + type_code_integer = None + + if type_code_integer in _SCALAR_DECODERS: + return _SCALAR_DECODERS[type_code_integer] + elif type_code_integer == _PROTO_TYPE_CODE: return lambda value_pb: _parse_proto(value_pb, column_info, field_name) - elif type_code == TypeCode.ENUM: + elif type_code_integer == _ENUM_TYPE_CODE: return lambda value_pb: _parse_proto_enum(value_pb, column_info, field_name) - elif type_code == TypeCode.ARRAY: + elif type_code_integer == _ARRAY_TYPE_CODE: element_decoder = _get_type_decoder( field_type.array_element_type, field_name, column_info ) return lambda value_pb: _parse_array(value_pb, element_decoder) - elif type_code == TypeCode.STRUCT: + elif type_code_integer == _STRUCT_TYPE_CODE: element_decoders = [ _get_type_decoder(item_field.type_, field_name, column_info) for item_field in field_type.struct_type.fields ] return lambda value_pb: _parse_struct(value_pb, element_decoders) - elif type_code == TypeCode.INTERVAL: - return _parse_interval else: raise ValueError("Unknown type: %s" % (field_type,)) @@ -702,6 +683,29 @@ def _parse_interval(value_pb): return Interval.from_str(value_pb) +# Note: STRING and BOOL use operator.attrgetter because direct attribute extraction +# is faster in Python. Other types require type transformation, so they use lambdas. +_SCALAR_DECODERS = { + int(TypeCode.STRING): operator.attrgetter("string_value"), + int(TypeCode.BYTES): lambda value_pb: value_pb.string_value.encode("utf8"), + int(TypeCode.BOOL): operator.attrgetter("bool_value"), + int(TypeCode.INT64): lambda value_pb: int(value_pb.string_value), + int(TypeCode.FLOAT64): _parse_float, + int(TypeCode.FLOAT32): _parse_float, + int(TypeCode.DATE): lambda value_pb: _date_fromisoformat(value_pb.string_value), + int(TypeCode.TIMESTAMP): _parse_timestamp, + int(TypeCode.NUMERIC): lambda value_pb: _Decimal(value_pb.string_value), + int(TypeCode.JSON): lambda value_pb: _json_from_str(value_pb.string_value), + int(TypeCode.UUID): lambda value_pb: _uuid_UUID(value_pb.string_value), + int(TypeCode.INTERVAL): _parse_interval, +} + +_PROTO_TYPE_CODE = int(TypeCode.PROTO) +_ENUM_TYPE_CODE = int(TypeCode.ENUM) +_ARRAY_TYPE_CODE = int(TypeCode.ARRAY) +_STRUCT_TYPE_CODE = int(TypeCode.STRUCT) + + class _SessionWrapper(object): """Base class for objects wrapping a session. diff --git a/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py b/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py index a92f008f5e32..e5a77b5ef050 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py @@ -108,44 +108,101 @@ def _merge_chunk(self, value): self._pending_chunk = None return merged + def _append_to_current_row(self, values): + """Append cells to the in-progress partial row.""" + if self._lazy_decode: + self._current_row.extend(values) + else: + decoders = self._decoders + start_column = len(self._current_row) + for column_offset, value in enumerate(values): + if value.HasField("null_value"): + self._current_row.append(None) + else: + self._current_row.append( + decoders[start_column + column_offset](value) + ) + + def _decode_lazy_rows(self, values, values_offset, batch_end, width): + """Slice raw protobuf values into rows for lazy decoding.""" + if width == 1: + self._rows.extend([[value] for value in values[values_offset:batch_end]]) + else: + self._rows.extend( + [ + values[row_start : row_start + width] + for row_start in range(values_offset, batch_end, width) + ] + ) + + def _decode_eager_rows(self, values, values_offset, batch_end, width): + """Decode complete row batches into typed Python values.""" + if width == 1: + decoder = self._decoders[0] + self._rows.extend( + [ + [None if value.HasField("null_value") else decoder(value)] + for value in values[values_offset:batch_end] + ] + ) + else: + decoders = self._decoders + rows_append = self._rows.append + column_indices = list(range(width)) + for row_start in range(values_offset, batch_end, width): + rows_append( + [ + None + if values[row_start + column_index].HasField("null_value") + else decoders[column_index](values[row_start + column_index]) + for column_index in column_indices + ] + ) + def _merge_values(self, values): """Merge values into rows. :type values: list of :class:`~google.protobuf.struct_pb2.Value` :param values: non-chunked values from partial result set.""" - decoders = self._decoders + if not values: + return + width = len(self.fields) - index = len(self._current_row) - current_row = self._current_row - rows = self._rows - current_row_append = current_row.append - rows_append = rows.append + if width == 0: + return + + values_offset = 0 + total_values = len(values) + + # 1. Complete pending partial row from previous chunk (if any) + if self._current_row: + needed = width - len(self._current_row) + fill_count = min(needed, total_values) + self._append_to_current_row(values[:fill_count]) + values_offset = fill_count + if len(self._current_row) == width: + self._rows.append(self._current_row) + self._current_row = [] + else: + return + + remaining_values = total_values - values_offset + if remaining_values == 0: + return + + row_count = remaining_values // width + full_values_count = row_count * width + batch_end = values_offset + full_values_count + + # 2. Batch-decode complete rows if self._lazy_decode: - for value in values: - current_row_append(value) - index += 1 - if index == width: - rows_append(current_row) - current_row = [] - current_row_append = current_row.append - index = 0 + self._decode_lazy_rows(values, values_offset, batch_end, width) else: - for value in values: - # Note: We manually check value.HasField("null_value") here instead of - # wrapping every decoder in _parse_nullable to avoid the overhead of - # an extra Python function call layer for every cell value decoded in this loop. - # If the nullable check logic is updated in _parse_nullable, update this check. - if value.HasField("null_value"): - current_row_append(None) - else: - current_row_append(decoders[index](value)) - index += 1 - if index == width: - rows_append(current_row) - current_row = [] - current_row_append = current_row.append - index = 0 - self._current_row = current_row + self._decode_eager_rows(values, values_offset, batch_end, width) + + # 3. Buffer trailing partial row remainder for the next chunk (if any) + if remaining_values > full_values_count: + self._append_to_current_row(values[batch_end:]) def _consume_next(self): """Consume the next partial result set from the stream. diff --git a/packages/google-cloud-spanner/tests/unit/_async/test_streamed.py b/packages/google-cloud-spanner/tests/unit/_async/test_streamed.py index f3ec2bb4d0cb..15114616fc26 100644 --- a/packages/google-cloud-spanner/tests/unit/_async/test_streamed.py +++ b/packages/google-cloud-spanner/tests/unit/_async/test_streamed.py @@ -1124,7 +1124,7 @@ async def test___iter___large_batch(self): self._make_scalar_field("name", TypeCode.STRING), ] metadata = self._make_result_set_metadata(fields) - expected_rows = [[i, f"name_{i}"] for i in range(500)] + expected_rows = [[index, f"name_{index}"] for index in range(500)] values = [self._make_value(cell) for row in expected_rows for cell in row] result_set = self._make_partial_result_set(values, metadata=metadata) @@ -1143,7 +1143,7 @@ async def test___iter___stepwise_consumption(self): self._make_scalar_field("name", TypeCode.STRING), ] metadata = self._make_result_set_metadata(fields) - expected_rows = [[i, f"name_{i}"] for i in range(20)] + expected_rows = [[index, f"name_{index}"] for index in range(20)] values = [self._make_value(cell) for row in expected_rows for cell in row] result_set = self._make_partial_result_set(values, metadata=metadata) @@ -1164,8 +1164,8 @@ async def test___iter___stepwise_across_chunks(self): self._make_scalar_field("name", TypeCode.STRING), ] metadata = self._make_result_set_metadata(fields) - chunk1_rows = [[i, f"name_{i}"] for i in range(10)] - chunk2_rows = [[i, f"name_{i}"] for i in range(10, 20)] + chunk1_rows = [[index, f"name_{index}"] for index in range(10)] + chunk2_rows = [[index, f"name_{index}"] for index in range(10, 20)] values1 = [self._make_value(cell) for row in chunk1_rows for cell in row] values2 = [self._make_value(cell) for row in chunk2_rows for cell in row] @@ -1190,7 +1190,7 @@ async def test___iter___early_break(self): self._make_scalar_field("name", TypeCode.STRING), ] metadata = self._make_result_set_metadata(fields) - expected_rows = [[i, f"name_{i}"] for i in range(10)] + expected_rows = [[index, f"name_{index}"] for index in range(10)] values = [self._make_value(cell) for row in expected_rows for cell in row] result_set = self._make_partial_result_set(values, metadata=metadata) @@ -1213,7 +1213,7 @@ async def test___iter___mid_stream_error(self): self._make_scalar_field("name", TypeCode.STRING), ] metadata = self._make_result_set_metadata(fields) - chunk1_rows = [[i, f"name_{i}"] for i in range(5)] + chunk1_rows = [[index, f"name_{index}"] for index in range(5)] values1 = [self._make_value(cell) for row in chunk1_rows for cell in row] result_set1 = self._make_partial_result_set(values1, metadata=metadata) @@ -1230,6 +1230,323 @@ async def mock_iterator(): self.assertEqual(consumed, chunk1_rows) self.assertIn("Stream error midway", str(context.exception)) + @CrossSync.pytest + async def test_decode_rows_direct_width_one(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [self._make_scalar_field("count", TypeCode.INT64)] + metadata = self._make_result_set_metadata(fields) + expected_rows = [[42]] + values = [self._make_value(42)] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator) + found = [row async for row in streamed] + self.assertEqual(found, expected_rows) + self.assertTrue(streamed._done) + + @CrossSync.pytest + async def test_decode_rows_direct_multi_column(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + self._make_scalar_field("active", TypeCode.BOOL), + ] + metadata = self._make_result_set_metadata(fields) + expected_rows = [[1, "alpha", True], [2, "beta", False]] + values = [self._make_value(cell) for row in expected_rows for cell in row] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator) + found = [row async for row in streamed] + self.assertEqual(found, expected_rows) + self.assertTrue(streamed._done) + + @CrossSync.pytest + async def test_decode_rows_direct_with_null_values(self): + from google.protobuf.struct_pb2 import NULL_VALUE, Value + + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + values = [ + self._make_value(1), + Value(null_value=NULL_VALUE), + self._make_value(2), + self._make_value("Alice"), + ] + expected_rows = [[1, None], [2, "Alice"]] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator) + found = [row async for row in streamed] + self.assertEqual(found, expected_rows) + + @CrossSync.pytest + async def test_decode_rows_direct_lazy_decode_width_one(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [self._make_scalar_field("id", TypeCode.INT64)] + metadata = self._make_result_set_metadata(fields) + raw_value = self._make_value(100) + values = [raw_value] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator, lazy_decode=True) + found = [row async for row in streamed] + self.assertEqual(len(found), 1) + self.assertEqual(found[0], [raw_value]) + self.assertEqual(streamed.decode_row(found[0]), [100]) + + @CrossSync.pytest + async def test_decode_rows_direct_lazy_decode_multi_column(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + value_id = self._make_value(1) + value_name = self._make_value("test") + values = [value_id, value_name] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator, lazy_decode=True) + found = [row async for row in streamed] + self.assertEqual(len(found), 1) + self.assertEqual(found[0], [value_id, value_name]) + self.assertEqual(streamed.decode_row(found[0]), [1, "test"]) + + @CrossSync.pytest + async def test_decode_rows_direct_trailing_partial_row(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + values1 = [ + self._make_value(1), + self._make_value("a"), + self._make_value(2), + ] + values2 = [ + self._make_value("b"), + self._make_value(3), + self._make_value("c"), + ] + + result_set1 = self._make_partial_result_set(values1, metadata=metadata) + result_set2 = self._make_partial_result_set(values2, last=True) + iterator = _MockCancellableIterator(result_set1, result_set2) + streamed = self._make_one(iterator) + found = [row async for row in streamed] + self.assertEqual( + found, + [ + [1, "a"], + [2, "b"], + [3, "c"], + ], + ) + + @CrossSync.pytest + async def test_decode_rows_direct_trailing_partial_row_lazy_decode(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + value1 = self._make_value(1) + value_a = self._make_value("a") + value2 = self._make_value(2) + value_b = self._make_value("b") + + result_set1 = self._make_partial_result_set( + [value1, value_a, value2], metadata=metadata + ) + result_set2 = self._make_partial_result_set([value_b], last=True) + iterator = _MockCancellableIterator(result_set1, result_set2) + streamed = self._make_one(iterator, lazy_decode=True) + found = [row async for row in streamed] + self.assertEqual(len(found), 2) + self.assertEqual(found[0], [value1, value_a]) + self.assertEqual(found[1], [value2, value_b]) + self.assertEqual(streamed.decode_row(found[0]), [1, "a"]) + self.assertEqual(streamed.decode_row(found[1]), [2, "b"]) + + @CrossSync.pytest + async def test_decode_rows_direct_trailing_partial_row_with_null(self): + from google.protobuf.struct_pb2 import NULL_VALUE, Value + + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + values1 = [ + self._make_value(1), + self._make_value("a"), + Value(null_value=NULL_VALUE), + ] + values2 = [ + self._make_value("b"), + ] + + result_set1 = self._make_partial_result_set(values1, metadata=metadata) + result_set2 = self._make_partial_result_set(values2, last=True) + iterator = _MockCancellableIterator(result_set1, result_set2) + streamed = self._make_one(iterator) + found = [row async for row in streamed] + self.assertEqual( + found, + [ + [1, "a"], + [None, "b"], + ], + ) + + @CrossSync.pytest + async def test_decode_rows_direct_prefix_partial_row_with_null(self): + from google.protobuf.struct_pb2 import NULL_VALUE, Value + + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + values1 = [ + self._make_value(1), + self._make_value("a"), + self._make_value(2), + ] + values2 = [ + Value(null_value=NULL_VALUE), + self._make_value(3), + self._make_value("c"), + ] + + result_set1 = self._make_partial_result_set(values1, metadata=metadata) + result_set2 = self._make_partial_result_set(values2, last=True) + iterator = _MockCancellableIterator(result_set1, result_set2) + streamed = self._make_one(iterator) + found = [row async for row in streamed] + self.assertEqual( + found, + [ + [1, "a"], + [2, None], + [3, "c"], + ], + ) + + @CrossSync.pytest + async def test_decode_rows_direct_width_one_with_null(self): + from google.protobuf.struct_pb2 import NULL_VALUE, Value + + from google.cloud.spanner_v1 import TypeCode + + fields = [self._make_scalar_field("id", TypeCode.INT64)] + metadata = self._make_result_set_metadata(fields) + values = [ + self._make_value(1), + Value(null_value=NULL_VALUE), + self._make_value(2), + ] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator) + found = [row async for row in streamed] + self.assertEqual(found, [[1], [None], [2]]) + + @CrossSync.pytest + async def test_decode_rows_direct_three_chunk_split_row(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + self._make_scalar_field("active", TypeCode.BOOL), + ] + metadata = self._make_result_set_metadata(fields) + prs1 = self._make_partial_result_set([self._make_value(1)], metadata=metadata) + prs2 = self._make_partial_result_set([self._make_value("alpha")]) + prs3 = self._make_partial_result_set([self._make_value(True)], last=True) + + iterator = _MockCancellableIterator(prs1, prs2, prs3) + streamed = self._make_one(iterator) + found = [row async for row in streamed] + self.assertEqual(found, [[1, "alpha", True]]) + + @CrossSync.pytest + async def test_decode_rows_direct_three_chunk_split_row_lazy_decode(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + self._make_scalar_field("active", TypeCode.BOOL), + ] + metadata = self._make_result_set_metadata(fields) + val1 = self._make_value(1) + val2 = self._make_value("alpha") + val3 = self._make_value(True) + prs1 = self._make_partial_result_set([val1], metadata=metadata) + prs2 = self._make_partial_result_set([val2]) + prs3 = self._make_partial_result_set([val3], last=True) + + iterator = _MockCancellableIterator(prs1, prs2, prs3) + streamed = self._make_one(iterator, lazy_decode=True) + found = [row async for row in streamed] + self.assertEqual(found, [[val1, val2, val3]]) + self.assertEqual(streamed.decode_row(found[0]), [1, "alpha", True]) + + @CrossSync.pytest + async def test_decode_rows_direct_empty_values(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [self._make_scalar_field("id", TypeCode.INT64)] + metadata = self._make_result_set_metadata(fields) + + result_set1 = self._make_partial_result_set([], metadata=metadata) + result_set2 = self._make_partial_result_set([self._make_value(1)], last=True) + iterator = _MockCancellableIterator(result_set1, result_set2) + streamed = self._make_one(iterator) + found = [row async for row in streamed] + self.assertEqual(found, [[1]]) + + @CrossSync.pytest + async def test_merge_values_zero_fields(self): + from google.cloud.spanner_v1 import ResultSetMetadata, StructType + + metadata = ResultSetMetadata(row_type=StructType(fields=[])) + result_set = self._make_partial_result_set([], metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator) + _ = [row async for row in streamed] + streamed._merge_values([self._make_value(1)]) + self.assertEqual(streamed._rows, []) + class _MockCancellableIterator(object): cancel_calls = 0 diff --git a/packages/google-cloud-spanner/tests/unit/test__helpers.py b/packages/google-cloud-spanner/tests/unit/test__helpers.py index 0a6e9594b167..e5d63584bcc9 100644 --- a/packages/google-cloud-spanner/tests/unit/test__helpers.py +++ b/packages/google-cloud-spanner/tests/unit/test__helpers.py @@ -1862,3 +1862,104 @@ def test_large_values(self): self.assertEqual(result.months, case["expected_months"]) self.assertEqual(result.days, case["expected_days"]) self.assertEqual(result.nanos, case["expected_nanos"]) + + +class Test_get_type_decoder(unittest.TestCase): + def _callFUT(self, *args, **kwargs): + from google.cloud.spanner_v1._helpers import _get_type_decoder + + return _get_type_decoder(*args, **kwargs) + + def test_scalar_decoders(self): + import datetime + import decimal + import uuid + + from google.protobuf.struct_pb2 import Value + + from google.cloud.spanner_v1 import Type, TypeCode + from google.cloud.spanner_v1._helpers import _SCALAR_DECODERS + from google.cloud.spanner_v1.data_types import Interval, JsonObject + + test_cases = [ + (TypeCode.STRING, Value(string_value="hello"), "hello"), + (TypeCode.BYTES, Value(string_value="bytes"), b"bytes"), + (TypeCode.BOOL, Value(bool_value=True), True), + (TypeCode.INT64, Value(string_value="42"), 42), + (TypeCode.FLOAT64, Value(string_value="3.14"), 3.14), + (TypeCode.FLOAT32, Value(string_value="2.5"), 2.5), + ( + TypeCode.DATE, + Value(string_value="2026-03-15"), + datetime.date(2026, 3, 15), + ), + ( + TypeCode.TIMESTAMP, + Value(string_value="2026-03-15T12:00:00Z"), + datetime.datetime(2026, 3, 15, 12, 0, tzinfo=datetime.timezone.utc), + ), + (TypeCode.NUMERIC, Value(string_value="99.99"), decimal.Decimal("99.99")), + (TypeCode.JSON, Value(string_value='{"a": 1}'), JsonObject({"a": 1})), + ( + TypeCode.UUID, + Value(string_value="12345678-1234-5678-1234-567812345678"), + uuid.UUID("12345678-1234-5678-1234-567812345678"), + ), + (TypeCode.INTERVAL, Value(string_value="P1Y"), Interval.from_str("P1Y")), + ] + for type_code, sample_value_pb, expected_result in test_cases: + field_type = Type(code=type_code) + decoder = self._callFUT(field_type, "column_name") + self.assertIs(decoder, _SCALAR_DECODERS[int(type_code)]) + self.assertEqual(decoder(sample_value_pb), expected_result) + + def test_proto_and_enum(self): + from google.protobuf.struct_pb2 import Value + + from google.cloud.spanner_v1 import Type, TypeCode + + proto_type = Type(code=TypeCode.PROTO) + proto_decoder = self._callFUT(proto_type, "proto_column") + self.assertTrue(callable(proto_decoder)) + + enum_type = Type(code=TypeCode.ENUM) + enum_decoder = self._callFUT(enum_type, "enum_column") + self.assertTrue(callable(enum_decoder)) + self.assertEqual(enum_decoder(Value(string_value="1")), 1) + + def test_array_and_struct(self): + from google.cloud.spanner_v1 import StructType, Type, TypeCode + + array_type = Type( + code=TypeCode.ARRAY, + array_element_type=Type(code=TypeCode.STRING), + ) + array_decoder = self._callFUT(array_type, "array_column") + self.assertTrue(callable(array_decoder)) + + struct_field = StructType.Field( + name="subfield", type_=Type(code=TypeCode.STRING) + ) + struct_type = Type( + code=TypeCode.STRUCT, + struct_type=StructType(fields=[struct_field]), + ) + struct_decoder = self._callFUT(struct_type, "struct_column") + self.assertTrue(callable(struct_decoder)) + + def test_unknown_and_unspecified_types(self): + from unittest import mock + + from google.cloud.spanner_v1 import Type, TypeCode + + unspecified_type = Type(code=TypeCode.TYPE_CODE_UNSPECIFIED) + with self.assertRaises(ValueError): + self._callFUT(unspecified_type, "unspecified") + + unknown_type = mock.Mock(code=999) + with self.assertRaises(ValueError): + self._callFUT(unknown_type, "unknown") + + invalid_code_type = mock.Mock(code="invalid") + with self.assertRaises(ValueError): + self._callFUT(invalid_code_type, "invalid") diff --git a/packages/google-cloud-spanner/tests/unit/test_streamed.py b/packages/google-cloud-spanner/tests/unit/test_streamed.py index 3d3ba709145d..4c8ff04d3193 100644 --- a/packages/google-cloud-spanner/tests/unit/test_streamed.py +++ b/packages/google-cloud-spanner/tests/unit/test_streamed.py @@ -1043,7 +1043,7 @@ def test___iter___large_batch(self): self._make_scalar_field("name", TypeCode.STRING), ] metadata = self._make_result_set_metadata(fields) - expected_rows = [[i, f"name_{i}"] for i in range(500)] + expected_rows = [[index, f"name_{index}"] for index in range(500)] values = [self._make_value(cell) for row in expected_rows for cell in row] result_set = self._make_partial_result_set(values, metadata=metadata) @@ -1061,7 +1061,7 @@ def test___iter___stepwise_consumption(self): self._make_scalar_field("name", TypeCode.STRING), ] metadata = self._make_result_set_metadata(fields) - expected_rows = [[i, f"name_{i}"] for i in range(20)] + expected_rows = [[index, f"name_{index}"] for index in range(20)] values = [self._make_value(cell) for row in expected_rows for cell in row] result_set = self._make_partial_result_set(values, metadata=metadata) @@ -1081,8 +1081,8 @@ def test___iter___stepwise_across_chunks(self): self._make_scalar_field("name", TypeCode.STRING), ] metadata = self._make_result_set_metadata(fields) - chunk1_rows = [[i, f"name_{i}"] for i in range(10)] - chunk2_rows = [[i, f"name_{i}"] for i in range(10, 20)] + chunk1_rows = [[index, f"name_{index}"] for index in range(10)] + chunk2_rows = [[index, f"name_{index}"] for index in range(10, 20)] values1 = [self._make_value(cell) for row in chunk1_rows for cell in row] values2 = [self._make_value(cell) for row in chunk2_rows for cell in row] @@ -1106,7 +1106,7 @@ def test___iter___early_break(self): self._make_scalar_field("name", TypeCode.STRING), ] metadata = self._make_result_set_metadata(fields) - expected_rows = [[i, f"name_{i}"] for i in range(10)] + expected_rows = [[index, f"name_{index}"] for index in range(10)] values = [self._make_value(cell) for row in expected_rows for cell in row] result_set = self._make_partial_result_set(values, metadata=metadata) @@ -1128,7 +1128,7 @@ def test___iter___mid_stream_error(self): self._make_scalar_field("name", TypeCode.STRING), ] metadata = self._make_result_set_metadata(fields) - chunk1_rows = [[i, f"name_{i}"] for i in range(5)] + chunk1_rows = [[index, f"name_{index}"] for index in range(5)] values1 = [self._make_value(cell) for row in chunk1_rows for cell in row] result_set1 = self._make_partial_result_set(values1, metadata=metadata) @@ -1145,6 +1145,309 @@ def mock_iterator(): self.assertEqual(consumed, chunk1_rows) self.assertIn("Stream error midway", str(context.exception)) + def test_decode_rows_direct_width_one(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [self._make_scalar_field("count", TypeCode.INT64)] + metadata = self._make_result_set_metadata(fields) + expected_rows = [[42]] + values = [self._make_value(42)] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator) + found = list(streamed) + self.assertEqual(found, expected_rows) + self.assertTrue(streamed._done) + + def test_decode_rows_direct_multi_column(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + self._make_scalar_field("active", TypeCode.BOOL), + ] + metadata = self._make_result_set_metadata(fields) + expected_rows = [[1, "alpha", True], [2, "beta", False]] + values = [self._make_value(cell) for row in expected_rows for cell in row] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator) + found = list(streamed) + self.assertEqual(found, expected_rows) + self.assertTrue(streamed._done) + + def test_decode_rows_direct_with_null_values(self): + from google.protobuf.struct_pb2 import NULL_VALUE, Value + + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + values = [ + self._make_value(1), + Value(null_value=NULL_VALUE), + self._make_value(2), + self._make_value("Alice"), + ] + expected_rows = [[1, None], [2, "Alice"]] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator) + found = list(streamed) + self.assertEqual(found, expected_rows) + + def test_decode_rows_direct_lazy_decode_width_one(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [self._make_scalar_field("id", TypeCode.INT64)] + metadata = self._make_result_set_metadata(fields) + raw_value = self._make_value(100) + values = [raw_value] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator, lazy_decode=True) + found = list(streamed) + self.assertEqual(len(found), 1) + self.assertEqual(found[0], [raw_value]) + self.assertEqual(streamed.decode_row(found[0]), [100]) + + def test_decode_rows_direct_lazy_decode_multi_column(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + value_id = self._make_value(1) + value_name = self._make_value("test") + values = [value_id, value_name] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator, lazy_decode=True) + found = list(streamed) + self.assertEqual(len(found), 1) + self.assertEqual(found[0], [value_id, value_name]) + self.assertEqual(streamed.decode_row(found[0]), [1, "test"]) + + def test_decode_rows_direct_trailing_partial_row(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + values1 = [ + self._make_value(1), + self._make_value("a"), + self._make_value(2), + ] + values2 = [ + self._make_value("b"), + self._make_value(3), + self._make_value("c"), + ] + + result_set1 = self._make_partial_result_set(values1, metadata=metadata) + result_set2 = self._make_partial_result_set(values2, last=True) + iterator = _MockCancellableIterator(result_set1, result_set2) + streamed = self._make_one(iterator) + found = list(streamed) + self.assertEqual( + found, + [ + [1, "a"], + [2, "b"], + [3, "c"], + ], + ) + + def test_decode_rows_direct_trailing_partial_row_lazy_decode(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + value1 = self._make_value(1) + value_a = self._make_value("a") + value2 = self._make_value(2) + value_b = self._make_value("b") + + result_set1 = self._make_partial_result_set( + [value1, value_a, value2], metadata=metadata + ) + result_set2 = self._make_partial_result_set([value_b], last=True) + iterator = _MockCancellableIterator(result_set1, result_set2) + streamed = self._make_one(iterator, lazy_decode=True) + found = list(streamed) + self.assertEqual(len(found), 2) + self.assertEqual(found[0], [value1, value_a]) + self.assertEqual(found[1], [value2, value_b]) + self.assertEqual(streamed.decode_row(found[0]), [1, "a"]) + self.assertEqual(streamed.decode_row(found[1]), [2, "b"]) + + def test_decode_rows_direct_trailing_partial_row_with_null(self): + from google.protobuf.struct_pb2 import NULL_VALUE, Value + + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + values1 = [ + self._make_value(1), + self._make_value("a"), + Value(null_value=NULL_VALUE), + ] + values2 = [ + self._make_value("b"), + ] + + result_set1 = self._make_partial_result_set(values1, metadata=metadata) + result_set2 = self._make_partial_result_set(values2, last=True) + iterator = _MockCancellableIterator(result_set1, result_set2) + streamed = self._make_one(iterator) + found = list(streamed) + self.assertEqual( + found, + [ + [1, "a"], + [None, "b"], + ], + ) + + def test_decode_rows_direct_prefix_partial_row_with_null(self): + from google.protobuf.struct_pb2 import NULL_VALUE, Value + + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + ] + metadata = self._make_result_set_metadata(fields) + values1 = [ + self._make_value(1), + self._make_value("a"), + self._make_value(2), + ] + values2 = [ + Value(null_value=NULL_VALUE), + self._make_value(3), + self._make_value("c"), + ] + + result_set1 = self._make_partial_result_set(values1, metadata=metadata) + result_set2 = self._make_partial_result_set(values2, last=True) + iterator = _MockCancellableIterator(result_set1, result_set2) + streamed = self._make_one(iterator) + found = list(streamed) + self.assertEqual( + found, + [ + [1, "a"], + [2, None], + [3, "c"], + ], + ) + + def test_decode_rows_direct_width_one_with_null(self): + from google.protobuf.struct_pb2 import NULL_VALUE, Value + + from google.cloud.spanner_v1 import TypeCode + + fields = [self._make_scalar_field("id", TypeCode.INT64)] + metadata = self._make_result_set_metadata(fields) + values = [ + self._make_value(1), + Value(null_value=NULL_VALUE), + self._make_value(2), + ] + + result_set = self._make_partial_result_set(values, metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator) + found = list(streamed) + self.assertEqual(found, [[1], [None], [2]]) + + def test_decode_rows_direct_three_chunk_split_row(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + self._make_scalar_field("active", TypeCode.BOOL), + ] + metadata = self._make_result_set_metadata(fields) + prs1 = self._make_partial_result_set([self._make_value(1)], metadata=metadata) + prs2 = self._make_partial_result_set([self._make_value("alpha")]) + prs3 = self._make_partial_result_set([self._make_value(True)], last=True) + + iterator = _MockCancellableIterator(prs1, prs2, prs3) + streamed = self._make_one(iterator) + found = list(streamed) + self.assertEqual(found, [[1, "alpha", True]]) + + def test_decode_rows_direct_three_chunk_split_row_lazy_decode(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [ + self._make_scalar_field("id", TypeCode.INT64), + self._make_scalar_field("name", TypeCode.STRING), + self._make_scalar_field("active", TypeCode.BOOL), + ] + metadata = self._make_result_set_metadata(fields) + val1 = self._make_value(1) + val2 = self._make_value("alpha") + val3 = self._make_value(True) + prs1 = self._make_partial_result_set([val1], metadata=metadata) + prs2 = self._make_partial_result_set([val2]) + prs3 = self._make_partial_result_set([val3], last=True) + + iterator = _MockCancellableIterator(prs1, prs2, prs3) + streamed = self._make_one(iterator, lazy_decode=True) + found = list(streamed) + self.assertEqual(found, [[val1, val2, val3]]) + self.assertEqual(streamed.decode_row(found[0]), [1, "alpha", True]) + + def test_decode_rows_direct_empty_values(self): + from google.cloud.spanner_v1 import TypeCode + + fields = [self._make_scalar_field("id", TypeCode.INT64)] + metadata = self._make_result_set_metadata(fields) + + result_set1 = self._make_partial_result_set([], metadata=metadata) + result_set2 = self._make_partial_result_set([self._make_value(1)], last=True) + iterator = _MockCancellableIterator(result_set1, result_set2) + streamed = self._make_one(iterator) + found = list(streamed) + self.assertEqual(found, [[1]]) + + def test_merge_values_zero_fields(self): + from google.cloud.spanner_v1 import ResultSetMetadata, StructType + + metadata = ResultSetMetadata(row_type=StructType(fields=[])) + result_set = self._make_partial_result_set([], metadata=metadata, last=True) + iterator = _MockCancellableIterator(result_set) + streamed = self._make_one(iterator) + list(streamed) + streamed._merge_values([self._make_value(1)]) + self.assertEqual(streamed._rows, []) + class _MockCancellableIterator(object): cancel_calls = 0