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.
Measured bytes using the real `_finalize_successful_query` /
`_serialize_payload` functions from master
d176d68451bee9200102f38e2725d5f1c8821276 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 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'`
gives 2,473 passed; staged pre-commit, including mypy, passes. The legacy
payload-limit tests also pass (3 selected tests, including the JSON
normalization control).
Siblings checked:
- superset/sql/execution/celery_task.py:250 — fixed in
e63443ac54bc300f61027b7cb50784b528473c8d
- superset/sql/execution/celery_task.py:178 — 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:517 — fixed in
e63443ac54bc300f61027b7cb50784b528473c8d
- superset/sql_lab.py:717 — 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:852 — not affected: the return-results branch converts
Arrow back to records before its pre-compression `sys.getsizeof` check.
- superset/sql/execution/celery_task.py:285 — not affected:
serialized/compressed byte sizes here are logging only and continue to report
actual storage sizes.
- superset/sql_lab.py:737 — 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:777 — 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]
