This is an automated email from the ASF dual-hosted git repository.
jorisvandenbossche pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/main by this push:
new 6521489333 GH-36045:[Python] Improve usability of pc.map_lookup /
MapLookupOptions (#36387)
6521489333 is described below
commit 65214893331fe6c16203337e719f75f751e5d8fc
Author: Junming Chen <[email protected]>
AuthorDate: Wed Jul 5 21:21:45 2023 +0800
GH-36045:[Python] Improve usability of pc.map_lookup / MapLookupOptions
(#36387)
### Rationale for this change
This PR is for #36045, which aims to improve usability of pc.map_lookup /
MapLookupOptions
### What changes are included in this PR?
For `query_key` which is not a subclass of `pyarrow.lib.Scalar`, we convert
it to `scalar`
### Are these changes tested?
Yes, add one test in tests/test_compute.py::test_map_lookup
* Closes: #36045
Lead-authored-by: Junming Chen <[email protected]>
Co-authored-by: Joris Van den Bossche <[email protected]>
Signed-off-by: Joris Van den Bossche <[email protected]>
---
python/pyarrow/_compute.pyx | 5 ++++-
python/pyarrow/tests/test_compute.py | 1 +
2 files changed, 5 insertions(+), 1 deletion(-)
diff --git a/python/pyarrow/_compute.pyx b/python/pyarrow/_compute.pyx
index 7e32ba48fb..ab63a7a19f 100644
--- a/python/pyarrow/_compute.pyx
+++ b/python/pyarrow/_compute.pyx
@@ -1565,7 +1565,7 @@ class MapLookupOptions(_MapLookupOptions):
Parameters
----------
- query_key : Scalar
+ query_key : Scalar or Object can be converted to Scalar
The key to search for.
occurrence : str
The occurrence(s) to return from the Map
@@ -1573,6 +1573,9 @@ class MapLookupOptions(_MapLookupOptions):
"""
def __init__(self, query_key, occurrence):
+ if not isinstance(query_key, lib.Scalar):
+ query_key = lib.scalar(query_key)
+
self._set_options(query_key, occurrence)
diff --git a/python/pyarrow/tests/test_compute.py
b/python/pyarrow/tests/test_compute.py
index d9209ada24..865fecc7b2 100644
--- a/python/pyarrow/tests/test_compute.py
+++ b/python/pyarrow/tests/test_compute.py
@@ -3093,6 +3093,7 @@ def test_map_lookup():
result_all = pa.array([[1], None, None, [5, 7], None],
type=pa.list_(pa.int32()))
+ assert pc.map_lookup(arr, 'one', 'first') == result_first
assert pc.map_lookup(arr, pa.scalar(
'one', type=pa.utf8()), 'first') == result_first
assert pc.map_lookup(arr, pa.scalar(
