aminghadersohi commented on code in PR #45348:
URL: https://github.com/apache/superset/pull/45348#discussion_r4288014576
##########
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:
Fixed in 98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d. The task carries the
original nested-column values alongside Arrow, and the reader restores them,
matching `to_pandas_df()` without guessing from ordinary strings.
Regression test:
`tests/unit_tests/sql/execution/test_async_results.py::test_async_json_nested_values_match_sync`
covers both transports, both reader settings, dictionaries/lists, nulls, and a
JSON-looking string in the same column. The MessagePack writer cases fail at
c310e0254be44e31cb44d5b398273319b253b829 and pass at
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d.
Validation: `pytest tests/unit_tests/sql/execution/test_async_results.py
tests/unit_tests/sql/execution/test_celery_task.py -q -m "not manual and not
llm_compliance"` gives 21 failed / 148 passed with source at
c310e0254be44e31cb44d5b398273319b253b829, and 169 passed at
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d. The SQL plus result-set suites give
2,461 passed; staged pre-commit and mypy pass.
Siblings checked:
- superset/sql/execution/celery_task.py:189 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/sql/execution/executor.py:456 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/result_set.py:380 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/result_set.py:368 — not affected: synchronous conversion already
restores the same original nested values.
- superset/sql/execution/executor.py:1179 — not affected: cached records
retain original Python objects; test_cache_hit_returns_dataframes exercises
cache reconstruction.
- superset/sql_lab.py:401 — not affected: SQL Lab retains its existing
browser-oriented string representation and separate writer.
##########
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)
+
+ columns = [
+ column.get("column_name", column.get("name", ""))
+ for column in stmt_data.get("columns", [])
+ ]
+ arrow_types = {}
+ if schema_data := stmt_data.get("arrow_schema") or
stmt_data.get("temporal_schema"):
+ import base64
+
+ schema =
pa.ipc.read_schema(pa.BufferReader(base64.b64decode(schema_data)))
+ arrow_types = {field.name: field.type for field in schema}
+ return _records_to_dataframe(data, columns, arrow_types)
+
+
+def _json_column_to_series(
+ name: str, values: list[Any], arrow_type: pa.DataType
Review Comment:
I used the shared Arrow approach in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d: MessagePack carries IPC bytes; JSON
carries those same bytes as Base64. The type-by-type JSON decoder, schema
sidecars, and binary/duration rewrites are removed.
Regression tests:
`tests/unit_tests/sql/execution/test_async_results.py::test_async_payload_uses_one_arrow_format`,
`::test_async_mixed_large_integers_match_sync`, and
`::test_async_month_day_nano_interval_matches_sync`. The payload-shape tests
and JSON scalar regressions fail at c310e0254be44e31cb44d5b398273319b253b829
and pass at 98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d.
Validation: `pytest tests/unit_tests/sql/execution/test_async_results.py
tests/unit_tests/sql/execution/test_celery_task.py -q -m "not manual and not
llm_compliance"` gives 21 failed / 148 passed with source at
c310e0254be44e31cb44d5b398273319b253b829, and 169 passed at
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d. The SQL plus result-set suites give
2,461 passed; staged pre-commit and mypy pass.
Siblings checked:
- superset/sql/execution/celery_task.py:305 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/sql/execution/executor.py:433 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/sql_lab.py:392 — not affected: SQL Lab uses its separate
top-level data/selected_columns payload; its JSON records are intended for
browser clients.
- superset/views/utils.py:650 — not affected: reads that separate SQL Lab
payload, not SQLExecutor statement data.
- superset/commands/sql_lab/results.py — not affected: the SQL Lab command
delegates to its existing top-level payload reader.
- superset/commands/sql_lab/export.py:123 — not affected: consumes the
separate top-level SQL Lab data/columns payload for CSV.
##########
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:
I removed the paragraph from the admin page and documented the Python
contract and legacy-result limitations in
`docs/developer_docs/extensions/mcp.md` in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d. I also narrowed the cache
description: an empty/all-null unparseable fixed-offset dtype is a miss.
Regression tests for the formerly overstated behavior:
`tests/unit_tests/sql/execution/test_async_results.py::test_timezone_cache_hit_dtype_matches_sync`,
`::test_async_mixed_large_integers_match_sync`, and
`::test_async_month_day_nano_interval_matches_sync`; the reported scenarios
fail at c310e0254be44e31cb44d5b398273319b253b829 and pass at
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d.
Validation: `pytest tests/unit_tests/sql/execution/test_async_results.py
tests/unit_tests/sql/execution/test_celery_task.py -q -m "not manual and not
llm_compliance"` gives 21 failed / 148 passed with source at
c310e0254be44e31cb44d5b398273319b253b829, and 169 passed at
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d. The SQL plus result-set suites give
2,461 passed; staged pre-commit and mypy pass.
Siblings checked:
- docs/admin_docs/configuration/async-queries-celery.mdx — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- docs/developer_docs/extensions/mcp.md:31 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/sql/execution/executor.py:433 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/sql/execution/executor.py:1179 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/sql_lab.py:392 — not affected: the documented Python contract
does not change SQL Lab's public JSON format.
--
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]
