rebenitez1802 commented on code in PR #45348:
URL: https://github.com/apache/superset/pull/45348#discussion_r4296623961
##########
superset/sql/execution/celery_task.py:
##########
@@ -240,6 +254,23 @@ def _finalize_successful_query(
payload["query"]["state"] = QueryStatus.SUCCESS.value
+def _get_payload_size(payload: dict[str, Any], serialized_payload: bytes) ->
int:
+ """Count raw Arrow and metadata, excluding duplicate nested sidecar
values."""
+ size = len(serialized_payload)
+ for statement in payload.get("statements", []):
+ if isinstance(data := statement.get("data"), str):
+ # Retain all envelope metadata without allocating another decoded
+ # copy of the Arrow data just to determine its size.
+ arrow_size = len(data) // 4 * 3 - data.endswith("=") -
data.endswith("==")
+ size -= len(data) - arrow_size
+ # Arrow retains the full nested column as text for older readers and
+ # unknown tags. Its typed copy must not charge the same values twice.
+ size -= sum(
+ len(value) for value in statement.get("nested_columns_msgpack",
{}).values()
+ )
+ return size
Review Comment:
**🟡 Medium — Nested binary values bypass `SQLLAB_PAYLOAD_MAX_MB`**
This subtracts every sidecar in full, assuming Arrow already holds the same
values as text. That's true for JSON-like values but not for bytes. When the
result set stringifies nested values, `base_json_conv`
(`superset/utils/json.py`) turns any bytes that aren't valid UTF-8 or UTF-16
into the literal `"[bytes]"`. The binary data then exists only in the sidecar
and is never counted.
Repro: `SQLLAB_PAYLOAD_MAX_MB = 1` and 6,000 rows of `[os.urandom(4097)]`.
Both transports return SUCCESS with about 23.5 MiB stored, against 0.09 MiB
counted. Random bytes don't compress. Drivers return nested bytes for types
like DuckDB `BLOB[]`, Trino `array(varbinary)` and BigQuery `ARRAY<BYTES>`.
Master stored only the placeholders, and d205c6f counted the sidecar.
The suggestion below changes two things:
- It only discounts the part of each sidecar that its Arrow column also
holds: `min(raw size, Arrow column nbytes)`.
- It counts sidecars at their raw size in both transports.
I measured it with the suggestion applied. The counted size is the same in
both transports:
| Result | Counted |
| --- | --- |
| 50k JSONB-like rows | 5.62 MiB (same as master) |
| 2k rows of 4 KB nested binary | 7.82 MiB (about 0.03 MiB at this head) |
| 20k rows of nested dates | 0.88 MiB |
The Base64 Arrow data is decoded only for JSON-mode statements that have
sidecars. `pyarrow` is imported locally only so the change fits in one
suggestion block; it can move to the module imports.
Follow-ups this suggestion can't cover:
- `test_store_results_in_backend_raw_arrow_size_boundary` builds a sidecar
for a `nested` column that isn't in its Arrow table, so its expected size needs
to count that sidecar.
- Please add a test that an over-limit nested binary column returns
`RESULT_TOO_LARGE`.
- Reword `development.md:275-276`.
```suggestion
def _get_payload_size(payload: dict[str, Any], serialized_payload: bytes) ->
int:
"""Count raw Arrow and metadata, excluding duplicate nested sidecar
values."""
import pyarrow as pa
def raw_size(value: bytes | str) -> int:
# Decoded Base64 length, without allocating a decoded copy.
if isinstance(value, bytes):
return len(value)
return len(value) // 4 * 3 - value.endswith("=") -
value.endswith("==")
size = len(serialized_payload)
for statement in payload.get("statements", []):
data = statement.get("data")
if isinstance(data, str):
size -= len(data) - raw_size(data)
if sidecars := statement.get("nested_columns_msgpack"):
# Arrow keeps nested columns as text for older readers and
unknown
# tags, so only the part of each typed copy that Arrow also
holds is
# a duplicate. Values Arrow stores as placeholders (such as
binary
# data) still count, at their raw size.
arrow = base64.b64decode(data) if isinstance(data, str) else data
table = pa.ipc.open_stream(pa.BufferReader(arrow)).read_all()
for name, value in sidecars.items():
duplicate = (
min(raw_size(value), table.column(name).nbytes)
if name in table.column_names
else 0
)
size -= len(value) - raw_size(value) + duplicate
return size
```
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
