perf(spanner): optimize query result decoding (#17375)

Work in progress.
Optimizes the decoding and reading of (large) result sets for Spanner.
<img width="1643" height="943" alt="image"
src="https://github.com/user-attachments/assets/81997b8d-f77f-4523-acb2-e44fdccb939b"
/>

Author: olavloite

packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py

@@ -129,16 +129,39 @@ def _merge_values(self, values):
 decoders=self._decoders
 width=len(self.fields)
 index=len(self._current_row)
-forvalueinvalues:
-ifself._lazy_decode:
-self._current_row.append(value)
-else:
-self._current_row.append(_parse_nullable(value, decoders[index]))
-index+=1
-ifindex==width:
-self._rows.append(self._current_row)
-self._current_row= []
-index=0
+current_row=self._current_row
+rows=self._rows
+
+current_row_append=current_row.append
+rows_append=rows.append
+
+ifself._lazy_decode:
+forvalueinvalues:
+current_row_append(value)
+index+=1
+ifindex==width:
+rows_append(current_row)
+current_row= []
+current_row_append=current_row.append
+index=0
+else:
+forvalueinvalues:
+# 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.
+ifvalue.HasField("null_value"):
+current_row_append(None)
+else:
+current_row_append(decoders[index](value))
+index+=1
+ifindex==width:
+rows_append(current_row)
+current_row= []
+current_row_append=current_row.append
+index=0
+
+self._current_row=current_row
 
 @CrossSync.convert
 asyncdef_consume_next(self):

packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py

@@ -19,14 +19,14 @@
 importdecimal
 importlogging
 importmath
+importoperator
 importthreading
 importtime
 importuuid
 fromcontextlibimportcontextmanager
 
 fromgoogle.api_coreimportdatetime_helpers
 fromgoogle.api_core.exceptionsimportAborted
-fromgoogle.cloud._helpersimport_date_from_iso8601_date
 fromgoogle.protobuf.internal.enum_type_wrapperimportEnumTypeWrapper
 fromgoogle.protobuf.messageimportDecodeError, Message
 fromgoogle.protobuf.struct_pb2importListValue, Value
@@ -465,6 +465,12 @@ def _parse_value_pb(value_pb, field_type, field_name, column_info=None):
 return_parse_nullable(value_pb, decoder)
 
 
+_date_fromisoformat=datetime.date.fromisoformat
+_Decimal=decimal.Decimal
+_json_from_str=JsonObject.from_str
+_uuid_UUID=uuid.UUID
+
+
 def_get_type_decoder(field_type, field_name, column_info=None):
 """Returns a function that converts a Value protobuf to cell data.
 
@@ -489,28 +495,30 @@ 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.
 iftype_code==TypeCode.STRING:
-return_parse_string
+returnoperator.attrgetter("string_value")
 eliftype_code==TypeCode.BYTES:
-return_parse_bytes
+returnlambdavalue_pb: value_pb.string_value.encode("utf8")
 eliftype_code==TypeCode.BOOL:
-return_parse_bool
+returnoperator.attrgetter("bool_value")
 eliftype_code==TypeCode.INT64:
-return_parse_int64
+returnlambdavalue_pb: int(value_pb.string_value)
 eliftype_code==TypeCode.FLOAT64:
 return_parse_float
 eliftype_code==TypeCode.FLOAT32:
 return_parse_float
 eliftype_code==TypeCode.DATE:
-return_parse_date
+returnlambdavalue_pb: _date_fromisoformat(value_pb.string_value)
 eliftype_code==TypeCode.TIMESTAMP:
 return_parse_timestamp
 eliftype_code==TypeCode.NUMERIC:
-return_parse_numeric
+returnlambdavalue_pb: _Decimal(value_pb.string_value)
 eliftype_code==TypeCode.JSON:
-return_parse_json
+returnlambdavalue_pb: _json_from_str(value_pb.string_value)
 eliftype_code==TypeCode.UUID:
-return_parse_uuid
+returnlambdavalue_pb: _uuid_UUID(value_pb.string_value)
 eliftype_code==TypeCode.PROTO:
 returnlambdavalue_pb: _parse_proto(value_pb, column_info, field_name)
 eliftype_code==TypeCode.ENUM:
@@ -553,48 +561,81 @@ def _parse_list_value_pbs(rows, row_type):
 returnresult
 
 
-def_parse_string(value_pb) ->str:
-returnvalue_pb.string_value
-
-
-def_parse_bytes(value_pb):
-returnvalue_pb.string_value.encode("utf8")
-
-
-def_parse_bool(value_pb) ->bool:
-returnvalue_pb.bool_value
-
-
-def_parse_int64(value_pb) ->int:
-returnint(value_pb.string_value)
-
-
 def_parse_float(value_pb) ->float:
-ifvalue_pb.HasField("string_value"):
-returnfloat(value_pb.string_value)
-else:
-returnvalue_pb.number_value
-
-
-def_parse_date(value_pb):
-return_date_from_iso8601_date(value_pb.string_value)
+# Note: Storing val = value_pb.string_value and doing a truthiness check is faster
+# than calling value_pb.HasField("string_value") because it avoids the C-extension
+# method lookup/call overhead and accesses the attribute only once.
+val=value_pb.string_value
+returnfloat(val) ifvalelsevalue_pb.number_value
+
+
+_POWERS_OF_10= (
+1,
+10,
+100,
+1000,
+10000,
+100000,
+1000000,
+10000000,
+100000000,
+1000000000,
+)
 
 
 def_parse_timestamp(value_pb):
-DatetimeWithNanoseconds=datetime_helpers.DatetimeWithNanoseconds
-returnDatetimeWithNanoseconds.from_rfc3339(value_pb.string_value)
-
-
-def_parse_numeric(value_pb):
-returndecimal.Decimal(value_pb.string_value)
-
-
-def_parse_json(value_pb):
-returnJsonObject.from_str(value_pb.string_value)
-
-
-def_parse_uuid(value_pb):
-returnuuid.UUID(value_pb.string_value)
+val=value_pb.string_value
+try:
+iflen(val) <20orval[10] !="T":
+raiseValueError()
+no_fraction=val[:19]
+bare=datetime.datetime.fromisoformat(no_fraction)
+ifval[19] ==".":
+ifval.endswith("Z"):
+offset="Z"
+fraction=val[20:-1]
+elifval[-6] in ("+", "-"):
+offset=val[-6:]
+fraction=val[20:-6]
+else:
+raiseValueError()
+ifnotfractionorlen(fraction) >9ornotfraction.isdigit():
+raiseValueError()
+scale=9-len(fraction)
+nanos=int(fraction) *_POWERS_OF_10[scale]
+else:
+nanos=0
+ifval.endswith("Z"):
+offset="Z"
+elifval[-6] in ("+", "-"):
+offset=val[-6:]
+else:
+raiseValueError()
+
+ifoffset!="Z":
+sign=offset[0]
+hours=int(offset[1:3])
+minutes=int(offset[4:6])
+ifoffset[3] !=":":
+raiseValueError()
+delta=datetime.timedelta(hours=hours, minutes=minutes)
+ifsign=="-":
+delta=-delta
+tzinfo=datetime.timezone(delta)
+bare=bare.replace(tzinfo=tzinfo).astimezone(datetime.timezone.utc)
+
+returndatetime_helpers.DatetimeWithNanoseconds(
+bare.year,
+bare.month,
+bare.day,
+bare.hour,
+bare.minute,
+bare.second,
+nanosecond=nanos,
+tzinfo=datetime.timezone.utc,
+ )
+except (IndexError, ValueError) ase:
+raiseValueError("Timestamp: {} does not match pattern".format(val)) frome
 
 
 def_parse_proto(value_pb, column_info, field_name):

packages/google-cloud-spanner/google/cloud/spanner_v1/data_types.py

@@ -99,6 +99,11 @@ def serialize(self):
 returnjson.dumps(self, sort_keys=True, separators=(",", ":"))
 
 
+_INTERVAL_PATTERN=re.compile(
+r"^P(-?\d+Y)?(-?\d+M)?(-?\d+D)?(T(-?\d+H)?(-?\d+M)?(-?((\d+([.,]\d{1,9})?)|([.,]\d{1,9}))S)?)?$"
+)
+
+
 @dataclass
 classInterval:
 """Represents a Spanner INTERVAL type.
@@ -187,8 +192,7 @@ def __str__(self) -> str:
 @classmethod
 deffrom_str(cls, s: str) ->"Interval":
 """Parse an ISO8601 duration format string into an Interval."""
-pattern=r"^P(-?\d+Y)?(-?\d+M)?(-?\d+D)?(T(-?\d+H)?(-?\d+M)?(-?((\d+([.,]\d{1,9})?)|([.,]\d{1,9}))S)?)?$"
-match=re.match(pattern, s)
+match=_INTERVAL_PATTERN.match(s)
 ifnotmatchorlen(s) ==1:
 raiseValueError(f"Invalid interval format: {s}")
 

packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py

@@ -35,8 +35,7 @@ class StreamedResultSet(object):
  instances.
 
  :type source: :class:`~google.cloud.spanner_v1.snapshot.Snapshot`
- :param source: Deprecated. Snapshot from which the result set was fetched.
- """
+ :param source: Deprecated. Snapshot from which the result set was fetched."""
 
 def__init__(
 self,
@@ -117,16 +116,36 @@ def _merge_values(self, values):
 decoders=self._decoders
 width=len(self.fields)
 index=len(self._current_row)
-forvalueinvalues:
-ifself._lazy_decode:
-self._current_row.append(value)
-else:
-self._current_row.append(_parse_nullable(value, decoders[index]))
-index+=1
-ifindex==width:
-self._rows.append(self._current_row)
-self._current_row= []
-index=0
+current_row=self._current_row
+rows=self._rows
+current_row_append=current_row.append
+rows_append=rows.append
+ifself._lazy_decode:
+forvalueinvalues:
+current_row_append(value)
+index+=1
+ifindex==width:
+rows_append(current_row)
+current_row= []
+current_row_append=current_row.append
+index=0
+else:
+forvalueinvalues:
+# 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.
+ifvalue.HasField("null_value"):
+current_row_append(None)
+else:
+current_row_append(decoders[index](value))
+index+=1
+ifindex==width:
+rows_append(current_row)
+current_row= []
+current_row_append=current_row.append
+index=0
+self._current_row=current_row
 
 def_consume_next(self):
 """Consume the next partial result set from the stream.