Skip to content

Commit

Permalink
SNOW-1865595: Add currently supported structured type functions (#2790)
Browse files Browse the repository at this point in the history
  • Loading branch information
sfc-gh-jrose authored Dec 20, 2024
1 parent 637905b commit 69c41a2
Show file tree
Hide file tree
Showing 3 changed files with 124 additions and 0 deletions.
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,11 @@
#### New Features

- Added support for the following functions in `functions.py`
- `array_reverse`
- `divnull`
- `map_cat`
- `map_contains_key`
- `map_keys`
- `nullifzero`
- `snowflake_cortex_sentiment`
- Added `Catalog` class to manage snowflake objects. It can be accessed via `Session.catalog`.
Expand Down
9 changes: 9 additions & 0 deletions docs/source/snowpark/functions.rst
Original file line number Diff line number Diff line change
Expand Up @@ -35,21 +35,26 @@ Functions
array_construct_compact
array_contains
array_distinct
array_except
array_flatten
array_generate_range
array_insert
array_intersection
array_join
array_max
array_min
array_position
array_prepend
array_remove
array_reverse
array_size
array_slice
array_sort
array_to_string
array_union
array_unique_agg
arrays_overlap
arrays_zip
as_array
as_binary
as_char
Expand Down Expand Up @@ -205,6 +210,10 @@ Functions
lpad
ltrim
make_interval
map_cat
map_concat
map_contains_key
map_keys
max
md5
mean
Expand Down
111 changes: 111 additions & 0 deletions src/snowflake/snowpark/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -4785,6 +4785,32 @@ def array_flatten(array: ColumnOrName, _emit_ast: bool = True) -> Column:
return builtin("array_flatten", _emit_ast=_emit_ast)(array)


@publicapi
def array_reverse(col: ColumnOrName, _emit_ast: bool = True) -> Column:
"""Returns an array with the elements of the input array in reverse order.
Args:
col: The source array.
Example::
>>> df = session.sql("select [1, 2, 3, 4] :: ARRAY(INT) as A")
>>> df.select(array_reverse("A")).show()
--------------------------
|"ARRAY_REVERSE(""A"")" |
--------------------------
|[ |
| 4, |
| 3, |
| 2, |
| 1 |
|] |
--------------------------
<BLANKLINE>
"""
array = _to_col_if_str(col, "array_reverse")
return builtin("array_reverse", _emit_ast=_emit_ast)(array)


@publicapi
def array_sort(
array: ColumnOrName,
Expand Down Expand Up @@ -6933,6 +6959,88 @@ def array_unique_agg(col: ColumnOrName, _emit_ast: bool = True) -> Column:
return _call_function("array_unique_agg", True, c, _emit_ast=_emit_ast)


@publicapi
def map_cat(col1: ColumnOrName, col2: ColumnOrName, _emit_ast: bool = True):
"""Returns the concatenatation of two MAPs.
Args:
col1: The source map
col2: The map to be appended to col1
Example::
>>> df = session.sql("select {'k1': 'v1'} :: MAP(STRING,STRING) as A, {'k2': 'v2'} :: MAP(STRING,STRING) as B")
>>> df.select(map_cat("A", "B")).show()
---------------------------
|"MAP_CAT(""A"", ""B"")" |
---------------------------
|{ |
| "k1": "v1", |
| "k2": "v2" |
|} |
---------------------------
<BLANKLINE>
"""
m1 = _to_col_if_str(col1, "map_cat")
m2 = _to_col_if_str(col2, "map_cat")
return builtin("map_cat", _emit_ast=_emit_ast)(m1, m2)


@publicapi
def map_contains_key(value: ColumnOrLiteral, col: ColumnOrName, _emit_ast: bool = True):
"""Determines whether the specified MAP contains the specified key.
Args:
value: The key to find.
col: The map to be searched.
Example 1::
>>> df = session.sql("select {'k1': 'v1'} :: MAP(STRING,STRING) as M, 'k1' as V")
>>> df.select(map_contains_key(col("V"), "M")).show()
------------------------------------
|"MAP_CONTAINS_KEY(""V"", ""M"")" |
------------------------------------
|True |
------------------------------------
<BLANKLINE>
Example 2::
>>> df = session.sql("select {'k1': 'v1'} :: MAP(STRING,STRING) as M")
>>> df.select(map_contains_key("k1", "M")).show()
-----------------------------------
|"MAP_CONTAINS_KEY('K1', ""M"")" |
-----------------------------------
|True |
-----------------------------------
<BLANKLINE>
"""
m = _to_col_if_str(col, "map_contains")
return builtin("map_contains_key", _emit_ast=_emit_ast)(value, m)


@publicapi
def map_keys(col: ColumnOrName, _emit_ast: bool = True):
"""Returns the keys in a MAP.
Args:
col: The input map.
Example 1::
>>> df = session.sql("select {'k1': 'v1', 'k2': 'v2'} :: MAP(STRING,STRING) as M")
>>> df.select(map_keys("M")).show()
---------------------
|"MAP_KEYS(""M"")" |
---------------------
|[ |
| "k1", |
| "k2" |
|] |
---------------------
<BLANKLINE>
"""
m = _to_col_if_str(col, "map_keys")
return builtin("map_keys", _emit_ast=_emit_ast)(m)


@publicapi
def size(col: ColumnOrName, _emit_ast: bool = True) -> Column:
"""Returns the size of the input ARRAY, OBJECT or MAP. Returns NULL if the
Expand Down Expand Up @@ -10142,6 +10250,9 @@ def sproc(
sort_array = array_sort
map_from_arrays = arrays_to_object
signum = sign
array_join = array_to_string
array_union = array_cat
map_concat = map_cat


@publicapi
Expand Down

0 comments on commit 69c41a2

Please sign in to comment.