Skip to content

Enabler

ensure_column(*columns)

Enables PySpark functions to accept either column names (as strings) or Column objects.

Parameters: columns (ColumnOrName): Column names (as strings) or Column objects to be converted.

Returns: tuple[Column]: A tuple of Column objects.

Examples: 

>>> ensure_column("col1", "col2", F.col("col3"))
(Column<b'col1'>, Column<b'col2'>, Column<b'col3'>)

Source code in pysparky/enabler.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
def ensure_column(*columns: ColumnOrName) -> tuple[Column, ...]:
    """
    Enables PySpark functions to accept either column names (as strings) or Column objects.

    Parameters:
    columns (ColumnOrName): Column names (as strings) or Column objects to be converted.

    Returns:
    tuple[Column]: A tuple of Column objects.

    Examples:
    ``` py
    >>> ensure_column("col1", "col2", F.col("col3"))
    (Column<b'col1'>, Column<b'col2'>, Column<b'col3'>)
    ```

    """
    return tuple(
        map(
            lambda column: F.col(column) if isinstance(column, str) else column, columns
        )
    )

ensure_list(single_or_list)

Ensures the input is returned as a list.

If the input is not already a list, it wraps the input in a list. If the input is already a list, it returns the input unchanged.

Parameters:

Name Type Description Default
single_or_list Union[Any, List[Any]]

The input which can be a single item or a list of items.

 required

Returns:

Type Description
list[Any]

List[Any]: A list containing the input item(s).

Examples:

>>> ensure_list(5)

>>> ensure_list([1, 2, 3])
[1, 2, 3]
>>> ensure_list("hello")
["hello"]
Source code in pysparky/enabler.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def ensure_list(single_or_list: Any | list[Any]) -> list[Any]:
    """
    Ensures the input is returned as a list.

    If the input is not already a list, it wraps the input in a list.
    If the input is already a list, it returns the input unchanged.

    Args:
        single_or_list (Union[Any, List[Any]]): The input which can be a single item or a list of items.

    Returns:
        List[Any]: A list containing the input item(s).

    Examples:
        ``` py
        >>> ensure_list(5)

        >>> ensure_list([1, 2, 3])
        [1, 2, 3]
        >>> ensure_list("hello")
        ["hello"]
        ```

    """
    return single_or_list if isinstance(single_or_list, list) else [single_or_list]