aminghadersohi commented on code in PR #45348:
URL: https://github.com/apache/superset/pull/45348#discussion_r4288018787
##########
superset/sql/execution/executor.py:
##########
@@ -404,6 +405,94 @@ 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", [])
+ ]
+ return _records_to_dataframe(data, columns)
+
+
+def _records_to_dataframe(records: list[Any], columns: list[Any]) ->
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 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)
+
+ def to_series(values: list[Any]) -> pd.Series:
+ try:
+ array = pa.array(values)
+ except (pa.ArrowInvalid, pa.ArrowTypeError):
+ return pd.Series(values, dtype=object)
+ if pa.types.is_nested(array.type):
+ return pd.Series(values, dtype=object)
+ return array.to_pandas(integer_object_nulls=True)
Review Comment:
I need to correct my previous broad follow-up: at
c310e0254be44e31cb44d5b398273319b253b829, mixed small/large integers and
month/day/nanosecond intervals still failed in JSON mode, nested MessagePack
values remained strings, and fixed-offset cache hits failed. The scalar-only
examples and named-zone test did not establish those broader guarantees.
These cases are fixed for newly written results in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d with shared Arrow IPC, original nested
values, and cache inference/miss handling. Regression tests:
`tests/unit_tests/sql/execution/test_async_results.py::test_async_mixed_large_integers_match_sync`,
`::test_async_month_day_nano_interval_matches_sync`,
`::test_async_json_nested_values_match_sync`,
`::test_timezone_cache_hit_dtype_matches_sync`, and
`::test_corrupt_cache_is_a_miss`. The reported cases fail at
c310e0254be44e31cb44d5b398273319b253b829 and pass at
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d.
Remaining: legacy JSON records cannot regain types already serialized as
strings; rerun those queries. Intermediate-only schema formats were removed as
requested; the developer documentation states the released legacy-record
limitation.
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.
- superset/sql/execution/executor.py:598 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/sql/execution/executor.py:730 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/sql/execution/executor.py:1179 — fixed in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d
- superset/models/core.py:1689 — not affected: Database.execute delegates to
the tested synchronous executor path.
- superset/models/core.py:1705 — not affected: Database.execute_async
delegates to the tested asynchronous executor path.
- superset/sql/execution/executor.py:433 — not affected: async
results-backend data now uses Arrow directly, not cached pandas dtype strings.
- superset/mcp_service/sql_lab/tool/execute_sql.py:236 — not affected:
unchanged proxy calls Database.execute; the cache regressions exercise that
same synchronous entry point.
--
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]
