native/core/src/lib.rs

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#![allow(incomplete_features)]
#![allow(non_camel_case_types)]
#![allow(clippy::upper_case_acronyms)]
#![allow(clippy::result_large_err)]
// For prost generated struct
#![allow(clippy::derive_partial_eq_without_eq)]
// The clippy throws an error if the reference clone not wrapped into `Arc::clone`
// The lint makes easier for code reader/reviewer separate references clones from more heavyweight ones
#![deny(clippy::clone_on_ref_ptr)]
extern crate core;

#[macro_use]
extern crate datafusion_comet_jni_bridge;

use jni::{
    objects::{JClass, JString},
    EnvUnowned,
};
use log::info;
use log4rs::{
    append::console::{ConsoleAppender, Target},
    config::{load_config_file, Appender, Deserializers, Root},
    encode::pattern::PatternEncoder,
    Config,
};

// Re-export from jvm-bridge crate for internal use
pub use datafusion_comet_jni_bridge::errors;
pub use datafusion_comet_jni_bridge::JAVA_VM;

/// Re-export jvm-bridge items under the `jvm_bridge` name for convenience.
pub mod jvm_bridge {
    pub use datafusion_comet_jni_bridge::*;
}

use errors::{try_unwrap_or_throw, CometError, CometResult};

pub mod alloc_accounting;
pub mod cloud;
pub mod comet_native_udf_bridge;
pub mod execution;
pub mod parquet;
// this module is for non release only. Intended for debugging/profiling purposes
#[cfg(debug_assertions)]
pub mod debug;

// Global allocator selection. `backend` names the allocator the feature set asks for: jemalloc
// where it builds, otherwise mimalloc, otherwise the system allocator. The three cfgs partition
// every feature combination, so exactly one `backend` exists and a combination matching none would
// fail to compile. Whichever it is, it is installed wrapped in the `AccountingAllocator`, which
// counts the bytes Rust code holds for the memory usage log and the `native_allocated` metric.

/// jemalloc, on targets where it builds, unless mimalloc was also requested.
#[cfg(all(
    not(target_env = "msvc"),
    feature = "jemalloc",
    not(feature = "mimalloc")
))]
mod backend {
    pub type Backend = tikv_jemallocator::Jemalloc;
    pub const BACKEND: Backend = tikv_jemallocator::Jemalloc;
    pub const NAME: &str = "jemalloc";
}

/// mimalloc, unless a usable jemalloc was also requested.
#[cfg(all(
    feature = "mimalloc",
    not(all(not(target_env = "msvc"), feature = "jemalloc"))
))]
mod backend {
    pub type Backend = mimalloc::MiMalloc;
    pub const BACKEND: Backend = mimalloc::MiMalloc;
    pub const NAME: &str = "mimalloc";
}

/// The system allocator: the complement of the two cases above, which covers neither feature, a
/// jemalloc request on MSVC, and both features together.
#[cfg(not(any(
    all(
        not(target_env = "msvc"),
        feature = "jemalloc",
        not(feature = "mimalloc")
    ),
    all(
        feature = "mimalloc",
        not(all(not(target_env = "msvc"), feature = "jemalloc"))
    )
)))]
mod backend {
    pub type Backend = std::alloc::System;
    pub const BACKEND: Backend = std::alloc::System;
    pub const NAME: &str = "system";
}

/// The allocator backend this build selected, its name (`"jemalloc"`, `"mimalloc"` or
/// `"system"`), and an instance of it. The selection is decided here and nowhere else, so the
/// `alloc_overhead` benchmark takes the backend from here to measure it with and without the
/// accounting wrapper.
#[doc(hidden)]
pub use backend::{
    Backend as AllocatorBackend, BACKEND as BACKEND_ALLOCATOR, NAME as ALLOCATOR_BACKEND,
};

#[global_allocator]
static GLOBAL: alloc_accounting::AccountingAllocator<backend::Backend> =
    alloc_accounting::AccountingAllocator::new(backend::BACKEND);

#[no_mangle]
pub extern "system" fn Java_org_apache_comet_NativeBase_init(
    e: EnvUnowned,
    _: JClass,
    log_conf_path: JString,
    log_level: JString,
) {
    // Initialize the error handling to capture panic backtraces
    errors::init();

    try_unwrap_or_throw(&e, |env| {
        let path: String = log_conf_path.try_to_string(env)?;

        // The log level only applies without a custom log4rs config file.
        let log_level: String = if path.is_empty() {
            match log_level.try_to_string(env) {
                Ok(level) => level,
                Err(_) => "info".to_string(),
            }
        } else {
            String::new()
        };

        init_logging(&path, &log_level)?;

        // Initialize the global Java VM
        let java_vm = env.get_java_vm()?;
        JAVA_VM.get_or_init(|| java_vm);

        let comet_version = env!("CARGO_PKG_VERSION");
        info!("Comet native library version {comet_version} initialized");
        Ok(())
    })
}

/// Initializes logging from the log4rs config file at `log_conf_path`, or, when it is empty, to the
/// console at `log_level`. Logging can only be initialized once per process; later calls keep the
/// first configuration. Core of `NativeBase.init`, except for capturing the `JavaVM`.
pub fn init_logging(log_conf_path: &str, log_level: &str) -> CometResult<()> {
    let _ = log4rs::init_config(logger_config(log_conf_path, log_level)?)
        .map_err(|err| CometError::Config(err.to_string()));
    Ok(())
}

/// The log4rs config [`init_logging`] installs.
fn logger_config(log_conf_path: &str, log_level: &str) -> CometResult<Config> {
    // empty path means there is no custom log4rs config file provided, so fallback to use
    // the default configuration
    if log_conf_path.is_empty() {
        default_logger_config(log_level)
    } else {
        load_config_file(log_conf_path, Deserializers::default())
            .map_err(|err| CometError::Config(err.to_string()))
    }
}

#[no_mangle]
/// Releases the global Tokio runtime used by Comet native execution.
pub extern "system" fn Java_org_apache_comet_NativeBase_release(_e: EnvUnowned, _class: JClass) {
    execution::jni_api::release_runtime();
}

const LOG_PATTERN: &str = "{d(%y/%m/%d %H:%M:%S)} {l} {f}: {m}{n}";

/// JNI method to check if a specific feature is enabled in the native Rust code.
/// # Arguments
/// * `feature_name` - The name of the feature to check. Supported features:
///   - "jemalloc" - tikv-jemallocator memory allocator
///   - "hdfs-opendal" - HDFS support via OpenDAL
/// # Returns
/// * `1` (true) if the feature is enabled
/// * `0` (false) if the feature is disabled or unknown
#[no_mangle]
pub extern "system" fn Java_org_apache_comet_NativeBase_isFeatureEnabled(
    env: EnvUnowned,
    _: JClass,
    feature_name: JString,
) -> jni::sys::jboolean {
    try_unwrap_or_throw(&env, |env| {
        let feature: String = feature_name.try_to_string(env)?;
        Ok(is_feature_enabled(&feature))
    })
}

/// Whether the native build enables `feature`; unknown features are disabled. Core of
/// `NativeBase.isFeatureEnabled`.
pub fn is_feature_enabled(feature: &str) -> bool {
    match feature {
        "jemalloc" => cfg!(feature = "jemalloc"),
        "hdfs-opendal" => cfg!(feature = "hdfs-opendal"),
        _ => false, // Unknown features return false
    }
}

/// JNI: can object_store build a store AND an object key for this URL?
///
/// Source of truth for the JVM planner's "can the native reader handle this filesystem?" check.
/// `prepare_object_store_with_configs` dispatches non-hdfs/non-s3 schemes to object_store's
/// `parse_url` (driven by `ObjectStoreScheme::parse`), so answering from the same parser lets the
/// planner decline early without hardcoding the supported set. (hdfs/libhdfs are handled JVM-side.)
///
/// `ObjectStoreScheme::parse` validates the PATH as well as the scheme -- it ends in
/// `Path::from_url_path` -- so a recognized scheme carrying a key object_store forbids (e.g. a
/// directory name with a newline) answers false. `CometScanRule` relies on both halves: it probes
/// a synthetic scheme-only URL for the cached scheme gate and the real URL for the path gate.
///
/// Alias schemes (e.g. `blob`) are NOT answered here: this gate has no config. Their opt-in lives
/// in the JVM `CometScanRule` via `fs.comet.s3Compliant.schemes`; natively they are rewritten to
/// `s3://` by `normalize_object_store_url`, the first step of
/// `prepare_object_store_with_configs`.
#[no_mangle]
pub extern "system" fn Java_org_apache_comet_NativeBase_isObjectStoreSchemeSupported(
    env: EnvUnowned,
    _: JClass,
    url: JString,
) -> jni::sys::jboolean {
    try_unwrap_or_throw(&env, |env| {
        let url_str: String = url.try_to_string(env)?;
        Ok(is_object_store_scheme_supported(&url_str))
    })
}

/// Whether object_store can build a store and an object key for `url`. Core of
/// `NativeBase.isObjectStoreSchemeSupported`.
pub fn is_object_store_scheme_supported(url: &str) -> bool {
    url::Url::parse(url)
        .ok()
        .map(|u| object_store::ObjectStoreScheme::parse(&u).is_ok())
        .unwrap_or(false)
}

/// JNI: the version of the IANA timezone database native code uses, such as `2025b`. chrono-tz
/// compiles the database into libcomet, so it can differ from the JVM's `tzdb.dat`, and then local
/// times computed natively can differ from Spark's.
#[no_mangle]
pub extern "system" fn Java_org_apache_comet_NativeBase_getTzdataVersion(
    env: EnvUnowned,
    _: JClass,
) -> jni::sys::jstring {
    try_unwrap_or_throw(&env, |env| Ok(env.new_string(tzdata_version())?.into_raw()))
}

/// The version of the IANA timezone database compiled into libcomet. Core of
/// `NativeBase.getTzdataVersion`.
pub fn tzdata_version() -> &'static str {
    chrono_tz::IANA_TZDB_VERSION
}

// Creates a default log4rs config, which logs to console with log level.
fn default_logger_config(log_level: &str) -> CometResult<Config> {
    let console_append = ConsoleAppender::builder()
        .target(Target::Stderr)
        .encoder(Box::new(PatternEncoder::new(LOG_PATTERN)))
        .build();
    let appender = Appender::builder().build("console", Box::new(console_append));
    let root = Root::builder().appender("console").build(
        log_level
            .parse()
            .map_err(|err| CometError::Config(format!("{err}")))?,
    );
    Config::builder()
        .appender(appender)
        .build(root)
        .map_err(|err| CometError::Config(err.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn feature_enabled() {
        assert_eq!(is_feature_enabled("jemalloc"), cfg!(feature = "jemalloc"));
        assert_eq!(
            is_feature_enabled("hdfs-opendal"),
            cfg!(feature = "hdfs-opendal")
        );
        assert!(!is_feature_enabled("no-such-feature"));
    }

    #[test]
    fn logger_config_from_level_or_file() {
        assert!(logger_config("", "debug").is_ok());
        assert!(matches!(
            logger_config("", "not-a-level"),
            Err(CometError::Config(_))
        ));
        assert!(matches!(
            logger_config("/no/such/log4rs.yaml", "info"),
            Err(CometError::Config(_))
        ));
    }

    #[test]
    fn tzdata_version_is_set() {
        assert!(!tzdata_version().is_empty());
    }

    #[test]
    fn object_store_scheme_supported() {
        assert!(is_object_store_scheme_supported("file:///tmp/data.parquet"));
        assert!(is_object_store_scheme_supported("s3://bucket/key"));
        assert!(!is_object_store_scheme_supported("unknown://bucket/key"));
        assert!(!is_object_store_scheme_supported("not a url"));
    }
}