Update value_count serialization/deserialization to be consistent with original schema

merlin/schema/io/tensorflow_metadata.py

@@ -20,14 +20,7 @@
 import numpy
 
 from merlin.schema.io import proto_utils, schema_bp
-from merlin.schema.io.schema_bp import (
-    Feature,
-    FeatureType,
-    FixedShape,
-    FixedShapeDim,
-    FloatDomain,
-    IntDomain,
-)
+from merlin.schema.io.schema_bp import Feature, FeatureType, FloatDomain, IntDomain
 from merlin.schema.io.schema_bp import Schema as ProtoSchema
 from merlin.schema.io.schema_bp import ValueCount
 from merlin.schema.schema import ColumnSchema
@@ -282,17 +275,11 @@ def _pb_feature(column_schema):
 
     feature = _set_feature_domain(feature, column_schema)
 
-    if column_schema.is_list:
-        value_count = column_schema.properties.get("value_count", {})
-        min_length = value_count.get("min")
-        max_length = value_count.get("max")
-
-        if min_length and max_length and min_length == max_length:
-            feature.shape = FixedShape([FixedShapeDim(size=min_length)])
-        elif min_length and max_length and min_length < max_length:
-            feature.value_count = ValueCount(min=min_length, max=max_length)
-        else:
-            feature.value_count = ValueCount(min=0, max=0)
+    value_count = column_schema.properties.get("value_count", {})
+    if value_count:
+        min_length = value_count.get("min", 0)
+        max_length = value_count.get("max", 0)
+        feature.value_count = ValueCount(min=min_length, max=max_length)
 
     feature.annotation.tag = _pb_tag(column_schema)
     feature.annotation.extra_metadata.append(_pb_extra_metadata(column_schema))
@@ -342,8 +329,7 @@ def _merlin_domain(feature):
 def _merlin_value_count(feature):
     if proto_utils.has_field(feature, "value_count"):
         value_count = feature.value_count
-        if value_count.min != value_count.max != 0:
-            return {"min": value_count.min, "max": value_count.max}
+        return {"min": value_count.min, "max": value_count.max}
 
 
 def _merlin_properties(feature):
@@ -363,14 +349,17 @@ def _merlin_properties(feature):
         properties = {}
 
     domain = _merlin_domain(feature)
+
     if domain:
         properties["domain"] = domain
 
     value_count = _merlin_value_count(feature)
+
     if value_count:
         properties["value_count"] = value_count
-        properties["is_list"] = True
+        properties["is_list"] = value_count.get("min") > 0
         properties["is_ragged"] = value_count.get("min") != value_count.get("max")
+
     return properties
 
 

merlin/schema/schema.py

@@ -47,7 +47,7 @@ class ColumnSchema:
     tags: Optional[TagSet] = field(default_factory=TagSet)
     properties: Optional[Dict] = field(default_factory=dict)
     dtype: Optional[object] = None
-    is_list: bool = False
+    is_list: Optional[bool] = None
     is_ragged: Optional[bool] = None
 
     def __post_init__(self):
@@ -75,15 +75,35 @@ def __post_init__(self):
 
         object.__setattr__(self, "dtype", dtype)
 
+        value_count = self.properties.get("value_count")
+        value_count_provided = value_count is not None
+        value_count = value_count or {"min": 0, "max": 0}
+
+        if self.is_list is None:
+            if value_count["max"] > 0:
+                object.__setattr__(self, "is_list", True)
+            else:
+                object.__setattr__(self, "is_list", False)
+
         if self.is_ragged is None:
-            object.__setattr__(self, "is_ragged", self.is_list)
+            if value_count["max"] > value_count["min"]:
+                object.__setattr__(self, "is_ragged", True)
+            else:
+                object.__setattr__(self, "is_ragged", False)
 
         if self.is_ragged and not self.is_list:
             raise ValueError(
                 "`is_ragged` is set to `True` but `is_list` is not. "
                 "Only list columns can set the `is_ragged` flag."
             )
 
+        if self.is_ragged and value_count_provided and value_count["max"] == value_count["min"]:
+            raise ValueError(
+                "`is_ragged` is set to `True` but `value_count.min` == `value_count.max`. "
+                "If value_count min/max are equal. "
+                "This is a fixed size list and `is_ragged` should be set to False. "
+            )
+
     @property
     def quantity(self):
         """

tests/unit/schema/test_column_schemas.py

@@ -23,7 +23,7 @@
 
 @pytest.mark.parametrize("d_types", [numpy.float32, numpy.float64, numpy.uint32, numpy.uint64])
 def test_dtype_column_schema(d_types):
-    column = ColumnSchema("name", tags=[], properties=[], dtype=d_types)
+    column = ColumnSchema("name", tags=[], properties={}, dtype=d_types)
     assert column.dtype == d_types
 
 
@@ -180,8 +180,8 @@ def test_list_column_attributes():
     col2_schema = ColumnSchema("col2", is_list=True)
 
     assert col2_schema.is_list
-    assert col2_schema.is_ragged
-    assert col2_schema.quantity == ColumnQuantity.RAGGED_LIST
+    assert not col2_schema.is_ragged
+    assert col2_schema.quantity == ColumnQuantity.FIXED_LIST
 
     col3_schema = ColumnSchema("col3", is_list=True, is_ragged=True)
 
@@ -197,3 +197,11 @@ def test_list_column_attributes():
 
     with pytest.raises(ValueError):
         ColumnSchema("col5", is_list=False, is_ragged=True)
+
+
+def test_value_count():
+    with pytest.raises(ValueError) as exc_info:
+        ColumnSchema("col", is_ragged=True, properties={"value_count": {"min": 2, "max": 2}})
+    assert "`is_ragged` is set to `True` but `value_count.min` == `value_count.max`" in str(
+        exc_info.value
+    )

tests/unit/schema/test_schema_io.py

@@ -75,6 +75,33 @@ def test_merlin_to_proto_to_json_to_merlin():
     assert output_schema == schema
 
 
+@pytest.mark.parametrize(
+    ["value_count", "expected_is_list", "expected_is_ragged"],
+    [
+        [{"min": 0, "max": 0}, False, False],
+        [{"min": 1, "max": 1}, True, False],
+        [{"min": 1, "max": 2}, True, True],
+    ],
+)
+def test_value_count(value_count, expected_is_list, expected_is_ragged):
+    schema = Schema(
+        [
+            ColumnSchema(
+                "example",
+                properties={
+                    "value_count": value_count,
+                },
+            )
+        ]
+    )
+    assert schema["example"].is_list == expected_is_list
+    assert schema["example"].is_ragged == expected_is_ragged
+
+    json_schema = TensorflowMetadata.from_merlin_schema(schema).to_json()
+    output_schema = TensorflowMetadata.from_json(json_schema).to_merlin_schema()
+    assert output_schema == schema
+
+
 def test_column_schema_protobuf_domain_check(tmpdir):
     # create a schema
     schema1 = ColumnSchema(