jupyter-duckdb 1.2.100__py3-none-any.whl → 1.2.102__py3-none-any.whl

duckdb_kernel/kernel.py

@@ -10,6 +10,7 @@ from typing import Optional, Dict, List, Tuple
 from ipykernel.kernelbase import Kernel
 
 from .db import Connection, DatabaseError, Table
+from .visualization.lib import *
 from .db.error import *
 from .magics import *
 from .parser import RAParser, DCParser, ParserError
@@ -54,6 +55,8 @@ class DuckDBKernel(Kernel):
             MagicCommand('all_dc').arg('value', '1').on(self._all_dc_magic),
             MagicCommand('auto_parser').disable('sql', 'ra', 'dc').code(True).on(self._auto_parser_magic),
             MagicCommand('guess_parser').arg('value', '1').on(self._guess_parser_magic),
+            MagicCommand('plotly').arg('type').arg('mapping').opt('title').result(True).on(self._plotly_magic),
+            MagicCommand('plotly_raw').opt('title').result(True).on(self._plotly_raw_magic)
         )
 
         # create placeholders for database and tests
@@ -116,6 +119,7 @@ class DuckDBKernel(Kernel):
                     self._db = Postgres(host, port, username, password, database_name)
                 except ImportError:
                     self.print('psycopg could not be found', name='stderr')
+                    return False
 
             # Otherwise the provided path is used to create an
             # in-process instance.
@@ -145,7 +149,7 @@ class DuckDBKernel(Kernel):
         else:
             return False
 
-    def _execute_stmt(self, query: str, silent: bool,
+    def _execute_stmt(self, name: str, query: str, silent: bool,
                       column_name_mapping: Dict[str, str],
                       max_rows: Optional[int]) -> Tuple[Optional[List[str]], Optional[List[List]]]:
         if self._db is None:
@@ -161,52 +165,57 @@ class DuckDBKernel(Kernel):
 
         et = time.time()
 
-        # return result if silent
-        if silent:
-            return columns, rows
+        # print result if not silent
+        if not silent:
+            # print EXPLAIN queries as raw text if using DuckDB
+            if query.strip().startswith('EXPLAIN') and self._db.plain_explain():
+                for ekey, evalue in rows:
+                    html = f'<b>{ekey}</b><br><pre>{evalue}</pre>'
+                    break
 
-        # print EXPLAIN queries as raw text if using DuckDB
-        if query.strip().startswith('EXPLAIN') and self._db.plain_explain():
-            for ekey, evalue in rows:
-                self.print_data(f'<b>{ekey}</b><br><pre>{evalue}</pre>')
+                return None, None
 
-            return None, None
+            # print every other query as a table
+            else:
+                if columns is not None:
+                    # table header
+                    mapped_columns = (column_name_mapping.get(c, c) for c in columns)
+                    table_header = ''.join(f'<th>{c}</th>' for c in mapped_columns)
+
+                    # table data
+                    if max_rows is not None and len(rows) > max_rows:
+                        table_data = f'''
+                            {rows_table(rows[:math.ceil(max_rows / 2)])}
+                            <tr>
+                                <td colspan="{len(columns)}" 
+                                    style="text-align: center"
+                                    title="{row_count(len(rows) - max_rows)} omitted">
+                                    ...
+                                </td>
+                            </tr>
+                            {rows_table(rows[-math.floor(max_rows // 2):])}
+                        '''
+                    else:
+                        table_data = rows_table(rows)
+
+                    # send to client
+                    html = (f'''
+                        <table class="duckdb-query-result-table">
+                            {table_header}
+                            {table_data}
+                        </table>
+                        
+                        {row_count(len(rows))} in {et - st:.3f}s
+                    ''')
 
-        # print every other query as a table
-        else:
-            if columns is not None:
-                # table header
-                mapped_columns = (column_name_mapping.get(c, c) for c in columns)
-                table_header = ''.join(f'<th>{c}</th>' for c in mapped_columns)
-
-                # table data
-                if max_rows is not None and len(rows) > max_rows:
-                    table_data = f'''
-                        {rows_table(rows[:math.ceil(max_rows / 2)])}
-                        <tr>
-                            <td colspan="{len(columns)}" 
-                                style="text-align: center"
-                                title="{row_count(len(rows) - max_rows)} omitted">
-                                ...
-                            </td>
-                        </tr>
-                        {rows_table(rows[-math.floor(max_rows // 2):])}
-                    '''
                 else:
-                    table_data = rows_table(rows)
-
-                # send to client
-                self.print_data(f'''
-                    <table class="duckdb-query-result">
-                        {table_header}
-                        {table_data}
-                    </table>
-                ''')
-
-                self.print_data(f'{row_count(len(rows))} in {et - st:.3f}s')
+                    html = f'statement executed without result in {et - st:.3f}s'
 
-            else:
-                self.print_data(f'statement executed without result in {et - st:.3f}s')
+            self.print_data(f'''
+                <div class="duckdb-query-result {name}">
+                    {html}
+                </div>
+            ''')
 
         return columns, rows
 
@@ -496,16 +505,23 @@ class DuckDBKernel(Kernel):
         # create and show visualization
         if analyze:
             vd = RATreeDrawer(self._db, root_node, tables)
-            svg = vd.to_svg(True)
-
-            self.print_data(svg)
+            svg = vd.to_interactive_svg()
+            data = {
+                'generated_code': {
+                    node_id: node.to_sql_with_renamed_columns(tables)
+                    for node_id, node in vd.nodes.items()
+                }
+            }
 
-        # generate sql
-        sql = root_node.to_sql_with_renamed_columns(tables)
+        else:
+            svg = ''
+            data = {
+                'generated_code': root_node.to_sql_with_renamed_columns(tables)
+            }
 
-        return {
-            'generated_code': sql
-        }
+        # return data
+        self.print_data(svg)
+        return data
 
     def _all_ra_magic(self, silent: bool, value: str):
         if value.lower() in ('1', 'on', 'true'):
@@ -577,6 +593,90 @@ class DuckDBKernel(Kernel):
             if e.depth > 0:
                 raise e
 
+    def _plotly_magic(self, silent: bool, cols: List, rows: List[Tuple], type: str, mapping: str, title: str = None):
+        # split mapping and handle asterisks
+        mapping = [m.strip() for m in mapping.split(',')]
+
+        for i in range(len(mapping)):
+            if mapping[i] == '*':
+                mapping = mapping[:i] + cols + mapping[i+1:]
+
+        # convert all column names to lower case
+        lower_cols = [c.lower() for c in cols]
+        lower_mapping = [m.lower() for m in mapping]
+
+        # map desired columns to indices
+        mapped_indices = {}
+        for ok, lk in zip(mapping, lower_mapping):
+            for i in range(len(lower_cols)):
+                if lk == lower_cols[i]:
+                    mapped_indices[ok] = i
+                    break
+            else:
+                raise ValueError(f'unknown column {ok}')
+
+        # map desired columns to value lists
+        mapped_values = {
+            m: [r[i] for r in rows]
+            for m, i in mapped_indices.items()
+        }
+        mapped_keys = iter(mapped_values.keys())
+
+        # get required chart type
+        match type.lower():
+            case 'scatter':
+                if len(lower_mapping) < 2: raise ValueError('scatter requires at least x and y values')
+                html = draw_scatter_chart(title,
+                                          mapped_values[next(mapped_keys)],
+                                          **{k: mapped_values[k] for k in mapped_keys})
+            case 'line':
+                if len(lower_mapping) < 2: raise ValueError('lines requires at least x and y values')
+                html = draw_line_chart(title,
+                                       mapped_values[next(mapped_keys)],
+                                       **{k: mapped_values[k] for k in mapped_keys})
+
+            case 'bar':
+                if len(lower_mapping) < 2: raise ValueError('bar requires at least x and y values')
+                html = draw_bar_chart(title,
+                                      mapped_values[next(mapped_keys)],
+                                      **{k: mapped_values[k] for k in mapped_keys})
+
+            case 'pie':
+                if len(lower_mapping) != 2: raise ValueError('pie requires labels and values')
+                html = draw_pie_chart(title,
+                                      mapped_values[next(mapped_keys)],
+                                      mapped_values[next(mapped_keys)])
+
+            case 'bubble':
+                if len(lower_mapping) != 4: raise ValueError('bubble requires x, y, size and color')
+                html = draw_bubble_chart(title,
+                                         mapped_values[next(mapped_keys)],
+                                         mapped_values[next(mapped_keys)],
+                                         mapped_values[next(mapped_keys)],
+                                         mapped_values[next(mapped_keys)])
+
+            case 'heatmap':
+                if len(lower_mapping) != 3: raise ValueError('heatmap requires x, y and z values')
+                html = draw_heatmap_chart(title,
+                                          mapped_values[next(mapped_keys)],
+                                          mapped_values[next(mapped_keys)],
+                                          mapped_values[next(mapped_keys)])
+
+            case _:
+                raise ValueError(f'unknown type: {type}')
+
+        # finally print the code
+        self.print_data(html, mime='text/html')
+
+    def _plotly_raw_magic(self, silent: bool, cols: List, rows: List[Tuple], title: str = None):
+        if len(cols) != 1 and len(rows) != 1:
+            raise ValueError(f'expected exactly one column and one row')
+
+        self.print_data(
+            draw_chart(title, rows[0][0]),
+            mime='text/html'
+        )
+
     # jupyter related functions
     def do_execute(self, code: str, silent: bool,
                    store_history: bool = True, user_expressions: dict = None, allow_stdin: bool = False,
@@ -603,10 +703,14 @@ class DuckDBKernel(Kernel):
                 execution_args['column_name_mapping'] = {}
 
             # execute statement if needed
-            if clean_code.strip():
-                cols, rows = self._execute_stmt(clean_code, silent, **execution_args)
-            else:
-                cols, rows = None, None
+            cols, rows = None, None
+
+            if not isinstance(clean_code, dict):
+                clean_code = {'default': clean_code}
+
+            for name, code in clean_code.items():
+                if code.strip():
+                    cols, rows = self._execute_stmt(name, code, silent, **execution_args)
 
             # execute magic command here if it does depend on query results
             for callback in post_query_callbacks:

duckdb_kernel/magics/MagicCommandHandler.py

@@ -1,5 +1,5 @@
 import re
-from typing import Dict, Tuple, List, Optional
+from typing import Dict, Tuple, List
 
 from . import MagicCommand, MagicCommandException, MagicCommandCallback
 from .StringWrapper import StringWrapper
@@ -62,6 +62,11 @@ class MagicCommandHandler:
             args = [group if group is not None else default
                     for group, (_, default, _) in zip(match.groups(), magic.args)]
 
+            args = [arg[1:-1]
+                    if arg is not None and (arg[0] == '"' and arg[-1] == '"' or arg[0] == "'" and arg[-1] == "'")
+                    else arg
+                    for arg in args]
+
             if any(arg is None for arg in args):
                 raise MagicCommandException(f'could not parse parameters for command "{command}"')
 
@@ -87,6 +92,9 @@ class MagicCommandHandler:
                 value = match.group(i + 2)
                 i += 3
 
+                if value is not None and (value[0] == '"' and value[-1] == '"' or value[0] == "'" and value[-1] == "'"):
+                    value = value[1:-1]
+
                 if name is not None:
                     optionals[name.lower()] = value
 

duckdb_kernel/visualization/Plotly.py

@@ -0,0 +1,144 @@
+import json
+from decimal import Decimal
+from typing import Dict, List, Optional
+from uuid import uuid4
+
+from .lib import init_plotly
+
+
+def __div_id() -> str:
+    return f'div-{str(uuid4())}'
+
+
+def __layout(title: Optional[str]):
+    layout = {
+        'dragmode': False,
+        'xaxis': {
+            'rangeselector': {
+                'visible': False
+            }
+        }
+    }
+
+    if title is not None:
+        layout['title'] = {
+            'text': title,
+            'font': {
+                'family': 'sans-serif',
+                'size': 32,
+                'color': 'rgb(0, 0, 0)'
+            },
+            'xanchor': 'center'
+        }
+
+    return layout
+
+
+def __config():
+    return {
+        'displayModeBar': False,
+        'scrollZoom': False
+    }
+
+
+def __fix_decimal(x: List):
+    return [float(x) if isinstance(x, Decimal) else x
+            for x in x]
+
+
+def draw_chart(title: Optional[str], traces: List[Dict] | Dict) -> str:
+    init = init_plotly()
+    div_id = __div_id()
+    layout = __layout(title)
+    config = __config()
+
+    if not isinstance(traces, str):
+        traces = json.dumps(traces)
+
+    return f'''
+        <script type="text/javascript">
+            {init}
+        </script>
+
+        <div id="{div_id}"></div>
+        <script type="text/javascript">
+            Plotly.newPlot('{div_id}', {traces}, {json.dumps(layout)}, {json.dumps(config)});
+        </script>
+    '''
+
+
+def draw_scatter_chart(title: Optional[str], x, **ys) -> str:
+    return draw_chart(title, [
+        {
+            'x': __fix_decimal(x),
+            'y': __fix_decimal(y),
+            'mode': 'markers',
+            'type': 'scatter',
+            'name': name
+        }
+        for name, y in ys.items()
+    ])
+
+
+def draw_line_chart(title: Optional[str], x, **ys) -> str:
+    return draw_chart(title, [
+        {
+            'x': __fix_decimal(x),
+            'y': __fix_decimal(y),
+            'mode': 'lines+markers',
+            'name': name
+        }
+        for name, y in ys.items()
+    ])
+
+
+def draw_bar_chart(title: Optional[str], x, **ys) -> str:
+    return draw_chart(title, [
+        {
+            'x': __fix_decimal(x),
+            'y': __fix_decimal(y),
+            'type': 'bar',
+            'name': name
+        }
+        for name, y in ys.items()
+    ])
+
+
+def draw_pie_chart(title: Optional[str], x, y) -> str:
+    return draw_chart(title, [{
+        'values': __fix_decimal(y),
+        'labels': __fix_decimal(x),
+        'type': 'pie'
+    }])
+
+
+def draw_bubble_chart(title: Optional[str], x, y, s, c) -> str:
+    return draw_chart(title, [{
+        'x': __fix_decimal(x),
+        'y': __fix_decimal(y),
+        'mode': 'markers',
+        'marker': {
+            'size': __fix_decimal(s),
+            'color': __fix_decimal(c)
+        }
+    }])
+
+
+def draw_heatmap_chart(title: Optional[str], x, y, z) -> str:
+    return draw_chart(title, [{
+        'x': __fix_decimal(x[0]),
+        'y': __fix_decimal(y[0]),
+        'z': [__fix_decimal(v) for v in z[0]],
+        'type': 'heatmap'
+    }])
+
+
+__all__ = [
+    'draw_chart',
+    'draw_scatter_chart',
+    'draw_line_chart',
+    'draw_bar_chart',
+    'draw_pie_chart',
+    'draw_bubble_chart',
+    'draw_heatmap_chart',
+]

duckdb_kernel/visualization/RATreeDrawer.py

@@ -1,12 +1,14 @@
-from typing import Dict
+from typing import Dict, Optional
 
 from graphviz import Digraph
+from uuid import uuid4
 
 from duckdb_kernel.db import Table
 from duckdb_kernel.parser.elements import RAElement
 from duckdb_kernel.util.formatting import row_count
 from .Drawer import Drawer
 from ..db import Connection
+from .lib import *
 
 
 class RATreeDrawer(Drawer):
@@ -15,6 +17,9 @@ class RATreeDrawer(Drawer):
         self.root_node: RAElement = root_node
         self.tables: Dict[str, Table] = tables
 
+        self.nodes: Dict[str, RAElement] = {}
+        self.root_node_id: Optional[str] = None
+
     def to_graph(self) -> Digraph:
         # create graph
         ps = Digraph('Schema',
@@ -31,7 +36,11 @@ class RATreeDrawer(Drawer):
 
     def __add_node(self, ps: Digraph, node: RAElement) -> str:
         # use id of node object as identifier
-        node_id = f'node_{id(node)}'
+        node_id = f'node_{str(uuid4()).replace("-", "_")}'
+
+        self.nodes[node_id] = node
+        if node == self.root_node:
+            self.root_node_id = node_id
 
         # generate child nodes
         child_ids = [self.__add_node(ps, child) for child in node.children]
@@ -69,3 +78,26 @@ class RATreeDrawer(Drawer):
 
         # return node identifier to generate edges
         return node_id
+
+    def to_interactive_svg(self) -> str:
+        div_id = f'div-{str(uuid4())}'
+
+        css = init_css()
+        ra = init_ra()
+        svg = self.to_svg(True)
+
+        return f'''
+            <style type="text/css">
+                {css}
+            </style>
+            
+            <div id="{div_id}">
+                {svg}
+            </div>
+            
+            <script type="text/javascript">
+                {ra}
+                
+                animate_ra('{div_id}', '{self.root_node_id}')
+            </script>
+        '''

duckdb_kernel/visualization/__init__.py

@@ -1,2 +1,3 @@
+from .Plotly import *
 from .RATreeDrawer import RATreeDrawer
 from .SchemaDrawer import SchemaDrawer

duckdb_kernel/visualization/lib/__init__.py

@@ -0,0 +1,53 @@
+import os
+
+__CSS_INITIALIZED = False
+__RA_INITIALIZED = False
+__PLOTLY_INITIALIZED = False
+
+__location = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))
+
+
+def init_css() -> str:
+    global __CSS_INITIALIZED
+
+    if not __CSS_INITIALIZED:
+        with open(os.path.join(__location, 'ra.css')) as ra_file:
+            css = ra_file.read()
+    else:
+        css = ''
+
+    __CSS_INITIALIZED = True
+    return css
+
+
+def init_ra() -> str:
+    global __RA_INITIALIZED
+
+    if not __RA_INITIALIZED:
+        with open(os.path.join(__location, 'ra.js')) as ra_file:
+            ra = ra_file.read()
+    else:
+        ra = ''
+
+    __RA_INITIALIZED = True
+    return ra
+
+
+def init_plotly() -> str:
+    global __PLOTLY_INITIALIZED
+
+    if not __PLOTLY_INITIALIZED:
+        with open(os.path.join(__location, 'plotly-3.0.1.min.js')) as plotly_file:
+            plotly = plotly_file.read()
+    else:
+        plotly = ''
+
+    __PLOTLY_INITIALIZED = True
+    return plotly
+
+
+__all__ = [
+    'init_css',
+    'init_ra',
+    'init_plotly',
+]