Update with_properties to enable changing existing properties on ColumnSchema (#157)

* Update `ColumnSchema.with_properties` to change properties correctly.

* Add tests for `ColumnSchema` equality and update methods `with_*`

Author: oliverholworthy

merlin/schema/schema.py

@@ -172,12 +172,12 @@ def with_properties(self, properties: dict) -> "ColumnSchema":
             raise TypeError("properties must be in dict format, key: value")
 
         # Using new dictionary to avoid passing old ref to new schema
-        properties.update(self.properties)
+        new_properties = {**self.properties, **properties}
 
         return ColumnSchema(
             self.name,
             tags=self.tags,
-            properties=properties,
+            properties=new_properties,
             dtype=self.dtype,
             is_list=self.is_list,
             is_ragged=self.is_ragged,

tests/unit/schema/test_column_schemas.py

@@ -27,25 +27,136 @@ def test_dtype_column_schema(d_types):
     assert column.dtype == d_types
 
 
-def test_column_schema_meta():
-    column = ColumnSchema("name", tags=["tag-1"], properties={"p1": "prop-1"})
-
-    assert column.name == "name"
-    assert "tag-1" in column.tags
-    assert column.with_name("a").name == "a"
-    assert set(column.with_tags("tag-2").tags) == set(["tag-1", "tag-2"])
-    assert column.with_properties({"p2": "prop-2"}).properties == {
-        "p1": "prop-1",
-        "p2": "prop-2",
-    }
-    assert column.with_tags("tag-2").properties == {"p1": "prop-1"}
-    assert set(column.with_properties({"p2": "prop-2"}).tags) == set(["tag-1"])
-
-    assert column == ColumnSchema("name", tags=["tag-1"], properties={"p1": "prop-1"})
-    # should not be the same no properties
-    assert column != ColumnSchema("name", tags=["tag-1"])
-    # should not be the same no tags
-    assert column != ColumnSchema("name", properties={"p1": "prop-1"})
+@pytest.mark.parametrize(
+    ["column_schema_a", "column_schema_b"],
+    [
+        [ColumnSchema("col"), ColumnSchema("col")],
+        [ColumnSchema("col_b", tags=["tag-1"]), ColumnSchema("col_b", tags=["tag-1"])],
+        [
+            ColumnSchema("col", dtype=numpy.int32, properties={"domain": {"min": 0, "max": 8}}),
+            ColumnSchema("col", dtype=numpy.int32, properties={"domain": {"min": 0, "max": 8}}),
+        ],
+        [
+            ColumnSchema(
+                "col",
+                dtype=numpy.float32,
+                tags=["tag-2", Tags.CONTINUOUS],
+                properties={"p1": "prop-1"},
+            ),
+            ColumnSchema(
+                "col",
+                dtype=numpy.float32,
+                tags=["tag-2", Tags.CONTINUOUS],
+                properties={"p1": "prop-1"},
+            ),
+        ],
+    ],
+)
+def test_equal(column_schema_a, column_schema_b):
+    assert column_schema_a == column_schema_b
+    assert column_schema_a.name == column_schema_b.name
+    assert column_schema_a.dtype == column_schema_b.dtype
+    assert column_schema_a.tags == column_schema_b.tags
+    assert column_schema_a.properties == column_schema_b.properties
+
+
+@pytest.mark.parametrize(
+    ["column_schema_a", "column_schema_b"],
+    [
+        [ColumnSchema("col_a"), ColumnSchema("col_b")],
+        [ColumnSchema("name"), ColumnSchema("name", tags=["tags-1"])],
+        [ColumnSchema("name"), ColumnSchema("name", properties={"p1": "prop-1"})],
+        [
+            ColumnSchema("name", tags=["tag-1"]),
+            ColumnSchema("name", properties={"p1": "prop-1"}),
+        ],
+        [
+            ColumnSchema("name", tags=["tag-1"], properties={"p1": "prop-1"}),
+            ColumnSchema("name", properties={"p1": "prop-1"}),
+        ],
+    ],
+)
+def test_not_equal(column_schema_a, column_schema_b):
+    assert column_schema_a != column_schema_b
+
+
+@pytest.mark.parametrize(
+    ["column_schema", "name", "expected_column_schema"],
+    [
+        [ColumnSchema("col_a"), "col_b", ColumnSchema("col_b")],
+        [ColumnSchema("feat", tags=["tag-1"]), "seq", ColumnSchema("seq", tags=["tag-1"])],
+        [
+            ColumnSchema(
+                "feat",
+                tags=["tag-1"],
+                dtype=numpy.float32,
+                properties={"domain": {"min": 0.0, "max": 6.0}},
+            ),
+            "feat_b",
+            ColumnSchema(
+                "feat_b",
+                tags=["tag-1"],
+                dtype=numpy.float32,
+                properties={"domain": {"min": 0.0, "max": 6.0}},
+            ),
+        ],
+    ],
+)
+def test_with_name(column_schema, name, expected_column_schema):
+    assert column_schema.with_name(name) == expected_column_schema
+
+
+@pytest.mark.parametrize(
+    ["column_schema", "tags", "expected_column_schema"],
+    [
+        [
+            ColumnSchema("example", tags=["tag-1"], properties={"p1": "prop-1"}),
+            "tag-2",
+            ColumnSchema("example", tags=["tag-1", "tag-2"], properties={"p1": "prop-1"}),
+        ],
+        [
+            ColumnSchema("example", tags=["tag-1"], dtype=numpy.float32),
+            ["tag-2", Tags.CONTINUOUS],
+            ColumnSchema("example", tags=["tag-1", "tag-2", Tags.CONTINUOUS], dtype=numpy.float32),
+        ],
+    ],
+)
+def test_with_tags(column_schema, tags, expected_column_schema):
+    assert column_schema.with_tags(tags) == expected_column_schema
+
+
+@pytest.mark.parametrize(
+    ["column_schema", "properties", "expected_column_schema"],
+    [
+        [
+            ColumnSchema("example", properties={"a": "old"}),
+            {"a": "new"},
+            ColumnSchema("example", properties={"a": "new"}),
+        ],
+        [
+            ColumnSchema("example", properties={"a": 1, "b": 2}),
+            {"a": 4, "c": 3},
+            ColumnSchema("example", properties={"a": 4, "b": 2, "c": 3}),
+        ],
+        [
+            ColumnSchema(
+                "example_col_2",
+                dtype=numpy.float32,
+                tags=[Tags.CONTINUOUS],
+                properties={"a": 1, "domain": {"min": 0, "max": 5}},
+            ),
+            {"a": 4, "c": 3, "domain": {"max": 8}},
+            ColumnSchema(
+                "example_col_2",
+                dtype=numpy.float32,
+                tags=[Tags.CONTINUOUS],
+                properties={"a": 4, "c": 3, "domain": {"max": 8}},
+            ),
+        ],
+    ],
+)
+def test_with_properties(column_schema, properties, expected_column_schema):
+    assert column_schema.with_properties(properties) == expected_column_schema
 
 
 def test_column_schema_tags_normalize():