Allow function reregistration (#161)

Author: nils-braun

dask_sql/context.py

@@ -224,6 +224,7 @@ def register_function(
         name: str,
         parameters: List[Tuple[str, type]],
         return_type: type,
+        replace: bool = False,
     ):
         """
         Register a custom function with the given name.
@@ -267,6 +268,7 @@ def f(x):
             parameters (:obj:`List[Tuple[str, type]]`): A list ot tuples of parameter name and parameter type.
                 Use `numpy dtypes <https://numpy.org/doc/stable/reference/arrays.dtypes.html>`_ if possible.
             return_type (:obj:`type`): The return type of the function
+            replace (:obj:`bool`): Do not raise an error if the function is already present
 
         See also:
             :func:`register_aggregation`
@@ -277,7 +279,7 @@ def f(x):
         )
 
         name = name.lower()
-        if name in self.functions:
+        if not replace and name in self.functions:
             if self.functions[name] != f:
                 raise ValueError(
                     "Registering different functions with the same name is not allowed"
@@ -290,6 +292,7 @@ def register_aggregation(
         name: str,
         parameters: List[Tuple[str, type]],
         return_type: type,
+        replace: bool = False,
     ):
         """
         Register a custom aggregation with the given name.
@@ -333,6 +336,7 @@ def register_aggregation(
             parameters (:obj:`List[Tuple[str, type]]`): A list ot tuples of parameter name and parameter type.
                 Use `numpy dtypes <https://numpy.org/doc/stable/reference/arrays.dtypes.html>`_ if possible.
             return_type (:obj:`type`): The return type of the function
+            replace (:obj:`bool`): Do not raise an error if the function is already present
 
         See also:
             :func:`register_function`
@@ -343,7 +347,7 @@ def register_aggregation(
         )
 
         name = name.lower()
-        if name in self.functions:
+        if not replace and name in self.functions:
             if self.functions[name] != f:
                 raise ValueError(
                     "Registering different functions with the same name is not allowed"

tests/integration/test_function.py

@@ -69,6 +69,9 @@ def f(x):
     with pytest.raises(ValueError):
         c.register_function(f, "f", [("x", np.float64)], np.float64)
 
+    # only if we replace it
+    c.register_function(f, "f", [("x", np.float64)], np.float64, replace=True)
+
     fagg = dd.Aggregation("f", lambda x: x.sum(), lambda x: x.sum())
     c.register_aggregation(fagg, "fagg", [("x", np.float64)], np.float64)
     c.register_aggregation(fagg, "fagg", [("x", np.int64)], np.int64)
@@ -77,3 +80,5 @@ def f(x):
 
     with pytest.raises(ValueError):
         c.register_aggregation(fagg, "fagg", [("x", np.float64)], np.float64)
+
+    c.register_aggregation(fagg, "fagg", [("x", np.float64)], np.float64, replace=True)