Skip to content

Commit

Permalink
remove bigdecimal dependency and use arrow's builtin decimal cast.
Browse files Browse the repository at this point in the history
implement better logging with stage and partition ids.

other minor code refactoring.
  • Loading branch information
zhangli20 committed Nov 25, 2024
1 parent 9a10161 commit ce474f9
Show file tree
Hide file tree
Showing 20 changed files with 452 additions and 591 deletions.
16 changes: 0 additions & 16 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

33 changes: 1 addition & 32 deletions native-engine/blaze/src/exec.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,25 +19,21 @@ use blaze_jni_bridge::{
jni_bridge::JavaClasses,
*,
};
use blaze_serde::protobuf::TaskDefinition;
use datafusion::{
common::Result,
error::DataFusionError,
execution::{
disk_manager::DiskManagerConfig,
runtime_env::{RuntimeConfig, RuntimeEnv},
},
physical_plan::{displayable, ExecutionPlan},
prelude::{SessionConfig, SessionContext},
};
use datafusion_ext_commons::df_execution_err;
use datafusion_ext_plans::memmgr::MemManager;
use jni::{
objects::{JClass, JObject},
JNIEnv,
};
use once_cell::sync::OnceCell;
use prost::Message;

use crate::{handle_unwinded_scope, logging::init_logging, rt::NativeExecutionRuntime};

Expand Down Expand Up @@ -81,38 +77,11 @@ pub extern "system" fn Java_org_apache_spark_sql_blaze_JniBridge_callNative(
})?;
let native_wrapper = jni_new_global_ref!(native_wrapper)?;

// decode plan
let raw_task_definition = jni_call!(
BlazeCallNativeWrapper(native_wrapper.as_obj())
.getRawTaskDefinition() -> JObject)?;
let task_definition = TaskDefinition::decode(
jni_convert_byte_array!(raw_task_definition.as_obj())?.as_slice(),
)
.or_else(|err| df_execution_err!("cannot decode execution plan: {err:?}"))?;

let task_id = &task_definition.task_id.expect("task_id is empty");
let plan = &task_definition.plan.expect("plan is empty");
drop(raw_task_definition);

// get execution plan
let execution_plan: Arc<dyn ExecutionPlan> = plan
.try_into()
.or_else(|err| df_execution_err!("cannot create execution plan: {err:?}"))?;
let execution_plan_displayable = displayable(execution_plan.as_ref())
.indent(true)
.to_string();
log::info!("Creating native execution plan succeeded");
log::info!(" task_id={task_id:?}");
log::info!(" execution plan:\n{execution_plan_displayable}");

// execute to stream
// create execution runtime
let runtime = Box::new(NativeExecutionRuntime::start(
native_wrapper,
execution_plan,
task_id.partition_id as usize,
SESSION.get().unwrap().task_ctx(),
)?);
log::info!("Blaze native thread created");

// returns runtime raw pointer
Ok::<_, DataFusionError>(Box::into_raw(runtime) as usize as i64)
Expand Down
11 changes: 9 additions & 2 deletions native-engine/blaze/src/logging.rs
Original file line number Diff line number Diff line change
Expand Up @@ -12,11 +12,16 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::time::Instant;
use std::{cell::Cell, time::Instant};

use log::{Level, LevelFilter, Log, Metadata, Record};
use once_cell::sync::OnceCell;

thread_local! {
pub static THREAD_STAGE_ID: Cell<usize> = Cell::new(0);
pub static THREAD_PARTITION_ID: Cell<usize> = Cell::new(0);
}

const MAX_LEVEL: Level = Level::Info;

pub fn init_logging() {
Expand All @@ -43,8 +48,10 @@ impl Log for SimpleLogger {
if self.enabled(record.metadata()) {
let elapsed = Instant::now() - self.start_instant;
let elapsed_sec = elapsed.as_secs_f64();
let stage_id = THREAD_STAGE_ID.get();
let partition_id = THREAD_PARTITION_ID.get();
eprintln!(
"(+{elapsed_sec:.3}s) [{}] Blaze - {}",
"(+{elapsed_sec:.3}s) [{}] (stage: {stage_id}, partition: {partition_id}) - {}",
record.level(),
record.args()
);
Expand Down
Loading

0 comments on commit ce474f9

Please sign in to comment.