Add serialization for distributions

Author: ghanse

dbldatagen/distributions/beta.py

@@ -35,6 +35,10 @@ def __init__(self, alpha=None, beta=None):
         self._alpha = alpha
         self._beta = beta
 
+    @classmethod
+    def getMapping(cls):
+        return {"alpha": "_alpha", "beta": "_beta"}
+
     @property
     def alpha(self):
         """ Return alpha parameter."""

dbldatagen/distributions/data_distribution.py

@@ -25,14 +25,20 @@
 import numpy as np
 import pyspark.sql.functions as F
 
+from ..serialization import Serializable
 
-class DataDistribution(ABC):
+
+class DataDistribution(Serializable, ABC):
     """ Base class for all distributions"""
 
     def __init__(self):
         self._rounding = False
         self._randomSeed = None
 
+    @classmethod
+    def getMapping(cls):
+        raise NotImplementedError("method not implemented")
+
     @staticmethod
     def get_np_random_generator(random_seed):
         """ Get numpy random number generator

dbldatagen/distributions/exponential_distribution.py

@@ -30,6 +30,10 @@ def __init__(self, rate=None):
         DataDistribution.__init__(self)
         self._rate = rate
 
+    @classmethod
+    def getMapping(cls):
+        return {"rate": "_rate"}
+
     def __str__(self):
         """ Return string representation"""
         return f"ExponentialDistribution(rate={self.rate}, randomSeed={self.randomSeed})"

dbldatagen/distributions/gamma.py

@@ -34,6 +34,10 @@ def __init__(self, shape, scale):
         self._shape = shape
         self._scale = scale
 
+    @classmethod
+    def getMapping(cls):
+        return {"shape": "_shape", "scale": "_scale"}
+
     @property
     def shape(self):
         """ Return shape parameter."""

dbldatagen/distributions/normal_distribution.py

@@ -27,6 +27,10 @@ def __init__(self, mean, stddev):
         self.mean = mean if mean is not None else 0.0
         self.stddev = stddev if stddev is not None else 1.0
 
+    @classmethod
+    def getMapping(cls):
+        return {"mean": "mean", "stddev": "stddev"}
+
     @staticmethod
     def normal_func(mean_series: pd.Series, std_dev_series: pd.Series, random_seed: pd.Series) -> pd.Series:
         """ Pandas / Numpy based function to generate normal / gaussian samples