aminghadersohi commented on code in PR #45348:
URL: https://github.com/apache/superset/pull/45348#discussion_r4288336737
##########
superset/sql/execution/executor.py:
##########
@@ -404,6 +406,147 @@ def execute_sql_with_cursor(
return results
+def _decode_async_results_payload(blob: bytes) -> dict[str, Any]:
+ """
+ Decode a results backend blob written by the async SQL execution task.
+
+ The task stores a zlib-compressed payload serialized either as MessagePack
+ or as UTF-8 JSON, depending on ``RESULTS_BACKEND_USE_MSGPACK``. The format
+ is detected from the payload itself (a JSON object starts with ``{``, which
+ is never the first byte of a MessagePack map), so results stay readable if
+ the setting differs between the writing worker and the reading process.
+
+ :param blob: Compressed payload read from the results backend
+ :returns: The deserialized payload
+ """
+ import msgpack
+
+ from superset.utils import json
+
+ decompressed = cast(bytes, utils.zlib_decompress(blob, decode=False))
+ if decompressed.lstrip()[:1] == b"{":
+ return json.loads(decompressed)
+ return msgpack.loads(decompressed, raw=False)
+
+
+def _decode_async_statement_data(stmt_data: dict[str, Any]) -> pd.DataFrame |
None:
+ """
+ Convert one stored statement's data into a DataFrame.
+
+ SELECT results are stored as Arrow IPC bytes (MessagePack payloads) or as
+ a list of records (JSON payloads); DML statements store ``None``. A
+ SELECT returning no rows yields an empty DataFrame that keeps its columns,
+ matching the synchronous execution path.
+
+ :param stmt_data: Statement entry from the stored payload
+ :returns: DataFrame of the statement rows, or ``None`` for DML
+ """
+ import pyarrow as pa
+
+ data = stmt_data.get("data")
+ if data is None:
+ return None
+
+ if isinstance(data, (bytes, bytearray, memoryview)):
+ from superset.result_set import SupersetResultSet
+
+ table = pa.ipc.open_stream(pa.BufferReader(data)).read_all()
+ return SupersetResultSet.convert_table_to_df(table)
Review Comment:
My earlier reply overstated the nested-type guarantee. Fixed the
oversized-integer crash and supported nested scalar type loss in
23be919b260402b6a56d554e2aea0cb682a882fc. A typed Base64 MessagePack sidecar
keeps arbitrary-sized integers, datetime (naive or datetime.timezone
fixed-offset), UUID, Decimal, and bytes out of the outer transport's lossy
serializer. Both backend modes retain these values and Python types.
Remaining: unsupported nested types use Arrow's string representation for
the entire affected column, rather than synchronous Python objects. This
includes sets, tuples, scalar subclasses, and datetimes using timezone
implementations such as zoneinfo.ZoneInfo. Existing untyped sidecars remain
readable but cannot recover types already lost. The dead arrow_schema removal
in the legacy test helper is also gone.
Regression tests added by this PR in
`tests/unit_tests/sql/execution/test_async_results.py`:
- `test_async_nested_oversized_integer_matches_sync`: before, both
MessagePack cases raised TypeError; after, all 4 cases pass, including negative
oversized integers and reader-setting changes.
- `test_async_nested_scalar_types_match_sync`: before, datetime/UUID/Decimal
failed in both transports and bytes failed in JSON; after, all 14 cases pass
with explicit value/type comparisons against synchronous execution.
- `test_async_unsupported_nested_types_use_arrow_strings`: all 8 cases
failed before and pass after, testing the documented whole-column string
fallback for sets, tuples, scalar subclasses, and named-timezone datetimes.
- `test_legacy_nested_oversized_integer_matches_sync`: failed before with
OverflowError; passes after.
I ran `.venv/bin/pytest tests/unit_tests/sql/execution/test_async_results.py
-q -k 'nested_oversized_integer or nested_scalar_types or
unsupported_nested_types or legacy_nested_sidecar' -m 'not manual and not
llm_compliance'` with implementation files restored to
`98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d`: **23 failed, 5 passed**; with this
commit: **28 passed**. The five previously passing cases are compatibility
controls, not claimed reproductions. `.venv/bin/pytest tests/unit_tests/sql/ -q
-m 'not manual and not llm_compliance'`: **2455 passed**. `pre-commit run` and
`pre-commit run --files <all branch-changed paths>` passed.
Siblings checked:
- superset/sql/execution/celery_task.py:193 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- superset/sql/execution/serialization.py:72 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- superset/sql/execution/serialization.py:95 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- superset/sql/execution/executor.py:434 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- superset/sql/execution/executor.py:473 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- superset/sql/execution/celery_task.py:122 — not affected: the outer
serializer receives Base64 sidecars, not raw nested scalars.
- superset/sql/execution/celery_task.py:308 — not affected: statement Arrow
IPC serialization is unchanged; original nested objects travel in the typed
sidecar.
- superset/sql/execution/executor.py:411 — not affected: outer transport
detection only decodes the payload map; the statement reader decodes the
sidecar.
- superset/sql/execution/executor.py:1185 — not affected: cache records with
object dtype use pandas Series directly, not the async results serializer.
- superset/sql/execution/executor.py:1238 — not affected: cache writes
retain Python records and dtype metadata, with no results-backend
JSON/MessagePack conversion.
- superset/result_set.py:368 — not affected: synchronous execution restores
original nested values directly, providing the tests' expected results.
- superset/sql_lab.py:392 — not affected: the separate SQL Lab writer uses
Arrow bytes for MessagePack and JSON records for JSON; it does not write this
nested sidecar.
- superset/views/utils.py:650 — not affected: the separate SQL Lab reader
decodes its top-level data contract, not SQLExecutor statement sidecars.
- superset/commands/sql_lab/results.py:170 — not affected: calls the
separate SQL Lab reader.
- superset/commands/sql_lab/export.py:121 — not affected: calls the separate
SQL Lab reader for CSV export.
- superset/sqllab/utils.py:76 — not affected: writes the existing Arrow
table, whose nested values are already strings.
- superset/connectors/sqla/utils.py:224 — not affected: expands column
metadata, not stored result values.
- superset/db_engine_specs/base.py:2788 — not affected: expands column
metadata, not stored result values.
- superset/db_engine_specs/bigquery.py:1125 — not affected: expands STRUCT
column metadata, not stored result values.
- superset/mcp_service/server.py:654 — not affected: serializes MCP
tool-search results, not SQLExecutor result payloads.
- superset/mcp_service/caching.py:30 — not affected: versions MCP tool cache
keys, not SQLExecutor result serialization.
##########
docs/admin_docs/configuration/async-queries-celery.mdx:
##########
@@ -76,6 +76,14 @@ For performance gains,
[MessagePack](https://github.com/msgpack/msgpack-python)
disabled by setting `RESULTS_BACKEND_USE_MSGPACK = False` in your
`superset_config.py`, should any
issues arise. Please clear your existing results cache store when upgrading an
existing environment.
+The Python `Database.execute_async()` result handle preserves date, time,
timestamp, interval,
+large integer, decimal, and binary column types with either serialization
format. JSON
+results include Arrow schema metadata; binary columns use Base64 encoding in
storage.
+Older JSON results without schema metadata still load but retain their
serialized strings;
+rerun those queries to obtain typed columns. Data cache hits also retain
timestamp units
+and timezones.
+This does not change the JSON representation returned by the SQL Lab results
API.
Review Comment:
My earlier reply and docs wording overstated what was restored. Corrected in
23be919b260402b6a56d554e2aea0cb682a882fc:
`docs/developer_docs/extensions/mcp.md` explicitly lists the supported nested
types and the fallback. Nested JSON scalars (including arbitrary-sized
integers), datetime (naive or datetime.timezone fixed-offset, including
microseconds), UUID, Decimal, and bytes retain their value and type in both
backend modes. The integer crash and nested scalar loss have regression
coverage.
Remaining: other nested types, including sets, tuples, scalar subclasses,
and named-timezone datetimes, use the Arrow string representation for the whole
affected column rather than matching synchronous Python objects. Old JSON
records and untyped sidecars cannot recover types already converted to
strings/floats; rerunning the query is required. The docs state these limits
and remain in the developer extension documentation.
Regression tests added by this PR in
`tests/unit_tests/sql/execution/test_async_results.py`:
- `test_async_nested_oversized_integer_matches_sync`: before, both
MessagePack cases raised TypeError; after, all 4 cases pass, including negative
oversized integers and reader-setting changes.
- `test_async_nested_scalar_types_match_sync`: before, datetime/UUID/Decimal
failed in both transports and bytes failed in JSON; after, all 14 cases pass
with explicit value/type comparisons against synchronous execution.
- `test_async_unsupported_nested_types_use_arrow_strings`: all 8 cases
failed before and pass after, testing the documented whole-column string
fallback for sets, tuples, scalar subclasses, and named-timezone datetimes.
- `test_legacy_nested_oversized_integer_matches_sync`: failed before with
OverflowError; passes after.
I ran `.venv/bin/pytest tests/unit_tests/sql/execution/test_async_results.py
-q -k 'nested_oversized_integer or nested_scalar_types or
unsupported_nested_types or legacy_nested_sidecar' -m 'not manual and not
llm_compliance'` with implementation files restored to
`98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d`: **23 failed, 5 passed**; with this
commit: **28 passed**. The five previously passing cases are compatibility
controls, not claimed reproductions. `.venv/bin/pytest tests/unit_tests/sql/ -q
-m 'not manual and not llm_compliance'`: **2455 passed**. `pre-commit run` and
`pre-commit run --files <all branch-changed paths>` passed.
Siblings checked:
- docs/developer_docs/extensions/mcp.md:31 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- docs/admin_docs/configuration/async-queries-celery.mdx — not affected: the
Python result-handle paragraph is absent from the admin deployment guide; `rg
-n 'Python SQL query results|Nested JSON values|Arrow IPC'
docs/admin_docs/configuration/async-queries-celery.mdx` returned no matches.
- superset/sql/execution/celery_task.py:193 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- superset/sql/execution/serialization.py:72 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- superset/sql/execution/serialization.py:95 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- superset/sql/execution/executor.py:434 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- superset/sql/execution/executor.py:473 — fixed in
23be919b260402b6a56d554e2aea0cb682a882fc
- superset/sql/execution/celery_task.py:122 — not affected: the outer
serializer receives Base64 sidecars, not raw nested scalars.
- superset/sql/execution/celery_task.py:308 — not affected: statement Arrow
IPC serialization is unchanged; original nested objects travel in the typed
sidecar.
- superset/sql/execution/executor.py:411 — not affected: outer transport
detection only decodes the payload map; the statement reader decodes the
sidecar.
- superset/sql/execution/executor.py:1185 — not affected: cache records with
object dtype use pandas Series directly, not the async results serializer.
- superset/sql/execution/executor.py:1238 — not affected: cache writes
retain Python records and dtype metadata, with no results-backend
JSON/MessagePack conversion.
- superset/result_set.py:368 — not affected: synchronous execution restores
original nested values directly, providing the tests' expected results.
- superset/sql_lab.py:392 — not affected: the separate SQL Lab writer uses
Arrow bytes for MessagePack and JSON records for JSON; it does not write this
nested sidecar.
- superset/views/utils.py:650 — not affected: the separate SQL Lab reader
decodes its top-level data contract, not SQLExecutor statement sidecars.
- superset/commands/sql_lab/results.py:170 — not affected: calls the
separate SQL Lab reader.
- superset/commands/sql_lab/export.py:121 — not affected: calls the separate
SQL Lab reader for CSV export.
- superset/sqllab/utils.py:76 — not affected: writes the existing Arrow
table, whose nested values are already strings.
- superset/connectors/sqla/utils.py:224 — not affected: expands column
metadata, not stored result values.
- superset/db_engine_specs/base.py:2788 — not affected: expands column
metadata, not stored result values.
- superset/db_engine_specs/bigquery.py:1125 — not affected: expands STRUCT
column metadata, not stored result values.
- superset/mcp_service/server.py:654 — not affected: serializes MCP
tool-search results, not SQLExecutor result payloads.
- superset/mcp_service/caching.py:30 — not affected: versions MCP tool cache
keys, not SQLExecutor result serialization.
--
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]
