Skip to content

perf(spanner): optimize query result decoding and type resolution #18330

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
olavloite wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(spanner): optimize query result decoding and type resolution #18330

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
116 changes: 85 additions & 31 deletions packages/google-cloud-spanner/google/cloud/spanner_v1/_async/streamed.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
)
Comment thread
olavloite marked this conversation as resolved.

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
]
)
Comment thread
olavloite marked this conversation as resolved.

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):
Expand Down
64 changes: 34 additions & 30 deletions packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
olavloite marked this conversation as resolved.
elif type_code == TypeCode.INTERVAL:
return _parse_interval
else:
raise ValueError("Unknown type: %s" % (field_type,))

Expand Down Expand Up @@ -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.

Expand Down
117 changes: 87 additions & 30 deletions packages/google-cloud-spanner/google/cloud/spanner_v1/streamed.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
)
Comment thread
olavloite marked this conversation as resolved.

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
]
)
Comment thread
olavloite marked this conversation as resolved.

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.
Expand Down
Loading
Loading