damccorm commented on code in PR #41059: URL: https://github.com/apache/beam/pull/41059#discussion_r4525280579
########## sdks/python/apache_beam/io/gcp/bigquery.py: ########## @@ -1522,6 +1523,39 @@ def BigQuerySink(*args, validate=False, **kwargs): _KNOWN_TABLES = set() +# Primitive scalar types that do not contain nested object references and can be +# sized directly via sys.getsizeof without gc.get_referents traversal. +_PRIMITIVE_SIZEOF_TYPES = (str, int, float, bool, bytes, type(None)) + + +def _fast_row_and_insert_id_size(row_and_insert_id): + """Fast-path byte size estimator for a ``(row_dict, insert_id)`` tuple. + + ``objsize.get_deep_size`` traverses the Python object graph via + ``gc.get_referents`` and tracks visited object IDs on every call. When the row + is a flat ``dict`` of primitive keys and values, we can sum ``sys.getsizeof`` + directly and fall back to ``get_deep_size`` only if nested containers or + non-primitive values are encountered. + """ + if type(row_and_insert_id) is tuple and len(row_and_insert_id) == 2: + row, insert_id = row_and_insert_id + if type(row) is dict and type(insert_id) in _PRIMITIVE_SIZEOF_TYPES: + getsizeof = sys.getsizeof + total = getsizeof(row_and_insert_id) + getsizeof(row) + # objsize.get_deep_size ignores None singletons (which are not tracked by + # the cyclic garbage collector). + if insert_id is not None: + total += getsizeof(insert_id) + for k, v in row.items(): + if not (type(k) in _PRIMITIVE_SIZEOF_TYPES and + type(v) in _PRIMITIVE_SIZEOF_TYPES): + return get_deep_size(row_and_insert_id) + # Avoid double-counting if a column value references the same object as + # insert_id, and skip None values to match get_deep_size behavior. + total += getsizeof(k) + ( + getsizeof(v) if v is not None and v is not insert_id else 0) + return total + return get_deep_size(row_and_insert_id) Review Comment: I'll note that if we reach this (or the other `get_deep_size` call in this function), and it returns something larger than self._max_insert_payload_size, then we will end up calling it twice. That is ok, though, since that is the rare error case and shouldn't be the normal hot path. -- 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]
