Add serialization for constraints

Author: ghanse

dbldatagen/constraints/chained_relation.py

@@ -38,6 +38,10 @@ def __init__(self, columns, relation):
         if not isinstance(self._columns, list) or len(self._columns) <= 1:
             raise ValueError("ChainedRelation constraints must be defined across more than one column")
 
+    @classmethod
+    def getMapping(cls):
+        return {"relation": "_relation", "columns": "_columns"}
+
     def _generateFilterExpression(self):
         """ Generated composite filter expression for chained set of filter expressions
 

dbldatagen/constraints/constraint.py

@@ -8,9 +8,10 @@
 import types
 from abc import ABC, abstractmethod
 from pyspark.sql import Column
+from ..serialization import Serializable
 
 
-class Constraint(ABC):
+class Constraint(Serializable, ABC):
     """ Constraint object - base class for predefined and custom constraints
 
     This class is meant for internal use only.

dbldatagen/constraints/literal_range_constraint.py

@@ -29,6 +29,10 @@ def __init__(self, columns, lowValue, highValue, strict=False):
         self._highValue = highValue
         self._strict = strict
 
+    @classmethod
+    def getMapping(cls):
+        return {"columns": "_columns", "lowValue": "_lowValue", "highValue": "_highValue", "strict": "_strict"}
+
     def _generateFilterExpression(self):
         """ Generate a SQL filter expression that may be used for filtering"""
         expressions = [F.col(colname) for colname in self._columns]

dbldatagen/constraints/literal_relation_constraint.py

@@ -29,6 +29,10 @@ def __init__(self, columns, relation, value):
         if relation not in self.SUPPORTED_OPERATORS:
             raise ValueError(f"Parameter `relation` should be one of the operators :{self.SUPPORTED_OPERATORS}")
 
+    @classmethod
+    def getMapping(cls):
+        return {"columns": "_columns", "relation": "_relation", "value": "_value"}
+
     def _generateFilterExpression(self):
         expressions = [F.col(colname) for colname in self._columns]
         literalValue = F.lit(self._value)

dbldatagen/constraints/negative_values.py

@@ -27,6 +27,10 @@ def __init__(self, columns, strict=False):
         self._columns = self._columnsFromListOrString(columns)
         self._strict = strict
 
+    @classmethod
+    def getMapping(cls):
+        return {"columns": "_columns", "strict": "_strict"}
+
     def _generateFilterExpression(self):
         expressions = [F.col(colname) for colname in self._columns]
         if self._strict:

dbldatagen/constraints/positive_values.py

@@ -27,6 +27,10 @@ def __init__(self, columns, strict=False):
         self._columns = self._columnsFromListOrString(columns)
         self._strict = strict
 
+    @classmethod
+    def getMapping(cls):
+        return {"columns": "_columns", "strict": "_strict"}
+
     def _generateFilterExpression(self):
         """ Generate a filter expression that may be used for filtering"""
         expressions = [F.col(colname) for colname in self._columns]

dbldatagen/constraints/ranged_values_constraint.py

@@ -28,6 +28,10 @@ def __init__(self, columns, lowValue, highValue, strict=False):
         self._highValue = highValue
         self._strict = strict
 
+    @classmethod
+    def getMapping(cls):
+        return {"columns": "_columns", "lowValue": "_lowValue", "highValue": "_highValue", "strict": "_strict"}
+
     def _generateFilterExpression(self):
         """ Generate a SQL filter expression that may be used for filtering"""
         expressions = [F.col(colname) for colname in self._columns]

dbldatagen/constraints/sql_expr.py

@@ -25,6 +25,10 @@ def __init__(self, expr: str):
         assert isinstance(expr, str) and len(expr.strip()) > 0, "Expression must be a valid SQL string"
         self._expr = expr
 
+    @classmethod
+    def getMapping(cls):
+        return {"expr": "_expr"}
+
     def _generateFilterExpression(self):
         """ Generate a SQL filter expression that may be used for filtering"""
         return F.expr(self._expr)

dbldatagen/constraints/unique_combinations.py

@@ -45,6 +45,10 @@ def __init__(self, columns=None):
         else:
             self._columns = None
 
+    @classmethod
+    def getMapping(cls):
+        return {"columns": "_columns"}
+
     def prepareDataGenerator(self, dataGenerator):
         """ Prepare the data generator to generate data that matches the constraint