andygrove commented on code in PR #4306:
URL: https://github.com/apache/datafusion-comet/pull/4306#discussion_r3229048622
##########
common/src/main/java/org/apache/comet/udf/CometUdfBridge.java:
##########
@@ -48,13 +50,45 @@ public class CometUdfBridge {
* @param inputSchemaPtrs addresses of pre-allocated FFI_ArrowSchema structs
(one per input)
* @param outArrayPtr address of pre-allocated FFI_ArrowArray for the result
* @param outSchemaPtr address of pre-allocated FFI_ArrowSchema for the
result
+ * @param numRows row count of the current batch. Mirrors DataFusion's {@code
+ * ScalarFunctionArgs.number_rows}; the only batch-size signal a
zero-input UDF (e.g. a
+ * zero-arg non-deterministic ScalaUDF) ever sees.
+ * @param taskContext propagated Spark {@link TaskContext} from the driving
Spark task thread, or
+ * {@code null} outside a Spark task. Installed as the thread-local for
the duration of the
+ * call when the current thread has none, so partition-sensitive
built-ins ({@code Rand},
+ * {@code Uuid}, {@code MonotonicallyIncreasingID}) work from Tokio
workers. Cleared in {@code
+ * finally} to avoid leaking across worker reuse.
*/
public static void evaluate(
String udfClassName,
long[] inputArrayPtrs,
long[] inputSchemaPtrs,
long outArrayPtr,
- long outSchemaPtr) {
+ long outSchemaPtr,
+ int numRows,
+ TaskContext taskContext) {
+ boolean installedTaskContext = false;
+ if (taskContext != null && TaskContext.get() == null) {
Review Comment:
Could you explain more about tthe `TaskContext.get() == null` check. Is it
possible that there could be a stale reference?
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
