aminghadersohi commented on code in PR #45348:
URL: https://github.com/apache/superset/pull/45348#discussion_r4288013203
##########
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:
Fixed in 98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d. Invalid pandas dtype
strings fall back to value inference, and cache read/decode errors become
misses for both entry points. Empty/all-null fixed-offset columns that cannot
be inferred are misses, not incorrectly typed cache hits.
Regression tests:
`tests/unit_tests/sql/execution/test_async_results.py::test_timezone_cache_hit_dtype_matches_sync`
(named/fixed-offset zones, populated/empty/null-only data) and
`::test_corrupt_cache_is_a_miss` (sync and async). The fixed-offset and
corrupt-cache 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/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.
##########
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:
Fixed in 98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d by removing the JSON
scalar reconstruction path: JSON carries Base64 Arrow IPC, so mixed integers
and month/day/nanosecond intervals never reach inference or a per-column cast.
Regression tests:
`tests/unit_tests/sql/execution/test_async_results.py::test_async_mixed_large_integers_match_sync`
and `::test_async_month_day_nano_interval_matches_sync`. JSON writer cases
fail at c310e0254be44e31cb44d5b398273319b253b829 and pass at
98a05a717c1ed67250f39b1ab7a0b0cbcf320a9d; both codecs and both reader settings
compare values and dtypes against synchronous execution.
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.
--
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]
