Skip to content

Commit

Permalink
[fix] Enforce the QueryResult.df to be a pandas.DataFrame (Phase II) (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
john-bodley committed Jan 13, 2020
1 parent 7ff18f0 commit 0fe0471
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 12 deletions.
15 changes: 7 additions & 8 deletions superset/common/query_context.py
Expand Up @@ -87,7 +87,7 @@ def get_query_result(self, query_object: QueryObject) -> Dict[str, Any]:
# 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.
Expand Down Expand Up @@ -129,15 +129,14 @@ def get_data( # pylint: disable=invalid-name,no-self-use
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]]:
Expand Down Expand Up @@ -174,7 +173,7 @@ def get_df_payload( # pylint: disable=too-many-locals,too-many-statements
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
Expand Down Expand Up @@ -241,5 +240,5 @@ def get_df_payload( # pylint: disable=too-many-locals,too-many-statements
"query": query,
"status": status,
"stacktrace": stacktrace,
"rowcount": len(df.index) if df is not None else 0,
"rowcount": len(df.index),
}
6 changes: 2 additions & 4 deletions superset/connectors/druid/models.py
Expand Up @@ -1375,11 +1375,9 @@ def query(self, query_obj: Dict) -> QueryResult:
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", []))
Expand Down

0 comments on commit 0fe0471

Please sign in to comment.