Repository: unionai-oss/pandera

Issue:
add `select_columns` method to DataFrameSchema class
this method should take a list of column names as input, and output a copy of the DataFrameSchema that only contains the columns specified in the input.


Discussion:

Resolving patch:
diff --git a/pandera/schemas.py b/pandera/schemas.py
--- a/pandera/schemas.py
+++ b/pandera/schemas.py
@@ -570,6 +570,41 @@ def update_column(self, column_name: str, **kwargs) -> "DataFrameSchema":
         schema_copy.columns.update({column_name: new_column})
         return schema_copy
 
+    def rename_columns(self, rename_dict: dict):
+        """Rename columns using a dictionary of key-value pairs.
+
+        :param rename_dict: dictionary of 'old_name': 'new_name' key-value
+            pairs.
+        :returns: dataframe schema (copy of original)
+        """
+
+        # We iterate over the existing columns dict and replace those keys
+        # that exist in the rename_dict
+        new_schema = copy.deepcopy(self)
+        new_columns = {
+            (
+                rename_dict[col_name]if col_name in rename_dict else col_name
+            ): col_attrs
+            for col_name, col_attrs in self.columns.items()
+        }
+
+        new_schema.columns = new_columns
+        return new_schema
+
+    def select_columns(self, columns: list):
+        """Select subset of columns in the schema.
+
+        :param columns: list of column names to select.
+        :returns: dataframe schema (copy of original)
+        """
+        new_schema = copy.deepcopy(self)
+        new_columns = {
+            col_name: column for col_name, column in self.columns.items()
+            if col_name in columns
+        }
+        new_schema.columns = new_columns
+        return new_schema
+
     @classmethod
     def from_yaml(cls, yaml_schema) -> "DataFrameSchema":
         """Create DataFrameSchema from yaml file.
@@ -591,25 +626,6 @@ def to_yaml(self, fp: Union[str, Path] = None):
         import pandera.io  # pylint: disable-all
         return pandera.io.to_yaml(self, fp)
 
-    def rename_columns(self, rename_dict: dict):
-        """Rename columns using a dictionary of key value pairs 
-
-        :param rename_dict: Dictionary of 'old_name':'new_name' key-value pairs.
-        :returns: dataframe schema (copy of original)
-        """
-
-        # We iterate over the existing columns dict and replace those keys
-        # that exist in the rename_dict
-        new_schema = copy.deepcopy(self)
-        new_columns = {
-            (rename_dict[col_name] if col_name in rename_dict else col_name): col_attrs
-            for col_name, col_attrs in self.columns.items()
-        }
-
-        new_schema.columns = new_columns
-
-        return new_schema
-
 
 class SeriesSchemaBase():
     """Base series validator object."""


Test patch:
diff --git a/tests/test_schema_components.py b/tests/test_schema_components.py
--- a/tests/test_schema_components.py
+++ b/tests/test_schema_components.py
@@ -439,29 +439,6 @@ def test_non_str_column_name_regex(column_key):
         )
 
 
-def test_rename_columns():
-    """Check that DataFrameSchema.rename_columns() method does it's job"""
-
-    rename_dict = {
-        'col1': 'col1_new_name',
-        'col2': 'col2_new_name'
-    }
-
-    schema_original = DataFrameSchema(
-        columns={
-            'col1': Column(Int),
-            'col2': Column(Float)
-        }
-    )
-
-    schema_renamed = schema_original.rename_columns(rename_dict)
-
-    # Check if new column names are indeed present in the new schema
-    assert all([col_name in rename_dict.values() for col_name in schema_renamed.columns])
-    # Check if original schema didn't change in the process
-    assert all([col_name in schema_original.columns for col_name in rename_dict])
-
-
 def test_column_type_can_be_set():
     """Test that the Column dtype can be edited during schema construction."""
 
diff --git a/tests/test_schemas.py b/tests/test_schemas.py
--- a/tests/test_schemas.py
+++ b/tests/test_schemas.py
@@ -669,6 +669,60 @@ def test_dataframe_schema_update_column(
     assertion_fn(schema, new_schema)
 
 
+def test_rename_columns():
+    """Check that DataFrameSchema.rename_columns() method does it's job"""
+
+    rename_dict = {"col1": "col1_new_name", "col2": "col2_new_name"}
+    schema_original = DataFrameSchema(
+        columns={"col1": Column(Int), "col2": Column(Float)}
+    )
+
+    schema_renamed = schema_original.rename_columns(rename_dict)
+
+    # Check if new column names are indeed present in the new schema
+    assert all(
+        [
+            col_name in rename_dict.values()
+            for col_name in schema_renamed.columns
+        ]
+    )
+    # Check if original schema didn't change in the process
+    assert all(
+        [col_name in schema_original.columns for col_name in rename_dict]
+    )
+
+
+@pytest.mark.parametrize(
+    "select_columns, schema", [
+        (
+            ["col1", "col2"], DataFrameSchema(
+                columns={
+                    "col1": Column(Int),
+                    "col2": Column(Int),
+                    "col3": Column(Int),
+                }
+            )
+        ),
+        (
+            [("col1", "col1b"), ("col2", "col2b")], DataFrameSchema(
+                columns={
+                    ("col1", "col1a"): Column(Int),
+                    ("col1", "col1b"): Column(Int),
+                    ("col2", "col2a"): Column(Int),
+                    ("col2", "col2b"): Column(Int),
+                }
+            )
+        )
+    ]
+)
+def test_select_columns(select_columns, schema):
+    """Check that select_columns method correctly creates new subset schema."""
+    original_columns = list(schema.columns)
+    schema_selected = schema.select_columns(select_columns)
+    assert all(x in select_columns for x in schema_selected.columns)
+    assert all(x in original_columns for x in schema.columns)
+
+
 def test_lazy_dataframe_validation_error():
     """Test exceptions on lazy dataframe validation."""
     schema = DataFrameSchema(