From 322654aeaf625172ba9c45d31bc2fe8a3f628afc Mon Sep 17 00:00:00 2001 From: Yicong-Huang <17627829+Yicong-Huang@users.noreply.github.com> Date: Tue, 12 May 2026 18:52:25 +0000 Subject: [PATCH] fix: pass ArrowBatchedUDF benchmark input_type via EvalConf --- python/benchmarks/bench_eval_type.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/python/benchmarks/bench_eval_type.py b/python/benchmarks/bench_eval_type.py index a4fda48dfcb16..14020d6f22eca 100644 --- a/python/benchmarks/bench_eval_type.py +++ b/python/benchmarks/bench_eval_type.py @@ -440,8 +440,8 @@ def peakmem_worker(self, *args): class _ArrowBatchedBenchMixin: """Provides ``_write_scenario`` for SQL_ARROW_BATCHED_UDF. - Writes the extra ``input_type`` (StructType JSON) that the wire protocol - requires before the UDF payload. + Passes the ``input_type`` (StructType JSON) through ``EvalConf``, which the + non-legacy and legacy paths both read via ``eval_conf.input_type``. """ # Row-by-row processing is ~100x slower than columnar Arrow UDFs, @@ -486,15 +486,12 @@ def _write_scenario(self, scenario, udf_name, buf): if ret_type is None: ret_type = schema.fields[0].dataType - def write_udf(b): - MockProtocolWriter.write_utf8(schema.json(), b) - MockProtocolWriter.write_udf_payload(udf_func, ret_type, arg_offsets, b) - MockProtocolWriter.write_worker_input( PythonEvalType.SQL_ARROW_BATCHED_UDF, - write_udf, + lambda b: MockProtocolWriter.write_udf_payload(udf_func, ret_type, arg_offsets, b), lambda b: MockProtocolWriter.write_data_payload(iter(batches), b), buf, + eval_conf={"input_type": schema.json()}, )