Add schema namespaces for UDFs (#1268)

## Description

Allows dataset and metric UDFs to be registered in specific schema
namespaces

```
@register_dataset_udf(["col1"], schema_name="bob")
def bob(x):
    return x["col1"]

@register_metric_udf("col1", schema_name="bob")
def rob(x):
    return x

@register_dataset_udf(["col1"], "add5")
def fob(x):
    return x["col1"] + 5

default_schema = udf_schema()
data = pd.DataFrame({"col1": [42, 12, 7]})
default_view = why.log(data, schema=default_schema).view()  # no bob or rob

bob_schema = udf_schema(schema_name="bob")
data = pd.DataFrame({"col1": [42, 12, 7]})  # original data frame stomped on by previous log() UDFs
bob_view = why.log(data, schema=bob_schema).view()  # bob & rob, but no add5
```

## Related

Relates to #1263

<!-- Reference related commits, issues and pull requests. Type `#` and
select from the list. -->

Closes [clickup task](https://app.clickup.com/t/866ab6rqn)


- [ ] I have reviewed the [Guidelines for Contributing](CONTRIBUTING.md)
and the [Code of Conduct](CODE_OF_CONDUCT.md).

Author: richard-rogers

python/tests/experimental/core/test_udf_schema.py

@@ -284,3 +284,34 @@ def test_udf_track() -> None:
     assert prod_summary["counts/n"] == 3
     div_summary = results.get_column("ratio").to_summary_dict()
     assert div_summary["distribution/n"] == 3
+
+
+@register_dataset_udf(["schema.col1"], schema_name="bob")
+def bob(x):
+    return x["schema.col1"]
+
+
+@register_metric_udf("schema.col1", schema_name="bob")
+def rob(x):
+    return x
+
+
+@register_dataset_udf(["schema.col1"], "add5")
+def fob(x):
+    return x["schema.col1"] + 5
+
+
+def test_schema_name() -> None:
+    default_schema = udf_schema()
+    data = pd.DataFrame({"schema.col1": [42, 12, 7]})
+    default_view = why.log(data, schema=default_schema).view()
+    assert "add5" in default_view.get_columns()
+    assert "bob" not in default_view.get_columns()
+    assert "udf" not in default_view.get_column("schema.col1").get_metric_names()
+
+    bob_schema = udf_schema(schema_name="bob")
+    data = pd.DataFrame({"schema.col1": [42, 12, 7]})
+    bob_view = why.log(data, schema=bob_schema).view()
+    assert "add5" not in bob_view.get_columns()
+    assert "bob" in bob_view.get_columns()
+    assert "udf" in bob_view.get_column("schema.col1").get_metric_names()

python/whylogs/experimental/core/metrics/udf_metric.py

@@ -177,13 +177,17 @@ def zero(cls, config: Optional[MetricConfig] = None) -> "UdfMetric":
 register_metric(UdfMetric)
 
 
-_col_name_submetrics: Dict[str, List[Tuple[str, Callable[[Any], Any]]]] = defaultdict(list)
-_col_name_submetric_schema: Dict[str, SubmetricSchema] = dict()
-_col_name_type_mapper: Dict[str, TypeMapper] = dict()
+_col_name_submetrics: Dict[str, Dict[str, List[Tuple[str, Callable[[Any], Any]]]]] = defaultdict(
+    lambda: defaultdict(list)
+)
+_col_name_submetric_schema: Dict[str, Dict[str, SubmetricSchema]] = defaultdict(dict)
+_col_name_type_mapper: Dict[str, Dict[str, TypeMapper]] = defaultdict(dict)
 
-_col_type_submetrics: Dict[DataType, List[Tuple[str, Callable[[Any], Any]]]] = defaultdict(list)
-_col_type_submetric_schema: Dict[DataType, SubmetricSchema] = dict()
-_col_type_type_mapper: Dict[DataType, TypeMapper] = dict()
+_col_type_submetrics: Dict[str, Dict[DataType, List[Tuple[str, Callable[[Any], Any]]]]] = defaultdict(
+    lambda: defaultdict(list)
+)
+_col_type_submetric_schema: Dict[str, Dict[DataType, SubmetricSchema]] = defaultdict(dict)
+_col_type_type_mapper: Dict[str, Dict[DataType, TypeMapper]] = defaultdict(dict)
 
 
 def register_metric_udf(
@@ -193,6 +197,7 @@ def register_metric_udf(
     submetric_schema: Optional[SubmetricSchema] = None,
     type_mapper: Optional[TypeMapper] = None,
     namespace: Optional[str] = None,
+    schema_name: str = "",
 ) -> Callable[[Any], Any]:
     """
     Decorator to easily configure UdfMetrics for your data set. Decorate your UDF
@@ -229,32 +234,32 @@ def decorator_register(func):
         subname = submetric_name or func.__name__
         subname = f"{namespace}.{subname}" if namespace else subname
         if col_name is not None:
-            _col_name_submetrics[col_name].append((subname, func))
+            _col_name_submetrics[schema_name][col_name].append((subname, func))
             if submetric_schema is not None:
-                if col_name in _col_name_submetric_schema:
+                if col_name in _col_name_submetric_schema[schema_name]:
                     logger.warn(f"Overwriting submetric schema for column {col_name}")
-                _col_name_submetric_schema[col_name] = submetric_schema
+                _col_name_submetric_schema[schema_name][col_name] = submetric_schema
             if type_mapper is not None:
-                if col_name in _col_name_type_mapper:
+                if col_name in _col_name_type_mapper[schema_name]:
                     logger.warn(f"Overwriting UdfMetric type mapper for column {col_name}")
-                _col_name_type_mapper[col_name] = type_mapper
+                _col_name_type_mapper[schema_name][col_name] = type_mapper
         else:
-            _col_type_submetrics[col_type].append((subname, func))
+            _col_type_submetrics[schema_name][col_type].append((subname, func))
             if submetric_schema is not None:
-                if col_type in _col_type_submetric_schema:
+                if col_type in _col_type_submetric_schema[schema_name]:
                     logger.warn(f"Overwriting submetric schema for column type {col_type}")
-                _col_type_submetric_schema[col_type] = submetric_schema
+                _col_type_submetric_schema[schema_name][col_type] = submetric_schema
             if type_mapper is not None:
-                if col_type in _col_type_type_mapper:
+                if col_type in _col_type_type_mapper[schema_name]:
                     logger.warn(f"Overwriting UdfMetric type mapper for column type {col_type}")
-                _col_type_type_mapper[col_type] = type_mapper
+                _col_type_type_mapper[schema_name][col_type] = type_mapper
 
         return func
 
     return decorator_register
 
 
-def generate_udf_resolvers() -> List[ResolverSpec]:
+def generate_udf_resolvers(schema_name: str = "") -> List[ResolverSpec]:
     """
     Generates a list of ResolverSpecs that implement the UdfMetrics specified
     by the @register_metric_udf decorators. The result only includes the UdfMetric,
@@ -283,25 +288,25 @@ def upper(x):
 
     resolvers: List[ResolverSpec] = list()
     udfs: Dict[str, Callable[[Any], Any]]
-    for col_name, submetrics in _col_name_submetrics.items():
+    for col_name, submetrics in _col_name_submetrics[schema_name].items():
         udfs = dict()
         for submetric in submetrics:
             udfs[submetric[0]] = submetric[1]
         config = UdfMetricConfig(
             udfs=udfs,
-            submetric_schema=_col_name_submetric_schema.get(col_name) or default_schema(),
-            type_mapper=_col_name_type_mapper.get(col_name) or StandardTypeMapper(),
+            submetric_schema=_col_name_submetric_schema[schema_name].get(col_name) or default_schema(),
+            type_mapper=_col_name_type_mapper[schema_name].get(col_name) or StandardTypeMapper(),
         )
         resolvers.append(ResolverSpec(col_name, None, [MetricSpec(UdfMetric, config)]))
 
-    for col_type, submetrics in _col_type_submetrics.items():
+    for col_type, submetrics in _col_type_submetrics[schema_name].items():
         udfs = dict()
         for submetric in submetrics:
             udfs[submetric[0]] = submetric[1]
         config = UdfMetricConfig(
             udfs=udfs,
-            submetric_schema=_col_type_submetric_schema.get(col_type) or default_schema(),
-            type_mapper=_col_type_type_mapper.get(col_type) or StandardTypeMapper(),
+            submetric_schema=_col_type_submetric_schema[schema_name].get(col_type) or default_schema(),
+            type_mapper=_col_type_type_mapper[schema_name].get(col_type) or StandardTypeMapper(),
         )
         resolvers.append(ResolverSpec(None, col_type, [MetricSpec(UdfMetric, config)]))
 
@@ -322,6 +327,7 @@ def udf_metric_schema(
     schema_based_automerge: bool = False,
     segments: Optional[Dict[str, SegmentationPartition]] = None,
     validators: Optional[Dict[str, List[Validator]]] = None,
+    schema_name: str = "",
 ) -> DeclarativeSchema:
     """
     Generates a DeclarativeSchema that implement the UdfMetrics specified
@@ -347,7 +353,7 @@ def upper(x):
     STANDARD_RESOLVER, the default metrics are also tracked for every column.
     """
 
-    resolvers = generate_udf_resolvers()
+    resolvers = generate_udf_resolvers(schema_name)
     non_udf_resolvers = non_udf_resolvers if non_udf_resolvers is not None else UDF_BASE_RESOLVER
 
     return DeclarativeSchema(

python/whylogs/experimental/core/udf_schema.py

@@ -1,4 +1,5 @@
 import logging
+from collections import defaultdict
 from copy import deepcopy
 from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple
@@ -112,15 +113,16 @@ def _run_udfs(
         return new_df, new_columns
 
 
-_multicolumn_udfs: List[UdfSpec] = []
-_resolver_specs: List[ResolverSpec] = []
+_multicolumn_udfs: Dict[str, List[UdfSpec]] = defaultdict(list)
+_resolver_specs: Dict[str, List[ResolverSpec]] = defaultdict(list)
 
 
 def register_dataset_udf(
     col_names: List[str],
     udf_name: Optional[str] = None,
     metrics: Optional[List[MetricSpec]] = None,
     namespace: Optional[str] = None,
+    schema_name: str = "",
 ) -> Callable[[Any], Any]:
     """
     Decorator to easily configure UDFs for your data set. Decorate your UDF
@@ -144,16 +146,16 @@ def decorator_register(func):
         global _multicolumn_udfs, _resolver_specs
         name = udf_name or func.__name__
         name = f"{namespace}.{name}" if namespace else name
-        _multicolumn_udfs.append(UdfSpec(col_names, {name: func}))
+        _multicolumn_udfs[schema_name].append(UdfSpec(col_names, {name: func}))
         if metrics:
-            _resolver_specs.append(ResolverSpec(name, None, deepcopy(metrics)))
+            _resolver_specs[schema_name].append(ResolverSpec(name, None, deepcopy(metrics)))
 
         return func
 
     return decorator_register
 
 
-def generate_udf_specs(other_udf_specs: Optional[List[UdfSpec]] = None) -> List[UdfSpec]:
+def generate_udf_specs(other_udf_specs: Optional[List[UdfSpec]] = None, schema_name: str = "") -> List[UdfSpec]:
     """
     Generates a list UdfSpecs that implement the UDFs specified
     by the @register_dataset_udf decorators. You can provide a list of
@@ -175,7 +177,7 @@ def add5(x):
     for every column.
     """
     specs = list(other_udf_specs) if other_udf_specs else []
-    specs += _multicolumn_udfs
+    specs += _multicolumn_udfs[schema_name]
     return specs
 
 
@@ -189,16 +191,17 @@ def udf_schema(
     schema_based_automerge: bool = False,
     segments: Optional[Dict[str, SegmentationPartition]] = None,
     validators: Optional[Dict[str, List[Validator]]] = None,
+    schema_name: str = "",
 ) -> UdfSchema:
     """
     Returns a UdfSchema that implements any registered UDFs, along with any
     other_udf_specs or resolvers passed in.
     """
     if resolvers is not None:
-        resolver_specs = resolvers + _resolver_specs
+        resolver_specs = resolvers + _resolver_specs[schema_name]
     else:
-        resolver_specs = UDF_BASE_RESOLVER + _resolver_specs
-    resolver_specs += generate_udf_resolvers()
+        resolver_specs = UDF_BASE_RESOLVER + _resolver_specs[schema_name]
+    resolver_specs += generate_udf_resolvers(schema_name)
     return UdfSchema(
         resolver_specs,
         types,
@@ -208,5 +211,5 @@ def udf_schema(
         schema_based_automerge,
         segments,
         validators,
-        generate_udf_specs(other_udf_specs),
+        generate_udf_specs(other_udf_specs, schema_name),
     )