Wrap df.

Author: Ghislain Fourny

src/jsoniq/sequence.py

@@ -1,11 +1,13 @@
 from pyspark import RDD
 from pyspark.sql import SparkSession
+from pyspark.sql import DataFrame
 import json
 
 class SequenceOfItems:
-    def __init__(self, sequence, sparkcontext):
+    def __init__(self, sequence, sparksession):
         self._jsequence = sequence
-        self._sparkcontext = sparkcontext
+        self._sparkcontext = sparksession.sparkContext
+        self._sparksession = sparksession
 
     def json(self):
         return [json.loads(l.serializeAsJSON()) for l in self._jsequence.items()]
@@ -15,6 +17,9 @@ def rdd(self):
         rdd = RDD(rdd, self._sparkcontext)
         return rdd.map(lambda l: json.loads(l))
 
+    def df(self):
+        return DataFrame(self._jsequence.getAsDataFrame(), self._sparksession)
+
     def nextJSON(self):
         return self._jsequence.next().serializeAsJSON()
 

src/jsoniq/session.py

@@ -140,7 +140,7 @@ def bindDataFrameAsVariable(self, name: str, df):
 
     def jsoniq(self, str):
         sequence = self._jrumblesession.runQuery(str);
-        return SequenceOfItems(sequence, self._sparksession.sparkContext);
+        return SequenceOfItems(sequence, self._sparksession);
 
     def __getattr__(self, item):
         return getattr(self._sparksession, item)