Initial commit of plugin engine-datafusion

Signed-off-by: Marc Handalian <marc.handalian@gmail.com>

Author: mch2

plugins/engine-datafusion/.gitignore

@@ -0,0 +1,38 @@
+# Gradle
+.gradle/
+build/
+
+# Java
+*.class
+*.jar
+*.war
+*.ear
+hs_err_pid*
+
+# IDE
+.idea/
+*.iml
+*.ipr
+*.iws
+.vscode/
+.settings/
+.project
+.classpath
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Rust
+jni/target/
+jni/Cargo.lock
+
+# Native libraries
+src/main/resources/native/
+
+# Logs
+*.log
+
+# Temporary files
+*.tmp
+*.temp

plugins/engine-datafusion/build.gradle

@@ -0,0 +1,111 @@
+/*
+ * 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.
+ */
+
+apply plugin: 'java'
+apply plugin: 'idea'
+apply plugin: 'opensearch.internal-cluster-test'
+apply plugin: 'opensearch.yaml-rest-test'
+apply plugin: 'opensearch.pluginzip'
+
+def pluginName = 'engine-datafusion'
+def pluginDescription = 'OpenSearch plugin providing access to DataFusion via JNI'
+def projectPath = 'org.opensearch'
+def pathToPlugin = 'datafusion.DataFusionPlugin'
+def pluginClassName = 'DataFusionPlugin'
+
+opensearchplugin {
+    name = pluginName
+    description = pluginDescription
+    classname = "${projectPath}.${pathToPlugin}"
+    licenseFile = rootProject.file('LICENSE.txt')
+    noticeFile = rootProject.file('NOTICE.txt')
+}
+
+dependencies {
+    implementation "org.apache.logging.log4j:log4j-api:${versions.log4j}"
+    implementation "org.apache.logging.log4j:log4j-core:${versions.log4j}"
+    testImplementation "junit:junit:${versions.junit}"
+    testImplementation "org.hamcrest:hamcrest:${versions.hamcrest}"
+    testImplementation "org.mockito:mockito-core:${versions.mockito}"
+}
+
+// Task to build the Rust JNI library
+task buildRustLibrary(type: Exec) {
+    description = 'Build the Rust JNI library using Cargo'
+    group = 'build'
+
+    workingDir file('jni')
+
+    // Determine the target directory and library name based on OS
+    def osName = System.getProperty('os.name').toLowerCase()
+    def libPrefix = osName.contains('windows') ? '' : 'lib'
+    def libExtension = osName.contains('windows') ? '.dll' : (osName.contains('mac') ? '.dylib' : '.so')
+
+    // Use debug build for development, release for production
+    def buildType = project.hasProperty('rustRelease') ? 'release' : 'debug'
+    def targetDir = "target/${buildType}"
+
+    def cargoArgs = ['cargo', 'build']
+    if (buildType == 'release') {
+        cargoArgs.add('--release')
+    }
+
+    if (osName.contains('windows')) {
+        commandLine cargoArgs
+    } else {
+        commandLine cargoArgs
+    }
+
+    // Set environment variables for cross-compilation if needed
+    environment 'CARGO_TARGET_DIR', file('jni/target').absolutePath
+
+    inputs.files fileTree('jni/src')
+    inputs.file 'jni/Cargo.toml'
+    outputs.files file("jni/${targetDir}/${libPrefix}opensearch_datafusion_jni${libExtension}")
+    System.out.println("Building Rust library in ${buildType} mode");
+}
+
+// Task to copy the native library to resources
+task copyNativeLibrary(type: Copy, dependsOn: buildRustLibrary) {
+    description = 'Copy the native library to Java resources'
+    group = 'build'
+
+    def osName = System.getProperty('os.name').toLowerCase()
+    def libPrefix = osName.contains('windows') ? '' : 'lib'
+    def libExtension = osName.contains('windows') ? '.dll' : (osName.contains('mac') ? '.dylib' : '.so')
+    def buildType = project.hasProperty('rustRelease') ? 'release' : 'debug'
+
+    from file("jni/target/${buildType}/${libPrefix}opensearch_datafusion_jni${libExtension}")
+    into file('src/main/resources/native')
+
+    // Rename to a standard name for Java to load
+    rename { filename ->
+        "libopensearch_datafusion_jni${libExtension}"
+    }
+}
+
+// Ensure native library is built before Java compilation
+compileJava.dependsOn copyNativeLibrary
+
+// Ensure processResources depends on copyNativeLibrary
+processResources.dependsOn copyNativeLibrary
+
+// Clean task should also clean Rust artifacts
+clean {
+    delete file('jni/target')
+    delete file('src/main/resources/native')
+}
+
+test {
+    // Set system property to help tests find the native library
+    systemProperty 'java.library.path', file('src/main/resources/native').absolutePath
+}
+
+yamlRestTest {
+    systemProperty 'tests.security.manager', 'false'
+}

plugins/engine-datafusion/jni/Cargo.toml

@@ -0,0 +1,43 @@
+[package]
+name = "opensearch-datafusion-jni"
+version = "0.1.0"
+edition = "2021"
+description = "JNI bindings for DataFusion integration with OpenSearch"
+license = "Apache-2.0"
+
+[lib]
+name = "opensearch_datafusion_jni"
+crate-type = ["cdylib"]
+
+[dependencies]
+datafusion = "49.0.0"
+arrow = "55.2"
+arrow-json = "55.2"
+
+# JNI dependencies
+jni = "0.21"
+
+# Async runtime
+tokio = { version = "1.0", features = ["rt", "rt-multi-thread", "macros"] }
+
+# Serialization
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+
+# Error handling
+anyhow = "1.0"
+thiserror = "1.0"
+
+# Logging
+log = "0.4"
+
+[profile.release]
+lto = true
+codegen-units = 1
+panic = "abort"
+
+[profile.dev]
+opt-level = 1          # Some optimization for reasonable performance
+lto = false           # Disable LTO for faster builds
+codegen-units = 16    # More parallel compilation
+incremental = true    # Enable incremental compilation

plugins/engine-datafusion/jni/src/lib.rs

@@ -0,0 +1,47 @@
+/*
+ * 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.
+ */
+
+use jni::objects::JClass;
+use jni::sys::{jlong, jstring};
+use jni::JNIEnv;
+
+use datafusion::execution::context::SessionContext;
+
+use datafusion::DATAFUSION_VERSION;
+use datafusion::prelude::SessionConfig;
+
+/// Create a new DataFusion session context
+#[no_mangle]
+pub extern "system" fn Java_org_opensearch_datafusion_DataFusionJNI_createContext(
+    _env: JNIEnv,
+    _class: JClass,
+) -> jlong {
+    let config = SessionConfig::new().with_repartition_aggregations(true);
+    let context = SessionContext::new_with_config(config);
+    let ctx = Box::into_raw(Box::new(context)) as jlong;
+    ctx
+}
+
+/// Close and cleanup a DataFusion context
+#[no_mangle]
+pub extern "system" fn Java_org_opensearch_datafusion_DataFusionJNI_closeContext(
+    _env: JNIEnv,
+    _class: JClass,
+    context_id: jlong,
+) {
+    let _ = unsafe { Box::from_raw(context_id as *mut SessionContext) };
+}
+
+/// Get version information
+#[no_mangle]
+pub extern "system" fn Java_org_opensearch_datafusion_DataFusionJNI_getVersion(
+    env: JNIEnv,
+    _class: JClass,
+) -> jstring {
+    env.new_string(DATAFUSION_VERSION).expect("Couldn't create Java string").as_raw()
+}

plugins/engine-datafusion/src/main/java/org/opensearch/datafusion/DataFusionJNI.java

@@ -0,0 +1,77 @@
+/*
+ * 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 org.opensearch.datafusion;
+
+import java.io.IOException;
+import java.io.InputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.StandardCopyOption;
+
+/**
+ * JNI wrapper for DataFusion operations
+ */
+public class DataFusionJNI {
+
+    private static boolean libraryLoaded = false;
+
+    static {
+        loadNativeLibrary();
+    }
+
+    /**
+     * Load the native library from resources
+     */
+    private static synchronized void loadNativeLibrary() {
+        if (libraryLoaded) {
+            return;
+        }
+
+        try {
+            String osName = System.getProperty("os.name").toLowerCase();
+            String libExtension;
+            String libName;
+
+            if (osName.contains("windows")) {
+                libExtension = ".dll";
+                libName = "libopensearch_datafusion_jni.dll";
+            } else if (osName.contains("mac")) {
+                libExtension = ".dylib";
+                libName = "libopensearch_datafusion_jni.dylib";
+            } else {
+                libExtension = ".so";
+                libName = "libopensearch_datafusion_jni.so";
+            }
+
+            // Try to load from resources first
+            InputStream libStream = DataFusionJNI.class.getResourceAsStream("/native/" + libName);
+            if (libStream != null) {
+                // Extract to temporary file and load
+                Path tempLib = Files.createTempFile("libopensearch_datafusion_jni", libExtension);
+                Files.copy(libStream, tempLib, StandardCopyOption.REPLACE_EXISTING);
+                tempLib.toFile().deleteOnExit();
+                System.load(tempLib.toAbsolutePath().toString());
+                libStream.close();
+            } else {
+                // Fallback to system library path
+                System.loadLibrary("opensearch_datafusion_jni");
+            }
+
+            libraryLoaded = true;
+        } catch (IOException | UnsatisfiedLinkError e) {
+            throw new RuntimeException("Failed to load DataFusion JNI library", e);
+        }
+    }
+
+    /**
+     * Get version information
+     * @return JSON string with version information
+     */
+    public static native String getVersion();
+}