From a6f04ec23e86134e890e9b3ddb4cb6da85bb33f9 Mon Sep 17 00:00:00 2001
From: Maarten Breddels <maartenbreddels@gmail.com>
Date: Tue, 19 Mar 2019 09:12:37 +0100
Subject: [PATCH] Make extending with new functions vaex more convenient (#188)

* Make extending with new functions vaex more convenient
---
 docs/source/api.rst                   |  25 +++-
 docs/source/conf.py                   |   4 +-
 packages/vaex-core/vaex/__init__.py   |   1 +
 packages/vaex-core/vaex/dataframe.py  |  64 +++++-----
 packages/vaex-core/vaex/expression.py |  63 +++-------
 packages/vaex-core/vaex/functions.py  | 174 +++++++++++++++++---------
 packages/vaex-core/vaex/scopes.py     |   2 +-
 packages/vaex-core/vaex/stat.py       |   2 +-
 8 files changed, 186 insertions(+), 149 deletions(-)

diff --git a/docs/source/api.rst b/docs/source/api.rst
index 56762ecd71..ab5fc8ddff 100644
--- a/docs/source/api.rst
+++ b/docs/source/api.rst
@@ -55,7 +55,7 @@ vaex-core
 ---------
 
 .. automodule:: vaex
-    :members: open, from_arrays, from_dict, from_items, from_arrow_table, from_csv, from_ascii, from_pandas, from_astropy_table, from_samp, open_many, server, example, app, delayed
+    :members: open, from_arrays, from_dict, from_items, from_arrow_table, from_csv, from_ascii, from_pandas, from_astropy_table, from_samp, open_many, register_function, server, example, app, delayed
     :undoc-members:
     :show-inheritance:
 
@@ -76,6 +76,27 @@ DataFrameLocal class
      :special-members:
 
 
+Expression class
+~~~~~~~~~~~~~~~~
+
+.. autoclass:: vaex.expression.Expression
+     :members:
+     :special-members:
+
+
+
+String operations
+~~~~~~~~~~~~~~~~
+
+.. autoclass:: vaex.expression.StringOperations
+     :members:
+     :special-members:
+
+.. autoclass:: vaex.expression.StringOperationsPandas
+     :members:
+     :special-members:
+
+
 vaex.stat module
 ~~~~~~~~~~~~~~~~
 
@@ -250,7 +271,7 @@ Boosted trees
 .. autoclass:: vaex.ml.lightgbm.LightGBMClassifier
      :members:
 
-.. autoclass:: vaex.ml.xgboost.XGBModel
+.. autoclass:: vaex.ml.xgboost.XGBoostModel
      :members:
 
 .. PyGBM support is in the incubator phase, which means support may disappear in future versions
diff --git a/docs/source/conf.py b/docs/source/conf.py
index caf9666710..aedf1fd3c8 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -69,9 +69,9 @@
 if 1:
 	import vaex
 	# The short X.Y version.
-	version = vaex.version
+	version = vaex.__version__
 	# The full version, including alpha/beta/rc tags.
-	release = vaex.version
+	release = vaex.__version__
 else:
 	print("failed finding vaex module, try finding version")
 	import sys
diff --git a/packages/vaex-core/vaex/__init__.py b/packages/vaex-core/vaex/__init__.py
index 78d766823d..b06720a07c 100644
--- a/packages/vaex-core/vaex/__init__.py
+++ b/packages/vaex-core/vaex/__init__.py
@@ -36,6 +36,7 @@
 import glob
 import vaex.dataframe
 import vaex.dataset
+from vaex.functions import register_function
 from . import stat
 # import vaex.file
 # import vaex.export
diff --git a/packages/vaex-core/vaex/dataframe.py b/packages/vaex-core/vaex/dataframe.py
index deeb9b25a2..0a0104dc6d 100644
--- a/packages/vaex-core/vaex/dataframe.py
+++ b/packages/vaex-core/vaex/dataframe.py
@@ -31,7 +31,7 @@
 import logging
 import vaex.kld
 from . import selections, tasks, scopes
-from .functions import expression_namespace
+from .expression import expression_namespace
 from .delayed import delayed, delayed_args, delayed_list
 from .column import Column, ColumnIndexed, ColumnSparse, ColumnString, ColumnConcatenatedLazy, str_type
 import vaex.events
@@ -209,6 +209,37 @@ def __getattr__(self, name):
         else:
             return object.__getattribute__(self, name)
 
+    @property
+    def func(self):
+        class Functions(object):
+            pass
+
+        functions = Functions()
+        for name, value in expression_namespace.items():
+            # f = vaex.expression.FunctionBuiltin(self, name)
+            def closure(name=name, value=value):
+                local_name = name
+                def wrap(*args, **kwargs):
+                    def myrepr(k):
+                        if isinstance(k, Expression):
+                            return str(k)
+                        else:
+                            return repr(k)
+                    arg_string = ", ".join([myrepr(k) for k in args] + ['{}={}'.format(name, myrepr(value)) for name, value in kwargs.items()])
+                    expression = "{}({})".format(local_name, arg_string)
+                    return vaex.expression.Expression(self, expression)
+                return wrap
+            f = closure()
+            try:
+                f = functools.wraps(value)(f)
+            except AttributeError:
+                pass # python2 quicks.. ?
+            setattr(functions, name, f)
+        for name, value in self.functions.items():
+            setattr(functions, name, value)
+
+        return functions
+
     @_hidden
     @vaex.utils.deprecated('use is_category')
     def iscategory(self, column):
@@ -4452,37 +4483,6 @@ class Datas(object):
             setattr(datas, name, array)
         return datas
 
-    @property
-    def func(self):
-        class Functions(object):
-            pass
-
-        functions = Functions()
-        for name, value in expression_namespace.items():
-            # f = vaex.expression.FunctionBuiltin(self, name)
-            def closure(name=name, value=value):
-                local_name = name
-                def wrap(*args, **kwargs):
-                    def myrepr(k):
-                        if isinstance(k, Expression):
-                            return str(k)
-                        else:
-                            return repr(k)
-                    arg_string = ", ".join([myrepr(k) for k in args] + ['{}={}'.format(name, myrepr(value)) for name, value in kwargs.items()])
-                    expression = "{}({})".format(local_name, arg_string)
-                    return vaex.expression.Expression(self, expression)
-                return wrap
-            f = closure()
-            try:
-                f = functools.wraps(value)(f)
-            except AttributeError:
-                pass # python2 quicks.. ?
-            setattr(functions, name, f)
-        for name, value in self.functions.items():
-            setattr(functions, name, value)
-
-        return functions
-
     def copy(self, column_names=None, virtual=True):
         df = DataFrameArrays()
         df._length_unfiltered = self._length_unfiltered
diff --git a/packages/vaex-core/vaex/expression.py b/packages/vaex-core/vaex/expression.py
index 5b17c4feb1..b1811d8562 100644
--- a/packages/vaex-core/vaex/expression.py
+++ b/packages/vaex-core/vaex/expression.py
@@ -8,7 +8,6 @@
 import numpy as np
 import tabulate
 
-from vaex.functions import expression_namespace, _scopes
 from vaex.utils import _ensure_strings_from_expressions, _ensure_string_from_expression
 from vaex.column import ColumnString
 import vaex.serialize
@@ -24,6 +23,10 @@
 default_shape = 128
 PRINT_MAX_COUNT = 10
 
+expression_namespace = {}
+expression_namespace['nan'] = np.nan
+
+
 _binary_ops = [
     dict(code="+", name='add', op=operator.add),
     dict(code="in", name='contains', op=operator.contains),
@@ -100,33 +103,11 @@ def f(a):
                     return Expression(self.ds, expression=expression)
                 attrs['__%s__' % op['name']] = f
             wrap(op)
-        for name, func_real in expression_namespace.items():
-            def wrap(name=name):
-                def f(*args, **kwargs):
-                    self = args[0]
-
-                    def to_expression(expression):
-                        if isinstance(expression, str):
-                            expression = repr(expression)
-                        if isinstance(expression, Expression):
-                            assert expression.ds == self.ds
-                            expression = expression.expression
-                        return expression
-                    expressions = [to_expression(e) for e in args]
-                    # print(name, expressions)
-                    expression = '{0}({1})'.format(name, ", ".join(expressions))
-                    return Expression(self.ds, expression=expression)
-                try:
-                    f = functools.wraps(func_real)(f)
-                except AttributeError:
-                    pass  # numpy ufuncs don't have a __module__, which may choke wraps
-
-                attrs['%s' % name] = f
-            if name not in attrs:
-                wrap(name)
         return type(future_class_name, future_class_parents, attrs)
 
+
 class DateTime(object):
+    """DateTime operations"""
     def __init__(self, expression):
         self.expression = expression
 
@@ -150,37 +131,21 @@ def hour(self):
     def weekofyear(self):
         return self.expression.ds.func.dt_weekofyear(self.expression)
 
+
 class StringOperations(object):
     """String operations"""
     def __init__(self, expression):
         self.expression = expression
 
-class PandasStringOperations(object):
+
+class StringOperationsPandas(object):
     """String operations using Pandas Series"""
     def __init__(self, expression):
         self.expression = expression
 
-for name, function in _scopes['str'].items():
-    full_name = 'str_' + name
-    def closure(name=name, full_name=full_name, function=function):
-        def wrapper(self, *args, **kwargs):
-            lazy_func = getattr(self.expression.ds.func, full_name)
-            args = (self.expression, ) + args
-            return lazy_func(*args, **kwargs)
-        return wrapper
-    setattr(StringOperations, name, closure())
-
-for name, function in _scopes['str_pandas'].items():
-    full_name = 'str_pandas_' + name
-    def closure(name=name, full_name=full_name, function=function):
-        def wrapper(self, *args, **kwargs):
-            lazy_func = getattr(self.expression.ds.func, full_name)
-            args = (self.expression, ) + args
-            return lazy_func(*args, **kwargs)
-        return wrapper
-    setattr(PandasStringOperations, name, closure())
 
 class Expression(with_metaclass(Meta)):
+    """Expression class"""
     def __init__(self, ds, expression):
         self.ds = ds
         if isinstance(expression, Expression):
@@ -199,7 +164,7 @@ def str(self):
     @property
     def str_pandas(self):
         """Gives access to string operations (using Pandas Series)"""
-        return PandasStringOperations(self)
+        return StringOperationsPandas(self)
 
     @property
     def values(self):
@@ -475,7 +440,7 @@ def jit_numba(self, verbose=False):
         import imp
         import hashlib
         names = []
-        funcs = set(vaex.dataset.expression_namespace.keys())
+        funcs = set(expression_namespace.keys())
         # if it's a virtual column, we probably want to optimize that
         # TODO: fully extract the virtual columns, i.e. depending ones?
         expression = self.expression
@@ -504,7 +469,7 @@ def jit_pythran(self, verbose=False):
             import hashlib
             # self._import_all(module)
             names = []
-            funcs = set(vaex.dataset.expression_namespace.keys())
+            funcs = set(expression_namespace.keys())
             expression = self.expression
             if expression in self.ds.virtual_columns:
                 expression = self.ds.virtual_columns[self.expression]
@@ -529,7 +494,7 @@ def f({0}):
 
             module = imp.load_dynamic(module_name, module_path)
             function_name = "f_" + m.hexdigest()
-            vaex.dataset.expression_namespace[function_name] = module.f
+            expression_namespace[function_name] = module.f
 
             return Expression(self.ds, "{0}({1})".format(function_name, argstring))
         finally:
diff --git a/packages/vaex-core/vaex/functions.py b/packages/vaex-core/vaex/functions.py
index 39791e80a8..ee4d92f769 100644
--- a/packages/vaex-core/vaex/functions.py
+++ b/packages/vaex-core/vaex/functions.py
@@ -4,13 +4,75 @@
 from vaex import column
 from vaex.column import _to_string_sequence
 import re
+import vaex.expression
+import functools
 
-# @vaex.serialize.register
+# @vaex.serialize.register_function
 # class Function(FunctionSerializable):
 
+scopes = {
+    'str': vaex.expression.StringOperations,
+    'str_pandas': vaex.expression.StringOperationsPandas,
+    'dt': vaex.expression.DateTime
+}
+
+def register_function(scope=None, as_property=False, name=None):
+    """Decorator to register a new function with vaex.
+
+    Example:
+
+    >>> import vaex
+    >>> df = vaex.example()
+    >>> @vaex.register_function()
+    >>> def invert(x):
+    >>>     return 1/x
+    >>> df.x.invert()
+
+
+    >>> import numpy as np
+    >>> df = vaex.from_arrays(departure=np.arange('2015-01-01', '2015-12-05', dtype='datetime64'))
+    >>> @vaex.register_function(as_property=True, scope='dt')
+    >>> def dt_relative_day(x):
+    >>>     return vaex.functions.dt_dayofyear(x)/365.
+    >>> df.departure.dt.relative_day
+    """
+    prefix = ''
+    if scope:
+        prefix = scope + "_"
+        if scope not in scopes:
+            raise KeyError("unknown scope")
+    def wrapper(f, name=name):
+        name = name or f.__name__
+        # remove possible prefix
+        if name.startswith(prefix):
+            name = name[len(prefix):]
+        full_name = prefix + name
+        if scope:
+            def closure(name=name, full_name=full_name, function=f):
+                def wrapper(self, *args, **kwargs):
+                    lazy_func = getattr(self.expression.ds.func, full_name)
+                    args = (self.expression, ) + args
+                    return lazy_func(*args, **kwargs)
+                return functools.wraps(function)(wrapper)
+            if as_property:
+                setattr(scopes[scope], name, property(closure()))
+            else:
+                setattr(scopes[scope], name, closure())
+        else:
+            def closure(name=name, full_name=full_name, function=f):
+                def wrapper(self, *args, **kwargs):
+                    lazy_func = getattr(self.ds.func, full_name)
+                    args = (self, ) + args
+                    return lazy_func(*args, **kwargs)
+                return functools.wraps(function)(wrapper)
+            setattr(vaex.expression.Expression, name, closure())
+        vaex.expression.expression_namespace[prefix + name] = f
+        return f  # we leave the original function as is
+    return wrapper
+
 # name maps to numpy function
 # <vaex name>:<numpy name>
-function_mapping = [name.strip().split(":") if ":" in name else (name, name) for name in """
+numpy_function_mapping = [name.strip().split(":") if ":" in name else (name, name) for name in """
 sinc
 sin
 cos
@@ -38,17 +100,25 @@
 minimum
 maximum
 clip
-nan
 searchsorted
 """.strip().split()]
-expression_namespace = {}
-for name, numpy_name in function_mapping:
+for name, numpy_name in numpy_function_mapping:
     if not hasattr(np, numpy_name):
         raise SystemError("numpy does not have: %s" % numpy_name)
     else:
-        expression_namespace[name] = getattr(np, numpy_name)
-
-
+        function = getattr(np, numpy_name)
+        def f(function=function):
+            def wrapper(*args, **kwargs):
+                return function(*args, **kwargs)
+            return wrapper
+        try:
+            function = functools.wraps(function)(f())
+        except AttributeError:
+            function = f()  # python 2 case
+        register_function(name=name)(function)
+
+
+@register_function()
 def fillna(ar, value, fill_nan=True, fill_masked=True):
     '''Returns an array where missing values are replaced by value.
 
@@ -74,67 +144,47 @@ def fillna(ar, value, fill_nan=True, fill_masked=True):
     return ar
 
 
-expression_namespace['fillna'] = fillna
-
+########## datetime operations ##########
 
+@register_function(scope='dt', as_property=True)
 def dt_dayofweek(x):
     import pandas as pd
     # x = x.astype("<M8[ns]")
     return pd.Series(x).dt.dayofweek.values
 
+@register_function(scope='dt', as_property=True)
 def dt_dayofyear(x):
     import pandas as pd
     # x = x.astype("<M8[ns]")
     return pd.Series(x).dt.dayofyear.values
 
+@register_function(scope='dt', as_property=True)
 def dt_year(x):
     import pandas as pd
     # x = x.astype("<M8[ns]")
     return pd.Series(x).dt.year.values
 
+@register_function(scope='dt', as_property=True)
 def dt_weekofyear(x):
     import pandas as pd
     # x = x.astype("<M8[ns]")
     return pd.Series(x).dt.weekofyear.values
 
+@register_function(scope='dt', as_property=True)
 def dt_hour(x):
     import pandas as pd
     # x = x.astype("<M8[ns]")
     return pd.Series(x).dt.hour.values
 
 
-expression_namespace["dt_dayofweek"] = dt_dayofweek
-expression_namespace["dt_dayofyear"] = dt_dayofyear
-expression_namespace["dt_year"] = dt_year
-expression_namespace["dt_weekofyear"] = dt_weekofyear
-expression_namespace["dt_hour"] = dt_hour
+########## string operations ##########
 
-_scopes = {}
-
-
-def register(scope=None):
-    prefix = ''
-    if scope:
-        prefix = scope + "_"
-        if scope not in _scopes:
-            _scopes[scope] = {}
-    def wrapper(f):
-        name = f.__name__
-        # remove possible prefix
-        if name.startswith(prefix):
-            name = name[len(prefix):]
-        if scope:
-            _scopes[scope][name] = f
-        expression_namespace[prefix + name] = f
-    return wrapper
-
-
-@register(scope='str')
+@register_function(scope='str')
 def str_capitalize(x):
     sl = _to_string_sequence(x).capitalize()
     return column.ColumnStringArrow(sl.bytes, sl.indices, sl.length, sl.offset, string_sequence=sl)
 
-@register(scope='str')
+@register_function(scope='str')
 def str_cat(x, other):
     sl1 = _to_string_sequence(x)
     sl2 = _to_string_sequence(other)
@@ -143,73 +193,73 @@ def str_cat(x, other):
 
 # TODO: center
 
-@register(scope='str')
+@register_function(scope='str')
 def str_contains(x, pattern, regex=True):
     return _to_string_sequence(x).search(pattern, regex)
 
 # TODO: default regex is False, which breaks with pandas
-@register(scope='str')
+@register_function(scope='str')
 def str_count(x, pat, regex=False):
     return _to_string_sequence(x).count(pat, regex)
 
 # TODO: what to do with decode and encode
 
-@register(scope='str')
+@register_function(scope='str')
 def str_endswith(x, pat):
     return _to_string_sequence(x).endswith(pat)
 
 # TODO: extract/extractall
 # TODO: find/findall/get/index/join
 
-@register(scope='str')
+@register_function(scope='str')
 def str_len(x):
     return _to_string_sequence(x).len()
 
-@register(scope='str')
+@register_function(scope='str')
 def str_byte_length(x):
     return _to_string_sequence(x).byte_length()
 
 # TODO: ljust
 
-@register(scope='str')
+@register_function(scope='str')
 def str_lower(x):
     sl = _to_string_sequence(x).lower()
     return column.ColumnStringArrow(sl.bytes, sl.indices, sl.length, sl.offset, string_sequence=sl)
 
 
-@register(scope='str')
+@register_function(scope='str')
 def str_lstrip(x, to_strip=None):
     # in c++ we give empty string the same meaning as None
     return _to_string_sequence(x).lstrip('' if to_strip is None else to_strip) if to_strip != '' else x
 
 # TODO: match, normalize, pad partition, repeat, replace, rfind, rindex, rjust, rpartition
 
-@register(scope='str')
+@register_function(scope='str')
 def str_rstrip(x, to_strip=None):
     # in c++ we give empty string the same meaning as None
     return _to_string_sequence(x).rstrip('' if to_strip is None else to_strip) if to_strip != '' else x
 
 # TODO: slice, slice_replace, 
 
-@register(scope='str')
+@register_function(scope='str')
 def str_split(x, pattern):  # TODO: support n
     sll = _to_string_sequence(x).split(pattern)
     return sll
 
 # TODO: rsplit
 
-@register(scope='str')
+@register_function(scope='str')
 def str_startswith(x, pat):
     return _to_string_sequence(x).startswith(pat)
 
-@register(scope='str')
+@register_function(scope='str')
 def str_strip(x, to_strip=None):
     # in c++ we give empty string the same meaning as None
     return _to_string_sequence(x).strip('' if to_strip is None else to_strip) if to_strip != '' else x
 
 # TODO: swapcase, title, translate
 
-@register(scope='str')
+@register_function(scope='str')
 def str_upper(x):
     sl = _to_string_sequence(x).upper()
     return column.ColumnStringArrow(sl.bytes, sl.indices, sl.length, sl.offset, string_sequence=sl)
@@ -221,18 +271,19 @@ def str_upper(x):
 # expression_namespace["str_capitalize"] = str_capitalize
 # expression_namespace["str_contains2"] = str_contains2
 
-@register(scope='str')
+@register_function(scope='str')
 def str_strip(x, chars=None):
+    """Removes leading and trailing characters."""
     # don't change the dtype, otherwise for each block the dtype may be different (string length)
     return np.char.strip(x, chars).astype(x.dtype)
 
-@register()
+@register_function()
 def to_string(x):
     # don't change the dtype, otherwise for each block the dtype may be different (string length)
     sl = vaex.strings.to_string(x)
     return column.ColumnStringArrow(sl.bytes, sl.indices, sl.length, sl.offset, string_sequence=sl)
 
-@register()
+@register_function()
 def format(x, format):
     """Uses http://www.cplusplus.com/reference/string/to_string/ for formatting"""
     # don't change the dtype, otherwise for each block the dtype may be different (string length)
@@ -240,9 +291,9 @@ def format(x, format):
     return column.ColumnStringArrow(sl.bytes, sl.indices, sl.length, sl.offset, string_sequence=sl)
 
 
-_scopes['str_pandas'] = {}
-
-for name, function in _scopes['str'].items():
+for name in dir(scopes['str']):
+    if name.startswith('__'):
+        continue
     def pandas_wrapper(name=name):
         def wrapper(*args, **kwargs):
             import pandas
@@ -258,23 +309,22 @@ def fix_arg(x):
             return method(*args, **kwargs)
         return wrapper
     wrapper = pandas_wrapper()
-    expression_namespace['str_pandas_' + name] = wrapper
-    _scopes['str_pandas'][name] = wrapper
+    wrapper.__doc__ = "Wrapper around pandas.Series.%s" % name
+    register_function(scope='str_pandas')(wrapper, name=name)
+    # expression_namespace['str_pandas_' + name] = wrapper
+    # _scopes['str_pandas'][name] = wrapper
 
 # expression_namespace['str_strip'] = str_strip
 
+@register_function(name='float')
 def _float(x):
     return x.astype(np.float64)
 
+@register_function(name='astype')
 def _astype(x, dtype):
     return x.astype(dtype)
 
 
-expression_namespace["float"] = _float
-expression_namespace["astype"] = _astype
-
-
-
 def add_geo_json(ds, json_or_file, column_name, longitude_expression, latitude_expresion, label=None, persist=True, overwrite=False, inplace=False, mapping=None):
     ds = ds if inplace else ds.copy()
     if not isinstance(json_or_file, (list, tuple)):
diff --git a/packages/vaex-core/vaex/scopes.py b/packages/vaex-core/vaex/scopes.py
index a966f2e209..72fc1d1ece 100644
--- a/packages/vaex-core/vaex/scopes.py
+++ b/packages/vaex-core/vaex/scopes.py
@@ -20,7 +20,7 @@
     as_flat_float,
     as_flat_array,
     _split_and_combine_mask)
-from .functions import expression_namespace
+from .expression import expression_namespace
 import vaex.expression
 
 logger = logging.getLogger('vaex.scopes')
diff --git a/packages/vaex-core/vaex/stat.py b/packages/vaex-core/vaex/stat.py
index 2915248ee4..9518953752 100644
--- a/packages/vaex-core/vaex/stat.py
+++ b/packages/vaex-core/vaex/stat.py
@@ -1,6 +1,6 @@
 from .expression import _unary_ops, _binary_ops, reversable
 from future.utils import with_metaclass
-from vaex.functions import expression_namespace
+from vaex.expression import expression_namespace
 from vaex.delayed import delayed