From 7f62c6a19dafec50f9fa1c2261e664be0ebff1d0 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Fri, 22 Jul 2022 07:27:03 -0700 Subject: [PATCH 1/4] initial --- dask_sql/context.py | 9 +++++++++ tests/integration/test_function.py | 12 ++++++++++++ 2 files changed, 21 insertions(+) diff --git a/dask_sql/context.py b/dask_sql/context.py index 98cc46e21..6033503d4 100644 --- a/dask_sql/context.py +++ b/dask_sql/context.py @@ -898,6 +898,15 @@ def _register_callable( row_udf: bool = False, ): """Helper function to do the function or aggregation registration""" + + # validate UDF metadata + try: + _ = python_to_sql_type(return_type) + for t in (param[1] for param in parameters): + _ = python_to_sql_type(t) + except NotImplementedError as e: + raise e + schema_name = schema_name or self.schema_name schema = self.schema[schema_name] diff --git a/tests/integration/test_function.py b/tests/integration/test_function.py index 92fc58b14..6adc08673 100644 --- a/tests/integration/test_function.py +++ b/tests/integration/test_function.py @@ -199,3 +199,15 @@ def f(x): c.register_aggregation(fagg, "fagg", [("x", np.float64)], np.float64) c.register_aggregation(fagg, "fagg", [("x", np.float64)], np.float64, replace=True) + + +@pytest.mark.parametrize("invalid_dtype", [np.timedelta64, None, "a string"]) +def test_unsupported_dtype(c): + def f(x): + return x**2 + + with pytest.raises(NotImplementedError): + c.register_function(f, "f", [("x", np.int64)], np.timedelta64) + + with pytest.raises(NotImplementedError): + c.register_function(f, "f", [("x", np.timedelta64)], np.int64) From 05d019ea9aced70d03e363f291bd9116a02999b1 Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Fri, 22 Jul 2022 09:15:41 -0700 Subject: [PATCH 2/4] improvements --- dask_sql/context.py | 8 ++------ tests/integration/test_function.py | 10 ++++++---- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/dask_sql/context.py b/dask_sql/context.py index 6033503d4..4103330ce 100644 --- a/dask_sql/context.py +++ b/dask_sql/context.py @@ -900,12 +900,8 @@ def _register_callable( """Helper function to do the function or aggregation registration""" # validate UDF metadata - try: - _ = python_to_sql_type(return_type) - for t in (param[1] for param in parameters): - _ = python_to_sql_type(t) - except NotImplementedError as e: - raise e + for dt in (*(param[1] for param in parameters), return_type): + _ = python_to_sql_type(dt) schema_name = schema_name or self.schema_name schema = self.schema[schema_name] diff --git a/tests/integration/test_function.py b/tests/integration/test_function.py index 6adc08673..ca5a3b70a 100644 --- a/tests/integration/test_function.py +++ b/tests/integration/test_function.py @@ -201,13 +201,15 @@ def f(x): c.register_aggregation(fagg, "fagg", [("x", np.float64)], np.float64, replace=True) -@pytest.mark.parametrize("invalid_dtype", [np.timedelta64, None, "a string"]) -def test_unsupported_dtype(c): +@pytest.mark.parametrize("dtype", [np.timedelta64, None, "a string"]) +def test_unsupported_dtype(c, dtype): def f(x): return x**2 + # test that an invalid return type raises with pytest.raises(NotImplementedError): - c.register_function(f, "f", [("x", np.int64)], np.timedelta64) + c.register_function(f, "f", [("x", np.int64)], dtype) + # test that an invalid param type raises with pytest.raises(NotImplementedError): - c.register_function(f, "f", [("x", np.timedelta64)], np.int64) + c.register_function(f, "f", [("x", dtype)], np.int64) From 32f8a9b2a8797554d76dcb13617786bab1f7036c Mon Sep 17 00:00:00 2001 From: brandon-b-miller Date: Fri, 22 Jul 2022 10:14:25 -0700 Subject: [PATCH 3/4] bugfixes --- dask_sql/context.py | 4 ---- dask_sql/datacontainer.py | 12 ++++++------ dask_sql/mappings.py | 5 +++++ tests/integration/test_function.py | 7 +------ 4 files changed, 12 insertions(+), 16 deletions(-) diff --git a/dask_sql/context.py b/dask_sql/context.py index 4103330ce..7feb5ffde 100644 --- a/dask_sql/context.py +++ b/dask_sql/context.py @@ -899,10 +899,6 @@ def _register_callable( ): """Helper function to do the function or aggregation registration""" - # validate UDF metadata - for dt in (*(param[1] for param in parameters), return_type): - _ = python_to_sql_type(dt) - schema_name = schema_name or self.schema_name schema = self.schema[schema_name] diff --git a/dask_sql/datacontainer.py b/dask_sql/datacontainer.py index db77c9dfc..eb25d0253 100644 --- a/dask_sql/datacontainer.py +++ b/dask_sql/datacontainer.py @@ -4,6 +4,8 @@ import dask.dataframe as dd import pandas as pd +from dask_sql.mappings import python_to_sql_type + ColumnType = Union[str, int] FunctionDescription = namedtuple( @@ -198,11 +200,10 @@ def __init__(self, func, row_udf: bool, params, return_type=None): self.names = [param[0] for param in params] - if return_type is None: - # These UDFs go through apply and without providing - # a return type, dask will attempt to guess it, and - # dask might be wrong. - raise ValueError("Return type must be provided") + # validate UDF metadata + for dt in (*(param[1] for param in params), return_type): + _ = python_to_sql_type(dt) + self.meta = (None, return_type) def __call__(self, *args, **kwargs): @@ -218,7 +219,6 @@ def __call__(self, *args, **kwargs): df = column_args[0].to_frame(self.names[0]) for name, col in zip(self.names[1:], column_args[1:]): df[name] = col - result = df.apply( self.func, axis=1, args=tuple(scalar_args), meta=self.meta ).astype(self.meta[1]) diff --git a/dask_sql/mappings.py b/dask_sql/mappings.py index e59025918..623c38a37 100644 --- a/dask_sql/mappings.py +++ b/dask_sql/mappings.py @@ -88,6 +88,11 @@ def python_to_sql_type(python_type): """Mapping between python and SQL types.""" + if python_type in (int, float): + python_type = np.dtype(python_type) + elif python_type is str: + python_type = np.dtype("object") + if isinstance(python_type, np.dtype): python_type = python_type.type diff --git a/tests/integration/test_function.py b/tests/integration/test_function.py index ca5a3b70a..d8ba40c0f 100644 --- a/tests/integration/test_function.py +++ b/tests/integration/test_function.py @@ -52,17 +52,12 @@ def f(row): @pytest.mark.parametrize( "retty", - [None, np.float64, np.float32, np.int64, np.int32, np.int16, np.int8, np.bool_], + [np.float64, np.float32, np.int64, np.int32, np.int16, np.int8, np.bool_], ) def test_custom_function_row_return_types(c, df, retty): def f(row): return row["x"] ** 2 - if retty is None: - with pytest.raises(ValueError): - c.register_function(f, "f", [("x", np.float64)], retty, row_udf=True) - return - c.register_function(f, "f", [("x", np.float64)], retty, row_udf=True) return_df = c.sql("SELECT F(a) AS a FROM df") From 4f9c3549e515375fc57f99c017fe08455cffbffb Mon Sep 17 00:00:00 2001 From: Charles Blackmon-Luca <20627856+charlesbluca@users.noreply.github.com> Date: Tue, 26 Jul 2022 07:57:49 -0700 Subject: [PATCH 4/4] Move UDF validation to registration, cache relevant info --- dask_sql/context.py | 21 ++++++++----- dask_sql/datacontainer.py | 6 ---- tests/unit/test_context.py | 62 +++++++++++++++++++++++--------------- 3 files changed, 52 insertions(+), 37 deletions(-) diff --git a/dask_sql/context.py b/dask_sql/context.py index 7feb5ffde..6dc1850f1 100644 --- a/dask_sql/context.py +++ b/dask_sql/context.py @@ -748,7 +748,7 @@ def _prepare_schemas(self): logger.debug("No custom functions defined.") for function_description in schema.function_lists: name = function_description.name - sql_return_type = python_to_sql_type(function_description.return_type) + sql_return_type = function_description.return_type if function_description.aggregation: logger.debug(f"Adding function '{name}' to schema as aggregation.") dask_function = DaskAggregateFunction(name, sql_return_type) @@ -771,10 +771,7 @@ def _prepare_schemas(self): @staticmethod def _add_parameters_from_description(function_description, dask_function): for parameter in function_description.parameters: - param_name, param_type = parameter - sql_param_type = python_to_sql_type(param_type) - - dask_function.addParameter(param_name, sql_param_type, False) + dask_function.addParameter(*parameter, False) return dask_function @@ -902,6 +899,12 @@ def _register_callable( schema_name = schema_name or self.schema_name schema = self.schema[schema_name] + # validate and cache UDF metadata + sql_parameters = [ + (name, python_to_sql_type(param_type)) for name, param_type in parameters + ] + sql_return_type = python_to_sql_type(return_type) + if not aggregation: f = UDF(f, row_udf, parameters, return_type) lower_name = name.lower() @@ -921,9 +924,13 @@ def _register_callable( ) schema.function_lists.append( - FunctionDescription(name.upper(), parameters, return_type, aggregation) + FunctionDescription( + name.upper(), sql_parameters, sql_return_type, aggregation + ) ) schema.function_lists.append( - FunctionDescription(name.lower(), parameters, return_type, aggregation) + FunctionDescription( + name.lower(), sql_parameters, sql_return_type, aggregation + ) ) schema.functions[lower_name] = f diff --git a/dask_sql/datacontainer.py b/dask_sql/datacontainer.py index eb25d0253..f81952e68 100644 --- a/dask_sql/datacontainer.py +++ b/dask_sql/datacontainer.py @@ -4,8 +4,6 @@ import dask.dataframe as dd import pandas as pd -from dask_sql.mappings import python_to_sql_type - ColumnType = Union[str, int] FunctionDescription = namedtuple( @@ -200,10 +198,6 @@ def __init__(self, func, row_udf: bool, params, return_type=None): self.names = [param[0] for param in params] - # validate UDF metadata - for dt in (*(param[1] for param in params), return_type): - _ = python_to_sql_type(dt) - self.meta = (None, return_type) def __call__(self, *args, **kwargs): diff --git a/tests/unit/test_context.py b/tests/unit/test_context.py index b8cfa6504..fca5c7454 100644 --- a/tests/unit/test_context.py +++ b/tests/unit/test_context.py @@ -6,6 +6,7 @@ from dask_sql import Context from dask_sql.datacontainer import Statistics +from dask_sql.mappings import python_to_sql_type from tests.utils import assert_eq try: @@ -198,6 +199,11 @@ def g(gpu=gpu): g(gpu=gpu) +int_sql_type = python_to_sql_type(int) +float_sql_type = python_to_sql_type(float) +str_sql_type = python_to_sql_type(str) + + def test_function_adding(): c = Context() @@ -211,12 +217,12 @@ def test_function_adding(): assert c.schema[c.schema_name].functions["f"].func == f assert len(c.schema[c.schema_name].function_lists) == 2 assert c.schema[c.schema_name].function_lists[0].name == "F" - assert c.schema[c.schema_name].function_lists[0].parameters == [("x", int)] - assert c.schema[c.schema_name].function_lists[0].return_type == float + assert c.schema[c.schema_name].function_lists[0].parameters == [("x", int_sql_type)] + assert c.schema[c.schema_name].function_lists[0].return_type == float_sql_type assert not c.schema[c.schema_name].function_lists[0].aggregation assert c.schema[c.schema_name].function_lists[1].name == "f" - assert c.schema[c.schema_name].function_lists[1].parameters == [("x", int)] - assert c.schema[c.schema_name].function_lists[1].return_type == float + assert c.schema[c.schema_name].function_lists[1].parameters == [("x", int_sql_type)] + assert c.schema[c.schema_name].function_lists[1].return_type == float_sql_type assert not c.schema[c.schema_name].function_lists[1].aggregation # Without replacement @@ -226,12 +232,16 @@ def test_function_adding(): assert c.schema[c.schema_name].functions["f"].func == f assert len(c.schema[c.schema_name].function_lists) == 4 assert c.schema[c.schema_name].function_lists[2].name == "F" - assert c.schema[c.schema_name].function_lists[2].parameters == [("x", float)] - assert c.schema[c.schema_name].function_lists[2].return_type == int + assert c.schema[c.schema_name].function_lists[2].parameters == [ + ("x", float_sql_type) + ] + assert c.schema[c.schema_name].function_lists[2].return_type == int_sql_type assert not c.schema[c.schema_name].function_lists[2].aggregation assert c.schema[c.schema_name].function_lists[3].name == "f" - assert c.schema[c.schema_name].function_lists[3].parameters == [("x", float)] - assert c.schema[c.schema_name].function_lists[3].return_type == int + assert c.schema[c.schema_name].function_lists[3].parameters == [ + ("x", float_sql_type) + ] + assert c.schema[c.schema_name].function_lists[3].return_type == int_sql_type assert not c.schema[c.schema_name].function_lists[3].aggregation # With replacement @@ -242,12 +252,12 @@ def test_function_adding(): assert c.schema[c.schema_name].functions["f"].func == f assert len(c.schema[c.schema_name].function_lists) == 2 assert c.schema[c.schema_name].function_lists[0].name == "F" - assert c.schema[c.schema_name].function_lists[0].parameters == [("x", str)] - assert c.schema[c.schema_name].function_lists[0].return_type == str + assert c.schema[c.schema_name].function_lists[0].parameters == [("x", str_sql_type)] + assert c.schema[c.schema_name].function_lists[0].return_type == str_sql_type assert not c.schema[c.schema_name].function_lists[0].aggregation assert c.schema[c.schema_name].function_lists[1].name == "f" - assert c.schema[c.schema_name].function_lists[1].parameters == [("x", str)] - assert c.schema[c.schema_name].function_lists[1].return_type == str + assert c.schema[c.schema_name].function_lists[1].parameters == [("x", str_sql_type)] + assert c.schema[c.schema_name].function_lists[1].return_type == str_sql_type assert not c.schema[c.schema_name].function_lists[1].aggregation @@ -264,12 +274,12 @@ def test_aggregation_adding(): assert c.schema[c.schema_name].functions["f"] == f assert len(c.schema[c.schema_name].function_lists) == 2 assert c.schema[c.schema_name].function_lists[0].name == "F" - assert c.schema[c.schema_name].function_lists[0].parameters == [("x", int)] - assert c.schema[c.schema_name].function_lists[0].return_type == float + assert c.schema[c.schema_name].function_lists[0].parameters == [("x", int_sql_type)] + assert c.schema[c.schema_name].function_lists[0].return_type == float_sql_type assert c.schema[c.schema_name].function_lists[0].aggregation assert c.schema[c.schema_name].function_lists[1].name == "f" - assert c.schema[c.schema_name].function_lists[1].parameters == [("x", int)] - assert c.schema[c.schema_name].function_lists[1].return_type == float + assert c.schema[c.schema_name].function_lists[1].parameters == [("x", int_sql_type)] + assert c.schema[c.schema_name].function_lists[1].return_type == float_sql_type assert c.schema[c.schema_name].function_lists[1].aggregation # Without replacement @@ -279,12 +289,16 @@ def test_aggregation_adding(): assert c.schema[c.schema_name].functions["f"] == f assert len(c.schema[c.schema_name].function_lists) == 4 assert c.schema[c.schema_name].function_lists[2].name == "F" - assert c.schema[c.schema_name].function_lists[2].parameters == [("x", float)] - assert c.schema[c.schema_name].function_lists[2].return_type == int + assert c.schema[c.schema_name].function_lists[2].parameters == [ + ("x", float_sql_type) + ] + assert c.schema[c.schema_name].function_lists[2].return_type == int_sql_type assert c.schema[c.schema_name].function_lists[2].aggregation assert c.schema[c.schema_name].function_lists[3].name == "f" - assert c.schema[c.schema_name].function_lists[3].parameters == [("x", float)] - assert c.schema[c.schema_name].function_lists[3].return_type == int + assert c.schema[c.schema_name].function_lists[3].parameters == [ + ("x", float_sql_type) + ] + assert c.schema[c.schema_name].function_lists[3].return_type == int_sql_type assert c.schema[c.schema_name].function_lists[3].aggregation # With replacement @@ -295,12 +309,12 @@ def test_aggregation_adding(): assert c.schema[c.schema_name].functions["f"] == f assert len(c.schema[c.schema_name].function_lists) == 2 assert c.schema[c.schema_name].function_lists[0].name == "F" - assert c.schema[c.schema_name].function_lists[0].parameters == [("x", str)] - assert c.schema[c.schema_name].function_lists[0].return_type == str + assert c.schema[c.schema_name].function_lists[0].parameters == [("x", str_sql_type)] + assert c.schema[c.schema_name].function_lists[0].return_type == str_sql_type assert c.schema[c.schema_name].function_lists[0].aggregation assert c.schema[c.schema_name].function_lists[1].name == "f" - assert c.schema[c.schema_name].function_lists[1].parameters == [("x", str)] - assert c.schema[c.schema_name].function_lists[1].return_type == str + assert c.schema[c.schema_name].function_lists[1].parameters == [("x", str_sql_type)] + assert c.schema[c.schema_name].function_lists[1].return_type == str_sql_type assert c.schema[c.schema_name].function_lists[1].aggregation