Repository: ibis-project/ibis

Issue:
bug: can't use functools.partial with Array.map()
### What happened?

small example:

```python
import ibis
from functools import partial

a = ibis.array([1,2,3])

def add_n(x, n):
    return x + n

add_2 = partial(add_n, n=2)
a.map(add_2)
```

causes `IbisTypeError: Callable has mandatory keyword-only arguments which cannot be specified`

What I'm actually trying to do is have a general function, but then specialize it for different values and have a different column for each one. eg:

```python
t = ibis.examples.penguins.fetch()
t2 = t.group_by("species").agg(
    years=_.year.collect()
)
m = {f"year_plus_{i}": _.years.map(partial(add_n, n=i)) for i in range(3)}
t2.mutate(**m)
```

I tried to do this:
```python
t = ibis.examples.penguins.fetch()
t2 = t.group_by("species").agg(
    years=_.year.collect()
)
m = {}
for i in range(3):
    def add_i(x):
        return x + i
    m[f"year_plus_{i}"] = _.years.map(add_i)
t2.mutate(**m)
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┓
┃ species   ┃ years                  ┃ year_plus_0            ┃ year_plus_1            ┃ year_plus_2            ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━┩
│ string    │ array<int64>           │ array<int64>           │ array<int64>           │ array<int64>           │
├───────────┼────────────────────────┼────────────────────────┼────────────────────────┼────────────────────────┤
│ Adelie    │ [2007, 2007, ... +150] │ [2009, 2009, ... +150] │ [2009, 2009, ... +150] │ [2009, 2009, ... +150] │
│ Gentoo    │ [2007, 2007, ... +122] │ [2009, 2009, ... +122] │ [2009, 2009, ... +122] │ [2009, 2009, ... +122] │
│ Chinstrap │ [2007, 2007, ... +66]  │ [2009, 2009, ... +66]  │ [2009, 2009, ... +66]  │ [2009, 2009, ... +66]  │
└───────────┴────────────────────────┴────────────────────────┴────────────────────────┴────────────────────────┘
```
but notice that the `i` is only using the last seen value of 2, which is to be expected with the semantics of how python doesn't bind that value at definition time. I can get it to work with a class:

```python
t = ibis.examples.penguins.fetch()
t2 = t.group_by("species").agg(
    years=_.year.collect()
)

class AddN:
    def __init__(self, n):
        self.n = n

    def __call__(self, x):
        return x + self.n

m = {}
for i in range(3):
    m[f"year_plus_{i}"] = _.years.map(AddN(i))
t2.mutate(**m)
```

but functools.partial is the canonical way to do this and I think it should be supported. IDK if there are other pathways that we should consider here.

### What version of ibis are you using?

6.1.0

### What backend(s) are you using, if any?

NA

### Relevant log output

_No response_

### Code of Conduct

- [X] I agree to follow this project's Code of Conduct


Discussion:
cc @kszucs, this is an interesting challenge: modeling an N argument function such that the type system accepts it as a single argument function if N - 1 arguments are already bound.
This *might* be a special case of handling keyword-only arguments, since that's how the partially bound arguments appear in the underlying `__signature__`.

Right now, we don't allow keyword-only arguments in our `Callable` processing.
This exception is raised by `ibis.common.patterns.CallableWith` mostly because I haven't had enough time to investigate how to properly type annotate and handle keyword only arguments. Strictly speaking this validation can be omitted.

I'd need to further investigate what kind of validation could/should we actually do in `CallableWith.match()`. 
I have a fix for this incoming - we can relax the check here since these keyword-only parameters have default values.

Resolving patch:
diff --git a/ibis/common/patterns.py b/ibis/common/patterns.py
--- a/ibis/common/patterns.py
+++ b/ibis/common/patterns.py
@@ -1520,20 +1520,18 @@ def match(self, value, context):
         fn = annotated(self.args, self.return_, value)
 
         has_varargs = False
-        positional, keyword_only = [], []
+        positional = []
         for p in fn.__signature__.parameters.values():
             if p.kind in (Parameter.POSITIONAL_ONLY, Parameter.POSITIONAL_OR_KEYWORD):
                 positional.append(p)
-            elif p.kind is Parameter.KEYWORD_ONLY:
-                keyword_only.append(p)
+            elif p.kind is Parameter.KEYWORD_ONLY and p.default is Parameter.empty:
+                raise MatchError(
+                    "Callable has mandatory keyword-only arguments which cannot be specified"
+                )
             elif p.kind is Parameter.VAR_POSITIONAL:
                 has_varargs = True
 
-        if keyword_only:
-            raise MatchError(
-                "Callable has mandatory keyword-only arguments which cannot be specified"
-            )
-        elif len(positional) > len(self.args):
+        if len(positional) > len(self.args):
             # Callable has more positional arguments than expected")
             return NoMatch
         elif len(positional) < len(self.args) and not has_varargs:
diff --git a/ibis/expr/operations/arrays.py b/ibis/expr/operations/arrays.py
--- a/ibis/expr/operations/arrays.py
+++ b/ibis/expr/operations/arrays.py
@@ -80,7 +80,7 @@ class ArrayApply(Value):
 
     @attribute
     def parameter(self):
-        (name,) = self.func.__signature__.parameters.keys()
+        name = next(iter(self.func.__signature__.parameters.keys()))
         return name
 
     @attribute
diff --git a/ibis/expr/types/arrays.py b/ibis/expr/types/arrays.py
--- a/ibis/expr/types/arrays.py
+++ b/ibis/expr/types/arrays.py
@@ -396,8 +396,8 @@ def map(self, func: Callable[[ir.Value], ir.Value]) -> ir.ArrayValue:
         """
 
         @functools.wraps(func)
-        def wrapped(x):
-            return func(x.to_expr())
+        def wrapped(x, **kwargs):
+            return func(x.to_expr(), **kwargs)
 
         return ops.ArrayMap(self, func=wrapped).to_expr()
 
@@ -444,8 +444,8 @@ def filter(
         """
 
         @functools.wraps(predicate)
-        def wrapped(x):
-            return predicate(x.to_expr())
+        def wrapped(x, **kwargs):
+            return predicate(x.to_expr(), **kwargs)
 
         return ops.ArrayFilter(self, func=wrapped).to_expr()
 


Test patch:
diff --git a/ibis/backends/tests/test_array.py b/ibis/backends/tests/test_array.py
--- a/ibis/backends/tests/test_array.py
+++ b/ibis/backends/tests/test_array.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import contextlib
+import functools
 import os
 
 import numpy as np
@@ -527,9 +528,14 @@ def test_array_slice(backend, start, stop):
 )
 def test_array_map(backend, con, input, output):
     t = ibis.memtable(input, schema=ibis.schema(dict(a="!array<int8>")))
+    expected = pd.DataFrame(output)
+
     expr = t.select(a=t.a.map(lambda x: x + 1))
     result = con.execute(expr)
-    expected = pd.DataFrame(output)
+    backend.assert_frame_equal(result, expected)
+
+    expr = t.select(a=t.a.map(functools.partial(lambda x, y: x + y, y=1)))
+    result = con.execute(expr)
     backend.assert_frame_equal(result, expected)
 
 
@@ -555,9 +561,14 @@ def test_array_map(backend, con, input, output):
 )
 def test_array_filter(backend, con, input, output):
     t = ibis.memtable(input, schema=ibis.schema(dict(a="!array<int8>")))
+    expected = pd.DataFrame(output)
+
     expr = t.select(a=t.a.filter(lambda x: x > 1))
     result = con.execute(expr)
-    expected = pd.DataFrame(output)
+    backend.assert_frame_equal(result, expected)
+
+    expr = t.select(a=t.a.filter(functools.partial(lambda x, y: x > y, y=1)))
+    result = con.execute(expr)
     backend.assert_frame_equal(result, expected)
 
 
diff --git a/ibis/common/tests/test_patterns.py b/ibis/common/tests/test_patterns.py
--- a/ibis/common/tests/test_patterns.py
+++ b/ibis/common/tests/test_patterns.py
@@ -527,7 +527,10 @@ def func_with_args(a, b, *args):
     def func_with_kwargs(a, b, c=1, **kwargs):
         return str(a) + b + str(c)
 
-    def func_with_mandatory_kwargs(*, c):
+    def func_with_optional_keyword_only_kwargs(a, *, c=1):
+        return a + c
+
+    def func_with_required_keyword_only_kwargs(*, c):
         return c
 
     p = CallableWith([InstanceOf(int), InstanceOf(str)])
@@ -535,7 +538,7 @@ def func_with_mandatory_kwargs(*, c):
 
     msg = "Callable has mandatory keyword-only arguments which cannot be specified"
     with pytest.raises(MatchError, match=msg):
-        p.match(func_with_mandatory_kwargs, context={})
+        p.match(func_with_required_keyword_only_kwargs, context={})
 
     # Callable has more positional arguments than expected
     p = CallableWith([InstanceOf(int)] * 2)
@@ -556,6 +559,10 @@ def func_with_mandatory_kwargs(*, c):
     with pytest.raises(ValidationError, match="2 doesn't match InstanceOf"):
         wrapped(1, 2)
 
+    p = CallableWith([InstanceOf(int)])
+    wrapped = p.match(func_with_optional_keyword_only_kwargs, context={})
+    assert wrapped(1) == 2
+
 
 def test_pattern_list():
     p = PatternSequence([1, 2, InstanceOf(int), ...])
diff --git a/ibis/tests/expr/test_value_exprs.py b/ibis/tests/expr/test_value_exprs.py
--- a/ibis/tests/expr/test_value_exprs.py
+++ b/ibis/tests/expr/test_value_exprs.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import functools
 import ipaddress
 import operator
 import uuid
@@ -1536,12 +1537,32 @@ def test_array_map():
     assert result_float.type() == dt.Array(dt.float64)
 
 
+def test_array_map_partial():
+    arr = ibis.array([1, 2, 3])
+
+    def add(x, y):
+        return x + y
+
+    result = arr.map(functools.partial(add, y=2))
+    assert result.type() == dt.Array(dt.int16)
+
+
 def test_array_filter():
     arr = ibis.array([1, 2, 3])
     result = arr.filter(is_negative)
     assert result.type() == arr.type()
 
 
+def test_array_filter_partial():
+    arr = ibis.array([1, 2, 3])
+
+    def equal(x, y):
+        return x == y
+
+    result = arr.filter(functools.partial(equal, y=2))
+    assert result.type() == arr.type()
+
+
 @pytest.mark.parametrize(
     ("func", "expected_type"),
     [