@@ -1454,6 +1454,29 @@ def evaluate_all(self, values: list[pa.Array], num_rows: int) -> pa.Array:
14541454 return pa .array ([first ] * num_rows )
14551455
14561456
1457+ def _total (name = "total" ):
1458+ """An aggregate function under a name the caller picks."""
1459+ return udaf (
1460+ _Total ,
1461+ pa .int64 (),
1462+ pa .int64 (),
1463+ [pa .int64 ()],
1464+ volatility = "stable" ,
1465+ name = name ,
1466+ )
1467+
1468+
1469+ def _first (name = "first_value_of" ):
1470+ """A window function under a name the caller picks."""
1471+ return udwf (
1472+ _First ,
1473+ pa .int64 (),
1474+ pa .int64 (),
1475+ volatility = "immutable" ,
1476+ name = name ,
1477+ )
1478+
1479+
14571480class _FunctionExtension :
14581481 """Contributes functions and nothing else.
14591482
@@ -1483,23 +1506,9 @@ def test_with_extensions_registers_a_declared_udf(ctx):
14831506
14841507def test_with_extensions_registers_udafs_and_udwfs (ctx ):
14851508 """The other two function kinds install the same way."""
1486- total = udaf (
1487- _Total ,
1488- pa .int64 (),
1489- pa .int64 (),
1490- [pa .int64 ()],
1491- volatility = "stable" ,
1492- name = "total" ,
1493- )
1494- first = udwf (
1495- _First ,
1496- pa .int64 (),
1497- pa .int64 (),
1498- volatility = "immutable" ,
1499- name = "first_value_of" ,
1509+ result = ctx .with_extensions (
1510+ _FunctionExtension (udafs = (_total (),), udwfs = (_first (),))
15001511 )
1501-
1502- result = ctx .with_extensions (_FunctionExtension (udafs = (total ,), udwfs = (first ,)))
15031512 result .from_pydict ({"a" : [1 , 2 , 3 ]}, name = "nums" )
15041513
15051514 assert result .sql ("SELECT total(a) FROM nums" ).collect ()[0 ].column (0 ) == pa .array (
@@ -1521,26 +1530,39 @@ def test_with_extensions_registers_on_the_shared_session(ctx):
15211530 assert ctx .udf ("double" ).name == "double"
15221531
15231532
1524- def test_with_extensions_rejects_a_name_two_extensions_claim (ctx ):
1533+ @pytest .mark .parametrize (
1534+ ("field" , "make" , "label" , "lookup" ),
1535+ [
1536+ ("udfs" , _doubler , "scalar function" , "udf" ),
1537+ ("udafs" , _total , "aggregate function" , "udaf" ),
1538+ ("udwfs" , _first , "window function" , "udwf" ),
1539+ ],
1540+ )
1541+ def test_with_extensions_rejects_a_name_two_extensions_claim (
1542+ ctx , field , make , label , lookup
1543+ ):
15251544 """Registrations have no fall-through, so a clash cannot be resolved by order.
15261545
15271546 Unlike codecs, which dispatch by id, a second function under one name would
1528- silently replace the first.
1547+ silently replace the first. Parametrized over the kinds to pin each
1548+ ``_FUNCTION_KINDS`` row's field and label wiring, not just the machinery.
15291549
15301550 A codec-carrying bundle rides along to pin the other half of the
15311551 transaction: resolution runs *after* the codec chains are built, so this
15321552 failure lands between the two steps, and the chains must not reach the
15331553 session either.
15341554 """
1535- with pytest .raises (ValueError , match = r"scalar function named 'double'" ):
1555+ name = make ().name
1556+
1557+ with pytest .raises (ValueError , match = rf"{ label } named '{ name } '" ):
15361558 ctx .with_extensions (
15371559 _CodecOnlyExtension (),
1538- _FunctionExtension (udfs = ( _doubler (),)),
1539- _FunctionExtension (udfs = ( _doubler (),)),
1560+ _FunctionExtension (** { field : ( make (),)} ),
1561+ _FunctionExtension (** { field : ( make (),)} ),
15401562 )
15411563
15421564 with pytest .raises (KeyError ):
1543- ctx . udf ( "double" )
1565+ getattr ( ctx , lookup )( name )
15441566 assert ctx .logical_extension_codec_ids () == []
15451567 assert ctx .physical_extension_codec_ids () == []
15461568
@@ -1560,6 +1582,8 @@ def test_with_extensions_rejects_one_extension_passed_twice(ctx):
15601582 ctx .with_extensions (extension , extension )
15611583
15621584 assert "rename" not in str (excinfo .value )
1585+ with pytest .raises (KeyError ):
1586+ ctx .udf ("double" )
15631587
15641588
15651589def test_with_extensions_rejects_a_name_one_extension_claims_twice (ctx ):
0 commit comments