| # |
| # 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. |
| # |
| from typing import Optional |
| |
| from pyspark.errors import PySparkValueError |
| from pyspark.sql.column import Column |
| from pyspark.sql.dataframe import DataFrame |
| from pyspark.sql.session import SparkSession |
| |
| __all__ = ["TableValuedFunction"] |
| |
| |
| class TableValuedFunction: |
| """ |
| Interface for invoking table-valued functions in Spark SQL. |
| """ |
| |
| def __init__(self, sparkSession: SparkSession): |
| self._sparkSession = sparkSession |
| |
| def range( |
| self, |
| start: int, |
| end: Optional[int] = None, |
| step: int = 1, |
| numPartitions: Optional[int] = None, |
| ) -> DataFrame: |
| """ |
| Create a :class:`DataFrame` with single :class:`pyspark.sql.types.LongType` column named |
| ``id``, containing elements in a range from ``start`` to ``end`` (exclusive) with |
| step value ``step``. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| start : int |
| the start value |
| end : int, optional |
| the end value (exclusive) |
| step : int, optional |
| the incremental step (default: 1) |
| numPartitions : int, optional |
| the number of partitions of the DataFrame |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| Examples |
| -------- |
| >>> spark.tvf.range(1, 7, 2).show() |
| +---+ |
| | id| |
| +---+ |
| | 1| |
| | 3| |
| | 5| |
| +---+ |
| |
| If only one argument is specified, it will be used as the end value. |
| |
| >>> spark.tvf.range(3).show() |
| +---+ |
| | id| |
| +---+ |
| | 0| |
| | 1| |
| | 2| |
| +---+ |
| """ |
| return self._sparkSession.range(start, end, step, numPartitions) |
| |
| def explode(self, collection: Column) -> DataFrame: |
| """ |
| Returns a :class:`DataFrame` containing a new row for each element |
| in the given array or map. |
| Uses the default column name `col` for elements in the array and |
| `key` and `value` for elements in the map unless specified otherwise. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| collection : :class:`~pyspark.sql.Column` |
| Target column to work on. |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| See Also |
| -------- |
| :meth:`pyspark.sql.functions.explode` |
| |
| Examples |
| -------- |
| Example 1: Exploding an array column |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.explode(sf.array(sf.lit(1), sf.lit(2), sf.lit(3))).show() |
| +---+ |
| |col| |
| +---+ |
| | 1| |
| | 2| |
| | 3| |
| +---+ |
| |
| Example 2: Exploding a map column |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.explode( |
| ... sf.create_map(sf.lit("a"), sf.lit("b"), sf.lit("c"), sf.lit("d")) |
| ... ).show() |
| +---+-----+ |
| |key|value| |
| +---+-----+ |
| | a| b| |
| | c| d| |
| +---+-----+ |
| |
| Example 3: Exploding an array of struct column |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.explode(sf.array( |
| ... sf.named_struct(sf.lit("a"), sf.lit(1), sf.lit("b"), sf.lit(2)), |
| ... sf.named_struct(sf.lit("a"), sf.lit(3), sf.lit("b"), sf.lit(4)) |
| ... )).select("col.*").show() |
| +---+---+ |
| | a| b| |
| +---+---+ |
| | 1| 2| |
| | 3| 4| |
| +---+---+ |
| |
| Example 4: Exploding an empty array column |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.explode(sf.array()).show() |
| +---+ |
| |col| |
| +---+ |
| +---+ |
| |
| Example 5: Exploding an empty map column |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.explode(sf.create_map()).show() |
| +---+-----+ |
| |key|value| |
| +---+-----+ |
| +---+-----+ |
| """ |
| return self._fn("explode", collection) |
| |
| def explode_outer(self, collection: Column) -> DataFrame: |
| """ |
| Returns a :class:`DataFrame` containing a new row for each element with position |
| in the given array or map. |
| Unlike explode, if the array/map is null or empty then null is produced. |
| Uses the default column name `col` for elements in the array and |
| `key` and `value` for elements in the map unless specified otherwise. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| collection : :class:`~pyspark.sql.Column` |
| target column to work on. |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| See Also |
| -------- |
| :meth:`pyspark.sql.functions.explode_outer` |
| |
| Examples |
| -------- |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.explode_outer(sf.array(sf.lit("foo"), sf.lit("bar"))).show() |
| +---+ |
| |col| |
| +---+ |
| |foo| |
| |bar| |
| +---+ |
| >>> spark.tvf.explode_outer(sf.array()).show() |
| +----+ |
| | col| |
| +----+ |
| |NULL| |
| +----+ |
| >>> spark.tvf.explode_outer(sf.create_map(sf.lit("x"), sf.lit(1.0))).show() |
| +---+-----+ |
| |key|value| |
| +---+-----+ |
| | x| 1.0| |
| +---+-----+ |
| >>> spark.tvf.explode_outer(sf.create_map()).show() |
| +----+-----+ |
| | key|value| |
| +----+-----+ |
| |NULL| NULL| |
| +----+-----+ |
| """ |
| return self._fn("explode_outer", collection) |
| |
| def inline(self, input: Column) -> DataFrame: |
| """ |
| Explodes an array of structs into a table. |
| |
| This function takes an input column containing an array of structs and returns a |
| new column where each struct in the array is exploded into a separate row. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| input : :class:`~pyspark.sql.Column` |
| Input column of values to explode. |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| See Also |
| -------- |
| :meth:`pyspark.sql.functions.inline` |
| |
| Examples |
| -------- |
| Example 1: Using inline with a single struct array |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.inline(sf.array( |
| ... sf.named_struct(sf.lit("a"), sf.lit(1), sf.lit("b"), sf.lit(2)), |
| ... sf.named_struct(sf.lit("a"), sf.lit(3), sf.lit("b"), sf.lit(4)) |
| ... )).show() |
| +---+---+ |
| | a| b| |
| +---+---+ |
| | 1| 2| |
| | 3| 4| |
| +---+---+ |
| |
| Example 2: Using inline with an empty struct array column |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.inline(sf.array().astype("array<struct<a:int,b:int>>")).show() |
| +---+---+ |
| | a| b| |
| +---+---+ |
| +---+---+ |
| |
| Example 3: Using inline with a struct array column containing null values |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.inline(sf.array( |
| ... sf.named_struct(sf.lit("a"), sf.lit(1), sf.lit("b"), sf.lit(2)), |
| ... sf.lit(None), |
| ... sf.named_struct(sf.lit("a"), sf.lit(3), sf.lit("b"), sf.lit(4)) |
| ... )).show() |
| +----+----+ |
| | a| b| |
| +----+----+ |
| | 1| 2| |
| |NULL|NULL| |
| | 3| 4| |
| +----+----+ |
| """ |
| return self._fn("inline", input) |
| |
| def inline_outer(self, input: Column) -> DataFrame: |
| """ |
| Explodes an array of structs into a table. |
| Unlike inline, if the array is null or empty then null is produced for each nested column. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| input : :class:`~pyspark.sql.Column` |
| input column of values to explode. |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| See Also |
| -------- |
| :meth:`pyspark.sql.functions.inline_outer` |
| |
| Examples |
| -------- |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.inline_outer(sf.array( |
| ... sf.named_struct(sf.lit("a"), sf.lit(1), sf.lit("b"), sf.lit(2)), |
| ... sf.named_struct(sf.lit("a"), sf.lit(3), sf.lit("b"), sf.lit(4)) |
| ... )).show() |
| +---+---+ |
| | a| b| |
| +---+---+ |
| | 1| 2| |
| | 3| 4| |
| +---+---+ |
| >>> spark.tvf.inline_outer(sf.array().astype("array<struct<a:int,b:int>>")).show() |
| +----+----+ |
| | a| b| |
| +----+----+ |
| |NULL|NULL| |
| +----+----+ |
| """ |
| return self._fn("inline_outer", input) |
| |
| def json_tuple(self, input: Column, *fields: Column) -> DataFrame: |
| """ |
| Creates a new row for a json column according to the given field names. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| input : :class:`~pyspark.sql.Column` |
| string column in json format |
| fields : :class:`~pyspark.sql.Column` |
| a field or fields to extract |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| See Also |
| -------- |
| :meth:`pyspark.sql.functions.json_tuple` |
| |
| Examples |
| -------- |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.json_tuple( |
| ... sf.lit('{"f1": "value1", "f2": "value2"}'), sf.lit("f1"), sf.lit("f2") |
| ... ).show() |
| +------+------+ |
| | c0| c1| |
| +------+------+ |
| |value1|value2| |
| +------+------+ |
| """ |
| from pyspark.sql.classic.column import _to_seq, _to_java_column |
| |
| if len(fields) == 0: |
| raise PySparkValueError( |
| errorClass="CANNOT_BE_EMPTY", |
| messageParameters={"item": "field"}, |
| ) |
| |
| sc = self._sparkSession.sparkContext |
| return DataFrame( |
| self._sparkSession._jsparkSession.tvf().json_tuple( |
| _to_java_column(input), _to_seq(sc, fields, _to_java_column) |
| ), |
| self._sparkSession, |
| ) |
| |
| def posexplode(self, collection: Column) -> DataFrame: |
| """ |
| Returns a :class:`DataFrame` containing a new row for each element with position |
| in the given array or map. |
| Uses the default column name `pos` for position, and `col` for elements in the |
| array and `key` and `value` for elements in the map unless specified otherwise. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| collection : :class:`~pyspark.sql.Column` |
| target column to work on. |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| See Also |
| -------- |
| :meth:`pyspark.sql.functions.posexplode` |
| |
| Examples |
| -------- |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.posexplode(sf.array(sf.lit(1), sf.lit(2), sf.lit(3))).show() |
| +---+---+ |
| |pos|col| |
| +---+---+ |
| | 0| 1| |
| | 1| 2| |
| | 2| 3| |
| +---+---+ |
| >>> spark.tvf.posexplode(sf.create_map(sf.lit("a"), sf.lit("b"))).show() |
| +---+---+-----+ |
| |pos|key|value| |
| +---+---+-----+ |
| | 0| a| b| |
| +---+---+-----+ |
| """ |
| return self._fn("posexplode", collection) |
| |
| def posexplode_outer(self, collection: Column) -> DataFrame: |
| """ |
| Returns a :class:`DataFrame` containing a new row for each element with position |
| in the given array or map. |
| Unlike posexplode, if the array/map is null or empty then the row (null, null) is produced. |
| Uses the default column name `pos` for position, and `col` for elements in the |
| array and `key` and `value` for elements in the map unless specified otherwise. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| collection : :class:`~pyspark.sql.Column` |
| target column to work on. |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| See Also |
| -------- |
| :meth:`pyspark.sql.functions.posexplode_outer` |
| |
| Examples |
| -------- |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.posexplode_outer(sf.array(sf.lit("foo"), sf.lit("bar"))).show() |
| +---+---+ |
| |pos|col| |
| +---+---+ |
| | 0|foo| |
| | 1|bar| |
| +---+---+ |
| >>> spark.tvf.posexplode_outer(sf.array()).show() |
| +----+----+ |
| | pos| col| |
| +----+----+ |
| |NULL|NULL| |
| +----+----+ |
| >>> spark.tvf.posexplode_outer(sf.create_map(sf.lit("x"), sf.lit(1.0))).show() |
| +---+---+-----+ |
| |pos|key|value| |
| +---+---+-----+ |
| | 0| x| 1.0| |
| +---+---+-----+ |
| >>> spark.tvf.posexplode_outer(sf.create_map()).show() |
| +----+----+-----+ |
| | pos| key|value| |
| +----+----+-----+ |
| |NULL|NULL| NULL| |
| +----+----+-----+ |
| """ |
| return self._fn("posexplode_outer", collection) |
| |
| def stack(self, n: Column, *fields: Column) -> DataFrame: |
| """ |
| Separates `col1`, ..., `colk` into `n` rows. Uses column names col0, col1, etc. by default |
| unless specified otherwise. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| n : :class:`~pyspark.sql.Column` |
| the number of rows to separate |
| fields : :class:`~pyspark.sql.Column` |
| input elements to be separated |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| See Also |
| -------- |
| :meth:`pyspark.sql.functions.stack` |
| |
| Examples |
| -------- |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.stack(sf.lit(2), sf.lit(1), sf.lit(2), sf.lit(3)).show() |
| +----+----+ |
| |col0|col1| |
| +----+----+ |
| | 1| 2| |
| | 3|NULL| |
| +----+----+ |
| """ |
| from pyspark.sql.classic.column import _to_seq, _to_java_column |
| |
| sc = self._sparkSession.sparkContext |
| return DataFrame( |
| self._sparkSession._jsparkSession.tvf().stack( |
| _to_java_column(n), _to_seq(sc, fields, _to_java_column) |
| ), |
| self._sparkSession, |
| ) |
| |
| def collations(self) -> DataFrame: |
| """ |
| Get all of the Spark SQL string collations. |
| |
| .. versionadded:: 4.0.0 |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| Examples |
| -------- |
| >>> spark.tvf.collations().show() |
| +-------+-------+-------------------+... |
| |CATALOG| SCHEMA| NAME|... |
| +-------+-------+-------------------+... |
| ... |
| +-------+-------+-------------------+... |
| """ |
| return self._fn("collations") |
| |
| def sql_keywords(self) -> DataFrame: |
| """ |
| Get Spark SQL keywords. |
| |
| .. versionadded:: 4.0.0 |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| Examples |
| -------- |
| >>> spark.tvf.sql_keywords().show() |
| +-------------+--------+ |
| | keyword|reserved| |
| +-------------+--------+ |
| ... |
| +-------------+--------+... |
| """ |
| return self._fn("sql_keywords") |
| |
| def variant_explode(self, input: Column) -> DataFrame: |
| """ |
| Separates a variant object/array into multiple rows containing its fields/elements. |
| |
| Its result schema is `struct<pos int, key string, value variant>`. `pos` is the position of |
| the field/element in its parent object/array, and `value` is the field/element value. |
| `key` is the field name when exploding a variant object, or is NULL when exploding a variant |
| array. It ignores any input that is not a variant array/object, including SQL NULL, variant |
| null, and any other variant values. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| input : :class:`~pyspark.sql.Column` |
| input column of values to explode. |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| Examples |
| -------- |
| Example 1: Using variant_explode with a variant array |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.variant_explode(sf.parse_json(sf.lit('["hello", "world"]'))).show() |
| +---+----+-------+ |
| |pos| key| value| |
| +---+----+-------+ |
| | 0|NULL|"hello"| |
| | 1|NULL|"world"| |
| +---+----+-------+ |
| |
| Example 2: Using variant_explode with a variant object |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.variant_explode(sf.parse_json(sf.lit('{"a": true, "b": 3.14}'))).show() |
| +---+---+-----+ |
| |pos|key|value| |
| +---+---+-----+ |
| | 0| a| true| |
| | 1| b| 3.14| |
| +---+---+-----+ |
| |
| Example 3: Using variant_explode with an empty variant array |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.variant_explode(sf.parse_json(sf.lit('[]'))).show() |
| +---+---+-----+ |
| |pos|key|value| |
| +---+---+-----+ |
| +---+---+-----+ |
| |
| Example 4: Using variant_explode with an empty variant object |
| |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.variant_explode(sf.parse_json(sf.lit('{}'))).show() |
| +---+---+-----+ |
| |pos|key|value| |
| +---+---+-----+ |
| +---+---+-----+ |
| """ |
| return self._fn("variant_explode", input) |
| |
| def variant_explode_outer(self, input: Column) -> DataFrame: |
| """ |
| Separates a variant object/array into multiple rows containing its fields/elements. |
| |
| Its result schema is `struct<pos int, key string, value variant>`. `pos` is the position of |
| the field/element in its parent object/array, and `value` is the field/element value. |
| `key` is the field name when exploding a variant object, or is NULL when exploding a variant |
| array. Unlike variant_explode, if the given variant is not a variant array/object, including |
| SQL NULL, variant null, and any other variant values, then NULL is produced. |
| |
| .. versionadded:: 4.0.0 |
| |
| Parameters |
| ---------- |
| input : :class:`~pyspark.sql.Column` |
| input column of values to explode. |
| |
| Returns |
| ------- |
| :class:`DataFrame` |
| |
| Examples |
| -------- |
| >>> import pyspark.sql.functions as sf |
| >>> spark.tvf.variant_explode_outer(sf.parse_json(sf.lit('["hello", "world"]'))).show() |
| +---+----+-------+ |
| |pos| key| value| |
| +---+----+-------+ |
| | 0|NULL|"hello"| |
| | 1|NULL|"world"| |
| +---+----+-------+ |
| >>> spark.tvf.variant_explode_outer(sf.parse_json(sf.lit('[]'))).show() |
| +----+----+-----+ |
| | pos| key|value| |
| +----+----+-----+ |
| |NULL|NULL| NULL| |
| +----+----+-----+ |
| >>> spark.tvf.variant_explode_outer(sf.parse_json(sf.lit('{"a": true, "b": 3.14}'))).show() |
| +---+---+-----+ |
| |pos|key|value| |
| +---+---+-----+ |
| | 0| a| true| |
| | 1| b| 3.14| |
| +---+---+-----+ |
| >>> spark.tvf.variant_explode_outer(sf.parse_json(sf.lit('{}'))).show() |
| +----+----+-----+ |
| | pos| key|value| |
| +----+----+-----+ |
| |NULL|NULL| NULL| |
| +----+----+-----+ |
| """ |
| return self._fn("variant_explode_outer", input) |
| |
| def _fn(self, functionName: str, *args: Column) -> DataFrame: |
| from pyspark.sql.classic.column import _to_java_column |
| |
| return DataFrame( |
| getattr(self._sparkSession._jsparkSession.tvf(), functionName)( |
| *(_to_java_column(arg) for arg in args) |
| ), |
| self._sparkSession, |
| ) |
| |
| |
| def _test() -> None: |
| import os |
| import doctest |
| import sys |
| import pyspark.sql.tvf |
| |
| os.chdir(os.environ["SPARK_HOME"]) |
| |
| globs = pyspark.sql.tvf.__dict__.copy() |
| globs["spark"] = SparkSession.builder.master("local[4]").appName("sql.tvf tests").getOrCreate() |
| (failure_count, test_count) = doctest.testmod( |
| pyspark.sql.tvf, |
| globs=globs, |
| optionflags=doctest.ELLIPSIS | doctest.NORMALIZE_WHITESPACE, |
| ) |
| globs["spark"].stop() |
| if failure_count: |
| sys.exit(-1) |
| |
| |
| if __name__ == "__main__": |
| _test() |
