Subclass column meta (#468)

* make it possible to subclass column meta * update * update documentation * remove extra notebook field in documentation
kaiko-ai · Aug 12, 2024 · f022d18 · f022d18
1 parent 32c4abf
commit f022d18
Show file tree

Hide file tree

Showing 21 changed files with 144 additions and 26 deletions.
diff --git a/docs/source/advanced.rst b/docs/source/advanced.rst
@@ -4,6 +4,7 @@ Advanced Topics
 .. toctree::
 
    subclassing_schemas
+   subclass_column_meta
    column_ambiguity
    advanced_linting_support
    dataset_implements
diff --git a/docs/source/advanced_linting_support.ipynb b/docs/source/advanced_linting_support.ipynb
@@ -187,7 +187,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/autocomplete_in_notebooks.ipynb b/docs/source/autocomplete_in_notebooks.ipynb
@@ -137,7 +137,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/column_ambiguity.ipynb b/docs/source/column_ambiguity.ipynb
@@ -292,7 +292,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/complex_datatypes.ipynb b/docs/source/complex_datatypes.ipynb
@@ -144,7 +144,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/create_empty_datasets.ipynb b/docs/source/create_empty_datasets.ipynb
@@ -229,7 +229,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/create_schema_in_notebook.ipynb b/docs/source/create_schema_in_notebook.ipynb
@@ -321,7 +321,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/dataset_implements.ipynb b/docs/source/dataset_implements.ipynb
@@ -186,7 +186,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/documentation.ipynb b/docs/source/documentation.ipynb
@@ -72,7 +72,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/loading_datasets_in_notebooks.ipynb b/docs/source/loading_datasets_in_notebooks.ipynb
@@ -309,7 +309,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 13,
    "id": "acfd55af",
    "metadata": {},
    "outputs": [],
@@ -327,7 +327,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 14,
    "id": "b0491ad1",
    "metadata": {},
    "outputs": [],
@@ -337,7 +337,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 15,
    "id": "f7f41731",
    "metadata": {},
    "outputs": [],
@@ -367,7 +367,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 16,
    "id": "264c2957",
    "metadata": {},
    "outputs": [
@@ -385,7 +385,7 @@
        "    age: Column[LongType]"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 16,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -408,7 +408,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 17,
    "id": "7199bee0",
    "metadata": {},
    "outputs": [
@@ -426,7 +426,7 @@
        "    age: Column[LongType]"
       ]
      },
-     "execution_count": 14,
+     "execution_count": 17,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -459,7 +459,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/schema_attributes.ipynb b/docs/source/schema_attributes.ipynb
@@ -98,7 +98,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/structtype_columns.ipynb b/docs/source/structtype_columns.ipynb
@@ -273,7 +273,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/structtypes_in_notebooks.ipynb b/docs/source/structtypes_in_notebooks.ipynb
@@ -191,7 +191,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/subclass_column_meta.ipynb b/docs/source/subclass_column_meta.ipynb
@@ -0,0 +1,85 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Defining your own ColumnMeta attributes\n",
+    "\n",
+    "In this notebook, we will see how to define your own `ColumnMeta` attributes. This is useful when you want to add some metadata to your columns that are not already defined in the `ColumnMeta` class."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'id': {'comment': 'Identifies the person', 'primary_key': True},\n",
+       " 'name': {},\n",
+       " 'age': {}}"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from dataclasses import dataclass\n",
+    "from typing import Annotated\n",
+    "from pyspark.sql.types import LongType, StringType\n",
+    "from typedspark import ColumnMeta, Schema\n",
+    "from typedspark._core.column import Column\n",
+    "\n",
+    "\n",
+    "@dataclass\n",
+    "class MyColumnMeta(ColumnMeta):\n",
+    "    primary_key: bool = False\n",
+    "\n",
+    "\n",
+    "class Persons(Schema):\n",
+    "    id: Annotated[\n",
+    "        Column[LongType],\n",
+    "        MyColumnMeta(\n",
+    "            comment=\"Identifies the person\",\n",
+    "            primary_key=True,\n",
+    "        ),\n",
+    "    ]\n",
+    "    name: Column[StringType]\n",
+    "    age: Column[LongType]\n",
+    "\n",
+    "\n",
+    "Persons.get_metadata()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "typedspark",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/docs/source/subclassing_schemas.ipynb b/docs/source/subclassing_schemas.ipynb
@@ -96,7 +96,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/transforming_datasets.ipynb b/docs/source/transforming_datasets.ipynb
@@ -213,7 +213,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/source/type_checking.ipynb b/docs/source/type_checking.ipynb
@@ -290,7 +290,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/docs/videos/ide.ipynb b/docs/videos/ide.ipynb
@@ -182,7 +182,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.2"
+   "version": "3.11.9"
   }
  },
  "nbformat": 4,

diff --git a/tests/_core/test_column.py b/tests/_core/test_column.py
@@ -1,10 +1,13 @@
+from dataclasses import dataclass
+from typing import Annotated
+
 import pandas as pd
 import pytest
 from pyspark.sql import SparkSession
 from pyspark.sql.functions import lit
 from pyspark.sql.types import LongType, StringType
 
-from typedspark import Column, Schema
+from typedspark import Column, ColumnMeta, Schema
 from typedspark._utils.create_dataset import create_partially_filled_dataset
 
 
@@ -41,3 +44,27 @@ def test_column_reference_without_spark_session():
 def test_column_with_deprecated_dataframe_param(spark: SparkSession):
     df = create_partially_filled_dataset(spark, A, {A.a: [1, 2, 3]})
     Column("a", dataframe=df)
+
+
+@dataclass
+class MyColumnMeta(ColumnMeta):
+    primary_key: bool = False
+
+
+class Persons(Schema):
+    id: Annotated[
+        Column[LongType],
+        MyColumnMeta(
+            comment="Identifies the person",
+            primary_key=True,
+        ),
+    ]
+    name: Column[StringType]
+    age: Column[LongType]
+
+
+def test_get_metadata():
+    assert Persons.get_metadata()["id"] == {
+        "comment": "Identifies the person",
+        "primary_key": True,
+    }
diff --git a/typedspark/_core/column_meta.py b/typedspark/_core/column_meta.py
@@ -1,6 +1,6 @@
 """Metadata for ``Column`` objects that can be accessed during runtime."""
 
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 from typing import Dict, Optional
 
 
@@ -23,4 +23,5 @@ class A(Schema):
 
     def get_metadata(self) -> Optional[Dict[str, str]]:
         """Returns the metadata of this column."""
-        return {"comment": self.comment} if self.comment else None
+        res = {k: v for k, v in asdict(self).items() if v is not None}
+        return res if len(res) > 0 else None
diff --git a/typedspark/_schema/schema.py b/typedspark/_schema/schema.py
@@ -198,6 +198,10 @@ def get_schema_name(cls):
         """Returns the name with which the schema was initialized."""
         return cls._original_name if cls._original_name else cls.__name__
 
+    def get_metadata(cls) -> dict[str, dict[str, Any]]:
+        """Returns the metadata of each of the columns in the schema."""
+        return {field.name: field.metadata for field in cls.get_structtype().fields}
+
 
 class Schema(Protocol, metaclass=MetaSchema):
     # pylint: disable=empty-docstring