@@ -231,7 +231,7 @@ fn prepare_output(
231231 array_addrs : jlongArray ,
232232 schema_addrs : jlongArray ,
233233 output_batch : RecordBatch ,
234- exec_context : & mut ExecutionContext ,
234+ validate : bool ,
235235) -> CometResult < jlong > {
236236 let array_address_array = unsafe { JLongArray :: from_raw ( array_addrs) } ;
237237 let num_cols = env. get_array_length ( & array_address_array) ? as usize ;
@@ -255,7 +255,7 @@ fn prepare_output(
255255 ) ) ) ;
256256 }
257257
258- if exec_context . debug_native {
258+ if validate {
259259 // Validate the output arrays.
260260 for array in results. iter ( ) {
261261 let array_data = array. to_data ( ) ;
@@ -275,9 +275,6 @@ fn prepare_output(
275275 i += 1 ;
276276 }
277277
278- // Update metrics
279- update_metrics ( env, exec_context) ?;
280-
281278 Ok ( num_rows as jlong )
282279}
283280
@@ -356,22 +353,22 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
356353 let next_item = exec_context. stream . as_mut ( ) . unwrap ( ) . next ( ) ;
357354 let poll_output = exec_context. runtime . block_on ( async { poll ! ( next_item) } ) ;
358355
356+ // Update metrics
357+ update_metrics ( & mut env, exec_context) ?;
358+
359359 match poll_output {
360360 Poll :: Ready ( Some ( output) ) => {
361+ // prepare output for FFI transfer
361362 return prepare_output (
362363 & mut env,
363364 array_addrs,
364365 schema_addrs,
365366 output?,
366- exec_context,
367+ exec_context. debug_native ,
367368 ) ;
368369 }
369370 Poll :: Ready ( None ) => {
370371 // Reaches EOF of output.
371-
372- // Update metrics
373- update_metrics ( & mut env, exec_context) ?;
374-
375372 if exec_context. explain_native {
376373 if let Some ( plan) = & exec_context. root_op {
377374 let formatted_plan_str =
@@ -391,9 +388,6 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan(
391388 // A poll pending means there are more than one blocking operators,
392389 // we don't need go back-forth between JVM/Native. Just keeping polling.
393390 Poll :: Pending => {
394- // Update metrics
395- update_metrics ( & mut env, exec_context) ?;
396-
397391 // Pull input batches
398392 pull_input_batches ( exec_context) ?;
399393
0 commit comments