55class DupeGrouper:
 56    """Top-level entrypoint for grouping duplicates
 57
 58    This class handles initialisation of a dataframe, dispatching appropriately
 59    given the supported dataframe libraries (e.g. Pandas). An instance of this
 60    class can then accept a variety of strategies for deduplication and
 61    grouping.
 62
 63    Upon initialisation, `DupeGrouper` sets a new column, usually `"group_id"`
 64    — but you can control this by setting an environment variable `GROUP_ID` at
 65    runtime. The group_id is a monotonically increasing, numeric id column
 66    starting at 1 to the length of the dataframe provided.
 67    """
 68
 69    def __init__(
 70        self,
 71        df: DataFrameLike,
 72        spark_session: SparkSession | None = None,
 73        id: str | None = None,
 74    ):
 75        self._df: WrappedDataFrame = _wrap(df, id)
 76        self._strategy_manager = _StrategyManager()
 77        self._spark_session = spark_session
 78        self._id = id
 79
 80    @singledispatchmethod
 81    def _call_strategy_deduper(
 82        self,
 83        strategy: DeduplicationStrategy | tuple[typing.Callable, typing.Any],
 84        attr: str,
 85    ):
 86        """Dispatch the appropriate strategy deduplication method.
 87
 88        If the strategy is an instance of a dupegrouper `DeduplicationStrategy`
 89        the strategy will have been added as such, with it's parameters. In the
 90        case of a custom implementation of a Callable, passed as a tuple, we
 91        pass this *directly* to the `Custom` class and initialise that.
 92
 93        Args:
 94            strategy: A `dupegrouper` deduplication strategy or a tuple
 95                containing a (customer) callable and its parameters.
 96            attr: The attribute used for deduplication.
 97
 98        Returns:
 99            A deduplicated dataframe
100
101        Raises:
102            NotImplementedError.
103        """
104        del attr  # Unused
105
106        raise NotImplementedError(f"Unsupported strategy: {type(strategy)}")
107
108    @_call_strategy_deduper.register(DeduplicationStrategy)
109    def _(self, strategy, attr) -> WrappedDataFrame:
110        return strategy.with_frame(self._df).dedupe(attr)
111
112    @_call_strategy_deduper.register(tuple)
113    def _(self, strategy: tuple[typing.Callable, typing.Any], attr) -> WrappedDataFrame:
114        func, kwargs = strategy
115        return Custom(func, attr, **kwargs).with_frame(self._df).dedupe()
116
117    @singledispatchmethod
118    def _dedupe(
119        self,
120        attr: str | None,
121        strategies: StrategyMapCollection,
122    ):
123        """Dispatch the appropriate deduplication logic.
124
125        If strategies have been added individually, they are stored under a
126        "default" key and retrived as such when the public `.dedupe` method is
127        called _with_ the attribute label. In the case of having added
128        strategies in one go with a direct dict (mapping) object, the attribute
129        label is first extracted from strategy collection dictionary keys.
130        Upon completing deduplication the strategy collection is wiped for
131        (any) subsequent deduplication.
132
133        Args:
134            attr: The attribute used for deduplication; or None in the case
135                of strategies being a mapping object
136
137        Returns:
138            None; internal `_df` attribute is updated.
139
140        Raises:
141            NotImplementedError.
142        """
143        del strategies  # Unused
144        raise NotImplementedError(f"Unsupported attribute type: {type(attr)}")
145
146    @_dedupe.register(str)
147    def _(self, attr, strategies):
148        for strategy in strategies["default"]:
149            self._df = self._call_strategy_deduper(strategy, attr)
150
151    @_dedupe.register(NoneType)
152    def _(self, attr, strategies):
153        del attr  # Unused
154        for attr, strategies in strategies.items():
155            for strategy in strategies:
156                self._df = self._call_strategy_deduper(strategy, attr)
157
158    def _dedupe_spark(self, attr: str | None, strategies: StrategyMapCollection):
159        """Spark specific deduplication helper
160
161        Maps dataframe partitions to be processed via the RDD API yielding low-
162        level list[Rows], which are then post-processed back to a dataframe.
163
164        Args:
165            attr: The attribute to deduplicate.
166            strategies: the collection of strategies
167        Retuns:
168            Instance's _df attribute is updated
169        """
170        id = typing.cast(str, self._id)
171        id_type = typing.cast(DataType, PYSPARK_TYPES.get(dict(self._df.dtypes).get(id)))  # type: ignore
172
173        deduped_rdd = self._df.rdd.mapPartitions(
174            lambda partition_iter: _process_partition(partition_iter, strategies, id, attr)
175        )
176
177        if GROUP_ID in self._df.columns:
178            schema = StructType(self._df.schema.fields)
179        else:
180            schema = StructType(self._df.schema.fields + [StructField(GROUP_ID, id_type, True)])
181
182        self._df = WrappedSparkDataFrame(
183            typing.cast(SparkSession, self._spark_session).createDataFrame(deduped_rdd, schema=schema), id
184        )
185
186    # PUBLIC API:
187
188    @singledispatchmethod
189    def add_strategy(self, strategy: DeduplicationStrategy | tuple | StrategyMapCollection):
190        """
191        Add a strategy to the strategy manager.
192
193        Instances of `DeduplicationStrategy` or tuple are added to the
194        "default" key. Mapping objects update the manager directly
195
196        Args:
197            strategy: A deduplication strategy, tuple, or strategy collection
198                (mapping) to add.
199
200        Returns:
201            self is updated
202
203        Raises:
204            NotImplementedError
205        """
206        raise NotImplementedError(f"Unsupported strategy: {type(strategy)}")
207
208    @add_strategy.register(DeduplicationStrategy)
209    @add_strategy.register(tuple)
210    def _(self, strategy):
211        self._strategy_manager.add("default", strategy)
212
213    @add_strategy.register(dict)
214    def _(self, strategy: StrategyMapCollection):
215        for attr, strat_list in strategy.items():
216            for strat in strat_list:
217                self._strategy_manager.add(attr, strat)
218
219    def dedupe(self, attr: str | None = None):
220        """dedupe, and group, the data based on the provided attribute
221
222        Args:
223            attr: The attribute to deduplicate. If strategies have been added
224                as a mapping object, this must not passed, as the keys of the
225                mapping object will be used instead
226        """
227        strategies = self._strategy_manager.get()
228
229        if isinstance(self._df, WrappedSparkDataFrame):
230            self._dedupe_spark(attr, strategies)
231        else:
232            self._dedupe(attr, strategies)
233
234        self._strategy_manager.reset()
235
236    @property
237    def strategies(self) -> None | tuple[str, ...] | dict[str, tuple[str, ...]]:
238        """
239        Returns the strategies currently stored in the strategy manager.
240
241        If no strategies are stored, returns `None`. Otherwise, returns a tuple
242        of strategy names or a dictionary mapping attributes to their
243        respective strategies.
244
245        Returns:
246            The stored strategies, formatted
247        """
248        strategies = self._strategy_manager.get()
249        if not strategies:
250            return None
251
252        def parse_strategies(dict_values):
253            return tuple(
254                [
255                    (vx[0].__name__ if isinstance(vx, tuple) else vx.__class__.__name__)
256                    #
257                    for vx in dict_values
258                ]
259            )
260
261        if "default" in strategies:
262            return tuple([parse_strategies(v) for _, v in strategies.items()])[0]
263        return {k: parse_strategies(v) for k, v in strategies.items()}
264
265    @property
266    def df(self) -> DataFrameLike:
267        return self._df.unwrap()