|
| 1 | +from jsoniq import RumbleSession |
| 2 | +from unittest import TestCase |
| 3 | +import json |
| 4 | +class TryTesting(TestCase): |
| 5 | + def test1(self): |
| 6 | + # The syntax to start a session is similar to that of Spark. |
| 7 | + # A RumbleSession is a SparkSession that additionally knows about RumbleDB. |
| 8 | + # All attributes and methods of SparkSession are also available on RumbleSession. |
| 9 | + rumble = RumbleSession.builder.appName("PyRumbleExample").getOrCreate(); |
| 10 | + # A more complex, standalone query |
| 11 | + |
| 12 | + seq = rumble.jsoniq(""" |
| 13 | + let $stores := |
| 14 | + [ |
| 15 | + { "store number" : 1, "state" : "MA" }, |
| 16 | + { "store number" : 2, "state" : "MA" }, |
| 17 | + { "store number" : 3, "state" : "CA" }, |
| 18 | + { "store number" : 4, "state" : "CA" } |
| 19 | + ] |
| 20 | + let $sales := [ |
| 21 | + { "product" : "broiler", "store number" : 1, "quantity" : 20 }, |
| 22 | + { "product" : "toaster", "store number" : 2, "quantity" : 100 }, |
| 23 | + { "product" : "toaster", "store number" : 2, "quantity" : 50 }, |
| 24 | + { "product" : "toaster", "store number" : 3, "quantity" : 50 }, |
| 25 | + { "product" : "blender", "store number" : 3, "quantity" : 100 }, |
| 26 | + { "product" : "blender", "store number" : 3, "quantity" : 150 }, |
| 27 | + { "product" : "socks", "store number" : 1, "quantity" : 500 }, |
| 28 | + { "product" : "socks", "store number" : 2, "quantity" : 10 }, |
| 29 | + { "product" : "shirt", "store number" : 3, "quantity" : 10 } |
| 30 | + ] |
| 31 | + let $join := |
| 32 | + for $store in $stores[], $sale in $sales[] |
| 33 | + where $store."store number" = $sale."store number" |
| 34 | + return { |
| 35 | + "nb" : $store."store number", |
| 36 | + "state" : $store.state, |
| 37 | + "sold" : $sale.product |
| 38 | + } |
| 39 | + return $join |
| 40 | + """); |
| 41 | + |
| 42 | + expected = ({'nb': 1, 'state': 'MA', 'sold': 'broiler'}, {'nb': 1, 'state': 'MA', 'sold': 'socks'}, {'nb': 2, 'state': 'MA', 'sold': 'toaster'}, {'nb': 2, 'state': 'MA', 'sold': 'toaster'}, {'nb': 2, 'state': 'MA', 'sold': 'socks'}, {'nb': 3, 'state': 'CA', 'sold': 'toaster'}, {'nb': 3, 'state': 'CA', 'sold': 'blender'}, {'nb': 3, 'state': 'CA', 'sold': 'blender'}, {'nb': 3, 'state': 'CA', 'sold': 'shirt'}) |
| 43 | + |
| 44 | + self.assertTrue(json.dumps(seq.json()) == json.dumps(expected)) |
| 45 | + |
| 46 | + self.assertIn("DataFrame", seq.availableOutputs()) |
| 47 | + |
| 48 | + seq.df().show() |
| 49 | + |
| 50 | + self.assertEqual(seq.df().count(), 9) |
0 commit comments