You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by vi...@apache.org on 2026/06/05 16:26:05 UTC
(superset) branch master updated: fix(dataset): preserve numeric column types when pydruid infers STRING from first-row value (#40677)
This is an automated email from the ASF dual-hosted git repository.
villebro pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new fa42b13eb83 fix(dataset): preserve numeric column types when pydruid infers STRING from first-row value (#40677)
fa42b13eb83 is described below
commit fa42b13eb83649ad4b9e5c702f56e0de253e8b13
Author: madhushreeag <ma...@gmail.com>
AuthorDate: Fri Jun 5 09:25:57 2026 -0700
fix(dataset): preserve numeric column types when pydruid infers STRING from first-row value (#40677)
Co-authored-by: madhushree agarwal <ma...@apple.com>
---
superset/db_engine_specs/base.py | 33 +++++++
superset/db_engine_specs/druid.py | 26 ++++++
superset/result_set.py | 25 +++---
tests/unit_tests/connectors/sqla/utils_test.py | 1 +
tests/unit_tests/db_engine_specs/test_base.py | 21 +++++
tests/unit_tests/db_engine_specs/test_druid.py | 116 +++++++++++++++++++++++++
6 files changed, 209 insertions(+), 13 deletions(-)
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 5c33a3fd1f0..1143abc3df6 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -540,6 +540,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
# Set this to True on any engine spec where at least one row must be
# fetched for cursor.description to be populated.
type_probe_needs_row: bool = False
+ requires_column_value_normalization: bool = False
try_remove_schema_from_table_name = True # pylint: disable=invalid-name
run_multiple_statements_as_one = False
custom_errors: dict[
@@ -1304,6 +1305,38 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
return type_code.upper()
return None
+ @classmethod
+ def normalize_column_values(cls, col_values: list[Any]) -> list[Any]:
+ """
+ Engine-specific hook to normalize column values before PyArrow conversion.
+
+ Called when the initial pa.array() conversion raises an exception, giving
+ the engine a chance to clean up values (e.g. replace sentinel strings with
+ None) before a second conversion attempt.
+
+ :param col_values: Raw Python values for one column
+ :return: Normalized values; return the input list unchanged by default
+ """
+ return col_values
+
+ @classmethod
+ def resolve_column_type(
+ cls, cursor_type: str | None, pa_mapped: str | None
+ ) -> str | None:
+ """
+ Choose the reported column type from the cursor description type and the
+ type inferred by PyArrow.
+
+ The default prefers the cursor description when available. Override in
+ engine specs where the cursor description is unreliable (e.g. pydruid
+ infers STRING from a None or special-float first row value).
+
+ :param cursor_type: Type string from the cursor description, or None
+ :param pa_mapped: Type string inferred by PyArrow, or None
+ :return: The type string to report for this column
+ """
+ return cursor_type or pa_mapped
+
@classmethod
@deprecated(deprecated_in="3.0")
def normalize_indexes(cls, indexes: list[dict[str, Any]]) -> list[dict[str, Any]]:
diff --git a/superset/db_engine_specs/druid.py b/superset/db_engine_specs/druid.py
index f80c1e1d353..61845909c2a 100644
--- a/superset/db_engine_specs/druid.py
+++ b/superset/db_engine_specs/druid.py
@@ -37,6 +37,8 @@ if TYPE_CHECKING:
logger = logging.getLogger()
+_DRUID_FLOAT_SPECIAL = frozenset({"NaN", "Infinity", "-Infinity"})
+
class DruidEngineSpec(BaseEngineSpec):
"""Engine spec for Druid.io"""
@@ -49,6 +51,7 @@ class DruidEngineSpec(BaseEngineSpec):
# pydruid builds cursor.description from the first returned row, so a
# WHERE FALSE query (zero rows) leaves description as None.
type_probe_needs_row = True
+ requires_column_value_normalization = True
metadata = {
"description": (
@@ -215,3 +218,26 @@ class DruidEngineSpec(BaseEngineSpec):
return {
requests_exceptions.ConnectionError: SupersetDBAPIConnectionError,
}
+
+ @classmethod
+ def normalize_column_values(cls, col_values: list[Any]) -> list[Any]:
+ # pydruid emits NaN, Infinity, and -Infinity as JSON strings because
+ # the JSON spec does not support those values. Replace them with None
+ # so PyArrow can keep the column numeric after the initial conversion
+ # attempt fails.
+ return [
+ None if isinstance(v, str) and v in _DRUID_FLOAT_SPECIAL else v
+ for v in col_values
+ ]
+
+ @classmethod
+ def resolve_column_type(
+ cls, cursor_type: str | None, pa_mapped: str | None
+ ) -> str | None:
+ # pydruid infers column types from the first row value. A None or
+ # special-float-string first value causes the column to be labelled
+ # STRING even when the actual data is numeric. When PyArrow infers a
+ # more specific type, prefer it over the cursor-description STRING.
+ if cursor_type == "STRING" and pa_mapped is not None and pa_mapped != "STRING":
+ return pa_mapped
+ return cursor_type or pa_mapped
diff --git a/superset/result_set.py b/superset/result_set.py
index f8945fb4f58..9eaae7a3680 100644
--- a/superset/result_set.py
+++ b/superset/result_set.py
@@ -171,11 +171,14 @@ class SupersetResultSet:
# only do expensive recasting if datatype is not standard list of tuples
if data and (not isinstance(data, list) or not isinstance(data[0], tuple)):
data = [tuple(row) for row in data]
- array = np.array(data, dtype=numpy_dtype)
+ columns = np.array(data, dtype=numpy_dtype)
for column in column_names:
+ col_values = columns[column].tolist()
+ if db_engine_spec.requires_column_value_normalization:
+ col_values = db_engine_spec.normalize_column_values(col_values)
try:
- pa_data.append(pa.array(array[column].tolist()))
+ pa_data.append(pa.array(col_values))
except (
pa.lib.ArrowInvalid,
pa.lib.ArrowTypeError,
@@ -185,7 +188,7 @@ class SupersetResultSet:
# https://issues.apache.org/jira/browse/ARROW-7855
):
# attempt serialization of values as strings
- stringified_arr = stringify_values(array[column])
+ stringified_arr = stringify_values(columns[column])
pa_data.append(pa.array(stringified_arr.tolist()))
if pa_data: # pylint: disable=too-many-nested-blocks
@@ -194,19 +197,19 @@ class SupersetResultSet:
# TODO: revisit nested column serialization once nested types
# are added as a natively supported column type in Superset
# (superset.utils.core.GenericDataType).
- stringified_arr = stringify_values(array[column])
+ stringified_arr = stringify_values(columns[column])
pa_data[i] = pa.array(stringified_arr.tolist())
elif pa.types.is_temporal(pa_data[i].type):
# workaround for bug converting
# `psycopg2.tz.FixedOffsetTimezone` tzinfo values.
# related: https://issues.apache.org/jira/browse/ARROW-5248
- sample = self.first_nonempty(array[column])
+ sample = self.first_nonempty(columns[column])
if sample and isinstance(sample, datetime.datetime):
try:
if sample.tzinfo:
tz = sample.tzinfo
- series = pd.Series(array[column])
+ series = pd.Series(columns[column])
series = pd.to_datetime(
series, utc=True, errors="coerce"
)
@@ -277,13 +280,9 @@ class SupersetResultSet:
def data_type(self, col_name: str, pa_dtype: pa.DataType) -> Optional[str]:
"""Given a pyarrow data type, Returns a generic database type"""
- if set_type := self._type_dict.get(col_name):
- return set_type
-
- if mapped_type := self.convert_pa_dtype(pa_dtype):
- return mapped_type
-
- return None
+ set_type = self._type_dict.get(col_name)
+ pa_mapped = self.convert_pa_dtype(pa_dtype)
+ return self.db_engine_spec.resolve_column_type(set_type, pa_mapped)
def to_pandas_df(self) -> pd.DataFrame:
return self.convert_table_to_df(self.table)
diff --git a/tests/unit_tests/connectors/sqla/utils_test.py b/tests/unit_tests/connectors/sqla/utils_test.py
index 6c6f5e38926..35e1eb34c3e 100644
--- a/tests/unit_tests/connectors/sqla/utils_test.py
+++ b/tests/unit_tests/connectors/sqla/utils_test.py
@@ -48,6 +48,7 @@ def test_returns_column_descriptions(mocker: MockerFixture) -> None:
database.mutate_sql_based_on_config.return_value = "SELECT * FROM table LIMIT 1"
db_engine_spec.fetch_data.return_value = [("col1", "col1", "STRING", None, False)]
db_engine_spec.get_datatype.return_value = "STRING"
+ db_engine_spec.resolve_column_type.return_value = "STRING"
db_engine_spec.get_column_spec.return_value.is_dttm = False
db_engine_spec.get_column_spec.return_value.generic_type = "STRING"
diff --git a/tests/unit_tests/db_engine_specs/test_base.py b/tests/unit_tests/db_engine_specs/test_base.py
index b0e416e6081..e46517a9927 100644
--- a/tests/unit_tests/db_engine_specs/test_base.py
+++ b/tests/unit_tests/db_engine_specs/test_base.py
@@ -1339,3 +1339,24 @@ def test_get_view_names_strips_schema_with_regex_metacharacters(
# "axbc.other" would match the old unescaped pattern ^a.b(c)\. and be
# incorrectly stripped — the escaped version correctly preserves it.
assert views == {"report", "axbc.other"}
+
+
+def test_normalize_column_values_is_identity() -> None:
+ """BaseEngineSpec.normalize_column_values must return the input unchanged."""
+ values = [1, None, "text", 3.14]
+ assert BaseEngineSpec.normalize_column_values(values) is values
+
+
+def test_resolve_column_type_prefers_cursor_type() -> None:
+ """cursor_type wins over pa_mapped when both are present."""
+ assert BaseEngineSpec.resolve_column_type("INTEGER", "FLOAT") == "INTEGER"
+
+
+def test_resolve_column_type_falls_back_to_pa_mapped() -> None:
+ """pa_mapped is used when cursor_type is absent."""
+ assert BaseEngineSpec.resolve_column_type(None, "FLOAT") == "FLOAT"
+
+
+def test_resolve_column_type_returns_none_when_both_absent() -> None:
+ """None is returned when neither source provides a type."""
+ assert BaseEngineSpec.resolve_column_type(None, None) is None
diff --git a/tests/unit_tests/db_engine_specs/test_druid.py b/tests/unit_tests/db_engine_specs/test_druid.py
index 4d6f84e1ee6..6a55c1e0b9a 100644
--- a/tests/unit_tests/db_engine_specs/test_druid.py
+++ b/tests/unit_tests/db_engine_specs/test_druid.py
@@ -18,9 +18,12 @@ from datetime import datetime
from typing import Optional
from unittest import mock
+import pandas as pd
import pytest
from sqlalchemy import column
+from superset.db_engine_specs.base import BaseEngineSpec
+from superset.result_set import SupersetResultSet
from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm
from tests.unit_tests.fixtures.common import dttm # noqa: F401
@@ -95,3 +98,116 @@ def test_extras_with_ssl() -> None:
connect_args = extras["engine_params"]["connect_args"]
assert connect_args["scheme"] == "https"
assert "ssl_verify_cert" in connect_args
+
+
+# ---------------------------------------------------------------------------
+# DruidEngineSpec column normalization tests
+#
+# pydruid infers column types from the first row value, which causes two
+# related problems:
+#
+# Case 1 – Mixed IEEE special-float strings and numbers:
+# Druid cannot represent NaN/Infinity in JSON, so pydruid emits them as
+# the strings "NaN", "Infinity", or "-Infinity". When these appear in a
+# numeric column, pa.array() raises ArrowInvalid on the mixed str/float
+# list and the column falls back to string serialisation.
+#
+# Case 2 – None as the first value:
+# pydruid's get_type(None) returns Type.STRING, so any nullable numeric
+# column whose first row is null gets labelled STRING in the cursor
+# description. pa.array() succeeds (producing float64) but
+# data_type() used to return STRING because the cursor description won.
+#
+# DruidEngineSpec overrides normalize_column_values and resolve_column_type
+# to handle both cases. BaseEngineSpec preserves the original behaviour.
+# ---------------------------------------------------------------------------
+
+
+def test_druid_ieee_special_floats_preserved_as_numeric() -> None:
+ """
+ Case 1, DruidEngineSpec: columns that mix IEEE special-float strings with
+ real numbers must keep their numeric type (specials become null).
+ """
+ from superset.db_engine_specs.druid import DruidEngineSpec
+
+ data = [("NaN",), (1.5,), ("Infinity",), (2.3,), ("-Infinity",), (None,)]
+ description = [("metric", "STRING", None, None, None, None, None)]
+ result_set = SupersetResultSet(data, description, DruidEngineSpec) # type: ignore
+
+ col = result_set.columns[0]
+ assert col["type"] == "FLOAT"
+
+ df = result_set.to_pandas_df()
+ assert pd.isna(df["metric"].iloc[0]) # "NaN" → null
+ assert df["metric"].iloc[1] == 1.5
+ assert pd.isna(df["metric"].iloc[2]) # "Infinity" → null
+ assert df["metric"].iloc[3] == 2.3
+ assert pd.isna(df["metric"].iloc[4]) # "-Infinity" → null
+ assert pd.isna(df["metric"].iloc[5]) # None → null
+
+
+def test_base_spec_ieee_special_floats_stringified() -> None:
+ """
+ Case 1, BaseEngineSpec: without Druid's override, columns with mixed
+ special-float strings and numbers fall through to string serialisation.
+ """
+ data = [("NaN",), (1.5,), ("Infinity",)]
+ description = [("metric", "STRING", None, None, None, None, None)]
+ result_set = SupersetResultSet(data, description, BaseEngineSpec) # type: ignore
+
+ col = result_set.columns[0]
+ assert col["type"] == "STRING"
+
+ df = result_set.to_pandas_df()
+ assert df["metric"].iloc[0] == "NaN"
+ assert df["metric"].iloc[1] == "1.5"
+ assert df["metric"].iloc[2] == "Infinity"
+
+
+def test_druid_none_first_value_reports_numeric_type() -> None:
+ """
+ Case 2, DruidEngineSpec: when the cursor description says STRING (pydruid's
+ first-row None inference) but PyArrow correctly infers float64, the column
+ must be reported as FLOAT, not STRING.
+ """
+ from superset.db_engine_specs.druid import DruidEngineSpec
+
+ data = [(None,), (1.5,), (2.3,), (None,), (4.7,)]
+ description = [("metric", "STRING", None, None, None, None, None)]
+ result_set = SupersetResultSet(data, description, DruidEngineSpec) # type: ignore
+
+ col = result_set.columns[0]
+ assert col["type"] == "FLOAT"
+
+ df = result_set.to_pandas_df()
+ assert pd.isna(df["metric"].iloc[0])
+ assert df["metric"].iloc[1] == 1.5
+ assert df["metric"].iloc[4] == 4.7
+
+
+def test_base_spec_none_first_value_reports_string_type() -> None:
+ """
+ Case 2, BaseEngineSpec: the cursor-description STRING type must continue
+ to win over PyArrow's float64 inference for non-Druid engines.
+ """
+ data = [(None,), (1.5,), (2.3,)]
+ description = [("metric", "STRING", None, None, None, None, None)]
+ result_set = SupersetResultSet(data, description, BaseEngineSpec) # type: ignore
+
+ col = result_set.columns[0]
+ assert col["type"] == "STRING"
+
+
+def test_non_string_cursor_type_unaffected_by_druid_spec() -> None:
+ """
+ Columns with a non-STRING cursor description type must not be affected by
+ DruidEngineSpec's resolve_column_type override.
+ """
+ from superset.db_engine_specs.druid import DruidEngineSpec
+
+ data = [(1,), (2,), (3,)]
+ description = [("count", "INT", None, None, None, None, None)]
+ result_set = SupersetResultSet(data, description, DruidEngineSpec) # type: ignore
+
+ col = result_set.columns[0]
+ assert col["type"] == "INT"