Skip to content

Commit 58dee73

Browse files
authored
minor: refactor prepare_output so that it does not require an ExecutionContext (#1194)
1 parent 639fa2f commit 58dee73

File tree

1 file changed

+7
-13
lines changed

1 file changed

+7
-13
lines changed

native/core/src/execution/jni_api.rs

Lines changed: 7 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)