wringlet 0.0.3.dev6__py3-none-any.whl

wringlet/__init__.py

@@ -0,0 +1,15 @@
+from .data_provenance import (
+    add_provenance_column,
+    data_provenance_enabled,
+    data_provenance_session_builder,
+    provenance_column_name,
+    remove_provenance_column,
+)
+
+__all__ = [
+    "provenance_column_name",
+    "add_provenance_column",
+    "remove_provenance_column",
+    "data_provenance_enabled",
+    "data_provenance_session_builder",
+]

wringlet/data_provenance.py

@@ -0,0 +1,156 @@
+import os
+import typing as T
+from contextlib import contextmanager
+
+from pyspark.sql import DataFrame, SparkSession
+
+from wringlet.py4j_utils import _get_provenance_jvm_function
+
+DataFrameOrView = str | DataFrame
+
+
+def provenance_column_name(spark: SparkSession) -> str:
+    """
+    Returns the name of the provenance column from the Spark configuration.
+    """
+    return _get_provenance_jvm_function("provenanceColumnName", spark)(spark._jsparkSession)
+
+
+@T.overload
+def add_provenance_column(df_or_view: str, spark: SparkSession) -> str: ...
+
+
+@T.overload
+def add_provenance_column(df_or_view: DataFrame, spark: SparkSession) -> DataFrame: ...
+
+
+def add_provenance_column(df_or_view: DataFrameOrView, spark: SparkSession) -> DataFrameOrView:
+    """
+    Adds a provenance column to the given DataFrame.
+    """
+    jfunction = _get_provenance_jvm_function("addProvenance", spark)
+    match df_or_view:
+        case str():
+            jfunction(spark._jsparkSession, df_or_view)
+            return df_or_view
+        case DataFrame():
+            return DataFrame(jfunction(df_or_view._jdf), spark)
+        case _:
+            raise TypeError("df_or_view must be a str or a DataFrame")
+
+
+@T.overload
+def remove_provenance_column(df_or_view: str, spark: SparkSession) -> str: ...
+
+
+@T.overload
+def remove_provenance_column(df_or_view: DataFrame, spark: SparkSession) -> DataFrame: ...
+
+
+def remove_provenance_column(df_or_view: DataFrameOrView, spark: SparkSession) -> DataFrameOrView:
+    """
+    Removes the provenance column from the given DataFrame.
+    """
+    jfunction = _get_provenance_jvm_function("removeProvenance", spark)
+    match df_or_view:
+        case str():
+            jfunction(spark._jsparkSession, df_or_view)
+            return df_or_view
+        case DataFrame():
+            return DataFrame(jfunction(df_or_view._jdf), spark)
+        case _:
+            raise TypeError("df_or_view must be a str or a DataFrame")
+
+
+@contextmanager
+def data_provenance_enabled(
+    spark: SparkSession, *args: DataFrameOrView
+) -> T.Generator[T.Tuple[DataFrameOrView, ...] | DataFrameOrView | None, None, None]:
+    """
+    Context manager to enable data provenance for the duration of a block of code.
+    When DataFrames or view names are provided, it adds a provenance column to them
+    and removes it after the block is executed.
+    """
+    # Remember the previous state in case these are nested
+    is_data_provenance_enabled = str(spark.conf.get("spark.provenance.enabled", "false"))
+    try:
+        # Turn data provenance on for this block
+        spark.conf.set("spark.provenance.enabled", "true")
+        dataframe_or_views_with_provenance = tuple(add_provenance_column(df, spark) for df in args)
+        if not args:
+            yield None
+        else:
+            yield dataframe_or_views_with_provenance if len(args) > 1 else dataframe_or_views_with_provenance[0]
+    finally:
+        # Revert to whatever it was before
+        spark.conf.set("spark.provenance.enabled", is_data_provenance_enabled)
+        # Remove the provenance column from the DataFrames/views
+        for df in args:
+            remove_provenance_column(df, spark)
+
+
+def data_provenance_session_builder(
+    provenance_builder: str = "display",
+) -> SparkSession.Builder:
+    """
+    Helper function to automatically find the bundled JAR
+    and initialize a SparkSession Builder with the plugin enabled.
+
+    Args:
+        provenance_builder: Builder strategy used by the Scala extension.
+            Supported values are: display, boolean, semi-why, full-why.
+    """
+    # 1. Find the path to the 'jars' folder dynamically
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    jar_path = os.path.join(current_dir, "jars", "dp-spark_2.13-0.0.1.jar")
+
+    # 2. Build and return the SparkSession
+    return (
+        SparkSession.builder.config("spark.jars", jar_path)
+        .config(
+            "spark.sql.extensions",
+            "org.dataprov.dp.wringlet.SparkProvenanceExtension",
+        )
+        .config("spark.provenance.builder", provenance_builder)
+    )
+
+
+def get_minimal_sources(
+    df: DataFrame, source_dfs: T.Sequence[DataFrame], spark: SparkSession, provenance_col: str | None = None
+) -> T.List[DataFrame]:
+    """
+    Extracts the minimal rows from source DataFrames that contributed to the given final DataFrame.
+    """
+    if provenance_col is None:
+        provenance_col = provenance_column_name(spark)
+
+    gw = spark._sc._gateway
+    if gw is None:
+        raise RuntimeError("Spark context gateway is not initialized")
+
+    java_source_dfs = gw.jvm.java.util.ArrayList()
+    for source in source_dfs:
+        java_source_dfs.add(source._jdf)
+
+    try:
+        j_function = _get_provenance_jvm_function("getMinimalSources", spark)
+        java_result_list = j_function(df._jdf, java_source_dfs, provenance_col)
+    except Exception:
+        j_extractor = gw.jvm.org.dataprov.dp.wringlet.ProvenanceExtractor
+        java_result_list = j_extractor.getMinimalSources(df._jdf, java_source_dfs, provenance_col)
+
+    python_result_dfs = []
+    for i in range(java_result_list.size()):
+        python_result_dfs.append(DataFrame(java_result_list.get(i), spark))
+
+    return python_result_dfs
+
+
+def _py_dataframe_extension_get_minimal_sources(
+    self: DataFrame, source_dfs: T.Sequence[DataFrame], provenance_col: str | None = None
+) -> T.List[DataFrame]:
+    """Helper attached to PySpark DataFrame class for fluent API usage."""
+    return get_minimal_sources(self, source_dfs, self.sparkSession, provenance_col)
+
+
+setattr(DataFrame, "get_minimal_sources", _py_dataframe_extension_get_minimal_sources)

wringlet/jobs/getting_started_job.py

@@ -0,0 +1,78 @@
+#
+# File: https://github.com/data-prov/wringlet/blob/main/pyspark-wringlet/src/wringlet/jobs/getting_started_job.py
+#
+# See also https://github.com/data-prov/wringlet/blob/main/pyspark-wringlet/notebooks/demo.ipynb
+#
+"""
+Pyspark script to test the data provenance concept
+"""
+
+import datetime
+
+from pyspark.sql import SparkSession
+
+import wringlet as dp
+
+JOB_NAME = "getting_started_job"
+today_date: str = datetime.date.today().strftime("%Y-%m-%d")
+now_datetime: datetime.datetime = datetime.datetime.now(datetime.UTC)
+
+
+def main():
+    """
+    The dp-spark JAR is expected to have been installed locally, for instance,
+    thanks to `sbt publishLocal publishM2` in the Scala-related folder
+    """
+    spark = (
+        SparkSession.builder.appName(JOB_NAME)
+        .config("spark.sql.extensions", "org.dataprov.dp.ProvenanceExtension")
+        .enableHiveSupport()
+        .getOrCreate()
+    )
+
+    # Showcase that the data_provenance_enabled() function works as expected
+    print(spark.conf.get("spark.provenance.enabled", "false"))
+    with dp.data_provenance_enabled(spark):
+        print(spark.conf.get("spark.provenance.enabled", "false"))
+    print(spark.conf.get("spark.provenance.enabled", "false"))
+
+    # Toy sample
+    df = spark.createDataFrame(
+        [
+            ("A", datetime.date(2026, 1, 15), 10.0, 90),
+            ("A", datetime.date(2026, 1, 16), 10.0, 120),
+            ("A", datetime.date(2026, 1, 17), 5.0, 300),
+            ("B", datetime.date(2026, 1, 15), 100.0, 20),
+            ("B", datetime.date(2026, 1, 16), 100.0, 30),
+            ("B", datetime.date(2026, 1, 17), 80.0, 60),
+        ],
+        ["product", "date", "price", "sales"],
+    )
+    df.printSchema()
+    df.show()
+
+    # Test with PySpark syntax
+
+    ## Without provenance
+    df2 = df.select("product")
+    print("Without provenance")
+    df2.show()
+
+    ## With provenance
+    df3 = None
+    with dp.data_provenance_enabled(spark):
+        df3 = df.select("product")
+    print("With provenance")
+    df3.show()
+
+    # Test with SQL syntax
+    result_df = None
+    df.createOrReplaceTempView("sales")
+    with dp.data_provenance_enabled(spark):
+        result_df = spark.sql("select * from sales")
+
+    result_df.show()
+
+
+if __name__ == "__main__":
+    main()

wringlet/py4j_utils.py

@@ -0,0 +1,21 @@
+import typing as T
+
+from py4j.java_gateway import JVMView
+from pyspark.sql import SparkSession
+
+
+def _get_jvm_from_spark(spark: SparkSession) -> JVMView:
+    """
+    Helper function to get the JVM view from a SparkSession.
+    """
+    assert spark._jvm is not None
+    return spark._jvm
+
+
+def _get_provenance_jvm_function(name: str, spark: SparkSession) -> T.Callable:
+    """
+    Retrieves JVM function identified by name from
+    Java gateway associated with Spark session.
+    """
+    jvm = _get_jvm_from_spark(spark)
+    return getattr(getattr(jvm, "org.dataprov.dp.wringlet.ProvenanceApi"), name)

wringlet-0.0.3.dev6.dist-info/METADATA

@@ -0,0 +1,135 @@
+Metadata-Version: 2.3
+Name: wringlet
+Version: 0.0.3.dev6
+Summary: PySpark helper utilities for Spark data provenance
+Author: data-prov
+Requires-Python: >=3.12
+Project-URL: repository, https://github.com/data-prov/wringlet
+Description-Content-Type: text/markdown
+
+# PySpark - Fine-grained data provenance - PySpark part
+
+## Table of Content (ToC)
+
+* [PySpark \- Fine\-grained data provenance \- PySpark part](#pyspark---fine-grained-data-provenance---pyspark-part)
+  * [Table of Content (ToC)](#table-of-content-toc)
+  * [Overview](#overview)
+  * [Quick Start](#quick-start)
+  * [Project Layout](#project-layout)
+  * [Common commands](#common-commands)
+  * [Development life\-cycle](#development-life-cycle)
+  * [PyPI package](#pypi-package)
+  * [CI/CD](#cicd)
+
+Created by [gh-md-toc](https://github.com/ekalinin/github-markdown-toc.go)
+
+## Overview
+
+Python package in this monorepo for enabling Spark data provenance features
+in PySpark jobs.
+
+This project is managed with `uv`.
+
+## Quick Start
+
+From the repository root:
+
+```bash
+cd pyspark-wringlet
+make init-uv-python
+make init
+make check
+make test
+make run
+```
+
+## Project Layout
+
+```text
+pyspark-wringlet/
+    src/wringlet/
+    tests/
+    main.py
+    pyproject.toml
+    Makefile
+```
+
+## Common commands
+
+* `make init` - Create/refresh lock file and sync dependencies
+* `make update` - Upgrade dependencies and sync environment
+* `make check` - Run lint and type checks
+* `make test` - Run unit tests
+* `make build` - Build package wheel
+* `make publish` - Publish package (PyPI)
+
+## Development life-cycle
+
+* In order to switch (bump) to a newer
+  [version of the Python package](https://github.com/data-prov/wringlet/blob/main/pyspark-wringlet/VERSION),
+  the following are the main options, ordered by the general probability of
+  occurrence in the development life-cycle, from the highest to the lowest)
+  * Increment the dev version (_e.g._, from `2.4.3.dev5` to `2.4.3.dev6` or
+  from `2.4.3` to `2.4.4.dev0`):
+  `make increment-dev-version`
+  * Bump to minor version (_e.g._, from `2.4.3.dev5` to `2.4.3`):
+  `make bump-to-minor-version`
+  * Bump to patch version (_e.g._, from `2.4.3.dev5` to `2.4.4`):
+  `make bump-to-patch-version`
+  * Bump to major version (_e.g._, from `2.4.3.dev5` to `3.0.0`):
+  `make bump-to-major-version`
+* Then, the version bump has to be cascaded to related files (typically,
+  [`pyproject.toml`](https://github.com/data-prov/wringlet/blob/main/pyspark-wringlet/pyproject.toml))
+
+## PyPI package
+
+* The
+  [Python package is published on PyPI](https://pypi.org/project/pyspark-data-provenance/)
+
+* Optionally, the Python wheel may be published manually onto
+  [Pypi.org](https://pypi.org/project/pyspark-data-provenance/)
+  * The `UV_PUBLISH_TOKEN` environment variable then needs to be specified
+  (see the
+  [uv documentation](https://docs.astral.sh/uv/guides/package/#publishing-your-package)
+  for further details)
+
+```bash
+make publish
+```
+
+* However, the easiest way to publish the Python wheel is through the CI/CD
+  pipeline, that is, the
+  [GitHub Actions Python publishing pipeline](https://github.com/data-prov/wringlet/actions/workflows/python-publish.yml).
+  * That CI/CD pipeline is automatically triggered when creating a release on
+  the Git repository
+  * It may also be triggered manually by contributors of the Git repository
+
+* Check the versions of the Python wheel on
+  [Pypi.org](https://pypi.org/project/pyspark-data-provenance/)
+
+* Install the Python wheel locally, in the main/global Python environment
+  (as `spark-submit` does not seem to work properly with uv):
+
+```bash
+make install-local
+```
+
+* For information, that actually install the Python wheel with pip:
+
+```bash
+python -mpip install -U pyspark-data-provenance
+```
+
+* Run the PySpark job locally, in the main/global Python environment
+  (as `spark-submit` does not seem to work properly with uv):
+
+```bash
+make run-local
+```
+
+## CI/CD
+
+Repository-level workflows are provided for:
+
+* CI: lint, type-check, and tests on pushes/PRs affecting this package
+* Publish: build and publish on release

wringlet-0.0.3.dev6.dist-info/RECORD

@@ -0,0 +1,9 @@
+wringlet/__init__.py,sha256=8teUbaWAJzmxAny_X0WnGDD_Dm3YrNz68tQ4sl2qzy0,360
+wringlet/data_provenance.py,sha256=i_FKaCMrjDVng-vXKuYaUaGo19o-cmb37KC5YguJD-0,5634
+wringlet/jars/.gitkeep,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+wringlet/jobs/getting_started_job.py,sha256=jN1VzKvJXiyCVn6xa3jJX1yZGT2KuMuIehTDcIQRjtw,2263
+wringlet/py4j_utils.py,sha256=TGvYnZJmb_SgjDFVdM7hWf-dA5_4ktHFvApQ99DOTnE,602
+wringlet-0.0.3.dev6.dist-info/WHEEL,sha256=GuAqCqoyQuys5_R4zkHUJFlKXw4RpRLNzo31-ui90WQ,81
+wringlet-0.0.3.dev6.dist-info/entry_points.txt,sha256=IiCZThoHflXIJmJ5_675wH_RyJj4wCX4ieDa4vrVcEQ,80
+wringlet-0.0.3.dev6.dist-info/METADATA,sha256=5KXEgs2Wfy-JI8nwEA5MbSORWsWObS65fjbswXSokFs,4070
+wringlet-0.0.3.dev6.dist-info/RECORD,,

wringlet-0.0.3.dev6.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.10.12
+Root-Is-Purelib: true
+Tag: py3-none-any

wringlet-0.0.3.dev6.dist-info/entry_points.txt

@@ -0,0 +1,3 @@
+[console_scripts]
+getting_started_job = wringlet.jobs.getting_started_job:main
+