matchbox-db 0.6.2__tar.gz → 0.6.3.dev53__tar.gz

{matchbox_db-0.6.2 → matchbox_db-0.6.3.dev53}/.github/workflows/prerelease.yml

@@ -7,8 +7,10 @@ env:
   IMAGE_NAME: ${{ github.repository }}
 
 jobs:
-  build-and-push-development-image:
+  build:
     runs-on: ubuntu-latest
+    outputs:
+      mb_version: ${{ steps.get_version.outputs.mb_version }}
     permissions:
       contents: write
       packages: write
@@ -34,11 +36,21 @@ jobs:
           python-version: "3.11"
 
       - name: Extract development version
+        id: get_version
         run: |
-          echo "MB_VERSION=$( \
-            SETUPTOOLS_SCM_PRETEND_METADATA='{distance=1}' \
-            uv run --frozen python -m setuptools_scm \
-          )" >> $GITHUB_ENV
+          echo "mb_version=$(uv run --frozen python -m setuptools_scm | sed 's/+.*//')" \
+          >> "$GITHUB_OUTPUT"
+
+      - name: Build package
+        env:
+          SETUPTOOLS_SCM_PRETEND_VERSION: ${{ steps.get_version.outputs.mb_version }}
+        run: uv build
+
+      - name: Upload package artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: package-dist
+          path: ./dist
 
       - name: Extract tag metadata for Docker
         id: meta
@@ -46,7 +58,7 @@ jobs:
         with:
           images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
           tags: |
-            type=raw,value=${{ env.MB_VERSION }}
+            type=raw,value=${{ steps.get_version.outputs.mb_version }}
             type=raw,value=development
 
       - name: Build and push Docker image
@@ -56,7 +68,7 @@ jobs:
           file: src/matchbox/server/Dockerfile
           push: true
           build-args: |
-            MB_VERSION=${{ env.MB_VERSION }}
+            MB_VERSION=${{ steps.get_version.outputs.mb_version }}
           tags: ${{ steps.meta.outputs.tags }}
 
       - name: Delete existing development release if it exists
@@ -76,5 +88,23 @@ jobs:
             
             May be unstable.
             
-            **Version:** ${{ env.MB_VERSION }}
+            **Version:** ${{ steps.get_version.outputs.mb_version }}
             **Commit:** ${{ github.sha }}
+
+  deploy-package:
+    needs: build
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+    
+    environment: pypi
+    
+    steps:
+      - name: Download package artifacts
+        uses: actions/download-artifact@v4
+        with:
+          name: package-dist
+          path: ./dist
+
+      - name: Publish package to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

{matchbox_db-0.6.2/src/matchbox_db.egg-info → matchbox_db-0.6.3.dev53}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: matchbox-db
-Version: 0.6.2
+Version: 0.6.3.dev53
 Summary: A framework for orchestrating and comparing data linking and deduplication methodologies.
 Author: Department for Business and Trade
 Project-URL: Documentation, https://uktrade.github.io/matchbox/

{matchbox_db-0.6.2 → matchbox_db-0.6.3.dev53}/src/matchbox/client/_handler.py

@@ -8,6 +8,7 @@ from importlib.metadata import version
 from io import BytesIO
 
 import httpx
+import polars as pl
 from pyarrow import Table
 from pyarrow.parquet import read_table
 from tenacity import (
@@ -301,13 +302,14 @@ def get_resolution(
 
 @http_retry
 def set_data(
-    name: ResolutionName, data: Table, validate_type: ResolutionType
+    name: ResolutionName, data: pl.DataFrame, validate_type: ResolutionType
 ) -> UploadStatus:
     """Upload source hashes or model results to server."""
     log_prefix = f"Resolution {name}"
     logger.debug("Uploading results", prefix=log_prefix)
 
-    buffer = table_to_buffer(table=data)
+    data_arrow = data.to_arrow() if isinstance(data, pl.DataFrame) else data
+    buffer = table_to_buffer(table=data_arrow)
 
     # Initialise upload
     metadata_res = CLIENT.post(
@@ -432,7 +434,7 @@ def download_eval_data() -> tuple[Table, Table]:
     check_schema(SCHEMA_JUDGEMENTS, judgements.schema)
     check_schema(SCHEMA_CLUSTER_EXPANSION, expansion.schema)
 
-    return judgements, expansion
+    return pl.from_arrow(judgements), pl.from_arrow(expansion)
 
 
 # Admin

{matchbox_db-0.6.2 → matchbox_db-0.6.3.dev53}/src/matchbox/client/eval/utils.py

@@ -128,11 +128,7 @@ class EvalData:
 
         threshold = int(threshold * 100)
 
-        root_leaf = (
-            results.root_leaf()
-            .rename({"root_id": "root", "leaf_id": "leaf"})
-            .to_arrow()
-        )
+        root_leaf = results.root_leaf().rename({"root_id": "root", "leaf_id": "leaf"})
         return precision_recall([root_leaf], self.judgements, self.expansion)[0]
 
     def pr_curve(self, results: Results) -> Figure:
@@ -140,7 +136,7 @@ class EvalData:
         all_p = []
         all_r = []
 
-        probs = pl.from_arrow(results.probabilities)
+        probs = results.probabilities
         thresholds = probs.select("probability").unique().to_series()
         for i, t in enumerate(sorted(thresholds)):
             float_thresh = t / 100

{matchbox_db-0.6.2 → matchbox_db-0.6.3.dev53}/src/matchbox/client/models/models.py

@@ -224,8 +224,8 @@ class Model:
         if for_validation:
             self.results = Results(
                 probabilities=results,
-                left_root_leaf=self.left_query.leaf_id.to_arrow(),
-                right_root_leaf=self.right_query.leaf_id.to_arrow()
+                left_root_leaf=self.left_query.leaf_id,
+                right_root_leaf=self.right_query.leaf_id
                 if right_df is not None
                 else None,
             )

{matchbox_db-0.6.2 → matchbox_db-0.6.3.dev53}/src/matchbox/client/results.py

@@ -4,8 +4,6 @@ from collections.abc import Hashable
 from typing import ParamSpec, TypeVar
 
 import polars as pl
-import pyarrow as pa
-import pyarrow.compute as pc
 from pydantic import ConfigDict
 
 from matchbox.common.arrow import SCHEMA_RESULTS
@@ -33,14 +31,14 @@ class Results:
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
-    probabilities: pa.Table
-    _clusters: pa.Table | None = None
+    probabilities: pl.DataFrame
+    _clusters: pl.DataFrame | None = None
 
     def __init__(
         self,
-        probabilities: pa.Table | pl.DataFrame,
-        left_root_leaf: pa.Table | None = None,
-        right_root_leaf: pa.Table | None = None,
+        probabilities: pl.DataFrame,
+        left_root_leaf: pl.DataFrame | None = None,
+        right_root_leaf: pl.DataFrame | None = None,
     ) -> None:
         """Initialises and validates results."""
         self.left_root_leaf = None
@@ -51,61 +49,49 @@ class Results:
         if right_root_leaf is not None:
             self.right_root_leaf = right_root_leaf
 
-        if isinstance(probabilities, pl.DataFrame):
-            probabilities = probabilities.to_arrow()
-
-        if not isinstance(probabilities, pa.Table):
-            raise ValueError("Expected a polars DataFrame or pyarrow Table.")
+        if not isinstance(probabilities, pl.DataFrame):
+            raise ValueError(f"Expected a polars DataFrame, got {type(probabilities)}.")
 
         expected_fields = set(SCHEMA_RESULTS.names)
-        if set(probabilities.column_names) != expected_fields:
+        if set(probabilities.columns) != expected_fields:
             raise ValueError(
                 f"Expected {expected_fields}.\nFound {set(probabilities.column_names)}."
             )
 
         # Handle empty tables
-        if probabilities.num_rows == 0:
-            empty_arrays = [pa.array([], type=field.type) for field in SCHEMA_RESULTS]
-            probabilities = pa.Table.from_arrays(
-                empty_arrays, names=[field.name for field in SCHEMA_RESULTS]
-            )
+        if probabilities.height == 0:
+            probabilities = pl.DataFrame(schema=pl.Schema(SCHEMA_RESULTS))
 
         # Process probability field if it contains floating-point or decimal values
-        probability_type = probabilities["probability"].type
-        if pa.types.is_floating(probability_type) or pa.types.is_decimal(
-            probability_type
-        ):
-            probability_uint8 = pc.cast(
-                pc.round(pc.multiply(probabilities["probability"], 100)),
-                options=pc.CastOptions(
-                    target_type=pa.uint8(),
-                    allow_float_truncate=True,
-                    allow_decimal_truncate=True,
-                ),
+        probability_type = probabilities["probability"].dtype
+        if probability_type.is_float() or probability_type.is_decimal():
+            probability_uint8 = pl.Series(
+                probabilities.select(
+                    pl.col("probability").mul(100).round(0).cast(pl.UInt8)
+                )
             )
 
             # Check max value only if the table is not empty
-            max_prob = pc.max(probability_uint8)
-            if max_prob is not None and max_prob.as_py() > 100:
-                p_max = pc.max(probabilities["probability"]).as_py()
-                p_min = pc.min(probabilities["probability"]).as_py()
+            max_prob = probability_uint8.max()
+            if max_prob is not None and max_prob > 100:
+                p_max = max_prob
+                p_min = probability_uint8.min()
                 raise ValueError(f"Probability range misconfigured: [{p_min}, {p_max}]")
 
-            probabilities = probabilities.set_column(
-                i=probabilities.schema.get_field_index("probability"),
-                field_="probability",
-                column=probability_uint8,
+            probabilities = probabilities.replace_column(
+                probabilities.get_column_index("probability"), probability_uint8
             )
 
-        self.probabilities = probabilities.cast(SCHEMA_RESULTS)
+        # need schema in format recognised by polars
+        self.probabilities = probabilities.cast(pl.Schema(SCHEMA_RESULTS))
 
     @property
     def clusters(self):
         """Retrieve new clusters implied by these results."""
-        if not self._clusters:
+        if self._clusters is None:
             im = IntMap()
             self._clusters = to_clusters(
-                results=self.probabilities, dtype=pa.int64, hash_func=im.index
+                results=self.probabilities, dtype=pl.Int64, hash_func=im.index
             )
         return self._clusters
 
@@ -146,7 +132,7 @@ class Results:
     ) -> pl.DataFrame:
         """Enriches the probability results with the source data."""
         return self._merge_with_source_data(
-            base_df=pl.from_arrow(self.probabilities),
+            base_df=self.probabilities,
             base_df_cols=["left_id", "right_id", "probability"],
             left_data=left_data,
             left_key=left_key,
@@ -165,7 +151,7 @@ class Results:
     ) -> pl.DataFrame:
         """Enriches the cluster results with the source data."""
         return self._merge_with_source_data(
-            base_df=pl.from_arrow(self.clusters),
+            base_df=self.clusters,
             base_df_cols=["parent", "child", "probability"],
             left_data=left_data,
             left_key=left_key,
@@ -182,20 +168,19 @@ class Results:
                 "This Results object wasn't instantiated for validation features."
             )
 
-        parents_root_leaf = pl.from_arrow(self.left_root_leaf.select(["id", "leaf_id"]))
+        parents_root_leaf = self.left_root_leaf.select(["id", "leaf_id"])
         if self.right_root_leaf is not None:
             parents_root_leaf = pl.concat(
                 [
                     parents_root_leaf,
-                    pl.from_arrow(self.right_root_leaf.select(["id", "leaf_id"])),
+                    self.right_root_leaf.select(["id", "leaf_id"]),
                 ]
             )
 
         # Go from parent-child (where child could be the root of another model)
         # to root-leaf, where leaf is a source cluster ID
         root_leaf_res = (
-            pl.from_arrow(self.clusters)
-            .rename({"parent": "root_id"})
+            self.clusters.rename({"parent": "root_id"})
             .join(parents_root_leaf, left_on="child", right_on="id")
             .select(["root_id", "leaf_id"])
             .unique()
@@ -205,7 +190,7 @@ class Results:
         unmerged_ids_rows = (
             parents_root_leaf.select("id", "leaf_id")
             .join(
-                pl.from_arrow(self.clusters).select("child"),
+                self.clusters.select("child"),
                 left_on="id",
                 right_on="child",
                 how="anti",

{matchbox_db-0.6.2 → matchbox_db-0.6.3.dev53}/src/matchbox/common/eval.py

@@ -4,7 +4,6 @@ from itertools import chain, combinations
 from typing import TypeAlias
 
 import polars as pl
-from pyarrow import Table
 from pydantic import BaseModel, Field, field_validator
 
 from matchbox.common.graph import ModelResolutionName
@@ -39,7 +38,9 @@ class Judgement(BaseModel):
 
 
 def precision_recall(
-    models_root_leaf: list[Table], judgements: Table, expansion: Table
+    models_root_leaf: list[pl.DataFrame],
+    judgements: pl.DataFrame,
+    expansion: pl.DataFrame,
 ) -> list[PrecisionRecall]:
     """From models and eval data, compute scores inspired by precision-recall.
 
@@ -78,10 +79,9 @@ def precision_recall(
     for root_leaf in models_root_leaf:
         if not len(root_leaf):
             raise ValueError("Model data cannot be empty.")
-        leaves_per_set.append(set(root_leaf["leaf"].to_pylist()))
+        leaves_per_set.append(set(root_leaf["leaf"].to_list()))
         clusters = (
-            pl.from_arrow(root_leaf)
-            .group_by("root")
+            root_leaf.group_by("root")
             .agg(pl.col("leaf").alias("leaves"))
             .select("leaves")
             .to_series()
@@ -93,7 +93,7 @@ def precision_recall(
         pairs_per_model.append(model_pairs)
 
     validation_pairs, validation_net_count, validation_leaves = process_judgements(
-        pl.from_arrow(judgements), pl.from_arrow(expansion)
+        judgements, expansion
     )
     leaves_per_set.append(validation_leaves)
 

{matchbox_db-0.6.2 → matchbox_db-0.6.3.dev53}/src/matchbox/common/factories/entities.py

@@ -521,7 +521,7 @@ def generate_entities(
 
 
 def probabilities_to_results_entities(
-    probabilities: pa.Table,
+    probabilities: pl.DataFrame,
     left_clusters: tuple[ClusterEntity, ...],
     right_clusters: tuple[ClusterEntity, ...] | None = None,
     threshold: float | int = 0,
@@ -547,7 +547,7 @@ def probabilities_to_results_entities(
             djs.add(entity)
 
     # Add edges to the disjoint set
-    for record in probabilities.to_pylist():
+    for record in probabilities.to_dicts():
         if record["probability"] >= threshold:
             djs.union(
                 left_lookup[record["left_id"]],

{matchbox_db-0.6.2 → matchbox_db-0.6.3.dev53}/src/matchbox/common/factories/models.py

@@ -76,12 +76,12 @@ add_model_class(MockDeduper)
 add_model_class(MockLinker)
 
 
-def component_report(all_nodes: list[Any], table: pa.Table) -> dict:
+def component_report(all_nodes: list[Any], table: pl.DataFrame) -> dict:
     """Fast reporting on connected components using rustworkx.
 
     Args:
         all_nodes: list of identities of inputs being matched
-        table: PyArrow table with 'left', 'right' columns
+        table: Polars dataframe with 'left', 'right' columns
 
     Returns:
         dictionary containing basic component statistics
@@ -252,7 +252,7 @@ def generate_dummy_probabilities(
     num_components: int,
     total_rows: int | None = None,
     seed: int = 42,
-) -> pa.Table:
+) -> pl.DataFrame:
     """Generate dummy Arrow probabilities data with guaranteed isolated components.
 
     While much of the factory system uses generate_entity_probabilities, this function
@@ -269,7 +269,7 @@ def generate_dummy_probabilities(
         seed: Random seed for reproducibility
 
     Returns:
-        PyArrow Table with 'left_id', 'right_id', and 'probability' columns
+        Polars dataframe with 'left_id', 'right_id', and 'probability' columns
     """
     # Validate inputs
     deduplicate = False
@@ -419,14 +419,9 @@ def generate_dummy_probabilities(
     # Convert to arrays
     lefts, rights, probs = zip(*all_edges, strict=True)
 
-    # Create PyArrow arrays
-    left_array = pa.array(lefts, type=pa.uint64())
-    right_array = pa.array(rights, type=pa.uint64())
-    prob_array = pa.array(probs, type=pa.uint8())
-
-    return pa.table(
-        [left_array, right_array, prob_array],
-        names=["left_id", "right_id", "probability"],
+    return pl.DataFrame(
+        {"left_id": lefts, "right_id": rights, "probability": probs},
+        schema={"left_id": pl.UInt64, "right_id": pl.UInt64, "probability": pl.UInt8},
     )
 
 
@@ -436,7 +431,7 @@ def generate_entity_probabilities(
     source_entities: frozenset[SourceEntity],
     prob_range: tuple[float, float] = (0.8, 1.0),
     seed: int = 42,
-) -> pa.Table:
+) -> pl.DataFrame:
     """Generate probabilities that will recover entity relationships.
 
     Compares ClusterEntity objects against ground truth SourceEntities by checking
@@ -530,27 +525,9 @@ def generate_entity_probabilities(
 
     # If no edges were generated, return empty table with correct schema
     if not edges:
-        return pa.table(
-            [
-                pa.array([], type=pa.uint64()),
-                pa.array([], type=pa.uint64()),
-                pa.array([], type=pa.uint8()),
-            ],
-            schema=SCHEMA_RESULTS,
-        )
+        return pl.DataFrame(schema=pl.Schema(SCHEMA_RESULTS))
 
-    # Convert to arrays
-    lefts, rights, probs = zip(*edges, strict=False)
-
-    # Create PyArrow arrays
-    left_array = pa.array(lefts, type=pa.uint64())
-    right_array = pa.array(rights, type=pa.uint64())
-    prob_array = pa.array(probs, type=pa.uint8())
-
-    return pa.table(
-        [left_array, right_array, prob_array],
-        schema=SCHEMA_RESULTS,
-    )
+    return pl.DataFrame(edges, orient="row", schema=pl.Schema(SCHEMA_RESULTS))
 
 
 class ModelTestkit(BaseModel):
@@ -565,7 +542,7 @@ class ModelTestkit(BaseModel):
     right_data: pa.Table | None
     right_query: Query | None
     right_clusters: dict[int, ClusterEntity] | None
-    probabilities: pa.Table
+    probabilities: pl.DataFrame
 
     _entities: tuple[ClusterEntity, ...]
     _threshold: int

{matchbox_db-0.6.2 → matchbox_db-0.6.3.dev53}/src/matchbox/common/factories/scenarios.py

@@ -6,6 +6,7 @@ from pathlib import Path
 from typing import Any, Literal
 
 import pyarrow as pa
+from polars.testing import assert_frame_equal
 from pydantic_settings import BaseSettings, SettingsConfigDict
 from sqlalchemy import Engine
 
@@ -171,7 +172,9 @@ def create_dedupe_scenario(
 
         # Add to backend and DAG
         backend.insert_resolution(resolution=model_testkit.model.to_resolution())
-        backend.insert_model_data(name=name, results=model_testkit.probabilities)
+        backend.insert_model_data(
+            name=name, results=model_testkit.probabilities.to_arrow()
+        )
         dag.add_model(model_testkit)
 
     return dag
@@ -216,7 +219,9 @@ def create_probabilistic_dedupe_scenario(
 
         # Add to backend and DAG
         backend.insert_resolution(resolution=model_testkit.model.to_resolution())
-        backend.insert_model_data(name=name, results=model_testkit.probabilities)
+        backend.insert_model_data(
+            name=name, results=model_testkit.probabilities.to_arrow()
+        )
         backend.set_model_truth(name=name, truth=50)
         dag.add_model(model_testkit)
 
@@ -271,7 +276,9 @@ def create_link_scenario(
 
     # Add to backend and DAG
     backend.insert_resolution(resolution=crn_duns_model.model.to_resolution())
-    backend.insert_model_data(name=crn_duns_name, results=crn_duns_model.probabilities)
+    backend.insert_model_data(
+        name=crn_duns_name, results=crn_duns_model.probabilities.to_arrow()
+    )
     dag.add_model(crn_duns_model)
 
     # Create CRN-CDMS link
@@ -299,7 +306,9 @@ def create_link_scenario(
 
     # Add to backend and DAG
     backend.insert_resolution(resolution=crn_cdms_model.model.to_resolution())
-    backend.insert_model_data(name=crn_cdms_name, results=crn_cdms_model.probabilities)
+    backend.insert_model_data(
+        name=crn_cdms_name, results=crn_cdms_model.probabilities.to_arrow()
+    )
     backend.set_model_truth(name=crn_cdms_name, truth=75)
     dag.add_model(crn_cdms_model)
 
@@ -344,7 +353,7 @@ def create_link_scenario(
     # Add to backend and DAG
     backend.insert_resolution(resolution=final_join_model.model.to_resolution())
     backend.insert_model_data(
-        name=final_join_name, results=final_join_model.probabilities
+        name=final_join_name, results=final_join_model.probabilities.to_arrow()
     )
     dag.add_model(final_join_model)
 
@@ -422,15 +431,17 @@ def create_alt_dedupe_scenario(
             seed=seed,
         )
 
-        assert model_testkit1.probabilities.num_rows > 0
-        assert model_testkit1.probabilities == model_testkit2.probabilities
+        assert len(model_testkit1.probabilities) > 0
+        assert_frame_equal(model_testkit1.probabilities, model_testkit2.probabilities)
 
         for model, threshold in ((model_testkit1, 50), (model_testkit2, 75)):
             model.threshold = threshold
 
             # Add both models to backend and DAG
             backend.insert_resolution(resolution=model.model.to_resolution())
-            backend.insert_model_data(name=model.name, results=model.probabilities)
+            backend.insert_model_data(
+                name=model.name, results=model.probabilities.to_arrow()
+            )
             backend.set_model_truth(name=model.name, truth=threshold)
 
             # Add to DAG
@@ -509,7 +520,7 @@ def create_convergent_scenario(
             seed=seed,
         )
 
-        assert model_testkit.probabilities.num_rows > 0
+        assert len(model_testkit.probabilities) > 0
 
         # Add to DAG
         dag.add_model(model_testkit)

{matchbox_db-0.6.2 → matchbox_db-0.6.3.dev53}/src/matchbox/common/factories/sources.py

@@ -202,7 +202,7 @@ class LinkedSourcesTestkit(BaseModel):
 
     def diff_results(
         self,
-        probabilities: pa.Table,
+        probabilities: pl.DataFrame,
         sources: list[SourceResolutionName],
         left_clusters: tuple[ClusterEntity, ...],
         right_clusters: tuple[ClusterEntity, ...] | None = None,