add list of df support (intel-analytics#3113)

dding3 · Mar 25, 2021 · 4012454 · 4012454
1 parent 94721cb
commit 4012454
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/python/dllib/src/bigdl/utils/common.py b/python/dllib/src/bigdl/utils/common.py
@@ -627,6 +627,11 @@ def _java2py(gateway, r, encoding="bytes"):
 
         if clsName in _picklable_classes:
             r = gateway.jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.dumps(r)
+        elif isinstance(r, (JavaArray, JavaList)) and len(r) != 0 \
+                and isinstance(r[0], JavaObject) \
+                and r[0].getClass().getSimpleName() in ['DataFrame', 'Dataset']:
+            spark = get_spark_sql_context(get_spark_context())
+            r = list(map(lambda x: DataFrame(x, spark), r))
         elif isinstance(r, (JavaArray, JavaList, JavaMap)):
             try:
                 r = gateway.jvm.org.apache.spark.bigdl.api.python.BigDLSerDe.dumps(