jupyter-duckdb 1.2.7__py3-none-any.whl → 1.2.101__py3-none-any.whl

duckdb_kernel/kernel.py

@@ -54,6 +54,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
@@ -577,6 +579,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,

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/parser/elements/binary/Difference.py

@@ -8,7 +8,7 @@ from ...util.RenamableColumnList import RenamableColumnList
 class Difference(RABinaryOperator):
     @staticmethod
     def symbols() -> Tuple[str, ...]:
-        return '\\',
+        return '-', '\\'
 
     def to_sql(self, tables: Dict[str, Table]) -> Tuple[str, RenamableColumnList]:
         # execute subqueries

duckdb_kernel/parser/elements/binary/Divide.py

@@ -8,7 +8,7 @@ class Divide(LogicOperator):
 
     @staticmethod
     def symbols() -> Tuple[str, ...]:
-        return '÷', '/'
+        return '/', '÷'
 
     @property
     def sql_symbol(self) -> str:

duckdb_kernel/parser/elements/binary/Division.py

@@ -27,10 +27,6 @@ class Division(RABinaryOperator):
         # inter_name_left = ', '.join(l.current_name for l, _ in inter_cols)
         inter_name_right = ', '.join(r.current_name for _, r in inter_cols)
 
-        print('-', diff_name)
-        print(inter_name)
-        print(inter_name_right)
-
         # create sql
         return f'''
             SELECT {diff_name}

duckdb_kernel/parser/elements/binary/FullOuterJoin.py

@@ -1,7 +1,9 @@
+from typing import Optional
 from typing import Tuple, Dict
 
 from duckdb_kernel.db import Table
 from ..RABinaryOperator import RABinaryOperator
+from ...util.RenamableColumn import RenamableColumn
 from ...util.RenamableColumnList import RenamableColumnList
 
 
@@ -10,6 +12,13 @@ class FullOuterJoin(RABinaryOperator):
     def symbols() -> Tuple[str, ...]:
         return chr(10199), 'fjoin', 'ojoin'
 
+    @staticmethod
+    def _coalesce(c1: RenamableColumn, c2: Optional[RenamableColumn]) -> str:
+        if c2 is not None:
+            return f'COALESCE({c1.current_name}, {c2.current_name}) AS {c1.current_name}'
+        else:
+            return c1.current_name
+
     def to_sql(self, tables: Dict[str, Table]) -> Tuple[str, RenamableColumnList]:
         # execute subqueries
         lq, lcols = self.left.to_sql(tables)
@@ -19,9 +28,7 @@ class FullOuterJoin(RABinaryOperator):
         join_cols, all_cols = lcols.intersect(rcols)
 
         replacements = {c1: c2 for c1, c2 in join_cols}
-        select_cols = [
-            f'COALESCE({c.current_name}, {replacements.get(c).current_name})' if c in replacements else c.current_name
-            for c in all_cols]
+        select_cols = [self._coalesce(c, replacements.get(c)) for c in all_cols]
         select_clause = ', '.join(select_cols)
 
         on_clause = ' AND '.join(f'{l.current_name} = {r.current_name}' for l, r in join_cols)

duckdb_kernel/tests/__init__.py

@@ -0,0 +1,76 @@
+import os
+from typing import Dict, List, Tuple
+
+from duckdb_kernel.db import Connection as DB, Table
+from duckdb_kernel.db.error import EmptyResultError
+from duckdb_kernel.parser.elements import RAElement
+from duckdb_kernel.parser.elements.binary import ConditionalSet
+
+with open('examples/tables.sql', 'r', encoding='utf-8') as file:
+    EXAMPLE_STMTS = [stmt
+                     for stmt in file.read().split(';')
+                     if stmt.strip()]
+
+
+class Connection:
+    def __enter__(self):
+        db_type = os.environ.get('DB_TYPE')
+        if db_type == 'postgres':
+            host = os.environ.get('POSTGRES_HOST', 'localhost')
+            port = int(os.environ.get('POSTGRES_PORT', 5432))
+            username = os.environ.get('POSTGRES_USER', 'postgres')
+            password = os.environ.get('POSTGRES_PASSWORD', 'postgres')
+
+            from duckdb_kernel.db.implementation.postgres import Connection as PostgreSQL
+            self.con: DB = PostgreSQL(host, port, username, password, None)
+        elif db_type == 'sqlite':
+            from duckdb_kernel.db.implementation.sqlite import Connection as SQLite
+            self.con: DB = SQLite(':memory:')
+        else:
+            from duckdb_kernel.db.implementation.duckdb import Connection as DuckDB
+            self.con: DB = DuckDB(':memory:')
+
+        for stmt in EXAMPLE_STMTS:
+            try:
+                self.con.execute(stmt)
+            except EmptyResultError:
+                pass
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.con.close()
+
+    @property
+    def tables(self) -> Dict[str, Table]:
+        return self.con.analyze()
+
+    def execute_sql_return_cols(self, query: str) -> Tuple[List, List]:
+        return self.con.execute(query)
+
+    def execute_sql(self, query: str) -> List:
+        _, rows = self.execute_sql_return_cols(query)
+        return rows
+
+    def execute_ra_return_cols(self, root: RAElement) -> Tuple[List, List]:
+        sql = root.to_sql_with_renamed_columns(self.tables)
+        cols, rows = self.execute_sql_return_cols(sql)
+
+        return cols, sorted(rows, key=lambda t: tuple(-1 if x is None else x for x in t))
+
+    def execute_ra(self, root: RAElement) -> List:
+        _, rows = self.execute_ra_return_cols(root)
+        return rows
+
+    def execute_dc_return_cols(self, root: ConditionalSet) -> Tuple[List, List]:
+        sql, cnm = root.to_sql_with_renamed_columns(self.tables)
+        cols, rows = self.execute_sql_return_cols(sql)
+
+        return [cnm.get(c, c) for c in cols], sorted(rows)
+
+    def execute_dc(self, root: ConditionalSet) -> List:
+        _, rows = self.execute_dc_return_cols(root)
+        return rows
+
+
+__all__ = ['Connection']