rtpsw commented on code in PR #14682:
URL: https://github.com/apache/arrow/pull/14682#discussion_r1064040644
##########
python/pyarrow/tests/test_udf.py:
##########
@@ -504,3 +504,112 @@ def test_input_lifetime(unary_func_fixture):
# Calling a UDF should not have kept `v` alive longer than required
v = None
assert proxy_pool.bytes_allocated() == 0
+
+
+def _record_batch_from_iters(schema, *iters):
+ arrays = [pa.array(list(v), type=schema[i].type)
+ for i, v in enumerate(iters)]
+ return pa.RecordBatch.from_arrays(arrays=arrays, schema=schema)
+
+
+def _record_batch_for_range(schema, n):
+ return _record_batch_from_iters(schema,
+ range(n, n + 10),
+ range(n + 1, n + 11))
+
+
+def make_udt_func(schema, batch_gen):
+ def udf_func(ctx):
+ class UDT:
+ def __init__(self):
+ self.caller = None
+
+ def __call__(self, ctx):
+ try:
+ if self.caller is None:
+ self.caller, ctx = batch_gen(ctx).send, None
+ batch = self.caller(ctx)
+ except StopIteration:
+ arrays = [pa.array([], type=field.type)
+ for field in schema]
+ batch = pa.RecordBatch.from_arrays(
+ arrays=arrays, schema=schema)
Review Comment:
The reason is that, with the current design here, a UDTF must return an
array (see [this
post](https://github.com/apache/arrow/pull/14682/#discussion_r1045132971)). I
don't know whether it's a good enough reason... If we'd like to allow a
`StopIteration`, then I'd prefer this be done in a separate issue.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
