update datapipeline for parallel map-style processing

Author: miraculixx

.gitignore

@@ -248,4 +248,4 @@ mlops-export/**
 mlruns/**
 .vscode/**
 .dccache
-/scripts/mongoinit.js
+omegaml/tests/**/*.db

omegaml/backends/sqlalchemy.py

@@ -1,16 +1,19 @@
+from logging import warning
+
 import logging
 import os
+import pandas as pd
+import sqlalchemy
 import string
+import threading
 import warnings
 from getpass import getuser
 from hashlib import sha256
-from logging import warning
+from sqlalchemy.exc import StatementError
 from urllib.parse import quote_plus
 
-import pandas as pd
-import sqlalchemy
 from omegaml.backends.basedata import BaseDataBackend
-from omegaml.util import ProcessLocal, KeepMissing
+from omegaml.util import ProcessLocal, KeepMissing, tqdm_if_interactive
 
 try:
     import snowflake
@@ -172,8 +175,14 @@ def drop(self, name, secrets=None, **kwargs):
             self.__CNX_CACHE.clear()
         return super().drop(name, **kwargs)
 
+    def sign(self, values):
+        # sign a set of values
+        # -- this is used to ensure the values are not tampered with
+        # -- when passed to a SQL query
+        return sha256((str(threading.get_ident()) + str(values)).encode('utf-8')).hexdigest()
+
     def get(self, name, sql=None, chunksize=None, raw=False, sqlvars=None,
-            secrets=None, index=True, keep=None, lazy=False, table=None, *args, **kwargs):
+            secrets=None, index=True, keep=None, lazy=False, table=None, trusted=False, *args, **kwargs):
         """ retrieve a stored connection or query data from connection
 
         Args:
@@ -183,6 +192,8 @@ def get(self, name, sql=None, chunksize=None, raw=False, sqlvars=None,
               default by setting om.defaults.SQLALCHEMY_ALWAYS_CACHE = False)
             table (str): the name of the table, will be prefixed with the
                store's bucket name unless the table is specified as ':name'
+            trusted (bool|str): if passed must be the value for store.sign(sqlvars or kwargs),
+               otherwise a warning is issued for any remaining variables in the sql statement
 
         Returns:
             connection
@@ -230,7 +241,7 @@ def get(self, name, sql=None, chunksize=None, raw=False, sqlvars=None,
         sqlvars = sqlvars or {}
         table = self._default_table(table or meta.kind_meta.get('table') or name)
         if not raw and not valid_sql(sql):
-            sql = f'select * from {table}'
+            sql = f'select * from :sqltable'
         chunksize = chunksize or meta.kind_meta.get('chunksize')
         _default_keep = getattr(self.data_store.defaults,
                                 'SQLALCHEMY_ALWAYS_CACHE',
@@ -242,8 +253,9 @@ def get(self, name, sql=None, chunksize=None, raw=False, sqlvars=None,
         else:
             raise ValueError('no connection string')
         if not raw and valid_sql(sql):
+            sql = sql.replace(':sqltable', table)
             index_cols = _meta_to_indexcols(meta) if index else kwargs.get('index_col')
-            stmt = self._sanitize_statement(sql, sqlvars)
+            stmt = self._sanitize_statement(sql, sqlvars, trusted=trusted)
             kwargs = meta.kind_meta.get('kwargs') or {}
             kwargs.update(kwargs)
             if not lazy:
@@ -428,15 +440,9 @@ def copy_from_sql(self, sql, connstr, name, chunksize=10000,
         connection = self._get_connection(name, connstr, secrets=secrets)
         chunksize = chunksize or 10000  # avoid None
         pditer = pd.read_sql(sql, connection, chunksize=chunksize, **kwargs)
-        try:
-            import tqdm
-        except:
+        with tqdm_if_interactive().tqdm(unit='rows') as pbar:
             meta = self._chunked_insert(pditer, name, append=append,
-                                        transform=transform)
-        else:
-            with tqdm.tqdm(unit='rows') as pbar:
-                meta = self._chunked_insert(pditer, name, append=append,
-                                            transform=transform, pbar=pbar)
+                                        transform=transform, pbar=pbar)
         connection.close()
         return meta
 
@@ -448,7 +454,7 @@ def _chunked_to_sql(self, df, table, connection, if_exists='append', chunksize=N
         def chunker(seq, size):
             return (seq.iloc[pos:pos + size] for pos in range(0, len(seq), size))
 
-        def to_sql(df, table, connection, pbar=False):
+        def to_sql(df, table, connection, pbar=None):
             for i, cdf in enumerate(chunker(df, chunksize)):
                 exists_action = if_exists if i == 0 else "append"
                 cdf.to_sql(table, con=connection, if_exists=exists_action, **kwargs)
@@ -457,15 +463,8 @@ def to_sql(df, table, connection, pbar=False):
                 else:
                     print("writing chunk {}".format(i))
 
-        try:
-            import tqdm
-            if pbar is False:
-                pbar = None
-        except Exception as e:
+        with tqdm_if_interactive().tqdm(total=len(df), unit='rows') as pbar:
             to_sql(df, table, connection, pbar=pbar)
-        else:
-            with tqdm.tqdm(total=len(df), unit='rows') as pbar:
-                to_sql(df, table, connection, pbar=pbar)
 
     def _chunked_insert(self, pditer, name, append=True, transform=None, pbar=None):
         # insert into om dataset
@@ -525,7 +524,7 @@ def _default_table(self, name):
             name = name[1:]
         return name
 
-    def _sanitize_statement(self, sql, sqlvars):
+    def _sanitize_statement(self, sql, sqlvars, trusted=False):
         # sanitize sql:string statement in two steps
         # -- step 1: replace all {} variables by :notation
         # -- step 2: replace all remaining {} variables from sqlvars
@@ -538,7 +537,7 @@ def _sanitize_statement(self, sql, sqlvars):
         # replace all {...} variables with bound parameters
         #    sql = "select * from foo where user={username}"
         #       => "select * from foo where user=:username"
-        placeholders = string.Formatter().parse(sql)
+        placeholders = list(string.Formatter().parse(sql))
         vars = [spec[1] for spec in placeholders if spec[1]]
         safe_replacements = {var: f':{var}' for var in vars}
         sql = sql.format(**safe_replacements)
@@ -562,15 +561,15 @@ def _sanitize_statement(self, sql, sqlvars):
             #    =>    "select a, b from foo where user=:username
             placeholders = list(string.Formatter().parse(sql))
             vars = [spec[1] for spec in placeholders if spec[1]]
-            if vars:
+            if vars and trusted != self.sign(sqlvars):
                 warnings.warn(f'Statement >{sql}< contains unsafe variables {vars}. Use :notation or sanitize input.')
             sql = sql.format(**{**sqlvars, **safe_replacements})
         except KeyError as e:
             raise KeyError('{e}, specify sqlvars= to build query >{sql}<'.format(**locals()))
         # prepare sql statement with bound variables
         try:
             stmt = sqlalchemy.sql.text(sql)
-        except sqlalchemy.exc.StatementError as exc:
+        except StatementError as exc:
             raise
         return stmt
 
@@ -657,7 +656,7 @@ def load_sql(om=None, kind=SQLAlchemyBackend.KIND):
     from unittest.mock import MagicMock
     from IPython import get_ipython
     import omegaml as om
-    from sql.connection import Connection
+    from sql.connection import Connection # noqa
 
     class ConnectionShim:
         # this is required to trick sql magic into accepting existing connection objects

omegaml/datapipeline.py

@@ -1,37 +1,151 @@
-from sklearn.utils import Parallel, delayed
+from joblib import Parallel, delayed
+from sqlalchemy.exc import OperationalError
+
+from omegaml.util import ProcessLocal
 
 
 class ParallelStep(object):
     def __init__(self, steps=None, agg=None, n_jobs=-1):
         self.steps = steps
-        self.agg = agg or np.mean
+        self.agg = agg or self._default_agg
         self.n_jobs = n_jobs
 
-    def aggregate(self, values):
-        return self.agg(values)
+    def aggregate(self, values, **kwargs):
+        return self.agg(values, **kwargs)
+
+    def _default_agg(self, values, **kwargs):
+        return list(values)
 
-    def __call__(self, value):
-        values = Parallel(n_jobs=self.n_jobs)(delayed(pfn)(value) for pfn in self.steps)
-        return self.aggregate(values)
+    def __call__(self, value, **kwargs):
+        values = Parallel(n_jobs=self.n_jobs)(delayed(pfn)(value, **kwargs) for pfn in self.steps)
+        return self.aggregate(values, **kwargs)
 
 
 class DataPipeline(object):
-    def __init__(self, steps):
+    """ A data pipeline that processes data through a series of steps
+
+    The pipeline is a sequence of steps, each of which is a callable that
+    processes the data.
+
+    .. versionchanged:: 0.17
+        The pipeline now automatically gets a context object that is passed
+        along to each step as a keyword argument.
+    """
+    def __init__(self, steps, *args, context=None, n_jobs=None, **kwargs):
         self.steps = steps
-        self.args = None
-        self.kwargs = None
+        self.args = args
+        self.kwargs = kwargs
+        self.context = context or ProcessLocal()
+        self.n_jobs = n_jobs or -1
 
     def set_params(self, *args, **kwargs):
         self.args = args
         self.kwargs = kwargs
         return self
 
-    def process(self, value=None):
+    def process(self, value=None, **kwargs):
         for stepfn in self.steps:
-            value = stepfn(value)
+            value = stepfn(value, context=self.context, **kwargs)
         return value
 
-    def __call__(self):
-        return self.process()
+    def __call__(self, value=None, **kwargs):
+        return self.process(value, **kwargs)
+
+    def map(self, values, **kwargs):
+        parallel_steps = self.steps[:-1] if len(self.steps) > 1 else self.steps
+        finalizer = self.steps[-1] if len(self.steps) > 1 else lambda values, **kwargs: values
+        pfn = DataPipeline(*self.args, steps=parallel_steps, context=self.context, **self.kwargs)
+        parallel = Parallel(n_jobs=self.n_jobs)
+        results = parallel(delayed(pfn)(**kwargs) for kwargs in values)
+        return finalizer(results, context=self.context, **kwargs)
+
+
+class Model:
+    dburl = 'sqlite:///:memory:'
+    name = ''
+    sql = 'select * from :sqltable'
+    table = None
+    chunksize = None
+    delete_sql = 'delete from :sqltable'
+    keys = None
+    join_sql = '''
+    with join_:sqltable as (
+       select {{_join_vars_a}}, {{_join_vars_b}}
+       from   :sqltable as a
+       join  {{_join_sqltable}} as b
+       on    1 = 1 {{_join_cond}}   
+    )
+
+    select *
+    from join_:sqltable
+    '''
+
+    def __init__(self, sql=None, om=None):
+        self.sql = sql or self.sql
+        self._om = om
+        self.setup()
+
+    @property
+    def om(self):
+        import omegaml as _baseom
+        self._om = self._om or _baseom
+        return self._om
+
+    @property
+    def store(self):
+        return self.om.datasets
+
+    def setup(self):
+        assert self.name, 'name must be set'
+        return self.store.put(self.dburl, self.name, sql=self.sql, table=self.table)
+
+    def query(self, *args, sql=None, chunksize=None, trusted=False, **vars):
+        chunksize = chunksize or self.chunksize
+        return self.store.get(self.name, chunksize=chunksize, sqlvars=vars, sql=sql, trusted=trusted)
+
+    def insert(self, data):
+        return self.store.put(data, self.name, index=False)
+
+    def transform(self, value, **kwargs):
+        return value
+
+    def delete(self, *args, sql=None, **kwargs):
+        sql = sql or self.delete_sql
+        try:
+            cursor = self.store.get(self.name, sql=sql, sqlvars=kwargs, lazy=True)
+        except OperationalError as e:
+            pass
+        else:
+            cursor.close()
+
+    def drop(self, force=False):
+        return self.store.drop(self.name, force=force)
+
+    def join(self, other, on=None, on_left=None, on_right=None, columns_left=None, columns_right=None, **kwargs):
+        join_sql = self.join_sql
+        join_keys = {
+            ka: kb for ka, kb in zip(on or on_left or [], on or on_right or [])
+        }
+        _join_cond = ' and '.join([f'a.{k} = b.{v}' for k, v in join_keys.items()])
+        _join_cond = ' and ' + _join_cond if _join_cond else ''
+        _left_cols = ','.join(columns_left or ['a.*'])
+        _right_cols = ','.join(columns_right or ['b.*'])
+        sqlvars = {
+            '_join_vars_a': _left_cols,
+            '_join_vars_b': _right_cols,
+            '_join_sqltable': self.store.get_backend(self.name)._default_table(other.table or other.name),
+            '_join_cond': _join_cond,
+        }
+        return self.query(sql=join_sql, trusted=self.store.get_backend(self.name).sign(sqlvars), **sqlvars)
 
+    def count(self, sql=None, raw=False, **vars):
+        if raw and not vars:
+            sql = sql or 'select count(*) from :sqltable'
+            count = self.query(sql=sql).values[-1]
+        else:
+            count = len(self.query(sql=sql, **vars))
+        return count
 
+    def __call__(self, value, **kwargs):
+        value = self.query(**kwargs)
+        return self.transform(value, **kwargs)