Fine-tune user experience.

Author: Ghislain Fourny

src/jsoniq/jars/rumbledb-1.24.0.jar

@@ -4,16 +4,26 @@
 import json
 
 class SequenceOfItems:
-    def __init__(self, sequence, sparksession):
+    def __init__(self, sequence, rumblesession):
         self._jsequence = sequence
-        self._sparkcontext = sparksession.sparkContext
-        self._sparksession = sparksession
+        self._rumblesession = rumblesession
+        self._sparksession = rumblesession._sparksession
+        self._sparkcontext = self._sparksession.sparkContext
+
+    def items(self):
+        return self.getAsList()
+
+    def take(self, n):
+        return tuple(self.getFirstItemsAsList(n))
+    
+    def first(self):
+        return tuple(self.getFirstItemsAsList(self._rumblesession.getRumbleConf().getResultSizeCap()))
 
     def json(self):
-        return tuple([json.loads(l.serializeAsJSON()) for l in self._jsequence.items()])
+        return tuple([json.loads(l.serializeAsJSON()) for l in self._jsequence.getAsList()])
 
     def rdd(self):
-        rdd = self._jsequence.getAsPickledStringRDD();
+        rdd = self._jsequence.getAsPickledStringRDD()
         rdd = RDD(rdd, self._sparkcontext)
         return rdd.map(lambda l: json.loads(l))
 
@@ -22,7 +32,10 @@ def df(self):
 
     def pdf(self):
         return self.df().toPandas()
-
+    
+    def count(self):
+        return self._jsequence.count()
+    
     def nextJSON(self):
         return self._jsequence.next().serializeAsJSON()
 

src/jsoniq/sequence.py

@@ -22,6 +22,9 @@ def __init__(self, spark_session: SparkSession):
         self._sparksession = spark_session
         self._jrumblesession = spark_session._jvm.org.rumbledb.api.Rumble(spark_session._jsparkSession)
 
+    def getRumbleConf(self):
+        return self._jrumblesession.getConfiguration()
+
     class Builder:
         def __init__(self):
 
@@ -166,7 +169,7 @@ def bindDataFrameAsVariable(self, name: str, df):
 
     def jsoniq(self, str):
         sequence = self._jrumblesession.runQuery(str);
-        return SequenceOfItems(sequence, self._sparksession);
+        return SequenceOfItems(sequence, self);
 
     def __getattr__(self, item):
         return getattr(self._sparksession, item)

src/jsoniq/session.py

@@ -5,4 +5,5 @@
 
 def load_ipython_extension(ipython):
     rumble = RumbleSession.builder.getOrCreate();
+    rumble.getRumbleConf().setResultSizeCap(10);
     ipython.register_magics(JSONiqMagic)

src/jsoniqmagic/__init__.py

@@ -35,10 +35,11 @@ def run(self, line, cell=None, timed=False):
         if ("DataFrame" in response.availableOutputs()):
             print(response.pdf())
         elif ("Local" in response.availableOutputs()):
-            count = response.getAsRDD().count()
-            if count > 200:
-                print("The query output %s items, which is too many to display. Displaying the first 200 items:" % count)
-                for e in response.first():
+            capplusone = response.take(rumble.getRumbleConf().getResultSizeCap() + 1)
+            if len(capplusone) > rumble.getRumbleConf().getResultSizeCap():
+                count = response.count()
+                print("The query output %s items, which is too many to display. Displaying the first %s items:" % (count, rumble.getRumbleConf().getResultSizeCap()))
+                for e in capplusone[:-1]:
                     print(json.dumps(json.loads(e.serializeAsJSON()), indent=2))
             else:
                 for e in response.json():