support json type

Author: Yibo-Chen13

README.rst

@@ -143,4 +143,63 @@ Big fan of Pandas? We too! You can mix SQL and Pandas API together. Also you can
                print(cur.pl()) # Polars DataFrame
 
                cur.execute('SELECT * FROM table(test)')
-               print(cur.arrow()) # Arrow Table
+               print(cur.arrow()) # Arrow Table
+
+Json Type
+------------
+You can control named tuple and nested types deserialization by setting ``namedtuple_as_json``.To interpret these columns (i.e. named tuple, nested) as a json set ``namedtuple_as_json`` to ``True``
+
+.. code-block:: python
+
+    from proton_driver import client
+    import time
+
+    if __name__ == "__main__":
+        c = client.Client(host='127.0.0.1', port=8463)
+
+        c.execute("DROP STREAM IF EXISTS test")
+        c.execute("""CREATE STREAM test (
+                        j json,
+                        t tuple(code int,data tuple(id int,price float)),
+                        n nested(name string, locations array(string))
+                    )""", settings=dict(flatten_nested=0))
+
+        # add some data
+        data = [(
+                {
+                    "id": 341436,
+                    "info": {
+                        "name": "timeplus",
+                        "timezone": "Asia/Shanghai",
+                    },
+                    "arr": [43, 65, 87, 456, 356],
+                },
+                (200, (341436, 699.98)),
+                [("proton", ["Ningbo", "Hangzhou", "Shanghai", "Beijing"])],
+        )]
+        c.execute("INSERT INTO test (j, t, n) VALUES", data)
+
+        # wait for 3 sec to make sure data available in historical store
+        time.sleep(3) 
+
+        # proton return json type as a named tuple by default.
+        rv = c.execute(
+            "SELECT j, t, n FROM table(test)",
+        )
+
+        rv_as_json = c.execute(
+            "SELECT j, t, n FROM table(test)",
+            # Set namedtuple_as_json to True,to convert named tuple and nested types to json
+            settings=dict(namedtuple_as_json=True) 
+        )
+
+        for record_as_json, record_as_tuple in zip(rv_as_json, rv):
+            print(f"colnmu j:\n - json:  {record_as_json[0]}\n - tuple: {record_as_tuple[0]}")
+            print(f"colnmu t:\n - json:  {record_as_json[1]}\n - tuple: {record_as_tuple[1]}")
+            print(f"colnmu n:\n - json:  {record_as_json[2]}\n - tuple: {record_as_tuple[2]}")
+
+        # For json type, you can cast it to json in SQL statement without namedtuple_as_json setting.
+        rv_json = c.execute(
+            "SELECT j::json FROM table(test)",
+        )
+        print(rv_json)

proton_driver/block.py

@@ -1,7 +1,7 @@
 from .reader import read_varint, read_binary_uint8, read_binary_int32
 from .varint import write_varint
 from .writer import write_binary_uint8, write_binary_int32
-from .columns import nestedcolumn
+from .columns.util import get_inner_columns_with_types
 
 
 class BlockInfo(object):
@@ -172,7 +172,7 @@ def _pure_mutate_dicts_to_rows(
         for name, type_ in columns_with_types:
             cwt = None
             if type_.startswith('nested'):
-                cwt = nestedcolumn.get_columns_with_types(type_)
+                cwt = get_inner_columns_with_types('nested', type_)
             columns_with_cwt.append((name, cwt))
 
         for i, row in enumerate(data):

proton_driver/client.py

@@ -49,6 +49,11 @@ class Client(object):
         * ``quota_key`` -- A string to differentiate quotas when the user have
                            keyed quotas configured on server.
                            New in version *0.2.3*.
+        * ``namedtuple_as_json`` -- Controls named tuple and nested types
+                           deserialization. To interpret these column as
+                           Python tuple set ``namedtuple_as_json``
+                           to ``False``. Default: False.
+                           New in version *0.2.12*.
     """
 
     available_client_settings = (
@@ -58,7 +63,8 @@ class Client(object):
         'use_numpy',
         'opentelemetry_traceparent',
         'opentelemetry_tracestate',
-        'quota_key'
+        'quota_key',
+        'namedtuple_as_json'
     )
 
     def __init__(self, *args, **kwargs):
@@ -85,6 +91,9 @@ def __init__(self, *args, **kwargs):
             ),
             'quota_key': self.settings.pop(
                 'quota_key', ''
+            ),
+            'namedtuple_as_json': self.settings.pop(
+                'namedtuple_as_json', False
             )
         }
 

proton_driver/columns/arraycolumn.py

@@ -28,25 +28,31 @@ class ArrayColumn(Column):
     py_types = (list, tuple)
 
     def __init__(self, nested_column, **kwargs):
-        self.size_column = UInt64Column()
+        self.init_kwargs = kwargs
+        self.size_column = UInt64Column(**kwargs)
         self.nested_column = nested_column
         self._write_depth_0_size = True
         super(ArrayColumn, self).__init__(**kwargs)
+        self.null_value = []
 
     def write_data(self, data, buf):
         # Column of Array(T) is stored in "compact" format and passed to server
         # wrapped into another Array without size of wrapper array.
-        self.nested_column = ArrayColumn(self.nested_column)
+        self.nested_column = ArrayColumn(
+            self.nested_column, **self.init_kwargs
+        )
         self.nested_column.nullable = self.nullable
         self.nullable = False
         self._write_depth_0_size = False
         self._write(data, buf)
 
-    def read_data(self, rows, buf):
-        self.nested_column = ArrayColumn(self.nested_column)
+    def read_data(self, n_rows, buf):
+        self.nested_column = ArrayColumn(
+            self.nested_column, **self.init_kwargs
+        )
         self.nested_column.nullable = self.nullable
         self.nullable = False
-        return self._read(rows, buf)[0]
+        return self._read(n_rows, buf)[0]
 
     def _write_sizes(self, value, buf):
         nulls_map = []
@@ -99,14 +105,19 @@ def _write_nulls_data(self, value, buf):
                 self.nested_column._write_nulls_map(value, buf)
 
     def _write(self, value, buf):
+        value = self.prepare_items(value)
         self._write_sizes(value, buf)
         self._write_nulls_data(value, buf)
         self._write_data(value, buf)
 
     def read_state_prefix(self, buf):
-        return self.nested_column.read_state_prefix(buf)
+        super(ArrayColumn, self).read_state_prefix(buf)
+
+        self.nested_column.read_state_prefix(buf)
 
     def write_state_prefix(self, buf):
+        super(ArrayColumn, self).write_state_prefix(buf)
+
         self.nested_column.write_state_prefix(buf)
 
     def _read(self, size, buf):
@@ -145,6 +156,6 @@ def _read(self, size, buf):
         return tuple(data)
 
 
-def create_array_column(spec, column_by_spec_getter):
+def create_array_column(spec, column_by_spec_getter, column_options):
     inner = spec[6:-1]
-    return ArrayColumn(column_by_spec_getter(inner))
+    return ArrayColumn(column_by_spec_getter(inner), **column_options)

proton_driver/columns/jsoncolumn.py

@@ -0,0 +1,39 @@
+from .base import Column
+from .stringcolumn import String
+from ..reader import read_binary_uint8, read_binary_str
+from ..util.compat import json
+from ..writer import write_binary_uint8
+
+
+class JsonColumn(Column):
+    py_types = (dict, )
+
+    # No NULL value actually
+    null_value = {}
+
+    def __init__(self, column_by_spec_getter, **kwargs):
+        self.column_by_spec_getter = column_by_spec_getter
+        self.string_column = String(**kwargs)
+        super(JsonColumn, self).__init__(**kwargs)
+
+    def write_state_prefix(self, buf):
+        # Read in binary format.
+        # Write in text format.
+        write_binary_uint8(1, buf)
+
+    def read_items(self, n_items, buf):
+        read_binary_uint8(buf)
+        spec = read_binary_str(buf)
+        col = self.column_by_spec_getter(
+            spec, dict(namedtuple_as_json=True)
+        )
+        col.read_state_prefix(buf)
+        return col.read_data(n_items, buf)
+
+    def write_items(self, items, buf):
+        items = [x if isinstance(x, str) else json.dumps(x) for x in items]
+        self.string_column.write_items(items, buf)
+
+
+def create_json_column(spec, column_by_spec_getter, column_options):
+    return JsonColumn(column_by_spec_getter, **column_options)

proton_driver/columns/nestedcolumn.py

@@ -1,73 +1,10 @@
 
 from .arraycolumn import create_array_column
+from .util import get_inner_spec
 
 
-def create_nested_column(spec, column_by_spec_getter):
+def create_nested_column(spec, column_by_spec_getter, column_options):
     return create_array_column(
-        'array(tuple({}))'.format(','.join(get_nested_columns(spec))),
-        column_by_spec_getter=column_by_spec_getter
+        'array(tuple({}))'.format(get_inner_spec('nested', spec)),
+        column_by_spec_getter, column_options
     )
-
-
-def get_nested_columns(spec):
-    brackets = 0
-    column_begin = 0
-
-    inner_spec = get_inner_spec(spec)
-    nested_columns = []
-    for i, x in enumerate(inner_spec + ','):
-        if x == ',':
-            if brackets == 0:
-                nested_columns.append(inner_spec[column_begin:i])
-                column_begin = i + 1
-        elif x == '(':
-            brackets += 1
-        elif x == ')':
-            brackets -= 1
-        elif x == ' ':
-            if brackets == 0:
-                column_begin = i + 1
-    return nested_columns
-
-
-def get_columns_with_types(spec):
-    brackets = 0
-    prev_comma = 0
-    prev_space = 0
-
-    inner_spec = get_inner_spec(spec)
-    columns_with_types = []
-
-    for i, x in enumerate(inner_spec + ','):
-        if x == ',':
-            if brackets == 0:
-                columns_with_types.append((
-                    inner_spec[prev_comma:prev_space].strip(),
-                    inner_spec[prev_space:i]
-                ))
-                prev_comma = i + 1
-        elif x == '(':
-            brackets += 1
-        elif x == ')':
-            brackets -= 1
-        elif x == ' ':
-            if brackets == 0:
-                prev_space = i + 1
-    return columns_with_types
-
-
-def get_inner_spec(spec):
-    brackets = 0
-    offset = len('nested')
-    i = offset
-    for i, ch in enumerate(spec[offset:], offset):
-        if ch == '(':
-            brackets += 1
-
-        elif ch == ')':
-            brackets -= 1
-
-        if brackets == 0:
-            break
-
-    return spec[offset + 1:i]

proton_driver/columns/service.py

@@ -32,6 +32,7 @@
     IntervalSecondColumn
 )
 from .ipcolumn import IPv4Column, IPv6Column
+from .jsoncolumn import create_json_column
 
 
 column_by_type = {c.ch_type: c for c in [
@@ -64,7 +65,11 @@ def get_column_by_spec(spec, column_options, use_numpy=None):
             logger.warning('NumPy support is not implemented for %s. '
                            'Using generic column', spec)
 
-    def create_column_with_options(x):
+    def create_column_with_options(x, settings=None):
+        if settings:
+            client_settings = column_options['context'].client_settings
+            client_settings.update(settings)
+            column_options['context'].client_settings = client_settings
         return get_column_by_spec(x, column_options, use_numpy=use_numpy)
 
     if spec == 'string' or spec.startswith('fixed_string'):
@@ -80,13 +85,23 @@ def create_column_with_options(x):
         return create_decimal_column(spec, column_options)
 
     elif spec.startswith('array'):
-        return create_array_column(spec, create_column_with_options)
+        return create_array_column(
+            spec, create_column_with_options, column_options
+        )
 
     elif spec.startswith('tuple'):
-        return create_tuple_column(spec, create_column_with_options)
+        return create_tuple_column(
+            spec, create_column_with_options, column_options
+        )
+    elif spec.startswith('json'):
+        return create_json_column(
+            spec, create_column_with_options, column_options
+        )
 
     elif spec.startswith('nested'):
-        return create_nested_column(spec, create_column_with_options)
+        return create_nested_column(
+            spec, create_column_with_options, column_options
+        )
 
     elif spec.startswith('nullable'):
         return create_nullable_column(spec, create_column_with_options)