sparkleframe 0.0.1__py3-none-any.whl

sparkleframe/__init__.py

@@ -0,0 +1,8 @@
+"""
+Sparkleframe
+"""
+
+import os
+
+with open(os.path.join(os.path.dirname(__file__), "version.txt"), "r", encoding="utf-8") as f:
+    __version__ = f.read()

sparkleframe/activate.py

@@ -0,0 +1,96 @@
+"""
+Sparkleframe
+"""
+
+# ruff: noqa: F401
+from __future__ import annotations
+
+import importlib
+import sys
+from contextlib import contextmanager
+from unittest.mock import MagicMock
+
+# Helps users to use `from sparkleframe import Engine` instead of `from sparkleframe.engine import Engine`
+from .engine import Engine
+
+
+NAME_TO_FILE_OVERRIDE = {
+    # "DataFrameNaFunctions": "dataframe",
+    # "DataFrameStatFunctions": "dataframe",
+    # "DataFrameReader": "readwriter",
+    # "DataFrameWriter": "readwriter",
+    # "GroupedData": "group",
+    "SparkSession": "session",
+    # "WindowSpec": "window",
+    # "UDFRegistration": "udf",
+}
+
+ACTIVATE_CONFIG = {}
+
+
+def activate() -> None:
+
+    pyspark_mock = MagicMock()
+    pyspark_mock.__file__ = "pyspark"
+    sys.modules["pyspark"] = pyspark_mock
+    # pyspark_mock.testing = testing
+    # sys.modules["pyspark.testing"] = testing
+
+    engine = "polarsdf"
+    prefix = "Polars"
+    engine_module = importlib.import_module(f"sparkleframe.{engine}")
+
+    sys.modules["pyspark.sql"] = engine_module
+    pyspark_mock.sql = engine_module
+    types = engine_module.__dict__.copy()
+    resolved_files = set()
+    for name, obj in types.items():
+        # print("name, obj, name.startswith(prefix): ",name, obj, name.startswith(prefix))
+        if name.startswith(prefix) or name in [
+            "Column",
+            "Window",
+            "WindowSpec",
+            "functions",
+            "types",
+            "Column",
+            "DataFrame",
+            "SparkSession",
+        ]:
+
+            name_without_prefix = name.replace(prefix, "")
+            if name_without_prefix == "Session":
+                name_without_prefix = "SparkSession"
+
+            setattr(engine_module, name_without_prefix, obj)
+
+            file = NAME_TO_FILE_OVERRIDE.get(name_without_prefix, name_without_prefix).lower()
+
+            engine_file = importlib.import_module(f"sparkleframe.{engine}.{file}")
+
+            if engine_file not in resolved_files:
+                sys.modules[f"pyspark.sql.{file}"] = engine_file
+                resolved_files.add(engine_file)
+
+            setattr(engine_file, name_without_prefix, obj)
+
+
+def deactivate() -> None:
+    pyspark_imports = [k for k in sys.modules if k.startswith("pyspark")]
+
+    for k, v in sys.modules.copy().items():
+        if k in pyspark_imports:
+            del sys.modules[k]
+    # Try importing the pyspark imports again and see if pyspark is installed and therefore available
+    # if not then nothing will change
+    for k in pyspark_imports:
+        try:
+            sys.modules[k] = importlib.import_module(k)
+        except ImportError:
+            pass
+    ACTIVATE_CONFIG.clear()
+
+
+@contextmanager
+def activate_context():
+    activate()
+    yield

sparkleframe/base/dataframe.py

@@ -0,0 +1,20 @@
+from typing import Any
+
+from pandas import DataFrame as PandasDataFrame
+
+
+class DataFrame:
+    def __init__(self, df: Any):
+        self._alias = None
+        self.df = df
+
+    def pandas_api(self) -> PandasDataFrame:
+        """Mimics Pandas On Spark dataframe"""
+        return self.toPandas()
+
+    def to_native_df(self) -> "DataFrame":
+        """Return the underlying Polars DataFrame."""
+        return self.df
+
+    def to_spark(self) -> "DataFrame":
+        return self

sparkleframe/conftest.py

@@ -0,0 +1,13 @@
+import sys
+
+import pytest
+
+# Skip writing pyc files on a readonly filesystem.
+sys.dont_write_bytecode = True
+
+
+@pytest.fixture(scope="function", autouse=False)
+def spark():
+    from sparkleframe.tests.spark import spark as spark_session
+
+    return spark_session

sparkleframe/context.py

@@ -0,0 +1,86 @@
+import importlib
+import sys
+import types
+from contextlib import contextmanager
+from typing import Callable
+
+from sparkleframe import Engine
+
+SPARKLE_MODULES = [
+    "functions",
+    "types",
+    "session",
+    "dataframe",
+]
+
+
+@contextmanager
+def mock_pyspark_modules(mock_map):
+    # Save original pyspark modules
+    original_modules = {k: sys.modules[k] for k in list(sys.modules) if k.startswith("pyspark")}
+
+    # Remove existing pyspark.* modules
+    for key in original_modules:
+        del sys.modules[key]
+
+    # Make sure root pyspark module exists
+    if "pyspark" not in sys.modules:
+        sys.modules["pyspark"] = types.ModuleType("pyspark")
+
+    # Add parent modules for each mock
+    for full_name in mock_map:
+        parts = full_name.split(".")
+        for i in range(1, len(parts)):
+            parent = ".".join(parts[:i])
+            if parent not in sys.modules:
+                sys.modules[parent] = types.ModuleType(parent)
+
+    # Inject mocks
+    sys.modules.update(mock_map)
+
+    try:
+        yield
+    finally:
+        # Clean up mocked modules
+        for key in mock_map:
+            pyspark_module_name = key
+            if pyspark_module_name in sys.modules:
+                del sys.modules[pyspark_module_name]
+
+            sparkle_module_name = mock_map[key].__name__
+            if sparkle_module_name in sys.modules:
+                del sys.modules[sparkle_module_name]
+
+        # Restore and reload original modules
+        for name, module in original_modules.items():
+            sys.modules[name] = module
+            importlib.import_module(name)
+
+
+def get_mock_map(engine: Engine) -> dict:
+
+    sparkleframe_module = f"sparkleframe.{engine.module}"
+    engine_module = importlib.import_module(sparkleframe_module)
+
+    mock_map = {"pyspark.sql": engine_module}
+    types = engine_module.__dict__.copy()
+    for name, obj in types.items():
+        name = engine.clean_class_name(name)
+
+        if name in SPARKLE_MODULES:
+            module_name = obj.__name__.split(".")[-1]
+            pyspark_module = f"pyspark.sql.{module_name}"
+            sparkle_module = f"{sparkleframe_module}.{module_name}"
+            mock_map[pyspark_module] = importlib.import_module(sparkle_module)
+
+    return mock_map
+
+
+def run_with_context(func: Callable, engine: Engine = None):
+    engine = engine or Engine.POLARS
+
+    with mock_pyspark_modules(get_mock_map(engine)):
+        for name in SPARKLE_MODULES:
+            importlib.import_module(f"sparkleframe.{engine.module}.{name}")
+
+        return func()

sparkleframe/engine.py

@@ -0,0 +1,16 @@
+from enum import Enum
+
+
+class Engine(Enum):
+    POLARS = ("polarsdf", "Polars")
+
+    @property
+    def module(self):
+        return self.value[0]
+
+    @property
+    def class_prefix(self):
+        return self.value[1]
+
+    def clean_class_name(self, klass):
+        return klass.replace(self.class_prefix, "") if klass.startswith(self.class_prefix) else klass

sparkleframe/gpt_utils.py

@@ -0,0 +1,91 @@
+"""
+this is my <file_name.py>
+
+<content of the file>
+
+implement the pyspark <function> function like as defined in <function_url> into my <file_name.py>
+
+"""
+
+import os
+
+
+def generate_pyspark_function_instruction(file_path: str, function_name: str, function_url: str) -> str:
+    """
+    Generates a formatted message including the contents of a Python file and an instruction
+    to implement a specific PySpark function based on its official documentation.
+
+    Args:
+        file_path (str): Path to the Python file.
+        function_name (str): Name of the PySpark function to implement.
+        function_url (str): URL to the official PySpark documentation for the function.
+
+    Returns:
+        str: Formatted message.
+    """
+    if not os.path.isfile(file_path):
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    with open(file_path, "r") as file:
+        file_contents = file.read()
+
+    file_name = os.path.basename(file_path)
+
+    message = f"""\
+this is the content of my {file_name}:
+
+{file_contents}
+
+implement the pyspark {function_name} function like as defined in {function_url} into my {file_name}
+"""
+    return message
+
+
+def generate_test_instruction(file_path: str, function_name: str, compare_output_spark=False) -> str:
+    """
+    Generates a formatted message including the contents of a Python file and an instruction
+    to implement a specific PySpark function based on its official documentation.
+
+    Args:
+        file_path (str): Path to the Python file.
+        function_name (str): Name of the PySpark function to implement.
+        function_url (str): URL to the official PySpark documentation for the function.
+
+    Returns:
+        str: Formatted message.
+    """
+    if not os.path.isfile(file_path):
+        raise FileNotFoundError(f"File not found: {file_path}")
+
+    with open(file_path, "r") as file:
+        file_contents = file.read()
+
+    file_name = os.path.basename(file_path)
+
+    spark_msg_output = (
+        ", create a spark dataframe build from PolarDataframe and use another spark dataframe to compare them "
+        if compare_output_spark
+        else ""
+    )
+
+    message = f"""\
+this is the content of my test file {file_name}:
+
+{file_contents}
+
+generate a pytest test case, using pytest.parametrize to test the logic `{function_name}` that has been just implemented{spark_msg_output}
+"""
+    return message
+
+
+file_path = "./polarsdf/dataframe.py"
+function_name = "fillna"
+pyspark_function_url = (
+    "https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/api/pyspark.sql.DataFrame.fillna.html"
+)
+
+msg = generate_pyspark_function_instruction(file_path, function_name, pyspark_function_url)
+
+msg = generate_test_instruction(file_path, function_name, compare_output_spark=True)
+
+print(msg)

sparkleframe/polarsdf/__init__.py

@@ -0,0 +1,12 @@
+# ruff: noqa: F401
+from sparkleframe.polarsdf.column import Column
+from sparkleframe.polarsdf.dataframe import DataFrame
+from sparkleframe.polarsdf.functions import col
+from sparkleframe.polarsdf.session import SparkSession
+from sparkleframe.polarsdf.types import StringType
+
+__all__ = [
+    "Column",
+    "SparkSession",
+    "DataFrame",
+]

sparkleframe/polarsdf/column.py

@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import polars as pl
+
+from sparkleframe.polarsdf.types import DataType
+
+
+class Column:
+    def __init__(self, expr_or_name):
+        if isinstance(expr_or_name, str):
+            self.expr = pl.col(expr_or_name)
+        else:
+            self.expr = expr_or_name
+
+    # Arithmetic operations
+    def __mul__(self, other):
+        return Column(self.expr * _to_expr(other))
+
+    def __add__(self, other):
+        return Column(self.expr + _to_expr(other))
+
+    def __sub__(self, other):
+        return Column(self.expr - _to_expr(other))
+
+    def __truediv__(self, other):
+        return Column(self.expr / _to_expr(other))
+
+    def __radd__(self, other):
+        return Column(_to_expr(other) + self.expr)
+
+    def __rsub__(self, other):
+        return Column(_to_expr(other) - self.expr)
+
+    def __rmul__(self, other):
+        return Column(_to_expr(other) * self.expr)
+
+    def __rtruediv__(self, other):
+        return Column(_to_expr(other) / self.expr)
+
+    # Comparison operations
+    def __eq__(self, other):
+        return Column(self.expr == _to_expr(other))
+
+    def __ne__(self, other):
+        return Column(self.expr != _to_expr(other))
+
+    def __lt__(self, other):
+        return Column(self.expr < _to_expr(other))
+
+    def __le__(self, other):
+        return Column(self.expr <= _to_expr(other))
+
+    def __gt__(self, other):
+        return Column(self.expr > _to_expr(other))
+
+    def __ge__(self, other):
+        return Column(self.expr >= _to_expr(other))
+
+    # Logical operations
+    def __and__(self, other):
+        return Column(self.expr & _to_expr(other))
+
+    def __rand__(self, other):
+        return Column(_to_expr(other) & self.expr)
+
+    def __or__(self, other):
+        return Column(self.expr | _to_expr(other))
+
+    def __ror__(self, other):
+        return Column(_to_expr(other) | self.expr)
+
+    def alias(self, name: str) -> Column:
+        """
+        Mimics pyspark.sql.Column.alias
+
+        Args:
+            name (str): Alias name for the column expression
+
+        Returns:
+            Column: A new Column with the alias applied
+        """
+        return Column(self.expr.alias(name))
+
+    def cast(self, data_type: DataType) -> Column:
+        """
+        Mimics pyspark.sql.Column.cast using Polars' cast().
+
+        Args:
+            data_type (DataType): A sparkleframe-defined DataType object.
+
+        Returns:
+            Column: A new Column with the expression casted.
+        """
+        if not isinstance(data_type, DataType):
+            raise TypeError(f"cast() expects a DataType, got {type(data_type)}")
+        return Column(self.expr.cast(data_type.to_native()))
+
+    def isin(self, *values) -> Column:
+        """
+        Mimics pyspark.sql.Column.isin and supports both:
+            col("x").isin("a", "b") and col("x").isin(["a", "b"])
+
+        Args:
+            *values: A list of values or individual arguments.
+
+        Returns:
+            Column: A Column representing a boolean expression.
+        """
+        # If a single iterable (non-str) is passed, use that directly
+        if len(values) == 1 and isinstance(values[0], Iterable) and not isinstance(values[0], str):
+            value_list = list(values[0])
+        else:
+            value_list = list(values)
+
+        return Column(self.expr.is_in(value_list))
+
+    def isNotNull(self) -> Column:
+        """
+        Mimics pyspark.sql.Column.isNotNull
+
+        Returns:
+            Column: A Column representing the non-null condition.
+        """
+        return Column(self.expr.is_not_null())
+
+    def rlike(self, pattern: str) -> Column:
+        """
+        Mimics pyspark.sql.Column.rlike using Polars' regex matching.
+
+        Args:
+            pattern (str): Regular expression pattern to match.
+
+        Returns:
+            Column: A new Column representing a boolean expression.
+        """
+        if not isinstance(pattern, str):
+            raise TypeError(f"rlike() expects a string pattern, got {type(pattern)}")
+
+        return Column(self.expr.str.contains(pattern))
+
+    def to_native(self) -> pl.Expr:
+        return self.expr
+
+
+def _to_expr(value):
+    if isinstance(value, Column):
+        return value.to_native()
+    elif isinstance(value, pl.Expr):
+        return value
+    else:
+        return pl.lit(value)