rebenitez1802 commented on code in PR #45348:
URL: https://github.com/apache/superset/pull/45348#discussion_r4287278853
##########
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
+) -> pd.Series:
+ """Restore a JSON column using its original Arrow scalar type."""
+ import base64
+ from datetime import time
+
+ import pandas as pd
+ import pyarrow as pa
+
+ from superset.result_set import SupersetResultSet
+
+ if pa.types.is_time(arrow_type):
+ values = [time.fromisoformat(v) if isinstance(v, str) else v for v in
values]
+ elif (
+ pa.types.is_binary(arrow_type)
+ or pa.types.is_large_binary(arrow_type)
+ or pa.types.is_fixed_size_binary(arrow_type)
+ ):
+ values = [base64.b64decode(v) if isinstance(v, str) else v for v in
values]
+ # Result sets retain nested/JSON values as Python objects even when
+ # their Arrow representation is a string column.
+ if any(isinstance(value, (dict, list)) for value in values):
+ return pd.Series(values, dtype=object)
+ if pa.types.is_duration(arrow_type):
+ array = pa.array(values, type=arrow_type)
+ else:
+ array = pa.array(values).cast(arrow_type)
Review Comment:
**🟡 Medium — One unconvertible column fails the whole JSON-mode result**
With `RESULTS_BACKEND_USE_MSGPACK = False`, `df_to_records(df,
convert_decimals=True)` (`celery_task.py:357`) turns only integers above 2^53
into strings. A column like `[1, 2**60]` then reaches this `cast` as mixed
int/str and raises. The outer handler turns the entire result into `FAILED:
Error loading results: Could not convert '1152921504606846976' with type str:
tried to convert to int64`, while `get_status()` reports SUCCESS. Arrow
`month_day_nano_interval` columns (e.g. Trino `INTERVAL YEAR TO MONTH`) fail
the same way. 64-bit hash columns, or large IDs mixed with sentinels like
`0`/`-1`, will hit this; the PR's first commit returned them (as strings)
instead of failing.
Fix: pass `convert_big_integers=False` in this writer (only Python reads
this payload), wrap `_json_column_to_series` in a per-column try/except that
falls back to type inference, and add a `[1, 2**60]` test.
##########
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
+) -> pd.Series:
+ """Restore a JSON column using its original Arrow scalar type."""
+ import base64
+ from datetime import time
+
+ import pandas as pd
+ import pyarrow as pa
+
+ from superset.result_set import SupersetResultSet
+
+ if pa.types.is_time(arrow_type):
+ values = [time.fromisoformat(v) if isinstance(v, str) else v for v in
values]
+ elif (
+ pa.types.is_binary(arrow_type)
+ or pa.types.is_large_binary(arrow_type)
+ or pa.types.is_fixed_size_binary(arrow_type)
+ ):
+ values = [base64.b64decode(v) if isinstance(v, str) else v for v in
values]
+ # Result sets retain nested/JSON values as Python objects even when
+ # their Arrow representation is a string column.
+ if any(isinstance(value, (dict, list)) for value in values):
+ return pd.Series(values, dtype=object)
+ if pa.types.is_duration(arrow_type):
+ array = pa.array(values, type=arrow_type)
+ else:
+ array = pa.array(values).cast(arrow_type)
+ return SupersetResultSet.convert_table_to_df(pa.table({name: array}))[name]
+
+
+def _records_to_dataframe(
+ records: list[Any],
+ columns: list[Any],
+ arrow_types: dict[str, pa.DataType] | None = None,
+ dtypes: list[str] | None = None,
+) -> pd.DataFrame:
+ """
+ Build a DataFrame from JSON-style records.
+
+ Values are converted through Arrow, as SELECT results are on the
+ synchronous path, so integer columns with nulls keep ``None`` instead of
+ becoming ``NaN`` floats; mixed or nested values are kept as Python
+ objects. The given columns are kept even when there are no records.
+
+ :param records: Rows as dicts keyed by column name
+ :param arrow_types: Original Arrow types for JSON-serialized columns
+ :param dtypes: Original pandas dtypes for cached DataFrames
+ :param columns: Column names, in order
+ :returns: DataFrame of the records
+ """
+ import pandas as pd
+ import pyarrow as pa
+
+ if not columns or not all(isinstance(row, dict) for row in records):
+ return pd.DataFrame(records, columns=columns or None)
+
+ column_dtypes = dict(zip(columns, dtypes, strict=True)) if dtypes else {}
+
+ def to_series(name: str, values: list[Any]) -> pd.Series:
+ # DataFrame.to_dict keeps pandas' missing timestamp sentinel on cache
hits.
+ values = [None if value is pd.NaT else value for value in values]
+ if name in column_dtypes:
+ return pd.Series(values, dtype=column_dtypes[name])
Review Comment:
**🔴 High — Cache hits fail for fixed-offset timezone columns (regression)**
`_store_in_cache` saves `str(dtype)` (`executor.py:1307-1311`) and this line
rebuilds each column with `pd.Series(values, dtype=...)`. pandas can't parse
its own fixed-offset timezone strings: `datetime64[ns, pytz.FixedOffset(-240)]`
raises `TypeError: Invalid datetime unit in metadata string`. Example with
pandas 2.3.3 / pyarrow 25: Postgres `timestamptz` in a non-UTC session. The
first `database.execute()` succeeds, every repeat returns FAILED until the
cache entry expires, and `execute_async()` raises `TypeError` to the caller
because `_try_get_cached_result` (`:763`) isn't in a try block. Master returns
SUCCESS here, so this is a regression on the sync path that MCP `execute_sql`
uses. `test_timezone_cache_hit_dtype_matches_sync` only covers named zones.
Fix: fall back to `pd.Series(values)` when the dtype string doesn't parse
(it infers the correct fixed-offset dtype), treat any cache decode error as a
cache miss, and add a `timezone(timedelta(hours=-4))` test case.
##########
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:
**🟡 Medium — Nested columns come back as JSON strings in the default
MessagePack mode**
This converts the Arrow data with `convert_table_to_df`, but the step that
turns nested values back into Python objects only runs in
`SupersetResultSet.to_pandas_df()`. So sync and JSON-mode async return
`[{'key': [1, 2]}]`, but MessagePack-mode async (the default) returns
`['{"key": [1, 2]}']`, although the docstring promises the same result as sync.
`test_async_json_nested_values_match_sync` (`test_async_results.py:721`) only
runs with `write_msgpack=False`.
Fix: run the test with both settings, then either store the nested column
names in the payload and restore them on read, or document the difference.
##########
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"):
Review Comment:
**🟢 Low — `temporal_schema` support is dead compatibility code**
This line and `test_async_legacy_temporal_schema`
(`test_async_results.py:668`) read a key that only this PR's intermediate
commit b522f99 ever wrote; no released version wrote it.
Fix: remove the `temporal_schema` fallback (suggestion below), that test,
and the `pop("temporal_schema")` at `test_async_results.py:113`.
```suggestion
if schema_data := stmt_data.get("arrow_schema"):
```
##########
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:
**🟢 Low — New docs paragraph overstates guarantees and sits in admin docs**
This says large integers and intervals are preserved in either format and
that cache hits keep timezones; the fixed-offset cache failure and the
JSON-mode large-integer/interval failures contradict both claims. It also
documents a Python extension API's return types on an admin deployment page.
Fix: drop the paragraph, or move it to the developer/extension docs once the
behavior is fixed.
##########
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:
**🟡 Medium — JSON mode hand-builds a second serialization format**
`celery_task.py:189-216` and `executor.py:469-497` rebuild Arrow types from
JSON one type at a time (Base64 binary, duration ticks, `time.fromisoformat`,
casts). The large-integer and interval failures are the first gaps, and any
type that isn't special-cased will drift. Only `SQLExecutor` reads these
payloads: `celery_task` is imported only by `executor.py`, and SQL Lab never
fetches these `results_key`s. So JSON mode could store the Arrow data as Base64
and reuse the MessagePack reader path. A quick prototype is about +14/−78 lines
against the PR head. It passes the scalar and temporal tests in all four
writer/reader setting combinations, fixes the large-integer case, and stores
less (~11 MB vs ~18 MB for 100k rows × 10 columns).
Fix: switch JSON mode to Base64 Arrow data, or land the decoding fix alone
and move type fidelity to a follow-up PR.
##########
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
+) -> pd.Series:
+ """Restore a JSON column using its original Arrow scalar type."""
+ import base64
Review Comment:
**🟢 Low — Imports inside functions and loops**
`import base64` is at module level in `celery_task.py` but imported locally
twice here (`:462`, `:473`). In `celery_task.py:195`, `import pyarrow as pa`
sits inside the per-statement loop, though `superset.result_set` already
imports pyarrow at module load.
Fix: move them to module level.
##########
superset/sql/execution/celery_task.py:
##########
@@ -185,6 +186,35 @@ def _finalize_successful_query(
"execution_time_ms": exec_time,
}
)
+ # JSON records lose scalar types. Retain the complete Arrow schema
+ # and encode binary columns losslessly instead of decoding as text.
+ statements_data[-1]["arrow_schema"] = base64.b64encode(
+ result_set.pa_table.schema.serialize().to_pybytes()
+ ).decode("ascii")
+ if isinstance(data, list):
+ import pyarrow as pa
+
Review Comment:
**🟢 Low — `arrow_schema` is also written in MessagePack mode**
This adds a Base64 schema to every SELECT, but in MessagePack mode the Arrow
data already carries it and the reader returns before looking at it
(`executor.py:450-454`). That is about 2.5 KB per 30-column statement, counted
against `SQLLAB_PAYLOAD_MAX_MB`.
Fix: only write it for JSON records:
```suggestion
if isinstance(data, list):
import pyarrow as pa
# JSON records lose scalar types. Retain the complete Arrow
schema
# and encode binary columns losslessly instead of decoding
as text.
statements_data[-1]["arrow_schema"] = base64.b64encode(
result_set.pa_table.schema.serialize().to_pybytes()
).decode("ascii")
```
##########
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
+) -> pd.Series:
+ """Restore a JSON column using its original Arrow scalar type."""
+ import base64
+ from datetime import time
+
+ import pandas as pd
+ import pyarrow as pa
+
+ from superset.result_set import SupersetResultSet
+
+ if pa.types.is_time(arrow_type):
+ values = [time.fromisoformat(v) if isinstance(v, str) else v for v in
values]
+ elif (
+ pa.types.is_binary(arrow_type)
+ or pa.types.is_large_binary(arrow_type)
+ or pa.types.is_fixed_size_binary(arrow_type)
+ ):
+ values = [base64.b64decode(v) if isinstance(v, str) else v for v in
values]
+ # Result sets retain nested/JSON values as Python objects even when
+ # their Arrow representation is a string column.
+ if any(isinstance(value, (dict, list)) for value in values):
+ return pd.Series(values, dtype=object)
+ if pa.types.is_duration(arrow_type):
+ array = pa.array(values, type=arrow_type)
+ else:
+ array = pa.array(values).cast(arrow_type)
+ return SupersetResultSet.convert_table_to_df(pa.table({name: array}))[name]
+
+
+def _records_to_dataframe(
+ records: list[Any],
+ columns: list[Any],
+ arrow_types: dict[str, pa.DataType] | None = None,
+ dtypes: list[str] | None = None,
+) -> pd.DataFrame:
+ """
+ Build a DataFrame from JSON-style records.
+
+ Values are converted through Arrow, as SELECT results are on the
+ synchronous path, so integer columns with nulls keep ``None`` instead of
+ becoming ``NaN`` floats; mixed or nested values are kept as Python
+ objects. The given columns are kept even when there are no records.
+
+ :param records: Rows as dicts keyed by column name
+ :param arrow_types: Original Arrow types for JSON-serialized columns
+ :param dtypes: Original pandas dtypes for cached DataFrames
+ :param columns: Column names, in order
Review Comment:
**🟢 Low — `_records_to_dataframe` docstring param order**
These `:param` lines list `records, arrow_types, dtypes, columns`; the
signature order is `records, columns, arrow_types, dtypes`.
```suggestion
:param records: Rows as dicts keyed by column name
:param columns: Column names, in order
:param arrow_types: Original Arrow types for JSON-serialized columns
:param dtypes: Original pandas dtypes for cached DataFrames
```
--
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]
