From 9ff1e22afca6fde0ddb898fcabb3091bd0b2c9a1 Mon Sep 17 00:00:00 2001 From: Murray Vanwyk Date: Fri, 9 Aug 2024 20:49:57 +0200 Subject: [PATCH 1/3] feat: segment stats calc now uses duckdb to improve performance --- poetry.lock | 57 +++++++++++++++++- pyproject.toml | 1 + pyretailscience/options.py | 6 ++ pyretailscience/segmentation.py | 103 ++++++++++++++++++++------------ tests/test_segmentation.py | 30 +++++++--- 5 files changed, 148 insertions(+), 49 deletions(-) diff --git a/poetry.lock b/poetry.lock index fe88daf5..1bb648c1 100644 --- a/poetry.lock +++ b/poetry.lock @@ -777,6 +777,61 @@ files = [ {file = "distlib-0.3.8.tar.gz", hash = "sha256:1530ea13e350031b6312d8580ddb6b27a104275a31106523b8f123787f494f64"}, ] +[[package]] +name = "duckdb" +version = "1.0.0" +description = "DuckDB in-process database" +optional = false +python-versions = ">=3.7.0" +files = [ + {file = "duckdb-1.0.0-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:4a8ce2d1f9e1c23b9bab3ae4ca7997e9822e21563ff8f646992663f66d050211"}, + {file = "duckdb-1.0.0-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:19797670f20f430196e48d25d082a264b66150c264c1e8eae8e22c64c2c5f3f5"}, + {file = "duckdb-1.0.0-cp310-cp310-macosx_12_0_x86_64.whl", hash = "sha256:b71c342090fe117b35d866a91ad6bffce61cd6ff3e0cff4003f93fc1506da0d8"}, + {file = "duckdb-1.0.0-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:25dd69f44ad212c35ae2ea736b0e643ea2b70f204b8dff483af1491b0e2a4cec"}, + {file = "duckdb-1.0.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8da5f293ecb4f99daa9a9352c5fd1312a6ab02b464653a0c3a25ab7065c45d4d"}, + {file = "duckdb-1.0.0-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3207936da9967ddbb60644ec291eb934d5819b08169bc35d08b2dedbe7068c60"}, + {file = "duckdb-1.0.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:1128d6c9c33e883b1f5df6b57c1eb46b7ab1baf2650912d77ee769aaa05111f9"}, + {file = "duckdb-1.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:02310d263474d0ac238646677feff47190ffb82544c018b2ff732a4cb462c6ef"}, + {file = "duckdb-1.0.0-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:75586791ab2702719c284157b65ecefe12d0cca9041da474391896ddd9aa71a4"}, + {file = "duckdb-1.0.0-cp311-cp311-macosx_12_0_universal2.whl", hash = "sha256:83bb415fc7994e641344f3489e40430ce083b78963cb1057bf714ac3a58da3ba"}, + {file = "duckdb-1.0.0-cp311-cp311-macosx_12_0_x86_64.whl", hash = "sha256:bee2e0b415074e84c5a2cefd91f6b5ebeb4283e7196ba4ef65175a7cef298b57"}, + {file = "duckdb-1.0.0-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:fa5a4110d2a499312609544ad0be61e85a5cdad90e5b6d75ad16b300bf075b90"}, + {file = "duckdb-1.0.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5fa389e6a382d4707b5f3d1bc2087895925ebb92b77e9fe3bfb23c9b98372fdc"}, + {file = "duckdb-1.0.0-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7ede6f5277dd851f1a4586b0c78dc93f6c26da45e12b23ee0e88c76519cbdbe0"}, + {file = "duckdb-1.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:0b88cdbc0d5c3e3d7545a341784dc6cafd90fc035f17b2f04bf1e870c68456e5"}, + {file = "duckdb-1.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:fd1693cdd15375156f7fff4745debc14e5c54928589f67b87fb8eace9880c370"}, + {file = "duckdb-1.0.0-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:c65a7fe8a8ce21b985356ee3ec0c3d3b3b2234e288e64b4cfb03356dbe6e5583"}, + {file = "duckdb-1.0.0-cp312-cp312-macosx_12_0_universal2.whl", hash = "sha256:e5a8eda554379b3a43b07bad00968acc14dd3e518c9fbe8f128b484cf95e3d16"}, + {file = "duckdb-1.0.0-cp312-cp312-macosx_12_0_x86_64.whl", hash = "sha256:a1b6acdd54c4a7b43bd7cb584975a1b2ff88ea1a31607a2b734b17960e7d3088"}, + {file = "duckdb-1.0.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a677bb1b6a8e7cab4a19874249d8144296e6e39dae38fce66a80f26d15e670df"}, + {file = "duckdb-1.0.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:752e9d412b0a2871bf615a2ede54be494c6dc289d076974eefbf3af28129c759"}, + {file = "duckdb-1.0.0-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3aadb99d098c5e32d00dc09421bc63a47134a6a0de9d7cd6abf21780b678663c"}, + {file = "duckdb-1.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:83b7091d4da3e9301c4f9378833f5ffe934fb1ad2b387b439ee067b2c10c8bb0"}, + {file = "duckdb-1.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:6a8058d0148b544694cb5ea331db44f6c2a00a7b03776cc4dd1470735c3d5ff7"}, + {file = "duckdb-1.0.0-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e40cb20e5ee19d44bc66ec99969af791702a049079dc5f248c33b1c56af055f4"}, + {file = "duckdb-1.0.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d7bce1bc0de9af9f47328e24e6e7e39da30093179b1c031897c042dd94a59c8e"}, + {file = "duckdb-1.0.0-cp37-cp37m-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8355507f7a04bc0a3666958f4414a58e06141d603e91c0fa5a7c50e49867fb6d"}, + {file = "duckdb-1.0.0-cp37-cp37m-musllinux_1_2_x86_64.whl", hash = "sha256:39f1a46f5a45ad2886dc9b02ce5b484f437f90de66c327f86606d9ba4479d475"}, + {file = "duckdb-1.0.0-cp37-cp37m-win_amd64.whl", hash = "sha256:a6d29ba477b27ae41676b62c8fae8d04ee7cbe458127a44f6049888231ca58fa"}, + {file = "duckdb-1.0.0-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:1bea713c1925918714328da76e79a1f7651b2b503511498ccf5e007a7e67d49e"}, + {file = "duckdb-1.0.0-cp38-cp38-macosx_12_0_universal2.whl", hash = "sha256:bfe67f3bcf181edbf6f918b8c963eb060e6aa26697d86590da4edc5707205450"}, + {file = "duckdb-1.0.0-cp38-cp38-macosx_12_0_x86_64.whl", hash = "sha256:dbc6093a75242f002be1d96a6ace3fdf1d002c813e67baff52112e899de9292f"}, + {file = "duckdb-1.0.0-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ba1881a2b11c507cee18f8fd9ef10100be066fddaa2c20fba1f9a664245cd6d8"}, + {file = "duckdb-1.0.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:445d0bb35087c522705c724a75f9f1c13f1eb017305b694d2686218d653c8142"}, + {file = "duckdb-1.0.0-cp38-cp38-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:224553432e84432ffb9684f33206572477049b371ce68cc313a01e214f2fbdda"}, + {file = "duckdb-1.0.0-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:d3914032e47c4e76636ad986d466b63fdea65e37be8a6dfc484ed3f462c4fde4"}, + {file = "duckdb-1.0.0-cp38-cp38-win_amd64.whl", hash = "sha256:af9128a2eb7e1bb50cd2c2020d825fb2946fdad0a2558920cd5411d998999334"}, + {file = "duckdb-1.0.0-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:dd2659a5dbc0df0de68f617a605bf12fe4da85ba24f67c08730984a0892087e8"}, + {file = "duckdb-1.0.0-cp39-cp39-macosx_12_0_universal2.whl", hash = "sha256:ac5a4afb0bc20725e734e0b2c17e99a274de4801aff0d4e765d276b99dad6d90"}, + {file = "duckdb-1.0.0-cp39-cp39-macosx_12_0_x86_64.whl", hash = "sha256:2c5a53bee3668d6e84c0536164589d5127b23d298e4c443d83f55e4150fafe61"}, + {file = "duckdb-1.0.0-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b980713244d7708b25ee0a73de0c65f0e5521c47a0e907f5e1b933d79d972ef6"}, + {file = "duckdb-1.0.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:21cbd4f9fe7b7a56eff96c3f4d6778770dd370469ca2212eddbae5dd63749db5"}, + {file = "duckdb-1.0.0-cp39-cp39-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ed228167c5d49888c5ef36f6f9cbf65011c2daf9dcb53ea8aa7a041ce567b3e4"}, + {file = "duckdb-1.0.0-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:46d8395fbcea7231fd5032a250b673cc99352fef349b718a23dea2c0dd2b8dec"}, + {file = "duckdb-1.0.0-cp39-cp39-win_amd64.whl", hash = "sha256:6ad1fc1a4d57e7616944166a5f9417bdbca1ea65c490797e3786e3a42e162d8a"}, + {file = "duckdb-1.0.0.tar.gz", hash = "sha256:a2a059b77bc7d5b76ae9d88e267372deff19c291048d59450c431e166233d453"}, +] + [[package]] name = "entrypoints" version = "0.4" @@ -4313,4 +4368,4 @@ pillow = "*" [metadata] lock-version = "2.0" python-versions = ">=3.10,<3.12" -content-hash = "4cae23465878937d41808ef42c3e0a962f9363c4c2717d62875043f7af295ea6" +content-hash = "631ecd192c79c770bbc7be8fbb8b4627d913ba03685fcaaf2cf0d252a55ff154" diff --git a/pyproject.toml b/pyproject.toml index 4c162ce7..520b9f33 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ scipy = "^1.13.0" scikit-learn = "^1.4.2" matplotlib-set-diagrams = "~0.0.2" toml = "^0.10.2" +duckdb = "^1.0.0" [tool.poetry.group.dev.dependencies] pytest = "^8.0.0" diff --git a/pyretailscience/options.py b/pyretailscience/options.py index d354a38c..15804f70 100644 --- a/pyretailscience/options.py +++ b/pyretailscience/options.py @@ -53,6 +53,9 @@ def __init__(self) -> None: # Calculated columns "column.calc.price_per_unit": "price_per_unit", "column.calc.units_per_transaction": "units_per_transaction", + "column.calc.spend_per_customer": "spend_per_customer", + "column.calc.spend_per_transaction": "spend_per_transaction", + "column.calc.transactions_per_customer": "transactions_per_customer", # Abbreviation suffix "column.suffix.count": "cnt", "column.suffix.percent": "pct", @@ -86,6 +89,9 @@ def __init__(self) -> None: # Calculated columns "column.calc.price_per_unit": "The name of the column containing the price per unit.", "column.calc.units_per_transaction": "The name of the column containing the units per transaction.", + "column.calc.spend_per_customer": "The name of the column containing the spend per customer.", + "column.calc.spend_per_transaction": "The name of the column containing the spend per transaction.", + "column.calc.transactions_per_customer": "The name of the column containing the transactions per customer.", # Abbreviation suffixes "column.suffix.count": "The suffix to use for count columns.", "column.suffix.percent": "The suffix to use for percentage columns.", diff --git a/pyretailscience/segmentation.py b/pyretailscience/segmentation.py index a020dfca..96842223 100644 --- a/pyretailscience/segmentation.py +++ b/pyretailscience/segmentation.py @@ -2,7 +2,9 @@ from typing import Literal +import duckdb import pandas as pd +from duckdb import DuckDBPyRelation from matplotlib.axes import Axes, SubplotBase import pyretailscience.style.graph_utils as gu @@ -204,12 +206,14 @@ def __init__( class SegTransactionStats: """Calculates transaction statistics by segment.""" - def __init__(self, df: pd.DataFrame, segment_col: str = "segment_id") -> None: + def __init__(self, data: pd.DataFrame | DuckDBPyRelation, segment_col: str = "segment_id") -> None: """Calculates transaction statistics by segment. Args: - df (pd.DataFrame): A dataframe with the transaction data. The dataframe must comply with the - TransactionItemLevelContract or the TransactionLevelContract. + data (pd.DataFrame | DuckDBPyRelation): The transaction data. The dataframe must contain the columns + customer_id, unit_spend and transaction_id. If the dataframe contains the column unit_quantity, then + the columns unit_spend and unit_quantity are used to calculate the price_per_unit and + units_per_transaction. segment_col (str, optional): The column to use for the segmentation. Defaults to "segment_id". Raises: @@ -223,54 +227,75 @@ def __init__(self, df: pd.DataFrame, segment_col: str = "segment_id") -> None: get_option("column.transaction_id"), segment_col, ] - if get_option("column.unit_quantity") in df.columns: + if get_option("column.unit_quantity") in data.columns: required_cols.append(get_option("column.unit_quantity")) - contract = CustomContract( - df, - basic_expectations=build_expected_columns(columns=required_cols), - extended_expectations=build_non_null_columns(columns=required_cols), - ) - if contract.validate() is False: - msg = f"The dataframe requires the columns {required_cols} and they must be non-null" + missing_cols = [col for col in required_cols if col not in data.columns] + + if len(missing_cols) > 0: + msg = f"The following columns are required but missing: {missing_cols}" raise ValueError(msg) self.segment_col = segment_col - self.df = self._calc_seg_stats(df, segment_col) + self.df = self._calc_seg_stats(data, segment_col) @staticmethod - def _calc_seg_stats(df: pd.DataFrame, segment_col: str) -> pd.DataFrame: - aggs = { - get_option("column.agg.unit_spend"): (get_option("column.unit_spend"), "sum"), - get_option("column.agg.transaction_id"): (get_option("column.transaction_id"), "nunique"), - get_option("column.agg.customer_id"): (get_option("column.customer_id"), "nunique"), - } - total_aggs = { - get_option("column.agg.unit_spend"): [df[get_option("column.unit_spend")].sum()], - get_option("column.agg.transaction_id"): [df[get_option("column.transaction_id")].nunique()], - get_option("column.agg.customer_id"): [df[get_option("column.customer_id")].nunique()], - } - if get_option("column.unit_quantity") in df.columns: - aggs[get_option("column.agg.unit_quantity")] = (get_option("column.unit_quantity"), "sum") - total_aggs[get_option("column.agg.unit_quantity")] = [df[get_option("column.unit_quantity")].sum()] - - stats_df = pd.concat( - [ - df.groupby(segment_col).agg(**aggs), - pd.DataFrame(total_aggs, index=["total"]), - ], - ) + def _calc_seg_stats(data: pd.DataFrame | DuckDBPyRelation, segment_col: str) -> pd.DataFrame: + """Calculates the transaction statistics by segment. + + Args: + data (DuckDBPyRelation): The transaction data. + segment_col (str): The column to use for the segmentation. - if get_option("column.unit_quantity") in df.columns: - stats_df[get_option("column.calc.price_per_unit")] = ( - stats_df[get_option("column.agg.unit_spend")] / stats_df[get_option("column.agg.unit_quantity")] + Returns: + pd.DataFrame: The transaction statistics by segment. + + """ + if isinstance(data, pd.DataFrame): + data = duckdb.from_df(data) + + base_aggs = [ + f"SUM({get_option('column.unit_spend')}) as {get_option('column.agg.unit_spend')},", + f"COUNT(DISTINCT {get_option('column.transaction_id')}) as {get_option('column.agg.transaction_id')},", + f"COUNT(DISTINCT {get_option('column.customer_id')}) as {get_option('column.agg.customer_id')},", + ] + + total_customers = data.aggregate("COUNT(DISTINCT customer_id)").fetchone()[0] + return_cols = [ + "*,", + f"{get_option('column.agg.unit_spend')} / {get_option('column.agg.customer_id')} ", + f"as {get_option('column.calc.spend_per_customer')},", + f"{get_option('column.agg.unit_spend')} / {get_option('column.agg.transaction_id')} ", + f"as {get_option('column.calc.spend_per_transaction')},", + f"{get_option('column.agg.transaction_id')} / {get_option('column.agg.customer_id')} ", + f"as {get_option('column.calc.transactions_per_customer')},", + f"{get_option('column.agg.customer_id')} / {total_customers}", + f"as customers_{get_option('column.suffix.percent')},", + ] + + if get_option("column.unit_quantity") in data.columns: + base_aggs.append( + f"SUM({get_option('column.unit_quantity')})::bigint as {get_option('column.agg.unit_quantity')},", ) - stats_df[get_option("column.calc.units_per_transaction")] = ( - stats_df[get_option("column.agg.unit_quantity")] / stats_df[get_option("column.agg.transaction_id")] + return_cols.extend( + [ + f"({get_option('column.agg.unit_spend')} / {get_option('column.agg.unit_quantity')}) ", + f"as {get_option('column.calc.price_per_unit')},", + f"({get_option('column.agg.unit_quantity')} / {get_option('column.agg.transaction_id')}) ", + f"as {get_option('column.calc.units_per_transaction')},", + ], ) - return stats_df + segment_stats = data.aggregate(f"{segment_col} as segment_name," + "".join(base_aggs)) + total_stats = data.aggregate("'Total' as segment_name," + "".join(base_aggs)) + final_stats_df = segment_stats.union(total_stats).select("".join(return_cols)).df() + final_stats_df = final_stats_df.set_index("segment_name").sort_index() + + # Make sure Total is the last row + desired_index_sort = final_stats_df.index.drop("Total").tolist() + ["Total"] # noqa: RUF005 + + return final_stats_df.reindex(desired_index_sort) def plot( self, diff --git a/tests/test_segmentation.py b/tests/test_segmentation.py index f9ec0292..a7f5354c 100644 --- a/tests/test_segmentation.py +++ b/tests/test_segmentation.py @@ -27,15 +27,19 @@ def test_correctly_calculates_revenue_transactions_customers_per_segment(self, b """Test that the method correctly calculates at the transaction-item level.""" expected_output = pd.DataFrame( { - get_option("column.agg.unit_spend"): [500, 500, 1000], + "segment_name": ["A", "B", "Total"], + get_option("column.agg.unit_spend"): [500.0, 500.0, 1000.0], get_option("column.agg.transaction_id"): [3, 2, 5], get_option("column.agg.customer_id"): [3, 2, 5], get_option("column.agg.unit_quantity"): [50, 50, 100], + get_option("column.calc.spend_per_customer"): [166.666667, 250.0, 200.0], + get_option("column.calc.spend_per_transaction"): [166.666667, 250.0, 200.0], + get_option("column.calc.transactions_per_customer"): [1.0, 1.0, 1.0], + f"customers_{get_option('column.suffix.percent')}": [0.6, 0.4, 1.0], get_option("column.calc.price_per_unit"): [10.0, 10.0, 10.0], get_option("column.calc.units_per_transaction"): [16.666667, 25.0, 20.0], }, - index=["A", "B", "total"], - ) + ).set_index("segment_name") segment_stats = SegTransactionStats._calc_seg_stats(base_df, "segment_id") pd.testing.assert_frame_equal(segment_stats, expected_output) @@ -53,12 +57,16 @@ def test_correctly_calculates_revenue_transactions_customers(self): expected_output = pd.DataFrame( { - get_option("column.agg.unit_spend"): [500, 500, 1000], + "segment_name": ["A", "B", "Total"], + get_option("column.agg.unit_spend"): [500.0, 500.0, 1000.0], get_option("column.agg.transaction_id"): [3, 2, 5], get_option("column.agg.customer_id"): [3, 2, 5], + get_option("column.calc.spend_per_customer"): [166.666667, 250.0, 200.0], + get_option("column.calc.spend_per_transaction"): [166.666667, 250.0, 200.0], + get_option("column.calc.transactions_per_customer"): [1.0, 1.0, 1.0], + f"customers_{get_option('column.suffix.percent')}": [0.6, 0.4, 1.0], }, - index=["A", "B", "total"], - ) + ).set_index("segment_name") segment_stats = SegTransactionStats._calc_seg_stats(df, "segment_id") pd.testing.assert_frame_equal(segment_stats, expected_output) @@ -77,15 +85,19 @@ def test_handles_dataframe_with_one_segment(self, base_df): expected_output = pd.DataFrame( { - get_option("column.agg.unit_spend"): [1000, 1000], + "segment_name": ["A", "Total"], + get_option("column.agg.unit_spend"): [1000.0, 1000.0], get_option("column.agg.transaction_id"): [5, 5], get_option("column.agg.customer_id"): [5, 5], get_option("column.agg.unit_quantity"): [100, 100], + get_option("column.calc.spend_per_customer"): [200.0, 200.0], + get_option("column.calc.spend_per_transaction"): [200.0, 200.0], + get_option("column.calc.transactions_per_customer"): [1.0, 1.0], + f"customers_{get_option('column.suffix.percent')}": [1.0, 1.0], get_option("column.calc.price_per_unit"): [10.0, 10.0], get_option("column.calc.units_per_transaction"): [20.0, 20.0], }, - index=["A", "total"], - ) + ).set_index("segment_name") segment_stats = SegTransactionStats._calc_seg_stats(df, "segment_id") pd.testing.assert_frame_equal(segment_stats, expected_output) From b5e7564d53a843e37ea7c59e68eb55ebb844ef6c Mon Sep 17 00:00:00 2001 From: mvanwyk <2493311+mvanwyk@users.noreply.github.com> Date: Fri, 9 Aug 2024 21:23:52 +0200 Subject: [PATCH 2/3] Update pyretailscience/segmentation.py Co-authored-by: codiumai-pr-agent-pro[bot] <151058649+codiumai-pr-agent-pro[bot]@users.noreply.github.com> --- pyretailscience/segmentation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyretailscience/segmentation.py b/pyretailscience/segmentation.py index 96842223..2634af0f 100644 --- a/pyretailscience/segmentation.py +++ b/pyretailscience/segmentation.py @@ -230,7 +230,7 @@ def __init__(self, data: pd.DataFrame | DuckDBPyRelation, segment_col: str = "se if get_option("column.unit_quantity") in data.columns: required_cols.append(get_option("column.unit_quantity")) - missing_cols = [col for col in required_cols if col not in data.columns] + missing_cols = set(required_cols) - set(data.columns) if len(missing_cols) > 0: msg = f"The following columns are required but missing: {missing_cols}" From 697f7616bd1bf8453379bb77b8ded5e3c1033fab Mon Sep 17 00:00:00 2001 From: mvanwyk <2493311+mvanwyk@users.noreply.github.com> Date: Fri, 9 Aug 2024 21:24:22 +0200 Subject: [PATCH 3/3] Update pyretailscience/segmentation.py Co-authored-by: codiumai-pr-agent-pro[bot] <151058649+codiumai-pr-agent-pro[bot]@users.noreply.github.com> --- pyretailscience/segmentation.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyretailscience/segmentation.py b/pyretailscience/segmentation.py index 2634af0f..4b418171 100644 --- a/pyretailscience/segmentation.py +++ b/pyretailscience/segmentation.py @@ -254,6 +254,8 @@ def _calc_seg_stats(data: pd.DataFrame | DuckDBPyRelation, segment_col: str) -> """ if isinstance(data, pd.DataFrame): data = duckdb.from_df(data) + elif not isinstance(data, DuckDBPyRelation): + raise TypeError("data must be either a pandas DataFrame or a DuckDBPyRelation") base_aggs = [ f"SUM({get_option('column.unit_spend')}) as {get_option('column.agg.unit_spend')},",