aminghadersohi commented on code in PR #45348:
URL: https://github.com/apache/superset/pull/45348#discussion_r4288016501
##########
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:
Removed in 98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d, together with
`arrow_schema` reconstruction. Neither intermediate schema format is written or
interpreted; the real pre-PR JSON-record fallback remains.
Regression tests:
`tests/unit_tests/sql/execution/test_async_results.py::test_async_payload_uses_one_arrow_format`
asserts both schema keys are absent and fails at
c310e0254be44e31cb44d5b398273319b253b829, passing at
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d.
`::test_async_legacy_payload_without_schema` separately exercises the released
records format. `rg -n temporal_schema superset/sql/execution` returns no
matches; the intermediate-only test and helper pop are absent.
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/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:
Removed in 98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d. Both transports carry
their schema exactly once inside Arrow IPC; neither writes an `arrow_schema`
sidecar.
Regression test:
`tests/unit_tests/sql/execution/test_async_results.py::test_async_payload_uses_one_arrow_format`
fails for both writer settings at c310e0254be44e31cb44d5b398273319b253b829 and
passes at 98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d; it opens the IPC stream and
verifies the data as well as the absence of both schema sidecars.
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:
##########
@@ -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:
The Base64 import is at module scope in both modules in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d. The per-statement PyArrow import and
the local type-reconstruction imports disappeared with the shared Arrow
rewrite; pandas/PyArrow used by executor helpers are module imports.
Runtime regression test for that rewrite:
`tests/unit_tests/sql/execution/test_async_results.py::test_async_payload_uses_one_arrow_format`
fails at c310e0254be44e31cb44d5b398273319b253b829 and passes at
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d. The import placement itself is a
source-only cleanup, not a behavior claimed to be proved by that test. `rg -n
"import base64|import pyarrow" superset/sql/execution` shows only module-level
imports.
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:
##########
@@ -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:
The signature and `:param` entries both read `records, columns, dtypes` in
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d; `arrow_types` disappeared with the
Arrow rewrite. I checked the source order directly. This is documentation-only;
I am not claiming a runtime test detects docstring ordering.
Regression test for the helper's associated cache repair:
`tests/unit_tests/sql/execution/test_async_results.py::test_corrupt_cache_is_a_miss`
fails at c310e0254be44e31cb44d5b398273319b253b829 and passes 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/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]
