damccorm commented on code in PR #41059: URL: https://github.com/apache/beam/pull/41059#discussion_r4531689824
########## 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) + ( Review Comment: Yeah, it could end up slightly shrinking batch sizes (computed here - https://github.com/damccorm/beam/blob/73cc912b136b3d9a9db3d0c06436bfceda181c6e/sdks/python/apache_beam/io/gcp/bigquery.py#L1821), but I think this is ok. The goal here is to be a fast/quick estimate of the true size, so if we're a little more accurate than the previous version it should be ok (and maybe even guard against running over the limit a little bit better) -- 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]
