Adding JNI abstractions for logging in rust through RustLoggerBridge

Signed-off-by: Raghuvansh Raj <[email protected]>

Author: raghuvanshraj

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/bridge/RustBridge.java

@@ -21,8 +21,13 @@ public class RustBridge {
 
     static {
         NativeLibraryLoader.load("parquet_dataformat_jni");
+
+        initLogger();
     }
 
+    // Logger initialization method
+    public static native void initLogger();
+
     // Enhanced native methods that handle validation and provide better error reporting
     public static native void createWriter(String file, long schemaAddress) throws IOException;
     public static native void write(String file, long arrayAddress, long schemaAddress) throws IOException;

modules/parquet-data-format/src/main/java/com/parquet/parquetdataformat/bridge/RustLoggerBridge.java

@@ -0,0 +1,34 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package com.parquet.parquetdataformat.bridge;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class RustLoggerBridge {
+
+    private static final Logger logger = LoggerFactory.getLogger(RustLoggerBridge.class);
+
+    // Instance methods for direct Java usage
+    public static void logInfo(String message) {
+        logger.info(message);
+    }
+
+    public static void logWarn(String message) {
+        logger.warn(message);
+    }
+
+    public static void logError(String message) {
+        logger.error(message);
+    }
+
+    public static void logDebug(String message) {
+        logger.debug(message);
+    }
+}

modules/parquet-data-format/src/main/rust/src/lib.rs

@@ -1,19 +1,17 @@
 use jni::objects::{JClass, JString};
 use jni::sys::{jint, jlong};
-use jni::JNIEnv;
+use jni::{JNIEnv, JavaVM};
 use dashmap::DashMap;
 use arrow::record_batch::RecordBatch;
 use parquet::arrow::ArrowWriter;
 use std::fs::File;
 use std::sync::{Arc, Mutex};
 use lazy_static::lazy_static;
 use arrow::ffi::{FFI_ArrowSchema, FFI_ArrowArray};
-use std::fs::OpenOptions;
-use std::io::Write;
-use chrono::Utc;
 use parquet::basic::{Compression, ZstdLevel};
 use parquet::file::properties::WriterProperties;
 
+pub mod logger;
 pub mod parquet_merge;
 pub use parquet_merge::*;
 
@@ -27,21 +25,15 @@ struct NativeParquetWriter;
 impl NativeParquetWriter {
 
     fn create_writer(filename: String, schema_address: i64) -> Result<(), Box<dyn std::error::Error>> {
-        let log_msg = format!("[RUST] create_writer called for file: {}, schema_address: {}\n", filename, schema_address);
-        println!("{}", log_msg.trim());
-        Self::log_to_file(&log_msg);
+        logger::log_info(&format!("[RUST] create_writer called for file: {}, schema_address: {}", filename, schema_address));
 
         let arrow_schema = unsafe { FFI_ArrowSchema::from_raw(schema_address as *mut _) };
         let schema = Arc::new(arrow::datatypes::Schema::try_from(&arrow_schema)?);
 
-        let schema_msg = format!("[RUST] Schema created with {} fields\n", schema.fields().len());
-        println!("{}", schema_msg.trim());
-        Self::log_to_file(&schema_msg);
+        logger::log_info(&format!("[RUST] Schema created with {} fields", schema.fields().len()));
 
         for (i, field) in schema.fields().iter().enumerate() {
-            let field_msg = format!("[RUST] Field {}: {} ({})\n", i, field.name(), field.data_type());
-            println!("{}", field_msg.trim());
-            Self::log_to_file(&field_msg);
+            logger::log_debug(&format!("[RUST] Field {}: {} ({})", i, field.name(), field.data_type()));
         }
 
         let file = File::create(&filename)?;
@@ -56,29 +48,21 @@ impl NativeParquetWriter {
     }
 
     fn write_data(filename: String, array_address: i64, schema_address: i64) -> Result<(), Box<dyn std::error::Error>> {
-        let log_msg = format!("[RUST] write_data called for file: {}, array_address: {}, schema_address: {}\n", filename, array_address, schema_address);
-        println!("{}", log_msg.trim());
-        Self::log_to_file(&log_msg);
+        logger::log_info(&format!("[RUST] write_data called for file: {}, array_address: {}, schema_address: {}", filename, array_address, schema_address));
 
         unsafe {
             let arrow_schema = FFI_ArrowSchema::from_raw(schema_address as *mut _);
             let arrow_array = FFI_ArrowArray::from_raw(array_address as *mut _);
 
             match arrow::ffi::from_ffi(arrow_array, &arrow_schema) {
                 Ok(array_data) => {
-                    let data_msg = format!("[RUST] Successfully imported array_data, length: {}\n", array_data.len());
-                    println!("{}", data_msg.trim());
-                    Self::log_to_file(&data_msg);
+                    logger::log_debug(&format!("[RUST] Successfully imported array_data, length: {}", array_data.len()));
 
                     let array: Arc<dyn arrow::array::Array> = arrow::array::make_array(array_data);
-                    let array_msg = format!("[RUST] Array type: {:?}, length: {}\n", array.data_type(), array.len());
-                    println!("{}", array_msg.trim());
-                    Self::log_to_file(&array_msg);
+                    logger::log_debug(&format!("[RUST] Array type: {:?}, length: {}", array.data_type(), array.len()));
 
                     if let Some(struct_array) = array.as_any().downcast_ref::<arrow::array::StructArray>() {
-                        let struct_msg = format!("[RUST] Successfully cast to StructArray with {} columns\n", struct_array.num_columns());
-                        println!("{}", struct_msg.trim());
-                        Self::log_to_file(&struct_msg);
+                        logger::log_debug(&format!("[RUST] Successfully cast to StructArray with {} columns", struct_array.num_columns()));
 
                         let schema = Arc::new(arrow::datatypes::Schema::new(
                             struct_array.fields().clone()
@@ -89,70 +73,50 @@ impl NativeParquetWriter {
                             struct_array.columns().to_vec(),
                         )?;
 
-                        let batch_msg = format!("[RUST] Created RecordBatch with {} rows and {} columns\n", record_batch.num_rows(), record_batch.num_columns());
-                        println!("{}", batch_msg.trim());
-                        Self::log_to_file(&batch_msg);
+                        logger::log_info(&format!("[RUST] Created RecordBatch with {} rows and {} columns", record_batch.num_rows(), record_batch.num_columns()));
 
                         if let Some(writer_arc) = WRITER_MANAGER.get(&filename) {
-                            let write_msg = "[RUST] Writing RecordBatch to file\n";
-                            println!("{}", write_msg.trim());
-                            Self::log_to_file(write_msg);
+                            logger::log_debug("[RUST] Writing RecordBatch to file");
                             let mut writer = writer_arc.lock().unwrap();
                             writer.write(&record_batch)?;
-                            let success_msg = "[RUST] Successfully wrote RecordBatch\n";
-                            println!("{}", success_msg.trim());
-                            Self::log_to_file(success_msg);
+                            logger::log_info("[RUST] Successfully wrote RecordBatch");
                         } else {
-                            let error_msg = format!("[RUST] ERROR: No writer found for file: {}\n", filename);
-                            println!("{}", error_msg.trim());
-                            Self::log_to_file(&error_msg);
+                            logger::log_error(&format!("[RUST] ERROR: No writer found for file: {}", filename));
                         }
                         Ok(())
                     } else {
-                        let error_msg = format!("[RUST] ERROR: Array is not a StructArray, type: {:?}\n", array.data_type());
-                        println!("{}", error_msg.trim());
-                        Self::log_to_file(&error_msg);
+                        logger::log_error(&format!("[RUST] ERROR: Array is not a StructArray, type: {:?}", array.data_type()));
                         Err("Expected struct array from VectorSchemaRoot".into())
                     }
                 }
                 Err(e) => {
-                    let error_msg = format!("[RUST] ERROR: Failed to import from FFI: {:?}\n", e);
-                    println!("{}", error_msg.trim());
-                    Self::log_to_file(&error_msg);
+                    logger::log_error(&format!("[RUST] ERROR: Failed to import from FFI: {:?}", e));
                     Err(e.into())
                 }
             }
         }
     }
 
     fn close_writer(filename: String) -> Result<(), Box<dyn std::error::Error>> {
-        let log_msg = format!("[RUST] close_writer called for file: {}\n", filename);
-        println!("{}", log_msg.trim());
-        Self::log_to_file(&log_msg);
+        logger::log_info(&format!("[RUST] close_writer called for file: {}", filename));
 
         if let Some((_, writer_arc)) = WRITER_MANAGER.remove(&filename) {
             match Arc::try_unwrap(writer_arc) {
                 Ok(mutex) => {
                     let mut writer = mutex.into_inner().unwrap();
                     match writer.close() {
                         Ok(_) => {
-                            let success_msg = format!("[RUST] Successfully closed writer for file: {}\n", filename);
-                            println!("{}", success_msg.trim());
-                            Self::log_to_file(&success_msg);
+                            logger::log_info(&format!("[RUST] Successfully closed writer for file: {}", filename));
                             Ok(())
                         }
                         Err(e) => {
-                            let error_msg = format!("[RUST] ERROR: Failed to close writer for file: {}\n", filename);
-                            println!("{}", error_msg.trim());
-                            Self::log_to_file(&error_msg);
+                            logger::log_error(&format!("[RUST] ERROR: Failed to close writer for file: {}", filename));
                             Err(e.into())
                         }
                     }
                 }
                 Err(_) => {
-                    let error_msg = format!("[RUST] ERROR: Writer still in use for file: {}\n", filename);
-                    println!("{}", error_msg.trim());
-                    Self::log_to_file(&error_msg);
+                    logger::log_error(&format!("[RUST] ERROR: Writer still in use for file: {}", filename));
                     Err("Writer still in use".into())
                 }
             }
@@ -162,37 +126,27 @@ impl NativeParquetWriter {
     }
 
     fn flush_to_disk(filename: String) -> Result<(), Box<dyn std::error::Error>> {
-        let log_msg = format!("[RUST] fsync_file called for file: {}\n", filename);
-        println!("{}", log_msg.trim());
-        Self::log_to_file(&log_msg);
+        logger::log_info(&format!("[RUST] fsync_file called for file: {}", filename));
 
         if let Some(file) = FILE_MANAGER.get_mut(&filename) {
             match file.sync_all() {
                 Ok(_) => {
-                    let success_msg = format!("[RUST] Successfully fsynced file: {}\n", filename);
-                    println!("{}", success_msg.trim());
-                    Self::log_to_file(&success_msg);
+                    logger::log_info(&format!("[RUST] Successfully fsynced file: {}", filename));
                     Ok(())
                 }
                 Err(e) => {
-                    let error_msg = format!("[RUST] ERROR: Failed to fsync file: {}\n", filename);
-                    println!("{}", error_msg.trim());
-                    Self::log_to_file(&error_msg);
+                    logger::log_error(&format!("[RUST] ERROR: Failed to fsync file: {}", filename));
                     Err(e.into())
                 }
             }
         } else {
-            let error_msg = format!("[RUST] ERROR: File not found for fsync: {}\n", filename);
-            println!("{}", error_msg.trim());
-            Self::log_to_file(&error_msg);
+            logger::log_error(&format!("[RUST] ERROR: File not found for fsync: {}", filename));
             Err("File not found".into())
         }
     }
 
     fn get_filtered_writer_memory_usage(path_prefix: String) -> Result<usize, Box<dyn std::error::Error>> {
-        let _log_msg = format!("[RUST] get_filtered_writer_memory_usage called with prefix: {}\n", path_prefix);
-//         println!("{}", _log_msg.trim());
-//         Self::log_to_file(&_log_msg);
+        logger::log_debug(&format!("[RUST] get_filtered_writer_memory_usage called with prefix: {}", path_prefix));
 
         let mut total_memory = 0;
         let mut writer_count = 0;
@@ -208,29 +162,24 @@ impl NativeParquetWriter {
                     total_memory += memory_usage;
                     writer_count += 1;
 
-                    let usage_msg = format!("[RUST] Filtered Writer {}: {} bytes\n", filename, memory_usage);
-                    println!("{}", usage_msg.trim());
-                    Self::log_to_file(&usage_msg);
+                    logger::log_debug(&format!("[RUST] Filtered Writer {}: {} bytes", filename, memory_usage));
                 }
             }
         }
 
-        let _total_msg = format!("[RUST] Total memory usage across {} filtered ArrowWriters (prefix: {}): {} bytes\n", writer_count, path_prefix, total_memory);
-        //println!("{}", _total_msg.trim());
-        //Self::log_to_file(&_total_msg);
+        logger::log_debug(&format!("[RUST] Total memory usage across {} filtered ArrowWriters (prefix: {}): {} bytes", writer_count, path_prefix, total_memory));
 
         Ok(total_memory)
     }
+}
 
-    fn log_to_file(message: &str) {
-        if let Ok(mut file) = OpenOptions::new()
-            .create(true)
-            .append(true)
-            .open("/tmp/rust_parquet_debug.log") {
-            let timestamp = Utc::now().format("%Y-%m-%d %H:%M:%S%.3f UTC");
-            let timestamped_message = format!("[{}] {}", timestamp, message);
-            let _ = file.write_all(timestamped_message.as_bytes());
-        }
+#[unsafe(no_mangle)]
+pub extern "system" fn Java_com_parquet_parquetdataformat_bridge_RustBridge_initLogger(
+    env: JNIEnv,
+    _class: JClass,
+) {
+    if let Ok(jvm) = env.get_java_vm() {
+        logger::init_logger(jvm);
     }
 }
 

modules/parquet-data-format/src/main/rust/src/logger.rs

@@ -0,0 +1,103 @@
+use jni::{JNIEnv, JavaVM};
+use std::sync::OnceLock;
+
+static JAVA_VM: OnceLock<JavaVM> = OnceLock::new();
+
+/// Initialize the logger with the JVM instance
+pub fn init_logger(jvm: JavaVM) {
+    JAVA_VM.set(jvm).ok();
+}
+
+/// Log an info message through JNI callback to Java
+pub fn log_info(message: &str) {
+    if let Some(jvm) = JAVA_VM.get() {
+        if let Ok(mut env) = jvm.attach_current_thread() {
+            call_java_logger(&mut env, "logInfo", message);
+        }
+    }
+}
+
+/// Log a warning message through JNI callback to Java
+pub fn log_warn(message: &str) {
+    if let Some(jvm) = JAVA_VM.get() {
+        if let Ok(mut env) = jvm.attach_current_thread() {
+            call_java_logger(&mut env, "logWarn", message);
+        }
+    }
+}
+
+/// Log an error message through JNI callback to Java
+pub fn log_error(message: &str) {
+    if let Some(jvm) = JAVA_VM.get() {
+        if let Ok(mut env) = jvm.attach_current_thread() {
+            call_java_logger(&mut env, "logError", message);
+        }
+    }
+}
+
+/// Log a debug message through JNI callback to Java
+pub fn log_debug(message: &str) {
+    if let Some(jvm) = JAVA_VM.get() {
+        if let Ok(mut env) = jvm.attach_current_thread() {
+            call_java_logger(&mut env, "logDebug", message);
+        }
+    }
+}
+
+/// Internal function to call the Java logger method
+fn call_java_logger(env: &mut JNIEnv, method_name: &str, message: &str) {
+    let result = (|| -> Result<(), Box<dyn std::error::Error>> {
+        // Find the RustLoggerBridge class
+        let class = env.find_class("com/parquet/parquetdataformat/bridge/RustLoggerBridge")?;
+
+        // Convert Rust string to Java string
+        let java_message = env.new_string(message)?;
+
+        // Call the static method
+        env.call_static_method(
+            class,
+            method_name,
+            "(Ljava/lang/String;)V",
+            &[(&java_message).into()],
+        )?;
+
+        Ok(())
+    })();
+
+    // If logging fails, fall back to println as last resort
+    if result.is_err() {
+        println!("[RUST_LOG_FALLBACK] {}: {}", method_name, message);
+    }
+}
+
+/// Macro for easy info logging
+#[macro_export]
+macro_rules! rust_log_info {
+    ($($arg:tt)*) => {
+        $crate::logger::log_info(&format!($($arg)*))
+    };
+}
+
+/// Macro for easy warning logging
+#[macro_export]
+macro_rules! rust_log_warn {
+    ($($arg:tt)*) => {
+        $crate::logger::log_warn(&format!($($arg)*))
+    };
+}
+
+/// Macro for easy error logging
+#[macro_export]
+macro_rules! rust_log_error {
+    ($($arg:tt)*) => {
+        $crate::logger::log_error(&format!($($arg)*))
+    };
+}
+
+/// Macro for easy debug logging
+#[macro_export]
+macro_rules! rust_log_debug {
+    ($($arg:tt)*) => {
+        $crate::logger::log_debug(&format!($($arg)*))
+    };
+}