| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| |
| """Spark-compatible function bindings. |
| |
| These functions mirror the semantics of their Apache Spark counterparts |
| exactly. Some override DataFusion built-ins (``substring`` is 1-indexed, |
| ``concat`` propagates NULL, ``round`` uses HALF_UP rounding, etc.), which is |
| why they live in a separate namespace rather than replacing the defaults. |
| |
| For DataFrame use, import this module and call functions directly. For SQL |
| use, call :py:meth:`datafusion.SessionContext.enable_spark_functions` to |
| register the Spark UDFs by name (overriding any built-ins with matching |
| names) before issuing SQL queries. |
| """ |
| |
| from __future__ import annotations |
| |
| from typing import TYPE_CHECKING, Any |
| |
| import pyarrow as pa |
| |
| from datafusion._internal import functions as _functions |
| from datafusion.expr import ( |
| Expr, |
| _to_raw_expr, |
| coerce_to_expr, |
| sort_list_to_raw_sort_list, |
| ) |
| |
| if TYPE_CHECKING: |
| from datafusion.common import NullTreatment |
| from datafusion.expr import SortKey |
| |
| _f = _functions.spark |
| |
| # Reused int32 literal so optional-arg defaults don't rebuild it per call. |
| _ZERO_I32 = Expr.literal(pa.scalar(0, type=pa.int32())) |
| |
| |
| def _filter_raw(filter: Expr | None) -> Any: |
| return filter.expr if filter is not None else None |
| |
| |
| def _coerce_i32(value: Expr | int | None) -> Expr | None: |
| """Coerce a native ``int`` to an int32 literal, passing ``Expr``/``None`` through. |
| |
| Several Spark datetime and interval builders require 32-bit integer |
| inputs, so a bare ``int`` must become an int32 literal rather than the |
| int64 default that :meth:`Expr.literal` would produce. |
| """ |
| if value is None or isinstance(value, Expr): |
| return value |
| return Expr.literal(pa.scalar(value, type=pa.int32())) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Aggregate functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def avg( |
| col: Expr, |
| distinct: bool | None = None, |
| filter: Expr | None = None, |
| order_by: list[SortKey] | SortKey | None = None, |
| null_treatment: NullTreatment | None = None, |
| ) -> Expr: |
| """Spark ``avg``: returns the mean of a numeric column. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"a": [1.0, 2.0, 3.0]}) |
| >>> r = df.aggregate( |
| ... [], [dfn.functions.spark.avg(dfn.col("a")).alias("v")]) |
| >>> r.collect_column("v")[0].as_py() |
| 2.0 |
| """ |
| return Expr( |
| _f.avg( |
| col.expr, |
| distinct=distinct, |
| filter=_filter_raw(filter), |
| order_by=sort_list_to_raw_sort_list(order_by), |
| null_treatment=null_treatment, |
| ) |
| ) |
| |
| |
| def try_sum( |
| col: Expr, |
| distinct: bool | None = None, |
| filter: Expr | None = None, |
| order_by: list[SortKey] | SortKey | None = None, |
| null_treatment: NullTreatment | None = None, |
| ) -> Expr: |
| """Spark ``try_sum``: sum that returns NULL on overflow. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"a": [1, 2, 3]}) |
| >>> r = df.aggregate( |
| ... [], [dfn.functions.spark.try_sum(dfn.col("a")).alias("v")]) |
| >>> r.collect_column("v")[0].as_py() |
| 6 |
| """ |
| return Expr( |
| _f.try_sum( |
| col.expr, |
| distinct=distinct, |
| filter=_filter_raw(filter), |
| order_by=sort_list_to_raw_sort_list(order_by), |
| null_treatment=null_treatment, |
| ) |
| ) |
| |
| |
| def collect_list( |
| col: Expr, |
| distinct: bool | None = None, |
| filter: Expr | None = None, |
| order_by: list[SortKey] | SortKey | None = None, |
| null_treatment: NullTreatment | None = None, |
| ) -> Expr: |
| """Spark ``collect_list``: collect values into an array (preserves dups). |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"a": [1, 2, 2]}) |
| >>> r = df.aggregate( |
| ... [], [dfn.functions.spark.collect_list(dfn.col("a")).alias("v")]) |
| >>> sorted(r.collect_column("v")[0].as_py()) |
| [1, 2, 2] |
| """ |
| return Expr( |
| _f.collect_list( |
| col.expr, |
| distinct=distinct, |
| filter=_filter_raw(filter), |
| order_by=sort_list_to_raw_sort_list(order_by), |
| null_treatment=null_treatment, |
| ) |
| ) |
| |
| |
| def collect_set( |
| col: Expr, |
| distinct: bool | None = None, |
| filter: Expr | None = None, |
| order_by: list[SortKey] | SortKey | None = None, |
| null_treatment: NullTreatment | None = None, |
| ) -> Expr: |
| """Spark ``collect_set``: collect distinct values into an array. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"a": [1, 2, 2, 3]}) |
| >>> r = df.aggregate( |
| ... [], [dfn.functions.spark.collect_set(dfn.col("a")).alias("v")]) |
| >>> sorted(r.collect_column("v")[0].as_py()) |
| [1, 2, 3] |
| """ |
| return Expr( |
| _f.collect_set( |
| col.expr, |
| distinct=distinct, |
| filter=_filter_raw(filter), |
| order_by=sort_list_to_raw_sort_list(order_by), |
| null_treatment=null_treatment, |
| ) |
| ) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Array functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def array_contains(col: Expr, value: Expr | Any) -> Expr: |
| """Spark ``array_contains``: true if the array contains the element. |
| |
| ``value`` accepts a native Python literal or an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.array_contains( |
| ... dfn.functions.spark.array(dfn.lit(1), dfn.lit(2)), 1 |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| True |
| """ |
| return Expr(_f.array_contains(col.expr, coerce_to_expr(value).expr)) |
| |
| |
| def array(*cols: Expr) -> Expr: |
| """Spark ``array``: builds an array from the given elements. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.array( |
| ... dfn.lit(1), dfn.lit(2), dfn.lit(3) |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| [1, 2, 3] |
| """ |
| return Expr(_f.array(*[c.expr for c in cols])) |
| |
| |
| def shuffle(col: Expr, seed: int | None = None) -> Expr: |
| """Spark ``shuffle``: returns a random permutation of the input array. |
| |
| ``seed`` is accepted for pyspark parity but is not yet wired through the |
| Rust binding; passing a non-``None`` value raises ``NotImplementedError``. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.shuffle( |
| ... dfn.functions.spark.array(dfn.lit(1), dfn.lit(2), dfn.lit(3)) |
| ... ).alias("v") |
| ... ) |
| >>> sorted(r.collect_column("v")[0].as_py()) |
| [1, 2, 3] |
| """ |
| if seed is not None: |
| msg = "shuffle(seed=...) is not yet supported by the Spark UDF binding" |
| raise NotImplementedError(msg) |
| return Expr(_f.shuffle(col.expr)) |
| |
| |
| def array_repeat(col: Expr, count: Expr | int) -> Expr: |
| """Spark ``array_repeat``: array of ``element`` repeated ``count`` times. |
| |
| ``count`` accepts a native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.array_repeat(dfn.lit("a"), 3).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| ['a', 'a', 'a'] |
| """ |
| return Expr(_f.array_repeat(col.expr, coerce_to_expr(count).expr)) |
| |
| |
| def slice(x: Expr, start: Expr | int, length: Expr | int) -> Expr: |
| """Spark ``slice``: subset of the array from 1-indexed ``start`` with ``length``. |
| |
| Negative ``start`` counts from the end. ``start`` and ``length`` accept |
| native ``int`` values or :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.slice( |
| ... dfn.functions.spark.array( |
| ... dfn.lit(1), dfn.lit(2), dfn.lit(3), dfn.lit(4)), |
| ... 2, 2, |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| [2, 3] |
| """ |
| return Expr( |
| _f.slice(x.expr, coerce_to_expr(start).expr, coerce_to_expr(length).expr) |
| ) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Bitmap functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def bitmap_count(col: Expr) -> Expr: |
| r"""Spark ``bitmap_count``: number of set bits in a bitmap. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.bitmap_count(dfn.lit(b"\xff")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 8 |
| """ |
| return Expr(_f.bitmap_count(col.expr)) |
| |
| |
| def bitmap_bit_position(col: Expr) -> Expr: |
| """Spark ``bitmap_bit_position``: bit position for a child expression. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.bitmap_bit_position(dfn.lit(15)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 14 |
| """ |
| return Expr(_f.bitmap_bit_position(col.expr)) |
| |
| |
| def bitmap_bucket_number(col: Expr) -> Expr: |
| """Spark ``bitmap_bucket_number``: bucket number for a child expression. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.bitmap_bucket_number(dfn.lit(15)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 1 |
| """ |
| return Expr(_f.bitmap_bucket_number(col.expr)) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Bitwise functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def bit_get(col: Expr, pos: Expr | str) -> Expr: |
| """Spark ``bit_get``: returns the bit (0 or 1) at ``pos``. |
| |
| A bare ``str`` ``pos`` is treated as a column name (matching pyspark), |
| not a literal; pass :func:`~datafusion.lit` for a literal position. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.bit_get(dfn.lit(5), dfn.lit(0)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 1 |
| """ |
| return Expr(_f.bit_get(col.expr, _to_raw_expr(pos))) |
| |
| |
| def bit_count(col: Expr) -> Expr: |
| """Spark ``bit_count``: number of bits set in the integer's binary form. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.bit_count(dfn.lit(7)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 3 |
| """ |
| return Expr(_f.bit_count(col.expr)) |
| |
| |
| def bitwise_not(col: Expr) -> Expr: |
| """Spark ``~``: bitwise NOT. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.bitwise_not(dfn.lit(0)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| -1 |
| """ |
| return Expr(_f.bitwise_not(col.expr)) |
| |
| |
| def shiftleft(col: Expr, numBits: Expr | int) -> Expr: # noqa: N803 |
| """Spark ``shiftleft``: ``value`` shifted left by ``shift`` bits. |
| |
| ``numBits`` accepts a native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.shiftleft(dfn.lit(1), 3).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 8 |
| """ |
| return Expr(_f.shiftleft(col.expr, coerce_to_expr(numBits).expr)) |
| |
| |
| def shiftright(col: Expr, numBits: Expr | int) -> Expr: # noqa: N803 |
| """Spark ``shiftright``: arithmetic right shift. |
| |
| ``numBits`` accepts a native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.shiftright(dfn.lit(8), 2).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 2 |
| """ |
| return Expr(_f.shiftright(col.expr, coerce_to_expr(numBits).expr)) |
| |
| |
| def shiftrightunsigned(col: Expr, numBits: Expr | int) -> Expr: # noqa: N803 |
| """Spark ``shiftrightunsigned``: logical (unsigned) right shift. |
| |
| ``numBits`` accepts a native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.shiftrightunsigned(dfn.lit(8), 2).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 2 |
| """ |
| return Expr(_f.shiftrightunsigned(col.expr, coerce_to_expr(numBits).expr)) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Collection / Conditional / Conversion |
| # --------------------------------------------------------------------------- |
| |
| |
| def size(col: Expr) -> Expr: |
| """Spark ``size``: length of an array or map. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.size( |
| ... dfn.functions.spark.array(dfn.lit(1), dfn.lit(2), dfn.lit(3)) |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| 3 |
| """ |
| return Expr(_f.size(col.expr)) |
| |
| |
| def if_(condition: Expr, if_true: Expr | Any, if_false: Expr | Any) -> Expr: |
| """Spark ``if``: returns ``if_true`` when ``condition`` is true, else ``if_false``. |
| |
| Exposed as ``if_`` because ``if`` is a Python keyword. ``if_true`` and |
| ``if_false`` accept native Python literals or :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.if_( |
| ... dfn.lit(2) > dfn.lit(1), "big", "small" |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| 'big' |
| """ |
| return Expr( |
| _f.if_( |
| condition.expr, coerce_to_expr(if_true).expr, coerce_to_expr(if_false).expr |
| ) |
| ) |
| |
| |
| def spark_cast(arg: Expr, type_str: Expr | str) -> Expr: |
| """Spark ``cast``: cast ``arg`` to the type named by ``type_str``. |
| |
| Uses Spark cast semantics (e.g. overflow returns NULL, not error). |
| ``type_str`` accepts a native ``str`` or an :class:`Expr`. |
| |
| Currently only supports casting numeric values to ``"timestamp"``. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.spark_cast( |
| ... dfn.lit(1579098645), "timestamp" |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py().isoformat() |
| '2020-01-15T14:30:45+00:00' |
| """ |
| return Expr(_f.spark_cast(arg.expr, coerce_to_expr(type_str).expr)) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Datetime functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def add_months(start: Expr, months: Expr | int) -> Expr: |
| """Spark ``add_months``: date + N months. |
| |
| ``months`` accepts a native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import date |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> d = dfn.lit(pa.scalar(date(2020, 1, 15), type=pa.date32())) |
| >>> r = df.select(dfn.functions.spark.add_months(d, 2).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.date(2020, 3, 15) |
| """ |
| return Expr(_f.add_months(start.expr, _coerce_i32(months).expr)) |
| |
| |
| def date_add(start: Expr, days: Expr | int) -> Expr: |
| """Spark ``date_add``: date + N days. |
| |
| ``days`` accepts a native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import date |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> d = dfn.lit(pa.scalar(date(2020, 1, 15), type=pa.date32())) |
| >>> r = df.select(dfn.functions.spark.date_add(d, 5).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.date(2020, 1, 20) |
| """ |
| return Expr(_f.date_add(start.expr, _coerce_i32(days).expr)) |
| |
| |
| def date_sub(start: Expr, days: Expr | int) -> Expr: |
| """Spark ``date_sub``: date - N days. |
| |
| ``days`` accepts a native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import date |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> d = dfn.lit(pa.scalar(date(2020, 1, 15), type=pa.date32())) |
| >>> r = df.select(dfn.functions.spark.date_sub(d, 5).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.date(2020, 1, 10) |
| """ |
| return Expr(_f.date_sub(start.expr, _coerce_i32(days).expr)) |
| |
| |
| def hour(col: Expr) -> Expr: |
| """Spark ``hour``: extract hour component of a timestamp. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import datetime |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> ts = dfn.lit( |
| ... pa.scalar(datetime(2020, 1, 15, 14, 30, 45), |
| ... type=pa.timestamp('us'))) |
| >>> r = df.select(dfn.functions.spark.hour(ts).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 14 |
| """ |
| return Expr(_f.hour(col.expr)) |
| |
| |
| def minute(col: Expr) -> Expr: |
| """Spark ``minute``: extract minute component of a timestamp. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import datetime |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> ts = dfn.lit( |
| ... pa.scalar(datetime(2020, 1, 15, 14, 30, 45), |
| ... type=pa.timestamp('us'))) |
| >>> r = df.select(dfn.functions.spark.minute(ts).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 30 |
| """ |
| return Expr(_f.minute(col.expr)) |
| |
| |
| def second(col: Expr) -> Expr: |
| """Spark ``second``: extract second component of a timestamp. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import datetime |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> ts = dfn.lit( |
| ... pa.scalar(datetime(2020, 1, 15, 14, 30, 45), |
| ... type=pa.timestamp('us'))) |
| >>> r = df.select(dfn.functions.spark.second(ts).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 45 |
| """ |
| return Expr(_f.second(col.expr)) |
| |
| |
| def last_day(col: Expr) -> Expr: |
| """Spark ``last_day``: last day of the month containing the date. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import date |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> d = dfn.lit(pa.scalar(date(2020, 1, 15), type=pa.date32())) |
| >>> r = df.select(dfn.functions.spark.last_day(d).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.date(2020, 1, 31) |
| """ |
| return Expr(_f.last_day(col.expr)) |
| |
| |
| def make_dt_interval( |
| days: Expr | int | None = None, |
| hours: Expr | int | None = None, |
| mins: Expr | int | None = None, |
| secs: Expr | float | None = None, |
| ) -> Expr: |
| """Spark ``make_dt_interval``: day-time interval from components. |
| |
| All parts are optional; omitted parts default to zero, matching pyspark. |
| Integer parts accept a native ``int`` and ``secs`` accepts a ``float``, |
| or any part may be an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.make_dt_interval().alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.timedelta(0) |
| |
| >>> r = df.select( |
| ... dfn.functions.spark.make_dt_interval( |
| ... days=1, hours=2, mins=3, secs=4.5 |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.timedelta(days=1, seconds=7384, microseconds=500000) |
| """ |
| return Expr( |
| _f.make_dt_interval( |
| (_coerce_i32(days) if days is not None else _ZERO_I32).expr, |
| (_coerce_i32(hours) if hours is not None else _ZERO_I32).expr, |
| (_coerce_i32(mins) if mins is not None else _ZERO_I32).expr, |
| (coerce_to_expr(secs) if secs is not None else Expr.literal(0.0)).expr, |
| ) |
| ) |
| |
| |
| def make_interval( |
| years: Expr | int | None = None, |
| months: Expr | int | None = None, |
| weeks: Expr | int | None = None, |
| days: Expr | int | None = None, |
| hours: Expr | int | None = None, |
| mins: Expr | int | None = None, |
| secs: Expr | float | None = None, |
| ) -> Expr: |
| """Spark ``make_interval``: interval from year/month/week/day/hour/min/sec parts. |
| |
| All parts are optional; omitted parts default to zero, matching pyspark. |
| Integer parts accept a native ``int`` and ``secs`` accepts a ``float``, |
| or any part may be an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.make_interval().alias("v")) |
| >>> r.collect_column("v")[0].as_py().months |
| 0 |
| |
| >>> r = df.select(dfn.functions.spark.make_interval(years=1).alias("v")) |
| >>> r.collect_column("v")[0].as_py().months |
| 12 |
| """ |
| return Expr( |
| _f.make_interval( |
| (_coerce_i32(years) if years is not None else _ZERO_I32).expr, |
| (_coerce_i32(months) if months is not None else _ZERO_I32).expr, |
| (_coerce_i32(weeks) if weeks is not None else _ZERO_I32).expr, |
| (_coerce_i32(days) if days is not None else _ZERO_I32).expr, |
| (_coerce_i32(hours) if hours is not None else _ZERO_I32).expr, |
| (_coerce_i32(mins) if mins is not None else _ZERO_I32).expr, |
| (coerce_to_expr(secs) if secs is not None else Expr.literal(0.0)).expr, |
| ) |
| ) |
| |
| |
| def next_day(date: Expr, dayOfWeek: Expr | str) -> Expr: # noqa: N803 |
| """Spark ``next_day``: first date after ``start_date`` named ``day_of_week``. |
| |
| ``dayOfWeek`` accepts a native ``str`` or an :class:`Expr`. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import date |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> d = dfn.lit(pa.scalar(date(2020, 1, 15), type=pa.date32())) |
| >>> r = df.select(dfn.functions.spark.next_day(d, "Mon").alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.date(2020, 1, 20) |
| """ |
| return Expr(_f.next_day(date.expr, coerce_to_expr(dayOfWeek).expr)) |
| |
| |
| def date_diff(end: Expr, start: Expr) -> Expr: |
| """Spark ``date_diff``: number of days from ``start_date`` to ``end_date``. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import date |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> d = dfn.lit(pa.scalar(date(2020, 1, 15), type=pa.date32())) |
| >>> end = dfn.lit(pa.scalar(date(2020, 1, 20), type=pa.date32())) |
| >>> r = df.select(dfn.functions.spark.date_diff(end, d).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 5 |
| """ |
| return Expr(_f.date_diff(end.expr, start.expr)) |
| |
| |
| def date_trunc(format: Expr | str, timestamp: Expr) -> Expr: |
| """Spark ``date_trunc``: truncate timestamp to unit ``fmt``. |
| |
| ``format`` accepts a native ``str`` or an :class:`Expr`. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import datetime |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> ts = dfn.lit( |
| ... pa.scalar(datetime(2020, 1, 15, 14, 30, 45), |
| ... type=pa.timestamp('us'))) |
| >>> r = df.select( |
| ... dfn.functions.spark.date_trunc("month", ts).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.datetime(2020, 1, 1, 0, 0) |
| """ |
| return Expr(_f.date_trunc(coerce_to_expr(format).expr, timestamp.expr)) |
| |
| |
| def time_trunc(unit: Expr | str, time: Expr) -> Expr: |
| """Spark ``time_trunc``: truncate time value to unit ``fmt``. |
| |
| A bare ``str`` ``unit`` is treated as a column name (matching pyspark), |
| not a literal; pass :func:`~datafusion.lit` for a literal unit. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import time |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> t = dfn.lit(pa.scalar(time(14, 30, 45), type=pa.time64('us'))) |
| >>> r = df.select( |
| ... dfn.functions.spark.time_trunc(dfn.lit("hour"), t).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.time(14, 0) |
| """ |
| return Expr(_f.time_trunc(_to_raw_expr(unit), time.expr)) |
| |
| |
| def trunc(date: Expr, format: Expr | str) -> Expr: |
| """Spark ``trunc``: truncate date to unit ``fmt``. |
| |
| ``format`` accepts a native ``str`` or an :class:`Expr`. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import date |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> d = dfn.lit(pa.scalar(date(2020, 1, 15), type=pa.date32())) |
| >>> r = df.select(dfn.functions.spark.trunc(d, "YEAR").alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.date(2020, 1, 1) |
| """ |
| return Expr(_f.trunc(date.expr, coerce_to_expr(format).expr)) |
| |
| |
| def date_part(field: Expr | str, source: Expr) -> Expr: |
| """Spark ``date_part``: extract ``field`` from a date/time/timestamp. |
| |
| ``field`` accepts a native ``str`` or an :class:`Expr`. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import date |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> d = dfn.lit(pa.scalar(date(2020, 1, 15), type=pa.date32())) |
| >>> r = df.select( |
| ... dfn.functions.spark.date_part("year", d).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 2020 |
| """ |
| return Expr(_f.date_part(coerce_to_expr(field).expr, source.expr)) |
| |
| |
| def from_utc_timestamp(timestamp: Expr, tz: Expr | str) -> Expr: |
| """Spark ``from_utc_timestamp``: interpret ``ts`` as UTC, convert to ``tz``. |
| |
| ``tz`` accepts a native ``str`` or an :class:`Expr`. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import datetime |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> ts = dfn.lit( |
| ... pa.scalar(datetime(2020, 1, 15, 14, 30, 45), |
| ... type=pa.timestamp('us'))) |
| >>> r = df.select( |
| ... dfn.functions.spark.from_utc_timestamp(ts, "UTC").alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.datetime(2020, 1, 15, 14, 30, 45) |
| """ |
| return Expr(_f.from_utc_timestamp(timestamp.expr, coerce_to_expr(tz).expr)) |
| |
| |
| def to_utc_timestamp(timestamp: Expr, tz: Expr | str) -> Expr: |
| """Spark ``to_utc_timestamp``: interpret ``ts`` as ``tz``, convert to UTC. |
| |
| ``tz`` accepts a native ``str`` or an :class:`Expr`. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import datetime |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> ts = dfn.lit( |
| ... pa.scalar(datetime(2020, 1, 15, 14, 30, 45), |
| ... type=pa.timestamp('us'))) |
| >>> r = df.select( |
| ... dfn.functions.spark.to_utc_timestamp(ts, "UTC").alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| datetime.datetime(2020, 1, 15, 14, 30, 45) |
| """ |
| return Expr(_f.to_utc_timestamp(timestamp.expr, coerce_to_expr(tz).expr)) |
| |
| |
| def unix_date(col: Expr) -> Expr: |
| """Spark ``unix_date``: days since 1970-01-01 for ``dt``. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import date |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> d = dfn.lit(pa.scalar(date(2020, 1, 15), type=pa.date32())) |
| >>> r = df.select(dfn.functions.spark.unix_date(d).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 18276 |
| """ |
| return Expr(_f.unix_date(col.expr)) |
| |
| |
| def unix_micros(col: Expr) -> Expr: |
| """Spark ``unix_micros``: microseconds since epoch for ``ts``. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import datetime |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> ts = dfn.lit( |
| ... pa.scalar(datetime(2020, 1, 15, 14, 30, 45), |
| ... type=pa.timestamp('us'))) |
| >>> r = df.select(dfn.functions.spark.unix_micros(ts).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 1579098645000000 |
| """ |
| return Expr(_f.unix_micros(col.expr)) |
| |
| |
| def unix_millis(col: Expr) -> Expr: |
| """Spark ``unix_millis``: milliseconds since epoch for ``ts``. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import datetime |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> ts = dfn.lit( |
| ... pa.scalar(datetime(2020, 1, 15, 14, 30, 45), |
| ... type=pa.timestamp('us'))) |
| >>> r = df.select(dfn.functions.spark.unix_millis(ts).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 1579098645000 |
| """ |
| return Expr(_f.unix_millis(col.expr)) |
| |
| |
| def unix_seconds(col: Expr) -> Expr: |
| """Spark ``unix_seconds``: seconds since epoch for ``ts``. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> from datetime import datetime |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> ts = dfn.lit( |
| ... pa.scalar(datetime(2020, 1, 15, 14, 30, 45), |
| ... type=pa.timestamp('us'))) |
| >>> r = df.select(dfn.functions.spark.unix_seconds(ts).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 1579098645 |
| """ |
| return Expr(_f.unix_seconds(col.expr)) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Hash functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def crc32(col: Expr) -> Expr: |
| """Spark ``crc32``: cyclic redundancy check value as a bigint. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"s": ["ABC"]}) |
| >>> r = df.select(dfn.functions.spark.crc32(dfn.col("s")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 2743272264 |
| """ |
| return Expr(_f.crc32(col.expr)) |
| |
| |
| def sha1(col: Expr) -> Expr: |
| """Spark ``sha1``: SHA-1 hash as a hex string. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"s": ["hello"]}) |
| >>> r = df.select(dfn.functions.spark.sha1(dfn.col("s")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'aaf4c61ddcc5e8a2dabede0f3b482cd9aea9434d' |
| """ |
| return Expr(_f.sha1(col.expr)) |
| |
| |
| def sha2(col: Expr, numBits: Expr | int) -> Expr: # noqa: N803 |
| """Spark ``sha2``: SHA-2 family hash (224, 256, 384, 512). Bit length 0 = 256. |
| |
| ``numBits`` accepts a native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"s": ["hello"]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.sha2(dfn.col("s"), 256).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| '2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824' |
| """ |
| return Expr(_f.sha2(col.expr, coerce_to_expr(numBits).expr)) |
| |
| |
| def xxhash64(*cols: Expr) -> Expr: |
| """Spark ``xxhash64``: 64-bit xxHash of the arguments. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.xxhash64(dfn.lit("hello")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| -4367754540140381902 |
| """ |
| return Expr(_f.xxhash64(*[c.expr for c in cols])) |
| |
| |
| # --------------------------------------------------------------------------- |
| # JSON functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def json_tuple(col: Expr, *fields: Expr | str) -> Expr: |
| """Spark ``json_tuple``: extract top-level fields from a JSON string. |
| |
| Each field name accepts a native ``str`` or an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.json_tuple( |
| ... dfn.lit('{"a":1,"b":"x"}'), "a", "b" |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| {'c0': '1', 'c1': 'x'} |
| """ |
| return Expr(_f.json_tuple(col.expr, *[coerce_to_expr(f).expr for f in fields])) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Map functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def map_from_arrays(col1: Expr, col2: Expr) -> Expr: |
| """Spark ``map_from_arrays``: build a map from parallel key/value arrays. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> keys = dfn.functions.spark.array(dfn.lit("a"), dfn.lit("b")) |
| >>> vals = dfn.functions.spark.array(dfn.lit(1), dfn.lit(2)) |
| >>> r = df.select( |
| ... dfn.functions.spark.map_from_arrays(keys, vals).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| [('a', 1), ('b', 2)] |
| """ |
| return Expr(_f.map_from_arrays(col1.expr, col2.expr)) |
| |
| |
| def map_from_entries(col: Expr) -> Expr: |
| """Spark ``map_from_entries``: build a map from an array of key/value structs. |
| |
| ``col`` must be an array whose elements are two-field structs; the first |
| field becomes the map key and the second the value. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> ctx = dfn.SessionContext() |
| >>> entry_type = pa.list_( |
| ... pa.struct([("key", pa.string()), ("value", pa.int64())])) |
| >>> entries = pa.array( |
| ... [[{"key": "a", "value": 1}, {"key": "b", "value": 2}]], |
| ... type=entry_type) |
| >>> df = ctx.from_arrow(pa.record_batch([entries], names=["e"])) |
| >>> r = df.select( |
| ... dfn.functions.spark.map_from_entries(dfn.col("e")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| [('a', 1), ('b', 2)] |
| """ |
| return Expr(_f.map_from_entries(col.expr)) |
| |
| |
| def str_to_map( |
| text: Expr, |
| pairDelim: Expr | str | None = None, # noqa: N803 |
| keyValueDelim: Expr | str | None = None, # noqa: N803 |
| ) -> Expr: |
| """Spark ``str_to_map``: split text into key/value pairs using delimiters. |
| |
| Delimiters default to ``","`` and ``":"`` when omitted, matching pyspark. |
| Parameter names match ``pyspark.sql.functions.str_to_map``; pyspark types |
| the delimiters as column-or-name, so a bare ``str`` is treated as a column |
| name. Pass :func:`~datafusion.lit` for a literal delimiter. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.str_to_map(dfn.lit("a:1,b:2")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| [('a', '1'), ('b', '2')] |
| |
| >>> r = df.select( |
| ... dfn.functions.spark.str_to_map( |
| ... dfn.lit("a=1;b=2"), |
| ... pairDelim=dfn.lit(";"), |
| ... keyValueDelim=dfn.lit("="), |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| [('a', '1'), ('b', '2')] |
| """ |
| pd = _to_raw_expr(pairDelim) if pairDelim is not None else Expr.literal(",").expr |
| kvd = ( |
| _to_raw_expr(keyValueDelim) |
| if keyValueDelim is not None |
| else Expr.literal(":").expr |
| ) |
| return Expr(_f.str_to_map(text.expr, pd, kvd)) |
| |
| |
| # --------------------------------------------------------------------------- |
| # Math functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def abs(col: Expr) -> Expr: |
| """Spark ``abs``: absolute value. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.abs(dfn.lit(-5)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 5 |
| """ |
| return Expr(_f.abs(col.expr)) |
| |
| |
| def ceil(col: Expr) -> Expr: |
| """Spark ``ceil``: smallest integer ≥ arg. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.ceil(dfn.lit(1.2)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 2 |
| """ |
| return Expr(_f.ceil(col.expr)) |
| |
| |
| def expm1(col: Expr) -> Expr: |
| """Spark ``expm1``: exp(arg) - 1. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.expm1(dfn.lit(0.0)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 0.0 |
| """ |
| return Expr(_f.expm1(col.expr)) |
| |
| |
| def factorial(col: Expr) -> Expr: |
| """Spark ``factorial``: n! for n in [0..20], else NULL. |
| |
| Examples: |
| >>> import pyarrow as pa |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.factorial( |
| ... dfn.lit(pa.scalar(5, type=pa.int32())) |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| 120 |
| """ |
| return Expr(_f.factorial(col.expr)) |
| |
| |
| def floor(col: Expr) -> Expr: |
| """Spark ``floor``: largest integer ≤ arg. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.floor(dfn.lit(1.8)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 1 |
| """ |
| return Expr(_f.floor(col.expr)) |
| |
| |
| def hex(col: Expr) -> Expr: |
| """Spark ``hex``: hexadecimal representation. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.hex(dfn.lit(255)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'FF' |
| """ |
| return Expr(_f.hex(col.expr)) |
| |
| |
| def modulus(dividend: Expr | float, divisor: Expr | float) -> Expr: |
| """Spark ``mod``: remainder of ``dividend / divisor`` (sign follows dividend). |
| |
| ``dividend`` and ``divisor`` accept native numbers or :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.modulus(dfn.lit(10), dfn.lit(3)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 1 |
| """ |
| return Expr(_f.modulus(coerce_to_expr(dividend).expr, coerce_to_expr(divisor).expr)) |
| |
| |
| def pmod(dividend: Expr | float, divisor: Expr | float) -> Expr: |
| """Spark ``pmod``: positive remainder of division. |
| |
| ``dividend`` and ``divisor`` accept native numbers or :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.pmod(dfn.lit(-1), dfn.lit(3)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 2 |
| """ |
| return Expr(_f.pmod(coerce_to_expr(dividend).expr, coerce_to_expr(divisor).expr)) |
| |
| |
| def rint(col: Expr) -> Expr: |
| """Spark ``rint``: round to nearest mathematical integer (as double). |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.rint(dfn.lit(2.5)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 2.0 |
| """ |
| return Expr(_f.rint(col.expr)) |
| |
| |
| def round(col: Expr, scale: Expr | int | None = None) -> Expr: |
| """Spark ``round``: round to ``scale`` decimal places, HALF_UP rounding. |
| |
| ``scale`` defaults to zero when omitted, matching pyspark, and accepts a |
| native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.round(dfn.lit(2.5)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 3.0 |
| |
| >>> r = df.select( |
| ... dfn.functions.spark.round(dfn.lit(2.345), scale=2).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 2.35 |
| """ |
| scale_expr = coerce_to_expr(scale) if scale is not None else _ZERO_I32 |
| return Expr(_f.round(col.expr, scale_expr.expr)) |
| |
| |
| def unhex(col: Expr) -> Expr: |
| r"""Spark ``unhex``: convert hexadecimal string to binary. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.unhex(dfn.lit("FF")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| b'\xff' |
| """ |
| return Expr(_f.unhex(col.expr)) |
| |
| |
| def width_bucket( |
| v: Expr, |
| min: Expr, |
| max: Expr, |
| numBucket: Expr | int, # noqa: N803 |
| ) -> Expr: |
| """Spark ``width_bucket``: bucket number for ``value`` in equi-width histogram. |
| |
| ``numBucket`` accepts a native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.width_bucket( |
| ... dfn.lit(5.0), dfn.lit(0.0), dfn.lit(10.0), 5 |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| 3 |
| """ |
| return Expr( |
| _f.width_bucket(v.expr, min.expr, max.expr, coerce_to_expr(numBucket).expr) |
| ) |
| |
| |
| def csc(col: Expr) -> Expr: |
| """Spark ``csc``: cosecant. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.csc(dfn.lit(1.5708)).alias("v")) |
| >>> f"{r.collect_column('v')[0].as_py():.4f}" |
| '1.0000' |
| """ |
| return Expr(_f.csc(col.expr)) |
| |
| |
| def sec(col: Expr) -> Expr: |
| """Spark ``sec``: secant. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.sec(dfn.lit(0.0)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 1.0 |
| """ |
| return Expr(_f.sec(col.expr)) |
| |
| |
| def negative(col: Expr) -> Expr: |
| """Spark ``negative``: unary minus. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.negative(dfn.lit(3)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| -3 |
| """ |
| return Expr(_f.negative(col.expr)) |
| |
| |
| def bin(col: Expr) -> Expr: |
| """Spark ``bin``: binary string representation of a long. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.bin(dfn.lit(7)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| '111' |
| """ |
| return Expr(_f.bin(col.expr)) |
| |
| |
| # --------------------------------------------------------------------------- |
| # String functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def ascii(col: Expr) -> Expr: |
| """Spark ``ascii``: code point of the first character. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.ascii(dfn.lit("A")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 65 |
| """ |
| return Expr(_f.ascii(col.expr)) |
| |
| |
| def base64(col: Expr) -> Expr: |
| """Spark ``base64``: encode binary as a base64 string. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.base64(dfn.lit(b"hi")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'aGk=' |
| """ |
| return Expr(_f.base64(col.expr)) |
| |
| |
| def char(col: Expr) -> Expr: |
| """Spark ``char``: ASCII character for a code point (mod 256). |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.char(dfn.lit(65)).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'A' |
| """ |
| return Expr(_f.char(col.expr)) |
| |
| |
| def concat(*cols: Expr) -> Expr: |
| """Spark ``concat``: concatenates strings; NULL if any input is NULL. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.concat(dfn.lit("a"), dfn.lit("b")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'ab' |
| """ |
| return Expr(_f.concat(*[c.expr for c in cols])) |
| |
| |
| def elt(*inputs: Expr) -> Expr: |
| """Spark ``elt``: returns the n-th input (1-indexed). |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.elt( |
| ... dfn.lit(2), dfn.lit("a"), dfn.lit("b") |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| 'b' |
| """ |
| return Expr(_f.elt(*[i.expr for i in inputs])) |
| |
| |
| def ilike( |
| str: Expr, |
| pattern: Expr | str, |
| escapeChar: str | None = None, # noqa: N803 |
| ) -> Expr: |
| """Spark ``ilike``: case-insensitive pattern match. |
| |
| A bare ``str`` ``pattern`` is treated as a column name (matching pyspark), |
| not a literal; pass :func:`~datafusion.lit` for a literal pattern. |
| ``escapeChar`` is accepted for pyspark parity but is not yet wired through |
| the Rust binding; passing a non-``None`` value raises ``NotImplementedError``. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.ilike(dfn.lit("HELLO"), dfn.lit("h%")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| True |
| """ |
| if escapeChar is not None: |
| msg = "ilike(escapeChar=...) is not yet supported by the Spark UDF binding" |
| raise NotImplementedError(msg) |
| return Expr(_f.ilike(str.expr, _to_raw_expr(pattern))) |
| |
| |
| def length(col: Expr) -> Expr: |
| """Spark ``length``: character length of a string, or byte length of binary. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.length(dfn.lit("hello")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 5 |
| """ |
| return Expr(_f.length(col.expr)) |
| |
| |
| def like( |
| str: Expr, |
| pattern: Expr | str, |
| escapeChar: str | None = None, # noqa: N803 |
| ) -> Expr: |
| """Spark ``like``: case-sensitive pattern match. |
| |
| A bare ``str`` ``pattern`` is treated as a column name (matching pyspark), |
| not a literal; pass :func:`~datafusion.lit` for a literal pattern. |
| ``escapeChar`` is accepted for pyspark parity but is not yet wired through |
| the Rust binding; passing a non-``None`` value raises ``NotImplementedError``. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.like(dfn.lit("hello"), dfn.lit("h%")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| True |
| """ |
| if escapeChar is not None: |
| msg = "like(escapeChar=...) is not yet supported by the Spark UDF binding" |
| raise NotImplementedError(msg) |
| return Expr(_f.like(str.expr, _to_raw_expr(pattern))) |
| |
| |
| def luhn_check(col: Expr) -> Expr: |
| """Spark ``luhn_check``: true if the digit string passes the Luhn check. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.luhn_check( |
| ... dfn.lit("4111111111111111") |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| True |
| """ |
| return Expr(_f.luhn_check(col.expr)) |
| |
| |
| def format_string(format: str | Expr, *cols: Expr) -> Expr: |
| """Spark ``format_string``: printf-style format string. |
| |
| ``format`` is the printf-style template (a plain ``str`` is auto-promoted |
| to a literal expression); remaining args are values to substitute. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.format_string( |
| ... "%d-%s", dfn.lit(42), dfn.lit("hi") |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| '42-hi' |
| """ |
| fmt_expr = format if isinstance(format, Expr) else Expr.literal(format) |
| return Expr(_f.format_string(fmt_expr.expr, *[c.expr for c in cols])) |
| |
| |
| def space(col: Expr | int) -> Expr: |
| """Spark ``space``: string of n spaces. |
| |
| ``col`` accepts a native ``int`` or an :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.space(3).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| ' ' |
| """ |
| return Expr(_f.space(_coerce_i32(col).expr)) |
| |
| |
| def substring(str: Expr, pos: Expr | int, len: Expr | int) -> Expr: |
| """Spark ``substring``: 1-indexed substring starting at ``pos`` of given ``length``. |
| |
| Negative ``pos`` counts from the end. ``pos`` and ``len`` accept native |
| ``int`` values or :class:`Expr`. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.substring(dfn.lit("hello"), 1, 3).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'hel' |
| """ |
| return Expr( |
| _f.substring(str.expr, coerce_to_expr(pos).expr, coerce_to_expr(len).expr) |
| ) |
| |
| |
| def unbase64(col: Expr) -> Expr: |
| """Spark ``unbase64``: decode a base64 string to binary. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.unbase64(dfn.lit("aGk=")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| b'hi' |
| """ |
| return Expr(_f.unbase64(col.expr)) |
| |
| |
| def soundex(col: Expr) -> Expr: |
| """Spark ``soundex``: Soundex phonetic code. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select(dfn.functions.spark.soundex(dfn.lit("Robert")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'R163' |
| """ |
| return Expr(_f.soundex(col.expr)) |
| |
| |
| def is_valid_utf8(str: Expr) -> Expr: |
| """Spark ``is_valid_utf8``: true if the string is valid UTF-8. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.is_valid_utf8(dfn.lit("hello")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| True |
| """ |
| return Expr(_f.is_valid_utf8(str.expr)) |
| |
| |
| def make_valid_utf8(str: Expr) -> Expr: |
| """Spark ``make_valid_utf8``: replace invalid UTF-8 bytes with U+FFFD. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.make_valid_utf8(dfn.lit("hello")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'hello' |
| """ |
| return Expr(_f.make_valid_utf8(str.expr)) |
| |
| |
| # --------------------------------------------------------------------------- |
| # URL functions |
| # --------------------------------------------------------------------------- |
| |
| |
| def parse_url( |
| url: Expr, |
| partToExtract: Expr | str, # noqa: N803 |
| key: Expr | str | None = None, |
| ) -> Expr: |
| """Spark ``parse_url``: extract a part from a URL; errors on invalid URLs. |
| |
| ``partToExtract`` is one of ``"HOST"``, ``"PATH"``, ``"QUERY"``, |
| ``"REF"``, ``"PROTOCOL"``, ``"FILE"``, ``"AUTHORITY"``, ``"USERINFO"``. |
| Pass ``key`` only with ``"QUERY"`` to extract a single parameter. Bare |
| ``str`` values for ``partToExtract``/``key`` are treated as column names |
| (matching pyspark); pass :func:`~datafusion.lit` for a literal. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.parse_url( |
| ... dfn.lit("http://example.com/path?q=1"), dfn.lit("HOST") |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| 'example.com' |
| |
| >>> r = df.select( |
| ... dfn.functions.spark.parse_url( |
| ... dfn.lit("http://example.com/path?q=1"), |
| ... dfn.lit("QUERY"), |
| ... key=dfn.lit("q"), |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| '1' |
| """ |
| if key is None: |
| return Expr(_f.parse_url(url.expr, _to_raw_expr(partToExtract))) |
| return Expr(_f.parse_url(url.expr, _to_raw_expr(partToExtract), _to_raw_expr(key))) |
| |
| |
| def try_parse_url( |
| url: Expr, |
| partToExtract: Expr | str, # noqa: N803 |
| key: Expr | str | None = None, |
| ) -> Expr: |
| """Spark ``try_parse_url``: like ``parse_url`` but returns NULL on invalid URLs. |
| |
| Bare ``str`` values for ``partToExtract``/``key`` are treated as column |
| names (matching pyspark); pass :func:`~datafusion.lit` for a literal. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.try_parse_url( |
| ... dfn.lit("http://example.com/"), dfn.lit("HOST") |
| ... ).alias("v") |
| ... ) |
| >>> r.collect_column("v")[0].as_py() |
| 'example.com' |
| """ |
| if key is None: |
| return Expr(_f.try_parse_url(url.expr, _to_raw_expr(partToExtract))) |
| return Expr( |
| _f.try_parse_url(url.expr, _to_raw_expr(partToExtract), _to_raw_expr(key)) |
| ) |
| |
| |
| def url_decode(str: Expr) -> Expr: |
| """Spark ``url_decode``: decode an application/x-www-form-urlencoded string. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.url_decode(dfn.lit("a%20b")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'a b' |
| """ |
| return Expr(_f.url_decode(str.expr)) |
| |
| |
| def try_url_decode(str: Expr) -> Expr: |
| """Spark ``try_url_decode``: like ``url_decode``; returns NULL on invalid input. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.try_url_decode(dfn.lit("a%20b")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'a b' |
| """ |
| return Expr(_f.try_url_decode(str.expr)) |
| |
| |
| def url_encode(str: Expr) -> Expr: |
| """Spark ``url_encode``: encode a string in application/x-www-form-urlencoded. |
| |
| Examples: |
| >>> ctx = dfn.SessionContext() |
| >>> df = ctx.from_pydict({"x": [1]}) |
| >>> r = df.select( |
| ... dfn.functions.spark.url_encode(dfn.lit("a b")).alias("v")) |
| >>> r.collect_column("v")[0].as_py() |
| 'a+b' |
| """ |
| return Expr(_f.url_encode(str.expr)) |
| |
| |
| __all__ = [ |
| # Math |
| "abs", |
| # Datetime |
| "add_months", |
| "array", |
| # Array |
| "array_contains", |
| "array_repeat", |
| # String |
| "ascii", |
| # Aggregate |
| "avg", |
| "base64", |
| "bin", |
| "bit_count", |
| # Bitwise |
| "bit_get", |
| "bitmap_bit_position", |
| "bitmap_bucket_number", |
| # Bitmap |
| "bitmap_count", |
| "bitwise_not", |
| "ceil", |
| "char", |
| "collect_list", |
| "collect_set", |
| "concat", |
| # Hash |
| "crc32", |
| "csc", |
| "date_add", |
| "date_diff", |
| "date_part", |
| "date_sub", |
| "date_trunc", |
| "elt", |
| "expm1", |
| "factorial", |
| "floor", |
| "format_string", |
| "from_utc_timestamp", |
| "hex", |
| "hour", |
| "if_", |
| "ilike", |
| "is_valid_utf8", |
| # JSON |
| "json_tuple", |
| "last_day", |
| "length", |
| "like", |
| "luhn_check", |
| "make_dt_interval", |
| "make_interval", |
| "make_valid_utf8", |
| # Map |
| "map_from_arrays", |
| "map_from_entries", |
| "minute", |
| "modulus", |
| "negative", |
| "next_day", |
| # URL |
| "parse_url", |
| "pmod", |
| "rint", |
| "round", |
| "sec", |
| "second", |
| "sha1", |
| "sha2", |
| "shiftleft", |
| "shiftright", |
| "shiftrightunsigned", |
| "shuffle", |
| # Collection / Conditional / Conversion |
| "size", |
| "slice", |
| "soundex", |
| "space", |
| "spark_cast", |
| "str_to_map", |
| "substring", |
| "time_trunc", |
| "to_utc_timestamp", |
| "trunc", |
| "try_parse_url", |
| "try_sum", |
| "try_url_decode", |
| "unbase64", |
| "unhex", |
| "unix_date", |
| "unix_micros", |
| "unix_millis", |
| "unix_seconds", |
| "url_decode", |
| "url_encode", |
| "width_bucket", |
| "xxhash64", |
| ] |
