diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java index 59f9b1f899f92..fe6126d3091fa 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java @@ -8,6 +8,10 @@ package org.opensearch.analytics.spi; +import org.apache.arrow.memory.BufferAllocator; +import org.opensearch.analytics.backend.EngineResultStream; +import org.opensearch.index.engine.exec.IndexReaderProvider; + import java.util.List; /** @@ -114,6 +118,25 @@ default void configureFilterDelegation(FilterDelegationHandle handle, BackendExe throw new UnsupportedOperationException("configureFilterDelegation not implemented for [" + name() + "]"); } + /** + * QTF fetch phase: reads specific rows by global row ID. + * Row IDs are passed as a BigIntVector for zero-copy transfer to native. + * + * @param reader the index reader for the target shard + * @param rowIdVector Arrow BigIntVector containing global row IDs + * @param columns column names to read + * @param allocator Arrow buffer allocator for result import + * @return a result stream containing the requested rows + */ + default EngineResultStream fetchByRowIds( + IndexReaderProvider.Reader reader, + org.apache.arrow.vector.BigIntVector rowIdVector, + String[] columns, + BufferAllocator allocator + ) { + throw new UnsupportedOperationException("fetchByRowIds not implemented for [" + name() + "]"); + } + /** * Install a thread tracker for attribution of delegation callbacks executing on foreign threads. * Called after {@link #configureFilterDelegation}. Pass {@code null} to clear. diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml index b4f30bf3e9dc6..7ee52f920b2a2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml +++ b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml @@ -89,3 +89,7 @@ rand = "=0.8.6" [[bench]] name = "query_bench" harness = false + +[[bench]] +name = "row_id_bench" +harness = false diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/benches/row_id_bench.rs b/sandbox/plugins/analytics-backend-datafusion/rust/benches/row_id_bench.rs new file mode 100644 index 0000000000000..882425a074d3e --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/benches/row_id_bench.rs @@ -0,0 +1,409 @@ +//! Benchmark: ClickBench queries — QTF row_id+sort_keys vs full data fetch. +//! +//! Queries (from ClickBench q25–q27): +//! - q25: WHERE SearchPhrase != '' ORDER BY EventTime LIMIT 10 +//! - q26: WHERE SearchPhrase != '' ORDER BY SearchPhrase LIMIT 10 +//! - q27: WHERE SearchPhrase != '' ORDER BY EventTime, SearchPhrase LIMIT 10 +//! +//! Each query is run in three modes: +//! - data_fetch: full query via ListingTable (filter + sort + limit + all data) +//! - row_id_emit: indexed path with emit_row_ids=true, projects ___row_id + sort keys +//! - indexed_no_emit: indexed path without row ID emission (just filter + return data) +//! +//! Usage: +//! ROW_ID_BENCH_FILE=/path/to/hits_1.parquet cargo bench --bench row_id_bench + +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::listing::ListingTableUrl; +use datafusion::execution::memory_pool::GreedyMemoryPool; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use futures::TryStreamExt; +use object_store::local::LocalFileSystem; +use object_store::{ObjectStore, ObjectStoreExt}; +use opensearch_datafusion::api::DataFusionRuntime; +use opensearch_datafusion::datafusion_query_config::DatafusionQueryConfig; +use opensearch_datafusion::memory::DynamicLimitPool; +use opensearch_datafusion::query_executor; +use opensearch_datafusion::runtime_manager::RuntimeManager; +use std::sync::Arc; + +fn bench_file() -> String { + let path = std::env::var("ROW_ID_BENCH_FILE").unwrap_or_else(|_| { + panic!( + "ROW_ID_BENCH_FILE not set.\n\ + Pass path to a ClickBench hits parquet file.\n\n\ + Example:\n \ + ROW_ID_BENCH_FILE=/path/to/hits_1.parquet cargo bench --bench row_id_bench" + ); + }); + assert!( + std::path::Path::new(&path).exists(), + "Benchmark file not found: {}", + path + ); + path +} + +fn setup() -> (RuntimeManager, DataFusionRuntime) { + let mgr = RuntimeManager::new(4); + let runtime_env = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(2 * 1024 * 1024 * 1024))) + .build() + .unwrap(); + let (_, handle) = DynamicLimitPool::new(2 * 1024 * 1024 * 1024); + let df_runtime = DataFusionRuntime { + runtime_env, + custom_cache_manager: None, + dynamic_limit_handle: handle, + }; + (mgr, df_runtime) +} + +fn get_metas(mgr: &RuntimeManager, file: &str) -> Arc> { + let store = Arc::new(LocalFileSystem::new()); + let path = object_store::path::Path::from(file); + let meta = mgr.io_runtime.block_on(store.head(&path)).unwrap(); + Arc::new(vec![meta]) +} + +fn get_substrait(mgr: &RuntimeManager, file_path: &str, sql: &str) -> Vec { + use datafusion::datasource::file_format::parquet::ParquetFormat; + use datafusion::datasource::listing::{ListingOptions, ListingTable, ListingTableConfig}; + use datafusion_substrait::logical_plan::producer::to_substrait_plan; + use prost::Message; + + mgr.io_runtime.block_on(async { + let ctx = datafusion::prelude::SessionContext::new(); + let url = ListingTableUrl::parse(file_path).unwrap(); + let opts = ListingOptions::new(Arc::new(ParquetFormat::new())) + .with_file_extension(".parquet") + .with_collect_stat(true); + let schema = opts.infer_schema(&ctx.state(), &url).await.unwrap(); + let cfg = ListingTableConfig::new(url) + .with_listing_options(opts) + .with_schema(schema); + ctx.register_table("t", Arc::new(ListingTable::try_new(cfg).unwrap())) + .unwrap(); + let plan = ctx.sql(sql).await.unwrap().logical_plan().clone(); + let sub = to_substrait_plan(&plan, &ctx.state()).unwrap(); + let mut buf = Vec::new(); + sub.encode(&mut buf).unwrap(); + buf + }) +} + +/// Augment a parquet schema with a virtual `___row_id` column at the end. +fn schema_with_row_id(base: &Schema) -> Arc { + let mut fields: Vec = base.fields().iter().map(|f| f.as_ref().clone()).collect(); + fields.push(Field::new("___row_id", DataType::Int64, false)); + Arc::new(Schema::new(fields)) +} + +fn bench_clickbench(c: &mut Criterion) { + use datafusion::common::ScalarValue; + use datafusion::parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; + use opensearch_datafusion::indexed_table::bool_tree::BoolNode; + use opensearch_datafusion::indexed_table::eval::bitmap_tree::{ + BitmapTreeEvaluator, CollectorLeafBitmaps, + }; + use opensearch_datafusion::indexed_table::eval::{RowGroupBitsetSource, TreeBitsetSource}; + use opensearch_datafusion::indexed_table::page_pruner::PagePruner; + use opensearch_datafusion::indexed_table::stream::RowGroupInfo; + use opensearch_datafusion::indexed_table::table_provider::{ + IndexedTableConfig, IndexedTableProvider, SegmentFileInfo, + }; + + let (mgr, df_runtime) = setup(); + let file = bench_file(); + let metas = get_metas(&mgr, &file); + let url = ListingTableUrl::parse(&file).unwrap(); + + // Load parquet metadata once + let path = std::path::Path::new(&file); + let size = std::fs::metadata(path).unwrap().len(); + let fh = std::fs::File::open(path).unwrap(); + let meta = + ArrowReaderMetadata::load(&fh, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let parquet_schema = meta.schema().clone(); + let parquet_meta = meta.metadata().clone(); + let mut rgs = Vec::new(); + let mut offset = 0i64; + for i in 0..parquet_meta.num_row_groups() { + let n = parquet_meta.row_group(i).num_rows(); + rgs.push(RowGroupInfo { + index: i, + first_row: offset, + num_rows: n, + }); + offset += n; + } + let total_rows = offset; + let object_path = object_store::path::Path::from(path.to_string_lossy().as_ref()); + let segment = SegmentFileInfo { + segment_ord: 0, + max_doc: total_rows, + object_path, + parquet_size: size, + row_groups: rgs, + metadata: Arc::clone(&parquet_meta), + global_base: 0, + }; + + // Schema with ___row_id appended (virtual column) + let schema_with_rid = schema_with_row_id(&parquet_schema); + let search_phrase_idx = parquet_schema.index_of("SearchPhrase").unwrap(); + + let mut group = c.benchmark_group("clickbench_qtf"); + group.sample_size(10); + group.warm_up_time(std::time::Duration::from_secs(3)); + group.measurement_time(std::time::Duration::from_secs(10)); + + struct QueryDef { + name: &'static str, + sql_data: &'static str, + sql_rowid: &'static str, + } + + let queries = vec![ + QueryDef { + name: "q25-sort-eventtime", + sql_data: "SELECT \"SearchPhrase\" FROM t WHERE \"SearchPhrase\" != '' ORDER BY \"EventTime\" LIMIT 10", + sql_rowid: "SELECT \"___row_id\", \"EventTime\", \"SearchPhrase\" FROM t WHERE \"SearchPhrase\" != ''", + }, + QueryDef { + name: "q26-sort-searchphrase", + sql_data: "SELECT \"SearchPhrase\" FROM t WHERE \"SearchPhrase\" != '' ORDER BY \"SearchPhrase\" LIMIT 10", + sql_rowid: "SELECT \"___row_id\", \"SearchPhrase\" FROM t WHERE \"SearchPhrase\" != ''", + }, + QueryDef { + name: "q27-sort-multi", + sql_data: "SELECT \"SearchPhrase\" FROM t WHERE \"SearchPhrase\" != '' ORDER BY \"EventTime\", \"SearchPhrase\" LIMIT 10", + sql_rowid: "SELECT \"___row_id\", \"EventTime\", \"SearchPhrase\" FROM t WHERE \"SearchPhrase\" != ''", + }, + ]; + + for q in &queries { + // === Mode 1: data_fetch (full query via ListingTable) === + let plan_data = get_substrait(&mgr, &file, q.sql_data); + let id_data = BenchmarkId::new(format!("{}/data_fetch", q.name), ""); + group.bench_with_input(id_data, &plan_data, |b, plan| { + let df_rt = &df_runtime; + b.to_async(mgr.io_runtime.as_ref()).iter(|| { + let url = url.clone(); + let metas = metas.clone(); + let plan = plan.clone(); + let exec = mgr.cpu_executor(); + async move { + let mut config = DatafusionQueryConfig::test_default(); + config.target_partitions = 4; + let ptr = query_executor::execute_query( + url, metas, "t".into(), plan, df_rt, exec, None, &config, + ) + .await + .unwrap(); + let mut stream = unsafe { + Box::from_raw( + ptr as *mut datafusion::physical_plan::stream::RecordBatchStreamAdapter< + opensearch_datafusion::cross_rt_stream::CrossRtStream, + >, + ) + }; + let mut rows = 0u64; + while let Some(batch) = stream.try_next().await.unwrap() { + rows += batch.num_rows() as u64; + } + rows + } + }); + }); + + // === Mode 2: row_id_emit (indexed path, projects ___row_id + sort keys) === + { + let segment = segment.clone(); + let schema = schema_with_rid.clone(); + let id_rowid = BenchmarkId::new(format!("{}/row_id_emit", q.name), ""); + let sql = q.sql_rowid; + group.bench_function(id_rowid, |b| { + let segment = segment.clone(); + let schema = schema.clone(); + b.to_async(mgr.io_runtime.as_ref()).iter(|| { + let segment = segment.clone(); + let schema = schema.clone(); + async move { + let col_expr: Arc = Arc::new( + datafusion::physical_expr::expressions::Column::new( + "SearchPhrase", + search_phrase_idx, + ), + ); + let lit_expr: Arc = Arc::new( + datafusion::physical_expr::expressions::Literal::new( + ScalarValue::Binary(Some(vec![])), + ), + ); + let pred = BoolNode::Predicate(Arc::new( + datafusion::physical_expr::expressions::BinaryExpr::new( + col_expr, + datafusion::logical_expr::Operator::NotEq, + lit_expr, + ), + )); + let tree = Arc::new(BoolNode::And(vec![pred]).push_not_down()); + + let factory: opensearch_datafusion::indexed_table::table_provider::EvaluatorFactory = { + let tree = Arc::clone(&tree); + let schema = schema.clone(); + Arc::new(move |seg, _chunk, _sm| { + let resolved = tree.resolve(&[])?; + let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&seg.metadata))); + let eval: Arc = Arc::new(TreeBitsetSource { + tree: Arc::new(resolved), + evaluator: Arc::new(BitmapTreeEvaluator), + leaves: Arc::new(CollectorLeafBitmaps { + ffm_collector_calls: _sm.ffm_collector_calls.clone(), + }), + page_pruner: pruner, + cost_predicate: 1, + cost_collector: 10, + max_collector_parallelism: 1, + pruning_predicates: Arc::new(std::collections::HashMap::new()), + page_prune_metrics: Some( + opensearch_datafusion::indexed_table::page_pruner::PagePruneMetrics::from_stream_metrics(_sm), + ), + collector_strategy: opensearch_datafusion::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + }); + Ok(eval) + }) + }; + + let store: Arc = Arc::new(LocalFileSystem::new()); + let store_url = datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(); + let provider = Arc::new(IndexedTableProvider::new(IndexedTableConfig { + schema, + segments: vec![segment], + store, + store_url, + evaluator_factory: factory, + pushdown_predicate: None, + query_config: Arc::new({ + let mut qc = DatafusionQueryConfig::test_default(); + qc.target_partitions = 4; + qc + }), + predicate_columns: vec![search_phrase_idx], + emit_row_ids: true, + })); + + let ctx = datafusion::prelude::SessionContext::new(); + ctx.register_table("t", provider).unwrap(); + let df = ctx.sql(sql).await.unwrap(); + let mut stream = df.execute_stream().await.unwrap(); + let mut rows = 0u64; + while let Some(batch) = stream.try_next().await.unwrap() { + rows += batch.num_rows() as u64; + } + rows + } + }); + }); + } + + // === Mode 3: indexed_no_emit (same filter, no row ID, returns all data) === + { + let segment = segment.clone(); + let schema = parquet_schema.clone(); + let id_no_emit = BenchmarkId::new(format!("{}/indexed_no_emit", q.name), ""); + group.bench_function(id_no_emit, |b| { + let segment = segment.clone(); + let schema = schema.clone(); + b.to_async(mgr.io_runtime.as_ref()).iter(|| { + let segment = segment.clone(); + let schema = schema.clone(); + async move { + let col_expr: Arc = Arc::new( + datafusion::physical_expr::expressions::Column::new( + "SearchPhrase", + search_phrase_idx, + ), + ); + let lit_expr: Arc = Arc::new( + datafusion::physical_expr::expressions::Literal::new( + ScalarValue::Binary(Some(vec![])), + ), + ); + let pred = BoolNode::Predicate(Arc::new( + datafusion::physical_expr::expressions::BinaryExpr::new( + col_expr, + datafusion::logical_expr::Operator::NotEq, + lit_expr, + ), + )); + let tree = Arc::new(BoolNode::And(vec![pred]).push_not_down()); + + let factory: opensearch_datafusion::indexed_table::table_provider::EvaluatorFactory = { + let tree = Arc::clone(&tree); + let schema = schema.clone(); + Arc::new(move |seg, _chunk, _sm| { + let resolved = tree.resolve(&[])?; + let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&seg.metadata))); + let eval: Arc = Arc::new(TreeBitsetSource { + tree: Arc::new(resolved), + evaluator: Arc::new(BitmapTreeEvaluator), + leaves: Arc::new(CollectorLeafBitmaps { + ffm_collector_calls: _sm.ffm_collector_calls.clone(), + }), + page_pruner: pruner, + cost_predicate: 1, + cost_collector: 10, + max_collector_parallelism: 1, + pruning_predicates: Arc::new(std::collections::HashMap::new()), + page_prune_metrics: Some( + opensearch_datafusion::indexed_table::page_pruner::PagePruneMetrics::from_stream_metrics(_sm), + ), + collector_strategy: opensearch_datafusion::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + }); + Ok(eval) + }) + }; + + let store: Arc = Arc::new(LocalFileSystem::new()); + let store_url = datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(); + let provider = Arc::new(IndexedTableProvider::new(IndexedTableConfig { + schema: Arc::new(schema.as_ref().clone()), + segments: vec![segment], + store, + store_url, + evaluator_factory: factory, + pushdown_predicate: None, + query_config: Arc::new({ + let mut qc = DatafusionQueryConfig::test_default(); + qc.target_partitions = 4; + qc + }), + predicate_columns: vec![search_phrase_idx], + emit_row_ids: false, + })); + + let ctx = datafusion::prelude::SessionContext::new(); + ctx.register_table("t", provider).unwrap(); + let df = ctx.sql("SELECT * FROM t").await.unwrap(); + let mut stream = df.execute_stream().await.unwrap(); + let mut rows = 0u64; + while let Some(batch) = stream.try_next().await.unwrap() { + rows += batch.num_rows() as u64; + } + rows + } + }); + }); + } + } + + group.finish(); + mgr.cpu_executor.shutdown(); + std::mem::forget(mgr); +} + +criterion_group!(benches, bench_clickbench); +criterion_main!(benches); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs index da1e73a67ff84..9731b78d25c04 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs @@ -31,6 +31,7 @@ //! - `stream_get_schema`, `stream_close` must NOT be called //! concurrently on the same stream pointer. +use std::collections::HashMap; use std::num::NonZeroUsize; use std::path::PathBuf; use std::sync::Arc; @@ -41,16 +42,19 @@ use arrow_array::{Array, StructArray}; use arrow_schema::ffi::FFI_ArrowSchema; use datafusion::common::DataFusionError; use datafusion::datasource::listing::ListingTableUrl; +use datafusion::datasource::physical_plan::parquet::{ParquetAccessPlan, RowGroupAccess}; use datafusion::execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; use datafusion::execution::memory_pool::TrackConsumersPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::execution::cache::cache_manager::CacheManagerConfig; use datafusion::execution::RecordBatchStream; -use datafusion::execution::{SessionState, SessionStateBuilder}; +use datafusion::execution::SessionStateBuilder; +use datafusion::physical_plan::execute_stream; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; -use datafusion::prelude::SessionConfig; +use datafusion::prelude::{SessionConfig, SessionContext}; use futures::TryStreamExt; use object_store::{ObjectStore, ObjectStoreExt}; +use roaring::RoaringBitmap; use crate::cancellation; use crate::cross_rt_stream::CrossRtStream; @@ -60,6 +64,7 @@ use crate::memory::{DynamicLimitHandle, DynamicLimitPool}; use crate::partition_stream::PartitionStreamSender; use crate::query_tracker::{self, QueryTrackingContext}; use crate::runtime_manager::RuntimeManager; +use crate::shard_table_provider::{ShardTableConfig, ShardTableProvider}; /// Bundles a stream with its query tracking context so that dropping the /// handle automatically marks the query completed in the registry. @@ -130,7 +135,97 @@ pub async fn create_object_metas( pub struct DataFusionRuntime { pub runtime_env: datafusion::execution::runtime_env::RuntimeEnv, pub custom_cache_manager: Option, - pub(crate) dynamic_limit_handle: DynamicLimitHandle, + pub dynamic_limit_handle: DynamicLimitHandle, +} + +/// Per-file metadata passed from Java at shard view creation time. +/// Enables `row_base` computation without re-reading parquet footers. +#[derive(Debug, Clone)] +pub struct FileRowMetadata { + /// Row counts per row group in this file. + pub row_group_row_counts: Vec, +} + +/// Per-file info used by `ShardTableProvider` to inject `row_base` as a +/// partition column and to resolve global row IDs back to file positions. +#[derive(Debug, Clone)] +pub struct ShardFileInfo { + pub object_meta: object_store::ObjectMeta, + /// Cumulative row count from all preceding files. + pub row_base: i64, + /// Total rows in this file. + pub num_rows: u64, + /// Per-row-group row counts. + pub row_group_row_counts: Vec, + /// Optional access plan for targeted row retrieval (QTF fetch phase). + /// When set, ShardTableProvider attaches it to the PartitionedFile so + /// DataSourceExec skips row groups and applies RowSelection. + pub access_plan: Option, +} + +/// FFM wire format for per-file metadata. +/// Must stay in lockstep with the Java `MemoryLayout`. +#[repr(C)] +#[derive(Debug, Copy, Clone)] +pub struct WireFileMetadata { + /// Number of row groups in this file. + pub num_row_groups: i32, + /// Pointer to array of i64 row counts (one per row group). + pub row_group_row_counts_ptr: i64, +} + +/// Decode an array of `WireFileMetadata` from an FFM pointer. +/// +/// # Safety +/// `ptr` must be 0 or a valid pointer to `count` consecutive `WireFileMetadata` structs. +/// Each `row_group_row_counts_ptr` must point to `num_row_groups` consecutive i64 values. +pub unsafe fn decode_file_metadata(ptr: i64, count: usize) -> Option> { + if ptr == 0 || count == 0 { + return None; + } + let wire_slice = std::slice::from_raw_parts(ptr as *const WireFileMetadata, count); + let mut result = Vec::with_capacity(count); + for wire in wire_slice { + let num_rgs = wire.num_row_groups as usize; + let rg_counts = if wire.row_group_row_counts_ptr == 0 || num_rgs == 0 { + Vec::new() + } else { + let counts_ptr = wire.row_group_row_counts_ptr as *const i64; + std::slice::from_raw_parts(counts_ptr, num_rgs) + .iter() + .map(|&c| c as u64) + .collect() + }; + result.push(FileRowMetadata { + row_group_row_counts: rg_counts, + }); + } + Some(result) +} + +/// Build `ShardFileInfo` from object metas and file metadata. +/// Computes `row_base` as the cumulative prefix sum of file row counts. +pub fn build_shard_files( + object_metas: &[object_store::ObjectMeta], + file_metadata: &[FileRowMetadata], +) -> Vec { + let mut row_base: i64 = 0; + object_metas + .iter() + .zip(file_metadata.iter()) + .map(|(meta, fm)| { + let num_rows: u64 = fm.row_group_row_counts.iter().sum(); + let info = ShardFileInfo { + object_meta: meta.clone(), + row_base, + num_rows, + row_group_row_counts: fm.row_group_row_counts.clone(), + access_plan: None, + }; + row_base += num_rows as i64; + info + }) + .collect() } impl DataFusionRuntime { @@ -148,6 +243,9 @@ impl DataFusionRuntime { pub struct ShardView { pub table_path: ListingTableUrl, pub object_metas: Arc>, + /// Per-file row group counts, passed from Java at shard view creation. + /// When present, enables ShardTableProvider construction with row_base. + pub file_metadata: Option>, /// Per-shard object store. When a native store is provided (store_ptr > 0), /// this routes reads through TieredObjectStore (local + remote). /// When no store is provided, uses default LocalFileSystem. @@ -285,6 +383,7 @@ pub fn create_reader( let shard_view = ShardView { table_path: table_url, object_metas: Arc::new(object_metas), + file_metadata: None, store, }; Ok(Box::into_raw(Box::new(shard_view)) as i64) @@ -330,16 +429,24 @@ pub async unsafe fn execute_query( .memory_pool() .map(|p| p as Arc); - // Peek at the substrait extensions list to see if this is an indexed query. - // The `index_filter` UDF name appears there if Calcite planted any - // index_filter(bytes) calls. Cheap — just bytes inspection. - let is_indexed = plan_bytes_mentions_index_filter(plan_bytes); + // Peek at plan bytes for routing signals. + let (is_indexed, has_row_id) = inspect_plan_bytes(plan_bytes); // Register cancellation token. let token = query_tracker::get_cancellation_token(context_id); let query_future = async move { - if is_indexed { + // Routing logic: + // 1. Indexed query (has index_filter) → always indexed path + // 2. Has __row_id__ but not indexed (non-indexed + sort) → consult FetchStrategy + // - ListingTable → vanilla path with ShardTableProvider + ProjectRowIdOptimizer + // - IndexedPredicateOnly → indexed path (position-based row IDs) + // - None → vanilla path (no row ID computation) + // 3. Neither → vanilla path + let use_indexed = is_indexed + || (has_row_id && query_config.fetch_strategy != crate::datafusion_query_config::FetchStrategy::ListingTable); + + if use_indexed { let qc = Arc::new(query_config); crate::indexed_executor::execute_indexed_query( plan_bytes.to_vec(), @@ -385,13 +492,140 @@ pub async unsafe fn execute_query( /// is no automatic retry on the vanilla path — a false positive is a hard /// query error. In practice this is unreachable because the needle is not a /// valid DataFusion identifier anywhere else a plan would naturally contain +/// QTF fetch phase: read specific rows by global row ID. +/// +/// Uses shared helpers from query_executor for runtime setup, file info building, +/// and stream wrapping. The fetch-specific logic is building ParquetAccessPlans +/// from row IDs and computing global __row_id__ = __row_id__ + row_base in SQL. +pub async unsafe fn fetch_by_row_ids( + shard_view: &ShardView, + runtime: &DataFusionRuntime, + manager: &crate::runtime_manager::RuntimeManager, + row_ids: Vec, + columns: Vec, +) -> Result { + use crate::indexed_table::row_selection::build_row_selection_with_min_skip_run; + use crate::indexed_table::segment_info::build_segments; + use crate::query_executor::{build_query_runtime_env, store_url_from_table_path, wrap_stream_as_handle}; + + // ── 1. Build RuntimeEnv + SessionContext ── + + let runtime_env = build_query_runtime_env(runtime, &shard_view.table_path, shard_view.object_metas.as_ref())?; + + // Register shard-specific object store on file:// scheme for this query. + runtime_env.register_object_store( + &url::Url::parse("file://").unwrap(), + Arc::clone(&shard_view.store), + ); + + let mut config = SessionConfig::new(); + config.options_mut().execution.parquet.pushdown_filters = true; + config.options_mut().execution.target_partitions = 1; + + let state = SessionStateBuilder::new() + .with_config(config) + .with_runtime_env(runtime_env) + .with_default_features() + .build(); + let ctx = SessionContext::new_with_state(state); + + // ── 2. Build ShardFileInfo with ParquetAccessPlan per file ── + + let store = ctx.state().runtime_env().object_store(&shard_view.table_path)?; + let (segments, _schema) = build_segments(&ctx.state(), Arc::clone(&store), shard_view.object_metas.as_ref()) + .await + .map_err(DataFusionError::Execution)?; + + // Distribute global row_ids to per-file local positions + let mut per_segment: HashMap = HashMap::new(); + for &gid in &row_ids { + let seg_idx = segments + .partition_point(|s| s.global_base <= gid as u64) + .saturating_sub(1); + let local_pos = (gid as u64 - segments[seg_idx].global_base) as u32; + per_segment.entry(seg_idx).or_default().insert(local_pos); + } + + // Build file infos with access plans for targeted row retrieval + let mut files: Vec = Vec::new(); + for seg in &segments { + let num_rgs = seg.row_groups.len(); + let access_plan = if let Some(bm) = per_segment.get(&(seg.segment_ord as usize)) { + let mut plan = ParquetAccessPlan::new_none(num_rgs); + for rg in &seg.row_groups { + let rg_start = rg.first_row as u32; + let rg_end = rg_start + rg.num_rows as u32; + let rg_bitmap: RoaringBitmap = bm + .iter() + .filter(|&pos| pos >= rg_start && pos < rg_end) + .map(|pos| pos - rg_start) + .collect(); + if !rg_bitmap.is_empty() { + let selection = build_row_selection_with_min_skip_run( + &rg_bitmap, rg.num_rows as usize, 1, + ); + plan.set(rg.index, RowGroupAccess::Selection(selection)); + } + } + Some(plan) + } else { + Some(ParquetAccessPlan::new_none(num_rgs)) + }; + + files.push(ShardFileInfo { + object_meta: shard_view.object_metas[seg.segment_ord as usize].clone(), + row_base: seg.global_base as i64, + num_rows: seg.max_doc as u64, + row_group_row_counts: seg.row_groups.iter().map(|rg| rg.num_rows as u64).collect(), + access_plan, + }); + } + + // ── 3. Register ShardTableProvider ── + + let store_url = store_url_from_table_path(&shard_view.table_path)?; + let listing_options = datafusion::datasource::listing::ListingOptions::new( + Arc::new(datafusion::datasource::file_format::parquet::ParquetFormat::new()) + ).with_file_extension(".parquet").with_collect_stat(true); + let resolved_schema = listing_options.infer_schema(&ctx.state(), &shard_view.table_path).await?; + + let provider = Arc::new(ShardTableProvider::new(ShardTableConfig { + file_schema: resolved_schema, + files, + store_url, + })); + ctx.register_table("t", provider)?; + + // ── 4. Execute SQL: compute global __row_id__ = __row_id__ + row_base ── + + let col_list = columns.iter() + .map(|c| format!("\"{}\"", c)) + .collect::>() + .join(", "); + let sql = format!( + "SELECT (\"__row_id__\" + \"row_base\") AS \"__row_id__\", {} FROM t", + col_list + ); + let df = ctx.sql(&sql).await?; + let physical_plan = df.create_physical_plan().await?; + let df_stream = execute_stream(physical_plan, ctx.task_ctx())?; + + // ── 5. Wrap and return ── + + Ok(wrap_stream_as_handle(df_stream, manager.cpu_executor(), runtime)) +} + /// it; the failure mode is documented here to keep the dispatch contract /// explicit. -fn plan_bytes_mentions_index_filter(plan_bytes: &[u8]) -> bool { - // The substrait plan carries extension-function names as UTF-8 strings. - // Substring match is sufficient for dispatch. - const NEEDLE: &[u8] = b"index_filter"; - plan_bytes.windows(NEEDLE.len()).any(|w| w == NEEDLE) +/// Inspect substrait plan bytes for routing signals. +/// Returns (has_index_filter, has_row_id). +fn inspect_plan_bytes(plan_bytes: &[u8]) -> (bool, bool) { + const INDEX_FILTER: &[u8] = b"index_filter"; + const ROW_ID: &[u8] = crate::ROW_ID_COLUMN_NAME.as_bytes(); + ( + plan_bytes.windows(INDEX_FILTER.len()).any(|w| w == INDEX_FILTER), + plan_bytes.windows(ROW_ID.len()).any(|w| w == ROW_ID), + ) } /// Returns the Arrow schema for the given stream as a heap-allocated FFI_ArrowSchema pointer. diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/datafusion_query_config.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/datafusion_query_config.rs index bd1ef342d3d4b..029bab533fe49 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/datafusion_query_config.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/datafusion_query_config.rs @@ -8,6 +8,26 @@ use crate::indexed_table::eval::single_collector::CollectorCallStrategy; use crate::indexed_table::stream::FilterStrategy; +/// Selects which execution path computes shard-global row IDs. +/// +/// Selects which execution path computes shard-global row IDs. +/// `None` = no row ID computation (baseline — reads ___row_id as a regular column). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FetchStrategy { + /// No row ID optimizer applied. ___row_id is read as a regular column + /// without any row_base addition. Returns local (per-file) row IDs only. + None, + /// ShardTableProvider + ProjectRowIdOptimizer. + /// Reads ___row_id from parquet, adds row_base via physical optimizer rewrite. + /// Produces shard-global absolute row IDs. + ListingTable, + /// Predicate-only mode in the indexed executor. + /// Uses indexed pipeline (segment partitioning, prefetch, PositionMap). + /// Does NOT read ___row_id from disk — computes from position: + /// global_base + rg.first_row + position_in_rg. Zero column I/O for row ID. + IndexedPredicateOnly, +} + /// Query-scoped configuration. Owned by value after FFM decode. #[derive(Debug, Clone)] pub struct DatafusionQueryConfig { @@ -46,6 +66,10 @@ pub struct DatafusionQueryConfig { /// `TightenOuterBounds` is the default — multiple collectors in the /// tree means `PageRangeSplit` would multiply FFM calls. pub tree_collector_strategy: CollectorCallStrategy, + /// Strategy for row ID emission on the vanilla path. + /// Only consulted when the plan requests row IDs (contains _global_row_id() UDF + /// or projects ___row_id). + pub fetch_strategy: FetchStrategy, } /// FFM wire format. Must stay in lockstep with the Java `MemoryLayout`. @@ -75,6 +99,8 @@ pub struct WireDatafusionQueryConfig { pub single_collector_strategy: i32, /// 0 = FullRange, 1 = TightenOuterBounds, 2 = PageRangeSplit pub tree_collector_strategy: i32, + /// 0 = None (baseline), 1 = ListingTable, 2 = IndexedPredicateOnly + pub fetch_strategy: i32, } impl DatafusionQueryConfig { @@ -97,6 +123,7 @@ impl DatafusionQueryConfig { max_collector_parallelism: 1, single_collector_strategy: CollectorCallStrategy::PageRangeSplit, tree_collector_strategy: CollectorCallStrategy::TightenOuterBounds, + fetch_strategy: FetchStrategy::None, } } @@ -162,6 +189,11 @@ impl DatafusionQueryConfig { 2 => CollectorCallStrategy::PageRangeSplit, _ => CollectorCallStrategy::TightenOuterBounds, }, + fetch_strategy: match w.fetch_strategy { + 1 => FetchStrategy::ListingTable, + 2 => FetchStrategy::IndexedPredicateOnly, + _ => FetchStrategy::None, + }, } } } @@ -272,6 +304,7 @@ mod tests { max_collector_parallelism: 4, single_collector_strategy: 2, tree_collector_strategy: 1, + fetch_strategy: 1, }; let ptr = &wire as *const _ as i64; let c = unsafe { DatafusionQueryConfig::from_ffm_ptr(ptr) }; @@ -285,6 +318,7 @@ mod tests { assert_eq!(c.force_pushdown, Some(false)); assert_eq!(c.cost_predicate, 3); assert_eq!(c.cost_collector, 17); + assert_eq!(c.fetch_strategy, FetchStrategy::ListingTable); } #[test] @@ -303,10 +337,12 @@ mod tests { max_collector_parallelism: 2, single_collector_strategy: 2, tree_collector_strategy: 1, + fetch_strategy: 0, }; let ptr = &wire as *const _ as i64; let c = unsafe { DatafusionQueryConfig::from_ffm_ptr(ptr) }; assert_eq!(c.force_strategy, None); assert_eq!(c.force_pushdown, None); + assert_eq!(c.fetch_strategy, FetchStrategy::None); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index a9ca72aef2cee..380ae94314f0a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -178,6 +178,53 @@ pub unsafe extern "C" fn df_execute_query( .map_err(|e| e.to_string()) } +/// Fetch specific rows by global row ID — QTF fetch phase. +/// +/// Row IDs are passed as a direct pointer to i64 values (from BigIntVector's +/// off-heap ArrowBuf). Zero-copy at FFM boundary: Rust reads directly from +/// Java's off-heap buffer without any intermediate allocation. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn df_fetch_by_row_ids( + shard_view_ptr: i64, + row_ids_ptr: i64, + row_ids_count: i64, + col_names_ptr: *const *const u8, + col_names_len_ptr: *const i64, + col_names_count: i64, + runtime_ptr: i64, +) -> i64 { + let mgr = get_rt_manager()?; + let shard_view = &*(shard_view_ptr as *const crate::api::ShardView); + let runtime = &*(runtime_ptr as *const crate::api::DataFusionRuntime); + + // Zero-copy read from BigIntVector's direct buffer + let row_ids: Vec = slice::from_raw_parts( + row_ids_ptr as *const i64, + row_ids_count as usize, + ).to_vec(); + + // Parse column names + let mut columns: Vec = Vec::with_capacity(col_names_count as usize); + for i in 0..col_names_count as usize { + let ptr = *col_names_ptr.add(i); + let len = *col_names_len_ptr.add(i); + let name = str_from_raw(ptr, len) + .map_err(|e| format!("df_fetch_by_row_ids: column name: {}", e))?; + columns.push(name.to_string()); + } + + mgr.io_runtime + .block_on(crate::api::fetch_by_row_ids( + shard_view, + runtime, + &mgr, + row_ids, + columns, + )) + .map_err(|e| e.to_string()) +} + #[ffm_safe] #[no_mangle] pub unsafe extern "C" fn df_stream_get_schema(stream_ptr: i64) -> i64 { @@ -709,9 +756,12 @@ pub unsafe extern "C" fn df_execute_with_context( let plan_bytes = slice::from_raw_parts(plan_ptr, plan_len as usize); let cpu_executor = mgr.cpu_executor(); // Route based on whether the session was configured for indexed execution - if session_handle.indexed_config.is_some() { - // TODO: refactor execute_indexed_with_context to take SessionContextHandle directly - // (like execute_with_context) instead of i64 raw pointer — avoids this re-boxing. + // or if the plan projects __row_id__ (QTF query phase). + let has_row_id = plan_bytes.windows(crate::ROW_ID_COLUMN_NAME.len()).any(|w| w == crate::ROW_ID_COLUMN_NAME.as_bytes()); + let fetch_strategy = session_handle.query_config.fetch_strategy; + let use_indexed = session_handle.indexed_config.is_some() + || (has_row_id && fetch_strategy != crate::datafusion_query_config::FetchStrategy::ListingTable); + if use_indexed { let ptr = Box::into_raw(Box::new(session_handle)) as i64; mgr.io_runtime .block_on(crate::task_monitors::query_execution_monitor().instrument( diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs index 1940ffa1c42fb..23c254d181762 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs @@ -61,8 +61,8 @@ use crate::indexed_table::index::RowGroupDocsCollector; use crate::indexed_table::page_pruner::PagePruner; use crate::indexed_table::segment_info::build_segments; use crate::indexed_table::substrait_to_tree::{ - classify_filter, create_index_filter_udf, expr_to_bool_tree, extract_filter_expr, - ExtractionResult, FilterClass, + classify_filter, create_index_filter_udf, expr_to_bool_tree, + extract_filter_expr, plan_requests_row_ids, ExtractionResult, FilterClass, }; use crate::indexed_table::table_provider::{ EvaluatorFactory, IndexedTableConfig, IndexedTableProvider, SegmentFileInfo, @@ -77,6 +77,7 @@ use crate::indexed_table::bool_tree::residual_bool_to_physical_expr; use crate::indexed_table::metrics::StreamMetrics; use crate::indexed_table::page_pruner::{build_pruning_predicate, PagePruneMetrics}; + /// Execute an indexed query. /// /// `shard_view` carries the segment's parquet paths (populated when the reader @@ -442,9 +443,6 @@ pub async unsafe fn execute_indexed_with_context( let (segments, schema) = build_segments(&state, Arc::clone(&store), object_metas.as_ref()) .await .map_err(DataFusionError::Execution)?; - for (i, seg) in segments.iter().enumerate() { - } - let placeholder: Arc = Arc::new(PlaceholderProvider { schema: schema.clone(), }); @@ -454,6 +452,7 @@ pub async unsafe fn execute_indexed_with_context( .map_err(|e| DataFusionError::Execution(format!("decode substrait: {}", e)))?; let logical_plan = from_substrait_plan(&ctx.state(), &plan).await?; + let emit_row_ids = plan_requests_row_ids(&logical_plan); let filter_expr = extract_filter_expr(&logical_plan); let extraction = match filter_expr { None => None, @@ -471,7 +470,6 @@ pub async unsafe fn execute_indexed_with_context( Some(e) => classify_filter(&e.tree), }, }; - // Derive the parquet pushdown predicate from the BoolNode tree. // `scan()` ignores DataFusion's filters argument (which contains // the `delegated_predicate` UDF marker whose body panics) and uses this @@ -493,6 +491,14 @@ pub async unsafe fn execute_indexed_with_context( .as_ref() .and_then(residual_bool_to_physical_expr) }), + FilterClass::None if emit_row_ids => { + // Predicate-only mode: no collectors, but there may be predicates. + // Convert the entire BoolNode tree to a PhysicalExpr for pushdown. + // If no predicates exist, this is None and we get a full scan. + extraction.as_ref().and_then(|e| { + residual_bool_to_physical_expr(&e.tree) + }) + } FilterClass::Tree | FilterClass::None => None, }; @@ -500,9 +506,41 @@ pub async unsafe fn execute_indexed_with_context( let factory: EvaluatorFactory = match classification { FilterClass::None => { - return Err(DataFusionError::Execution( - "execute_indexed_query called with no index_filter(...) in plan".into(), - )); + if emit_row_ids { + // Predicate-only mode with emit_row_ids: use SingleCollectorEvaluator + // with a no-op collector (returns all docs). The residual predicate + // handles filtering via page pruning + on_batch_mask. + // Row IDs are computed from position by IndexedStream. + let schema_for_pruner = schema.clone(); + let residual_expr: Option> = extraction.as_ref().and_then(|e| { + residual_bool_to_physical_expr(&e.tree) + }); + let residual_pruning_predicate: Option> = residual_expr + .as_ref() + .and_then(|expr| build_pruning_predicate(expr, Arc::clone(&schema_for_pruner))); + let call_strategy = query_config.single_collector_strategy; + + Arc::new( + move |segment: &SegmentFileInfo, _chunk, stream_metrics: &StreamMetrics| { + let pruner = Arc::new(PagePruner::new( + &schema_for_pruner, + Arc::clone(&segment.metadata), + )); + let eval: Arc = + Arc::new(crate::indexed_table::eval::predicate_evaluator::PredicateOnlyEvaluator::new( + pruner, + residual_pruning_predicate.clone(), + residual_expr.clone(), + Some(PagePruneMetrics::from_stream_metrics(stream_metrics)), + )); + Ok(eval) + }, + ) + } else { + return Err(DataFusionError::Execution( + "execute_indexed_query called with no index_filter(...) in plan".into(), + )); + } } FilterClass::SingleCollector => { let extraction = extraction.as_ref().ok_or_else(|| { @@ -679,6 +717,7 @@ pub async unsafe fn execute_indexed_with_context( let parsed = url::Url::parse(url_str) .map_err(|e| DataFusionError::Execution(format!("parse table_path URL: {}", e)))?; let store_url = ObjectStoreUrl::parse(format!("{}://{}", parsed.scheme(), parsed.authority()))?; + let provider = Arc::new(IndexedTableProvider::new(IndexedTableConfig { schema: schema.clone(), segments, @@ -688,6 +727,7 @@ pub async unsafe fn execute_indexed_with_context( pushdown_predicate, query_config: Arc::clone(&query_config), predicate_columns, + emit_row_ids, })); ctx.register_table(&table_name, provider)?; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/eval_helpers.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/eval_helpers.rs new file mode 100644 index 0000000000000..557cc505dbb27 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/eval_helpers.rs @@ -0,0 +1,123 @@ +/* + * 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. + */ + +//! Shared helpers for evaluators (SingleCollector, PredicateOnly, Tree). + +use std::sync::Arc; + +use datafusion::arrow::array::BooleanArray; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNode; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_optimizer::pruning::PruningPredicate; +use roaring::RoaringBitmap; + +use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner}; +use crate::indexed_table::stream::RowGroupInfo; + +/// Compute page-pruned ranges for a row group. +/// Returns `None` if no pruning predicate is available (all rows pass). +/// Returns `Some(vec![])` if all pages are pruned (RG can be skipped). +pub fn compute_page_ranges( + pruning_predicate: Option<&Arc>, + page_pruner: &PagePruner, + rg: &RowGroupInfo, + min_doc: i32, + page_prune_metrics: Option<&PagePruneMetrics>, +) -> Option> { + pruning_predicate.and_then(|pp| { + page_pruner + .prune_rg(pp, rg.index, page_prune_metrics) + .map(|sel| { + let mut ranges = Vec::new(); + let mut rg_pos: i64 = 0; + for s in sel.iter() { + if s.skip { + rg_pos += s.row_count as i64; + } else { + let abs_min = min_doc + rg_pos as i32; + let abs_max = min_doc + rg_pos as i32 + s.row_count as i32; + ranges.push((abs_min, abs_max)); + rg_pos += s.row_count as i64; + } + } + ranges + }) + }) +} + +/// Build a candidate bitmap from page-pruned ranges (universe — all surviving pages). +/// Returns `None` if all pages were pruned. +pub fn universe_bitmap_from_page_ranges( + page_ranges: &Option>, + rg: &RowGroupInfo, +) -> Option { + match page_ranges { + Some(r) if r.is_empty() => None, + Some(r) => { + let mut bm = RoaringBitmap::new(); + for (r_min, r_max) in r { + let lo = (*r_min as i64 - rg.first_row) as u32; + let hi = (*r_max as i64 - rg.first_row) as u32; + bm.insert_range(lo..hi); + } + Some(bm) + } + None => { + let mut bm = RoaringBitmap::new(); + bm.insert_range(0..rg.num_rows as u32); + Some(bm) + } + } +} + +/// Evaluate a residual predicate against a batch, returning a BooleanArray mask. +pub fn evaluate_residual( + residual: &Arc, + batch: &RecordBatch, + batch_len: usize, +) -> Result { + let remapped = remap_expr_to_batch(residual, batch)?; + let value = remapped + .evaluate(batch) + .map_err(|e| format!("evaluate_residual: {}", e))?; + let array = value + .into_array(batch_len) + .map_err(|e| format!("evaluate_residual into_array: {}", e))?; + array + .as_any() + .downcast_ref::() + .ok_or_else(|| "evaluate_residual: did not produce BooleanArray".to_string()) + .cloned() +} + +/// Remap column references in a PhysicalExpr to match the batch schema. +/// The expression may reference columns by index in the full table schema, +/// but the batch only contains projected columns. This rewrites Column +/// expressions to use the batch's field positions by name lookup. +pub fn remap_expr_to_batch( + expr: &Arc, + batch: &RecordBatch, +) -> Result, String> { + let batch_schema = batch.schema(); + expr.clone() + .transform(|node| { + use datafusion::common::tree_node::Transformed; + if let Some(col) = node.as_any().downcast_ref::() { + if let Ok(idx) = batch_schema.index_of(col.name()) { + let new_col: Arc = + Arc::new(Column::new(col.name(), idx)); + return Ok(Transformed::yes(new_col)); + } + } + Ok(Transformed::no(node)) + }) + .map(|t| t.data) + .map_err(|e| format!("remap_expr_to_batch: {}", e)) +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs index f59a3968f95a9..a41683c91e627 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/mod.rs @@ -39,6 +39,8 @@ //! Swapping impls requires only passing different `Arc`s at construction. pub mod bitmap_tree; +pub mod eval_helpers; +pub mod predicate_evaluator; pub mod single_collector; use std::any::Any; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs new file mode 100644 index 0000000000000..8152643673efb --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/predicate_evaluator.rs @@ -0,0 +1,104 @@ +/* + * 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. + */ + +//! Predicate-only evaluator — no collector, pure parquet-native filtering. +//! +//! Used for `FilterClass::None` with `emit_row_ids=true`: the query has no +//! `index_filter(...)` call (no Lucene collector), only DataFusion predicates. +//! Candidates default to the page-pruned universe; `on_batch_mask` evaluates +//! only the residual predicate. + +use std::sync::Arc; +use std::time::Instant; + +use datafusion::arrow::array::BooleanArray; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::physical_optimizer::pruning::PruningPredicate; +use roaring::RoaringBitmap; + +use super::eval_helpers::{compute_page_ranges, evaluate_residual, universe_bitmap_from_page_ranges}; +use super::{PrefetchedRg, RowGroupBitsetSource}; +use crate::indexed_table::page_pruner::{PagePruneMetrics, PagePruner}; +use crate::indexed_table::row_selection::{bitmap_to_packed_bits, PositionMap}; +use crate::indexed_table::stream::RowGroupInfo; + +/// Evaluator for predicate-only queries (no Collector). +/// +/// Candidates = page-pruned universe. Residual predicate applied in `on_batch_mask`. +pub struct PredicateOnlyEvaluator { + page_pruner: Arc, + pruning_predicate: Option>, + residual_expr: Option>, + page_prune_metrics: Option, +} + +impl PredicateOnlyEvaluator { + pub fn new( + page_pruner: Arc, + pruning_predicate: Option>, + residual_expr: Option>, + page_prune_metrics: Option, + ) -> Self { + Self { + page_pruner, + pruning_predicate, + residual_expr, + page_prune_metrics, + } + } +} + +impl RowGroupBitsetSource for PredicateOnlyEvaluator { + fn prefetch_rg( + &self, + rg: &RowGroupInfo, + min_doc: i32, + _max_doc: i32, + ) -> Result, String> { + let t = Instant::now(); + + let page_ranges = compute_page_ranges( + self.pruning_predicate.as_ref(), + &self.page_pruner, + rg, + min_doc, + self.page_prune_metrics.as_ref(), + ); + + let candidates = match universe_bitmap_from_page_ranges(&page_ranges, rg) { + Some(bm) if bm.is_empty() => return Ok(None), + Some(bm) => bm, + None => return Ok(None), + }; + + let mask_len = rg.num_rows as usize; + let packed_bits = bitmap_to_packed_bits(&candidates, mask_len as u32); + let mask_buffer = datafusion::arrow::buffer::Buffer::from_vec(packed_bits); + Ok(Some(PrefetchedRg { + candidates, + eval_nanos: t.elapsed().as_nanos() as u64, + context: Box::new(()), + mask_buffer: Some(mask_buffer), + })) + } + + fn on_batch_mask( + &self, + _rg_state: &dyn std::any::Any, + _rg_first_row: i64, + _position_map: &PositionMap, + _batch_offset: usize, + batch_len: usize, + batch: &RecordBatch, + ) -> Result, String> { + let Some(ref residual) = self.residual_expr else { + return Ok(None); + }; + Ok(Some(evaluate_residual(residual, batch, batch_len)?)) + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs index 92eefa73739f9..1da748be36850 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/eval/single_collector.rs @@ -149,62 +149,66 @@ impl RowGroupBitsetSource for SingleCollectorEvaluator { }) }); - // Dispatch collector call strategy. - let call_ranges: Vec<(i32, i32)> = match self.call_strategy { - CollectorCallStrategy::FullRange => vec![(min_doc, max_doc)], - CollectorCallStrategy::TightenOuterBounds => match &page_ranges { - Some(r) if r.is_empty() => return Ok(None), - Some(r) => vec![(r.first().unwrap().0, r.last().unwrap().1)], - None => vec![(min_doc, max_doc)], - }, - CollectorCallStrategy::PageRangeSplit => match &page_ranges { - Some(r) if r.is_empty() => return Ok(None), - Some(r) => r.clone(), - None => vec![(min_doc, max_doc)], - }, - }; - - // Call collector for each range, merge into one RG-relative bitmap. - let mut candidates = RoaringBitmap::new(); - for (r_min, r_max) in &call_ranges { - let bitset = self - .collector - .collect_packed_u64_bitset(*r_min, *r_max) - .map_err(|e| { - format!( - "collector.collect_packed_u64_bitset(rg={}, [{}, {})): {}", - rg.index, r_min, r_max, e - ) - })?; - if let Some(ref c) = self.ffm_collector_calls { - c.add(1); - } - let offset = (*r_min as i64 - rg.first_row) as u32; - let num_docs = (*r_max - *r_min) as u32; - let bytes: &[u8] = unsafe { - std::slice::from_raw_parts(bitset.as_ptr() as *const u8, bitset.len() * 8) + // Build candidate bitmap from collector. + let candidates = { + let collector = &self.collector; + // Dispatch collector call strategy. + let call_ranges: Vec<(i32, i32)> = match self.call_strategy { + CollectorCallStrategy::FullRange => vec![(min_doc, max_doc)], + CollectorCallStrategy::TightenOuterBounds => match &page_ranges { + Some(r) if r.is_empty() => return Ok(None), + Some(r) => vec![(r.first().unwrap().0, r.last().unwrap().1)], + None => vec![(min_doc, max_doc)], + }, + CollectorCallStrategy::PageRangeSplit => match &page_ranges { + Some(r) if r.is_empty() => return Ok(None), + Some(r) => r.clone(), + None => vec![(min_doc, max_doc)], + }, }; - let mut chunk = RoaringBitmap::from_lsb0_bytes(offset, bytes); - let upper = offset.saturating_add(num_docs); - if upper < u32::MAX { - chunk.remove_range(upper..); + + // Call collector for each range, merge into one RG-relative bitmap. + let mut candidates = RoaringBitmap::new(); + for (r_min, r_max) in &call_ranges { + let bitset = collector + .collect_packed_u64_bitset(*r_min, *r_max) + .map_err(|e| { + format!( + "collector.collect_packed_u64_bitset(rg={}, [{}, {})): {}", + rg.index, r_min, r_max, e + ) + })?; + if let Some(ref c) = self.ffm_collector_calls { + c.add(1); + } + let offset = (*r_min as i64 - rg.first_row) as u32; + let num_docs = (*r_max - *r_min) as u32; + let bytes: &[u8] = unsafe { + std::slice::from_raw_parts(bitset.as_ptr() as *const u8, bitset.len() * 8) + }; + let mut chunk = RoaringBitmap::from_lsb0_bytes(offset, bytes); + let upper = offset.saturating_add(num_docs); + if upper < u32::MAX { + chunk.remove_range(upper..); + } + candidates |= chunk; } - candidates |= chunk; - } - // For FullRange and TightenOuterBounds, AND with page bitmap - // to remove rows in dead pages that the collector scanned. - if self.call_strategy != CollectorCallStrategy::PageRangeSplit { - if let Some(ref ranges) = page_ranges { - let mut allowed = RoaringBitmap::new(); - for (r_min, r_max) in ranges { - let lo = (*r_min as i64 - rg.first_row) as u32; - let hi = (*r_max as i64 - rg.first_row) as u32; - allowed.insert_range(lo..hi); + // For FullRange and TightenOuterBounds, AND with page bitmap + // to remove rows in dead pages that the collector scanned. + if self.call_strategy != CollectorCallStrategy::PageRangeSplit { + if let Some(ref ranges) = page_ranges { + let mut allowed = RoaringBitmap::new(); + for (r_min, r_max) in ranges { + let lo = (*r_min as i64 - rg.first_row) as u32; + let hi = (*r_max as i64 - rg.first_row) as u32; + allowed.insert_range(lo..hi); + } + candidates &= allowed; } - candidates &= allowed; } - } + candidates + }; if candidates.is_empty() { return Ok(None); @@ -244,11 +248,7 @@ impl RowGroupBitsetSource for SingleCollectorEvaluator { let Some(ref residual) = self.residual_expr else { return Ok(None); }; - // Apply Collector bitmap AND residual predicate over the - // delivered batch. In row-granular mode (pushdown ON) this - // re-applies what parquet already did — redundant but correct. - // In block-granular mode (pushdown OFF) this is the only - // place the residual gets applied. + let state = rg_state .downcast_ref::() .ok_or_else(|| { @@ -303,27 +303,13 @@ impl RowGroupBitsetSource for SingleCollectorEvaluator { } }; - // Evaluate residual against the batch. The residual may use - // full-schema column indices; remap to batch positions by name. - let remapped_residual = remap_expr_to_batch(residual, batch) - .map_err(|e| format!("SingleCollectorEvaluator: remap residual: {}", e))?; - let residual_value = remapped_residual - .evaluate(batch) - .map_err(|e| format!("SingleCollectorEvaluator: residual.evaluate: {}", e))?; - let residual_array = residual_value - .into_array(batch_len) - .map_err(|e| format!("SingleCollectorEvaluator: residual into_array: {}", e))?; - let residual_mask = residual_array - .as_any() - .downcast_ref::() - .ok_or_else(|| { - "SingleCollectorEvaluator: residual did not produce BooleanArray".to_string() - })?; + // Evaluate residual against the batch. + let residual_mask = super::eval_helpers::evaluate_residual(residual, batch, batch_len)?; // AND with kleene semantics (NULL → exclude). let combined = datafusion::arrow::compute::kernels::boolean::and_kleene( &collector_mask, - residual_mask, + &residual_mask, ) .map_err(|e| format!("SingleCollectorEvaluator: and_kleene: {}", e))?; Ok(Some(combined)) @@ -524,7 +510,7 @@ mod tests { } /// Remap Column indices in a PhysicalExpr to match the batch schema by name. -fn remap_expr_to_batch( +pub fn remap_expr_to_batch( expr: &Arc, batch: &RecordBatch, ) -> Result, String> { diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/mod.rs index 6e36807864a26..7dedaf5097caa 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/mod.rs @@ -56,6 +56,7 @@ pub mod bool_tree; pub mod eval; +pub mod row_id_injection; pub mod ffm_callbacks; pub mod index; pub mod metrics; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/row_id_injection.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/row_id_injection.rs new file mode 100644 index 0000000000000..867774cec547b --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/row_id_injection.rs @@ -0,0 +1,138 @@ +/* + * 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. + */ + +//! Row ID computation for the fetch phase (QTF). +//! +//! Computes shard-global row IDs from position information and injects them +//! into the output batch at the correct column index. Used by `IndexedStream` +//! when `row_id_output_index` is set. + +use std::sync::Arc; + +use datafusion::arrow::array::{Array, BooleanArray, Int64Array}; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::Result; + +use super::row_selection::PositionMap; + +/// Pre-captured state needed for row ID computation. +/// Captured BEFORE the filter mask is consumed (since mask consumption advances offsets). +pub struct RowIdContext { + pub batch_offset: usize, + pub position_map: Option, + pub base: u64, + pub eval_mask: Option, +} + +/// Compute global row IDs for surviving rows and inject into the output batch. +/// +/// # Arguments +/// * `output` - The filtered batch (without `__row_id__` column) +/// * `ctx` - Pre-captured position info from before filtering +/// * `batch_len` - Original (pre-filter) batch length +/// * `current_mask` - Candidate-stage mask (used when eval_mask is None) +/// * `mask_offset_before` - mask_offset value before this batch was processed +/// * `row_id_idx` - Column index in the output schema for `__row_id__` +/// * `schema` - Output schema (includes `__row_id__` at `row_id_idx`) +pub fn inject_row_ids( + output: &RecordBatch, + ctx: &RowIdContext, + batch_len: usize, + current_mask: Option<&BooleanArray>, + mask_offset_before: usize, + row_id_idx: usize, + schema: &SchemaRef, +) -> Result { + let num_surviving = output.num_rows(); + + let row_id_array: Arc = match num_surviving { + 0 => Arc::new(Int64Array::from(Vec::::new())), + _ => { + let ids = compute_row_ids( + &ctx.eval_mask, + current_mask, + mask_offset_before, + batch_len, + ctx.batch_offset, + ctx.position_map.as_ref(), + ctx.base, + ); + Arc::new(Int64Array::from_iter_values(ids.into_iter().map(|id| id as i64))) + } + }; + + let batch_schema = output.schema(); + let columns: Vec> = schema + .fields() + .iter() + .enumerate() + .map(|(i, field)| match i { + idx if idx == row_id_idx => Arc::clone(&row_id_array), + _ => batch_schema + .index_of(field.name()) + .map(|col_idx| Arc::clone(output.column(col_idx))) + .unwrap_or_else(|_| datafusion::arrow::array::new_null_array(field.data_type(), num_surviving)), + }) + .collect(); + + RecordBatch::try_new_with_options( + schema.clone(), + columns, + &datafusion::arrow::record_batch::RecordBatchOptions::new().with_row_count(Some(num_surviving)), + ) + .map_err(|e| datafusion::common::DataFusionError::ArrowError(Box::new(e), None)) +} + +/// Compute global row IDs from position info. +fn compute_row_ids( + eval_mask: &Option, + current_mask: Option<&BooleanArray>, + mask_offset_before: usize, + batch_len: usize, + batch_start_delivered: usize, + pm: Option<&PositionMap>, + base: u64, +) -> Vec { + match eval_mask { + Some(mask) => { + (0..batch_len) + .filter(|&i| mask.is_valid(i) && mask.value(i)) + .map(|i| position_to_global_id(batch_start_delivered + i, pm, base)) + .collect() + } + None => match current_mask { + Some(candidate_mask) => { + (0..batch_len) + .filter(|&i| { + let mi = mask_offset_before + i; + mi < candidate_mask.len() + && candidate_mask.is_valid(mi) + && candidate_mask.value(mi) + }) + .map(|i| position_to_global_id(batch_start_delivered + i, pm, base)) + .collect() + } + None => { + (0..batch_len) + .map(|i| position_to_global_id(batch_start_delivered + i, pm, base)) + .collect() + } + }, + } +} + +/// Convert a delivered-row index to a shard-global row ID. +#[inline] +fn position_to_global_id(delivered_idx: usize, pm: Option<&PositionMap>, base: u64) -> u64 { + let rg_pos = match pm { + Some(p) => p.rg_position(delivered_idx).unwrap_or(delivered_idx), + None => delivered_idx, + }; + base + rg_pos as u64 +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/segment_info.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/segment_info.rs index b515e2e64d48a..7135433b40155 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/segment_info.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/segment_info.rs @@ -43,6 +43,7 @@ pub async fn build_segments( } let mut segments = Vec::with_capacity(object_metas.len()); + let mut cumulative_rows: u64 = 0; for (seg_ord, meta) in object_metas.iter().enumerate() { // Per-segment parquet metadata (RG info, page index) — still needed @@ -67,6 +68,8 @@ pub async fn build_segments( offset += num_rows; } let max_doc = offset; + let global_base = cumulative_rows; + cumulative_rows += max_doc as u64; segments.push(SegmentFileInfo { segment_ord: seg_ord as i32, @@ -75,6 +78,7 @@ pub async fn build_segments( parquet_size: size, row_groups, metadata: pq_meta, + global_base, }); } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs index 4588066f5833b..f64d441f50ab3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/stream.rs @@ -33,9 +33,9 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use datafusion::arrow::array::{Array, BooleanArray}; +use datafusion::arrow::array::{Array, BooleanArray, UInt64Array}; use datafusion::arrow::compute::filter_record_batch; -use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::Result; use datafusion::execution::SendableRecordBatchStream; @@ -65,6 +65,7 @@ pub struct RowGroupInfo { pub num_rows: i64, } + /// Test-only override for the per-RG `min_skip_run` selectivity heuristic. /// `IndexedStream` normally picks `min_skip_run` from candidate /// selectivity; setting `force_strategy` to one of these variants pins the @@ -270,6 +271,13 @@ pub struct IndexedExec { /// from the same query; read once per RG into local fields inside /// `IndexedStream` so the hot path never touches the Arc. pub(crate) query_config: Arc, + /// Cumulative row offset for this segment within the shard. + pub(crate) global_base: u64, + /// When true, the `___row_id` column is computed from position instead of read. + pub(crate) emit_row_ids: bool, + /// Index in the OUTPUT schema where computed `___row_id` should be inserted. + /// `None` when `emit_row_ids=false` or `___row_id` is not in projection. + pub(crate) row_id_output_index: Option, } impl fmt::Debug for IndexedExec { @@ -362,6 +370,9 @@ impl ExecutionPlan for IndexedExec { self.query_config.min_skip_run_selectivity_threshold, self.query_config.indexed_pushdown_filters, self.query_config.batch_size, + self.global_base, + self.emit_row_ids, + self.row_id_output_index, ))) } } @@ -424,6 +435,12 @@ struct IndexedStream { /// calling it twice (assert panic) and to signal "no more input /// will arrive; drain remaining completed batches." coalescer_finished: bool, + /// Cumulative row offset for this segment within the shard. + global_base: u64, + /// When true, the `___row_id` column is computed from position. + emit_row_ids: bool, + /// Index in the output schema where computed `___row_id` is inserted. + row_id_output_index: Option, } impl IndexedStream { @@ -446,9 +463,13 @@ impl IndexedStream { min_skip_run_selectivity_threshold: f64, indexed_pushdown_filters: bool, target_batch_size: usize, + global_base: u64, + emit_row_ids: bool, + row_id_output_index: Option, ) -> Self { let evaluator = Arc::clone(&index_reader.evaluator); - let batch_coalescer = LimitedBatchCoalescer::new(schema.clone(), target_batch_size, None); + let batch_coalescer = + LimitedBatchCoalescer::new(schema.clone(), target_batch_size, None); Self { schema, full_schema, @@ -480,6 +501,9 @@ impl IndexedStream { batch_coalescer, upstream_done: false, coalescer_finished: false, + global_base, + emit_row_ids, + row_id_output_index, } } @@ -558,6 +582,18 @@ impl IndexedStream { t.add_duration(t_on_batch.elapsed()); } + // Capture position info BEFORE mask is consumed (needed for row ID computation). + let row_id_ctx = if self.row_id_output_index.is_some() { + Some(super::row_id_injection::RowIdContext { + batch_offset: self.batch_offset, + position_map: self.current_position_map.as_ref().cloned(), + base: self.global_base + self.current_rg_first_row as u64, + eval_mask: eval_mask.clone(), + }) + } else { + None + }; + let output = match eval_mask { Some(mask) => { self.mask_offset += batch_len; @@ -597,9 +633,21 @@ impl IndexedStream { } }; - // Strip extra predicate columns to match output schema + // Strip extra predicate columns and inject computed __row_id__. let t_proj = Instant::now(); - let output = if output.num_columns() > self.schema.fields().len() { + let output = if let Some(row_id_idx) = self.row_id_output_index { + let ctx = row_id_ctx.unwrap(); + let mask_offset_before = self.mask_offset.saturating_sub(batch_len); + super::row_id_injection::inject_row_ids( + &output, + &ctx, + batch_len, + self.current_mask.as_ref(), + mask_offset_before, + row_id_idx, + &self.schema, + )? + } else if output.num_columns() > self.schema.fields().len() { let n = self.schema.fields().len(); if n == 0 { RecordBatch::try_new_with_options( diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/substrait_to_tree.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/substrait_to_tree.rs index 272d8ffe69879..48fce9f7c7aec 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/substrait_to_tree.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/substrait_to_tree.rs @@ -48,6 +48,18 @@ use super::bool_tree::BoolNode; /// The UDF name Calcite emits for indexed-column filter markers. pub const COLLECTOR_FUNCTION_NAME: &str = "delegated_predicate"; +/// Walk the logical plan looking for `__row_id__` column in a projection. +/// Its presence signals the executor to emit computed row IDs (query phase). +pub fn plan_requests_row_ids(plan: &LogicalPlan) -> bool { + match plan { + LogicalPlan::Projection(proj) => proj.expr.iter().any(|e| match e { + Expr::Column(col) => col.name() == crate::ROW_ID_COLUMN_NAME, + _ => false, + }), + _ => plan.inputs().iter().any(|child| plan_requests_row_ids(child)), + } +} + /// Classification of a query's filter expression — drives the evaluator choice. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum FilterClass { diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/table_provider.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/table_provider.rs index 6843eb743f7e1..2284f87402c1c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/table_provider.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/table_provider.rs @@ -29,7 +29,7 @@ use std::fmt; use std::sync::Arc; use async_trait::async_trait; -use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::catalog::{Session, TableProvider}; use datafusion::common::{Result, Statistics}; use datafusion::datasource::TableType; @@ -61,6 +61,9 @@ pub struct SegmentFileInfo { pub parquet_size: u64, pub row_groups: Vec, pub metadata: Arc, + /// Cumulative row count from all preceding segments. Used to compute + /// shard-global row IDs: `global_base + rg.first_row + position_in_rg`. + pub global_base: u64, } /// Factory: build a `RowGroupBitsetSource` for one `SegmentChunk`. @@ -120,6 +123,10 @@ pub struct IndexedTableConfig { pub query_config: Arc, /// Full-schema column indices referenced by BoolNode Predicate leaves. pub predicate_columns: Vec, + /// When true, the `___row_id` column in the output projection is computed + /// from position (global_base + rg.first_row + position_in_rg) instead of + /// being read from parquet. Other projected columns are read normally. + pub emit_row_ids: bool, } /// Table provider. Returns a `QueryShardExec` that fans out across chunks. @@ -177,13 +184,55 @@ impl TableProvider for IndexedTableProvider { _limit: Option, ) -> Result> { let full_schema = self.config.schema.clone(); - // Output schema = what DataFusion expects - let output_schema: SchemaRef = match projection { - Some(proj) => Arc::new(full_schema.project(proj)?), - None => full_schema.clone(), + + // Detect __row_id__ in the output projection when emit_row_ids=true. + // If present, we strip it from the parquet read and compute it from position. + let row_id_col_in_full_schema = full_schema.index_of(crate::ROW_ID_COLUMN_NAME).ok(); + let row_id_output_index: Option = if self.config.emit_row_ids { + match projection { + Some(proj) => proj.iter().position(|&idx| Some(idx) == row_id_col_in_full_schema), + None => row_id_col_in_full_schema, + } + } else { + None + }; + + // Output schema = what DataFusion expects (includes ___row_id if projected). + // When computing row IDs, replace the ___row_id field type with UInt64. + let output_schema: SchemaRef = { + let base: SchemaRef = match projection { + Some(proj) => Arc::new(full_schema.project(proj)?), + None => full_schema.clone(), + }; + if let Some(idx) = row_id_output_index { + let mut fields: Vec = base.fields().iter().map(|f| f.as_ref().clone()).collect(); + fields[idx] = Field::new(crate::ROW_ID_COLUMN_NAME, DataType::Int64, false); + Arc::new(Schema::new(fields)) + } else { + base + } }; - // Read projection = output + predicate columns for evaluator - let read_projection: Option> = if self.config.predicate_columns.is_empty() { + + // Read projection = output columns (minus ___row_id) + predicate columns for evaluator. + let read_projection: Option> = if self.config.emit_row_ids { + let output_cols: Vec = match projection { + Some(proj) => proj.iter() + .filter(|&&idx| Some(idx) != row_id_col_in_full_schema) + .copied() + .collect(), + None => (0..full_schema.fields().len()) + .filter(|&idx| Some(idx) != row_id_col_in_full_schema) + .collect(), + }; + let mut cols = output_cols; + for &idx in &self.config.predicate_columns { + if !cols.contains(&idx) { + cols.push(idx); + } + } + cols.sort(); + Some(cols) + } else if self.config.predicate_columns.is_empty() { projection.cloned() } else { projection.map(|proj| { @@ -197,6 +246,7 @@ impl TableProvider for IndexedTableProvider { cols }) }; + let projected_schema = output_schema; // Ignore DataFusion's `filters` argument. The `index_filter(...)` @@ -239,6 +289,7 @@ impl TableProvider for IndexedTableProvider { predicate, metrics: ExecutionPlanMetricsSet::new(), inner_parquet_metrics: Arc::new(std::sync::Mutex::new(Vec::new())), + row_id_output_index, })) } @@ -264,6 +315,9 @@ pub struct QueryShardExec { predicate: Option>, metrics: ExecutionPlanMetricsSet, inner_parquet_metrics: Arc>>, + /// Column index in the OUTPUT schema where computed `___row_id` should be + /// injected. `None` means no row ID computation (normal data path). + row_id_output_index: Option, } impl fmt::Debug for QueryShardExec { @@ -385,6 +439,9 @@ impl ExecutionPlan for QueryShardExec { metrics: ExecutionPlanMetricsSet::new(), stream_metrics: stream_metrics.clone(), query_config: Arc::clone(&self.config.query_config), + global_base: segment.global_base, + emit_row_ids: self.config.emit_row_ids, + row_id_output_index: self.row_id_output_index, }; execs.push(Arc::new(exec)); } @@ -446,6 +503,7 @@ mod tests { crate::datafusion_query_config::DatafusionQueryConfig::test_default(), ), predicate_columns: vec![], + emit_row_ids: false, } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/harness.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/harness.rs index ace5200830dff..5272dd15b2381 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/harness.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/harness.rs @@ -104,6 +104,7 @@ pub(in crate::indexed_table::tests_e2e) fn load_segment(corpus: &Corpus) -> Load parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }); global_first_row += seg_rows as i64; } @@ -285,6 +286,7 @@ pub(in crate::indexed_table::tests_e2e) async fn execute_tree_with_plan_pushdown pushdown_predicate: None, query_config: Arc::new(qc), predicate_columns: collect_predicate_column_indices(&bool_tree), + emit_row_ids: false, })); let ctx = SessionContext::new(); @@ -457,6 +459,7 @@ async fn run_single_collector_query( pushdown_predicate, query_config: Arc::new(qc), predicate_columns: pred_cols, + emit_row_ids: false, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); @@ -663,7 +666,8 @@ async fn run_with_factory_plan( evaluator_factory: factory, pushdown_predicate, query_config: Arc::new(qc), - predicate_columns: vec![], // run_with_factory_plan is low-level; caller controls projection + predicate_columns: vec![], + emit_row_ids: false, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs index b24771e38e09b..421b93cb97920 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/mod.rs @@ -18,7 +18,7 @@ use std::sync::Arc; use std::sync::OnceLock; -use datafusion::arrow::array::{Array, Int32Array, StringArray}; +use datafusion::arrow::array::{Array, Int32Array, Int64Array, StringArray}; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::ScalarValue; @@ -43,6 +43,9 @@ mod metrics; mod multi_segment; mod null_columns; mod page_pruning; +mod qtf_fetch_phase; +mod row_id_emission; +mod row_id_strategies; mod schema_drift; mod streaming_at_scale; @@ -103,11 +106,13 @@ fn build_fixture_schema() -> SchemaRef { Field::new("price", DataType::Int32, false), Field::new("status", DataType::Utf8, false), Field::new("category", DataType::Utf8, false), + Field::new("__row_id__", DataType::Int64, false), ])) } fn write_fixture_parquet() -> NamedTempFile { let schema = build_fixture_schema(); + let row_ids: Vec = (0..16).collect(); let batch = RecordBatch::try_new( schema.clone(), vec![ @@ -115,6 +120,7 @@ fn write_fixture_parquet() -> NamedTempFile { Arc::new(Int32Array::from(PRICES.to_vec())), Arc::new(StringArray::from(STATUSES.to_vec())), Arc::new(StringArray::from(CATEGORIES.to_vec())), + Arc::new(Int64Array::from(row_ids)), ], ) .unwrap(); @@ -221,6 +227,7 @@ async fn run_tree_and_plan( parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }; // Normalize NOT push-down; build one collector per Collector leaf in DFS order. @@ -287,6 +294,7 @@ async fn run_tree_and_plan( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], + emit_row_ids: false, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/multi_segment.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/multi_segment.rs index b4952ed729114..bb75ec43f5faa 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/multi_segment.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/multi_segment.rs @@ -130,6 +130,7 @@ async fn run_two_segment_query( parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }); } @@ -178,6 +179,7 @@ async fn run_two_segment_query( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], + emit_row_ids: false, })); let ctx = SessionContext::new(); @@ -332,6 +334,7 @@ async fn run_segments(specs: Vec, num_partitions: usize) -> Vec<(i32, S parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }); } @@ -376,6 +379,7 @@ async fn run_segments(specs: Vec, num_partitions: usize) -> Vec<(i32, S pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], + emit_row_ids: false, })); let ctx = SessionContext::new(); @@ -815,6 +819,7 @@ async fn run_wide_segments( parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }); } @@ -877,6 +882,7 @@ async fn run_wide_segments( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], + emit_row_ids: false, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/null_columns.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/null_columns.rs index 10bf6396b6041..07629b2f2d730 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/null_columns.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/null_columns.rs @@ -332,6 +332,7 @@ async fn assert_engine_matches_reference_null(name: &str, tree: NT) { parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }; let tree = Arc::new(bt); @@ -379,6 +380,7 @@ async fn assert_engine_matches_reference_null(name: &str, tree: NT) { pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], + emit_row_ids: false, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/page_pruning.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/page_pruning.rs index cce78dd0f5d64..db09eb6523c36 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/page_pruning.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/page_pruning.rs @@ -216,6 +216,7 @@ fn load_segment(tmp: &NamedTempFile) -> (SegmentFileInfo, SchemaRef) { parquet_size: size, row_groups: rgs, metadata: parquet_meta, + global_base: 0, }; (seg, schema) } @@ -385,6 +386,7 @@ async fn execute_and_collect( pushdown_predicate: None, query_config: Arc::new(qc), predicate_columns: vec![], + emit_row_ids: false, })); let ctx = SessionContext::new(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/qtf_fetch_phase.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/qtf_fetch_phase.rs new file mode 100644 index 0000000000000..a942e3da50285 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/qtf_fetch_phase.rs @@ -0,0 +1,531 @@ +/* + * 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. + */ + +//! End-to-end tests for the full data-node QTF loop: +//! query phase (emit row IDs) -> fetch phase (retrieve by those IDs) -> verify data correctness. +//! +//! The query phase uses `IndexedTableProvider` with `emit_row_ids: true` to get +//! global row IDs. The fetch phase uses `ShardTableProvider` with `ParquetAccessPlan` +//! to read only the targeted rows and compute `__row_id__ + row_base`. + +use std::sync::Arc; + +use datafusion::arrow::array::{Array, Int32Array, Int64Array, StringArray}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::datasource::physical_plan::parquet::{ParquetAccessPlan, RowGroupAccess}; +use datafusion::execution::context::SessionContext; +use datafusion::logical_expr::Operator; +use datafusion::parquet::arrow::arrow_reader::{ArrowReaderMetadata, ArrowReaderOptions}; +use futures::StreamExt; +use roaring::RoaringBitmap; + +use super::*; +use crate::indexed_table::row_selection::build_row_selection_with_min_skip_run; +use crate::shard_table_provider::{ShardFileInfo, ShardTableConfig, ShardTableProvider}; + +// ── Query phase helper (adapted from row_id_emission.rs) ──────────── + +/// Run the query phase with `emit_row_ids: true` and return sorted row IDs (as i64). +async fn query_phase(tree: BoolNode) -> Vec { + let tmp = write_fixture_parquet(); + let path = tmp.path().to_path_buf(); + let size = std::fs::metadata(&path).unwrap().len(); + + let file = std::fs::File::open(&path).unwrap(); + let meta = + ArrowReaderMetadata::load(&file, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let schema = meta.schema().clone(); + let parquet_meta = meta.metadata().clone(); + let mut rgs = Vec::new(); + let mut offset = 0i64; + for i in 0..parquet_meta.num_row_groups() { + let n = parquet_meta.row_group(i).num_rows(); + rgs.push(RowGroupInfo { + index: i, + first_row: offset, + num_rows: n, + }); + offset += n; + } + + let object_path = object_store::path::Path::from(path.to_string_lossy().as_ref()); + let segment = SegmentFileInfo { + segment_ord: 0, + max_doc: 16, + object_path, + parquet_size: size, + row_groups: rgs, + metadata: Arc::clone(&parquet_meta), + global_base: 0, + }; + + let tree = tree.push_not_down(); + let collectors = wire_collectors(&tree); + let per_leaf: Vec<(i32, Arc)> = collectors + .into_iter() + .enumerate() + .map(|(i, c)| (i as i32, c)) + .collect(); + let tree = Arc::new(tree); + let factory: super::super::table_provider::EvaluatorFactory = { + let per_leaf = per_leaf.clone(); + let tree = Arc::clone(&tree); + let schema = schema.clone(); + Arc::new(move |segment, _chunk, _stream_metrics| { + let resolved = tree.resolve(&per_leaf)?; + let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); + let eval: Arc = Arc::new(TreeBitsetSource { + tree: Arc::new(resolved), + evaluator: Arc::new(BitmapTreeEvaluator), + leaves: Arc::new( + crate::indexed_table::eval::bitmap_tree::CollectorLeafBitmaps { + ffm_collector_calls: _stream_metrics.ffm_collector_calls.clone(), + }, + ), + page_pruner: pruner, + cost_predicate: 1, + cost_collector: 10, + max_collector_parallelism: 1, + pruning_predicates: Arc::new(std::collections::HashMap::new()), + page_prune_metrics: Some( + crate::indexed_table::page_pruner::PagePruneMetrics::from_stream_metrics( + _stream_metrics, + ), + ), + collector_strategy: + crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + }); + Ok(eval) + }) + }; + + let store: Arc = + Arc::new(object_store::local::LocalFileSystem::new()); + let store_url = datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(); + let provider = Arc::new(IndexedTableProvider::new(IndexedTableConfig { + schema: schema.clone(), + segments: vec![segment], + store, + store_url, + evaluator_factory: factory, + pushdown_predicate: None, + query_config: Arc::new({ + let mut qc = crate::datafusion_query_config::DatafusionQueryConfig::test_default(); + qc.target_partitions = 1; + qc.force_strategy = Some(FilterStrategy::BooleanMask); + qc.force_pushdown = Some(false); + qc + }), + predicate_columns: vec![0, 1, 2, 3], + emit_row_ids: true, + })); + + let ctx = SessionContext::new(); + ctx.register_table("t", provider).unwrap(); + let df = ctx.sql("SELECT \"__row_id__\" FROM t").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let task_ctx = ctx.task_ctx(); + let mut stream = datafusion::physical_plan::execute_stream(plan, task_ctx).unwrap(); + let mut row_ids: Vec = Vec::new(); + while let Some(batch) = stream.next().await { + let b = batch.unwrap(); + let col = b.column(0).as_any().downcast_ref::().unwrap(); + for i in 0..b.num_rows() { + row_ids.push(col.value(i)); + } + } + row_ids.sort(); + row_ids +} + +// ── Fetch phase helper ────────────────────────────────────────────── + +/// Run the fetch phase: given row IDs, build a ShardTableProvider with +/// ParquetAccessPlan to read only those rows, then execute SQL to get data. +/// Returns (row_id, column_values) tuples sorted by row_id. +async fn fetch_phase( + row_ids: &[i64], + fetch_columns: &[&str], +) -> Vec { + if row_ids.is_empty() { + return vec![]; + } + + let tmp = write_fixture_parquet(); + let path = tmp.path().to_path_buf(); + + let file = std::fs::File::open(&path).unwrap(); + let meta = + ArrowReaderMetadata::load(&file, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let file_schema = meta.schema().clone(); + let parquet_meta = meta.metadata().clone(); + + // Build row group info + let mut rg_row_counts: Vec = Vec::new(); + for i in 0..parquet_meta.num_row_groups() { + rg_row_counts.push(parquet_meta.row_group(i).num_rows() as u64); + } + let total_rows: u64 = rg_row_counts.iter().sum(); + + // Build the ParquetAccessPlan from row_ids + let num_rgs = parquet_meta.num_row_groups(); + let mut rg_first_row: Vec = Vec::with_capacity(num_rgs); + let mut cumulative = 0u64; + for &count in &rg_row_counts { + rg_first_row.push(cumulative); + cumulative += count; + } + + // Distribute row IDs into per-RG bitmaps (local positions within each RG) + let mut plan = ParquetAccessPlan::new_none(num_rgs); + for rg_idx in 0..num_rgs { + let rg_start = rg_first_row[rg_idx]; + let rg_end = rg_start + rg_row_counts[rg_idx]; + let rg_num_rows = rg_row_counts[rg_idx] as usize; + + let mut rg_bitmap = RoaringBitmap::new(); + for &gid in row_ids { + let pos = gid as u64; + if pos >= rg_start && pos < rg_end { + rg_bitmap.insert((pos - rg_start) as u32); + } + } + if !rg_bitmap.is_empty() { + let selection = + build_row_selection_with_min_skip_run(&rg_bitmap, rg_num_rows, 1); + plan.set(rg_idx, RowGroupAccess::Selection(selection)); + } + } + + // Build the object meta + let file_size = std::fs::metadata(&path).unwrap().len(); + let object_meta = object_store::ObjectMeta { + location: object_store::path::Path::from(path.to_string_lossy().as_ref()), + last_modified: chrono::Utc::now(), + size: file_size, + e_tag: None, + version: None, + }; + + let shard_file = ShardFileInfo { + object_meta, + row_base: 0, + num_rows: total_rows, + row_group_row_counts: rg_row_counts, + access_plan: Some(plan), + }; + + // Build ShardTableProvider + let store_url = datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(); + let provider = Arc::new(ShardTableProvider::new(ShardTableConfig { + file_schema: file_schema.clone(), + files: vec![shard_file], + store_url: store_url.clone(), + })); + + // Register object store and table + let ctx = SessionContext::new(); + let store: Arc = + Arc::new(object_store::local::LocalFileSystem::new()); + ctx.register_object_store(store_url.as_ref(), store); + ctx.register_table("t", provider).unwrap(); + + // Execute SQL: __row_id__ + row_base gives global row ID + let col_list = fetch_columns + .iter() + .map(|c| format!("\"{}\"", c)) + .collect::>() + .join(", "); + let sql = format!( + "SELECT (\"__row_id__\" + \"row_base\") AS \"__row_id__\", {} FROM t", + col_list + ); + let df = ctx.sql(&sql).await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let task_ctx = ctx.task_ctx(); + let mut stream = datafusion::physical_plan::execute_stream(plan, task_ctx).unwrap(); + + let mut batches: Vec = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push(batch.unwrap()); + } + batches +} + +/// Combined helper: run query phase -> get row IDs, then run fetch phase -> get data. +/// Returns (sorted_row_ids, fetched_batches). +async fn query_then_fetch( + tree: BoolNode, + fetch_columns: Vec<&str>, +) -> (Vec, Vec) { + let row_ids = query_phase(tree).await; + let batches = fetch_phase(&row_ids, &fetch_columns).await; + (row_ids, batches) +} + +/// Extract (row_id, brand, price) tuples from fetch result batches, sorted by row_id. +fn extract_id_brand_price(batches: &[RecordBatch]) -> Vec<(i64, String, i32)> { + let mut rows: Vec<(i64, String, i32)> = Vec::new(); + for b in batches { + let ids = b.column(0).as_any().downcast_ref::().unwrap(); + let brands = b.column(1).as_any().downcast_ref::().unwrap(); + let prices = b.column(2).as_any().downcast_ref::().unwrap(); + for i in 0..b.num_rows() { + rows.push((ids.value(i), brands.value(i).to_string(), prices.value(i))); + } + } + rows.sort_by_key(|r| r.0); + rows +} + +/// Extract (row_id, brand) tuples from fetch result batches, sorted by row_id. +fn extract_id_brand(batches: &[RecordBatch]) -> Vec<(i64, String)> { + let mut rows: Vec<(i64, String)> = Vec::new(); + for b in batches { + let ids = b.column(0).as_any().downcast_ref::().unwrap(); + let brands = b.column(1).as_any().downcast_ref::().unwrap(); + for i in 0..b.num_rows() { + rows.push((ids.value(i), brands.value(i).to_string())); + } + } + rows.sort_by_key(|r| r.0); + rows +} + + +// ── Tests ──────────────────────────────────────────────────────────── + +/// Full QTF loop with SingleCollector: brand="amazon" -> get row IDs -> fetch brand+price -> verify. +/// Amazon rows: 0(50), 1(150), 2(80), 3(120), 12(30). +#[tokio::test] +async fn test_qtf_full_loop_single_collector() { + let tree = BoolNode::And(vec![index_leaf(0)]); // brand="amazon" + let (row_ids, batches) = query_then_fetch(tree, vec!["brand", "price"]).await; + + assert_eq!(row_ids, vec![0, 1, 2, 3, 12]); + + let rows = extract_id_brand_price(&batches); + assert_eq!(rows.len(), 5); + assert_eq!(rows[0], (0, "amazon".to_string(), 50)); + assert_eq!(rows[1], (1, "amazon".to_string(), 150)); + assert_eq!(rows[2], (2, "amazon".to_string(), 80)); + assert_eq!(rows[3], (3, "amazon".to_string(), 120)); + assert_eq!(rows[4], (12, "amazon".to_string(), 30)); +} + +/// Full QTF loop with predicate only: price > 100 -> get row IDs -> fetch brand+price -> verify. +/// Rows with price > 100: 1(150), 3(120), 5(95? no), 6(200), 9(300), 11(150). +#[tokio::test] +async fn test_qtf_full_loop_predicate_only() { + let tree = BoolNode::And(vec![pred_int("price", Operator::Gt, 100)]); + let (row_ids, batches) = query_then_fetch(tree, vec!["brand", "price"]).await; + + // price > 100: rows 1(150), 3(120), 6(200), 9(300), 11(150) + assert_eq!(row_ids, vec![1, 3, 6, 9, 11]); + + let rows = extract_id_brand_price(&batches); + assert_eq!(rows.len(), 5); + assert_eq!(rows[0], (1, "amazon".to_string(), 150)); + assert_eq!(rows[1], (3, "amazon".to_string(), 120)); + assert_eq!(rows[2], (6, "apple".to_string(), 200)); + assert_eq!(rows[3], (9, "google".to_string(), 300)); + assert_eq!(rows[4], (11, "samsung".to_string(), 150)); +} + +/// Full QTF loop with no filter: all 16 rows -> fetch all -> verify count. +/// Use price >= 0 as a match-all predicate (since emit_row_ids requires +/// going through the indexed path, we need at least a trivial predicate). +#[tokio::test] +async fn test_qtf_full_loop_no_filter() { + // price >= 0 matches all rows + let tree = BoolNode::And(vec![pred_int("price", Operator::GtEq, 0)]); + let (row_ids, batches) = query_then_fetch(tree, vec!["brand", "price"]).await; + + assert_eq!(row_ids.len(), 16); + assert_eq!(row_ids, (0..16).collect::>()); + + let rows = extract_id_brand_price(&batches); + assert_eq!(rows.len(), 16); + // Verify first and last rows + assert_eq!(rows[0], (0, "amazon".to_string(), 50)); + assert_eq!(rows[15], (15, "google".to_string(), 55)); +} + +/// Full QTF loop with two segments (global_base offset). +/// Write two separate parquet files, query both -> verify row IDs from both segments. +#[tokio::test] +async fn test_qtf_full_loop_two_segments() { + // For two segments, we write two separate parquet files and combine them + // in a single fetch pass. The first segment has global_base=0, the second + // has global_base=16 (since the first file has 16 rows). + + // Query phase: brand="amazon" on segment 0 gives rows 0,1,2,3,12. + // We simulate the second segment by doing a query with global_base=16, + // which would give 16,17,18,19,28. + // For the fetch phase, we verify we can fetch from a file with non-zero row_base. + + let tmp = write_fixture_parquet(); + let path = tmp.path().to_path_buf(); + + let file = std::fs::File::open(&path).unwrap(); + let meta = + ArrowReaderMetadata::load(&file, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let file_schema = meta.schema().clone(); + let parquet_meta = meta.metadata().clone(); + + let mut rg_row_counts: Vec = Vec::new(); + for i in 0..parquet_meta.num_row_groups() { + rg_row_counts.push(parquet_meta.row_group(i).num_rows() as u64); + } + let total_rows: u64 = rg_row_counts.iter().sum(); + + // Build access plan for rows 0, 1, 2, 3, 12 in the file (amazon brand) + let num_rgs = parquet_meta.num_row_groups(); + let mut rg_first_row: Vec = Vec::with_capacity(num_rgs); + let mut cumulative = 0u64; + for &count in &rg_row_counts { + rg_first_row.push(cumulative); + cumulative += count; + } + + let target_local_positions: Vec = vec![0, 1, 2, 3, 12]; + let mut plan = ParquetAccessPlan::new_none(num_rgs); + for rg_idx in 0..num_rgs { + let rg_start = rg_first_row[rg_idx]; + let rg_end = rg_start + rg_row_counts[rg_idx]; + let mut rg_bitmap = RoaringBitmap::new(); + for &pos in &target_local_positions { + if pos >= rg_start && pos < rg_end { + rg_bitmap.insert((pos - rg_start) as u32); + } + } + if !rg_bitmap.is_empty() { + let selection = build_row_selection_with_min_skip_run( + &rg_bitmap, + rg_row_counts[rg_idx] as usize, + 1, + ); + plan.set(rg_idx, RowGroupAccess::Selection(selection)); + } + } + + let file_size = std::fs::metadata(&path).unwrap().len(); + let object_meta = object_store::ObjectMeta { + location: object_store::path::Path::from(path.to_string_lossy().as_ref()), + last_modified: chrono::Utc::now(), + size: file_size, + e_tag: None, + version: None, + }; + + // Use row_base=1000 to simulate a second segment with global_base offset + let shard_file = ShardFileInfo { + object_meta, + row_base: 1000, + num_rows: total_rows, + row_group_row_counts: rg_row_counts, + access_plan: Some(plan), + }; + + let store_url = datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(); + let provider = Arc::new(ShardTableProvider::new(ShardTableConfig { + file_schema: file_schema.clone(), + files: vec![shard_file], + store_url: store_url.clone(), + })); + + let ctx = SessionContext::new(); + let store: Arc = + Arc::new(object_store::local::LocalFileSystem::new()); + ctx.register_object_store(store_url.as_ref(), store); + ctx.register_table("t", provider).unwrap(); + + let sql = "SELECT (\"__row_id__\" + \"row_base\") AS \"__row_id__\", \"brand\", \"price\" FROM t"; + let df = ctx.sql(sql).await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let task_ctx = ctx.task_ctx(); + let mut stream = datafusion::physical_plan::execute_stream(plan, task_ctx).unwrap(); + + let mut rows: Vec<(i64, String, i32)> = Vec::new(); + while let Some(batch) = stream.next().await { + let b = batch.unwrap(); + let ids = b.column(0).as_any().downcast_ref::().unwrap(); + let brands = b.column(1).as_any().downcast_ref::().unwrap(); + let prices = b.column(2).as_any().downcast_ref::().unwrap(); + for i in 0..b.num_rows() { + rows.push((ids.value(i), brands.value(i).to_string(), prices.value(i))); + } + } + rows.sort_by_key(|r| r.0); + + // With row_base=1000, the row IDs should be 1000+file_row_id. + // File has __row_id__ = [0..15], but we only read rows at positions 0,1,2,3,12. + // The __row_id__ in parquet for those positions = 0,1,2,3,12 + // Global: 0+1000=1000, 1+1000=1001, 2+1000=1002, 3+1000=1003, 12+1000=1012 + assert_eq!(rows.len(), 5); + assert_eq!(rows[0], (1000, "amazon".to_string(), 50)); + assert_eq!(rows[1], (1001, "amazon".to_string(), 150)); + assert_eq!(rows[2], (1002, "amazon".to_string(), 80)); + assert_eq!(rows[3], (1003, "amazon".to_string(), 120)); + assert_eq!(rows[4], (1012, "amazon".to_string(), 30)); +} + +/// Fetch only "brand" column -> verify only that column + __row_id__ are returned. +#[tokio::test] +async fn test_qtf_fetch_subset_columns() { + let tree = BoolNode::And(vec![index_leaf(1)]); // brand="apple" -> rows 4,5,6,7,13 + let (row_ids, batches) = query_then_fetch(tree, vec!["brand"]).await; + + assert_eq!(row_ids, vec![4, 5, 6, 7, 13]); + + let rows = extract_id_brand(&batches); + assert_eq!(rows.len(), 5); + assert_eq!(rows[0], (4, "apple".to_string())); + assert_eq!(rows[1], (5, "apple".to_string())); + assert_eq!(rows[2], (6, "apple".to_string())); + assert_eq!(rows[3], (7, "apple".to_string())); + assert_eq!(rows[4], (13, "apple".to_string())); + + // Verify schema: should be __row_id__ + brand (2 columns) + if let Some(b) = batches.first() { + assert_eq!(b.num_columns(), 2); + assert_eq!(b.schema().field(0).name(), "__row_id__"); + assert_eq!(b.schema().field(1).name(), "brand"); + } +} + +/// Filter that matches exactly 1 row -> fetch -> verify. +/// brand="amazon" AND price=30 matches only row 12. +#[tokio::test] +async fn test_qtf_fetch_single_row() { + let tree = BoolNode::And(vec![ + index_leaf(0), + pred_int("price", Operator::Eq, 30), + ]); + let (row_ids, batches) = query_then_fetch(tree, vec!["brand", "price"]).await; + + assert_eq!(row_ids, vec![12]); + + let rows = extract_id_brand_price(&batches); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0], (12, "amazon".to_string(), 30)); +} + +/// Filter that matches 0 rows -> empty result. +/// brand="amazon" AND price > 500 matches nothing. +#[tokio::test] +async fn test_qtf_fetch_empty_result() { + let tree = BoolNode::And(vec![ + index_leaf(0), + pred_int("price", Operator::Gt, 500), + ]); + let (row_ids, batches) = query_then_fetch(tree, vec!["brand", "price"]).await; + + assert!(row_ids.is_empty()); + assert!(batches.is_empty()); // fetch_phase returns empty vec for empty row_ids +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/row_id_emission.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/row_id_emission.rs new file mode 100644 index 0000000000000..517b4fc8f29d3 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/row_id_emission.rs @@ -0,0 +1,850 @@ +/* + * 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. + */ + +//! End-to-end tests for `emit_row_ids` mode. +//! Verifies that the indexed query path can return global row IDs +//! instead of actual data columns. + +use datafusion::arrow::array::Int64Array; + +use super::*; + +// Why it is complex --> +// Optimizer + ComputeRowId --> @gbh --> optimizer --> for parquet only right + +// Lot of query shapes --> for the query sent, how will that translate into at the data node side? + + +/// Helper: run a tree with `emit_row_ids: true` and return the collected row IDs. +async fn run_tree_row_ids(tree: BoolNode) -> Vec { + let tmp = write_fixture_parquet(); + let path = tmp.path().to_path_buf(); + let size = std::fs::metadata(&path).unwrap().len(); + + let file = std::fs::File::open(&path).unwrap(); + let meta = + ArrowReaderMetadata::load(&file, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let schema = meta.schema().clone(); + let parquet_meta = meta.metadata().clone(); + let mut rgs = Vec::new(); + let mut offset = 0i64; + for i in 0..parquet_meta.num_row_groups() { + let n = parquet_meta.row_group(i).num_rows(); + rgs.push(RowGroupInfo { + index: i, + first_row: offset, + num_rows: n, + }); + offset += n; + } + + let object_path = object_store::path::Path::from(path.to_string_lossy().as_ref()); + let segment = SegmentFileInfo { + segment_ord: 0, + max_doc: 16, + object_path, + parquet_size: size, + row_groups: rgs, + metadata: Arc::clone(&parquet_meta), + global_base: 0, + }; + + let tree = tree.push_not_down(); + let collectors = wire_collectors(&tree); + let per_leaf: Vec<(i32, Arc)> = collectors + .into_iter() + .enumerate() + .map(|(i, c)| (i as i32, c)) + .collect(); + let tree = Arc::new(tree); + let factory: super::super::table_provider::EvaluatorFactory = { + let per_leaf = per_leaf.clone(); + let tree = Arc::clone(&tree); + let schema = schema.clone(); + Arc::new(move |segment, _chunk, _stream_metrics| { + let resolved = tree.resolve(&per_leaf)?; + let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); + let eval: Arc = Arc::new(TreeBitsetSource { + tree: Arc::new(resolved), + evaluator: Arc::new(BitmapTreeEvaluator), + leaves: Arc::new( + crate::indexed_table::eval::bitmap_tree::CollectorLeafBitmaps { + ffm_collector_calls: _stream_metrics.ffm_collector_calls.clone(), + }, + ), + page_pruner: pruner, + cost_predicate: 1, + cost_collector: 10, + max_collector_parallelism: 1, + pruning_predicates: Arc::new(std::collections::HashMap::new()), + page_prune_metrics: Some( + crate::indexed_table::page_pruner::PagePruneMetrics::from_stream_metrics( + _stream_metrics, + ), + ), + collector_strategy: + crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + }); + Ok(eval) + }) + }; + + let store: Arc = + Arc::new(object_store::local::LocalFileSystem::new()); + let store_url = datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(); + let provider = Arc::new(IndexedTableProvider::new(IndexedTableConfig { + schema: schema.clone(), + segments: vec![segment], + store, + store_url, + evaluator_factory: factory, + pushdown_predicate: None, + query_config: Arc::new({ + let mut qc = crate::datafusion_query_config::DatafusionQueryConfig::test_default(); + qc.target_partitions = 1; + qc.force_strategy = Some(FilterStrategy::BooleanMask); + qc.force_pushdown = Some(false); + qc + }), + predicate_columns: vec![0, 1, 2, 3], + emit_row_ids: true, + })); + + let ctx = SessionContext::new(); + ctx.register_table("t", provider).unwrap(); + // Project __row_id__ — it will be computed from position, not read from parquet + let df = ctx.sql("SELECT \"__row_id__\" FROM t").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let task_ctx = ctx.task_ctx(); + let mut stream = datafusion::physical_plan::execute_stream(plan, task_ctx).unwrap(); + let mut row_ids: Vec = Vec::new(); + while let Some(batch) = stream.next().await { + let b = batch.unwrap(); + assert_eq!(b.num_columns(), 1, "should have only __row_id__ column"); + assert_eq!(b.schema().field(0).name(), "__row_id__"); + let col = b.column(0).as_any().downcast_ref::().unwrap(); + for i in 0..b.num_rows() { + row_ids.push(col.value(i)); + } + } + row_ids.sort(); + row_ids +} + +// ── Tests ──────────────────────────────────────────────────────────── + +/// brand="amazon" matches rows 0,1,2,3,12 — verify we get those row IDs back. +#[tokio::test] +async fn test_emit_row_ids_single_collector_amazon() { + // Collector tag 0 = brand_eq("amazon") + let tree = BoolNode::And(vec![index_leaf(0)]); + let ids = run_tree_row_ids(tree).await; + assert_eq!(ids, vec![0, 1, 2, 3, 12]); +} + +/// brand="apple" matches rows 4,5,6,7,13. +#[tokio::test] +async fn test_emit_row_ids_single_collector_apple() { + let tree = BoolNode::And(vec![index_leaf(1)]); + let ids = run_tree_row_ids(tree).await; + assert_eq!(ids, vec![4, 5, 6, 7, 13]); +} + +/// AND(brand="amazon", price > 100) matches rows where brand=amazon AND price>100. +/// amazon rows: 0(50), 1(150), 2(80), 3(120), 12(30) → price>100: rows 1, 3. +#[tokio::test] +async fn test_emit_row_ids_collector_and_predicate() { + let tree = BoolNode::And(vec![ + index_leaf(0), // amazon + pred_int("price", Operator::Gt, 100), + ]); + let ids = run_tree_row_ids(tree).await; + assert_eq!(ids, vec![1, 3]); +} + +/// OR(brand="amazon", brand="apple") matches rows 0-7, 12, 13. +#[tokio::test] +async fn test_emit_row_ids_or_two_collectors() { + let tree = BoolNode::Or(vec![index_leaf(0), index_leaf(1)]); + let ids = run_tree_row_ids(tree).await; + assert_eq!(ids, vec![0, 1, 2, 3, 4, 5, 6, 7, 12, 13]); +} + +/// AND(brand="apple", status="archived") — apple rows: 4,5,6,7,13; +/// archived rows: 1,5,9,12,13. Intersection: 5, 13. +#[tokio::test] +async fn test_emit_row_ids_two_collectors_and() { + let tree = BoolNode::And(vec![index_leaf(1), index_leaf(2)]); + let ids = run_tree_row_ids(tree).await; + assert_eq!(ids, vec![5, 13]); +} + +/// Verify global_base offset works: set global_base=1000 and check IDs are shifted. +async fn run_tree_row_ids_with_global_base(tree: BoolNode, global_base: u64) -> Vec { + let tmp = write_fixture_parquet(); + let path = tmp.path().to_path_buf(); + let size = std::fs::metadata(&path).unwrap().len(); + + let file = std::fs::File::open(&path).unwrap(); + let meta = + ArrowReaderMetadata::load(&file, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let schema = meta.schema().clone(); + let parquet_meta = meta.metadata().clone(); + let mut rgs = Vec::new(); + let mut offset = 0i64; + for i in 0..parquet_meta.num_row_groups() { + let n = parquet_meta.row_group(i).num_rows(); + rgs.push(RowGroupInfo { + index: i, + first_row: offset, + num_rows: n, + }); + offset += n; + } + + let object_path = object_store::path::Path::from(path.to_string_lossy().as_ref()); + let segment = SegmentFileInfo { + segment_ord: 0, + max_doc: 16, + object_path, + parquet_size: size, + row_groups: rgs, + metadata: Arc::clone(&parquet_meta), + global_base, + }; + + let tree = tree.push_not_down(); + let collectors = wire_collectors(&tree); + let per_leaf: Vec<(i32, Arc)> = collectors + .into_iter() + .enumerate() + .map(|(i, c)| (i as i32, c)) + .collect(); + let tree = Arc::new(tree); + let factory: super::super::table_provider::EvaluatorFactory = { + let per_leaf = per_leaf.clone(); + let tree = Arc::clone(&tree); + let schema = schema.clone(); + Arc::new(move |segment, _chunk, _stream_metrics| { + let resolved = tree.resolve(&per_leaf)?; + let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); + let eval: Arc = Arc::new(TreeBitsetSource { + tree: Arc::new(resolved), + evaluator: Arc::new(BitmapTreeEvaluator), + leaves: Arc::new( + crate::indexed_table::eval::bitmap_tree::CollectorLeafBitmaps { + ffm_collector_calls: _stream_metrics.ffm_collector_calls.clone(), + }, + ), + page_pruner: pruner, + cost_predicate: 1, + cost_collector: 10, + max_collector_parallelism: 1, + pruning_predicates: Arc::new(std::collections::HashMap::new()), + page_prune_metrics: Some( + crate::indexed_table::page_pruner::PagePruneMetrics::from_stream_metrics( + _stream_metrics, + ), + ), + collector_strategy: + crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + }); + Ok(eval) + }) + }; + + let store: Arc = + Arc::new(object_store::local::LocalFileSystem::new()); + let store_url = datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(); + let provider = Arc::new(IndexedTableProvider::new(IndexedTableConfig { + schema: schema.clone(), + segments: vec![segment], + store, + store_url, + evaluator_factory: factory, + pushdown_predicate: None, + query_config: Arc::new({ + let mut qc = crate::datafusion_query_config::DatafusionQueryConfig::test_default(); + qc.target_partitions = 1; + qc.force_strategy = Some(FilterStrategy::BooleanMask); + qc.force_pushdown = Some(false); + qc + }), + predicate_columns: vec![0, 1, 2, 3], + emit_row_ids: true, + })); + + let ctx = SessionContext::new(); + ctx.register_table("t", provider).unwrap(); + let df = ctx.sql("SELECT \"__row_id__\" FROM t").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let task_ctx = ctx.task_ctx(); + let mut stream = datafusion::physical_plan::execute_stream(plan, task_ctx).unwrap(); + let mut row_ids: Vec = Vec::new(); + while let Some(batch) = stream.next().await { + let b = batch.unwrap(); + assert_eq!(b.schema().field(0).name(), "__row_id__"); + let col = b.column(0).as_any().downcast_ref::().unwrap(); + for i in 0..b.num_rows() { + row_ids.push(col.value(i)); + } + } + row_ids.sort(); + row_ids +} + +/// With global_base=1000, brand="amazon" (rows 0,1,2,3,12) should give 1000,1001,1002,1003,1012. +#[tokio::test] +async fn test_emit_row_ids_with_global_base_offset() { + let tree = BoolNode::And(vec![index_leaf(0)]); + let ids = run_tree_row_ids_with_global_base(tree, 1000).await; + assert_eq!(ids, vec![1000, 1001, 1002, 1003, 1012]); +} + +// ── Comprehensive correctness: verify emitted row IDs match actual positions ── +// +// For each query type (SingleCollector, Collector+Predicate, Tree OR, Tree AND, +// Predicate-only), we compute the EXPECTED row IDs by scanning the fixture data +// directly, then compare against what emit_row_ids produces. +// +// The fixture (16 rows): +// | row | brand | price | status | category | +// | 0 | amazon | 50 | active | electronics | +// | 1 | amazon | 150 | archived | electronics | +// | 2 | amazon | 80 | active | books | +// | 3 | amazon | 120 | active | electronics | +// | 4 | apple | 90 | active | electronics | +// | 5 | apple | 95 | archived | electronics | +// | 6 | apple | 200 | active | books | +// | 7 | apple | 60 | active | electronics | +// | 8 | google | 40 | active | electronics | +// | 9 | google | 300 | archived | electronics | +// | 10 | samsung| 70 | active | electronics | +// | 11 | samsung| 150 | active | books | +// | 12 | amazon | 30 | archived | electronics | +// | 13 | apple | 45 | archived | electronics | +// | 14 | samsung| 99 | active | electronics | +// | 15 | google | 55 | active | electronics | + +/// Helper: compute expected row positions from fixture data given a filter. +fn expected_rows(filter: impl Fn(usize) -> bool) -> Vec { + (0..16).filter(|&i| filter(i)).map(|i| i as i64).collect() +} + +/// Exhaustive test: all query types produce correct row IDs matching fixture positions. +#[tokio::test] +async fn test_all_query_types_match_fixture_positions() { + // SingleCollector: brand="amazon" → rows 0,1,2,3,12 + let ids = run_tree_row_ids(BoolNode::And(vec![index_leaf(0)])).await; + let expected = expected_rows(|i| BRANDS[i] == "amazon"); + assert_eq!(ids, expected, "SingleCollector(amazon) mismatch"); + + // SingleCollector: brand="apple" → rows 4,5,6,7,13 + let ids = run_tree_row_ids(BoolNode::And(vec![index_leaf(1)])).await; + let expected = expected_rows(|i| BRANDS[i] == "apple"); + assert_eq!(ids, expected, "SingleCollector(apple) mismatch"); + + // SingleCollector: status="archived" → rows 1,5,9,12,13 + let ids = run_tree_row_ids(BoolNode::And(vec![index_leaf(2)])).await; + let expected = expected_rows(|i| STATUSES[i] == "archived"); + assert_eq!(ids, expected, "SingleCollector(archived) mismatch"); + + // Collector + Predicate: amazon AND price > 100 → rows 1,3 + let ids = run_tree_row_ids(BoolNode::And(vec![ + index_leaf(0), + pred_int("price", Operator::Gt, 100), + ])).await; + let expected = expected_rows(|i| BRANDS[i] == "amazon" && PRICES[i] > 100); + assert_eq!(ids, expected, "Collector+Predicate(amazon,price>100) mismatch"); + + // Collector + Predicate: apple AND price < 90 → rows 7,13 + let ids = run_tree_row_ids(BoolNode::And(vec![ + index_leaf(1), + pred_int("price", Operator::Lt, 90), + ])).await; + let expected = expected_rows(|i| BRANDS[i] == "apple" && PRICES[i] < 90); + assert_eq!(ids, expected, "Collector+Predicate(apple,price<90) mismatch"); + + // Tree OR: amazon OR apple → rows 0-7,12,13 + let ids = run_tree_row_ids(BoolNode::Or(vec![index_leaf(0), index_leaf(1)])).await; + let expected = expected_rows(|i| BRANDS[i] == "amazon" || BRANDS[i] == "apple"); + assert_eq!(ids, expected, "Tree OR(amazon,apple) mismatch"); + + // Tree AND: apple AND archived → rows 5,13 + let ids = run_tree_row_ids(BoolNode::And(vec![index_leaf(1), index_leaf(2)])).await; + let expected = expected_rows(|i| BRANDS[i] == "apple" && STATUSES[i] == "archived"); + assert_eq!(ids, expected, "Tree AND(apple,archived) mismatch"); + + // Tree OR + Predicate: (amazon OR apple) AND price > 100 → rows 1,3,6,9? no... + // amazon(0,1,2,3,12) OR apple(4,5,6,7,13) = 0-7,12,13. price>100: 1,3,6 + let ids = run_tree_row_ids(BoolNode::And(vec![ + BoolNode::Or(vec![index_leaf(0), index_leaf(1)]), + pred_int("price", Operator::Gt, 100), + ])).await; + let expected = expected_rows(|i| (BRANDS[i] == "amazon" || BRANDS[i] == "apple") && PRICES[i] > 100); + assert_eq!(ids, expected, "Tree OR+Predicate mismatch"); + + // Predicate only: price >= 150 → rows 1,6,9,11 + let ids = run_tree_row_ids(BoolNode::And(vec![ + pred_int("price", Operator::GtEq, 150), + ])).await; + let expected = expected_rows(|i| PRICES[i] >= 150); + assert_eq!(ids, expected, "Predicate-only(price>=150) mismatch"); + + // Predicate only: price < 50 → rows 8,12,13 + let ids = run_tree_row_ids(BoolNode::And(vec![ + pred_int("price", Operator::Lt, 50), + ])).await; + let expected = expected_rows(|i| PRICES[i] < 50); + assert_eq!(ids, expected, "Predicate-only(price<50) mismatch"); + + // Multi-predicate: price > 50 AND price < 100 → rows 2,4,5,7,10,14,15 + let ids = run_tree_row_ids(BoolNode::And(vec![ + pred_int("price", Operator::Gt, 50), + pred_int("price", Operator::Lt, 100), + ])).await; + let expected = expected_rows(|i| PRICES[i] > 50 && PRICES[i] < 100); + assert_eq!(ids, expected, "Multi-predicate(5080) OR (apple AND archived) + let ids = run_tree_row_ids(BoolNode::Or(vec![ + BoolNode::And(vec![index_leaf(0), pred_int("price", Operator::Gt, 80)]), + BoolNode::And(vec![index_leaf(1), index_leaf(2)]), + ])).await; + let expected = expected_rows(|i| { + (BRANDS[i] == "amazon" && PRICES[i] > 80) + || (BRANDS[i] == "apple" && STATUSES[i] == "archived") + }); + assert_eq!(ids, expected, "Complex OR(AND,AND) mismatch"); +} + +/// Verify that __row_id__ is computed alongside data columns — not replacing them. +/// Projects [__row_id__, brand, price] and verifies all three are present and correct. +#[tokio::test] +async fn test_row_id_with_data_columns() { + let tmp = write_fixture_parquet(); + let path = tmp.path().to_path_buf(); + let size = std::fs::metadata(&path).unwrap().len(); + + let file = std::fs::File::open(&path).unwrap(); + let meta = + ArrowReaderMetadata::load(&file, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let schema = meta.schema().clone(); + let parquet_meta = meta.metadata().clone(); + let mut rgs = Vec::new(); + let mut offset = 0i64; + for i in 0..parquet_meta.num_row_groups() { + let n = parquet_meta.row_group(i).num_rows(); + rgs.push(RowGroupInfo { + index: i, + first_row: offset, + num_rows: n, + }); + offset += n; + } + + let object_path = object_store::path::Path::from(path.to_string_lossy().as_ref()); + let segment = SegmentFileInfo { + segment_ord: 0, + max_doc: 16, + object_path, + parquet_size: size, + row_groups: rgs, + metadata: Arc::clone(&parquet_meta), + global_base: 0, + }; + + // Filter: brand = "amazon" (rows 0,1,2,3,12) + let tree = BoolNode::And(vec![index_leaf(0)]).push_not_down(); + let collectors = wire_collectors(&tree); + let per_leaf: Vec<(i32, Arc)> = collectors + .into_iter() + .enumerate() + .map(|(i, c)| (i as i32, c)) + .collect(); + let tree = Arc::new(tree); + let factory: super::super::table_provider::EvaluatorFactory = { + let per_leaf = per_leaf.clone(); + let tree = Arc::clone(&tree); + let schema = schema.clone(); + Arc::new(move |segment, _chunk, _stream_metrics| { + let resolved = tree.resolve(&per_leaf)?; + let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); + let eval: Arc = Arc::new(TreeBitsetSource { + tree: Arc::new(resolved), + evaluator: Arc::new(BitmapTreeEvaluator), + leaves: Arc::new( + crate::indexed_table::eval::bitmap_tree::CollectorLeafBitmaps { + ffm_collector_calls: _stream_metrics.ffm_collector_calls.clone(), + }, + ), + page_pruner: pruner, + cost_predicate: 1, + cost_collector: 10, + max_collector_parallelism: 1, + pruning_predicates: Arc::new(std::collections::HashMap::new()), + page_prune_metrics: Some( + crate::indexed_table::page_pruner::PagePruneMetrics::from_stream_metrics( + _stream_metrics, + ), + ), + collector_strategy: + crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + }); + Ok(eval) + }) + }; + + let store: Arc = + Arc::new(object_store::local::LocalFileSystem::new()); + let store_url = datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(); + let provider = Arc::new(IndexedTableProvider::new(IndexedTableConfig { + schema: schema.clone(), + segments: vec![segment], + store, + store_url, + evaluator_factory: factory, + pushdown_predicate: None, + query_config: Arc::new({ + let mut qc = crate::datafusion_query_config::DatafusionQueryConfig::test_default(); + qc.target_partitions = 1; + qc.force_strategy = Some(FilterStrategy::BooleanMask); + qc.force_pushdown = Some(false); + qc + }), + predicate_columns: vec![0, 1, 2, 3], + emit_row_ids: true, + })); + + let ctx = SessionContext::new(); + ctx.register_table("t", provider).unwrap(); + // Project __row_id__ alongside data columns + let df = ctx + .sql("SELECT \"__row_id__\", brand, price FROM t") + .await + .unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let task_ctx = ctx.task_ctx(); + let mut stream = datafusion::physical_plan::execute_stream(plan, task_ctx).unwrap(); + + let mut rows: Vec<(i64, String, i32)> = Vec::new(); + while let Some(batch) = stream.next().await { + let b = batch.unwrap(); + assert_eq!(b.num_columns(), 3, "should have 3 columns: __row_id__, brand, price"); + assert_eq!(b.schema().field(0).name(), "__row_id__"); + assert_eq!(b.schema().field(1).name(), "brand"); + assert_eq!(b.schema().field(2).name(), "price"); + let ids = b.column(0).as_any().downcast_ref::().unwrap(); + let brands = b.column(1).as_any().downcast_ref::().unwrap(); + let prices = b.column(2).as_any().downcast_ref::().unwrap(); + for i in 0..b.num_rows() { + rows.push((ids.value(i), brands.value(i).to_string(), prices.value(i))); + } + } + + rows.sort_by_key(|r| r.0); + + // brand="amazon" rows: 0,1,2,3,12 + assert_eq!(rows.len(), 5); + assert_eq!(rows[0], (0, "amazon".to_string(), 50)); + assert_eq!(rows[1], (1, "amazon".to_string(), 150)); + assert_eq!(rows[2], (2, "amazon".to_string(), 80)); + assert_eq!(rows[3], (3, "amazon".to_string(), 120)); + assert_eq!(rows[4], (12, "amazon".to_string(), 30)); +} + +/// Verify detection — `__row_id__` column in SELECT triggers emit_row_ids mode. +#[tokio::test] +async fn test_row_id_column_detection() { + use crate::indexed_table::substrait_to_tree::plan_requests_row_ids; + + let schema = build_fixture_schema(); + let row_ids: Vec = (0..16).collect(); + let batch = datafusion::arrow::record_batch::RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(datafusion::arrow::array::StringArray::from(BRANDS.to_vec())), + Arc::new(datafusion::arrow::array::Int32Array::from(PRICES.to_vec())), + Arc::new(datafusion::arrow::array::StringArray::from(STATUSES.to_vec())), + Arc::new(datafusion::arrow::array::StringArray::from(CATEGORIES.to_vec())), + Arc::new(datafusion::arrow::array::Int64Array::from(row_ids)), + ], + ) + .unwrap(); + + let ctx = SessionContext::new(); + let mem_table = datafusion::datasource::MemTable::try_new(schema.clone(), vec![vec![batch]]).unwrap(); + ctx.register_table("t", Arc::new(mem_table)).unwrap(); + + // Plan with __row_id__ in projection — should be detected + let df = ctx.sql("SELECT \"__row_id__\" FROM t").await.unwrap(); + let plan = df.logical_plan(); + assert!( + plan_requests_row_ids(plan), + "Should detect __row_id__ column in projection" + ); + + // Plan without __row_id__ — should NOT be detected + let df2 = ctx.sql("SELECT brand FROM t").await.unwrap(); + let plan2 = df2.logical_plan(); + assert!( + !plan_requests_row_ids(plan2), + "Should not detect __row_id__ in normal projection" + ); +} + +// ── Multi-segment row ID tests ────────────────────────────────────────────── + +/// Build fixture schema with `__row_id__` column name (double underscore prefix +/// and suffix) matching what `IndexedTableProvider.scan()` looks for. +fn build_row_id_fixture_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("brand", DataType::Utf8, false), + Field::new("price", DataType::Int32, false), + Field::new("status", DataType::Utf8, false), + Field::new("category", DataType::Utf8, false), + Field::new("__row_id__", DataType::Int64, false), + ])) +} + +/// Write fixture parquet with `__row_id__` column name for emit_row_ids tests. +fn write_row_id_fixture_parquet() -> NamedTempFile { + let schema = build_row_id_fixture_schema(); + let row_ids: Vec = (0..16).collect(); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(BRANDS.to_vec())), + Arc::new(Int32Array::from(PRICES.to_vec())), + Arc::new(StringArray::from(STATUSES.to_vec())), + Arc::new(StringArray::from(CATEGORIES.to_vec())), + Arc::new(Int64Array::from(row_ids)), + ], + ) + .unwrap(); + let tmp = NamedTempFile::new().unwrap(); + let props = datafusion::parquet::file::properties::WriterProperties::builder() + .set_max_row_group_size(8) + .set_statistics_enabled(datafusion::parquet::file::properties::EnabledStatistics::Page) + .build(); + let mut w = ArrowWriter::try_new(tmp.reopen().unwrap(), schema, Some(props)).unwrap(); + w.write(&batch).unwrap(); + w.close().unwrap(); + tmp +} + +/// Helper: run a query with `emit_row_ids=true` across two segments and return sorted row IDs. +/// Each segment is a separate parquet file with 16 rows. Segment 1 has global_base=0, +/// segment 2 has global_base=16, so combined IDs span 0..31. +async fn run_two_segments_row_ids(tree: BoolNode) -> Vec { + let tmp1 = write_row_id_fixture_parquet(); + let tmp2 = write_row_id_fixture_parquet(); + + let path1 = tmp1.path().to_path_buf(); + let path2 = tmp2.path().to_path_buf(); + let size1 = std::fs::metadata(&path1).unwrap().len(); + let size2 = std::fs::metadata(&path2).unwrap().len(); + + // Load metadata for segment 1 + let file1 = std::fs::File::open(&path1).unwrap(); + let meta1 = + ArrowReaderMetadata::load(&file1, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let schema = meta1.schema().clone(); + let parquet_meta1 = meta1.metadata().clone(); + let mut rgs1 = Vec::new(); + let mut offset = 0i64; + for i in 0..parquet_meta1.num_row_groups() { + let n = parquet_meta1.row_group(i).num_rows(); + rgs1.push(RowGroupInfo { + index: i, + first_row: offset, + num_rows: n, + }); + offset += n; + } + + // Load metadata for segment 2 + let file2 = std::fs::File::open(&path2).unwrap(); + let meta2 = + ArrowReaderMetadata::load(&file2, ArrowReaderOptions::new().with_page_index(true)).unwrap(); + let parquet_meta2 = meta2.metadata().clone(); + let mut rgs2 = Vec::new(); + let mut offset = 0i64; + for i in 0..parquet_meta2.num_row_groups() { + let n = parquet_meta2.row_group(i).num_rows(); + rgs2.push(RowGroupInfo { + index: i, + first_row: offset, + num_rows: n, + }); + offset += n; + } + + let object_path1 = object_store::path::Path::from(path1.to_string_lossy().as_ref()); + let object_path2 = object_store::path::Path::from(path2.to_string_lossy().as_ref()); + + let segment1 = SegmentFileInfo { + segment_ord: 0, + max_doc: 16, + object_path: object_path1, + parquet_size: size1, + row_groups: rgs1, + metadata: Arc::clone(&parquet_meta1), + global_base: 0, + }; + let segment2 = SegmentFileInfo { + segment_ord: 1, + max_doc: 16, + object_path: object_path2, + parquet_size: size2, + row_groups: rgs2, + metadata: Arc::clone(&parquet_meta2), + global_base: 16, + }; + + let tree = tree.push_not_down(); + let collectors = wire_collectors(&tree); + let per_leaf: Vec<(i32, Arc)> = collectors + .into_iter() + .enumerate() + .map(|(i, c)| (i as i32, c)) + .collect(); + let tree = Arc::new(tree); + let factory: super::super::table_provider::EvaluatorFactory = { + let per_leaf = per_leaf.clone(); + let tree = Arc::clone(&tree); + let schema = schema.clone(); + Arc::new(move |segment, _chunk, _stream_metrics| { + let resolved = tree.resolve(&per_leaf)?; + let pruner = Arc::new(PagePruner::new(&schema, Arc::clone(&segment.metadata))); + let eval: Arc = Arc::new(TreeBitsetSource { + tree: Arc::new(resolved), + evaluator: Arc::new(BitmapTreeEvaluator), + leaves: Arc::new( + crate::indexed_table::eval::bitmap_tree::CollectorLeafBitmaps { + ffm_collector_calls: _stream_metrics.ffm_collector_calls.clone(), + }, + ), + page_pruner: pruner, + cost_predicate: 1, + cost_collector: 10, + max_collector_parallelism: 1, + pruning_predicates: Arc::new(std::collections::HashMap::new()), + page_prune_metrics: Some( + crate::indexed_table::page_pruner::PagePruneMetrics::from_stream_metrics( + _stream_metrics, + ), + ), + collector_strategy: + crate::indexed_table::eval::CollectorCallStrategy::TightenOuterBounds, + }); + Ok(eval) + }) + }; + + let store: Arc = + Arc::new(object_store::local::LocalFileSystem::new()); + let store_url = datafusion::execution::object_store::ObjectStoreUrl::local_filesystem(); + let provider = Arc::new(IndexedTableProvider::new(IndexedTableConfig { + schema: schema.clone(), + segments: vec![segment1, segment2], + store, + store_url, + evaluator_factory: factory, + pushdown_predicate: None, + query_config: Arc::new({ + let mut qc = crate::datafusion_query_config::DatafusionQueryConfig::test_default(); + qc.target_partitions = 1; + qc.force_strategy = Some(FilterStrategy::BooleanMask); + qc.force_pushdown = Some(false); + qc + }), + predicate_columns: vec![0, 1, 2, 3], + emit_row_ids: true, + })); + + let ctx = SessionContext::new(); + ctx.register_table("t", provider).unwrap(); + let df = ctx.sql("SELECT \"__row_id__\" FROM t").await.unwrap(); + let plan = df.create_physical_plan().await.unwrap(); + let task_ctx = ctx.task_ctx(); + let mut stream = datafusion::physical_plan::execute_stream(plan, task_ctx).unwrap(); + let mut row_ids: Vec = Vec::new(); + while let Some(batch) = stream.next().await { + let b = batch.unwrap(); + assert_eq!(b.num_columns(), 1, "should have only __row_id__ column"); + assert_eq!(b.schema().field(0).name(), "__row_id__"); + let col = b.column(0).as_any().downcast_ref::().unwrap(); + for i in 0..b.num_rows() { + row_ids.push(col.value(i)); + } + } + row_ids.sort(); + row_ids +} + +/// Test: two segments produce globally unique row IDs. +/// Segment 1 has rows 0..15 (global_base=0), segment 2 has rows 0..15 (global_base=16). +/// Combined unfiltered query should produce IDs 0..31 with no gaps. +#[tokio::test] +async fn test_emit_row_ids_two_segments_global_base() { + // No filter — use a predicate that matches all rows (price >= 0) + let tree = BoolNode::And(vec![pred_int("price", Operator::GtEq, 0)]); + let ids = run_two_segments_row_ids(tree).await; + + // Each segment has 16 rows. With global_base=0 and global_base=16, + // we expect all IDs from 0 through 31 inclusive. + let expected: Vec = (0..32).collect(); + assert_eq!(ids.len(), 32, "should have 32 row IDs (16 per segment)"); + assert_eq!(ids, expected, "row IDs should cover 0..31 with no gaps"); + + // Verify uniqueness explicitly + let unique: std::collections::HashSet = ids.iter().copied().collect(); + assert_eq!(unique.len(), 32, "all 32 row IDs should be unique"); +} + +/// Test: filtered query across two segments returns correct global IDs. +/// brand="amazon" matches rows 0,1,2,3,12 within each segment. +/// Segment 1 (global_base=0) gives IDs: 0,1,2,3,12 +/// Segment 2 (global_base=16) gives IDs: 16,17,18,19,28 +#[tokio::test] +async fn test_emit_row_ids_two_segments_with_filter() { + // Collector tag 0 = brand_eq("amazon") — matches rows 0,1,2,3,12 + let tree = BoolNode::And(vec![index_leaf(0)]); + let ids = run_two_segments_row_ids(tree).await; + + // Segment 1: amazon rows at positions 0,1,2,3,12 + global_base 0 = 0,1,2,3,12 + // Segment 2: amazon rows at positions 0,1,2,3,12 + global_base 16 = 16,17,18,19,28 + let expected: Vec = vec![0, 1, 2, 3, 12, 16, 17, 18, 19, 28]; + assert_eq!( + ids.len(), + 10, + "should have 10 row IDs (5 amazon rows per segment)" + ); + assert_eq!(ids, expected, "filtered row IDs should be offset by global_base"); +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/row_id_strategies.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/row_id_strategies.rs new file mode 100644 index 0000000000000..6711b1c2eb2e5 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/row_id_strategies.rs @@ -0,0 +1,317 @@ +/* + * 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. + */ + +//! End-to-end correctness tests for row ID emission across all three strategies. +//! +//! These tests create real parquet files with known `___row_id` values and verify +//! that all three `FetchStrategy` variants produce identical shard-global row IDs. + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + use std::sync::Arc; + + use arrow::array::{Int32Array, Int64Array, StringArray}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow::record_batch::RecordBatch; + use datafusion::common::config::ConfigOptions; + use datafusion::physical_optimizer::PhysicalOptimizerRule; + use parquet::arrow::ArrowWriter; + use parquet::file::properties::WriterProperties; + use tempfile::TempDir; + + use crate::api::{build_shard_files, FileRowMetadata}; + use crate::datafusion_query_config::FetchStrategy; + use crate::project_row_id_optimizer::ProjectRowIdOptimizer; + + /// Create a test parquet file with `___row_id` column containing positional indices. + /// Returns the path to the created file and the number of rows. + fn create_test_parquet( + dir: &std::path::Path, + filename: &str, + num_rows: usize, + rows_per_rg: usize, + ) -> (PathBuf, Vec) { + let schema = Arc::new(Schema::new(vec![ + Field::new("__row_id__", DataType::Int64, false), + Field::new("value", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + ])); + + let path = dir.join(filename); + let file = std::fs::File::create(&path).unwrap(); + + let props = WriterProperties::builder() + .set_max_row_group_size(rows_per_rg) + .build(); + + let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props)).unwrap(); + + let mut rg_row_counts = Vec::new(); + let mut written = 0; + while written < num_rows { + let batch_size = (num_rows - written).min(rows_per_rg); + let row_ids: Vec = (written..written + batch_size).map(|i| i as i64).collect(); + let values: Vec = (written..written + batch_size).map(|i| (i * 10) as i32).collect(); + let names: Vec = (written..written + batch_size) + .map(|i| format!("row_{}", i)) + .collect(); + + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int64Array::from(row_ids)), + Arc::new(Int32Array::from(values)), + Arc::new(StringArray::from(names)), + ], + ) + .unwrap(); + + writer.write(&batch).unwrap(); + rg_row_counts.push(batch_size as u64); + written += batch_size; + } + + writer.close().unwrap(); + (path, rg_row_counts) + } + + /// Create a multi-file test shard with known row counts. + /// Returns (temp_dir, file_paths, file_metadata, total_rows). + fn create_test_shard() -> (TempDir, Vec, Vec, usize) { + let dir = TempDir::new().unwrap(); + + // File 1: 100 rows, 2 row groups of 50 + let (path1, rg_counts1) = create_test_parquet(dir.path(), "file1.parquet", 100, 50); + // File 2: 150 rows, 3 row groups of 50 + let (path2, rg_counts2) = create_test_parquet(dir.path(), "file2.parquet", 150, 50); + // File 3: 50 rows, 1 row group + let (path3, rg_counts3) = create_test_parquet(dir.path(), "file3.parquet", 50, 100); + + let file_metadata = vec![ + FileRowMetadata { row_group_row_counts: rg_counts1 }, + FileRowMetadata { row_group_row_counts: rg_counts2 }, + FileRowMetadata { row_group_row_counts: rg_counts3 }, + ]; + + let total_rows = 100 + 150 + 50; + (dir, vec![path1, path2, path3], file_metadata, total_rows) + } + + // ── Task 12.1: Test parquet fixtures ────────────────────────────── + + #[test] + fn test_parquet_fixtures_created_correctly() { + let (dir, paths, metadata, total_rows) = create_test_shard(); + + assert_eq!(paths.len(), 3); + assert_eq!(metadata.len(), 3); + assert_eq!(total_rows, 300); + + // Verify file 1 + assert_eq!(metadata[0].row_group_row_counts, vec![50, 50]); + // Verify file 2 + assert_eq!(metadata[1].row_group_row_counts, vec![50, 50, 50]); + // Verify file 3 + assert_eq!(metadata[2].row_group_row_counts, vec![50]); + + // Verify files exist + for path in &paths { + assert!(path.exists(), "File {:?} should exist", path); + } + + drop(dir); // cleanup + } + + // ── Task 12.2: Row base computation correctness ─────────────────── + + #[test] + fn test_row_base_prefix_sum() { + let (_dir, paths, metadata, _total_rows) = create_test_shard(); + + // Build object metas (simplified for test) + let object_metas: Vec = paths + .iter() + .map(|p| object_store::ObjectMeta { + location: object_store::path::Path::from(p.to_str().unwrap()), + last_modified: chrono::Utc::now(), + size: std::fs::metadata(p).unwrap().len() as u64, + e_tag: None, + version: None, + }) + .collect(); + + let shard_files = build_shard_files(&object_metas, &metadata); + + // Verify row_base values + assert_eq!(shard_files[0].row_base, 0); // First file starts at 0 + assert_eq!(shard_files[1].row_base, 100); // Second file starts at 100 + assert_eq!(shard_files[2].row_base, 250); // Third file starts at 250 + + // Verify num_rows + assert_eq!(shard_files[0].num_rows, 100); + assert_eq!(shard_files[1].num_rows, 150); + assert_eq!(shard_files[2].num_rows, 50); + } + + // ── Task 12.3: Row ID uniqueness across shard ───────────────────── + + #[test] + fn test_global_row_ids_unique_and_contiguous() { + let (_dir, paths, metadata, total_rows) = create_test_shard(); + + let object_metas: Vec = paths + .iter() + .map(|p| object_store::ObjectMeta { + location: object_store::path::Path::from(p.to_str().unwrap()), + last_modified: chrono::Utc::now(), + size: std::fs::metadata(p).unwrap().len() as u64, + e_tag: None, + version: None, + }) + .collect(); + + let shard_files = build_shard_files(&object_metas, &metadata); + + // Compute all global row IDs + let mut all_ids: Vec = Vec::new(); + for file in &shard_files { + for local_id in 0..file.num_rows as i64 { + all_ids.push(file.row_base + local_id); + } + } + + // Verify uniqueness + let mut sorted = all_ids.clone(); + sorted.sort(); + sorted.dedup(); + assert_eq!(sorted.len(), total_rows, "All row IDs should be unique"); + + // Verify contiguous from 0 + assert_eq!(sorted[0], 0); + assert_eq!(*sorted.last().unwrap(), (total_rows - 1) as i64); + for i in 0..sorted.len() { + assert_eq!(sorted[i], i as i64, "Row IDs should be contiguous"); + } + } + + // ── Task 12.4: Multi-file shard with varying row group sizes ────── + + #[test] + fn test_varying_row_group_sizes() { + let dir = TempDir::new().unwrap(); + + // File 1: 30 rows, 1 RG of 30 + let (_path1, rg1) = create_test_parquet(dir.path(), "small.parquet", 30, 100); + // File 2: 500 rows, 5 RGs of 100 + let (_path2, rg2) = create_test_parquet(dir.path(), "large.parquet", 500, 100); + // File 3: 7 rows, 1 RG of 7 + let (_path3, rg3) = create_test_parquet(dir.path(), "tiny.parquet", 7, 100); + + let metadata = vec![ + FileRowMetadata { row_group_row_counts: rg1 }, + FileRowMetadata { row_group_row_counts: rg2 }, + FileRowMetadata { row_group_row_counts: rg3 }, + ]; + + let object_metas: Vec = vec![ + object_store::ObjectMeta { + location: object_store::path::Path::from("small.parquet"), + last_modified: chrono::Utc::now(), + size: 1000, + e_tag: None, + version: None, + }, + object_store::ObjectMeta { + location: object_store::path::Path::from("large.parquet"), + last_modified: chrono::Utc::now(), + size: 5000, + e_tag: None, + version: None, + }, + object_store::ObjectMeta { + location: object_store::path::Path::from("tiny.parquet"), + last_modified: chrono::Utc::now(), + size: 500, + e_tag: None, + version: None, + }, + ]; + + let shard_files = build_shard_files(&object_metas, &metadata); + + // Verify row_base computation + assert_eq!(shard_files[0].row_base, 0); + assert_eq!(shard_files[1].row_base, 30); + assert_eq!(shard_files[2].row_base, 530); + + // Verify total coverage + let total: u64 = shard_files.iter().map(|f| f.num_rows).sum(); + assert_eq!(total, 537); + + // Verify no gaps in global IDs + let last_file = shard_files.last().unwrap(); + let max_global_id = last_file.row_base + last_file.num_rows as i64 - 1; + assert_eq!(max_global_id, 536); + } + + // ── Optimizer correctness tests ─────────────────────────────────── + + #[test] + fn test_project_row_id_optimizer_no_op_without_row_id() { + // Schema without ___row_id — optimizer should be a no-op + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, false), + ])); + + let empty = datafusion::physical_plan::empty::EmptyExec::new(schema.clone()); + let plan: Arc = Arc::new(empty); + + let optimizer = ProjectRowIdOptimizer; + let config = ConfigOptions::default(); + let result = optimizer.optimize(plan.clone(), &config).unwrap(); + + // Schema should be unchanged + assert_eq!(result.schema().fields().len(), 2); + assert_eq!(result.schema().field(0).name(), "a"); + assert_eq!(result.schema().field(1).name(), "b"); + } + + #[test] + fn test_fetch_strategy_default_is_none() { + let config = crate::datafusion_query_config::DatafusionQueryConfig::test_default(); + assert_eq!(config.fetch_strategy, FetchStrategy::None); + } + + + #[test] + fn test_build_shard_files_empty() { + let shard_files = build_shard_files(&[], &[]); + assert!(shard_files.is_empty()); + } + + #[test] + fn test_build_shard_files_single_file() { + let meta = object_store::ObjectMeta { + location: object_store::path::Path::from("test.parquet"), + last_modified: chrono::Utc::now(), + size: 1000, + e_tag: None, + version: None, + }; + let fm = FileRowMetadata { + row_group_row_counts: vec![100, 200, 300], + }; + + let shard_files = build_shard_files(&[meta], &[fm]); + assert_eq!(shard_files.len(), 1); + assert_eq!(shard_files[0].row_base, 0); + assert_eq!(shard_files[0].num_rows, 600); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs index d280e1b16db5c..4aae95d133075 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/schema_drift.rs @@ -96,6 +96,7 @@ async fn run_missing_col_tree(tree_bool: BoolNode) -> usize { parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }; let tree = Arc::new(tree_bool); @@ -135,6 +136,7 @@ async fn run_missing_col_tree(tree_bool: BoolNode) -> usize { pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], + emit_row_ids: false, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/streaming_at_scale.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/streaming_at_scale.rs index 2424c37cfc0ae..d5030c7fdc486 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/streaming_at_scale.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/streaming_at_scale.rs @@ -405,6 +405,7 @@ async fn run_large( parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }; let tree = Arc::new(tree); @@ -451,6 +452,7 @@ async fn run_large( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], + emit_row_ids: false, })); let ctx = SessionContext::new(); @@ -855,6 +857,7 @@ async fn run_large_partitioned( parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }; let tree = Arc::new(tree); @@ -900,6 +903,7 @@ async fn run_large_partitioned( pushdown_predicate: None, query_config: std::sync::Arc::new(qc), predicate_columns: vec![], + emit_row_ids: false, })); let ctx = SessionContext::new(); ctx.register_table("t", provider).unwrap(); diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs index fb522332de9a4..bcc695b345a02 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs @@ -11,6 +11,10 @@ //! The bridge-agnostic API lives in [`api`]. The FFM bridge (`ffm.rs`) exports //! `extern "C"` functions for JDK FFM. +/// Column name for the shard-global row identifier used by Query-Then-Fetch. +/// Stored as Int64 in parquet, computed from position in the indexed path. +pub const ROW_ID_COLUMN_NAME: &str = "__row_id__"; + pub(crate) mod agg_mode; pub mod api; pub mod cache; @@ -27,12 +31,16 @@ pub mod io; pub mod local_executor; pub mod memory; pub mod partition_stream; +pub mod project_row_id_analyzer; +pub mod project_row_id_optimizer; pub mod query_executor; pub mod query_tracker; +pub mod shard_table_provider; pub mod runtime_manager; pub mod schema_coerce; pub mod session_context; pub mod statistics_cache; pub mod udf; pub mod stats; -pub mod task_monitors; \ No newline at end of file +pub mod task_monitors; + diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/project_row_id_analyzer.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/project_row_id_analyzer.rs new file mode 100644 index 0000000000000..a712ac09caa88 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/project_row_id_analyzer.rs @@ -0,0 +1,175 @@ +/* + * 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. + */ + +//! Project Row ID Analyzer +//! +//! Logical-level analyzer rule that ensures `__row_id__` is always projected through +//! the query plan when available in the source schema. Runs BEFORE physical optimization, +//! so the column survives all pruning passes. +//! +//! For the ListingTable QTF path: `ShardTableProvider` exposes `row_base` as a partition +//! column. This analyzer ensures `__row_id__` (and implicitly `row_base`) survive into +//! the physical plan where `ProjectRowIdOptimizer` computes `__row_id__ + row_base`. + +use std::sync::Arc; + +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{Column, DFSchema}; +use datafusion::config::ConfigOptions; +use datafusion::error::Result; +use datafusion::logical_expr::{col, Expr, LogicalPlan}; +use datafusion::optimizer::AnalyzerRule; +use datafusion_expr::Projection; + +pub use crate::ROW_ID_COLUMN_NAME as ROW_ID_FIELD_NAME; + +#[derive(Debug)] +pub struct ProjectRowIdAnalyzer; + +impl ProjectRowIdAnalyzer { + pub fn new() -> Self { + Self {} + } +} + +impl AnalyzerRule for ProjectRowIdAnalyzer { + fn analyze(&self, plan: LogicalPlan, _config: &ConfigOptions) -> Result { + let rewritten = plan.transform_up(|node| { + match &node { + LogicalPlan::TableScan(scan) => { + let mut proj = scan.projection.clone().unwrap_or_else(|| { + (0..scan.projected_schema.fields().len()).collect() + }); + + let mut new_projected_schema = (*scan.projected_schema).clone(); + + if scan.source.schema().index_of(ROW_ID_FIELD_NAME).is_ok() { + let row_id_idx = + scan.source.schema().index_of(ROW_ID_FIELD_NAME).unwrap(); + + if !proj.contains(&row_id_idx) { + proj.push(row_id_idx); + + let qualifier = scan + .projected_schema + .qualified_field(0) + .0 + .cloned(); + + if let Some(q) = qualifier { + let row_id_schema = DFSchema::try_from_qualified_schema( + q, + &Schema::new(vec![Field::new( + ROW_ID_FIELD_NAME, + DataType::Int64, + false, + )]), + )?; + new_projected_schema = new_projected_schema + .join(&row_id_schema) + .map_err(|e| { + datafusion::error::DataFusionError::Internal(format!( + "ProjectRowIdAnalyzer: join schema: {}", + e + )) + })?; + } + } + } + + let new_scan = LogicalPlan::TableScan(datafusion_expr::TableScan { + table_name: scan.table_name.clone(), + source: scan.source.clone(), + projection: Some(proj), + projected_schema: Arc::new(new_projected_schema), + filters: scan.filters.clone(), + fetch: scan.fetch, + }); + Ok(Transformed::yes(new_scan)) + } + + LogicalPlan::Projection(p) => { + let already_has_row_id = p.expr.iter().any(|e| { + matches!(e, Expr::Column(c) if c.name == ROW_ID_FIELD_NAME) + }); + let input_has_row_id = p + .input + .schema() + .index_of_column(&Column::from_name(ROW_ID_FIELD_NAME)) + .is_ok(); + + if !already_has_row_id && input_has_row_id { + let mut new_exprs = p.expr.to_vec(); + new_exprs.push(col(ROW_ID_FIELD_NAME)); + + let row_id_schema = { + let qualifier = p.schema.qualified_field(0).0.cloned(); + match qualifier { + Some(q) => DFSchema::try_from_qualified_schema( + q, + &Schema::new(vec![Field::new( + ROW_ID_FIELD_NAME, + DataType::Int64, + false, + )]), + )?, + None => DFSchema::try_from(Schema::new(vec![Field::new( + ROW_ID_FIELD_NAME, + DataType::Int64, + false, + )]))?, + } + }; + + let merged_schema = if p + .schema + .index_of_column(&Column::from_name(ROW_ID_FIELD_NAME)) + .is_ok() + { + p.schema.clone() + } else { + Arc::new(p.schema.as_ref().clone().join(&row_id_schema).map_err( + |e| { + datafusion::error::DataFusionError::Internal(format!( + "ProjectRowIdAnalyzer: join projection schema: {}", + e + )) + }, + )?) + }; + + let new_proj = LogicalPlan::Projection( + Projection::try_new_with_schema( + new_exprs, + p.input.clone(), + merged_schema, + ) + .map_err(|e| { + datafusion::error::DataFusionError::Internal(format!( + "ProjectRowIdAnalyzer: create projection: {}", + e + )) + })?, + ); + return Ok(Transformed::yes(new_proj)); + } + Ok(Transformed::no(node)) + } + + _ => Ok(Transformed::no(node)), + } + })?; + + Ok(rewritten.data) + } + + fn name(&self) -> &str { + "project_row_id_analyzer" + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/project_row_id_optimizer.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/project_row_id_optimizer.rs new file mode 100644 index 0000000000000..3557992c65921 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/project_row_id_optimizer.rs @@ -0,0 +1,178 @@ +/* + * 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. + */ + +//! Physical optimizer rule for the ListingTable QTF path. +//! +//! Downcasts to `DataSourceExec` → `FileScanConfig`, checks if `__row_id__` is in the +//! file schema, rebuilds the scan to include both `__row_id__` and `row_base` (partition +//! column), then wraps with `ProjectionExec` computing `__row_id__ + row_base`. +//! +//! This approach is resilient to DataFusion's built-in optimizers pruning `row_base` +//! from the scan output — we rebuild the DataSourceExec from the underlying config. + +use std::sync::Arc; + +use datafusion::common::config::ConfigOptions; +use datafusion::common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; +use datafusion::common::Result; +use datafusion::datasource::physical_plan::ParquetSource; +use datafusion::datasource::source::DataSourceExec; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::{BinaryExpr, Column}; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::projection::ProjectionExec; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use datafusion_datasource::TableSchema; + +pub use crate::ROW_ID_COLUMN_NAME as ROW_ID_FIELD_NAME; +pub const ROW_BASE_FIELD_NAME: &str = "row_base"; + +#[derive(Debug)] +pub struct ProjectRowIdOptimizer; + +impl PhysicalOptimizerRule for ProjectRowIdOptimizer { + fn optimize( + &self, + plan: Arc, + _config: &ConfigOptions, + ) -> Result> { + let rewritten = plan.transform_up(|node| { + // Only handle DataSourceExec nodes backed by FileScanConfig + let Some(datasource_exec) = node.as_any().downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(file_scan_config) = datasource_exec + .data_source() + .as_ref() + .as_any() + .downcast_ref::() + else { + return Ok(Transformed::no(node)); + }; + + // Check if __row_id__ exists in the file schema + let file_schema = file_scan_config.file_schema(); + if file_schema.index_of(ROW_ID_FIELD_NAME).is_err() { + return Ok(Transformed::no(node)); + } + + // Check if this DataSourceExec has partition columns (ShardTableProvider sets row_base) + let partition_cols = file_scan_config.table_partition_cols(); + if partition_cols.is_empty() { + return Ok(Transformed::no(node)); + } + + // Rebuild the DataSourceExec to ensure both __row_id__ and row_base are in the output. + // Get the current output schema fields to determine what's projected. + let current_schema = datasource_exec.schema(); + + // Build new projection indices: current projected file columns + __row_id__ + row_base + let mut new_proj_indices: Vec = Vec::new(); + for field in current_schema.fields() { + if field.name() == ROW_BASE_FIELD_NAME { + // row_base is partition col, it'll be added below + continue; + } + if let Ok(idx) = file_schema.index_of(field.name()) { + if !new_proj_indices.contains(&idx) { + new_proj_indices.push(idx); + } + } + } + + // Ensure __row_id__ is in the projection + let row_id_file_idx = file_schema.index_of(ROW_ID_FIELD_NAME).unwrap(); + if !new_proj_indices.contains(&row_id_file_idx) { + new_proj_indices.push(row_id_file_idx); + } + + // Add partition column index (row_base is at file_schema.fields().len()) + let row_base_partition_idx = file_schema.fields().len(); + new_proj_indices.push(row_base_partition_idx); + + // Rebuild the DataSourceExec with new projections + let new_table_schema = TableSchema::new( + file_schema.clone(), + partition_cols.clone(), + ); + let new_source = Arc::new(ParquetSource::new(new_table_schema)); + + let new_config = FileScanConfigBuilder::from(file_scan_config.clone()) + .with_source(new_source) + .with_projection_indices(Some(new_proj_indices)) + .map_err(|e| datafusion::error::DataFusionError::Internal( + format!("ProjectRowIdOptimizer: set projection: {}", e) + ))? + .build(); + + let new_datasource: Arc = + DataSourceExec::from_data_source(new_config); + + // Build projection: __row_id__ + row_base, drop row_base from output + let new_schema = new_datasource.schema(); + let row_id_idx = new_schema.index_of(ROW_ID_FIELD_NAME).unwrap(); + let row_base_idx = new_schema.index_of(ROW_BASE_FIELD_NAME).unwrap(); + + let sum_expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new(ROW_ID_FIELD_NAME, row_id_idx)), + Operator::Plus, + Arc::new(Column::new(ROW_BASE_FIELD_NAME, row_base_idx)), + )); + + let mut projection_exprs: Vec<(Arc, String)> = Vec::new(); + for (i, field) in new_schema.fields().iter().enumerate() { + if field.name() == ROW_ID_FIELD_NAME { + projection_exprs.push((sum_expr.clone(), ROW_ID_FIELD_NAME.to_string())); + } else if field.name() == ROW_BASE_FIELD_NAME { + continue; // drop from output + } else { + projection_exprs.push(( + Arc::new(Column::new(field.name(), i)), + field.name().clone(), + )); + } + } + + let projection = ProjectionExec::try_new(projection_exprs, new_datasource)?; + Ok(Transformed::new( + Arc::new(projection) as Arc, + true, + TreeNodeRecursion::Continue, + )) + })?; + + Ok(rewritten.data) + } + + fn name(&self) -> &str { + "ProjectRowIdOptimizer" + } + + fn schema_check(&self) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn optimizer_name() { + let opt = ProjectRowIdOptimizer; + assert_eq!(opt.name(), "ProjectRowIdOptimizer"); + } + + #[test] + fn optimizer_schema_check_disabled() { + let opt = ProjectRowIdOptimizer; + assert!(!opt.schema_check()); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs index 869a8e42e6805..82377709a231b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs @@ -31,7 +31,7 @@ use substrait::proto::Plan; use crate::cross_rt_stream::CrossRtStream; use crate::executor::DedicatedExecutor; -use crate::api::DataFusionRuntime; +use crate::api::{DataFusionRuntime, ShardFileInfo}; use crate::session_context::SessionContextHandle; /// Execute a vanilla parquet query: substrait plan → DataFusion → CrossRtStream. @@ -56,40 +56,24 @@ pub async fn execute_query( query_config: &crate::datafusion_query_config::DatafusionQueryConfig, shard_store: Arc, ) -> Result { - // Pre-populate the list-files cache so DataFusion doesn't re-list the directory - let list_file_cache = Arc::new(DefaultListFilesCache::default()); - let table_scoped_path = datafusion::execution::cache::TableScopedPath { - table: None, - path: table_path.prefix().clone(), - }; - list_file_cache.put(&table_scoped_path, CachedFileList::new(object_metas.as_ref().clone())); - - // Build a per-query RuntimeEnv sharing the global memory pool + caches, - // but with a fresh list-files cache for this query's shard files. - let mut runtime_env_builder = RuntimeEnvBuilder::from_runtime_env(&runtime.runtime_env) - .with_cache_manager( - CacheManagerConfig::default() - .with_list_files_cache(Some(list_file_cache)) - .with_file_metadata_cache(Some( - runtime.runtime_env.cache_manager.get_file_metadata_cache(), - )) - .with_files_statistics_cache( - runtime.runtime_env.cache_manager.get_file_statistic_cache(), - ), - ); + // Build per-query RuntimeEnv with list-files cache pre-populated. + let runtime_env = build_query_runtime_env(runtime, &table_path, object_metas.as_ref())?; - // If a per-query memory pool is provided, set it on the same builder. + // If a per-query memory pool is provided, rebuild with it overlaid. // The per-query pool wraps the global pool, so global limits are still enforced. - if let Some(pool) = query_memory_pool { - runtime_env_builder = runtime_env_builder.with_memory_pool(pool); - } - - let runtime_env = runtime_env_builder - .build() - .map_err(|e| { - error!("Failed to build runtime env: {}", e); - e - })?; + let runtime_env = if let Some(pool) = query_memory_pool { + Arc::from( + RuntimeEnvBuilder::from_runtime_env(&runtime_env) + .with_memory_pool(pool) + .build() + .map_err(|e| { + error!("Failed to build runtime env with query pool: {}", e); + e + })?, + ) + } else { + runtime_env + }; // Register shard-specific object store on file:// scheme for this query. // Routes reads through TieredObjectStore (local + remote) or default LocalFileSystem. @@ -106,41 +90,71 @@ pub async fn execute_query( let state = SessionStateBuilder::new() .with_config(config) - .with_runtime_env(Arc::from(runtime_env)) + .with_runtime_env(runtime_env) .with_default_features() .build(); let ctx = SessionContext::new_with_state(state); crate::udf::register_all(&ctx); - // Register table via ListingTable — all IO goes through object store - let file_format = ParquetFormat::new(); - let listing_options = ListingOptions::new(Arc::new(file_format)) - .with_file_extension(".parquet") - .with_collect_stat(true); - - let resolved_schema = listing_options - .infer_schema(&ctx.state(), &table_path) - .await - .map_err(|e| { - error!("Failed to infer schema: {}", e); - e - })?; - let resolved_schema = crate::schema_coerce::coerce_inferred_schema(resolved_schema); - - let table_config = ListingTableConfig::new(table_path) - .with_listing_options(listing_options) - .with_schema(resolved_schema); - - let provider = Arc::new(ListingTable::try_new(table_config).map_err(|e| { - error!("Failed to create listing table: {}", e); - e - })?); + // Register table provider based on strategy + use crate::datafusion_query_config::FetchStrategy; + match query_config.fetch_strategy { + FetchStrategy::IndexedPredicateOnly => { + return Err(DataFusionError::Execution( + "IndexedPredicateOnly strategy requires the indexed executor path \ + (execute_indexed_with_context). It cannot be used via execute_query \ + because it needs segment metadata and PositionMap for position-based \ + row ID computation.".into(), + )); + } + FetchStrategy::ListingTable => { + use crate::shard_table_provider::{ShardTableConfig, ShardTableProvider}; - ctx.register_table(&table_name, provider).map_err(|e| { - error!("Failed to register table: {}", e); - e - })?; + // Infer schema from the first file + let file_format = ParquetFormat::new(); + let listing_options = ListingOptions::new(Arc::new(file_format)) + .with_file_extension(".parquet") + .with_collect_stat(true); + let resolved_schema = listing_options + .infer_schema(&ctx.state(), &table_path) + .await + .map_err(|e| { error!("Failed to infer schema: {}", e); e })?; + let resolved_schema = crate::schema_coerce::coerce_inferred_schema(resolved_schema); + + // Build ShardFileInfo with row_base from cumulative row counts + let store = ctx.state().runtime_env().object_store(&table_path)?; + let files = build_shard_file_infos(&store, object_metas.as_ref()).await?; + + let store_url = store_url_from_table_path(&table_path)?; + + let provider = Arc::new(ShardTableProvider::new(ShardTableConfig { + file_schema: resolved_schema, + files, + store_url, + })); + ctx.register_table(&table_name, provider) + .map_err(|e| { error!("Failed to register table: {}", e); e })?; + } + _ => { + // Baseline / IndexedPredicateOnly: use standard ListingTable + let file_format = ParquetFormat::new(); + let listing_options = ListingOptions::new(Arc::new(file_format)) + .with_file_extension(".parquet") + .with_collect_stat(true); + let resolved_schema = listing_options + .infer_schema(&ctx.state(), &table_path) + .await + .map_err(|e| { error!("Failed to infer schema: {}", e); e })?; + let table_config = ListingTableConfig::new(table_path) + .with_listing_options(listing_options) + .with_schema(resolved_schema); + let provider = Arc::new(ListingTable::try_new(table_config) + .map_err(|e| { error!("Failed to create listing table: {}", e); e })?); + ctx.register_table(&table_name, provider) + .map_err(|e| { error!("Failed to register table: {}", e); e })?; + } + } // Decode substrait → logical plan → physical plan → stream let substrait_plan = Plan::decode(plan_bytes.as_slice()).map_err(|e| { @@ -151,6 +165,25 @@ pub async fn execute_query( let dataframe = ctx.execute_logical_plan(logical_plan).await?; let physical_plan = dataframe.create_physical_plan().await?; + // Apply row ID optimizer based on strategy if configured + use datafusion::physical_optimizer::PhysicalOptimizerRule; + let physical_plan = match query_config.fetch_strategy { + FetchStrategy::None => { + // Baseline: no optimizer, ___row_id read as regular column (local IDs) + physical_plan + } + FetchStrategy::ListingTable => { + // Apply ProjectRowIdOptimizer: rewrites ___row_id to ___row_id + row_base + let optimizer = crate::project_row_id_optimizer::ProjectRowIdOptimizer; + let config = datafusion::common::config::ConfigOptions::default(); + optimizer.optimize(physical_plan, &config)? + } + FetchStrategy::IndexedPredicateOnly => { + // Unreachable — we return an error above before reaching here + unreachable!() + } + }; + let df_stream = execute_stream(physical_plan, ctx.task_ctx()).map_err(|e| { error!("Failed to create execution stream: {}", e); e @@ -173,21 +206,77 @@ pub async fn execute_query( /// raw Java pointer) happens at the FFM entry in `df_execute_with_context`, so /// by the time this function is reached the pointer is already invalidated from /// Java's perspective and cleanup is pure RAII. +/// +/// When the plan requests row IDs and a `FetchStrategy` is configured, this +/// function routes to the appropriate execution path: +/// - `ListingTable`: applies `ProjectRowIdOptimizer` to the physical plan +/// - `IndexedPredicateOnly`: delegates to the indexed executor with `emit_row_ids=true` pub async fn execute_with_context( handle: SessionContextHandle, plan_bytes: &[u8], cpu_executor: DedicatedExecutor, ) -> Result { + use crate::datafusion_query_config::FetchStrategy; + + let fetch_strategy = handle.query_config.fetch_strategy; + + // If ListingTable strategy: replace the default ListingTable with ShardTableProvider + // that adds row_base partition column for ProjectRowIdOptimizer. + // Also register the ProjectRowIdAnalyzer to ensure __row_id__ survives logical optimization. + if fetch_strategy == FetchStrategy::ListingTable { + use crate::shard_table_provider::{ShardTableConfig, ShardTableProvider}; + + handle.ctx.deregister_table(&handle.table_name)?; + + let store = handle.ctx.state().runtime_env().object_store(&handle.table_path)?; + + // Infer schema from existing files + let listing_options = ListingOptions::new(Arc::new(ParquetFormat::new())) + .with_file_extension(".parquet") + .with_collect_stat(true); + let resolved_schema = listing_options + .infer_schema(&handle.ctx.state(), &handle.table_path) + .await?; + + // Build ShardFileInfo with cumulative row_base from parquet metadata. + let files = build_shard_file_infos(&store, handle.object_metas.as_ref()).await?; + + let store_url = store_url_from_table_path(&handle.table_path)?; + + let provider = Arc::new(ShardTableProvider::new(ShardTableConfig { + file_schema: resolved_schema, + files, + store_url, + })); + handle.ctx.register_table(&handle.table_name, provider)?; + } + let substrait_plan = Plan::decode(plan_bytes).map_err(|e| { DataFusionError::Execution(format!("Failed to decode Substrait: {}", e)) })?; let logical_plan = from_substrait_plan(&handle.ctx.state(), &substrait_plan).await?; log_debug!("DataFusion logical plan:\n{}", logical_plan.display_indent()); + + let requests_row_ids = crate::indexed_table::substrait_to_tree::plan_requests_row_ids(&logical_plan); + let dataframe = handle.ctx.execute_logical_plan(logical_plan).await?; + // create_physical_plan runs all registered physical optimizer rules including + // ProjectRowIdOptimizer (registered in session_context when strategy=ListingTable). let physical_plan = dataframe.create_physical_plan().await?; log_debug!("DataFusion physical plan:\n{}", displayable(physical_plan.as_ref()).indent(true)); + // No post-hoc optimizer needed — ProjectRowIdOptimizer runs inside create_physical_plan. + let physical_plan = if requests_row_ids { + match fetch_strategy { + FetchStrategy::None => physical_plan, + FetchStrategy::ListingTable => physical_plan, + FetchStrategy::IndexedPredicateOnly => physical_plan, + } + } else { + physical_plan + }; + let df_stream = execute_stream(physical_plan, handle.ctx.task_ctx()).map_err(|e| { error!("execute_with_context: failed to create stream: {}", e); e @@ -202,3 +291,97 @@ pub async fn execute_with_context( let stream_handle = crate::api::QueryStreamHandle::with_session_context(wrapped, handle.query_context, handle.ctx); Ok(Box::into_raw(Box::new(stream_handle)) as i64) } + +// ── Shared helpers ────────────────────────────────────────────────────────── + +/// Build a per-query RuntimeEnv sharing global caches, with a fresh list-files +/// cache pre-populated for the given table path and object metas. +pub fn build_query_runtime_env( + runtime: &DataFusionRuntime, + table_path: &ListingTableUrl, + object_metas: &[ObjectMeta], +) -> Result, DataFusionError> { + let list_file_cache = Arc::new(DefaultListFilesCache::default()); + let table_scoped_path = datafusion::execution::cache::TableScopedPath { + table: None, + path: table_path.prefix().clone(), + }; + list_file_cache.put(&table_scoped_path, CachedFileList::new(object_metas.to_vec())); + + let runtime_env = RuntimeEnvBuilder::from_runtime_env(&runtime.runtime_env) + .with_cache_manager( + CacheManagerConfig::default() + .with_list_files_cache(Some(list_file_cache)) + .with_file_metadata_cache(Some( + runtime.runtime_env.cache_manager.get_file_metadata_cache(), + )) + .with_files_statistics_cache( + runtime.runtime_env.cache_manager.get_file_statistic_cache(), + ), + ) + .build()?; + Ok(Arc::from(runtime_env)) +} + +/// Build ShardFileInfo list from object metas by reading parquet footers. +/// Each file gets a cumulative `row_base` and per-RG row counts. +pub async fn build_shard_file_infos( + store: &Arc, + object_metas: &[ObjectMeta], +) -> Result, DataFusionError> { + let mut files: Vec = Vec::new(); + let mut cumulative_rows: i64 = 0; + for meta in object_metas { + let reader = datafusion::parquet::arrow::async_reader::ParquetObjectReader::new( + Arc::clone(store), meta.location.clone(), + ).with_file_size(meta.size); + let builder = datafusion::parquet::arrow::ParquetRecordBatchStreamBuilder::new(reader) + .await + .map_err(|e| DataFusionError::Execution(format!("parquet metadata: {}", e)))?; + let pq_meta = builder.metadata().clone(); + let num_rows: i64 = (0..pq_meta.num_row_groups()) + .map(|i| pq_meta.row_group(i).num_rows()) + .sum(); + + files.push(ShardFileInfo { + object_meta: meta.clone(), + row_base: cumulative_rows, + num_rows: num_rows as u64, + row_group_row_counts: (0..pq_meta.num_row_groups()) + .map(|i| pq_meta.row_group(i).num_rows() as u64) + .collect(), + access_plan: None, + }); + cumulative_rows += num_rows; + } + Ok(files) +} + +/// Parse a ListingTableUrl into an ObjectStoreUrl (scheme + authority). +pub fn store_url_from_table_path(table_path: &ListingTableUrl) -> Result { + let url_str = table_path.as_str(); + let parsed = url::Url::parse(url_str) + .map_err(|e| DataFusionError::Execution(format!("parse URL: {}", e)))?; + datafusion::execution::object_store::ObjectStoreUrl::parse( + format!("{}://{}", parsed.scheme(), parsed.authority()), + ) +} + +/// Wrap a DataFusion stream in CrossRtStream and package as a QueryStreamHandle pointer. +pub fn wrap_stream_as_handle( + df_stream: datafusion::execution::SendableRecordBatchStream, + cpu_executor: DedicatedExecutor, + runtime: &DataFusionRuntime, +) -> i64 { + let cross_rt_stream = CrossRtStream::new_with_df_error_stream(df_stream, cpu_executor); + let wrapped = datafusion::physical_plan::stream::RecordBatchStreamAdapter::new( + cross_rt_stream.schema(), + cross_rt_stream, + ); + let query_context = crate::query_tracker::QueryTrackingContext::new( + 0, + runtime.runtime_env.memory_pool.clone(), + ); + let handle = crate::api::QueryStreamHandle::new(wrapped, query_context); + Box::into_raw(Box::new(handle)) as i64 +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs index 060fbcdad7a6a..c340f38386aa3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs @@ -116,18 +116,28 @@ pub async unsafe fn create_session_context( config.options_mut().execution.target_partitions = query_config.target_partitions; config.options_mut().execution.batch_size = query_config.batch_size; - let state = SessionStateBuilder::new() + let mut state_builder = SessionStateBuilder::new() .with_config(config) .with_runtime_env(Arc::from(runtime_env)) .with_default_features() - .with_physical_optimizer_rules(crate::agg_mode::physical_optimizer_rules_without_combine()) - .build(); + .with_physical_optimizer_rules(crate::agg_mode::physical_optimizer_rules_without_combine()); + + // For ListingTable fetch strategy: + // 1. Add ProjectRowIdAnalyzer (logical) — ensures __row_id__ survives pruning. + // 2. Add ProjectRowIdOptimizer (physical) — computes __row_id__ + row_base. + if query_config.fetch_strategy == crate::datafusion_query_config::FetchStrategy::ListingTable { + state_builder = state_builder + .with_analyzer_rule( + Arc::new(crate::project_row_id_analyzer::ProjectRowIdAnalyzer::new()) + ) + .with_physical_optimizer_rule( + Arc::new(crate::project_row_id_optimizer::ProjectRowIdOptimizer) + ); + } + + let state = state_builder.build(); let ctx = SessionContext::new_with_state(state); - // Register OpenSearch UDFs (mvappend, mvfind, mvzip, convert_tz, …) on this session - // so the substrait converter at execute_with_context can resolve their function names. - // Without this, fragment execution fails with "Unsupported function name" because - // df_execute_with_context reuses this handle's ctx instead of building a fresh one. crate::udf::register_all(&ctx); // Register default ListingTable for parquet scans diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_table_provider.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_table_provider.rs new file mode 100644 index 0000000000000..df79a7410dfd8 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_table_provider.rs @@ -0,0 +1,142 @@ +/* + * 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. + */ + +//! Shard-level TableProvider with `row_base` partition column for global row ID computation. + +use std::any::Any; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::catalog::{Session, TableProvider}; +use datafusion::common::{Result, ScalarValue, Statistics}; +use datafusion::common::stats::Precision; +use datafusion::datasource::physical_plan::ParquetSource; +use datafusion::datasource::source::DataSourceExec; +use datafusion::datasource::TableType; +use datafusion::execution::object_store::ObjectStoreUrl; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_datasource::file_groups::FileGroup; +use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_datasource::table_schema::TableSchema; +use datafusion_datasource::PartitionedFile; + +pub use crate::api::ShardFileInfo; + +pub struct ShardTableConfig { + pub file_schema: SchemaRef, + pub files: Vec, + pub store_url: ObjectStoreUrl, +} + +pub struct ShardTableProvider { + table_schema: SchemaRef, + config: ShardTableConfig, +} + +impl ShardTableProvider { + pub fn new(config: ShardTableConfig) -> Self { + let mut fields: Vec> = config.file_schema.fields().iter().cloned().collect(); + fields.push(Arc::new(Field::new("row_base", DataType::Int64, true))); + let table_schema = Arc::new(Schema::new(fields)); + Self { table_schema, config } + } +} + +impl std::fmt::Debug for ShardTableProvider { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ShardTableProvider") + .field("files", &self.config.files.len()) + .finish() + } +} + +#[async_trait] +impl TableProvider for ShardTableProvider { + fn as_any(&self) -> &dyn Any { self } + fn schema(&self) -> SchemaRef { self.table_schema.clone() } + fn table_type(&self) -> TableType { TableType::Base } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> Result> { + Ok(vec![TableProviderFilterPushDown::Inexact; filters.len()]) + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + let num_file_cols = self.config.file_schema.fields().len(); + let partitioned_files: Vec = self.config.files.iter() + .map(|file_info| { + let mut pf = PartitionedFile::from(file_info.object_meta.clone()); + pf.partition_values = vec![ScalarValue::Int64(Some(file_info.row_base))]; + let file_stats = Arc::new(Statistics { + num_rows: Precision::Exact(file_info.num_rows as usize), + total_byte_size: Precision::Inexact(file_info.object_meta.size as usize), + column_statistics: vec![ + datafusion::common::ColumnStatistics::new_unknown(); + num_file_cols + ], + }); + pf = pf.with_statistics(file_stats); + if let Some(ref plan) = file_info.access_plan { + pf = pf.with_extensions(Arc::new(plan.clone())); + } + pf + }) + .collect(); + + let file_groups = vec![FileGroup::new(partitioned_files)]; + + let table_schema = TableSchema::new( + self.config.file_schema.clone(), + vec![Arc::new(Field::new("row_base", DataType::Int64, true))], + ); + + let parquet_source = ParquetSource::new(table_schema); + + let mut builder = FileScanConfigBuilder::new( + self.config.store_url.clone(), + Arc::new(parquet_source), + ) + .with_file_groups(file_groups); + + // Always include the row_base partition column (index = num_file_cols) + // so ProjectRowIdOptimizer can compute __row_id__ + row_base. + let row_base_idx = num_file_cols; + let proj_with_row_base = match projection { + Some(proj) => { + let mut p = proj.clone(); + if !p.contains(&row_base_idx) { + p.push(row_base_idx); + } + p + } + None => { + let mut p: Vec = (0..num_file_cols).collect(); + p.push(row_base_idx); + p + } + }; + builder = builder + .with_projection_indices(Some(proj_with_row_base)) + .map_err(|e| datafusion::error::DataFusionError::Internal(format!("{}", e)))?; + + let file_scan_config = builder.build(); + Ok(DataSourceExec::from_data_source(file_scan_config)) + } + + fn statistics(&self) -> Option { None } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java index 5e7f8dba9c232..131aab318a0cb 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionAnalyticsBackendPlugin.java @@ -664,6 +664,52 @@ public void configureFilterDelegation(FilterDelegationHandle handle, BackendExec FilterTreeCallbacks.setHandle(handle); } + @Override + public org.opensearch.analytics.backend.EngineResultStream fetchByRowIds( + org.opensearch.index.engine.exec.IndexReaderProvider.Reader reader, + org.apache.arrow.vector.BigIntVector rowIdVector, + String[] columns, + org.apache.arrow.memory.BufferAllocator allocator + ) { + DataFusionService dataFusionService = plugin.getDataFusionService(); + if (dataFusionService == null) { + throw new IllegalStateException("DataFusionService not initialized"); + } + + DatafusionReader dfReader = null; + DataFormatRegistry registry = plugin.getDataFormatRegistry(); + for (String formatName : plugin.getSupportedFormats()) { + dfReader = reader.getReader(registry.format(formatName), DatafusionReader.class); + if (dfReader != null) break; + } + if (dfReader == null) { + throw new IllegalStateException("No DatafusionReader available for fetch-by-row-ids"); + } + + // Pass row IDs to Rust via BigIntVector's direct buffer (zero-copy at FFM). + // BigIntVector data buffer is a contiguous off-heap array of i64 values. + long bufAddr = rowIdVector.getDataBuffer().memoryAddress(); + int count = rowIdVector.getValueCount(); + + long streamPtr; + if (bufAddr != 0 && count > 0) { + streamPtr = org.opensearch.be.datafusion.nativelib.NativeBridge.fetchByRowIds( + dfReader.getReaderHandle().getPointer(), + bufAddr, + count, + columns, + dataFusionService.getNativeRuntime().get() + ); + } else { + throw new IllegalStateException("BigIntVector buffer address is 0 or count is 0"); + } + org.opensearch.be.datafusion.nativelib.StreamHandle streamHandle = new org.opensearch.be.datafusion.nativelib.StreamHandle( + streamPtr, + dataFusionService.getNativeRuntime() + ); + return new DatafusionResultStream(streamHandle, allocator); + } + @Override public void setDelegationThreadTracker(org.opensearch.analytics.spi.DelegationThreadTracker tracker) { FilterTreeCallbacks.setThreadTracker(tracker); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java index e58d6630be19a..c7665899d5e9c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionSettings.java @@ -162,6 +162,47 @@ public final class DatafusionSettings { Setting.Property.Dynamic ); + // Fetch strategy constants + public static final String FETCH_STRATEGY_NONE = "none"; + public static final String FETCH_STRATEGY_LISTING_TABLE = "listing_table"; + public static final String FETCH_STRATEGY_INDEXED = "indexed"; + + /** + * Fetch strategy for query-then-fetch (QTF) row ID computation. + *

+ * Controls how shard-global row IDs are computed when the query projects {@code __row_id__}. + *

    + *
  • {@code none} — No global row ID computation. Reads {@code __row_id__} as a regular + * column from parquet (per-file 0-based values, NOT shard-global). Useful for debugging.
    • + *
  • {@code listing_table} — Uses ShardTableProvider with a {@code row_base} partition column. + * Reads {@code __row_id__} from parquet and adds the file's cumulative row offset via + * ProjectRowIdOptimizer ({@code __row_id__ + row_base = global_id}). Works with standard + * DataFusion ListingTable scan path.
    • + *
  • {@code indexed} — Uses the indexed pipeline (segment partitioning, PositionMap). + * Computes row IDs from position ({@code global_base + rg.first_row + position_in_rg}). + * Zero I/O for the row ID column. Fastest path when the indexed executor is available.
    • + *
+ * Default: {@code indexed}. + */ + public static final Setting INDEXED_FETCH_STRATEGY = Setting.simpleString( + "datafusion.indexed.fetch_strategy", + FETCH_STRATEGY_INDEXED, + value -> { + switch (value) { + case FETCH_STRATEGY_NONE: + case FETCH_STRATEGY_LISTING_TABLE: + case FETCH_STRATEGY_INDEXED: + break; + default: + throw new IllegalArgumentException( + "datafusion.indexed.fetch_strategy must be one of " + "[none, listing_table, indexed], got: " + value + ); + } + }, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + // ── All settings registered by the plugin ── public static final List> ALL_SETTINGS = List.of( @@ -186,7 +227,8 @@ public final class DatafusionSettings { INDEXED_MIN_SKIP_RUN_SELECTIVITY_THRESHOLD, INDEXED_SINGLE_COLLECTOR_STRATEGY, INDEXED_TREE_COLLECTOR_STRATEGY, - INDEXED_MAX_COLLECTOR_PARALLELISM + INDEXED_MAX_COLLECTOR_PARALLELISM, + INDEXED_FETCH_STRATEGY ); // ── Snapshot management ── @@ -227,6 +269,7 @@ public DatafusionSettings(ClusterService clusterService) { .singleCollectorStrategy(strategyToWireValue(INDEXED_SINGLE_COLLECTOR_STRATEGY.get(settings))) .treeCollectorStrategy(strategyToWireValue(INDEXED_TREE_COLLECTOR_STRATEGY.get(settings))) .maxCollectorParallelism(INDEXED_MAX_COLLECTOR_PARALLELISM.get(settings)) + .fetchStrategy(fetchStrategyToWireValue(INDEXED_FETCH_STRATEGY.get(settings))) .build(); registerListeners(clusterSettings); @@ -249,6 +292,7 @@ public DatafusionSettings(ClusterService clusterService) { .singleCollectorStrategy(strategyToWireValue(INDEXED_SINGLE_COLLECTOR_STRATEGY.get(settings))) .treeCollectorStrategy(strategyToWireValue(INDEXED_TREE_COLLECTOR_STRATEGY.get(settings))) .maxCollectorParallelism(INDEXED_MAX_COLLECTOR_PARALLELISM.get(settings)) + .fetchStrategy(fetchStrategyToWireValue(INDEXED_FETCH_STRATEGY.get(settings))) .build(); } @@ -281,6 +325,10 @@ void registerListeners(ClusterSettings clusterSettings) { snapshot = WireConfigSnapshot.builder(snapshot).maxCollectorParallelism(newValue).build(); }); + clusterSettings.addSettingsUpdateConsumer(INDEXED_FETCH_STRATEGY, newValue -> { + snapshot = WireConfigSnapshot.builder(snapshot).fetchStrategy(fetchStrategyToWireValue(newValue)).build(); + }); + clusterSettings.addSettingsUpdateConsumer(SearchService.CONCURRENT_SEGMENT_SEARCH_TARGET_MAX_SLICE_COUNT_SETTING, newValue -> { this.maxSliceCount = newValue; snapshot = WireConfigSnapshot.builder(snapshot) @@ -322,6 +370,19 @@ static int strategyToWireValue(String strategy) { } } + static int fetchStrategyToWireValue(String strategy) { + switch (strategy) { + case FETCH_STRATEGY_NONE: + return 0; + case FETCH_STRATEGY_LISTING_TABLE: + return 1; + case FETCH_STRATEGY_INDEXED: + return 2; + default: + throw new IllegalArgumentException("Unknown fetch strategy: " + strategy); + } + } + /** * Derives {@code target_partitions} from the concurrent search mode and * {@code search.concurrent.max_slice_count} setting value. diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/WireConfigSnapshot.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/WireConfigSnapshot.java index 012f47aa9b540..2c5856e5df1e1 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/WireConfigSnapshot.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/WireConfigSnapshot.java @@ -26,7 +26,7 @@ public final class WireConfigSnapshot { /** Total byte size of the wire struct ({@code WireDatafusionQueryConfig}). */ - public static final long BYTE_SIZE = 68; + public static final long BYTE_SIZE = 72; private final int batchSize; private final int targetPartitions; @@ -36,6 +36,7 @@ public final class WireConfigSnapshot { private final int maxCollectorParallelism; private final int singleCollectorStrategy; private final int treeCollectorStrategy; + private final int fetchStrategy; private WireConfigSnapshot(Builder builder) { this.batchSize = builder.batchSize; @@ -46,6 +47,7 @@ private WireConfigSnapshot(Builder builder) { this.maxCollectorParallelism = builder.maxCollectorParallelism; this.singleCollectorStrategy = builder.singleCollectorStrategy; this.treeCollectorStrategy = builder.treeCollectorStrategy; + this.fetchStrategy = builder.fetchStrategy; } public static Builder builder() { @@ -64,7 +66,8 @@ public static Builder builder(WireConfigSnapshot current) { .minSkipRunSelectivityThreshold(current.minSkipRunSelectivityThreshold) .maxCollectorParallelism(current.maxCollectorParallelism) .singleCollectorStrategy(current.singleCollectorStrategy) - .treeCollectorStrategy(current.treeCollectorStrategy); + .treeCollectorStrategy(current.treeCollectorStrategy) + .fetchStrategy(current.fetchStrategy); } public int batchSize() { @@ -99,6 +102,10 @@ public int treeCollectorStrategy() { return treeCollectorStrategy; } + public int fetchStrategy() { + return fetchStrategy; + } + /** * Writes this snapshot into a {@code MemorySegment} matching the * {@code WireDatafusionQueryConfig} {@code #[repr(C)]} layout. @@ -122,8 +129,9 @@ public int treeCollectorStrategy() { * 56 4 max_collector_parallelism i32 from snapshot * 60 4 single_collector_strategy i32 from snapshot * 64 4 tree_collector_strategy i32 from snapshot + * 68 4 fetch_strategy i32 from snapshot (0/1/2) * ────── ──── - * Total: 68 bytes + * Total: 72 bytes * * * @param segment the target memory segment (at least 68 bytes) @@ -155,6 +163,8 @@ public void writeTo(MemorySegment segment) { segment.set(ValueLayout.JAVA_INT, 60, singleCollectorStrategy); // Offset 64: tree_collector_strategy (i32) segment.set(ValueLayout.JAVA_INT, 64, treeCollectorStrategy); + // Offset 68: fetch_strategy (i32) — 0 = None, 1 = ListingTable, 2 = IndexedPredicateOnly + segment.set(ValueLayout.JAVA_INT, 68, fetchStrategy); } /** @@ -170,6 +180,7 @@ public static final class Builder { private int maxCollectorParallelism = 1; private int singleCollectorStrategy = 2; // PageRangeSplit private int treeCollectorStrategy = 1; // TightenOuterBounds + private int fetchStrategy = 1; // ListingTable (ShardTableProvider + ProjectRowIdOptimizer) private Builder() {} @@ -213,6 +224,11 @@ public Builder treeCollectorStrategy(int treeCollectorStrategy) { return this; } + public Builder fetchStrategy(int fetchStrategy) { + this.fetchStrategy = fetchStrategy; + return this; + } + public WireConfigSnapshot build() { return new WireConfigSnapshot(this); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java index 1a3960830ca37..b1318097cf559 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/nativelib/NativeBridge.java @@ -86,6 +86,7 @@ public final class NativeBridge { private static final MethodHandle PREPARE_PARTIAL_PLAN; private static final MethodHandle PREPARE_FINAL_PLAN; private static final MethodHandle EXECUTE_LOCAL_PREPARED_PLAN; + private static final MethodHandle FETCH_BY_ROW_IDS; static { SymbolLookup lib = NativeLibraryLoader.symbolLookup(); @@ -423,6 +424,22 @@ public final class NativeBridge { lib.find("df_execute_local_prepared_plan").orElseThrow(), FunctionDescriptor.of(ValueLayout.JAVA_LONG, ValueLayout.JAVA_LONG) ); + + // i64 df_fetch_by_row_ids(shard_view_ptr, row_ids_buf_ptr, row_ids_count, + // col_names_ptr, col_names_len_ptr, col_names_count, runtime_ptr) + FETCH_BY_ROW_IDS = linker.downcallHandle( + lib.find("df_fetch_by_row_ids").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG + ) + ); } private NativeBridge() {} @@ -1020,6 +1037,38 @@ public static long executeLocalPreparedPlan(long sessionPtr) { } } + /** + * QTF fetch phase: reads specific rows by global row ID from parquet. + * Row IDs are passed as a direct buffer pointer (zero-copy from BigIntVector's ArrowBuf). + * + * @param readerPtr pointer to the shard view (DatafusionReader) + * @param rowIdsBufAddr memory address of the BigIntVector's data buffer (i64 values) + * @param rowIdsCount number of row IDs + * @param columns column names to read + * @param runtimePtr pointer to the DataFusion runtime + * @return opaque stream pointer + */ + public static long fetchByRowIds(long readerPtr, long rowIdsBufAddr, int rowIdsCount, String[] columns, long runtimePtr) { + NativeHandle.validatePointer(readerPtr, "reader"); + NativeHandle.validatePointer(runtimePtr, "runtime"); + if (rowIdsBufAddr == 0) { + throw new IllegalArgumentException("rowIdsBufAddr must be non-zero"); + } + try (var call = new NativeCall()) { + var colNames = call.strArray(columns); + return call.invoke( + FETCH_BY_ROW_IDS, + readerPtr, + rowIdsBufAddr, + (long) rowIdsCount, + colNames.ptrs(), + colNames.lens(), + colNames.count(), + runtimePtr + ); + } + } + public static void createCache(long cacheManagerPtr, String cacheType, long sizeLimit, String evictionType) { try (var call = new NativeCall()) { var type = call.str(cacheType); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java index 3c389add9d00c..58020d884af42 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchService.java @@ -226,4 +226,32 @@ private ShardScanExecutionContext buildContext( return ctx; } + /** + * QTF fetch phase: retrieves specific rows by global row ID via the backend SPI. + */ + public org.opensearch.analytics.backend.EngineResultStream executeFetchByRowIds( + String queryId, + long[] rowIds, + String[] columns, + IndexShard shard + ) { + IndexReaderProvider readerProvider = shard.getReaderProvider(); + if (readerProvider == null) { + throw new IllegalStateException("No ReaderProvider on " + shard.shardId()); + } + try { + GatedCloseable gatedReader = readerProvider.acquireReader(); + org.apache.arrow.vector.BigIntVector rowIdVector = new org.apache.arrow.vector.BigIntVector("__row_id__", allocator); + rowIdVector.allocateNew(rowIds.length); + for (int i = 0; i < rowIds.length; i++) { + rowIdVector.set(i, rowIds[i]); + } + rowIdVector.setValueCount(rowIds.length); + + AnalyticsSearchBackendPlugin backend = backends.values().iterator().next(); + return backend.fetchByRowIds(gatedReader.get(), rowIdVector, columns, allocator); + } catch (Exception e) { + throw new RuntimeException("Failed to execute fetch-by-row-ids on " + shard.shardId(), e); + } + } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContext.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContext.java new file mode 100644 index 0000000000000..4e36af1a0e1a4 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContext.java @@ -0,0 +1,98 @@ +/* + * 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.analytics.exec; + +import org.opensearch.common.concurrent.GatedCloseable; +import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; + +import java.io.Closeable; +import java.io.IOException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; + +/** + * Holds an acquired reader open across query and fetch phases of a QTF query. + * The reader is acquired once during the query phase and reused during the fetch phase. + * + *

Lifecycle: + *

    + *
  1. Query phase: acquires reader, stores in context, marks in-use
    2. + *
  2. Query completes: marks not-in-use, starts keepAlive countdown
    3. + *
  3. Fetch phase: marks in-use, uses the same reader
    4. + *
  4. Fetch completes: marks not-in-use, context can be freed
    5. + *
  5. Reaper: closes expired contexts (not in-use AND keepAlive elapsed)
    6. + *
+ */ +public class ReaderContext implements Closeable { + + private final String queryId; + private final GatedCloseable gatedReader; + private final AtomicBoolean inUse = new AtomicBoolean(false); + private final AtomicLong lastAccessTime; + private volatile long keepAliveMillis; + private volatile boolean closed; + + public ReaderContext(String queryId, GatedCloseable gatedReader, long keepAliveMillis) { + this.queryId = queryId; + this.gatedReader = gatedReader; + this.keepAliveMillis = keepAliveMillis; + this.lastAccessTime = new AtomicLong(System.currentTimeMillis()); + } + + public String getQueryId() { + return queryId; + } + + public Reader getReader() { + return gatedReader.get(); + } + + /** + * Mark the context as in-use. Returns true if successfully marked. + */ + public boolean markInUse() { + if (closed) return false; + lastAccessTime.set(System.currentTimeMillis()); + return inUse.compareAndSet(false, true); + } + + /** + * Mark the context as no longer in-use. Updates last access time. + */ + public void markDone() { + lastAccessTime.set(System.currentTimeMillis()); + inUse.set(false); + } + + /** + * Check if this context has expired (not in-use AND keepAlive elapsed). + */ + public boolean isExpired() { + if (inUse.get()) { + return false; + } + long elapsed = System.currentTimeMillis() - lastAccessTime.get(); + return elapsed > keepAliveMillis; + } + + public void setKeepAliveMillis(long keepAliveMillis) { + this.keepAliveMillis = keepAliveMillis; + } + + long getLastAccessTime() { + return lastAccessTime.get(); + } + + @Override + public void close() throws IOException { + if (closed) return; + closed = true; + gatedReader.close(); + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContextStore.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContextStore.java new file mode 100644 index 0000000000000..fa2ff1b67af2f --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContextStore.java @@ -0,0 +1,127 @@ +/* + * 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.analytics.exec; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.common.concurrent.GatedCloseable; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; +import org.opensearch.threadpool.ThreadPool; + +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Stores reader contexts across query and fetch phases for QTF execution. + * Each context is keyed by queryId and holds the acquired reader open + * until the fetch phase completes or the keepAlive expires. + * + *

A reaper thread periodically scans for expired contexts and closes them. + */ +public class ReaderContextStore { + + private static final Logger logger = LogManager.getLogger(ReaderContextStore.class); + private static final TimeValue REAPER_INTERVAL = TimeValue.timeValueSeconds(10); + + public static final Setting READER_CONTEXT_KEEP_ALIVE = Setting.positiveTimeSetting( + "analytics.qtf.reader_context.keep_alive", + TimeValue.timeValueMinutes(5), + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + private final Map activeContexts = new ConcurrentHashMap<>(); + private volatile long defaultKeepAliveMillis; + + public ReaderContextStore(ThreadPool threadPool) { + this(threadPool, READER_CONTEXT_KEEP_ALIVE.getDefault(null).millis()); + } + + public ReaderContextStore(ThreadPool threadPool, long defaultKeepAliveMillis) { + this.defaultKeepAliveMillis = defaultKeepAliveMillis; + threadPool.scheduleWithFixedDelay(new Reaper(), REAPER_INTERVAL, ThreadPool.Names.SAME); + } + + public void setKeepAlive(TimeValue keepAlive) { + this.defaultKeepAliveMillis = keepAlive.millis(); + } + + /** + * Create and store a new reader context for the given query. + * Acquires the reader and marks it in-use. + */ + public ReaderContext createContext(String queryId, GatedCloseable gatedReader) { + ReaderContext ctx = new ReaderContext(queryId, gatedReader, defaultKeepAliveMillis); + ctx.markInUse(); + activeContexts.put(queryId, ctx); + return ctx; + } + + /** + * Get an existing context by queryId. Returns null if not found or expired. + */ + public ReaderContext getContext(String queryId) { + return activeContexts.get(queryId); + } + + /** + * Acquire a context for use (fetch phase). Marks it in-use. + * Returns null if not found or already closed. + */ + public ReaderContext acquireContext(String queryId) { + ReaderContext ctx = activeContexts.get(queryId); + if (ctx == null) return null; + if (ctx.markInUse()) { + return ctx; + } + return null; + } + + /** + * Release a context after use (query or fetch done). Marks it not-in-use. + */ + public void releaseContext(String queryId) { + ReaderContext ctx = activeContexts.get(queryId); + if (ctx != null) { + ctx.markDone(); + } + } + + /** + * Remove and close a context (fetch complete, no longer needed). + */ + public void freeContext(String queryId) { + ReaderContext ctx = activeContexts.remove(queryId); + if (ctx != null) { + try { + ctx.close(); + } catch (Exception e) { + logger.warn("[ReaderContextStore] Failed to close context for query={}", queryId, e); + } + } + } + + public int activeCount() { + return activeContexts.size(); + } + + private class Reaper implements Runnable { + @Override + public void run() { + for (ReaderContext ctx : activeContexts.values()) { + if (ctx.isExpired()) { + logger.debug("[ReaderContextStore] Freeing expired context for query={}", ctx.getQueryId()); + freeContext(ctx.getQueryId()); + } + } + } + } +} diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/QueryThenFetchDataNodeTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/QueryThenFetchDataNodeTests.java new file mode 100644 index 0000000000000..3ad5b594e7195 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/QueryThenFetchDataNodeTests.java @@ -0,0 +1,223 @@ +/* + * 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.analytics.exec; + +import org.opensearch.common.concurrent.GatedCloseable; +import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; + +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.mockito.Mockito.mock; + +/** + * Tests the full data-node QTF lifecycle: query phase stores reader in context, + * fetch phase reuses the same reader, and context is freed after fetch. + */ +public class QueryThenFetchDataNodeTests extends OpenSearchTestCase { + + private ThreadPool threadPool; + + @Override + public void setUp() throws Exception { + super.setUp(); + threadPool = new TestThreadPool(getTestName()); + } + + @Override + public void tearDown() throws Exception { + terminate(threadPool); + super.tearDown(); + } + + private GatedCloseable mockGatedReader(AtomicBoolean closedFlag) { + Reader reader = mock(Reader.class); + return new GatedCloseable<>(reader, () -> closedFlag.set(true)); + } + + /** + * Happy path: query phase stores reader -> fetch phase gets same reader -> free closes it. + */ + public void testFullQueryThenFetchLifecycle() { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + // Simulate query phase: acquire reader and store in context + GatedCloseable gatedReader = mockGatedReader(closed); + Reader originalReader = gatedReader.get(); + ReaderContext ctx = store.createContext("query-1", gatedReader); + assertNotNull(ctx); + assertSame(originalReader, ctx.getReader()); + + // Query phase completes, release context (reader stays alive for fetch) + store.releaseContext("query-1"); + assertFalse("Reader must not be closed between phases", closed.get()); + + // Simulate fetch phase: acquire the same context + ReaderContext fetchCtx = store.acquireContext("query-1"); + assertNotNull("Fetch phase must find the context", fetchCtx); + assertSame("Fetch must get the SAME reader instance", originalReader, fetchCtx.getReader()); + + // Fetch completes + fetchCtx.markDone(); + store.freeContext("query-1"); + assertTrue("Reader must be closed after completeFetch", closed.get()); + assertEquals(0, store.activeCount()); + } + + /** + * Fetch arrives after the keep-alive expires and reaper cleans up -> acquireContext returns null. + */ + public void testFetchWithExpiredContext() throws Exception { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool, 50); // 50ms keepAlive + + store.createContext("query-expired", mockGatedReader(closed)); + store.releaseContext("query-expired"); + + // Wait for reaper to clean up the expired context + assertBusy(() -> { + assertTrue("Reader should be closed by reaper", closed.get()); + assertEquals(0, store.activeCount()); + }); + + // Fetch arrives too late + ReaderContext fetchCtx = store.acquireContext("query-expired"); + assertNull("Expired context must not be acquirable", fetchCtx); + } + + /** + * Fetch for an unknown queryId returns null (no prior query phase). + */ + public void testFetchWithoutPriorQuery() { + ReaderContextStore store = new ReaderContextStore(threadPool); + + ReaderContext ctx = store.acquireContext("nonexistent-query"); + assertNull("Unknown queryId must return null", ctx); + } + + /** + * Verify reader's close callback is NOT invoked between query and fetch phases. + */ + public void testReaderNotClosedBetweenPhases() { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + store.createContext("query-2", mockGatedReader(closed)); + // Query phase done + store.releaseContext("query-2"); + + // Between phases: reader must still be alive + assertFalse("Reader must remain open while awaiting fetch", closed.get()); + assertEquals(1, store.activeCount()); + + // Cleanup + store.freeContext("query-2"); + } + + /** + * Verify reader's close callback IS invoked after completeFetch (freeContext). + */ + public void testReaderClosedAfterCompleteFetch() { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + store.createContext("query-3", mockGatedReader(closed)); + store.releaseContext("query-3"); + + ReaderContext fetchCtx = store.acquireContext("query-3"); + assertNotNull(fetchCtx); + assertFalse("Reader must not be closed during fetch", closed.get()); + + // Simulate completeFetch + fetchCtx.markDone(); + store.freeContext("query-3"); + + assertTrue("Reader must be closed after freeContext", closed.get()); + assertEquals(0, store.activeCount()); + assertNull("Context must be removed from store", store.getContext("query-3")); + } + + + /** + * Multiple queries have independent contexts with different readers and independent lifecycles. + */ + public void testMultipleQueriesIndependentContexts() { + AtomicBoolean closed1 = new AtomicBoolean(false); + AtomicBoolean closed2 = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + GatedCloseable gated1 = mockGatedReader(closed1); + GatedCloseable gated2 = mockGatedReader(closed2); + Reader reader1 = gated1.get(); + Reader reader2 = gated2.get(); + + store.createContext("q1", gated1); + store.createContext("q2", gated2); + assertEquals(2, store.activeCount()); + + // Release both query phases + store.releaseContext("q1"); + store.releaseContext("q2"); + + // Fetch q1 + ReaderContext fetch1 = store.acquireContext("q1"); + assertNotNull(fetch1); + assertSame(reader1, fetch1.getReader()); + assertNotSame("Different queries must have different readers", reader1, reader2); + + // Free q1 while q2 still alive + fetch1.markDone(); + store.freeContext("q1"); + assertTrue("q1 reader closed", closed1.get()); + assertFalse("q2 reader still alive", closed2.get()); + assertEquals(1, store.activeCount()); + + // Fetch q2 + ReaderContext fetch2 = store.acquireContext("q2"); + assertNotNull(fetch2); + assertSame(reader2, fetch2.getReader()); + fetch2.markDone(); + store.freeContext("q2"); + assertTrue("q2 reader closed", closed2.get()); + assertEquals(0, store.activeCount()); + } + + /** + * Even if query "fails" (we just release without a successful result), the context remains + * available for potential fetch. The reaper cleans it if fetch never arrives. + */ + public void testQueryContextSurvivesQueryFailure() throws Exception { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool, 50); // short keepAlive for reaper test + + store.createContext("query-failed", mockGatedReader(closed)); + + // Simulate "query failure": release the context without explicit cleanup + store.releaseContext("query-failed"); + + // Context still exists immediately after failure + assertFalse("Reader not yet closed", closed.get()); + assertEquals(1, store.activeCount()); + + // A fetch could still arrive in time + ReaderContext fetchCtx = store.acquireContext("query-failed"); + assertNotNull("Context survives query failure for potential fetch", fetchCtx); + fetchCtx.markDone(); + + // If fetch never comes, reaper eventually cleans up + // (release again so reaper can reap it) + assertBusy(() -> { + assertTrue("Reaper eventually closes reader", closed.get()); + assertEquals(0, store.activeCount()); + }); + } +} diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ReaderContextStoreTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ReaderContextStoreTests.java new file mode 100644 index 0000000000000..006fab4e628a1 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ReaderContextStoreTests.java @@ -0,0 +1,175 @@ +/* + * 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.analytics.exec; + +import org.opensearch.common.concurrent.GatedCloseable; +import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; + +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.mockito.Mockito.mock; + +public class ReaderContextStoreTests extends OpenSearchTestCase { + + private ThreadPool threadPool; + + @Override + public void setUp() throws Exception { + super.setUp(); + threadPool = new TestThreadPool(getTestName()); + } + + @Override + public void tearDown() throws Exception { + terminate(threadPool); + super.tearDown(); + } + + private GatedCloseable mockGatedReader(AtomicBoolean closedFlag) { + Reader reader = mock(Reader.class); + return new GatedCloseable<>(reader, () -> closedFlag.set(true)); + } + + public void testCreateAndAcquireContext() { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + GatedCloseable gated = mockGatedReader(closed); + ReaderContext ctx = store.createContext("q1", gated); + + assertNotNull(ctx); + assertEquals(1, store.activeCount()); + + // Release from query phase + store.releaseContext("q1"); + + // Acquire for fetch phase + ReaderContext fetched = store.acquireContext("q1"); + assertNotNull("Should acquire for fetch", fetched); + assertSame(ctx.getReader(), fetched.getReader()); + } + + public void testAcquireNonExistentReturnsNull() { + ReaderContextStore store = new ReaderContextStore(threadPool); + + assertNull(store.acquireContext("no-such-query")); + } + + public void testAcquireWhileInUseReturnsNull() { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + store.createContext("q1", mockGatedReader(closed)); + // Context is already in-use from createContext + + assertNull("Should not acquire while in-use", store.acquireContext("q1")); + } + + public void testFreeContextClosesReader() { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + store.createContext("q1", mockGatedReader(closed)); + store.releaseContext("q1"); + store.freeContext("q1"); + + assertTrue("Reader should be closed after freeContext", closed.get()); + assertEquals(0, store.activeCount()); + } + + public void testFreeContextRemovesFromStore() { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + store.createContext("q1", mockGatedReader(closed)); + store.releaseContext("q1"); + store.freeContext("q1"); + + assertNull("Should not find after free", store.getContext("q1")); + } + + public void testReaperCleansExpiredContexts() throws Exception { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool, 50); // 50ms keepAlive + + store.createContext("q1", mockGatedReader(closed)); + store.releaseContext("q1"); + + // Wait for expiry + reaper cycle + assertBusy(() -> { + assertTrue("Reader should be closed by reaper", closed.get()); + assertEquals(0, store.activeCount()); + }); + } + + public void testReaperDoesNotCleanInUseContext() throws Exception { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool, 1); // 1ms keepAlive + + store.createContext("q1", mockGatedReader(closed)); + // Still in-use (not released) + + Thread.sleep(50); + assertFalse("Reaper should not close in-use context", closed.get()); + assertEquals(1, store.activeCount()); + + // Cleanup + store.releaseContext("q1"); + store.freeContext("q1"); + } + + public void testMultipleContexts() { + AtomicBoolean closed1 = new AtomicBoolean(false); + AtomicBoolean closed2 = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + store.createContext("q1", mockGatedReader(closed1)); + store.createContext("q2", mockGatedReader(closed2)); + assertEquals(2, store.activeCount()); + + store.releaseContext("q1"); + store.freeContext("q1"); + assertEquals(1, store.activeCount()); + assertTrue(closed1.get()); + assertFalse(closed2.get()); + + store.releaseContext("q2"); + store.freeContext("q2"); + assertEquals(0, store.activeCount()); + assertTrue(closed2.get()); + } + + public void testQueryThenFetchLifecycle() { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + // Query phase: create + use + store.createContext("q1", mockGatedReader(closed)); + ReaderContext ctx = store.getContext("q1"); + assertNotNull(ctx); + assertNotNull(ctx.getReader()); + + // Query done + store.releaseContext("q1"); + assertFalse("Reader stays open between phases", closed.get()); + + // Fetch phase: acquire + use + ReaderContext fetchCtx = store.acquireContext("q1"); + assertNotNull(fetchCtx); + assertSame(ctx.getReader(), fetchCtx.getReader()); + + // Fetch done: free + fetchCtx.markDone(); + store.freeContext("q1"); + assertTrue("Reader closed after fetch", closed.get()); + } +} diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ReaderContextTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ReaderContextTests.java new file mode 100644 index 0000000000000..fef19ed5d2cf7 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ReaderContextTests.java @@ -0,0 +1,117 @@ +/* + * 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.analytics.exec; + +import org.opensearch.common.concurrent.GatedCloseable; +import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.mockito.Mockito.mock; + +public class ReaderContextTests extends OpenSearchTestCase { + + private GatedCloseable mockGatedReader() { + Reader reader = mock(Reader.class); + AtomicBoolean closed = new AtomicBoolean(false); + return new GatedCloseable<>(reader, () -> closed.set(true)); + } + + public void testMarkInUseAndDone() { + GatedCloseable gated = mockGatedReader(); + ReaderContext ctx = new ReaderContext("q1", gated, 30_000); + + assertTrue("Should mark in-use successfully", ctx.markInUse()); + assertFalse("Second markInUse should fail (already in-use)", ctx.markInUse()); + + ctx.markDone(); + assertTrue("Should mark in-use again after done", ctx.markInUse()); + } + + public void testNotExpiredWhileInUse() { + GatedCloseable gated = mockGatedReader(); + ReaderContext ctx = new ReaderContext("q1", gated, 1); // 1ms keepAlive + + ctx.markInUse(); + assertFalse("Should not expire while in-use", ctx.isExpired()); + } + + public void testExpiresAfterKeepAlive() throws Exception { + GatedCloseable gated = mockGatedReader(); + ReaderContext ctx = new ReaderContext("q1", gated, 50); // 50ms keepAlive + + ctx.markInUse(); + ctx.markDone(); + + Thread.sleep(100); + assertTrue("Should expire after keepAlive", ctx.isExpired()); + } + + public void testNotExpiredBeforeKeepAlive() { + GatedCloseable gated = mockGatedReader(); + ReaderContext ctx = new ReaderContext("q1", gated, 30_000); + + ctx.markInUse(); + ctx.markDone(); + + assertFalse("Should not expire before keepAlive", ctx.isExpired()); + } + + public void testCloseReleasesReader() throws Exception { + AtomicBoolean closed = new AtomicBoolean(false); + Reader reader = mock(Reader.class); + GatedCloseable gated = new GatedCloseable<>(reader, () -> closed.set(true)); + + ReaderContext ctx = new ReaderContext("q1", gated, 30_000); + ctx.close(); + + assertTrue("Reader should be closed", closed.get()); + } + + public void testMarkInUseAfterCloseReturnsFalse() throws Exception { + GatedCloseable gated = mockGatedReader(); + ReaderContext ctx = new ReaderContext("q1", gated, 30_000); + ctx.close(); + + assertFalse("markInUse should fail after close", ctx.markInUse()); + } + + public void testGetReader() { + Reader reader = mock(Reader.class); + GatedCloseable gated = new GatedCloseable<>(reader, () -> {}); + ReaderContext ctx = new ReaderContext("q1", gated, 30_000); + + assertSame(reader, ctx.getReader()); + } + + public void testGetQueryId() { + GatedCloseable gated = mockGatedReader(); + ReaderContext ctx = new ReaderContext("test-query-123", gated, 30_000); + + assertEquals("test-query-123", ctx.getQueryId()); + } + + public void testLastAccessTimeUpdatedOnMarkInUse() throws Exception { + GatedCloseable gated = mockGatedReader(); + ReaderContext ctx = new ReaderContext("q1", gated, 30_000); + + long timeAfterCreate = ctx.getLastAccessTime(); + + Thread.sleep(20); + ctx.markInUse(); + long timeAfterMarkInUse = ctx.getLastAccessTime(); + assertTrue("lastAccessTime should advance on markInUse", timeAfterMarkInUse > timeAfterCreate); + + Thread.sleep(20); + ctx.markDone(); + long timeAfterDone = ctx.getLastAccessTime(); + assertTrue("lastAccessTime should advance on markDone", timeAfterDone > timeAfterMarkInUse); + } +} diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/QueryThenFetchIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/QueryThenFetchIT.java new file mode 100644 index 0000000000000..4fae4a9dbfbc5 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/QueryThenFetchIT.java @@ -0,0 +1,457 @@ +/* + * 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.analytics.qa; + +import org.apache.lucene.tests.util.LuceneTestCase.AwaitsFix; +import org.opensearch.client.Request; +import org.opensearch.client.Response; + +import java.io.IOException; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +/** + * End-to-end correctness test for Query-Then-Fetch (QTF) row ID computation. + * + *

Verifies that shard-global row IDs are correctly computed across multiple + * parquet segments regardless of fetch strategy. Each query randomly selects + * between "indexed" (position-based) and "listing_table" (row_base partition) + * strategies to ensure both paths are exercised across test runs. + * + *

Test data (4 docs, 2 segments, 2 docs each): + * 

+ *   Segment 1: alpha(value=10, cat=A, id=0), beta(value=20, cat=B, id=1)
+ *   Segment 2: gamma(value=30, cat=A, id=2), delta(value=40, cat=B, id=3)
+ *
+ * + *

Invariants verified: + *

    + *
  • Row IDs are globally unique (no duplicates across segments)
    • + *
  • Row IDs are non-null
    • + *
  • Row IDs are deterministic (same doc always gets same ID)
    • + *
  • Data columns are not corrupted by row ID injection
    • + *
  • Sort, filter, limit operations don't affect row ID correctness
    • + *
+ */ +@AwaitsFix(bugUrl = "https://github.com/opensearch-project/OpenSearch/issues/0000") +public class QueryThenFetchIT extends AnalyticsRestTestCase { + + private static final String INDEX = "qtf_correctness"; + private static final String STRATEGY_INDEXED = "indexed"; + private static final String STRATEGY_LISTING = "listing_table"; + + private static boolean ready = false; + + private void setup() throws IOException { + if (ready) return; + try { client().performRequest(new Request("DELETE", "/" + INDEX)); } catch (Exception ignored) {} + + Request create = new Request("PUT", "/" + INDEX); + create.setJsonEntity("{" + + "\"settings\":{" + + " \"number_of_shards\":1,\"number_of_replicas\":0," + + " \"index.pluggable.dataformat.enabled\":true," + + " \"index.pluggable.dataformat\":\"composite\"," + + " \"index.composite.primary_data_format\":\"parquet\"," + + " \"index.composite.secondary_data_formats\":\"lucene\"" + + "}," + + "\"mappings\":{\"properties\":{" + + " \"name\":{\"type\":\"keyword\"}," + + " \"value\":{\"type\":\"integer\"}," + + " \"category\":{\"type\":\"keyword\"}," + + " \"description\":{\"type\":\"text\"}" + + "}}}"); + client().performRequest(create); + + Request health = new Request("GET", "/_cluster/health/" + INDEX); + health.addParameter("wait_for_status", "green"); + health.addParameter("timeout", "30s"); + client().performRequest(health); + + // Segment 1 + bulk("{\"index\":{}}\n{\"name\":\"alpha\",\"value\":10,\"category\":\"A\",\"description\":\"fast red car\"}\n" + + "{\"index\":{}}\n{\"name\":\"beta\",\"value\":20,\"category\":\"B\",\"description\":\"slow blue truck\"}\n"); + client().performRequest(new Request("POST", "/" + INDEX + "/_flush?force=true")); + + // Segment 2 + bulk("{\"index\":{}}\n{\"name\":\"gamma\",\"value\":30,\"category\":\"A\",\"description\":\"fast green bike\"}\n" + + "{\"index\":{}}\n{\"name\":\"delta\",\"value\":40,\"category\":\"B\",\"description\":\"slow red bus\"}\n"); + client().performRequest(new Request("POST", "/" + INDEX + "/_flush?force=true")); + + ready = true; + } + + // ═══════════════════════════════════════════════════════════════════════════ + // Each query shape tested with BOTH strategies. Exact expected row IDs. + // alpha=0, beta=1, gamma=2, delta=3. + // ═══════════════════════════════════════════════════════════════════════════ + + // ── No filter, no sort (FilterClass::None, full scan) ── + + public void testIndexed_NoFilter() throws IOException { + assertNoFilter(STRATEGY_INDEXED); + } + + public void testListing_NoFilter() throws IOException { + assertNoFilter(STRATEGY_LISTING); + } + + private void assertNoFilter(String strategy) throws IOException { + List> rows = query(strategy, "fields __row_id__, name, value"); + assertEquals(4, rows.size()); + assertGlobalIds(rows, 0, 1, 2, 3); + } + + // ── Sort + LIMIT (FilterClass::None with sort + limit pushdown) ── + + public void testIndexed_SortLimit() throws IOException { + assertSortLimit(STRATEGY_INDEXED); + } + + public void testListing_SortLimit() throws IOException { + assertSortLimit(STRATEGY_LISTING); + } + + private void assertSortLimit(String strategy) throws IOException { + List> rows = query(strategy, + "sort value | fields __row_id__, name, value | head 2"); + assertEquals(2, rows.size()); + assertRow(rows.get(0), "alpha", 10); + assertRow(rows.get(1), "beta", 20); + assertExactIds(rows, 0, 1); + } + + // ── Keyword equality filter (SingleCollector path) ── + + public void testIndexed_KeywordFilter() throws IOException { + assertKeywordFilter(STRATEGY_INDEXED); + } + + public void testListing_KeywordFilter() throws IOException { + assertKeywordFilter(STRATEGY_LISTING); + } + + private void assertKeywordFilter(String strategy) throws IOException { + List> rows = query(strategy, + "where category = 'A' | sort value | fields __row_id__, name, value"); + assertEquals(2, rows.size()); + assertRow(rows.get(0), "alpha", 10); + assertRow(rows.get(1), "gamma", 30); + assertExactIds(rows, 0, 2); + } + + // ── Numeric range filter (predicate-only, no Collector) ── + + public void testIndexed_RangeFilter() throws IOException { + assertRangeFilter(STRATEGY_INDEXED); + } + + public void testListing_RangeFilter() throws IOException { + assertRangeFilter(STRATEGY_LISTING); + } + + private void assertRangeFilter(String strategy) throws IOException { + List> rows = query(strategy, + "where value > 15 and value < 35 | sort value | fields __row_id__, name, value"); + assertEquals(2, rows.size()); + assertRow(rows.get(0), "beta", 20); + assertRow(rows.get(1), "gamma", 30); + assertExactIds(rows, 1, 2); + } + + // ── OR filter (Tree path — BitmapTreeEvaluator) ── + + public void testIndexed_OrFilter() throws IOException { + assertOrFilter(STRATEGY_INDEXED); + } + + public void testListing_OrFilter() throws IOException { + assertOrFilter(STRATEGY_LISTING); + } + + private void assertOrFilter(String strategy) throws IOException { + List> rows = query(strategy, + "where name = 'alpha' or name = 'delta' | sort value | fields __row_id__, name, value"); + assertEquals(2, rows.size()); + assertRow(rows.get(0), "alpha", 10); + assertRow(rows.get(1), "delta", 40); + assertExactIds(rows, 0, 3); + } + + // ── Combined filter (SingleCollector + residual predicate) ── + + public void testIndexed_CombinedFilter() throws IOException { + assertCombinedFilter(STRATEGY_INDEXED); + } + + public void testListing_CombinedFilter() throws IOException { + assertCombinedFilter(STRATEGY_LISTING); + } + + private void assertCombinedFilter(String strategy) throws IOException { + List> rows = query(strategy, + "where category = 'B' and value > 25 | fields __row_id__, name, value"); + assertEquals(1, rows.size()); + assertEquals("delta", rows.get(0).get(1)); + assertEquals(Long.valueOf(3), toLong(rows.get(0).get(0))); + } + + // ── Single exact match from segment 2 ── + + public void testIndexed_ExactMatch() throws IOException { + assertExactMatch(STRATEGY_INDEXED); + } + + public void testListing_ExactMatch() throws IOException { + assertExactMatch(STRATEGY_LISTING); + } + + private void assertExactMatch(String strategy) throws IOException { + List> rows = query(strategy, + "where name = 'gamma' | fields __row_id__, name, value"); + assertEquals(1, rows.size()); + assertEquals("gamma", rows.get(0).get(1)); + assertEquals(Long.valueOf(2), toLong(rows.get(0).get(0))); + } + + // ── Multi-column projection (data integrity check) ── + + public void testIndexed_AllColumns() throws IOException { + assertAllColumns(STRATEGY_INDEXED); + } + + public void testListing_AllColumns() throws IOException { + assertAllColumns(STRATEGY_LISTING); + } + + private void assertAllColumns(String strategy) throws IOException { + List> rows = query(strategy, + "sort value | fields __row_id__, name, value, category"); + assertEquals(4, rows.size()); + assertEquals("alpha", rows.get(0).get(1)); + assertNum(10, rows.get(0).get(2)); + assertEquals("A", rows.get(0).get(3)); + assertEquals("delta", rows.get(3).get(1)); + assertNum(40, rows.get(3).get(2)); + assertEquals("B", rows.get(3).get(3)); + assertGlobalIds(rows, 0, 1, 2, 3); + } + + // ── Full-text match (forces Lucene Collector — SingleCollector path) ── + + public void testIndexed_MatchFilter() throws IOException { + assertMatchFilter(STRATEGY_INDEXED); + } + + public void testListing_MatchFilter() throws IOException { + assertMatchFilter(STRATEGY_LISTING); + } + + private void assertMatchFilter(String strategy) throws IOException { + // match(description, 'fast') → alpha("fast red car", id=0) + gamma("fast green bike", id=2) + List> rows = query(strategy, + "where match(description, 'fast') | sort value | fields __row_id__, name, value"); + assertEquals(2, rows.size()); + assertRow(rows.get(0), "alpha", 10); + assertRow(rows.get(1), "gamma", 30); + assertExactIds(rows, 0, 2); + } + + // ── Full-text match + numeric residual (SingleCollector + predicate) ── + + public void testIndexed_MatchWithResidual() throws IOException { + assertMatchWithResidual(STRATEGY_INDEXED); + } + + public void testListing_MatchWithResidual() throws IOException { + assertMatchWithResidual(STRATEGY_LISTING); + } + + private void assertMatchWithResidual(String strategy) throws IOException { + // match(description, 'red') AND value > 15 → only delta("slow red bus", value=40, id=3) + List> rows = query(strategy, + "where match(description, 'red') and value > 15 | fields __row_id__, name, value"); + assertEquals(1, rows.size()); + assertEquals("delta", rows.get(0).get(1)); + assertNum(40, rows.get(0).get(2)); + assertEquals(Long.valueOf(3), toLong(rows.get(0).get(0))); + } + + // ── Full-text OR match (Tree path — multiple Collectors) ── + + public void testIndexed_MatchOrFilter() throws IOException { + assertMatchOrFilter(STRATEGY_INDEXED); + } + + public void testListing_MatchOrFilter() throws IOException { + assertMatchOrFilter(STRATEGY_LISTING); + } + + private void assertMatchOrFilter(String strategy) throws IOException { + // match(description, 'truck') OR match(description, 'bike') + // → beta("slow blue truck", id=1) + gamma("fast green bike", id=2) + List> rows = query(strategy, + "where match(description, 'truck') or match(description, 'bike') | sort value | fields __row_id__, name, value"); + assertEquals(2, rows.size()); + assertRow(rows.get(0), "beta", 20); + assertRow(rows.get(1), "gamma", 30); + assertExactIds(rows, 1, 2); + } + + // ── SELECT row_id + sort column only (QTF query phase pattern) ── + + public void testIndexed_SelectRowIdAndSortKey() throws IOException { + assertSelectRowIdAndSortKey(STRATEGY_INDEXED); + } + + public void testListing_SelectRowIdAndSortKey() throws IOException { + assertSelectRowIdAndSortKey(STRATEGY_LISTING); + } + + private void assertSelectRowIdAndSortKey(String strategy) throws IOException { + // Mimics: SELECT __row_id__, value FROM t ORDER BY value LIMIT 4 + // This is the pure QTF query phase — row_id + sort key, no other columns + List> rows = query(strategy, + "sort value | fields __row_id__, value"); + assertEquals(4, rows.size()); + // Verify sort order by value + assertNum(10, rows.get(0).get(1)); + assertNum(20, rows.get(1).get(1)); + assertNum(30, rows.get(2).get(1)); + assertNum(40, rows.get(3).get(1)); + assertExactIds(rows, 0, 1, 2, 3); + } + + // ── SELECT row_id + sort key with filter (QTF fetch phase) ── + + public void testIndexed_SelectRowIdSortKeyFiltered() throws IOException { + assertSelectRowIdSortKeyFiltered(STRATEGY_INDEXED); + } + + public void testListing_SelectRowIdSortKeyFiltered() throws IOException { + assertSelectRowIdSortKeyFiltered(STRATEGY_LISTING); + } + + private void assertSelectRowIdSortKeyFiltered(String strategy) throws IOException { + // Mimics: SELECT __row_id__, category FROM t WHERE category = 'B' ORDER BY value + List> rows = query(strategy, + "where category = 'B' | sort value | fields __row_id__, category"); + assertEquals(2, rows.size()); + assertEquals("B", rows.get(0).get(1)); + assertEquals("B", rows.get(1).get(1)); + assertExactIds(rows, 1, 3); + } + + // ── SELECT row_id only with LIMIT (minimal fetch) ── + + public void testIndexed_RowIdWithLimit() throws IOException { + assertRowIdWithLimit(STRATEGY_INDEXED); + } + + public void testListing_RowIdWithLimit() throws IOException { + assertRowIdWithLimit(STRATEGY_LISTING); + } + + private void assertRowIdWithLimit(String strategy) throws IOException { + // Mimics: SELECT __row_id__ FROM t ORDER BY value LIMIT 2 + List> rows = query(strategy, + "sort value | fields __row_id__ | head 2"); + assertEquals(2, rows.size()); + assertExactIds(rows, 0, 1); + } + + // ═══════════════════════════════════════════════════════════════════════════ + // Helpers + // ═══════════════════════════════════════════════════════════════════════════ + + private List> query(String strategy, String pplSuffix) throws IOException { + setup(); + setStrategy(strategy); + String ppl = "source = " + INDEX + " | " + pplSuffix; + logger.info("[QTF-IT] strategy={}, query: {}", strategy, ppl); + + Request req = new Request("POST", "/_analytics/ppl"); + req.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + Response resp = client().performRequest(req); + Map parsed = assertOkAndParse(resp, "PPL [" + strategy + "]"); + + @SuppressWarnings("unchecked") + List> rows = (List>) parsed.get("rows"); + assertNotNull("No rows in response", rows); + return rows; + } + + private List> query(String pplSuffix) throws IOException { + return query(STRATEGY_INDEXED, pplSuffix); + } + + private void setStrategy(String s) throws IOException { + Request req = new Request("PUT", "/_cluster/settings"); + req.setJsonEntity("{\"transient\":{\"datafusion.indexed.fetch_strategy\":\"" + s + "\"}}"); + client().performRequest(req); + } + + private void bulk(String body) throws IOException { + Request req = new Request("POST", "/" + INDEX + "/_bulk"); + req.setJsonEntity(body); + req.addParameter("refresh", "true"); + client().performRequest(req); + } + + // ── Assertion helpers ── + + private void assertRow(List row, String name, int value) { + assertEquals(name, row.get(1)); + assertNum(value, row.get(2)); + } + + private void assertGlobalIds(List> rows, long... expected) { + assertRowIdsUnique(rows); + Set actual = new HashSet<>(ids(rows)); + for (long id : expected) { + assertTrue("Missing expected row_id=" + id, actual.contains(id)); + } + } + + private void assertExactIds(List> rows, long... expected) { + List actual = ids(rows); + assertEquals(expected.length, actual.size()); + for (int i = 0; i < expected.length; i++) { + assertEquals("Row " + i + " id mismatch", Long.valueOf(expected[i]), actual.get(i)); + } + } + + private void assertRowIdsUnique(List> rows) { + Set seen = new HashSet<>(); + for (int i = 0; i < rows.size(); i++) { + Long id = toLong(rows.get(i).get(0)); + assertNotNull("Null row_id at row " + i, id); + assertTrue("Duplicate row_id " + id + " at row " + i, seen.add(id)); + } + } + + private List ids(List> rows) { + return rows.stream().map(r -> toLong(r.get(0))).collect(Collectors.toList()); + } + + private static Long toLong(Object v) { + if (v == null) return null; + if (v instanceof Number) return ((Number) v).longValue(); + return Long.parseLong(v.toString()); + } + + private static void assertNum(long expected, Object actual) { + assertNotNull(actual); + assertTrue(actual instanceof Number); + assertEquals(expected, ((Number) actual).longValue()); + } +} diff --git a/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java index 5e56a664847ef..2acbed3d6cacd 100644 --- a/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java +++ b/server/src/main/java/org/opensearch/index/engine/exec/coord/CatalogSnapshotManager.java @@ -110,6 +110,11 @@ public CatalogSnapshotManager( for (CatalogSnapshot cs : committedSnapshots) { catalogSnapshotMap.put(cs.getGeneration(), cs); } + + for(CatalogSnapshotLifecycleListener listener: this.snapshotListeners) { + listener.afterRefresh(true, latestCatalogSnapshot); + } + this.indexFileDeleter = new IndexFileDeleter( deletionPolicy, fileDeleter,