"""Power Query ``Text.*`` column expressions."""
from __future__ import annotations
import re
from typing import Sequence, Union
from pyspark.sql import Column
from pyspark.sql import functions as F
from fabrictools.transform.text import norm_text
__all__ = ["Text"]
def _as_column(expr: Union[Column, str]) -> Column:
return F.lit(expr) if isinstance(expr, str) else expr
[docs]
class Text:
"""Namespace for Power Query ``Text.*`` functions.
All methods return Spark ``Column`` expressions for use in
:py:meth:`fabrictools.powerquery.table.Table.AddColumn` or
:py:meth:`fabrictools.powerquery.table.Table.TransformColumns`.
"""
[docs]
@staticmethod
def Clean(expr: Union[Column, str]) -> Column:
"""Lowercase string with control chars stripped and spaces removed (Power Query ``Text.Clean``).
Delegates to :py:func:`fabrictools.norm_text`.
:param expr: Spark column or string literal.
:type expr: ~pyspark.sql.Column | str
:returns: Transformed column expression.
:rtype: ~pyspark.sql.Column
.. rubric:: Example
>>> from fabrictools import Table, Text # doctest: +SKIP
>>> from pyspark.sql import functions as F # doctest: +SKIP
>>> df = Table.AddColumn(df, "Project No. 2", Text.Clean(F.col("Project No."))) # doctest: +SKIP
"""
return norm_text(expr)
[docs]
@staticmethod
def Select(expr: Union[Column, str], allowed: Sequence[str]) -> Column:
"""Keep only characters present in ``allowed`` (Power Query ``Text.Select``).
:param expr: Spark column or string literal.
:param allowed: Characters to keep (e.g. digits and punctuation for amounts).
:type expr: ~pyspark.sql.Column | str
:type allowed: collections.abc.Sequence[str]
:returns: Filtered string column expression.
:rtype: ~pyspark.sql.Column
.. rubric:: Example
>>> from fabrictools import Text # doctest: +SKIP
>>> from pyspark.sql import functions as F # doctest: +SKIP
>>> digits = Text.Select(F.col("amount_raw"), list("0123456789,.-")) # doctest: +SKIP
"""
import re
c = _as_column(expr)
as_str = F.coalesce(c.cast("string"), F.lit(""))
if not allowed:
return F.lit("")
escaped = "".join(re.escape(ch) for ch in allowed)
return F.regexp_replace(as_str, f"[^{escaped}]", "")
[docs]
@staticmethod
def Trim(expr: Union[Column, str]) -> Column:
"""Remove leading and trailing whitespace (Power Query ``Text.Trim``).
:param expr: Spark column or string literal.
:type expr: ~pyspark.sql.Column | str
:returns: Trimmed string column expression.
:rtype: ~pyspark.sql.Column
.. rubric:: Example
>>> from fabrictools import Text # doctest: +SKIP
>>> from pyspark.sql import functions as F # doctest: +SKIP
>>> df.withColumn("name", Text.Trim(F.col("name"))) # doctest: +SKIP
"""
return F.trim(_as_column(expr).cast("string"))
[docs]
@staticmethod
def Lower(expr: Union[Column, str]) -> Column:
"""Convert to lowercase (Power Query ``Text.Lower``).
:param expr: Spark column or string literal.
:type expr: ~pyspark.sql.Column | str
:returns: Lowercase column expression.
:rtype: ~pyspark.sql.Column
.. rubric:: Example
>>> from fabrictools import Text # doctest: +SKIP
>>> from pyspark.sql import functions as F # doctest: +SKIP
>>> df.withColumn("code", Text.Lower(F.col("code"))) # doctest: +SKIP
"""
return F.lower(_as_column(expr).cast("string"))
[docs]
@staticmethod
def Upper(expr: Union[Column, str]) -> Column:
"""Convert to uppercase (Power Query ``Text.Upper``).
:param expr: Spark column or string literal.
:type expr: ~pyspark.sql.Column | str
:returns: Uppercase column expression.
:rtype: ~pyspark.sql.Column
.. rubric:: Example
>>> from fabrictools import Text # doctest: +SKIP
>>> from pyspark.sql import functions as F # doctest: +SKIP
>>> df.withColumn("code", Text.Upper(F.col("code"))) # doctest: +SKIP
"""
return F.upper(_as_column(expr).cast("string"))
[docs]
@staticmethod
def Proper(expr: Union[Column, str]) -> Column:
"""Title-case each word (Power Query ``Text.Proper``).
:param expr: Spark column or string literal.
:type expr: ~pyspark.sql.Column | str
:returns: Title-cased column expression.
:rtype: ~pyspark.sql.Column
.. rubric:: Example
>>> from fabrictools import Text # doctest: +SKIP
>>> from pyspark.sql import functions as F # doctest: +SKIP
>>> df.withColumn("client", Text.Proper(F.col("client"))) # doctest: +SKIP
"""
return F.initcap(_as_column(expr).cast("string"))
[docs]
@staticmethod
def Combine(texts: Sequence[Union[Column, str]], separator: str = "") -> Column:
"""Concatenate text values with a separator (Power Query ``Text.Combine``).
:param texts: Column expressions or string literals to join.
:param separator: Delimiter between parts (default empty string).
:type texts: collections.abc.Sequence[~pyspark.sql.Column | str]
:type separator: str
:returns: Concatenated string column expression.
:rtype: ~pyspark.sql.Column
.. rubric:: Example
>>> from fabrictools import Text # doctest: +SKIP
>>> from pyspark.sql import functions as F # doctest: +SKIP
>>> full = Text.Combine([F.col("first"), F.col("last")], " ") # doctest: +SKIP
"""
cols = [_as_column(t).cast("string") for t in texts]
return F.concat_ws(separator, *cols)
[docs]
@staticmethod
def From(expr: Union[Column, str]) -> Column:
"""Cast a value to text (Power Query ``Text.From``).
:param expr: Spark column or string literal.
:type expr: ~pyspark.sql.Column | str
:returns: Non-null string column expression (null becomes ``""``).
:rtype: ~pyspark.sql.Column
.. rubric:: Example
>>> from fabrictools import Text # doctest: +SKIP
>>> from pyspark.sql import functions as F # doctest: +SKIP
>>> df.withColumn("id_text", Text.From(F.col("id"))) # doctest: +SKIP
"""
return F.coalesce(_as_column(expr).cast("string"), F.lit(""))
[docs]
@staticmethod
def StartsWith(expr: Union[Column, str], substring: str) -> Column:
"""Test whether text starts with a prefix (Power Query ``Text.StartsWith``).
:param expr: Spark column or string literal.
:param substring: Prefix to look for (case-sensitive).
:type expr: ~pyspark.sql.Column | str
:type substring: str
:returns: Boolean column expression (null when ``expr`` is null).
:rtype: ~pyspark.sql.Column
.. rubric:: Example
>>> from fabrictools import Table, Text # doctest: +SKIP
>>> from pyspark.sql import functions as F # doctest: +SKIP
>>> df = Table.SelectRows(df, ~Text.StartsWith(F.col("Article"), "SG")) # doctest: +SKIP
"""
return _as_column(expr).cast("string").startswith(substring)
[docs]
@staticmethod
def AfterDelimiter(
expr: Union[Column, str],
delimiter: str,
index: int = 0,
) -> Column:
"""Return the substring after a delimiter occurrence (Power Query ``Text.AfterDelimiter``).
``index`` is 0-based from the start: ``0`` is after the first delimiter,
``1`` after the second, and so on. If the occurrence is missing, returns
an empty string. Null input stays null.
:param expr: Spark column or string literal.
:param delimiter: Delimiter to search for.
:param index: 0-based occurrence index (Power Query ``index``).
:type expr: ~pyspark.sql.Column | str
:type delimiter: str
:type index: int
:returns: String column expression.
:rtype: ~pyspark.sql.Column
.. rubric:: Example
>>> from fabrictools import Table, Text # doctest: +SKIP
>>> from pyspark.sql import functions as F # doctest: +SKIP
>>> df = Table.AddColumn( # doctest: +SKIP
... df, "Chef de projet", Text.AfterDelimiter(F.col("Chef projet"), " ", 1),
... )
"""
c = _as_column(expr)
occurrence = int(index)
if occurrence < 0:
raise ValueError("Text.AfterDelimiter index must be >= 0")
limit = occurrence + 2
rest = F.element_at(
F.split(c.cast("string"), re.escape(delimiter), limit),
limit,
)
return F.when(c.isNull(), F.lit(None).cast("string")).otherwise(
F.coalesce(rest, F.lit(""))
)