Optional dependencies. Better log for polars joins

Author: raffaelemancuso

docs/source/examples.rst

@@ -52,14 +52,16 @@ Changed 2/344 (0.58%) values.
 polars
 ------
 
-Example of usage
+Let' import the libraries and create a logger:
 
 >>> import polars as pl
 >>> import polars.selectors as cs
 >>> import raffalib
 >>> import raffalib.polars
->>> 
 >>> logger = raffalib.create_logger(rich=False, fmt="{message}")
+
+Let's load a dataset
+
 >>> df = pl.read_csv("https://raw.githubusercontent.com/allisonhorst/palmerpenguins/refs/heads/main/inst/extdata/penguins.csv")
 >>> df = df.raffa.startlog(clone=True).raffa.replace_string_with_null("NA").raffa.endlog(timeit=False)
 Changed 19/2,752 (0.69%) values.
@@ -77,13 +79,42 @@ shape: (5, 8)
 │ Adelie  ┆ Torgersen ┆ null           ┆ null          ┆ null              ┆ null        ┆ null   ┆ 2,007 │
 │ Adelie  ┆ Torgersen ┆ 36.700001      ┆ 19.299999     ┆ 193.0             ┆ 3,450.0     ┆ female ┆ 2,007 │
 └─────────┴───────────┴────────────────┴───────────────┴───────────────────┴─────────────┴────────┴───────┘
+
+Let's see some data wrangling on this data: remove rows with null values, filtering certain values, or selecting certain columns:
+
 >>> _ = df.raffa.startlog().drop_nulls(subset=["bill_depth_mm"]).raffa.endlog(timeit=False)
 Removed 2/344 (0.58%) rows.
 >>> _ = df.raffa.startlog().filter(pl.col("species")=="Adelie").raffa.endlog(timeit=False)
 Removed 192/344 (55.81%) rows.
 >>> _ = df.raffa.startlog().select(pl.exclude(["bill_length_mm", "bill_depth_mm"])).raffa.endlog(timeit=False)
 Removed 2/8 (25.00%) columns.
+
+Operations that do not change the shape of the DataFrame but only its values requires `clone=True`
+in the `startlog` call to clone the entire initial dataframe for further comparison:
+
 >>> _ = df.raffa.startlog().fill_null(0).raffa.endlog(timeit=False)
 Shape is the same. No value-level comparison done because clone=False was used in startlog().
 >>> _ = df.raffa.startlog(clone=True).fill_null("0").raffa.endlog(timeit=False)
-Changed 11/2,752 (0.40%) values.
+Changed 11/2,752 (0.40%) values.
+
+Let's see an example with joins:
+
+>>> df1 = pl.DataFrame({"left_a": ["A", "B", "B", "C", "D"], "left_b": ["a", "b1", "b2", "c", "d"]})
+>>> df2 = pl.DataFrame({"right_a": ["A", "A", "A", "B", "E"], "right_d": ["alpha1", "alpha2", "alpha3", "beta", "gamma"]})
+>>> _ = df1.raffa.join(df2, left_on="left_a", right_on="right_a", how="full", keep_row_index=False)
+Total rows in output table: 8
+From left only: 2/8 (25.00%)
+From right only: 1/8 (12.50%)
+From both: 5/8 (62.50%) (left dups 3, right dups 2)
+>>> _ = df1.raffa.join(df2, left_on="left_a", right_on="right_a", how="left", keep_row_index=False)
+Total rows in output table: 7
+From left only: 2/7 (28.57%)
+From right only: 0/7 (0.00%)
+From both: 5/7 (71.43%) (left dups 3, right dups 2)
+
+Filtering joins are automatically detected:
+
+>>> _ = df1.raffa.join(df2, left_on="left_a", right_on="right_a", how="semi", keep_row_index=False)
+Detected filtering join. Rows variation -2/5 (-40.00%), total rows after join: 3/5 (60.00%)
+>>> _ = df1.raffa.join(df2, left_on="left_a", right_on="right_a", how="anti", keep_row_index=False)
+Detected filtering join. Rows variation -3/5 (-60.00%), total rows after join: 2/5 (40.00%)

pyproject.toml

@@ -18,17 +18,14 @@ dependencies = [
     "tqdm>=4.67.1",
 ]
 
-[build-system]
-requires = ["uv_build"]
-build-backend = "uv_build"
-
-[dependency-groups]
+[project.optional-dependencies]
 pandas = [
 	"pandas>=3.0.1",
 ]
 polars = [
 	"polars>=1.38.1",
-    "polars-config-meta>=0.3.1",
+	"polars-config-meta>=0.3.1",
+	"polars-permute>=0.1.6",
 ]
 bibliometrics = [
     "beautifulsoup4>=4.14.3",
@@ -54,3 +51,7 @@ docs = [
 web = [
     "selenium>=4.41.0",
 ]
+
+[build-system]
+requires = ["uv_build"]
+build-backend = "uv_build"

src/raffalib/export_docx.py

@@ -9,13 +9,13 @@
 
 @dataclass
 class DocxOptions:
-    page_height = 210
-    page_width = 297
-    landscape = False
-    left_margin = 25.4
-    right_margin = 25.4
-    top_margin = 25.4
-    bottom_margin = 25.4
+    page_height:float = 210
+    page_width:float = 297
+    landscape:bool = False
+    left_margin:float = 25.4
+    right_margin:float = 25.4
+    top_margin:float = 25.4
+    bottom_margin:float = 25.4
     heading_text:str|None = None
     heading_font_name:str = "Aptos"
     heading_font_size = 12
@@ -77,30 +77,17 @@ def prepare_docx(options:DocxOptions):
         paragraph_format = style.paragraph_format
         paragraph_format.space_before = Pt(options.heading_space_before)
         paragraph_format.space_after = Pt(options.heading_space_after)
-    paragraph_format.line_spacing_rule = WD_LINE_SPACING.SINGLE
+        paragraph_format.line_spacing_rule = WD_LINE_SPACING.SINGLE
 
     return doc
 
 # Additional workaround to make autofit actually work
-# See: https://github.com/python-openxml/python-docx/issues/209#issuecomment-566128709
-def set_autofit(doc: Document) -> Document:
-    """
-    Hotfix for autofit.
-    """
-    for t_idx, table in enumerate(doc.tables):
-        doc.tables[t_idx].autofit = True
-        doc.tables[t_idx].allow_autofit = True
-        doc.tables[t_idx]._tblPr.xpath("./w:tblW")[0].attrib[
-            "{http://schemas.openxmlformats.org/wordprocessingml/2006/main}type"
-        ] = "auto"
-        for row_idx, r_val in enumerate(doc.tables[t_idx].rows):
-            for cell_idx, c_val in enumerate(
-                doc.tables[t_idx].rows[row_idx].cells
-            ):
-                doc.tables[t_idx].rows[row_idx].cells[
-                    cell_idx
-                ]._tc.tcPr.tcW.type = "auto"
-                doc.tables[t_idx].rows[row_idx].cells[
-                    cell_idx
-                ]._tc.tcPr.tcW.w = 0
-    return doc
+# See: https://github.com/python-openxml/python-docx/issues/209#issuecomment-566128709 
+def table_autofit_hotfix(table):
+    for column in table.columns:
+        for cell in column.cells:
+            tc = cell._tc
+            tcPr = tc.get_or_add_tcPr()
+            tcW = tcPr.get_or_add_tcW()
+            tcW.type = 'auto'
+    return table

src/raffalib/pandas.py

@@ -22,7 +22,7 @@
 import logging
 from natsort import natsorted
 from pathlib import Path
-from .export_docx import prepare_docx, set_autofit, DocxOptions
+from .export_docx import prepare_docx, DocxOptions, table_autofit_hotfix
 
 logger = logging.getLogger(__name__)
 
@@ -292,7 +292,7 @@ def to_docx(
                     t.cell(i + 1, j).text = str(df.values[i, j])
 
         if autofit:
-            doc = set_autofit(doc)
+            t = table_autofit_hotfix(t)
 
         # Set table style
         # See: https://github.com/python-openxml/python-docx/issues/9
@@ -301,10 +301,3 @@ def to_docx(
         # save the doc
         doc.save(outfp)
 
-
-if __name__ == "__main__":
-    # Run with `uv run python -P pandas.py -v`
-    import doctest
-
-    # doctest.testmod(optionflags=doctest.REPORT_NDIFF)
-    doctest.testmod()

src/raffalib/polars.py

@@ -17,6 +17,7 @@
 import polars as pl
 import polars.selectors as cs
 import polars_config_meta
+import polars_permute
 from enum import Enum, auto
 import logging
 import copy
@@ -136,6 +137,7 @@ def startlog(self, clone=False) -> pl.DataFrame:
         :return: The same DataFrame for piping.
         :rtype: pl.DataFrame
         """
+        
         self._df.config_meta.set(initial_shape=self._df.shape, start_time=time.perf_counter_ns())
         if clone:
             self._df.config_meta.set(initial_df=self._df.clone())
@@ -243,12 +245,29 @@ def crosstab(
 
         return ct.with_columns(values / options[perc] * 100)
 
-    def join(self, df2, *args, keep_source: bool = False, log_col_changes=False, **kwargs):
+    def join(self, df2:pl.DataFrame, *args, keep_row_index: bool = False, log_col_changes=False, **kwargs):
+        """
+        Wrapper around `pl.DataFrame.join` to log join operations.
+
+        :param df2: The dataframe on the right of the join
+        :type df2: pl.DataFrame
+        :param *args: Additional position arguments passed to `pl.DataFrame.join`
+        :type *args: Any
+        :param keep_row_index: Whether to keep columns indicating the row index of the source table in the output table
+        :type keep_row_index: bool
+        :param log_col_changes: Whether to log column changes
+        :type log_col_changes: bool
+        :param *kwargs: Additional keyword arguments passed to `pl.DataFrame.join`
+        :type *kwargs: Any
+        :return: The joined DataFrame
+        :rtype: pl.DataFrame
+        """
+        
         left_col = "source_left"
         right_col = "source_right"
         # Get DataFrame to join, add source column
-        df1 = self._df.with_columns(pl.lit(True).alias(left_col))
-        df2 = df2.with_columns(pl.lit(True).alias(right_col))
+        df1 = self._df.with_row_index(left_col)
+        df2 = df2.with_row_index(right_col)
         # Join DataFrames
         joined = df1.join(df2, *args, **kwargs)
         n_rows_joined = joined.shape[0]
@@ -269,21 +288,21 @@ def join(self, df2, *args, keep_source: bool = False, log_col_changes=False, **k
             joined = joined.drop([left_col])
             return joined
         # Detect how many rows in the output table are present in the input tables
-        joined = joined.with_columns(
-            pl.col([left_col, right_col]).fill_null(False)
-        )
-        n_both = joined.filter(pl.col(left_col), pl.col(right_col)).shape[0]
+        joined_both = joined.filter(~pl.col(left_col).is_null(), ~pl.col(right_col).is_null())
+        n_both = joined_both.shape[0]
+        n_left_dups = joined_both.get_column(left_col).is_duplicated().sum()
+        n_right_dups = joined_both.get_column(right_col).is_duplicated().sum()
         n_left_only = joined.filter(
-            pl.col(left_col), ~pl.col(right_col)
+            ~pl.col(left_col).is_null(), pl.col(right_col).is_null()
         ).shape[0]
         n_right_only = joined.filter(
-            ~pl.col(left_col), pl.col(right_col)
+            pl.col(left_col).is_null(), ~pl.col(right_col).is_null()
         ).shape[0]
         # Log rows information
         msg = f"Total rows in output table: {n_rows_joined:,d}\n"
-        msg += f"\tFrom left: {n_left_only:,d} ({n_left_only / n_rows_joined:.2%})\n"
-        msg += f"\tFrom right: {n_right_only:,d} ({n_right_only / n_rows_joined:.2%})\n"
-        msg += f"\tFrom both: {n_both:,d} ({n_both / n_rows_joined:.2%})\n"
+        msg += f"From left only: {n_left_only:,d}/{n_rows_joined:,d} ({n_left_only / n_rows_joined:.2%})\n"
+        msg += f"From right only: {n_right_only:,d}/{n_rows_joined:,d} ({n_right_only / n_rows_joined:.2%})\n"
+        msg += f"From both: {n_both:,d}/{n_rows_joined:,d} ({n_both / n_rows_joined:.2%}) (left dups {n_left_dups}, right dups {n_right_dups})\n"
         # Detect added and removed columns
         if log_col_changes:
             cols_out = set(joined.columns)
@@ -292,21 +311,11 @@ def join(self, df2, *args, keep_source: bool = False, log_col_changes=False, **k
             msg += f"Columns in output table not in left table: {cols_out - cols_left})\n"
             msg += f"Columns in output table not in right table: {cols_out - cols_right})"
         logger.info(msg)
-        # Add a column that indicate where that row comes from
-        if keep_source:
-            joined = joined.with_columns(
-                pl.when(pl.col.source_left & pl.col.source_right)
-                .then(pl.lit("both"))
-                .when(pl.col.source_left & ~pl.col.source_right)
-                .then(pl.lit("left_only"))
-                .when(~pl.col.source_left & pl.col.source_right)
-                .then(pl.lit("right_only"))
-                .otherwise(pl.lit("error"))
-                .alias("source")
-            )
-            joined = joined.permute.append(["source"])
-        # Drop source columns and return
-        joined = joined.drop([left_col, right_col])
+        # Drop row indices
+        if not keep_row_index:
+            joined = joined.drop([left_col, right_col])
+        else:
+            joined = joined.permute.append([left_col, right_col])
         return joined
 
 if __name__ == "__main__":