Feature/rep rec

gradle/run.gradle

@@ -40,6 +40,26 @@ testClusters {
     testDistribution = 'archive'
     if (numZones > 1) numberOfZones = numZones
     if (numNodes > 1) numberOfNodes = numNodes
+    // S3 repository configuration
+    if (findProperty("enableS3")) {
+      plugin(':plugins:repository-s3')
+      if (findProperty("s3Endpoint")) {
+        setting 's3.client.default.endpoint', findProperty("s3Endpoint")
+      }
+      setting 's3.client.default.region', findProperty("s3Region") ?: 'us-east-1'
+      keystore 's3.client.default.access_key', findProperty("s3AccessKey") ?: System.getenv("AWS_ACCESS_KEY_ID") ?: 'test'
+      keystore 's3.client.default.secret_key', findProperty("s3SecretKey") ?: System.getenv("AWS_SECRET_ACCESS_KEY") ?: 'test'
+
+
+      // Remote store configuration
+      setting 'node.attr.remote_store.segment.repository', 'my-s3-repo'
+      setting 'node.attr.remote_store.translog.repository', 'my-s3-repo'
+      setting 'node.attr.remote_store.state.repository', 'my-s3-repo'
+      setting 'cluster.remote_store.state.enabled', 'true'
+      setting 'node.attr.remote_store.repository.my-s3-repo.type', 's3'
+      setting 'node.attr.remote_store.repository.my-s3-repo.settings.bucket', 'local-opensearch-bucket'
+      setting 'node.attr.remote_store.repository.my-s3-repo.settings.base_path', 'raghraaj-local-1230'
+    }
     if (findProperty("installedPlugins")) {
       installedPlugins = Eval.me(installedPlugins)
 

modules/parquet-data-format/src/internalClusterTest/java/com/parquet/parquetdataformat/ParquetSegmentReplicationIT.java

@@ -314,80 +314,80 @@ public void testFormatAwareMetadataReplication() throws Exception {
     /**
      * Tests that replica can recover from remote store with Parquet files.
      */
-//    public void testReplicaRecoveryWithParquetFiles() throws Exception {
-//        internalCluster().startClusterManagerOnlyNode();
-//        internalCluster().startDataOnlyNodes(2);
-//        createReplicationIndex(INDEX_NAME, 1);
-//
-//        // Index documents
-//        for (int i = 0; i < 20; i++) {
-//            client().prepareIndex(INDEX_NAME)
-//                .setId(String.valueOf(i))
-//                .setSource("id", String.valueOf(i), "field", "recovery" + i, "value", (long) i)
-//                .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
-//                .get();
-//        }
-//
-//        String primaryNode = getPrimaryNodeName(INDEX_NAME);
-//        String replicaNode = getReplicaNodeName(INDEX_NAME);
-//
-//        // Wait for initial replication
-//        assertBusy(() -> {
-//            IndexShard primaryShard = getIndexShard(primaryNode, INDEX_NAME);
-//            IndexShard replicaShard = getIndexShard(replicaNode, INDEX_NAME);
-//            assertEquals(
-//                primaryShard.getLatestReplicationCheckpoint().getSegmentInfosVersion(),
-//                replicaShard.getLatestReplicationCheckpoint().getSegmentInfosVersion()
-//            );
-//        }, 30, TimeUnit.SECONDS);
-//
-//        // Stop replica node to simulate failure
-//        internalCluster().restartNode(replicaNode, new InternalTestCluster.RestartCallback() {
-//            @Override
-//            public Settings onNodeStopped(String nodeName) throws Exception {
-//                // Index more documents on primary while replica is down
-//                try {
-//                    for (int i = 20; i < 40; i++) {
-//                        client().prepareIndex(INDEX_NAME)
-//                            .setId(String.valueOf(i))
-//                            .setSource("id", String.valueOf(i), "field", "after_failure" + i, "value", (long) i)
-//                            .get();
-//                    }
-//                    client().admin().indices().prepareRefresh(INDEX_NAME).get();
-//                } catch (Exception e) {
-//                    throw new RuntimeException(e);
-//                }
-//                return super.onNodeStopped(nodeName);
-//            }
-//        });
-//
-//        ensureGreen(INDEX_NAME);
-//
-//        // Verify replica recovered with Parquet files
-//        assertBusy(() -> {
-//            IndexShard primaryShard = getIndexShard(primaryNode, INDEX_NAME);
-//            IndexShard replicaShard = getIndexShard(replicaNode, INDEX_NAME);
-//
-//            // Verify checkpoints match after recovery
-//            assertEquals(
-//                "Replica should catch up after recovery",
-//                primaryShard.getLatestReplicationCheckpoint().getSegmentInfosVersion(),
-//                replicaShard.getLatestReplicationCheckpoint().getSegmentInfosVersion()
-//            );
-//
-//            // Verify replica has Parquet files
-//            RemoteSegmentStoreDirectory replicaRemoteDir = replicaShard.getRemoteDirectory();
-//            Map<String, UploadedSegmentMetadata> replicaSegments =
-//                replicaRemoteDir.getSegmentsUploadedToRemoteStore();
-//
-//            Set<String> formats = replicaSegments.keySet().stream()
-//                .map(file -> new FileMetadata(file).dataFormat())
-//                .collect(Collectors.toSet());
-//
-//            assertTrue("Recovered replica should have Parquet files", formats.contains("parquet"));
-//
-//        }, 60, TimeUnit.SECONDS);
-//    }
+    public void testReplicaRecoveryWithParquetFiles() throws Exception {
+        internalCluster().startClusterManagerOnlyNode();
+        internalCluster().startDataOnlyNodes(2);
+        createReplicationIndex(INDEX_NAME, 1);
+
+        // Index documents
+        for (int i = 0; i < 20; i++) {
+            client().prepareIndex(INDEX_NAME)
+                .setId(String.valueOf(i))
+                .setSource("id", String.valueOf(i), "field", "recovery" + i, "value", (long) i)
+                .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE)
+                .get();
+        }
+
+        String primaryNode = getPrimaryNodeName(INDEX_NAME);
+        String replicaNode = getReplicaNodeName(INDEX_NAME);
+
+        // Wait for initial replication
+        assertBusy(() -> {
+            IndexShard primaryShard = getIndexShard(primaryNode, INDEX_NAME);
+            IndexShard replicaShard = getIndexShard(replicaNode, INDEX_NAME);
+            assertEquals(
+                primaryShard.getLatestReplicationCheckpoint().getSegmentInfosVersion(),
+                replicaShard.getLatestReplicationCheckpoint().getSegmentInfosVersion()
+            );
+        }, 30, TimeUnit.SECONDS);
+
+        // Stop replica node to simulate failure
+        internalCluster().restartNode(replicaNode, new InternalTestCluster.RestartCallback() {
+            @Override
+            public Settings onNodeStopped(String nodeName) throws Exception {
+                // Index more documents on primary while replica is down
+                try {
+                    for (int i = 20; i < 40; i++) {
+                        client().prepareIndex(INDEX_NAME)
+                            .setId(String.valueOf(i))
+                            .setSource("id", String.valueOf(i), "field", "after_failure" + i, "value", (long) i)
+                            .get();
+                    }
+                    client().admin().indices().prepareRefresh(INDEX_NAME).get();
+                } catch (Exception e) {
+                    throw new RuntimeException(e);
+                }
+                return super.onNodeStopped(nodeName);
+            }
+        });
+
+        ensureGreen(INDEX_NAME);
+
+        // Verify replica recovered with Parquet files
+        assertBusy(() -> {
+            IndexShard primaryShard = getIndexShard(primaryNode, INDEX_NAME);
+            IndexShard replicaShard = getIndexShard(replicaNode, INDEX_NAME);
+
+            // Verify checkpoints match after recovery
+            assertEquals(
+                "Replica should catch up after recovery",
+                primaryShard.getLatestReplicationCheckpoint().getSegmentInfosVersion(),
+                replicaShard.getLatestReplicationCheckpoint().getSegmentInfosVersion()
+            );
+
+            // Verify replica has Parquet files
+            RemoteSegmentStoreDirectory replicaRemoteDir = replicaShard.getRemoteDirectory();
+            Map<String, UploadedSegmentMetadata> replicaSegments =
+                replicaRemoteDir.getSegmentsUploadedToRemoteStore();
+
+            Set<String> formats = replicaSegments.keySet().stream()
+                .map(file -> new FileMetadata(file).dataFormat())
+                .collect(Collectors.toSet());
+
+            assertTrue("Recovered replica should have Parquet files", formats.contains("parquet"));
+
+        }, 60, TimeUnit.SECONDS);
+    }
 
     /**
      * Tests that ReplicationCheckpoint contains format-aware metadata.

plugins/engine-datafusion/build.gradle

@@ -176,6 +176,12 @@ test {
   systemProperty 'java.library.path', file('src/main/resources/native').absolutePath
 }
 
+internalClusterTest {
+    // Add same JVM arguments for integration tests
+    jvmArgs += ["--add-opens", "java.base/java.nio=org.apache.arrow.memory.core,ALL-UNNAMED"]
+    systemProperty 'java.library.path', file('src/main/resources/native').absolutePath
+}
+
 yamlRestTest {
     systemProperty 'tests.security.manager', 'false'
     // Disable yamlRestTest since this plugin doesn't have REST API endpoints

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

@@ -8,7 +8,7 @@ use std::num::NonZeroUsize;
  * compatible open source license.
  */
 use std::ptr::addr_of_mut;
-use jni::objects::{JByteArray, JClass, JObject};
+use jni::objects::{JByteArray, JClass, JMap, JObject};
 use jni::objects::JLongArray;
 use jni::sys::{jboolean, jbyteArray, jint, jlong, jstring};
 use jni::{JNIEnv, JavaVM};
@@ -51,7 +51,7 @@ pub mod logger;
 use vectorized_exec_spi::{log_info, log_error, log_debug};
 
 use crate::custom_cache_manager::CustomCacheManager;
-use crate::util::{create_file_meta_from_filenames, parse_string_arr, set_action_listener_error, set_action_listener_error_global, set_action_listener_ok, set_action_listener_ok_global};
+use crate::util::{create_file_meta_from_filenames, parse_string_arr, set_action_listener_error, set_action_listener_error_global, set_action_listener_ok, set_action_listener_ok_global, set_action_listener_ok_global_with_map};
 use datafusion::execution::memory_pool::{GreedyMemoryPool, TrackConsumersPool};
 
 use crate::statistics_cache::CustomStatisticsCache;
@@ -482,6 +482,29 @@ impl CustomFileMeta {
     }
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct FileStats {
+    /// Total file size in bytes
+    pub size: u64,
+
+    /// Total number of rows in the file
+    pub num_rows: i64,
+}
+
+impl FileStats {
+    pub fn new(size: u64, num_rows: i64) -> Self {
+        Self { size, num_rows }
+    }
+
+    pub fn size(&self) -> u64 {
+        self.size
+    }
+
+    pub fn num_rows(&self) -> i64 {
+        self.num_rows
+    }
+}
+
 #[no_mangle]
 pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_executeQueryPhaseAsync(
     mut env: JNIEnv,
@@ -574,6 +597,55 @@ pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_executeQu
     });
 }
 
+#[no_mangle]
+pub extern "system" fn Java_org_opensearch_datafusion_jni_NativeBridge_fetchSegmentStats(
+    mut env: JNIEnv,
+    _class: JClass,
+    shard_view_ptr: jlong,
+    listener: JObject,
+) {
+    let manager = match TOKIO_RUNTIME_MANAGER.get() {
+        Some(m) => m,
+        None => {
+            log_info!("Runtime manager not initialized");
+            set_action_listener_error(&mut env, listener,
+                                    &DataFusionError::Execution("Runtime manager not initialized".to_string()));
+            return;
+        }
+    };
+
+    // Convert listener to GlobalRef (thread-safe)
+    let listener_ref = match env.new_global_ref(&listener) {
+        Ok(r) => r,
+        Err(e) => {
+            log_error!("Failed to create global ref: {}", e);
+            set_action_listener_error(&mut env, listener,
+                                    &DataFusionError::Execution(format!("Failed to create global ref: {}", e)));
+            return;
+        }
+    };
+    let io_runtime = manager.io_runtime.clone();
+
+    let shard_view = unsafe { &*(shard_view_ptr as *const ShardView) };
+    let files_meta = shard_view.files_metadata();
+
+    io_runtime.block_on(async move {
+        let file_stats = util::fetch_segment_statistics(files_meta).await;
+        match file_stats {
+            Ok(map) => {
+                with_jni_env(|env| {
+                    set_action_listener_ok_global_with_map(env, &listener_ref, &map);
+                });
+            }
+            Err(e) => {
+                with_jni_env(|env| {
+                    log_error!("Collecting file stats failed: {}", e);
+                    set_action_listener_error_global(env, &listener_ref, &e);
+                });
+            }
+        }
+    });
+}
 
 
 #[no_mangle]

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

@@ -8,6 +8,7 @@
 
 use std::sync::Arc;
 use std::collections::{BTreeSet, HashMap, HashSet};
+use datafusion::common::stats::Precision;
 use jni::sys::jlong;
 use datafusion::{
     common::DataFusionError,
@@ -51,7 +52,7 @@ use crate::listing_table::{ListingOptions, ListingTable, ListingTableConfig};
 use crate::partial_agg_optimizer::PartialAggregationOptimizer;
 use crate::executor::DedicatedExecutor;
 use crate::cross_rt_stream::CrossRtStream;
-use crate::CustomFileMeta;
+use crate::{CustomFileMeta, FileStats};
 use crate::DataFusionRuntime;
 use crate::project_row_id_analyzer::ProjectRowIdAnalyzer;
 use crate::absolute_row_id_optimizer::{AbsoluteRowIdOptimizer, ROW_BASE_FIELD_NAME, ROW_ID_FIELD_NAME};