preserve properties when inheriting schemas (#9435)

GitOrigin-RevId: 9e5f19481e5558041d2f09fa2891206927e78b33

Author: zxqfd555

CHANGELOG.md

@@ -7,6 +7,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ### Fixed
 - Endpoints created by `pw.io.http.rest_connector` now accept requests both with and without a trailing slash. For example, `/endpoint/` and `/endpoint` are now treated equivalently.
+- Schemas that inherit from other schemas now automatically preserve all properties from their parent schemas.
 
 ## [0.26.4] - 2025-10-16
 

python/pathway/internals/schema.py

@@ -178,13 +178,21 @@ def schema_add(*schemas: type[Schema]) -> type[Schema]:
 
 
 def _create_column_definitions(
-    schema: SchemaMetaclass, schema_properties: SchemaProperties
+    schema: SchemaMetaclass,
+    bases: tuple[type],
+    schema_properties: SchemaProperties,
 ) -> dict[str, ColumnSchema]:
     localns = locals()
     #  Update locals to handle recursive Schema definitions
     localns[schema.__name__] = schema
     annotations = get_type_hints(schema, localns=localns)
     fields = _cls_fields(schema)
+    for base in bases:
+        if not isinstance(base, SchemaMetaclass):
+            continue
+        for column_name, column_schema in base.__columns__.items():
+            if column_name not in fields:
+                fields[column_name] = column_schema.to_definition()
 
     columns = {}
 
@@ -274,15 +282,17 @@ class SchemaMetaclass(type):
     @trace.trace_user_frame
     def __init__(
         self,
-        *args,
+        name: str,
+        bases: tuple[type],
+        namespace: dict[str, Any],
         append_only: bool | None = None,
         id_dtype: dt.DType = dt.ANY_POINTER,
         id_append_only: bool | None = None,
         **kwargs,
     ) -> None:
-        super().__init__(*args, **kwargs)
+        super().__init__(name, bases, namespace, **kwargs)
         schema_properties = SchemaProperties(append_only=append_only)
-        self.__columns__ = _create_column_definitions(self, schema_properties)
+        self.__columns__ = _create_column_definitions(self, bases, schema_properties)
         pk_dtypes = [col.dtype for col in self.__columns__.values() if col.primary_key]
         if len(pk_dtypes) > 0:
             derived_type = dt.Pointer(*pk_dtypes)
@@ -588,6 +598,7 @@ def assert_matches_schema(
         allow_superset: bool = True,
         ignore_primary_keys: bool = True,
         allow_subtype: bool = True,
+        ignore_properties: bool = True,
     ) -> None:
         self_dict = self._dtypes()
         other_dict = other._dtypes()
@@ -619,6 +630,12 @@ def assert_matches_schema(
                 f"primary keys in the schemas do not match - they are {self.primary_key_columns()} in {self.__name__}",
                 f" and {other.primary_key_columns()} in {other.__name__}",
             )
+        if not ignore_properties:
+            self_columns = self.columns()
+            other_columns = other.columns()
+            for column_name, column_schema in self_columns.items():
+                other_column_schema = other_columns[column_name]
+                assert column_schema == other_column_schema
 
 
 def _schema_builder(

python/pathway/tests/test_schema.py

@@ -10,12 +10,13 @@
 from typing import Any
 
 import numpy as np
+import pandas as pd
 import pytest
 from dateutil import tz
 
 import pathway as pw
 from pathway.internals.schema import Schema
-from pathway.tests.utils import write_csv
+from pathway.tests.utils import write_csv, write_lines
 
 
 def assert_same_schema(left: type[Schema], right: type[Schema]):
@@ -433,3 +434,65 @@ class InputSchema(pw.Schema):
         DeserializedSchema.to_json_serializable_dict(), sort_keys=True
     )
     assert serialized_schema_json == roundtrip_schema_json
+
+
+@pytest.mark.parametrize("option", ["inheritance_1", "inheritance_2", "disjunction"])
+def test_schemas_composition(option, tmp_path):
+    input_path = tmp_path / "input.jsonl"
+
+    class BaseSchema_1(pw.Schema):
+        key: int = pw.column_definition(primary_key=True)
+        value: str = pw.column_definition(
+            default_value="default_value",
+            description="test description",
+            example="some example",
+        )
+
+    class BaseSchema_2(pw.Schema):
+        value2: bool
+
+    if option == "inheritance_1":
+
+        class InputSchema(BaseSchema_1):
+            value2: bool
+
+    elif option == "inheritance_2":
+
+        class InputSchema(BaseSchema_1, BaseSchema_2):
+            pass
+
+    elif option == "disjunction":
+        InputSchema = BaseSchema_1 | BaseSchema_2
+    else:
+        raise ValueError("unexpected option: {option}")
+
+    class ExpectedSchema(pw.Schema):
+        key: int = pw.column_definition(primary_key=True)
+        value: str = pw.column_definition(
+            default_value="default_value",
+            description="test description",
+            example="some example",
+        )
+        value2: bool
+
+    InputSchema.assert_matches_schema(
+        ExpectedSchema,
+        allow_subtype=False,
+        allow_superset=False,
+        ignore_primary_keys=False,
+        ignore_properties=False,
+    )
+
+    input_data = """
+        {"key": 0, "value": "hello", "value2": true}
+        {"key": 1, "value2": false}
+    """
+    write_lines(input_path, input_data)
+
+    table = pw.io.jsonlines.read(input_path, mode="static", schema=InputSchema)
+
+    output_table = pw.debug.table_to_pandas(table)
+    row_0 = output_table.loc[output_table["key"] == 0, ["value", "value2"]].iloc[0]
+    assert (row_0 == pd.Series({"value": "hello", "value2": True})).all()
+    row_1 = output_table.loc[output_table["key"] == 1, ["value", "value2"]].iloc[0]
+    assert (row_1 == pd.Series({"value": "default_value", "value2": False})).all()