def _serialize_and_expand_data(
cdf: SupersetDataFrame,
db_engine_spec: BaseEngineSpec,
use_msgpack: Optional[bool] = False,
expand_data: bool = False,
) -> Tuple[Union[bytes, str], list, list, list]:
selected_columns: list = cdf.columns or []
expanded_columns: list
if use_msgpack:
with stats_timing("sqllab.query.results_backend_pa_serialization",
stats_logger):
data = (pa.default_serialization_context().serialize(
cdf.raw_df).to_buffer().to_pybytes())
# expand when loading data from results backend
all_columns, expanded_columns = (selected_columns, [])
else:
data = cdf.data or []
if expand_data:
all_columns, data, expanded_columns = db_engine_spec.expand_data(
selected_columns, data)
else:
all_columns = selected_columns
expanded_columns = []
return (data, selected_columns, all_columns, expanded_columns)
def _serialize_and_expand_data(
result_set: SupersetResultSet,
db_engine_spec: BaseEngineSpec,
use_msgpack: Optional[bool] = False,
expand_data: bool = False,
) -> Tuple[Union[bytes, str], List[Any], List[Any], List[Any]]:
selected_columns = result_set.columns
all_columns: List[Any]
expanded_columns: List[Any]
if use_msgpack:
with stats_timing("sqllab.query.results_backend_pa_serialization",
stats_logger):
data = (pa.default_serialization_context().serialize(
result_set.pa_table).to_buffer().to_pybytes())
# expand when loading data from results backend
all_columns, expanded_columns = (selected_columns, [])
else:
df = result_set.to_pandas_df()
data = df_to_records(df) or []
if expand_data:
all_columns, data, expanded_columns = db_engine_spec.expand_data(
selected_columns, data)
else:
all_columns = selected_columns
expanded_columns = []
return (data, selected_columns, all_columns, expanded_columns)
def test_get_datatype(self):
self.assertEquals('STRING', PrestoEngineSpec.get_datatype('string'))
self.assertEquals('TINY', MySQLEngineSpec.get_datatype(1))
self.assertEquals('VARCHAR', MySQLEngineSpec.get_datatype(15))
self.assertEquals('VARCHAR', BaseEngineSpec.get_datatype('VARCHAR'))
def test_get_datatype(self):
self.assertEquals('STRING', PrestoEngineSpec.get_datatype('string'))
self.assertEquals('TINY', MySQLEngineSpec.get_datatype(1))
self.assertEquals('VARCHAR', MySQLEngineSpec.get_datatype(15))
self.assertEquals('VARCHAR', BaseEngineSpec.get_datatype('VARCHAR'))
