parent
731c19b630
commit
7090725de9
|
|
@ -106,11 +106,13 @@ class SupersetDataFrame(object):
|
|||
self.column_names = column_names
|
||||
|
||||
if dtype:
|
||||
# put data in a 2D array so we can efficiently access each column;
|
||||
# the reshape ensures the shape is 2D in case data is empty
|
||||
array = np.array(data, dtype="object").reshape(-1, len(column_names))
|
||||
# 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, dtype="object")
|
||||
data = {
|
||||
column: pd.Series(array[:, i], dtype=dtype[column])
|
||||
for i, column in enumerate(column_names)
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@
|
|||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
|
||||
from superset.dataframe import dedup, SupersetDataFrame
|
||||
from superset.db_engine_specs import BaseEngineSpec
|
||||
|
|
@ -135,3 +136,23 @@ class SupersetDataFrameTestCase(SupersetTestCase):
|
|||
cursor_descr = [("ds", "timestamp", None, None, None, None, True)]
|
||||
cdf = SupersetDataFrame(data, cursor_descr, PrestoEngineSpec)
|
||||
self.assertEqual(cdf.raw_df.dtypes[0], np.dtype("<M8[ns]"))
|
||||
|
||||
def test_no_type_coercion(self):
|
||||
data = [("a", 1), ("b", 2)]
|
||||
cursor_descr = [
|
||||
("one", "varchar", None, None, None, None, True),
|
||||
("two", "integer", None, None, None, None, True),
|
||||
]
|
||||
cdf = SupersetDataFrame(data, cursor_descr, PrestoEngineSpec)
|
||||
self.assertEqual(cdf.raw_df.dtypes[0], np.dtype("O"))
|
||||
self.assertEqual(cdf.raw_df.dtypes[1], pd.Int64Dtype())
|
||||
|
||||
def test_empty_data(self):
|
||||
data = []
|
||||
cursor_descr = [
|
||||
("one", "varchar", None, None, None, None, True),
|
||||
("two", "integer", None, None, None, None, True),
|
||||
]
|
||||
cdf = SupersetDataFrame(data, cursor_descr, PrestoEngineSpec)
|
||||
self.assertEqual(cdf.raw_df.dtypes[0], np.dtype("O"))
|
||||
self.assertEqual(cdf.raw_df.dtypes[1], pd.Int64Dtype())
|
||||
|
|
|
|||
Loading…
Reference in New Issue