GH-36045:[Python] Improve usability of pc.map_lookup / MapLookupOptio…

…ns (#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]>
apache · Jul 5, 2023 · 6521489 · 6521489
1 parent b679958
commit 6521489
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/python/pyarrow/_compute.pyx b/python/pyarrow/_compute.pyx
@@ -1565,14 +1565,17 @@ 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
         Accepted values are "first", "last", or "all".
     """
 
     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
@@ -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(