Browse Source

[sparksql] Show result for complex data types in correct format (#3097)

Harsh Gupta 3 years ago
parent
commit
4b8fadf114
1 changed files with 44 additions and 2 deletions
  1. 44 2
      desktop/libs/notebook/src/notebook/connectors/spark_shell.py

+ 44 - 2
desktop/libs/notebook/src/notebook/connectors/spark_shell.py

@@ -269,9 +269,9 @@ class SparkApi(Api):
           images = []
           images = []
         if 'application/json' in data:
         if 'application/json' in data:
           result = data['application/json']
           result = data['application/json']
-          data = result['data']
-          meta = [{'name': field['name'], 'type': field['type'], 'comment': ''} for field in result['schema']['fields']]
           type = 'table'
           type = 'table'
+          meta, is_complex_type = self._handle_result_meta(result)
+          data = self._handle_result_data(result, is_complex_type)
         else:
         else:
           data = [[data['text/plain']]]
           data = [[data['text/plain']]]
           meta = [{'name': 'Header', 'type': 'STRING_TYPE', 'comment': ''}]
           meta = [{'name': 'Header', 'type': 'STRING_TYPE', 'comment': ''}]
@@ -302,6 +302,48 @@ class SparkApi(Api):
       raise QueryError(msg)
       raise QueryError(msg)
 
 
 
 
+  def _handle_result_data(self, result, is_complex_type=False):
+    data = []
+
+    if is_complex_type:
+      for row in result['data']:
+        row_data = []
+        for ele in row:
+          if isinstance(ele, dict):
+            row_schema = []
+            for val in ele['schema']:
+              row_schema.append(val['name'])
+            row_data.append(dict(zip(row_schema, ele['values'])))
+          else:
+            row_data.append(ele)
+
+        data.append(row_data)
+    else:
+      data = result['data']
+    
+    return data
+
+
+  def _handle_result_meta(self, result):
+    meta = []
+    is_complex_type = False
+
+    for f in result['schema']['fields']:
+      if isinstance(f.get('type'), dict):
+        is_complex_type = True
+
+        if f['type']['type'] == 'struct':
+          complex_type = 'struct'
+        elif f['type']['type'] in ('array', 'map'):
+          complex_type = 'string'
+
+        meta.append({'name': f['name'], 'type': complex_type, 'comment': ''})
+      else:
+        meta.append({'name': f['name'], 'type': f['type'], 'comment': ''})
+    
+    return meta, is_complex_type
+
+
   def cancel(self, notebook, snippet):
   def cancel(self, notebook, snippet):
     api = self.get_api()
     api = self.get_api()
     session = _get_snippet_session(notebook, snippet)
     session = _get_snippet_session(notebook, snippet)