[fix] Enforce the QueryResult.df to be a pandas.DataFrame (Phase II) (#8948)

This commit is contained in:
John Bodley
2020-01-12 21:35:44 -08:00
committed by GitHub
parent 7ff18f060c
commit 0fe047171b
2 changed files with 9 additions and 12 deletions

View File

@@ -87,7 +87,7 @@ class QueryContext:
# be considered as the default ISO date format
# If the datetime format is unix, the parse will use the corresponding
# parsing logic
if df is not None and not df.empty:
if not df.empty:
if DTTM_ALIAS in df.columns:
if timestamp_format in ("epoch_s", "epoch_ms"):
# Column has already been formatted as a timestamp.
@@ -129,15 +129,14 @@ class QueryContext:
def get_single_payload(self, query_obj: QueryObject) -> Dict[str, Any]:
"""Returns a payload of metadata and data"""
payload = self.get_df_payload(query_obj)
df = payload.get("df")
status = payload.get("status")
df = payload["df"]
status = payload["status"]
if status != utils.QueryStatus.FAILED:
if df is None or df.empty:
if df.empty:
payload["error"] = "No data"
else:
payload["data"] = self.get_data(df)
if "df" in payload:
del payload["df"]
del payload["df"]
return payload
def get_payload(self) -> List[Dict[str, Any]]:
@@ -174,7 +173,7 @@ class QueryContext:
logging.info("Cache key: %s", cache_key)
is_loaded = False
stacktrace = None
df = None
df = pd.DataFrame()
cached_dttm = datetime.utcnow().isoformat().split(".")[0]
cache_value = None
status = None
@@ -241,5 +240,5 @@ class QueryContext:
"query": query,
"status": status,
"stacktrace": stacktrace,
"rowcount": len(df.index) if df is not None else 0,
"rowcount": len(df.index),
}

View File

@@ -1375,11 +1375,9 @@ class DruidDatasource(Model, BaseDatasource):
query_str = self.get_query_str(client=client, query_obj=query_obj, phase=2)
df = client.export_pandas()
if df is None or df.size == 0:
if df.empty:
return QueryResult(
df=pd.DataFrame(),
query=query_str,
duration=datetime.now() - qry_start_dttm,
df=df, query=query_str, duration=datetime.now() - qry_start_dttm
)
df = self.homogenize_types(df, query_obj.get("groupby", []))