Method isin

dask/dataframe/dask_expr/_collection.py:725–777

(self, values)

    @derived_from(pd.DataFrame)
    def isin(self, values):
        if isinstance(self, DataFrame):
            # DataFrame.isin does weird alignment stuff
            bad_types = (FrameBase, pd.Series, pd.DataFrame)
        else:
            bad_types = (FrameBase,)
        if isinstance(values, bad_types):
            if (
                isinstance(values, FrameBase)
                and values.ndim == 1
                and values.npartitions == 1
            ):
                # Can broadcast
                return new_collection(expr.Isin(self, values=values))
            raise NotImplementedError("Passing a %r to `isin`" % typename(type(values)))

        # We wrap values in a delayed for two reasons:
        # - avoid serializing data in every task
        # - avoid cost of traversal of large list in optimizations
        if isinstance(values, list):
            # Motivated by https://github.com/dask/dask/issues/9411.  This appears to be
            # caused by https://github.com/dask/distributed/issues/6368, and further
            # exacerbated by the fact that the list contains duplicates.  This is a patch until
            # we can create a better fix for Serialization.
            try:
                values = list(set(values))
            except TypeError:
                pass
            if not any(is_dask_collection(v) for v in values):
                # Avoid always passing a numpy array of object dtype
                inferred_type = pd.api.types.infer_dtype(values, skipna=False)
                object_like = {
                    "mixed-integer",
                    "decimal",
                    "categorical",
                    "time",
                    "period",
                    "mixed",
                    "unknown-array",
                }
                if inferred_type in object_like:
                    values = np.fromiter(values, dtype=object, count=len(values))
                else:
                    values = np.asarray(values)

        return new_collection(
            expr.Isin(
                self,
                values=expr._DelayedExpr(
                    delayed(values, name="delayed-" + _tokenize_deterministic(values))
                ),
            )
        )

    def _partitions(self, index):
        # Used by `partitions` for partition-wise slicing