mirror of https://github.com/apache/superset.git
Handle int64 columns with missing data in SQL Lab (#8226)
* Handle int64 columns with missing data in SQL Lab * Fix docstring * Add unit test * Small fix * Small fixes * Fix cursor description update * Better fix * Fix unit test, black * Fix nan comparison in unit test
This commit is contained in:
parent
1ad17936ba
commit
b9be01fcd8
|
@ -62,6 +62,12 @@ def dedup(l, suffix="__", case_sensitive=True):
|
||||||
return new_l
|
return new_l
|
||||||
|
|
||||||
|
|
||||||
|
def is_numeric(dtype):
|
||||||
|
if hasattr(dtype, "_is_numeric"):
|
||||||
|
return dtype._is_numeric
|
||||||
|
return np.issubdtype(dtype, np.number)
|
||||||
|
|
||||||
|
|
||||||
class SupersetDataFrame(object):
|
class SupersetDataFrame(object):
|
||||||
# Mapping numpy dtype.char to generic database types
|
# Mapping numpy dtype.char to generic database types
|
||||||
type_map = {
|
type_map = {
|
||||||
|
@ -80,21 +86,45 @@ class SupersetDataFrame(object):
|
||||||
}
|
}
|
||||||
|
|
||||||
def __init__(self, data, cursor_description, db_engine_spec):
|
def __init__(self, data, cursor_description, db_engine_spec):
|
||||||
column_names = []
|
|
||||||
if cursor_description:
|
|
||||||
column_names = [col[0] for col in cursor_description]
|
|
||||||
|
|
||||||
self.column_names = dedup(column_names)
|
|
||||||
|
|
||||||
data = data or []
|
data = data or []
|
||||||
self.df = pd.DataFrame(list(data), columns=self.column_names).infer_objects()
|
|
||||||
|
column_names = []
|
||||||
|
dtype = None
|
||||||
|
if cursor_description:
|
||||||
|
# get deduped list of column names
|
||||||
|
column_names = dedup([col[0] for col in cursor_description])
|
||||||
|
|
||||||
|
# fix cursor descriptor with the deduped names
|
||||||
|
cursor_description = [
|
||||||
|
tuple([column_name, *list(description)[1:]])
|
||||||
|
for column_name, description in zip(column_names, cursor_description)
|
||||||
|
]
|
||||||
|
|
||||||
|
# get type for better type casting, if possible
|
||||||
|
dtype = db_engine_spec.get_pandas_dtype(cursor_description)
|
||||||
|
|
||||||
|
self.column_names = column_names
|
||||||
|
|
||||||
|
if dtype:
|
||||||
|
# convert each column in data into a Series of the proper dtype; we
|
||||||
|
# need to do this because we can not specify a mixed dtype when
|
||||||
|
# instantiating the DataFrame, and this allows us to have different
|
||||||
|
# dtypes for each column.
|
||||||
|
array = np.array(data)
|
||||||
|
data = {
|
||||||
|
column: pd.Series(array[:, i], dtype=dtype[column])
|
||||||
|
for i, column in enumerate(column_names)
|
||||||
|
}
|
||||||
|
self.df = pd.DataFrame(data, columns=column_names)
|
||||||
|
else:
|
||||||
|
self.df = pd.DataFrame(list(data), columns=column_names).infer_objects()
|
||||||
|
|
||||||
self._type_dict = {}
|
self._type_dict = {}
|
||||||
try:
|
try:
|
||||||
# The driver may not be passing a cursor.description
|
# The driver may not be passing a cursor.description
|
||||||
self._type_dict = {
|
self._type_dict = {
|
||||||
col: db_engine_spec.get_datatype(cursor_description[i][1])
|
col: db_engine_spec.get_datatype(cursor_description[i][1])
|
||||||
for i, col in enumerate(self.column_names)
|
for i, col in enumerate(column_names)
|
||||||
if cursor_description
|
if cursor_description
|
||||||
}
|
}
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
@ -183,7 +213,7 @@ class SupersetDataFrame(object):
|
||||||
if (
|
if (
|
||||||
hasattr(dtype, "type")
|
hasattr(dtype, "type")
|
||||||
and issubclass(dtype.type, np.generic)
|
and issubclass(dtype.type, np.generic)
|
||||||
and np.issubdtype(dtype, np.number)
|
and is_numeric(dtype)
|
||||||
):
|
):
|
||||||
return "sum"
|
return "sum"
|
||||||
return None
|
return None
|
||||||
|
|
|
@ -275,6 +275,12 @@ class BaseEngineSpec:
|
||||||
return type_code.upper()
|
return type_code.upper()
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def get_pandas_dtype(
|
||||||
|
cls, cursor_description: List[tuple]
|
||||||
|
) -> Optional[Dict[str, str]]:
|
||||||
|
return None
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def extra_table_metadata(
|
def extra_table_metadata(
|
||||||
cls, database, table_name: str, schema_name: str
|
cls, database, table_name: str, schema_name: str
|
||||||
|
|
|
@ -39,6 +39,21 @@ from superset.utils import core as utils
|
||||||
|
|
||||||
QueryStatus = utils.QueryStatus
|
QueryStatus = utils.QueryStatus
|
||||||
|
|
||||||
|
# map between Presto types and Pandas
|
||||||
|
pandas_dtype_map = {
|
||||||
|
"boolean": "bool",
|
||||||
|
"tinyint": "Int64", # note: capital "I" means nullable int
|
||||||
|
"smallint": "Int64",
|
||||||
|
"integer": "Int64",
|
||||||
|
"bigint": "Int64",
|
||||||
|
"real": "float64",
|
||||||
|
"double": "float64",
|
||||||
|
"varchar": "object",
|
||||||
|
"timestamp": "datetime64",
|
||||||
|
"date": "datetime64",
|
||||||
|
"varbinary": "object",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
class PrestoEngineSpec(BaseEngineSpec):
|
class PrestoEngineSpec(BaseEngineSpec):
|
||||||
engine = "presto"
|
engine = "presto"
|
||||||
|
@ -1052,3 +1067,9 @@ class PrestoEngineSpec(BaseEngineSpec):
|
||||||
if df.empty:
|
if df.empty:
|
||||||
return ""
|
return ""
|
||||||
return df.to_dict()[field_to_return][0]
|
return df.to_dict()[field_to_return][0]
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def get_pandas_dtype(cls, cursor_description: List[tuple]) -> Dict[str, str]:
|
||||||
|
return {
|
||||||
|
col[0]: pandas_dtype_map.get(col[1], "object") for col in cursor_description
|
||||||
|
}
|
||||||
|
|
|
@ -18,6 +18,7 @@ import numpy as np
|
||||||
|
|
||||||
from superset.dataframe import dedup, SupersetDataFrame
|
from superset.dataframe import dedup, SupersetDataFrame
|
||||||
from superset.db_engine_specs import BaseEngineSpec
|
from superset.db_engine_specs import BaseEngineSpec
|
||||||
|
from superset.db_engine_specs.presto import PrestoEngineSpec
|
||||||
from .base_tests import SupersetTestCase
|
from .base_tests import SupersetTestCase
|
||||||
|
|
||||||
|
|
||||||
|
@ -108,3 +109,23 @@ class SupersetDataFrameTestCase(SupersetTestCase):
|
||||||
cursor_descr = (("a", "string"), ("a", "string"))
|
cursor_descr = (("a", "string"), ("a", "string"))
|
||||||
cdf = SupersetDataFrame(data, cursor_descr, BaseEngineSpec)
|
cdf = SupersetDataFrame(data, cursor_descr, BaseEngineSpec)
|
||||||
self.assertListEqual(cdf.column_names, ["a", "a__1"])
|
self.assertListEqual(cdf.column_names, ["a", "a__1"])
|
||||||
|
|
||||||
|
def test_int64_with_missing_data(self):
|
||||||
|
data = [(None,), (1239162456494753670,), (None,), (None,), (None,), (None,)]
|
||||||
|
cursor_descr = [("user_id", "bigint", None, None, None, None, True)]
|
||||||
|
|
||||||
|
# the base engine spec does not provide a dtype based on the cursor
|
||||||
|
# description, so the column is inferred as float64 because of the
|
||||||
|
# missing data
|
||||||
|
cdf = SupersetDataFrame(data, cursor_descr, BaseEngineSpec)
|
||||||
|
np.testing.assert_array_equal(
|
||||||
|
cdf.raw_df.values.tolist(),
|
||||||
|
[[np.nan], [1.2391624564947538e18], [np.nan], [np.nan], [np.nan], [np.nan]],
|
||||||
|
)
|
||||||
|
|
||||||
|
# currently only Presto provides a dtype based on the cursor description
|
||||||
|
cdf = SupersetDataFrame(data, cursor_descr, PrestoEngineSpec)
|
||||||
|
np.testing.assert_array_equal(
|
||||||
|
cdf.raw_df.values.tolist(),
|
||||||
|
[[np.nan], [1239162456494753670], [np.nan], [np.nan], [np.nan], [np.nan]],
|
||||||
|
)
|
||||||
|
|
Loading…
Reference in New Issue