Source code for xpark.dataset.processors.text_cleaner

from __future__ import annotations

import re
from typing import TYPE_CHECKING

from xpark.dataset.datatype import DataType
from xpark.dataset.expressions import BatchColumnClassProtocol, udf
from xpark.dataset.import_utils import lazy_import

if TYPE_CHECKING:
    import pyarrow as pa
else:
    pa = lazy_import("pyarrow", rename="pa")

#: Matches URLs (any scheme), ``www.`` prefixed hosts, and bare domains with a path.
#:
#: Adapted from data-juicer's ``clean_links_mapper``:
#: https://github.com/modelscope/data-juicer/blob/main/data_juicer/ops/mapper/clean_links_mapper.py
#:
#: Key properties:
#: - Supports balanced parentheses inside the URL (one level of nesting).
#: - The trailing character must NOT be sentence punctuation (ASCII or CJK),
#:   so punctuation that follows a URL in natural-language text is preserved.
PATTERN_URI = (
    r"\b("
    r"(?:[a-z][\w-]+:(?:\/{1,3}|[a-z0-9%])"
    r"|www\d{0,3}[.]"
    r"|[a-z0-9.\-]+[.][a-z]{2,4}\/)"
    r"(?:[^\s()<>]+|\(([^\s()<>]+|(\([^\s()<>]+\)))*\))+"
    r"(?:\(([^\s()<>]+|(\([^\s()<>]+\)))*\)"
    r"|[^\s`!()\[\]{};:\'\".,<>?\u00ab\u00bb\u201c\u201d\u2018\u2019])"
    r")"
)

#: Matches common email addresses
PATTERN_EMAIL = r"\b[a-zA-Z0-9._%+\-]+@[a-zA-Z0-9.\-]+\.[a-zA-Z]{2,}\b"

#: Matches common phone numbers (international +xx format and 11-digit mobile numbers)
PATTERN_PHONE = (
    r"\b(?:\+?\d{1,3}[\s\-]?)?"
    r"(?:\(?\d{2,4}\)?[\s\-]?)?"
    r"\d{3,4}[\s\-]?\d{4}\b"
)

#: Matches IPv4 addresses with each octet in 0-255
PATTERN_IPV4 = (
    r"\b(?:(?:25[0-5]|2[0-4]\d|1\d{2}|[1-9]?\d)\.){3}"
    r"(?:25[0-5]|2[0-4]\d|1\d{2}|[1-9]?\d)\b"
)

# Mapping from built-in pattern name to regex string
_BUILTIN_PATTERNS: dict[str, str] = {
    "uri": PATTERN_URI,
    "email": PATTERN_EMAIL,
    "phone": PATTERN_PHONE,
    "ipv4": PATTERN_IPV4,
}




@udf(return_dtype=DataType.string())
class TextPatternCleaner(BatchColumnClassProtocol):
    """
    Supports two usage modes:

    1. **Built-in patterns**: Pass a built-in pattern name (string); use a list for multiple.
       Available names: ``"uri"``, ``"email"``, ``"phone"``, ``"ipv4"``.

    2. **Custom regex**: Pass a regex string or list directly. All patterns are merged
       into a single ``|``-joined regex and applied in one substitution pass.

    Both modes can be mixed together.

    Args:
        patterns: Built-in pattern name(s) or custom regex string(s); accepts a single
            string or a list.
        replacement: Replacement string. Defaults to ``""`` (delete matched text).
        flags: Flags passed to :func:`re.compile`. Defaults to ``re.IGNORECASE``.

    Examples:
        .. code-block:: python

            from xpark.dataset.expressions import col
            from xpark.dataset import from_items
            from xpark.dataset.processors.text_cleaner import TextPatternCleaner

            ds = from_items(["Visit https://example.com, email: foo@bar.com"])

            ds = ds.with_column(
                "cleaned",
                TextPatternCleaner(patterns=["uri", "email"], replacement="[REDACTED]")
                .options(num_workers={"CPU": 1}, batch_size=10)
                .with_column(col("item")),
            )

            ds = ds.with_column(
                "cleaned",
                TextPatternCleaner(patterns=r"\\d{4}-\\d{4}-\\d{4}-\\d{4}", replacement="[CARD]")
                .options(num_workers={"CPU": 1}, batch_size=10)
                .with_column(col("item")),
            )

            print(ds.take(1))
    """

    def __init__(
        self,
        patterns: str | list[str],
        replacement: str = "",
        flags: int = re.IGNORECASE,
    ):
        if isinstance(patterns, str):
            patterns = [patterns]

        # Resolve built-in names; pass through unknown names as raw regex
        resolved = [_BUILTIN_PATTERNS.get(p, p) for p in patterns]
        merged = "|".join(f"(?:{p})" for p in resolved)

        self.pattern = merged
        self._pattern = re.compile(merged, flags)
        self.replacement = replacement

    def __call__(self, batch: pa.ChunkedArray) -> pa.Array:
        texts = batch.to_pylist()
        cleaned = [self._pattern.sub(self.replacement, t) if t is not None else None for t in texts]
        return pa.array(cleaned, type=pa.string())