Fix UDF deserialization in column profiles to support UDFs on pyspark (#1271)

## Description

udf metrics are multimetrics but have a different init signature which
leads to problems in deserializing them in pyspark. We need to have a
consistent construction that preserves the udfs submetrics on
deserialization, and works in the pyspark where a column profile is
deserialized from bytes.

## Changes

- support multimetric init shape in udf metric
- add tests to cover missing deserialization path

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

Author: jamie256

python/tests/api/logger/test_segments.py

@@ -2,6 +2,7 @@
 import pickle
 import tempfile
 from glob import glob
+from logging import getLogger
 from typing import Any
 
 import numpy as np
@@ -26,6 +27,8 @@
 from whylogs.core.view.dataset_profile_view import DatasetProfileView
 from whylogs.migration.converters import read_v0_to_view
 
+TEST_LOGGER = getLogger(__name__)
+
 
 def test_single_row_segment() -> None:
     segment_column = "col3"
@@ -282,8 +285,8 @@ def test_multi_column_segment_serialization_roundtrip_v0(tmp_path: Any) -> None:
     for file_path in paths:
         roundtrip_profiles.append(read_v0_to_view(file_path))
     assert len(roundtrip_profiles) == input_rows
-    print(roundtrip_profiles)
-    print(roundtrip_profiles[15])
+    TEST_LOGGER.info(roundtrip_profiles)
+    TEST_LOGGER.info(roundtrip_profiles[15])
 
     post_deserialization_view = roundtrip_profiles[15]
     assert post_deserialization_view is not None

python/tests/api/pyspark/experimental/test_profiler_function.py

@@ -174,6 +174,24 @@ def resolve(self, name, why_type, column_schema):
         assert profile_view.get_column("2").get_metric_names() == []
         assert profile_view.get_column("3").get_metric_names() == []
 
+    def test_collect_dataset_profile_view_with_udf_schema(self, input_df):
+        from whylogs.core.datatypes import Fractional
+        from whylogs.experimental.core.metrics.udf_metric import (
+            generate_udf_resolvers,
+            register_metric_udf,
+        )
+
+        def frob(x):
+            return x * x
+
+        register_metric_udf(col_type=Fractional, submetric_name="square", schema_name="pyspark_test")(frob)
+        test_schema = DeclarativeSchema(resolvers=generate_udf_resolvers(schema_name="pyspark_test"))
+        profile_view = collect_dataset_profile_view(input_df=input_df, schema=test_schema)
+
+        assert isinstance(profile_view, DatasetProfileView)
+        assert len(profile_view.get_columns()) > 0
+        assert "udf" in profile_view.get_column("0").get_metric_names()
+
     def test_map_vectors(self, embeddings_df):
         from pyspark.ml.functions import vector_to_array
 

python/tests/core/view/test_dataset_profile_view.py

@@ -233,10 +233,10 @@ def test_segmented_round_trip_metadata(tmp_path: str) -> None:
     segment_column = "A"
     results = why.log(df, schema=DatasetSchema(segments=segment_on_column(segment_column)))
     status = results.writer("local").write(dest=output_file)
-    TEST_LOGGER.info("serialized segmented profile to {output_file}" f" has status: {status}")
+    TEST_LOGGER.info(f"serialized segmented profile to {output_file} has status: {status}")
 
     view = DatasetProfileView.read(output_file)
-    TEST_LOGGER.info("round trip serialized and deserialized segmented profile" f" has metadata: {view._metadata}")
+    TEST_LOGGER.info(f"round trip serialized and deserialized segmented profile has metadata: {view._metadata}")
     assert view._metadata is not None
     segment_tag_metadata_key = _TAG_PREFIX + segment_column
     assert segment_tag_metadata_key in view._metadata

python/tests/experimental/core/metrics/test_udf_metric.py

@@ -1,3 +1,5 @@
+from logging import getLogger
+
 import pandas as pd
 
 import whylogs as why
@@ -10,6 +12,8 @@
     udf_metric_schema,
 )
 
+logger = getLogger(__name__)
+
 
 def test_udf_metric() -> None:
     config = UdfMetricConfig(
@@ -39,6 +43,37 @@ def test_udf_metric() -> None:
     assert "foo:frequent_items/frequent_strings" in summary
 
 
+def test_udf_metric_from_to_protobuf() -> None:
+    config = UdfMetricConfig(
+        udfs={
+            "fortytwo": lambda x: 42,
+            "foo": lambda x: "bar",
+        },
+    )
+    metric = UdfMetric.zero(config)
+    metric.columnar_update(PreprocessedColumn.apply([0]))
+    udf_metric_message = metric.to_protobuf()
+    logger.debug(f"serialized msg: {udf_metric_message}")
+    deserialied_udf_metric = UdfMetric.from_protobuf(udf_metric_message)
+    summary = deserialied_udf_metric.to_summary_dict()
+    logger.debug(f"serialized summary {summary}")
+    assert summary["fortytwo:counts/n"] == 1
+    assert summary["fortytwo:types/integral"] == 1
+    assert summary["fortytwo:types/string"] == 0
+    assert summary["fortytwo:cardinality/est"] == 1
+    assert summary["fortytwo:distribution/n"] == 1
+    assert summary["fortytwo:distribution/mean"] == 42
+    assert summary["fortytwo:ints/max"] == 42
+    assert summary["fortytwo:ints/min"] == 42
+    assert "fortytwo:frequent_items/frequent_strings" in summary
+
+    assert summary["foo:counts/n"] == 1
+    assert summary["foo:types/integral"] == 0
+    assert summary["foo:types/string"] == 1
+    assert summary["foo:cardinality/est"] == 1
+    assert "foo:frequent_items/frequent_strings" in summary
+
+
 def test_udf_throws() -> None:
     n = 0
 

python/tests/experimental/core/test_udf_schema.py

@@ -135,7 +135,9 @@ def test_multicolumn_udf_pandas() -> None:
     assert "counts/n" in sqr_summary
     div_summary = results.get_column("ratio").to_summary_dict()
     assert div_summary["distribution/n"] == 3
-    assert len(results.get_column("ratio").get_metrics()) == 2  # Integral -> counts plus registered distribution
+    # Integral -> counts plus registered distribution
+    assert results.get_column("ratio").get_metric("counts") is not None
+    assert results.get_column("ratio").get_metric("distribution") is not None
 
 
 def test_multicolumn_udf_row() -> None:
@@ -172,7 +174,9 @@ def test_multicolumn_udf_row() -> None:
     assert "counts/n" in sqr_summary
     div_summary = results.get_column("ratio").to_summary_dict()
     assert div_summary["distribution/n"] == 1
-    assert len(results.get_column("ratio").get_metrics()) == 2  # Integral -> counts plus registered distribution
+    # Integral -> counts plus registered distribution
+    assert results.get_column("ratio").get_metric("counts") is not None
+    assert results.get_column("ratio").get_metric("distribution") is not None
 
 
 n: int = 0

python/whylogs/core/view/column_profile_view.py

@@ -129,7 +129,6 @@ def from_protobuf(cls, msg: ColumnMessage) -> "ColumnProfileView":
 
             c_key = full_path[len(metric_name) + 1 :]
             metric_components[c_key] = c_msg
-
         for m_name, metric_components in metric_messages.items():
             m_enum = StandardMetric.__members__.get(m_name)
             if m_enum is None:
@@ -142,9 +141,11 @@ def from_protobuf(cls, msg: ColumnMessage) -> "ColumnProfileView":
 
             m_msg = MetricMessage(metric_components=metric_components)
             try:
-                result_metrics[m_name] = metric_class.from_protobuf(m_msg)
-            except:  # noqa
-                raise DeserializationError(f"Failed to deserialize metric: {m_name}")
+                deserialized_metric = metric_class.from_protobuf(m_msg)
+                result_metrics[m_name] = deserialized_metric
+            except Exception as error:  # noqa
+                raise DeserializationError(f"Failed to deserialize metric: {m_name}:{error}")
+
         return ColumnProfileView(metrics=result_metrics)
 
     @classmethod

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

@@ -103,14 +103,15 @@ class UdfMetric(MultiMetric):
 
     def __init__(
         self,
-        udfs: Dict[str, Callable[[Any], Any]],
+        submetrics: Dict[str, Dict[str, Metric]],
+        udfs: Optional[Dict[str, Callable[[Any], Any]]] = None,
         # discover these with resolver  submetrics: Dict[str, Dict[str, Metric]],  # feature name -> (namespace -> metric)
         submetric_schema: Optional[SubmetricSchema] = None,
         type_mapper: Optional[TypeMapper] = None,
         fi_disabled: bool = False,
     ):
-        super().__init__(dict())  # submetrics)
-        self._udfs = udfs
+        super().__init__(submetrics)
+        self._udfs = udfs or dict()
         self._submetric_schema = submetric_schema or default_schema()
         self._type_mapper = type_mapper or StandardTypeMapper()
         self._fi_disabled = fi_disabled
@@ -120,7 +121,7 @@ def namespace(self) -> str:
         return "udf"
 
     def merge(self, other: "UdfMetric") -> "UdfMetric":
-        merged = UdfMetric(self._udfs, self._submetric_schema, self._type_mapper, self._fi_disabled)
+        merged = UdfMetric(self.submetrics, self._udfs, self._submetric_schema, self._type_mapper, self._fi_disabled)
         merged.submetrics = self.merge_submetrics(other)
         return merged
 
@@ -166,10 +167,11 @@ def zero(cls, config: Optional[MetricConfig] = None) -> "UdfMetric":
             config = UdfMetricConfig()
 
         return UdfMetric(
-            config.udfs,
-            config.submetric_schema,
-            config.type_mapper,
-            config.fi_disabled,
+            dict(),
+            udfs=config.udfs,
+            submetric_schema=config.submetric_schema,
+            type_mapper=config.type_mapper,
+            fi_disabled=config.fi_disabled,
         )