From c4be5de3f795038df051d2d134030d94e5017234 Mon Sep 17 00:00:00 2001
From: John Bodley <john.bodley@airbnb.com>
Date: Fri, 10 Jan 2020 18:19:35 -0800
Subject: [PATCH] [fix] Enforce the QueryResult.df to be a pandas.DataFrame
 (Phase II)

---
 superset/common/query_context.py    | 15 +++++++--------
 superset/connectors/druid/models.py |  6 ++----
 2 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/superset/common/query_context.py b/superset/common/query_context.py
index b234c14773928..ae57a63bf7335 100644
--- a/superset/common/query_context.py
+++ b/superset/common/query_context.py
@@ -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.
@@ -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]]:
@@ -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
@@ -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),
         }
diff --git a/superset/connectors/druid/models.py b/superset/connectors/druid/models.py
index eacf5dd8e133b..8889ea62606f2 100644
--- a/superset/connectors/druid/models.py
+++ b/superset/connectors/druid/models.py
@@ -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", []))