Add json dump and load support for FeatureColumn (#2794)

* add json dump load support

* update vocabulary type

* update

* update

* update

Author: sneaxiy

python/runtime/feature/column.py

@@ -11,6 +11,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
+
 import six
 from runtime.feature.field_desc import DataType, FieldDesc
 
@@ -56,9 +58,10 @@ def to_dict(cls, feature_column):
         Returns:
             A Python dict which represents the FeatureColumn object.
         """
-        d = feature_column._to_dict()
-        d["type"] = type(feature_column).__name__
-        return d
+        return {
+            "type": type(feature_column).__name__,
+            "value": feature_column._to_dict(),
+        }
 
     def _to_dict(self):
         """
@@ -70,19 +73,25 @@ def _to_dict(self):
         raise NotImplementedError()
 
     @classmethod
-    def from_dict(cls, d):
+    def from_dict_or_feature_column(cls, obj):
         """
-        Create a FeatureColumn object from a Python dict. It can
-        be used to deserialize a FeatureColumn object from a JSON string.
+        If obj is of type dict, create a FeatureColumn object from a Python
+        dict. If obj is of type FeatureColumn, return itself. This method
+        can be used to deserialize a FeatureColumn object from a JSON string.
 
         Args:
-            d (dict): a Python dict object.
+            obj (dict|FeatureColumn): a Python dict or FeatureColumn object.
 
         Returns:
             A FeatureColumn object.
         """
-        typ = d.get("type")
-        return eval(typ)._from_dict(d)
+        if isinstance(obj, dict):
+            typ = obj.get("type")
+            return eval(typ)._from_dict(obj.get("value"))
+        elif isinstance(obj, FeatureColumn):
+            return obj
+        else:
+            raise TypeError("not supported type %s" % type(obj))
 
     @classmethod
     def _from_dict(self, d):
@@ -169,14 +178,14 @@ def num_class(self):
 
     def _to_dict(self):
         return {
-            "type": "BucketColumn",
             "source_column": FeatureColumn.to_dict(self.source_column),
             "boundaries": self.boundaries,
         }
 
     @classmethod
     def _from_dict(cls, d):
-        source_column = FeatureColumn.from_dict(d["source_column"])
+        source_column = FeatureColumn.from_dict_or_feature_column(
+            d["source_column"])
         boundaries = d["boundaries"]
         return BucketColumn(source_column, boundaries)
 
@@ -240,7 +249,6 @@ def num_class(self):
 
     def _to_dict(self):
         return {
-            "type": "CategoryHashColumn",
             "field_desc": self.field_desc.to_dict(),
             "bucket_size": self.bucket_size,
         }
@@ -343,7 +351,7 @@ def _from_dict(cls, d):
             if isinstance(k, six.string_types):
                 keys.append(k)
             else:
-                keys.append(FeatureColumn.from_dict(k))
+                keys.append(FeatureColumn.from_dict_or_feature_column(k))
 
         hash_bucket_size = d["hash_bucket_size"]
         return CrossColumn(keys, hash_bucket_size)
@@ -414,7 +422,8 @@ def _to_dict(self):
     def _from_dict(cls, d):
         category_column = d["category_column"]
         if category_column is not None:
-            category_column = FeatureColumn.from_dict(category_column)
+            category_column = FeatureColumn.from_dict_or_feature_column(
+                category_column)
 
         return EmbeddingColumn(category_column=category_column,
                                dimension=d["dimension"],
@@ -469,6 +478,75 @@ def _to_dict(self):
     def _from_dict(cls, d):
         category_column = d["category_column"]
         if category_column is not None:
-            category_column = FeatureColumn.from_dict(category_column)
+            category_column = FeatureColumn.from_dict_or_feature_column(
+                category_column)
 
         return IndicatorColumn(category_column=category_column, name=d["name"])
+
+
+class JSONEncoderWithFeatureColumn(json.JSONEncoder):
+    """
+    A helper class to serialize FeatureColumn objects to JSON string.
+    """
+    def default(self, obj):
+        """
+        Convert obj to an object that `json.dumps` accepts.
+        If obj is of type FeatureColumn, convert it to a Python
+        dict.
+
+        Args:
+            obj: any Python object.
+
+        Returns:
+            A Python object that `json.dumps` accepts.
+        """
+        if isinstance(obj, FeatureColumn):
+            return FeatureColumn.to_dict(obj)
+
+        # Use the default JSONEncoder if obj is not FeatureColumn
+        return json.JSONEncoder.default(self, obj)
+
+
+SUPPORTED_CONCRETE_FEATURE_COLUMNS = [
+    'NumericColumn',
+    'BucketColumn',
+    'CategoryIDColumn',
+    'CategoryHashColumn',
+    'SeqCategoryIDColumn',
+    'CrossColumn',
+    'EmbeddingColumn',
+    'IndicatorColumn',
+]
+
+
+def feature_column_json_hook(obj):
+    """
+    An object hook method that json.JSONDecoder accepts.
+    It is used to convert a Python dict to FeatureColumn object
+    if possible. See https://docs.python.org/3/library/json.html
+    for the usage of object hook.
+
+    Args:
+        obj: any Python object.
+
+    Returns:
+        If obj can be converted to a FeatureColumn object, convert
+        it. Otherwise, return itself.
+    """
+    if isinstance(obj, dict):
+        typ = obj.get("type")
+        if typ in SUPPORTED_CONCRETE_FEATURE_COLUMNS:
+            return FeatureColumn.from_dict_or_feature_column(obj)
+
+    return obj
+
+
+class JSONDecoderWithFeatureColumn(json.JSONDecoder):
+    """
+    A helper class to deserialize JSON string to FeatureColumn objects.
+    """
+    def __init__(self, *args, **kwargs):
+        # See here: https://docs.python.org/3/library/json.html
+        # for the usage of object_hook
+        kwargs['object_hook'] = feature_column_json_hook
+        super(JSONDecoderWithFeatureColumn, self).__init__(*args, **kwargs)

python/runtime/feature/column_test.py

@@ -29,11 +29,21 @@ def new_field_desc(self):
                             vocabulary=["a", "b", "c"])
         return desc
 
-    def check_from_dict(self, feature_column):
+    def check_serialize(self, feature_column):
         d = fc.FeatureColumn.to_dict(feature_column)
-        new_fc = fc.FeatureColumn.from_dict(d)
+        new_fc = fc.FeatureColumn.from_dict_or_feature_column(d)
+        new_d = fc.FeatureColumn.to_dict(new_fc)
+        typ = type(feature_column)
+        self.assertEqual(typ, type(new_fc))
+        self.assertEqual(typ.__name__, d["type"])
+        self.assertEqual(d, new_d)
+
+        dump_json = json.dumps(feature_column,
+                               cls=fc.JSONEncoderWithFeatureColumn)
+        new_fc = json.loads(dump_json, cls=fc.JSONDecoderWithFeatureColumn)
         new_d = fc.FeatureColumn.to_dict(new_fc)
         self.assertEqual(type(feature_column), type(new_fc))
+        self.assertEqual(typ, type(new_fc))
         self.assertEqual(d, new_d)
 
     def test_field_desc(self):
@@ -45,7 +55,8 @@ def test_field_desc(self):
         self.assertEqual(json_desc["format"], desc.format)
         self.assertEqual(json_desc["shape"], desc.shape)
         self.assertEqual(json_desc["is_sparse"], desc.is_sparse)
-        self.assertEqual(json_desc["vocabulary"], desc.vocabulary)
+        vocab = set(json_desc["vocabulary"])
+        self.assertEqual(vocab, desc.vocabulary)
         self.assertEqual(json_desc["max_id"], desc.max_id)
 
     def test_feature_column_subclass(self):
@@ -73,13 +84,13 @@ def test_numeric_column(self):
 
         d1 = fc.FeatureColumn.to_dict(nc1)
         self.assertEqual(d1["type"], "NumericColumn")
-        self.assertEqual(d1["field_desc"], desc1.to_dict())
-        self.check_from_dict(nc1)
+        self.assertEqual(d1["value"]["field_desc"], desc1.to_dict())
+        self.check_serialize(nc1)
 
         d2 = fc.FeatureColumn.to_dict(nc2)
         self.assertEqual(d2["type"], "NumericColumn")
-        self.assertEqual(d2["field_desc"], desc2.to_dict())
-        self.check_from_dict(nc2)
+        self.assertEqual(d2["value"]["field_desc"], desc2.to_dict())
+        self.check_serialize(nc2)
 
     def test_bucket_column(self):
         desc = self.new_field_desc()
@@ -93,10 +104,11 @@ def test_bucket_column(self):
         self.assertEqual(bc.get_field_desc()[0].to_json(), desc.to_json())
         d = fc.FeatureColumn.to_dict(bc)
         self.assertEqual(d["type"], "BucketColumn")
-        self.assertEqual(d["boundaries"], boundaries)
-        self.assertEqual(d["source_column"]["type"], "NumericColumn")
-        self.assertEqual(d["source_column"]["field_desc"], desc.to_dict())
-        self.check_from_dict(bc)
+        self.assertEqual(d["value"]["boundaries"], boundaries)
+        self.assertEqual(d["value"]["source_column"]["type"], "NumericColumn")
+        self.assertEqual(d["value"]["source_column"]["value"]["field_desc"],
+                         desc.to_dict())
+        self.check_serialize(bc)
 
         bc = bc.new_feature_column_from(desc)
         self.assertTrue(isinstance(bc, fc.BucketColumn))
@@ -106,10 +118,11 @@ def test_bucket_column(self):
         self.assertEqual(bc.get_field_desc()[0].to_json(), desc.to_json())
         d = fc.FeatureColumn.to_dict(bc)
         self.assertEqual(d["type"], "BucketColumn")
-        self.assertEqual(d["boundaries"], boundaries)
-        self.assertEqual(d["source_column"]["type"], "NumericColumn")
-        self.assertEqual(d["source_column"]["field_desc"], desc.to_dict())
-        self.check_from_dict(bc)
+        self.assertEqual(d["value"]["boundaries"], boundaries)
+        self.assertEqual(d["value"]["source_column"]["type"], "NumericColumn")
+        self.assertEqual(d["value"]["source_column"]["value"]["field_desc"],
+                         desc.to_dict())
+        self.check_serialize(bc)
 
     def test_category_column(self):
         desc = self.new_field_desc()
@@ -126,9 +139,9 @@ def test_category_column(self):
 
             d = fc.FeatureColumn.to_dict(cc)
             self.assertEqual(d["type"], fc_class.__name__)
-            self.assertEqual(d["field_desc"], desc.to_dict())
-            self.assertEqual(d["bucket_size"], bucket_size)
-            self.check_from_dict(cc)
+            self.assertEqual(d["value"]["field_desc"], desc.to_dict())
+            self.assertEqual(d["value"]["bucket_size"], bucket_size)
+            self.check_serialize(cc)
 
             cc = cc.new_feature_column_from(desc)
             self.assertTrue(isinstance(cc, fc_class))
@@ -138,9 +151,9 @@ def test_category_column(self):
 
             d = fc.FeatureColumn.to_dict(cc)
             self.assertEqual(d["type"], fc_class.__name__)
-            self.assertEqual(d["field_desc"], desc.to_dict())
-            self.assertEqual(d["bucket_size"], bucket_size)
-            self.check_from_dict(cc)
+            self.assertEqual(d["value"]["field_desc"], desc.to_dict())
+            self.assertEqual(d["value"]["bucket_size"], bucket_size)
+            self.check_serialize(cc)
 
     def test_cross_column(self):
         desc = self.new_field_desc()
@@ -155,13 +168,13 @@ def test_cross_column(self):
 
         d = fc.FeatureColumn.to_dict(cc)
         self.assertEqual(d["type"], "CrossColumn")
-        keys = d["keys"]
+        keys = d["value"]["keys"]
         self.assertEqual(len(keys), 2)
         self.assertEqual(keys[0]["type"], "NumericColumn")
-        self.assertEqual(keys[0]["field_desc"], desc.to_dict())
+        self.assertEqual(keys[0]["value"]["field_desc"], desc.to_dict())
         self.assertEqual(keys[1], "cross_feature_2")
-        self.assertEqual(d["hash_bucket_size"], hash_bucket_size)
-        self.check_from_dict(cc)
+        self.assertEqual(d["value"]["hash_bucket_size"], hash_bucket_size)
+        self.check_serialize(cc)
 
     def test_embedding_and_indicator_column(self):
         desc = self.new_field_desc()
@@ -179,13 +192,15 @@ def test_embedding_and_indicator_column(self):
 
             d = fc.FeatureColumn.to_dict(fc1)
             self.assertEqual(d["type"], fc_class.__name__)
-            self.assertEqual(d["name"], "")
-            self.assertEqual(d["category_column"]["type"],
+            self.assertEqual(d["value"]["name"], "")
+            self.assertEqual(d["value"]["category_column"]["type"],
                              "CategoryHashColumn")
-            self.assertEqual(d["category_column"]["field_desc"],
-                             desc.to_dict())
-            self.assertEqual(d["category_column"]["bucket_size"], 4096)
-            self.check_from_dict(fc1)
+            self.assertEqual(
+                d["value"]["category_column"]["value"]["field_desc"],
+                desc.to_dict())
+            self.assertEqual(
+                d["value"]["category_column"]["value"]["bucket_size"], 4096)
+            self.check_serialize(fc1)
 
             fc2 = fc_class(category_column=None, name="my_category_column")
             fc2_descs = fc2.get_field_desc()
@@ -197,9 +212,9 @@ def test_embedding_and_indicator_column(self):
 
             d = fc.FeatureColumn.to_dict(fc2)
             self.assertEqual(d["type"], fc_class.__name__)
-            self.assertEqual(d["name"], "my_category_column")
-            self.assertEqual(d["category_column"], None)
-            self.check_from_dict(fc2)
+            self.assertEqual(d["value"]["name"], "my_category_column")
+            self.assertEqual(d["value"]["category_column"], None)
+            self.check_serialize(fc2)
 
 
 if __name__ == '__main__':

python/runtime/feature/derivation_test.py

@@ -11,8 +11,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import unittest
 
+import runtime.feature.column as fc
 import runtime.feature.derivation as fd
 import runtime.testing as testing
 from runtime.feature.column import (CategoryIDColumn, CrossColumn,
@@ -82,6 +84,14 @@ def test_infer_index(self):
 @unittest.skipUnless(testing.get_driver() in ["mysql", "hive"],
                      "skip non MySQL and Hive tests")
 class TestFeatureDerivationWithMockedFeatures(unittest.TestCase):
+    def check_json_dump(self, features):
+        dump_json = json.dumps(features, cls=fc.JSONEncoderWithFeatureColumn)
+        new_features = json.loads(dump_json,
+                                  cls=fc.JSONDecoderWithFeatureColumn)
+        new_dump_json = json.dumps(new_features,
+                                   cls=fc.JSONEncoderWithFeatureColumn)
+        self.assertEqual(dump_json, new_dump_json)
+
     def test_without_cross(self):
         features = {
             'feature_columns': [
@@ -108,6 +118,9 @@ def test_without_cross(self):
         features, label = fd.infer_feature_columns(conn, select, features,
                                                    label)
 
+        self.check_json_dump(features)
+        self.check_json_dump(label)
+
         self.assertEqual(len(features), 1)
         self.assertTrue("feature_columns" in features)
         features = features["feature_columns"]
@@ -230,6 +243,9 @@ def test_with_cross(self):
         features, label = fd.infer_feature_columns(conn, select, features,
                                                    label)
 
+        self.check_json_dump(features)
+        self.check_json_dump(label)
+
         self.assertEqual(len(features), 1)
         self.assertTrue("feature_columns" in features)
         features = features["feature_columns"]
@@ -322,6 +338,9 @@ def test_no_column_clause(self):
         features, label = fd.infer_feature_columns(conn, select, features,
                                                    label)
 
+        self.check_json_dump(features)
+        self.check_json_dump(label)
+
         self.assertEqual(len(features), 1)
         self.assertTrue("feature_columns" in features)
         features = features["feature_columns"]