aminghadersohi commented on code in PR #45348:
URL: https://github.com/apache/superset/pull/45348#discussion_r4289169223
##########
superset/sql/execution/celery_task.py:
##########
@@ -298,33 +307,31 @@ def _store_results_in_backend(
def _serialize_result_set(
result_set: SupersetResultSet,
-) -> tuple[bytes | list[Any], list[Any]]:
+) -> tuple[bytes | str, list[Any]]:
"""
Serialize result set based on RESULTS_BACKEND_USE_MSGPACK config.
- When msgpack is enabled, uses Apache Arrow IPC format for efficiency.
- Otherwise, falls back to JSON-serializable records.
+ Both transports use Apache Arrow IPC to retain the result types. JSON
+ wraps the IPC bytes in Base64; MessagePack carries them directly.
:param result_set: Query result set to serialize
:returns: Tuple of (serialized_data, columns)
"""
from superset import results_backend_use_msgpack
- from superset.dataframe import df_to_records
- if results_backend_use_msgpack:
- if has_app_context():
- stats_logger = app.config["STATS_LOGGER"]
- with stats_timing(
- "sqllab.query.results_backend_pa_serialization", stats_logger
- ):
- data: bytes | list[Any] = write_ipc_buffer(
- result_set.pa_table
- ).to_pybytes()
- else:
- data = write_ipc_buffer(result_set.pa_table).to_pybytes()
+ if has_app_context():
+ stats_logger = app.config["STATS_LOGGER"]
+ with stats_timing(
+ "sqllab.query.results_backend_pa_serialization", stats_logger
+ ):
+ arrow_data = write_ipc_buffer(result_set.pa_table).to_pybytes()
else:
- df = result_set.to_pandas_df()
- data = df_to_records(df, convert_decimals=True) or []
+ arrow_data = write_ipc_buffer(result_set.pa_table).to_pybytes()
+ data = (
+ arrow_data
+ if results_backend_use_msgpack
+ else base64.b64encode(arrow_data).decode("ascii")
Review Comment:
You're right: the Base64 storage format is intentional, but charging its
encoding overhead against the result-size limit was not. I reproduced the 10
MiB-limit regression and corrected the accounting in
e63443ac54bc300f61027b7cb50784b528473c8d, retained after merging master in
d205c6fd3ac016d58b688b3eec250f974491eb0c.
Measured bytes using the real `_finalize_successful_query` /
`_serialize_payload` functions from master
13bd0cd932632522ddc28d193a1057f8f487691c and the reviewed head
f873c74a4a9fc2f3340b81d88bf190bda4e46b87 (PyArrow 25.0.1; identical result sets
and fixed query metadata):
| Result / transport | Master, before compression | Reviewed head, before
compression | Master / reviewed head, after zlib |
| --- | ---: | ---: | ---: |
| 100,000 mixed rows, JSON | 20,473,282 | 15,635,314 | 1,449,670 / 1,613,220
|
| 100,000 mixed rows, MessagePack | 11,726,451 | 11,726,451 | 1,248,333 /
1,248,333 |
| One 8 MiB text value, JSON | 8,389,027 | 11,185,615 | 8,436 / 11,317 |
| One 8 MiB text value, MessagePack | 8,389,217 | 8,389,217 | 8,552 / 8,552 |
The mixed rows contain an integer ID, an 81-character string, float,
boolean, timestamp, and nullable integer. The text case is one string of `8 *
1024 * 1024` ASCII characters. Its raw IPC is 8,388,904 bytes; Base64 is
11,185,208 bytes. The mixed case is smaller than master, but the text case
confirms that compression does not prevent the rejection: the check runs before
compression.
At the reviewed head, `superset/sql/execution/celery_task.py:250-265`
compares `len(serialized_payload)` with `SQLLAB_PAYLOAD_MAX_MB * 1024 * 1024`.
The new `_get_payload_size` calculation substitutes raw IPC length for each
Base64 statement value, retaining the envelope, metadata, and nested sidecars.
It does not change the stored bytes or use compressed size. The text case is
charged 8,389,311 bytes in JSON mode, below 10 MiB. Arrow/schema/envelope sizes
can still differ from the former record format; this removes the Base64 penalty
rather than promising byte-identical serialization.
Regression tests added by this PR:
-
`tests/unit_tests/sql/execution/test_async_results.py::test_async_text_payload_limit_excludes_base64_overhead`:
actual 8 MiB writer/reader round-trip under a 10 MiB limit, one/two SELECT
statements and both transports. Both JSON cases fail at f873c74a and pass at
e63443ac.
-
`tests/unit_tests/sql/execution/test_celery_task.py::test_store_results_in_backend_raw_arrow_size_boundary`:
exact-limit acceptance, one-byte-over rejection without writing, all three
Base64 padding lengths, DML, multiple SELECTs, metadata/sidecar accounting, and
unchanged stored bytes. Three JSON acceptance cases fail before the change and
pass afterward; MessagePack and rejection cases are controls.
The focused command (`pytest
tests/unit_tests/sql/execution/test_async_results.py
tests/unit_tests/sql/execution/test_celery_task.py -k 'text_payload_limit or
raw_arrow_size_boundary'`) gives 5 failed / 11 passed before and 16 passed
after. `pytest tests/unit_tests/sql/ -m 'not manual and not llm_compliance'` at
the merged head gives 2,634 passed / 2 skipped (the optional Impala driver is
not installed); pre-commit on the branch's changed files, including mypy,
passes. The legacy payload-limit tests also pass (4 selected tests, including
the JSON normalization control).
Siblings checked:
- superset/sql/execution/celery_task.py:252 — fixed in
d205c6fd3ac016d58b688b3eec250f974491eb0c
- superset/sql/execution/celery_task.py:294 — fixed in
e63443ac54bc300f61027b7cb50784b528473c8d
- superset/sql/execution/celery_task.py:200 — not affected: every SELECT
generated here reaches the shared storage check; the regression covers one and
multiple statements, including DML at the boundary.
- superset/sql/execution/celery_task.py:665 — fixed in
d205c6fd3ac016d58b688b3eec250f974491eb0c
- superset/sql_lab.py:714 — not affected: legacy backend storage uses JSON
records or raw MessagePack Arrow, not Base64 Arrow; its existing
pre-compression `sys.getsizeof` comparison is unchanged.
- superset/sql_lab.py:803 — not affected: the return-results branch converts
Arrow back to records before its pre-compression `sys.getsizeof` check.
- superset/sql/execution/celery_task.py:320 — not affected:
serialized/compressed byte sizes here are logging only and continue to report
actual storage sizes.
- superset/sql_lab.py:734 — not affected: serialized/compressed byte sizes
here are logging only, not additional limit checks.
- superset/config.py:2007 — not affected: the opt-in limit/default is
unchanged.
- superset/errors.py:90 — not affected: error enum and its code mapping at
line 196 do not measure payloads.
- superset/utils/error_sanitization.py:63 — not affected: error
classification does not measure payloads.
- docs/static/resources/openapi.json:121 — not affected: error-enum
documentation only.
-
docs/developer_docs_versioned_docs/version-6.1.0/api/get-the-result-of-a-sql-query-execution.StatusCodes.json:179
— not affected: error-enum documentation only.
- tests/unit_tests/sql_lab_test.py:183 — not affected: legacy within-limit
and over-limit controls pass; no Base64 producer in that path.
- tests/unit_tests/sql/execution/test_celery_task.py:779 — fixed in
e63443ac54bc300f61027b7cb50784b528473c8d
- tests/unit_tests/sql/execution/test_async_results.py:137 — fixed in
e63443ac54bc300f61027b7cb50784b528473c8d
- docs/developer_docs/extensions/mcp.md:47 — fixed in
e63443ac54bc300f61027b7cb50784b528473c8d
Search: `rg -n
'SQLLAB_PAYLOAD_MAX_MB|RESULT_TOO_LARGE_ERROR|serialized_payload_size|(?:len|(?:sys\.)?getsizeof)\((?:serialized_payload|compressed)\)'
superset tests docs`, plus searches for `_store_results_in_backend`,
`_serialize_result_set`, and `_serialize_and_expand_data` callers.
--
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]
