Fix no data in Presto (#8268)

* Fix no data in Presto

* Fix test
This commit is contained in:
Beto Dealmeida 2019-09-20 14:31:13 +02:00 committed by GitHub
parent 731c19b630
commit 7090725de9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 24 additions and 1 deletions

View File

@ -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)

View File

@ -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())