polarfrost 0.1.0__tar.gz

polarfrost-0.1.0/PKG-INFO

@@ -0,0 +1,86 @@
+Metadata-Version: 2.4
+Name: polarfrost
+Version: 0.1.0
+Summary: A fast k-anonymity implementation using Polars and PySpark
+Home-page: https://github.com/rglew/polarfrost
+Author: Richard Glew
+Author-email: richard.glew@hotmail.com
+Keywords: anonymization,privacy,polars,k-anonymity,data-privacy
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Security
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: polars>=0.13.0
+Requires-Dist: pandas>=1.3.0
+Requires-Dist: numpy>=1.21.0
+Provides-Extra: spark
+Requires-Dist: pyspark>=3.0.0; extra == "spark"
+Provides-Extra: dev
+Requires-Dist: pytest>=6.0; extra == "dev"
+Requires-Dist: pytest-cov>=2.0; extra == "dev"
+Requires-Dist: black>=21.0; extra == "dev"
+Requires-Dist: isort>=5.0; extra == "dev"
+Requires-Dist: mypy>=0.900; extra == "dev"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Polarfrost
+
+A fast k-anonymity implementation using Polars, featuring both Mondrian and Clustering algorithms for efficient privacy-preserving data analysis.
+
+## Features
+
+- 🚀 Blazing fast k-anonymity using Polars
+- 🧊 Supports both local (Polars) and distributed (PySpark) processing
+- 📊 Preserves data utility while ensuring privacy
+- 🐍 Simple Python API
+
+## Installation
+
+```bash
+pip install polarfrost
+```
+
+## Quick Start
+
+```python
+import polars as pl
+from polarfrost import mondrian_k_anonymity
+
+# Load your data
+df = pl.read_csv("your_data.csv")
+
+# Apply k-anonymity
+anonymized = mondrian_k_anonymity(
+    df,
+    quasi_identifiers=["age", "gender", "zipcode"],
+    sensitive_column="income",
+    k=3,
+    categorical=["gender", "zipcode"]
+)
+
+print(anonymized)
+```
+
+## License
+
+MIT

polarfrost-0.1.0/README.md

@@ -0,0 +1,41 @@
+# Polarfrost
+
+A fast k-anonymity implementation using Polars, featuring both Mondrian and Clustering algorithms for efficient privacy-preserving data analysis.
+
+## Features
+
+- 🚀 Blazing fast k-anonymity using Polars
+- 🧊 Supports both local (Polars) and distributed (PySpark) processing
+- 📊 Preserves data utility while ensuring privacy
+- 🐍 Simple Python API
+
+## Installation
+
+```bash
+pip install polarfrost
+```
+
+## Quick Start
+
+```python
+import polars as pl
+from polarfrost import mondrian_k_anonymity
+
+# Load your data
+df = pl.read_csv("your_data.csv")
+
+# Apply k-anonymity
+anonymized = mondrian_k_anonymity(
+    df,
+    quasi_identifiers=["age", "gender", "zipcode"],
+    sensitive_column="income",
+    k=3,
+    categorical=["gender", "zipcode"]
+)
+
+print(anonymized)
+```
+
+## License
+
+MIT

polarfrost-0.1.0/polarfrost/__init__.py

@@ -0,0 +1,26 @@
+"""
+PolarFrost: Fast k-anonymity implementation using Polars and PySpark.
+
+This package provides efficient implementations of k-anonymity algorithms,
+including the Mondrian algorithm, with support for both local (Polars)
+and distributed (PySpark) processing.
+"""
+
+__version__ = "0.1.0"
+
+# Import main functions
+try:
+    from .mondrian import (
+        mondrian_k_anonymity,
+        mondrian_k_anonymity_polars,
+        mondrian_k_anonymity_spark
+    )
+    __all__ = [
+        'mondrian_k_anonymity',
+        'mondrian_k_anonymity_polars',
+        'mondrian_k_anonymity_spark'
+    ]
+except ImportError as e:
+    import warnings
+    warnings.warn(f"Could not import mondrian: {e}")
+    __all__ = []

polarfrost-0.1.0/polarfrost/clustering.py

@@ -0,0 +1,29 @@
+"""
+Clustering-based k-anonymity implementation using Polars.
+"""
+from typing import List, Optional, Union
+import polars as pl
+
+def clustering_k_anonymity(
+    df: Union[pl.DataFrame, pl.LazyFrame],
+    quasi_identifiers: List[str],
+    sensitive_column: str,
+    k: int,
+    categorical: Optional[List[str]] = None,
+    method: str = "fcbg",
+) -> pl.DataFrame:
+    """
+    Perform clustering-based k-anonymity using Polars.
+    
+    Args:
+        df: Input DataFrame or LazyFrame
+        quasi_identifiers: List of column names to use for clustering
+        sensitive_column: Column containing sensitive information
+        k: Minimum group size for k-anonymity
+        categorical: List of categorical column names
+        method: Clustering method ('fcbg', 'rsc', or 'random')
+        
+    Returns:
+        Anonymized DataFrame with generalized quasi-identifiers
+    """
+    raise NotImplementedError("Clustering k-anonymity will be implemented soon")

polarfrost-0.1.0/polarfrost/mondrian.py

@@ -0,0 +1,256 @@
+"""
+Efficient Mondrian k-Anonymity implementation using Polars and PySpark.
+Compatible with local (Polars) and Databricks/Spark (PySpark) environments.
+"""
+
+from typing import List, Optional, Union, Dict, Any, cast, TYPE_CHECKING
+import polars as pl
+
+if TYPE_CHECKING:
+    from pyspark.sql import DataFrame as SparkDataFrame
+    from pyspark.sql.types import StructType
+
+# ------------------------- POLARS VERSION -------------------------
+def mondrian_k_anonymity_polars(
+    df: "pl.DataFrame | pl.LazyFrame",
+    quasi_identifiers: List[str],
+    sensitive_column: str,
+    k: int,
+    categorical: Optional[List[str]] = None,
+) -> pl.DataFrame:
+    """
+    Perform Mondrian k-anonymity using Polars LazyFrame for local processing.
+    Accepts either DataFrame or LazyFrame as input.
+    """
+    if categorical is None:
+        categorical = []
+    
+    # Convert to LazyFrame if not already
+    if isinstance(df, pl.DataFrame):
+        df = df.lazy()
+    elif not isinstance(df, pl.LazyFrame):
+        raise ValueError("Input must be a Polars DataFrame or LazyFrame")
+    
+    # Initialize partitions with the full dataset
+    partitions = [df]
+    result = []
+    
+    # Process partitions until none left
+    while partitions:
+        part = partitions.pop()
+        
+        # Get partition size (lazy evaluation)
+        n_rows = part.select(pl.len()).collect().item(0, 0)
+        
+        # If partition is too small to split, add to results
+        if n_rows < 2 * k:
+            result.append(part)
+            continue
+            
+        # Compute spans for each quasi-identifier
+        spans: Dict[str, Any] = {}
+        for col in quasi_identifiers:
+            if col in categorical:
+                # For categorical, use number of unique values as span
+                n_unique = part.select(pl.col(col).n_unique()).collect().item()
+                spans[col] = n_unique
+            else:
+                # For numerical, use range as span
+                stats = part.select([
+                    pl.col(col).min().alias("min"), 
+                    pl.col(col).max().alias("max")
+                ]).collect()
+                col_min = stats[0, "min"]
+                col_max = stats[0, "max"]
+                spans[col] = col_max - col_min if col_max is not None and col_min is not None else 0
+        
+        # Find the attribute with maximum span
+        split_col = max(spans, key=spans.get)  # type: ignore
+        
+        # If no split possible, add to results
+        if spans[split_col] == 0:
+            result.append(part)
+            continue
+            
+        # Split the partition
+        if split_col in categorical:
+            # For categorical, split on unique values
+            uniq_vals = part.select(pl.col(split_col).unique()).collect().to_series().to_list()
+            mid = len(uniq_vals) // 2
+            left_vals = set(uniq_vals[:mid])
+            right_vals = set(uniq_vals[mid:])
+            left = part.filter(pl.col(split_col).is_in(left_vals))
+            right = part.filter(pl.col(split_col).is_in(right_vals))
+        else:
+            # For numerical, split on median
+            median = part.select(pl.col(split_col).median()).collect().item()
+            left = part.filter(pl.col(split_col) <= median)
+            right = part.filter(pl.col(split_col) > median)
+        
+        # Check if both partitions satisfy k-anonymity
+        left_n = left.select(pl.len()).collect().item(0, 0)
+        right_n = right.select(pl.len()).collect().item(0, 0)
+        
+        if left_n >= k and right_n >= k:
+            # Both partitions are valid, continue splitting
+            partitions.extend([left, right])
+        else:
+            # At least one partition is too small, keep as is
+            result.append(part)
+    
+    # Aggregate each partition
+    agg_rows = []
+    for part in result:
+        # Collect only the columns we need
+        part_df = part.select(quasi_identifiers + [sensitive_column]).collect()
+        row = {}
+        
+        # Generalize quasi-identifiers
+        for col in quasi_identifiers:
+            if col in categorical:
+                # For categorical, use set of unique values
+                row[col] = ','.join(sorted(map(str, part_df[col].unique())))
+            else:
+                # For numerical, use range
+                row[col] = f"{part_df[col].min()}-{part_df[col].max()}"
+        
+        # Add sensitive values and count
+        row[sensitive_column] = ','.join(sorted(map(str, part_df[sensitive_column].unique())))
+        row['count'] = part_df.height
+        agg_rows.append(row)
+    
+    return pl.DataFrame(agg_rows)
+
+# ------------------------- PYSPARK VERSION -------------------------
+def mondrian_k_anonymity_spark(
+    df: "SparkDataFrame",
+    quasi_identifiers: List[str],
+    sensitive_column: str,
+    k: int,
+    categorical: Optional[List[str]] = None,
+    schema: Optional["StructType"] = None,
+) -> "SparkDataFrame":
+    """
+    Perform Mondrian k-anonymity using PySpark for distributed processing.
+    """
+    import pandas as pd
+    from pyspark.sql.functions import pandas_udf, PandasUDFType
+    
+    if categorical is None:
+        categorical = []
+    
+    @pandas_udf(schema, PandasUDFType.GROUPED_MAP)
+    def mondrian_partition(pdf: pd.DataFrame) -> pd.DataFrame:
+        partitions = [pdf]
+        result = []
+        
+        while partitions:
+            part = partitions.pop()
+            
+            # If partition is too small to split, add to results
+            if len(part) < 2 * k:
+                result.append(part)
+                continue
+                
+            # Compute spans for each quasi-identifier
+            spans = {}
+            for col in quasi_identifiers:
+                if col in categorical:
+                    spans[col] = part[col].nunique()
+                else:
+                    col_min = part[col].min()
+                    col_max = part[col].max()
+                    spans[col] = col_max - col_min if pd.notnull(col_max) and pd.notnull(col_min) else 0
+            
+            # Find the attribute with maximum span
+            split_col = max(spans, key=spans.get)
+            
+            # If no split possible, add to results
+            if spans[split_col] == 0:
+                result.append(part)
+                continue
+                
+            # Split the partition
+            if split_col in categorical:
+                # For categorical, split on unique values
+                uniq_vals = part[split_col].unique()
+                mid = len(uniq_vals) // 2
+                left_vals = set(uniq_vals[:mid])
+                right_vals = set(uniq_vals[mid:])
+                left = part[part[split_col].isin(left_vals)]
+                right = part[part[split_col].isin(right_vals)]
+            else:
+                # For numerical, split on median
+                median = part[split_col].median()
+                left = part[part[split_col] <= median]
+                right = part[part[split_col] > median]
+            
+            # Check if both partitions satisfy k-anonymity
+            if len(left) >= k and len(right) >= k:
+                # Both partitions are valid, continue splitting
+                partitions.extend([left, right])
+            else:
+                # At least one partition is too small, keep as is
+                result.append(part)
+        
+        # Aggregate the results
+        agg_rows = []
+        for part in result:
+            row = {}
+            
+            # Generalize quasi-identifiers
+            for col in quasi_identifiers:
+                if col in categorical:
+                    # For categorical, use set of unique values
+                    row[col] = ','.join(sorted(map(str, part[col].unique())))
+                else:
+                    # For numerical, use range
+                    row[col] = f"{part[col].min()}-{part[col].max()}"
+            
+            # Add sensitive values and count
+            row[sensitive_column] = ','.join(sorted(map(str, part[sensitive_column].unique())))
+            row['count'] = len(part)
+            agg_rows.append(row)
+        
+        return pd.DataFrame(agg_rows)
+    
+    # Apply the function to the entire DataFrame
+    if schema is not None:
+        return df.groupBy().applyInPandas(mondrian_partition, schema=schema)
+    else:
+        return df.groupBy().applyInPandas(mondrian_partition)
+
+# ------------------------- DISPATCHER -------------------------
+def mondrian_k_anonymity(
+    df: Union[pl.DataFrame, pl.LazyFrame, "SparkDataFrame"],
+    quasi_identifiers: List[str],
+    sensitive_column: str,
+    k: int,
+    categorical: Optional[List[str]] = None,
+    schema: Optional["StructType"] = None,
+) -> Union[pl.DataFrame, "SparkDataFrame"]:
+    """
+    Dispatcher: Use Polars or PySpark Mondrian k-anonymity depending on input type.
+    
+    Args:
+        df: Input DataFrame (Polars or PySpark)
+        quasi_identifiers: List of column names that are quasi-identifiers
+        sensitive_column: Name of the sensitive column
+        k: Anonymity parameter (minimum group size)
+        categorical: List of categorical column names
+        schema: Schema for PySpark output (required for PySpark)
+        
+    Returns:
+        Anonymized DataFrame with generalized quasi-identifiers
+    """
+    try:
+        from pyspark.sql import DataFrame as SparkDataFrame
+        if isinstance(df, SparkDataFrame):
+            return mondrian_k_anonymity_spark(df, quasi_identifiers, sensitive_column, k, categorical, schema)
+    except ImportError:
+        pass
+    
+    if isinstance(df, (pl.DataFrame, pl.LazyFrame)):
+        return mondrian_k_anonymity_polars(df, quasi_identifiers, sensitive_column, k, categorical)
+    
+    raise ValueError("Input df must be a polars.DataFrame, polars.LazyFrame, or pyspark.sql.DataFrame")

polarfrost-0.1.0/polarfrost/py.typed

@@ -0,0 +1,2 @@
+# This file indicates that this package is typed according to PEP 561
+# It allows type checkers to recognize the package as type-annotated

polarfrost-0.1.0/polarfrost.egg-info/PKG-INFO

@@ -0,0 +1,86 @@
+Metadata-Version: 2.4
+Name: polarfrost
+Version: 0.1.0
+Summary: A fast k-anonymity implementation using Polars and PySpark
+Home-page: https://github.com/rglew/polarfrost
+Author: Richard Glew
+Author-email: richard.glew@hotmail.com
+Keywords: anonymization,privacy,polars,k-anonymity,data-privacy
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Security
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: polars>=0.13.0
+Requires-Dist: pandas>=1.3.0
+Requires-Dist: numpy>=1.21.0
+Provides-Extra: spark
+Requires-Dist: pyspark>=3.0.0; extra == "spark"
+Provides-Extra: dev
+Requires-Dist: pytest>=6.0; extra == "dev"
+Requires-Dist: pytest-cov>=2.0; extra == "dev"
+Requires-Dist: black>=21.0; extra == "dev"
+Requires-Dist: isort>=5.0; extra == "dev"
+Requires-Dist: mypy>=0.900; extra == "dev"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Polarfrost
+
+A fast k-anonymity implementation using Polars, featuring both Mondrian and Clustering algorithms for efficient privacy-preserving data analysis.
+
+## Features
+
+- 🚀 Blazing fast k-anonymity using Polars
+- 🧊 Supports both local (Polars) and distributed (PySpark) processing
+- 📊 Preserves data utility while ensuring privacy
+- 🐍 Simple Python API
+
+## Installation
+
+```bash
+pip install polarfrost
+```
+
+## Quick Start
+
+```python
+import polars as pl
+from polarfrost import mondrian_k_anonymity
+
+# Load your data
+df = pl.read_csv("your_data.csv")
+
+# Apply k-anonymity
+anonymized = mondrian_k_anonymity(
+    df,
+    quasi_identifiers=["age", "gender", "zipcode"],
+    sensitive_column="income",
+    k=3,
+    categorical=["gender", "zipcode"]
+)
+
+print(anonymized)
+```
+
+## License
+
+MIT

polarfrost-0.1.0/polarfrost.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+README.md
+setup.py
+polarfrost/__init__.py
+polarfrost/clustering.py
+polarfrost/mondrian.py
+polarfrost/py.typed
+polarfrost.egg-info/PKG-INFO
+polarfrost.egg-info/SOURCES.txt
+polarfrost.egg-info/dependency_links.txt
+polarfrost.egg-info/not-zip-safe
+polarfrost.egg-info/requires.txt
+polarfrost.egg-info/top_level.txt
+polarfrost/tests/__init__.py
+tests/test_imports.py
+tests/test_mondrian.py
+tests/test_mondrian_implementation.py

polarfrost-0.1.0/polarfrost.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

polarfrost-0.1.0/polarfrost.egg-info/not-zip-safe

@@ -0,0 +1 @@
+

polarfrost-0.1.0/polarfrost.egg-info/requires.txt

@@ -0,0 +1,13 @@
+polars>=0.13.0
+pandas>=1.3.0
+numpy>=1.21.0
+
+[dev]
+pytest>=6.0
+pytest-cov>=2.0
+black>=21.0
+isort>=5.0
+mypy>=0.900
+
+[spark]
+pyspark>=3.0.0

polarfrost-0.1.0/polarfrost.egg-info/top_level.txt

@@ -0,0 +1 @@
+polarfrost

polarfrost-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

polarfrost-0.1.0/setup.py

@@ -0,0 +1,54 @@
+from setuptools import setup, find_packages
+import pathlib
+
+# Read the contents of README.md
+this_directory = pathlib.Path(__file__).parent
+long_description = (this_directory / "README.md").read_text()
+
+setup(
+    name="polarfrost",
+    version="0.1.0",
+    packages=find_packages(exclude=["tests*"]),
+    python_requires=">=3.8",
+    install_requires=[
+        "polars>=0.13.0",
+        "pandas>=1.3.0",
+        "numpy>=1.21.0",
+    ],
+    extras_require={
+        "spark": ["pyspark>=3.0.0"],
+        "dev": [
+            "pytest>=6.0",
+            "pytest-cov>=2.0",
+            "black>=21.0",
+            "isort>=5.0",
+            "mypy>=0.900",
+        ],
+    },
+    author="Richard Glew",
+    author_email="richard.glew@hotmail.com",
+    description="A fast k-anonymity implementation using Polars and PySpark",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url="https://github.com/rglew/polarfrost",
+    keywords=["anonymization", "privacy", "polars", "k-anonymity", "data-privacy"],
+    classifiers=[
+        "Development Status :: 3 - Alpha",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Science/Research",
+        "License :: OSI Approved :: MIT License",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Topic :: Scientific/Engineering",
+        "Topic :: Security",
+        "Topic :: Software Development :: Libraries :: Python Modules",
+    ],
+    package_data={
+        "polarfrost": ["py.typed"],
+    },
+    include_package_data=True,
+    zip_safe=False,
+)

polarfrost-0.1.0/tests/test_imports.py

@@ -0,0 +1,9 @@
+"""Test basic imports."""
+
+def test_import():
+    """Test that the package can be imported and has the expected attributes."""
+    import polarfrost
+    assert polarfrost.__version__ == "0.1.0"
+    assert hasattr(polarfrost, 'mondrian_k_anonymity')
+    assert hasattr(polarfrost, 'mondrian_k_anonymity_polars')
+    assert hasattr(polarfrost, 'mondrian_k_anonymity_spark')

polarfrost-0.1.0/tests/test_mondrian.py

@@ -0,0 +1,28 @@
+import os
+import polars as pl
+from polarfrost.mondrian import mondrian_k_anonymity
+
+def test_mondrian_basic():
+    """Test basic k-anonymity functionality."""
+    # Sample data
+    data = {
+        "age": [25, 25, 35, 35, 45, 45],
+        "gender": ["M", "M", "F", "F", "M", "M"],
+        "zipcode": ["12345", "12345", "12345", "12345", "67890", "67890"],
+        "income": [50000, 55000, 60000, 65000, 70000, 75000]
+    }
+    df = pl.DataFrame(data)
+    
+    # Apply k-anonymity with k=2
+    result = mondrian_k_anonymity(
+        df,
+        quasi_identifiers=["age", "gender", "zipcode"],
+        sensitive_column="income",
+        k=2,
+        categorical=["gender", "zipcode"]
+    )
+    
+    # Basic assertions
+    assert len(result) > 0
+    assert "count" in result.columns
+    assert all(result["count"] >= 2)

polarfrost-0.1.0/tests/test_mondrian_implementation.py

@@ -0,0 +1,88 @@
+"""
+Test the Mondrian k-anonymity implementation with sample data.
+"""
+import os
+import polars as pl
+import pytest
+from polarfrost import mondrian_k_anonymity
+
+def test_mondrian_basic():
+    """Test basic Mondrian k-anonymity with a small dataset."""
+    # Create a small test dataset
+    data = {
+        "age": [25, 25, 35, 35, 45, 45, 55, 55],
+        "gender": ["M", "M", "F", "F", "M", "M", "F", "F"],
+        "zipcode": ["12345", "12345", "12345", "12345", "67890", "67890", "67890", "67890"],
+        "income": [50000, 55000, 60000, 65000, 70000, 75000, 80000, 85000]
+    }
+    df = pl.DataFrame(data)
+    
+    # Apply k-anonymity with k=2
+    quasi_identifiers = ["age", "gender", "zipcode"]
+    sensitive_column = "income"
+    categorical = ["gender", "zipcode"]
+    k = 2
+    
+    anon_df = mondrian_k_anonymity(
+        df,
+        quasi_identifiers=quasi_identifiers,
+        sensitive_column=sensitive_column,
+        k=k,
+        categorical=categorical
+    )
+    
+    # Verify the output
+    assert isinstance(anon_df, pl.DataFrame)
+    assert len(anon_df) > 0  # Should have at least one group
+    assert all(count >= k for count in anon_df["count"])  # All groups should satisfy k-anonymity
+    
+    # Check that all quasi-identifiers are generalized
+    for col in quasi_identifiers:
+        assert col in anon_df.columns
+    
+    # Check that the sensitive column is included
+    assert sensitive_column in anon_df.columns
+    assert "count" in anon_df.columns
+
+def test_mondrian_with_lazyframe():
+    """Test that the function works with LazyFrames."""
+    data = {
+        "age": [25, 25, 35, 35, 45, 45],
+        "gender": ["M", "M", "F", "F", "M", "M"],
+        "zipcode": ["12345", "12345", "12345", "12345", "67890", "67890"],
+        "income": [50000, 55000, 60000, 65000, 70000, 75000]
+    }
+    df = pl.LazyFrame(data)
+    
+    anon_df = mondrian_k_anonymity(
+        df,
+        quasi_identifiers=["age", "gender", "zipcode"],
+        sensitive_column="income",
+        k=2,
+        categorical=["gender", "zipcode"]
+    )
+    
+    assert isinstance(anon_df, pl.DataFrame)
+    assert len(anon_df) > 0
+
+def test_mondrian_invalid_input():
+    """Test that invalid inputs raise appropriate errors."""
+    df = pl.DataFrame({"age": [1, 2, 3], "income": [10, 20, 30]})
+    
+    # Test with k larger than dataset size - should return a single group
+    result = mondrian_k_anonymity(df, ["age"], "income", k=5)
+    assert len(result) == 1  # Should return a single group with all records
+    assert result["count"][0] == 3  # All records should be in one group
+    
+    # Test with invalid column names
+    with pytest.raises(pl.exceptions.ColumnNotFoundError):
+        mondrian_k_anonymity(df, ["invalid"], "income", k=2)
+    
+    with pytest.raises(pl.exceptions.ColumnNotFoundError):
+        mondrian_k_anonymity(df, ["age"], "invalid", k=2)
+
+if __name__ == "__main__":
+    test_mondrian_basic()
+    test_mondrian_with_lazyframe()
+    test_mondrian_invalid_input()
+    print("All tests passed!")