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 4e9c3a5bb8b7a..e7c6125df7320 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,11 @@ package org.opensearch.analytics.spi; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.opensearch.analytics.backend.EngineResultStream; +import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; + import java.util.Collections; import java.util.List; import java.util.Map; @@ -134,6 +139,20 @@ default Map getTopQueriesByMemory() { return Collections.emptyMap(); } + /** + * 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(Reader reader, 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/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSink.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSink.java index 43d21c7371a9c..ea98e0e74d701 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSink.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ExchangeSink.java @@ -32,6 +32,24 @@ public interface ExchangeSink { */ void feed(VectorSchemaRoot batch); + /** + * Ingest an Arrow batch with a per-source ordinal — the index of the + * producer task within its stage's resolved target list (e.g. + * {@link org.opensearch.analytics.spi.ExchangeSink} consumed via + * {@code ShardExecutionTarget.ordinal()}). + * + *

Default implementation drops the ordinal and falls through to + * {@link #feed(VectorSchemaRoot)}. Sinks that need to discriminate + * batches by producer (e.g. Late Materialization, where the ordinal is + * stamped onto each batch as a column) override this method. + * + *

Producers that have a meaningful per-task ordinal call this overload; + * producers without one continue to call {@link #feed(VectorSchemaRoot)}. + */ + default void feed(VectorSchemaRoot batch, int sourceOrdinal) { + feed(batch); + } + /** * Signal that no more batches will be fed. Releases resources. */ diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldStorageInfo.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldStorageInfo.java index 9fd96c235a15b..0b0d52e27c725 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldStorageInfo.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FieldStorageInfo.java @@ -10,6 +10,7 @@ import org.apache.calcite.sql.type.SqlTypeName; +import java.util.LinkedHashSet; import java.util.List; /** @@ -28,6 +29,7 @@ public class FieldStorageInfo { private final List indexFormats; private final List storedFieldFormats; private final boolean derived; + private final LinkedHashSet dependsOnPhysicalCols; public FieldStorageInfo( String fieldName, @@ -37,6 +39,21 @@ public FieldStorageInfo( List indexFormats, List storedFieldFormats, boolean derived + ) { + // Default: no physical-col dependencies. Physical fields aren't "derived from" + // anything; derived fields' deps are supplied by the caller via the 8-arg ctor. + this(fieldName, mappingType, fieldType, docValueFormats, indexFormats, storedFieldFormats, derived, new LinkedHashSet<>()); + } + + public FieldStorageInfo( + String fieldName, + String mappingType, + FieldType fieldType, + List docValueFormats, + List indexFormats, + List storedFieldFormats, + boolean derived, + LinkedHashSet dependsOnPhysicalCols ) { this.fieldName = fieldName; this.mappingType = mappingType; @@ -45,11 +62,21 @@ public FieldStorageInfo( this.indexFormats = indexFormats; this.storedFieldFormats = storedFieldFormats; this.derived = derived; + this.dependsOnPhysicalCols = dependsOnPhysicalCols; } - /** Creates a derived column (agg result, expression) with no physical storage. - * FieldType inferred from SqlTypeName. */ + /** Creates a derived column (agg result, expression) with no physical storage and no deps. + * FieldType inferred from SqlTypeName. Use {@link #derivedColumn(String, SqlTypeName, LinkedHashSet)} + * when the caller can supply the underlying physical-column dependencies. */ public static FieldStorageInfo derivedColumn(String fieldName, SqlTypeName sqlTypeName) { + return derivedColumn(fieldName, sqlTypeName, new LinkedHashSet<>()); + } + + /** Creates a derived column with explicit physical-column dependencies — the + * TableScan-level fields whose values flow into this column's computation, in + * first-appearance order. {@link LinkedHashSet} makes both the ordering invariant + * and the no-duplicates invariant explicit at the type level. */ + public static FieldStorageInfo derivedColumn(String fieldName, SqlTypeName sqlTypeName, LinkedHashSet dependsOnPhysicalCols) { return new FieldStorageInfo( fieldName, sqlTypeName.getName(), @@ -57,7 +84,8 @@ public static FieldStorageInfo derivedColumn(String fieldName, SqlTypeName sqlTy List.of(), List.of(), List.of(), - true + true, + dependsOnPhysicalCols ); } @@ -88,6 +116,29 @@ public boolean isDerived() { return derived; } + /** + * Names of the TableScan-level (physical) columns whose values flow into this column's + * computation, in first-appearance order. Empty for physical columns (they aren't + * "derived from" anything — their identity is their own field name) and for derived + * columns with no physical inputs (e.g. {@code COUNT(*)}, pure-literal projects). + * For other derived columns it is the union of underlying physical cols across the + * expression / agg-call tree. + * + *

Used by the QTF rewriter to derive the fetch list off the topmost operator's FSI + * without re-walking RexNodes from scratch. {@link LinkedHashSet} makes both the + * ordering invariant and the no-duplicates invariant explicit at the type level. + * + *

TODO: today we use string field names because they stay stable across narrowed-scan + * rewrites and (future) Join/Union plans where int ordinals across multiple TableScans + * become ambiguous. For very large plans the per-FSI string set can become a memory + * hotspot — consider switching to int ordinals (rooted to the originating TableScan's + * rowType) when single-scan throughput dominates and stability across rewrites can be + * traded off. + */ + public LinkedHashSet getDependsOnPhysicalCols() { + return dependsOnPhysicalCols; + } + /** * Resolves a field by index from a fieldStorageInfos list, validating bounds and field type. * Throws if the index is out of bounds or the field type is unrecognized. diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentConvertor.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentConvertor.java index 2dd7c7cf0cc6c..48564ea59d678 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentConvertor.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentConvertor.java @@ -9,6 +9,7 @@ package org.opensearch.analytics.spi; import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.type.RelDataType; /** * Fragment conversion API for backend plugins. @@ -81,4 +82,19 @@ default byte[] attachPartialAggOnTop(RelNode partialAggFragment, byte[] innerByt default byte[] attachFragmentOnTop(RelNode fragment, byte[] innerBytes) { throw new UnsupportedOperationException("attachFragmentOnTop not implemented for this backend"); } + + /** + * Builds a schema-only stub plan describing a stage's output partition: a single + * {@code Read { named_table: "input-"; base_schema: rowType }}, no + * operators above. Used for stages whose runtime is non-Substrait (e.g. QTF + * scatter-gather) but whose parent reduce sink still needs a partition schema for + * {@code registerPartitionStream}. + * + * @param childStageId stage id whose output schema this stub describes + * @param rowType the partition's row type + * @return serialized plan bytes + */ + default byte[] convertSchemaOnlyRead(int childStageId, RelDataType rowType) { + throw new UnsupportedOperationException("convertSchemaOnlyRead not implemented for this backend"); + } } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentInstructionHandlerFactory.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentInstructionHandlerFactory.java index f40d7472c2d4d..993f8a1c2f766 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentInstructionHandlerFactory.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/FragmentInstructionHandlerFactory.java @@ -25,8 +25,11 @@ public interface FragmentInstructionHandlerFactory { // ── Coordinator-side: create instruction nodes ── - /** Creates a shard scan instruction node. */ - Optional createShardScanNode(); + /** + * Creates a shard scan instruction node. {@code requestsRowIds} signals that the scan + * must emit shard-global {@code __row_id__} values (QTF query phase). + */ + Optional createShardScanNode(boolean requestsRowIds); /** Creates a filter delegation instruction node with the given delegation metadata. */ Optional createFilterDelegationNode( @@ -35,8 +38,17 @@ Optional createFilterDelegationNode( List delegatedQueries ); - /** Creates a shard scan with delegation instruction node — combines scan setup with delegation config. */ - Optional createShardScanWithDelegationNode(FilterTreeShape treeShape, int delegatedPredicateCount); + /** + * Creates a shard scan with delegation instruction node — combines scan setup with + * delegation config. {@code requestsRowIds} signals that the scan must emit shard-global + * {@code __row_id__} values (QTF query phase). Backends that don't support QTF should + * return {@link Optional#empty()} when {@code requestsRowIds} is true. + */ + Optional createShardScanWithDelegationNode( + FilterTreeShape treeShape, + int delegatedPredicateCount, + boolean requestsRowIds + ); /** Creates a partial aggregate instruction node. */ Optional createPartialAggregateNode(); diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanInstructionNode.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanInstructionNode.java index 8000d34f68844..121e597de7217 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanInstructionNode.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanInstructionNode.java @@ -10,21 +10,37 @@ import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.common.io.stream.Writeable; import java.io.IOException; /** * Instruction node for base shard scan setup — reader acquisition, SessionContext creation, - * default table provider registration. + * table provider registration. {@code requestsRowIds} signals that the shard scan needs to + * emit shard-global {@code __row_id__} values (QTF query phase). Inherited by + * {@link ShardScanWithDelegationInstructionNode} so the same flag applies whether or not + * filter delegation is in play — QTF and delegation are orthogonal concerns. * * @opensearch.internal */ -public class ShardScanInstructionNode implements InstructionNode { +public class ShardScanInstructionNode implements InstructionNode, Writeable { - public ShardScanInstructionNode() {} + private final boolean requestsRowIds; + + public ShardScanInstructionNode() { + this(false); + } + + public ShardScanInstructionNode(boolean requestsRowIds) { + this.requestsRowIds = requestsRowIds; + } public ShardScanInstructionNode(StreamInput in) throws IOException { - // No fields to read + this.requestsRowIds = in.readBoolean(); + } + + public boolean requestsRowIds() { + return requestsRowIds; } @Override @@ -34,6 +50,6 @@ public InstructionType type() { @Override public void writeTo(StreamOutput out) throws IOException { - // No fields to write + out.writeBoolean(requestsRowIds); } } diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanWithDelegationInstructionNode.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanWithDelegationInstructionNode.java index 18af354e02355..8feb8d55fe23b 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanWithDelegationInstructionNode.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ShardScanWithDelegationInstructionNode.java @@ -27,6 +27,11 @@ public class ShardScanWithDelegationInstructionNode extends ShardScanInstruction private final int delegatedPredicateCount; public ShardScanWithDelegationInstructionNode(FilterTreeShape treeShape, int delegatedPredicateCount) { + this(treeShape, delegatedPredicateCount, false); + } + + public ShardScanWithDelegationInstructionNode(FilterTreeShape treeShape, int delegatedPredicateCount, boolean requestsRowIds) { + super(requestsRowIds); this.treeShape = treeShape; this.delegatedPredicateCount = delegatedPredicateCount; } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml index 65f640aefc6c1..2ac8a892d0c8a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml +++ b/sandbox/plugins/analytics-backend-datafusion/rust/Cargo.toml @@ -107,3 +107,7 @@ harness = false [[bench]] name = "cross_rt_throughput_bench" harness = false + +[[bench]] +name = "row_id_bench" +harness = false diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/benches/query_bench.rs b/sandbox/plugins/analytics-backend-datafusion/rust/benches/query_bench.rs index 6408a32715799..30145e020bd8c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/benches/query_bench.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/benches/query_bench.rs @@ -5,7 +5,7 @@ use datafusion::execution::memory_pool::GreedyMemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use futures::TryStreamExt; use object_store::local::LocalFileSystem; -use object_store::ObjectStore; +use object_store::{ObjectStore, ObjectStoreExt}; use opensearch_datafusion::api::DataFusionRuntime; use opensearch_datafusion::query_executor; use opensearch_datafusion::runtime_manager::RuntimeManager; @@ -104,7 +104,16 @@ fn bench_execute_query(c: &mut Criterion) { let exec = mgr.cpu_executor(); async { let ptr = query_executor::execute_query( - url, metas, "t".into(), plan, &df_runtime, exec, None, &opensearch_datafusion::datafusion_query_config::DatafusionQueryConfig::test_default(), + url, + metas, + "t".into(), + plan, + &df_runtime, + exec, + None, + &opensearch_datafusion::datafusion_query_config::DatafusionQueryConfig::test_default(), + 0, + Arc::new(LocalFileSystem::new()) as Arc, ).await.unwrap(); // Consume and free the stream let mut stream = unsafe { @@ -151,6 +160,8 @@ fn bench_stream_next(c: &mut Criterion) { None, &opensearch_datafusion::datafusion_query_config::DatafusionQueryConfig::test_default( ), + 0, + Arc::new(LocalFileSystem::new()) as Arc, ) .await .unwrap(); @@ -199,6 +210,8 @@ fn bench_aggregation(c: &mut Criterion) { None, &opensearch_datafusion::datafusion_query_config::DatafusionQueryConfig::test_default( ), + 0, + Arc::new(LocalFileSystem::new()) as Arc, ) .await .unwrap(); 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..105e16a6ef026 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/benches/row_id_bench.rs @@ -0,0 +1,418 @@ +//! 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 { + writer_generation: 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, + 0, + Arc::new(LocalFileSystem::new()) as Arc, + ) + .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/proptest-regressions/native_node_stats.txt b/sandbox/plugins/analytics-backend-datafusion/rust/proptest-regressions/native_node_stats.txt new file mode 100644 index 0000000000000..356c96dbd910c --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/proptest-regressions/native_node_stats.txt @@ -0,0 +1,7 @@ +# Seeds for failure cases proptest has generated in the past. It is +# automatically read and these particular cases re-run before any +# novel cases are generated. +# +# It is recommended to check this file in to source control so that +# everyone who runs the test benefits from these saved cases. +cc 8e47fb8fb1db24906b9af022c409e9ea7c4518a4590689e586d3a6ca9bb3fb27 # shrinks to search_task_current_incs = 51, search_task_current_decs = 22, search_task_total_incs = 95, shard_task_current_incs = 13, shard_task_current_decs = 76, shard_task_total_incs = 25 diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs index aa0b74a1854c0..14944ff4ca1df 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; @@ -43,16 +44,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; @@ -62,6 +66,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. @@ -150,7 +155,102 @@ 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 { + debug_assert_eq!( + object_metas.len(), + file_metadata.len(), + "build_shard_files: object_metas and file_metadata must have matching length" + ); + 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 { @@ -174,6 +274,9 @@ pub struct ShardView { /// footers in production. Footer-kv reads, when they happen, are debug-only /// assertions. pub writer_generations: 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. @@ -358,6 +461,7 @@ pub fn create_reader( table_path: table_url, object_metas: Arc::new(object_metas), writer_generations: Arc::new(writer_generations), + file_metadata: None, store, }; Ok(Box::into_raw(Box::new(shard_view)) as i64) @@ -439,10 +543,10 @@ pub async unsafe fn execute_query( (cfg, corrector) }; - // 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. + // - is_indexed: index_filter UDF present (indexed query path) + // - has_row_id: __row_id__ column requested (QTF query phase) + let (is_indexed, has_row_id) = inspect_plan_bytes(plan_bytes); // Register cancellation token. let token = query_tracker::get_cancellation_token(context_id); @@ -452,8 +556,18 @@ pub async unsafe fn execute_query( // which are always data-node work. Keeping the coordinator ungated avoids deadlock // in single-JVM test topologies where coordinator and data node share a gate. - let query_future = async { - if is_indexed { + let query_future = async move { + // Routing logic: + // 1. Indexed query (has index_filter) → always indexed path + // 2. Has __row_id__ but not indexed (non-indexed + sort) → consult QueryStrategy + // - 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 && effective_config.query_strategy != crate::datafusion_query_config::QueryStrategy::ListingTable); + + if use_indexed { let qc = Arc::new(effective_config); crate::indexed_executor::execute_indexed_query( plan_bytes.to_vec(), @@ -501,6 +615,160 @@ 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 metadata_cache = ctx.state().runtime_env().cache_manager.get_file_metadata_cache(); + let (segments, _schema) = build_segments( + &ctx.state(), + Arc::clone(&store), + shard_view.object_metas.as_ref(), + shard_view.writer_generations.as_ref(), + metadata_cache, + ) + .await + .map_err(DataFusionError::Execution)?; + + // Distribute global row_ids to per-file local positions. + // Note: Java validates non-empty + ascending row_ids before the FFM call; we don't repeat that here. + debug_assert!(!segments.is_empty(), "fetch_by_row_ids: build_segments returned empty for non-empty shard view"); + let mut per_segment: HashMap = HashMap::new(); + for &gid in &row_ids { + debug_assert!(gid >= 0, "fetch_by_row_ids: negative row id {}", gid); + let seg_idx = segments + .partition_point(|s| s.global_base <= gid as u64) + .saturating_sub(1); + let seg = &segments[seg_idx]; + debug_assert!( + (gid as u64) >= seg.global_base && (gid as u64) < seg.global_base + seg.max_doc as u64, + "fetch_by_row_ids: row id {} out of bounds for segment {} (base={}, max_doc={})", + gid, seg_idx, seg.global_base, seg.max_doc + ); + let local_pos = (gid as u64 - seg.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_ord, seg) in segments.iter().enumerate() { + let num_rgs = seg.row_groups.len(); + let access_plan = if let Some(bm) = per_segment.get(&seg_ord) { + 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_ord].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 ── + // + // Caller-supplied `columns` includes __row_id__ in its desired position. We project + // each column verbatim except __row_id__, which we replace with the synthesized + // expression. This preserves caller column order and guarantees a single __row_id__ + // column in the result. + let projection = columns.iter() + .map(|c| { + if c == crate::ROW_ID_COLUMN_NAME { + format!("(\"{}\" + \"row_base\") AS \"{}\"", crate::ROW_ID_COLUMN_NAME, crate::ROW_ID_COLUMN_NAME) + } else { + format!("\"{}\"", c) + } + }) + .collect::>() + .join(", "); + let sql = format!("SELECT {} FROM t", projection); + let df = ctx.sql(&sql).await?; + let physical_plan = df.create_physical_plan().await?; + let df_stream = execute_stream(physical_plan, ctx.task_ctx())?; + + // Post-condition: returned stream schema must contain __row_id__ plus every requested column. + // Catches drift if SQL synthesis or the optimizer ever drops a projection silently. + debug_assert!(assert_fetch_result_schema(df_stream.schema().as_ref(), &columns)); + + // ── 5. Wrap and return ── + + // In debug builds, interpose an adapter that asserts __row_id__ values are + // monotonically nondecreasing across the entire stream. target_partitions=1 + // means a single ordered execution, so the check is global, not per-batch only. + let df_stream = ascending_row_id_check_stream(df_stream); + Ok(wrap_stream_as_handle(df_stream, manager.cpu_executor(), runtime)) +} + /// it; the failure mode is documented here to keep the dispatch contract /// explicit. /// Resolve the dynamic spill limit based on available disk space. @@ -531,9 +799,15 @@ fn resolve_dynamic_spill_limit(spill_dir: &str) -> u64 { } } -fn plan_bytes_mentions_index_filter(plan_bytes: &[u8]) -> bool { - 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), + ) } /// Best-effort budget acquisition from cached parquet metadata. @@ -1504,3 +1778,67 @@ pub unsafe fn register_memtable( session.register_memtable(input_id, table_schema, batches)?; Ok(schema_ipc) } + +// ── QTF fetch-phase assertion helpers (kept at the bottom of the file) ──────── + +/// Wraps a stream in an adapter that asserts each emitted batch's `__row_id__` column +/// is monotonically nondecreasing, including across batch boundaries. No-op in release. +fn ascending_row_id_check_stream( + stream: datafusion::execution::SendableRecordBatchStream, +) -> datafusion::execution::SendableRecordBatchStream { + if !cfg!(debug_assertions) { + return stream; + } + use arrow_array::Int64Array; + use futures::StreamExt; + let schema = stream.schema(); + let row_id_idx = match schema.column_with_name(crate::ROW_ID_COLUMN_NAME) { + Some((idx, _)) => idx, + None => return stream, // schema-presence checked elsewhere; nothing to validate here + }; + let mut last_seen: Option = None; + let checked = stream.map(move |batch_res| { + let batch = match batch_res { + Ok(b) => b, + Err(e) => return Err(e), + }; + let col = batch.column(row_id_idx); + let arr = col + .as_any() + .downcast_ref::() + .expect("ascending_row_id_check_stream: __row_id__ column must be Int64"); + for i in 0..arr.len() { + if arr.is_null(i) { + continue; + } + let v = arr.value(i); + if let Some(prev) = last_seen { + if v < prev { + panic!( + "fetch_by_row_ids: __row_id__ not ascending — prev={}, next={} at row {}", + prev, v, i + ); + } + } + last_seen = Some(v); + } + Ok(batch) + }); + Box::pin(RecordBatchStreamAdapter::new(schema, checked)) +} + +/// Verify the fetch-result schema carries `__row_id__` plus every requested column. +/// Body only runs under `debug_assertions` (called from a `debug_assert!`). +fn assert_fetch_result_schema(schema: &datafusion::arrow::datatypes::Schema, columns: &[String]) -> bool { + if schema.column_with_name(crate::ROW_ID_COLUMN_NAME).is_none() { + let names: Vec = schema.fields().iter().map(|f| f.name().clone()).collect(); + panic!("fetch_by_row_ids: result schema missing {}, got {:?}", crate::ROW_ID_COLUMN_NAME, names); + } + for col in columns { + if schema.column_with_name(col).is_none() { + let names: Vec = schema.fields().iter().map(|f| f.name().clone()).collect(); + panic!("fetch_by_row_ids: result schema missing requested column {}, got {:?}", col, names); + } + } + true +} 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..a50fd0b175268 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 QueryStrategy { + /// 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 query_strategy: QueryStrategy, } /// 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 query_strategy: i32, } impl DatafusionQueryConfig { @@ -97,6 +123,7 @@ impl DatafusionQueryConfig { max_collector_parallelism: 1, single_collector_strategy: CollectorCallStrategy::PageRangeSplit, tree_collector_strategy: CollectorCallStrategy::TightenOuterBounds, + query_strategy: QueryStrategy::None, } } @@ -162,6 +189,11 @@ impl DatafusionQueryConfig { 2 => CollectorCallStrategy::PageRangeSplit, _ => CollectorCallStrategy::TightenOuterBounds, }, + query_strategy: match w.query_strategy { + 1 => QueryStrategy::ListingTable, + 2 => QueryStrategy::IndexedPredicateOnly, + _ => QueryStrategy::None, + }, } } } @@ -272,6 +304,7 @@ mod tests { max_collector_parallelism: 4, single_collector_strategy: 2, tree_collector_strategy: 1, + query_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.query_strategy, QueryStrategy::ListingTable); } #[test] @@ -303,10 +337,12 @@ mod tests { max_collector_parallelism: 2, single_collector_strategy: 2, tree_collector_strategy: 1, + query_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.query_strategy, QueryStrategy::None); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index 858369c0cc988..7f83d64021fbc 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -262,6 +262,67 @@ 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 { + // Hard FFM-boundary checks (UB risk if violated): pointers must be non-zero before any deref. + // Always-on `assert!` (not debug_assert!) — these protect against use-after-close from Java. + assert!(shard_view_ptr != 0, "df_fetch_by_row_ids: shard_view_ptr is null"); + assert!(runtime_ptr != 0, "df_fetch_by_row_ids: runtime_ptr is null"); + assert!(row_ids_count >= 0, "df_fetch_by_row_ids: negative row_ids_count {}", row_ids_count); + assert!(col_names_count >= 0, "df_fetch_by_row_ids: negative col_names_count {}", col_names_count); + if row_ids_count > 0 { + assert!(row_ids_ptr != 0, "df_fetch_by_row_ids: row_ids_ptr is null but count={}", row_ids_count); + } + if col_names_count > 0 { + assert!(!col_names_ptr.is_null(), "df_fetch_by_row_ids: col_names_ptr is null but count={}", col_names_count); + assert!(!col_names_len_ptr.is_null(), "df_fetch_by_row_ids: col_names_len_ptr is null but count={}", col_names_count); + } + + 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 { @@ -460,16 +521,14 @@ pub unsafe extern "C" fn df_execute_local_plan( // The IO runtime still drives the outer block_on (bridging the synchronous FFI // call to the async spawn handle). timed_block_on(&mgr.io_runtime, "execute_local_plan", crate::task_monitors::coordinator_reduce_monitor().instrument(async move { - // Acquire coordinator gate on IO runtime BEFORE spawning on CPU. - // This blocks the Java search thread when the gate is full. - let coord_gate = mgr_for_spawn.coordinator_gate().clone(); - let partition_weight = (num_cpus::get() as u32).max(1); - let permit = coord_gate.acquire_many(partition_weight.min(coord_gate.max_permits())).await; - - + // No coordinator-gate acquire here. The QTF coordinator-reduce code path runs + // synchronously inside the SEARCH-thread FFM call (DatafusionReduceSink.); + // gating it would deadlock when the gate is contended because the SEARCH thread + // is blocked waiting for permits its own work would release. Keep the gate + // exclusively on the data-node FFM entry points. let inner_fut = async move { unsafe { - api::execute_local_plan(session_ptr, &bytes_vec, &mgr_for_inner, context_id, Some(permit)) + api::execute_local_plan(session_ptr, &bytes_vec, &mgr_for_inner, context_id, None) .await } }; @@ -686,6 +745,7 @@ pub unsafe extern "C" fn df_create_session_context_indexed( context_id: i64, tree_shape: i32, delegated_predicate_count: i32, + requests_row_ids: u8, query_config_ptr: i64, plan_ptr: *const u8, plan_len: i64, @@ -703,7 +763,15 @@ pub unsafe extern "C" fn df_create_session_context_indexed( mgr.io_runtime .block_on(crate::task_monitors::plan_setup_monitor().instrument( crate::session_context::create_session_context_indexed( - runtime_ptr, shard_view_ptr, table_name, context_id, tree_shape, delegated_predicate_count, query_config, plan_bytes, + runtime_ptr, + shard_view_ptr, + table_name, + context_id, + tree_shape, + delegated_predicate_count, + requests_row_ids != 0, + query_config, + plan_bytes, ) )) .map_err(|e| e.to_string()) @@ -869,8 +937,16 @@ pub unsafe extern "C" fn df_execute_with_context( let cpu_for_cross = cpu_executor.clone(); let mgr_for_spawn = Arc::clone(&mgr); - // Route based on whether the session was configured for indexed execution - if session_handle.indexed_config.is_some() { + // Route based on whether the session was configured for indexed execution, + // or if the plan projects __row_id__ (QTF query phase) under a non-ListingTable + // fetch strategy. + let has_row_id = plan_bytes + .windows(crate::ROW_ID_COLUMN_NAME.len()) + .any(|w| w == crate::ROW_ID_COLUMN_NAME.as_bytes()); + let query_strategy = session_handle.query_config.query_strategy; + let use_indexed = session_handle.indexed_config.is_some() + || (has_row_id && query_strategy != crate::datafusion_query_config::QueryStrategy::ListingTable); + if use_indexed { // Extract target_partitions BEFORE boxing into raw pointer (session_handle is consumed). let partition_weight = session_handle.query_config.target_partitions.max(1) as u32; // TODO: refactor execute_indexed_with_context to take SessionContextHandle directly @@ -1068,14 +1144,8 @@ pub unsafe extern "C" fn df_execute_local_prepared_plan( context_id: i64, ) -> i64 { let mgr = get_rt_manager()?; - // Acquire coordinator concurrency gate before executing the prepared plan. - // Gate is acquired on the IO runtime (block_on) so the Java search thread - // blocks here when the gate is full — creating backpressure at the threadpool level. - let partition_weight = (num_cpus::get() as u32).max(1); - let coord_gate = mgr.coordinator_gate().clone(); - let permit = mgr.io_runtime.block_on( - coord_gate.acquire_many(partition_weight.min(coord_gate.max_permits())) - ); - - api::execute_local_prepared_plan(session_ptr, &mgr, context_id, Some(permit)).map_err(|e| e.to_string()) + // No coordinator-gate acquire here — see df_execute_local_plan for the rationale + // (the QTF coordinator-reduce path runs synchronously inside the SEARCH-thread FFM + // call and gating it can deadlock). + api::execute_local_prepared_plan(session_ptr, &mgr, context_id, None).map_err(|e| e.to_string()) } 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 59ee129b2a563..66f43d483f96e 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, 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 @@ -440,6 +441,8 @@ pub async unsafe fn execute_indexed_with_context( // spawning on the CPU runtime, so the Java search thread blocks at the // gate when it is full — creating backpressure at the Java threadpool level. + // Java-side QTF signal: scan must emit __row_id__. Captured before consuming indexed_config below. + let requests_row_ids = handle.indexed_config.as_ref().is_some_and(|c| c.requests_row_ids); let classification_override = handle.indexed_config.map(|config| { // FilterTreeShape: 1 = CONJUNCTIVE → SingleCollector, 2 = INTERLEAVED → Tree. match (config.tree_shape, config.delegated_predicate_count) { @@ -481,8 +484,6 @@ pub async unsafe fn execute_indexed_with_context( .await .map_err(DataFusionError::Execution)?; let schema = crate::schema_coerce::coerce_inferred_schema(schema); - for (i, seg) in segments.iter().enumerate() { - } let placeholder: Arc = Arc::new(PlaceholderProvider { schema: schema.clone(), @@ -493,6 +494,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 = requests_row_ids; let filter_expr = extract_filter_expr(&logical_plan); let extraction = match filter_expr { None => None, @@ -510,7 +512,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 @@ -532,6 +533,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, }; @@ -539,9 +548,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(|| { @@ -749,6 +790,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, @@ -758,6 +800,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 5114561257cea..57aaf2af94cb6 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 cbdcf66380db7..d7a407d83a649 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 @@ -460,11 +460,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(|| { @@ -519,27 +515,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 = super::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)) 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..08ff4ebb6f82f --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/row_id_injection.rs @@ -0,0 +1,140 @@ +/* + * 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, + }; + let id = base + rg_pos as u64; + debug_assert!(id >= base, "position_to_global_id: underflow base={} rg_pos={}", base, rg_pos); + id +} 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 da891f7edfdc7..d6cb97b343c6b 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 @@ -66,6 +66,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 @@ -93,6 +94,8 @@ pub async fn build_segments( offset += num_rows; } let max_doc = offset; + let global_base = cumulative_rows; + cumulative_rows += max_doc as u64; let writer_generation = writer_generations[seg_ord]; @@ -109,6 +112,7 @@ pub async fn build_segments( parquet_size: size, row_groups, metadata: pq_meta, + global_base, }); } @@ -121,10 +125,9 @@ pub async fn build_segments( // primitives (recursively merged for Struct/List/Union). // 5. Apply `binary_as_string` and `force_view_types` transforms // if configured. - // Disable Utf8View: the indexed path uses the file's physical schema for - // ParquetSource (to avoid column reordering issues), so the inferred table - // schema must also use Utf8 (not Utf8View) to stay consistent. - let format = ParquetFormat::default().with_force_view_types(false); + // Use Utf8View — ParquetOpener's apply_file_schema_type_coercions keeps the file/table + // schemas aligned, so QTF's coordinator-declared Utf8View matches the produced batches. + let format = ParquetFormat::default().with_force_view_types(true); let schema = FileFormat::infer_schema(&format, state, &store, object_metas) .await .map_err(|e| format!("infer_schema union: {}", e))?; 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 156270580f335..25316c0a529bb 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 @@ -284,6 +285,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 { @@ -376,6 +384,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, ))) } } @@ -438,6 +449,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 { @@ -460,9 +477,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, @@ -494,6 +515,9 @@ impl IndexedStream { batch_coalescer, upstream_done: false, coalescer_finished: false, + global_base, + emit_row_ids, + row_id_output_index, } } @@ -572,6 +596,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; @@ -611,11 +647,23 @@ impl IndexedStream { } }; - // Reorder/strip columns to match output schema. The parquet reader - // delivers columns in the file's physical order which may differ from - // the table schema order (e.g. when infer_schema sorted alphabetically). + // Inject computed __row_id__, or reorder/strip columns to match output schema. + // The parquet reader delivers columns in the file's physical order which may + // differ from the table schema order (e.g. when infer_schema sorted alphabetically). let t_proj = Instant::now(); - let output = if output.schema().as_ref() == self.schema.as_ref() { + 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.schema().as_ref() == self.schema.as_ref() { output } else { let n = self.schema.fields().len(); 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 68e4ba16ab8af..f4735f4dc0252 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 @@ -54,6 +54,18 @@ pub const COLLECTOR_FUNCTION_NAME: &str = "delegated_predicate"; /// DF's own pruning isn't selective enough for a row group). pub const DELEGATION_POSSIBLE_FUNCTION_NAME: &str = "delegation_possible"; +/// 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 35702bb9d6311..4682042111a8b 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; @@ -65,6 +65,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`. @@ -124,6 +127,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. @@ -181,13 +188,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| { @@ -201,6 +250,7 @@ impl TableProvider for IndexedTableProvider { cols }) }; + let projected_schema = output_schema; // Ignore DataFusion's `filters` argument. The `index_filter(...)` @@ -243,6 +293,7 @@ impl TableProvider for IndexedTableProvider { predicate, metrics: ExecutionPlanMetricsSet::new(), inner_parquet_metrics: Arc::new(std::sync::Mutex::new(Vec::new())), + row_id_output_index, })) } @@ -268,6 +319,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 { @@ -389,6 +443,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)); } @@ -450,6 +507,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/delegation.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/delegation.rs index 764594af9eab3..599971432c362 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/delegation.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_table/tests_e2e/fuzz/delegation.rs @@ -468,6 +468,7 @@ pub(in crate::indexed_table::tests_e2e) async fn execute_delegation_tree( pushdown_predicate: Some(Arc::clone(&residual_physical)), query_config: Arc::new(qc), predicate_columns: pred_cols, + emit_row_ids: false, })); let ctx = SessionContext::new(); 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 c12790956a00b..1a7f8583022cb 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; } @@ -289,6 +290,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(); @@ -468,6 +470,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(); @@ -675,7 +678,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 dce47e5033477..681fc6c266008 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 b175641622553..cb9b53840a772 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, }); } @@ -181,6 +182,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(); @@ -335,6 +337,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, }); } @@ -382,6 +385,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(); @@ -821,6 +825,7 @@ async fn run_wide_segments( parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }); } @@ -883,6 +888,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 d411542653dfa..9674000c51b11 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 @@ -333,6 +333,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); @@ -380,6 +381,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 763ff1cfb7d4c..8415fdd2a3ecc 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) } @@ -390,6 +391,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..35d71c045ee13 --- /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 { + writer_generation: 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..a01c0b26aed90 --- /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 { + writer_generation: 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 { + writer_generation: 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 { + writer_generation: 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 { + writer_generation: 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 { + writer_generation: 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..7a7d8dc9fea53 --- /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 `QueryStrategy` 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::QueryStrategy; + 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_query_strategy_default_is_none() { + let config = crate::datafusion_query_config::DatafusionQueryConfig::test_default(); + assert_eq!(config.query_strategy, QueryStrategy::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 481207b53f970..e35fed2e83443 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(); @@ -401,6 +403,7 @@ async fn query_with_mismatched_schema( parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }; let tree = Arc::new(tree_bool); let factory: super::super::table_provider::EvaluatorFactory = { @@ -440,6 +443,7 @@ async fn query_with_mismatched_schema( 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 b3b916497d47e..d70ed285e5db9 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 @@ -406,6 +406,7 @@ async fn run_large( parquet_size: size, row_groups: rgs, metadata: Arc::clone(&parquet_meta), + global_base: 0, }; let tree = Arc::new(tree); @@ -452,6 +453,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(); @@ -856,6 +858,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); @@ -901,6 +904,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 31ee9a9a426d0..29901e5c09287 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs @@ -11,6 +11,15 @@ //! 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. +/// +/// TODO: source this from Java — the canonical column name should be defined +/// once on the coordinator side and passed across FFM rather than hardcoded +/// here. Today both sides hardcode "__row_id__" and rely on the strings +/// matching by convention. +pub const ROW_ID_COLUMN_NAME: &str = "__row_id__"; + pub(crate) mod agg_mode; pub mod api; pub mod cache; @@ -30,10 +39,13 @@ pub mod memory_guard; pub mod native_error; pub mod partition_stream; pub mod phantom_corrector; +pub mod project_row_id_analyzer; +pub mod project_row_id_optimizer; pub mod query_budget; pub mod query_executor; pub mod query_tracker; pub mod relabel_exec; +pub mod shard_table_provider; pub mod runtime_manager; pub mod schema_coerce; pub mod session_context; 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 57a04c16c6ee0..8ed7843276955 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. @@ -54,43 +54,24 @@ pub async fn execute_query( shard_store: Arc, phantom_corrector: Option>, ) -> 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_metadata_cache_limit( - runtime.runtime_env.cache_manager.get_metadata_cache_limit(), - ) - .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. @@ -107,7 +88,7 @@ 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(); @@ -115,40 +96,62 @@ pub async fn execute_query( crate::udf::register_all(&ctx); crate::udaf::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); - - // Wire the global statistics cache into the ListingTable. - let stats_cache = runtime.runtime_env.cache_manager.get_file_statistic_cache(); - let provider = Arc::new( - ListingTable::try_new(table_config) - .map_err(|e| { - error!("Failed to create listing table: {}", e); - e - })? - .with_cache(stats_cache), - ); - - ctx.register_table(&table_name, provider).map_err(|e| { - error!("Failed to register table: {}", e); - e - })?; + // Register table provider based on strategy. + // + // Note: api::execute_query only routes to this function when the plan does NOT + // request row IDs (otherwise it dispatches to the indexed executor). The strategy + // therefore matters only for distinguishing the ShardTableProvider rewrite + // (ListingTable) from the plain ListingTable scan (None / IndexedPredicateOnly). + use crate::datafusion_query_config::QueryStrategy; + match query_config.query_strategy { + QueryStrategy::ListingTable => { + use crate::shard_table_provider::{ShardTableConfig, ShardTableProvider}; + + // 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: 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 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 })?); + 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| { @@ -166,6 +169,19 @@ pub async fn execute_query( let target_schema = crate::schema_coerce::coerce_inferred_schema(physical_plan.schema()); let physical_plan = crate::relabel_exec::wrap_if_relabel_needed(physical_plan, target_schema)?; + // Apply row ID optimizer when ShardTableProvider injected `row_base`. + // For other strategies the vanilla scan output is already what the plan expects. + use datafusion::physical_optimizer::PhysicalOptimizerRule; + let physical_plan = match query_config.query_strategy { + QueryStrategy::ListingTable => { + // 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)? + } + _ => physical_plan, + }; + let df_stream = execute_stream(physical_plan, ctx.task_ctx()).map_err(|e| { error!("Failed to create execution stream: {}", e); e @@ -199,6 +215,11 @@ 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 `QueryStrategy` 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], @@ -208,10 +229,45 @@ pub async fn execute_with_context( // Permit was acquired by the caller (ffm.rs) on the IO runtime before // spawning on the CPU runtime, so the Java search thread blocks at the // gate when it is full — creating backpressure at the Java threadpool level. + use crate::datafusion_query_config::QueryStrategy; let context_id = handle.query_context.context_id(); let token = crate::query_tracker::get_cancellation_token(context_id); + let query_strategy = handle.query_config.query_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 query_strategy == QueryStrategy::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?; + let resolved_schema = crate::schema_coerce::coerce_inferred_schema(resolved_schema); + + // 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 query_future = async { let substrait_plan = Plan::decode(plan_bytes).map_err(|e| { DataFusionError::Execution(format!("Failed to decode Substrait: {}", e)) @@ -221,6 +277,8 @@ pub async fn execute_with_context( let logical_plan = from_substrait_plan(&handle.ctx.state(), &substrait_plan).await?; log_debug!("DataFusion logical plan:\n{}", logical_plan.display_indent()); 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?; let target_schema = crate::schema_coerce::coerce_inferred_schema(physical_plan.schema()); let physical_plan = crate::relabel_exec::wrap_if_relabel_needed(physical_plan, target_schema)?; @@ -257,3 +315,97 @@ pub async fn execute_with_context( let stream_handle = crate::api::QueryStreamHandle::with_session_context(stream, handle.query_context, handle.ctx, Some(permit)); 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, None); + 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 1b280c451ab06..36b30956ec322 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/session_context.rs @@ -60,6 +60,8 @@ pub struct SessionContextHandle { pub struct IndexedExecutionConfig { pub tree_shape: i32, pub delegated_predicate_count: i32, + /// QTF query phase: scan must emit shard-global `__row_id__`. + pub requests_row_ids: bool, } /// Widens `inferred` to the plan's `base_schema` (for index-pattern / alias scans) so the @@ -200,12 +202,26 @@ pub async unsafe fn create_session_context( config.options_mut().execution.target_partitions = effective_partitions; config.options_mut().execution.batch_size = effective_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 query strategy: + // 1. Add ProjectRowIdAnalyzer (logical) — ensures __row_id__ survives pruning. + // 2. Add ProjectRowIdOptimizer (physical) — computes __row_id__ + row_base. + if query_config.query_strategy == crate::datafusion_query_config::QueryStrategy::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 (parse, item, mvappend, mvfind, mvzip, convert_tz, …) @@ -316,6 +332,7 @@ pub async unsafe fn create_session_context_indexed( context_id: i64, tree_shape: i32, delegated_predicate_count: i32, + requests_row_ids: bool, query_config: DatafusionQueryConfig, plan_bytes: &[u8], ) -> Result { @@ -328,6 +345,7 @@ pub async unsafe fn create_session_context_indexed( handle.indexed_config = Some(IndexedExecutionConfig { tree_shape, delegated_predicate_count, + requests_row_ids, }); Ok(ptr) 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..be8c60177181d --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/shard_table_provider.rs @@ -0,0 +1,148 @@ +/* + * 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> { + // Invariant: files are ordered by ascending row_base (build_shard_files contract). + // ProjectRowIdOptimizer relies on the partition value matching the file's true offset. + debug_assert!( + self.config.files.windows(2).all(|w| w[0].row_base <= w[1].row_base), + "ShardTableProvider: files not ordered by row_base — ProjectRowIdOptimizer would compute wrong global IDs" + ); + 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 725f2ffa7758e..e4299ab66fa86 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 @@ -8,15 +8,19 @@ package org.opensearch.be.datafusion; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.BigIntVector; import org.apache.calcite.sql.SqlOperator; import org.apache.calcite.sql.fun.SqlLibraryOperators; import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.spi.AbstractNameMappingAdapter; import org.opensearch.analytics.spi.AggregateCapability; import org.opensearch.analytics.spi.AggregateFunction; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.BackendCapabilityProvider; import org.opensearch.analytics.spi.BackendExecutionContext; +import org.opensearch.analytics.spi.DelegationThreadTracker; import org.opensearch.analytics.spi.DelegationType; import org.opensearch.analytics.spi.EngineCapability; import org.opensearch.analytics.spi.ExchangeSink; @@ -39,9 +43,12 @@ import org.opensearch.analytics.spi.WindowCapability; import org.opensearch.analytics.spi.WindowFunction; import org.opensearch.be.datafusion.indexfilter.FilterTreeCallbacks; +import org.opensearch.be.datafusion.nativelib.NativeBridge; +import org.opensearch.be.datafusion.nativelib.StreamHandle; import org.opensearch.be.datafusion.planner.adapter.NumericConversionFunctionAdapter; import org.opensearch.be.datafusion.planner.adapter.TimeConversionFunctionAdapter; import org.opensearch.index.engine.dataformat.DataFormatRegistry; +import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; import java.util.HashSet; import java.util.Map; @@ -810,7 +817,45 @@ public Map getTopQueriesByMemory() { } @Override - public void setDelegationThreadTracker(org.opensearch.analytics.spi.DelegationThreadTracker tracker) { + public EngineResultStream fetchByRowIds(Reader reader, BigIntVector rowIdVector, String[] columns, 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 = NativeBridge.fetchByRowIds( + dfReader.getReaderHandle().getPointer(), + bufAddr, + count, + columns, + dataFusionService.getNativeRuntime().get() + ); + } else { + throw new IllegalStateException("BigIntVector buffer address is 0 or count is 0"); + } + StreamHandle streamHandle = new StreamHandle(streamPtr, dataFusionService.getNativeRuntime()); + return new DatafusionResultStream(streamHandle, allocator); + } + + @Override + public void setDelegationThreadTracker(DelegationThreadTracker tracker) { FilterTreeCallbacks.setThreadTracker(tracker); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java index 4795bdb005eac..c378243f50f6c 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionFragmentConvertor.java @@ -55,6 +55,7 @@ import io.substrait.expression.Expression; import io.substrait.expression.FunctionArg; import io.substrait.expression.ImmutableAggregateFunctionInvocation; +import io.substrait.extension.ExtensionCollector; import io.substrait.extension.SimpleExtension; import io.substrait.isthmus.ConverterProvider; import io.substrait.isthmus.SubstraitRelVisitor; @@ -66,6 +67,8 @@ import io.substrait.plan.Plan; import io.substrait.plan.PlanProtoConverter; import io.substrait.plan.ProtoPlanConverter; +import io.substrait.proto.PlanRel; +import io.substrait.proto.ReadRel; import io.substrait.relation.Aggregate; import io.substrait.relation.Fetch; import io.substrait.relation.Filter; @@ -73,7 +76,9 @@ import io.substrait.relation.Project; import io.substrait.relation.Rel; import io.substrait.relation.Sort; +import io.substrait.type.NamedStruct; import io.substrait.type.Type; +import io.substrait.type.proto.TypeProtoConverter; /** * Converts Calcite RelNode fragments to Substrait protobuf bytes @@ -401,6 +406,15 @@ public class DataFusionFragmentConvertor implements FragmentConvertor { FunctionMappings.s(LOCAL_LIST_MERGE_DISTINCT_OP, "list_merge_distinct") ); + /** + * Shared {@link TypeProtoConverter} for schema-only conversions. Safe as a singleton + * because schema-only Reads convert primitive Calcite types to primitive Substrait + * protos — no functions or user-defined types touch the inner {@link ExtensionCollector}, + * so it never accumulates per-call state. Avoids re-allocating both objects on every + * {@link #convertSchemaOnlyRead} call. + */ + private static final TypeProtoConverter SCHEMA_ONLY_TYPE_PROTO_CONVERTER = new TypeProtoConverter(new ExtensionCollector()); + private final SimpleExtension.ExtensionCollection extensions; public DataFusionFragmentConvertor(SimpleExtension.ExtensionCollection extensions) { @@ -431,6 +445,43 @@ public byte[] attachPartialAggOnTop(RelNode partialAggFragment, byte[] innerByte return serializePlan(SubstraitPlanRewriter.rewrite(rewired)); } + /** + * Builds a schema-only stub plan directly via Substrait protos — no isthmus, no + * Calcite RelNode round-trip. Output: + * 

+     *   Plan { relations: [PlanRel { Root { input: Rel { Read { named_table: "input-<id>";
+     *                                                          base_schema: rowType } },
+     *                                 names: rowType.fieldNames }}] }
+     *
+ * + *

Used by the LM stage path: LM runs Java-only scatter/gather/stitch and emits no + * Substrait compute, but the parent reduce sink (Stage 3) still calls + * {@code registerPartitionStream} which needs the partition's named-table id and base + * schema. This stub is the minimum proto that satisfies that path. Bypassing isthmus + * avoids unnecessary {@code SubstraitRelVisitor} setup and keeps the produced bytes + * tightly scoped to the schema we care about. + */ + @Override + public byte[] convertSchemaOnlyRead(int childStageId, RelDataType rowType) { + // Fully-qualified names below: io.substrait.proto.{Plan,Rel,NamedStruct,RelRoot} clash with already-imported single-name imports. + NamedStruct ns = TypeConverter.DEFAULT.toNamedStruct(rowType); + io.substrait.proto.NamedStruct nsProto = ns.toProto(SCHEMA_ONLY_TYPE_PROTO_CONVERTER); + + ReadRel readRel = ReadRel.newBuilder() + .setNamedTable(ReadRel.NamedTable.newBuilder().addNames("input-" + childStageId).build()) + .setBaseSchema(nsProto) + .build(); + + io.substrait.proto.Rel inputRel = io.substrait.proto.Rel.newBuilder().setRead(readRel).build(); + PlanRel planRel = PlanRel.newBuilder() + .setRoot(io.substrait.proto.RelRoot.newBuilder().setInput(inputRel).addAllNames(rowType.getFieldNames()).build()) + .build(); + + byte[] bytes = io.substrait.proto.Plan.newBuilder().addRelations(planRel).build().toByteArray(); + LOGGER.debug("Schema-only Read for stage [{}]: {} bytes", childStageId, bytes.length); + return bytes; + } + @Override public byte[] attachFragmentOnTop(RelNode fragment, byte[] innerBytes) { LOGGER.debug("Attaching generic fragment [{}] on top of {} inner bytes", fragment.getClass().getSimpleName(), innerBytes.length); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionInstructionHandlerFactory.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionInstructionHandlerFactory.java index 737a0540b531e..2ab4bb1a0f8ac 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionInstructionHandlerFactory.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DataFusionInstructionHandlerFactory.java @@ -39,8 +39,8 @@ public DataFusionInstructionHandlerFactory(DataFusionPlugin plugin) { // ── Coordinator: create instruction nodes ── @Override - public Optional createShardScanNode() { - return Optional.of(new ShardScanInstructionNode()); + public Optional createShardScanNode(boolean requestsRowIds) { + return Optional.of(new ShardScanInstructionNode(requestsRowIds)); } @Override @@ -53,8 +53,12 @@ public Optional createFilterDelegationNode( } @Override - public Optional createShardScanWithDelegationNode(FilterTreeShape treeShape, int delegatedPredicateCount) { - return Optional.of(new ShardScanWithDelegationInstructionNode(treeShape, delegatedPredicateCount)); + public Optional createShardScanWithDelegationNode( + FilterTreeShape treeShape, + int delegatedPredicateCount, + boolean requestsRowIds + ) { + return Optional.of(new ShardScanWithDelegationInstructionNode(treeShape, delegatedPredicateCount, requestsRowIds)); } @Override diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java index 180059f22d614..1a47b903329f2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionReduceSink.java @@ -43,6 +43,13 @@ * *

Cleanup ownership lives in {@link #reduce}'s {@code finally} (via {@link SinkState}), * not {@link #close}, so a close call from another thread never races a parked drain. + * + *

TODO abstraction leak: this class implements {@link MultiInputExchangeSink} unconditionally + * even when only one child stage feeds it. The marker is meant for genuine multi-input shapes + * (Union/Join), and callers like {@code ReduceStageExecution.inputSink} have to dispatch on + * the logical child-stage count instead of the marker. Either split into a single-input + * subclass and a multi-input subclass, or drop the marker and let the caller always go through + * {@code feed()} when there's one child. Current behaviour is correct but the typing lies. */ public class DatafusionReduceSink extends AbstractDatafusionReduceSink implements MultiInputExchangeSink { 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 b3f8b7ebd7b0d..2e58008934f7c 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 @@ -182,6 +182,47 @@ public final class DatafusionSettings { Setting.Property.NodeScope ); + // Query strategy constants + public static final String QUERY_STRATEGY_NONE = "none"; + public static final String QUERY_STRATEGY_LISTING_TABLE = "listing_table"; + public static final String QUERY_STRATEGY_INDEXED = "indexed"; + + /** + * Query strategy for query-then-fetch (QTF) row ID computation. + *

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

+ * Default: {@code indexed}. + */ + public static final Setting INDEXED_QUERY_STRATEGY = Setting.simpleString( + "datafusion.indexed.query_strategy", + QUERY_STRATEGY_INDEXED, + value -> { + switch (value) { + case QUERY_STRATEGY_NONE: + case QUERY_STRATEGY_LISTING_TABLE: + case QUERY_STRATEGY_INDEXED: + break; + default: + throw new IllegalArgumentException( + "datafusion.indexed.query_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( @@ -215,7 +256,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_QUERY_STRATEGY ); // ── Snapshot management ── @@ -256,6 +298,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)) + .queryStrategy(queryStrategyToWireValue(INDEXED_QUERY_STRATEGY.get(settings))) .build(); registerListeners(clusterSettings); @@ -278,6 +321,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)) + .queryStrategy(queryStrategyToWireValue(INDEXED_QUERY_STRATEGY.get(settings))) .build(); } @@ -310,6 +354,10 @@ void registerListeners(ClusterSettings clusterSettings) { snapshot = WireConfigSnapshot.builder(snapshot).maxCollectorParallelism(newValue).build(); }); + clusterSettings.addSettingsUpdateConsumer(INDEXED_QUERY_STRATEGY, newValue -> { + snapshot = WireConfigSnapshot.builder(snapshot).queryStrategy(queryStrategyToWireValue(newValue)).build(); + }); + clusterSettings.addSettingsUpdateConsumer(SearchService.CONCURRENT_SEGMENT_SEARCH_TARGET_MAX_SLICE_COUNT_SETTING, newValue -> { this.maxSliceCount = newValue; snapshot = WireConfigSnapshot.builder(snapshot) @@ -351,6 +399,19 @@ static int strategyToWireValue(String strategy) { } } + static int queryStrategyToWireValue(String strategy) { + switch (strategy) { + case QUERY_STRATEGY_NONE: + return 0; + case QUERY_STRATEGY_LISTING_TABLE: + return 1; + case QUERY_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/ShardScanInstructionHandler.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ShardScanInstructionHandler.java index 4dd4c801050f7..9c6efd1d66870 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ShardScanInstructionHandler.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ShardScanInstructionHandler.java @@ -11,6 +11,7 @@ import org.opensearch.analytics.backend.ShardScanExecutionContext; import org.opensearch.analytics.spi.BackendExecutionContext; import org.opensearch.analytics.spi.CommonExecutionContext; +import org.opensearch.analytics.spi.FilterTreeShape; import org.opensearch.analytics.spi.FragmentInstructionHandler; import org.opensearch.analytics.spi.ShardScanInstructionNode; import org.opensearch.be.datafusion.nativelib.NativeBridge; @@ -60,15 +61,34 @@ public BackendExecutionContext apply( try (Arena arena = Arena.ofConfined()) { MemorySegment segment = arena.allocate(WireConfigSnapshot.BYTE_SIZE); snapshot.writeTo(segment); - // Plan bytes let Rust widen the schema for multi-index queries (null-fill missing columns). - SessionContextHandle sessionCtxHandle = NativeBridge.createSessionContext( - readerPtr, - runtimePtr, - tableName, - contextId, - segment.address(), - context.getFragmentBytes() - ); + SessionContextHandle sessionCtxHandle; + if (node.requestsRowIds()) { + // QTF query phase — narrowed scan emits __row_id__. Use the indexed session + // context so the IndexedTableProvider injects shard-global row ids during scan. + // No delegated predicates here (delegation goes through ShardScanWithDelegationHandler), + // so treeShape=NO_DELEGATION and delegatedPredicateCount=0. + sessionCtxHandle = NativeBridge.createSessionContextForIndexedExecution( + readerPtr, + runtimePtr, + tableName, + contextId, + FilterTreeShape.NO_DELEGATION.ordinal(), + 0, + true, + segment.address(), + context.getFragmentBytes() + ); + } else { + // Plan bytes let Rust widen the schema for multi-index queries (null-fill missing columns). + sessionCtxHandle = NativeBridge.createSessionContext( + readerPtr, + runtimePtr, + tableName, + contextId, + segment.address(), + context.getFragmentBytes() + ); + } return new DataFusionSessionState(sessionCtxHandle); } } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ShardScanWithDelegationHandler.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ShardScanWithDelegationHandler.java index 8b972bffe8936..fc48f5e8a4e01 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ShardScanWithDelegationHandler.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/ShardScanWithDelegationHandler.java @@ -72,6 +72,7 @@ public BackendExecutionContext apply( contextId, treeShape.ordinal(), delegatedPredicateCount, + node.requestsRowIds(), segment.address(), context.getFragmentBytes() ); 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..d43a9c840fbc8 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 queryStrategy; 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.queryStrategy = builder.queryStrategy; } 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) + .queryStrategy(current.queryStrategy); } public int batchSize() { @@ -99,6 +102,10 @@ public int treeCollectorStrategy() { return treeCollectorStrategy; } + public int queryStrategy() { + return queryStrategy; + } + /** * Writes this snapshot into a {@code MemorySegment} matching the * {@code WireDatafusionQueryConfig} {@code #[repr(C)]} layout. @@ -122,11 +129,12 @@ 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 query_strategy i32 from snapshot (0/1/2) * ────── ──── - * Total: 68 bytes + * Total: 72 bytes * * - * @param segment the target memory segment (at least 68 bytes) + * @param segment the target memory segment (at least {@link #BYTE_SIZE} bytes) */ public void writeTo(MemorySegment segment) { // Offset 0: batch_size (i64) @@ -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: query_strategy (i32) — 0 = None, 1 = ListingTable, 2 = IndexedPredicateOnly + segment.set(ValueLayout.JAVA_INT, 68, queryStrategy); } /** @@ -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 queryStrategy = 2; // IndexedPredicateOnly (matches DatafusionSettings default "indexed") private Builder() {} @@ -213,6 +224,11 @@ public Builder treeCollectorStrategy(int treeCollectorStrategy) { return this; } + public Builder queryStrategy(int queryStrategy) { + this.queryStrategy = queryStrategy; + 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 c5de18b144191..2a457956c5e51 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 @@ -114,6 +114,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(); @@ -386,9 +387,10 @@ public final class NativeBridge { ValueLayout.JAVA_LONG, ValueLayout.JAVA_INT, ValueLayout.JAVA_INT, - ValueLayout.JAVA_LONG, - ValueLayout.ADDRESS, - ValueLayout.JAVA_LONG + ValueLayout.JAVA_BYTE, // requestsRowIds (0/1) — QTF query phase signal + ValueLayout.JAVA_LONG, // queryConfigPtr + ValueLayout.ADDRESS, // planBytes (multi-index schema widening) + ValueLayout.JAVA_LONG // planLen ) ); @@ -502,6 +504,22 @@ public final class NativeBridge { lib.find("df_execute_local_prepared_plan").orElseThrow(), FunctionDescriptor.of(ValueLayout.JAVA_LONG, 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() {} @@ -1154,6 +1172,7 @@ public static SessionContextHandle createSessionContextForIndexedExecution( long contextId, int treeShapeOrdinal, int delegatedPredicateCount, + boolean requestsRowIds, long queryConfigPtr, byte[] planBytes ) { @@ -1173,6 +1192,7 @@ public static SessionContextHandle createSessionContextForIndexedExecution( contextId, treeShapeOrdinal, delegatedPredicateCount, + (byte) (requestsRowIds ? 1 : 0), queryConfigPtr, planSegment, planLen @@ -1279,6 +1299,38 @@ public static long executeLocalPreparedPlan(long sessionPtr, long contextId) { } } + /** + * 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-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java index 76c127808f450..82b44c3aeb0a7 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionFragmentConvertorTests.java @@ -221,7 +221,14 @@ public void testAttachPartialAggOnTop_WrapsInner() throws Exception { public void testConvertFinalAggFragment_WithStageInputScanLeaf() throws Exception { RelDataType stageRowType = rowType("A"); int childStageId = 7; - RelNode stageInput = new OpenSearchStageInputScan(cluster, cluster.traitSet(), childStageId, stageRowType, List.of("datafusion")); + RelNode stageInput = new OpenSearchStageInputScan( + cluster, + cluster.traitSet(), + childStageId, + stageRowType, + List.of("datafusion"), + List.of() + ); LogicalAggregate finalAgg = buildSumAggregate(stageInput, 0); byte[] bytes = newConvertor().convertFragment(finalAgg); @@ -253,7 +260,14 @@ public void testAttachFragmentOnTop_Sort() throws Exception { // Inner: final-agg over stage-input. RelDataType stageRowType = rowType("A"); int childStageId = 3; - RelNode stageInput = new OpenSearchStageInputScan(cluster, cluster.traitSet(), childStageId, stageRowType, List.of("datafusion")); + RelNode stageInput = new OpenSearchStageInputScan( + cluster, + cluster.traitSet(), + childStageId, + stageRowType, + List.of("datafusion"), + List.of() + ); LogicalAggregate finalAgg = buildSumAggregate(stageInput, 0); byte[] innerBytes = convertor.convertFragment(finalAgg); @@ -318,11 +332,25 @@ public void testAttachFragmentOnTop_AggregateOverMultiColumnInner_PlanRootNamesM // Inner: a final-agg fragment whose StageInputScan rowType is intentionally wide // (3 columns). The aggregate above narrows it to 1 column. RelDataType wideStageRowType = rowType("A", "B", "C"); - RelNode stageInput = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, wideStageRowType, List.of("datafusion")); + RelNode stageInput = new OpenSearchStageInputScan( + cluster, + cluster.traitSet(), + 0, + wideStageRowType, + List.of("datafusion"), + List.of() + ); // For this regression, the inner doesn't need to be a final-agg — a bare scan-shaped // plan with 3-column rowType is enough to surface the wrapper-vs-inner names mismatch. // Use convertFragment so the inner Plan.Root.names is the 3-column scan list. - RelNode innerStageScan = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, wideStageRowType, List.of("datafusion")); + RelNode innerStageScan = new OpenSearchStageInputScan( + cluster, + cluster.traitSet(), + 0, + wideStageRowType, + List.of("datafusion"), + List.of() + ); // Wrap it in a no-op aggregate so the convertor accepts it as a final-agg fragment shape. // The inner's Plan.Root.names then carries the agg-output (1 col, "sum_col"), but the // *wrapper* we attach above has its own output rowType. @@ -364,9 +392,9 @@ public void testMultisearchShape_SortOverAggregateOverThreeWayUnion_PlanRootName // Inner: Union(Sin, Sin, Sin) — three branches, each 6 columns wide. RelDataType branchRowType = rowType("a", "b", "c", "d", "e", "f"); - RelNode sin1 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 1, branchRowType, List.of("datafusion")); - RelNode sin2 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 2, branchRowType, List.of("datafusion")); - RelNode sin3 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 3, branchRowType, List.of("datafusion")); + RelNode sin1 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 1, branchRowType, List.of("datafusion"), List.of()); + RelNode sin2 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 2, branchRowType, List.of("datafusion"), List.of()); + RelNode sin3 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 3, branchRowType, List.of("datafusion"), List.of()); LogicalUnion union = LogicalUnion.create(List.of(sin1, sin2, sin3), true); byte[] unionBytes = convertor.convertFragment(union); @@ -404,9 +432,9 @@ public void testMultisearchShape_SystemLimitOverSortOverAggregateOverUnion_Names // Inner: Union(Sin, Sin, Sin) — 6-column rows. RelDataType branchRowType = rowType("a", "b", "c", "d", "e", "f"); - RelNode sin1 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 1, branchRowType, List.of("datafusion")); - RelNode sin2 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 2, branchRowType, List.of("datafusion")); - RelNode sin3 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 3, branchRowType, List.of("datafusion")); + RelNode sin1 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 1, branchRowType, List.of("datafusion"), List.of()); + RelNode sin2 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 2, branchRowType, List.of("datafusion"), List.of()); + RelNode sin3 = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 3, branchRowType, List.of("datafusion"), List.of()); LogicalUnion union = LogicalUnion.create(List.of(sin1, sin2, sin3), true); byte[] unionBytes = convertor.convertFragment(union); @@ -651,7 +679,14 @@ public void testAttachFragmentOnTop_PreservesLiftedWindowProjectLayer() throws E DataFusionFragmentConvertor convertor = newConvertor(); RelDataType inputRowType = rowType("a"); - RelNode innerStageScan = new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, inputRowType, List.of("datafusion")); + RelNode innerStageScan = new OpenSearchStageInputScan( + cluster, + cluster.traitSet(), + 0, + inputRowType, + List.of("datafusion"), + List.of() + ); byte[] innerBytes = convertor.convertFragment(innerStageScan); RelNode placeholderInput = buildTableScan("__placeholder__", "a"); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java index 2daa8030a58e4..daa9baa19e0a2 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DataFusionPluginSettingsTests.java @@ -102,6 +102,7 @@ public void testGetSettingsReturnsAllIndexedSettings() { assertTrue(settingKeys.contains("datafusion.indexed.single_collector_strategy")); assertTrue(settingKeys.contains("datafusion.indexed.tree_collector_strategy")); assertTrue(settingKeys.contains("datafusion.indexed.max_collector_parallelism")); + assertTrue(settingKeys.contains("datafusion.indexed.query_strategy")); } catch (Exception e) { throw new AssertionError(e); } @@ -110,7 +111,7 @@ public void testGetSettingsReturnsAllIndexedSettings() { public void testGetSettingsReturnsTotalExpectedCount() { try (DataFusionPlugin plugin = new DataFusionPlugin()) { List> settings = plugin.getSettings(); - assertEquals(23, settings.size()); + assertEquals(24, settings.size()); } catch (Exception e) { throw new AssertionError(e); } diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java index d68ed055578a6..e8b2e37c8719d 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionSettingsTests.java @@ -69,7 +69,7 @@ public void testMaxCollectorParallelismSettingDefinition() { } public void testAllSettingsContainsAllExpectedSettings() { - assertEquals(23, DatafusionSettings.ALL_SETTINGS.size()); + assertEquals(24, DatafusionSettings.ALL_SETTINGS.size()); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DatafusionSettings.INDEXED_BATCH_SIZE)); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DatafusionSettings.INDEXED_PARQUET_PUSHDOWN_FILTERS)); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DatafusionSettings.INDEXED_MIN_SKIP_RUN_DEFAULT)); @@ -77,6 +77,7 @@ public void testAllSettingsContainsAllExpectedSettings() { assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DatafusionSettings.INDEXED_SINGLE_COLLECTOR_STRATEGY)); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DatafusionSettings.INDEXED_TREE_COLLECTOR_STRATEGY)); assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DatafusionSettings.INDEXED_MAX_COLLECTOR_PARALLELISM)); + assertTrue(DatafusionSettings.ALL_SETTINGS.contains(DatafusionSettings.INDEXED_QUERY_STRATEGY)); } public void testDefaultSnapshotValuesMatchDefaults() { @@ -91,6 +92,7 @@ public void testDefaultSnapshotValuesMatchDefaults() { assertEquals(1, snapshot.treeCollectorStrategy()); // tighten_outer_bounds assertEquals(1, snapshot.maxCollectorParallelism()); assertEquals(DEFAULT_PARALLELISM, snapshot.targetPartitions()); + assertEquals(2, snapshot.queryStrategy()); // indexed } public void testTargetPartitionsPassthroughWhenNonZero() { @@ -135,6 +137,25 @@ public void testStrategyToWireValueMapping() { expectThrows(IllegalArgumentException.class, () -> DatafusionSettings.strategyToWireValue("invalid")); } + public void testQueryStrategySettingDefinition() { + assertEquals("datafusion.indexed.query_strategy", DatafusionSettings.INDEXED_QUERY_STRATEGY.getKey()); + assertEquals("indexed", DatafusionSettings.INDEXED_QUERY_STRATEGY.get(Settings.EMPTY)); + assertTrue(DatafusionSettings.INDEXED_QUERY_STRATEGY.isDynamic()); + assertTrue(DatafusionSettings.INDEXED_QUERY_STRATEGY.hasNodeScope()); + } + + public void testQueryStrategyToWireValueMapping() { + assertEquals(0, DatafusionSettings.queryStrategyToWireValue("none")); + assertEquals(1, DatafusionSettings.queryStrategyToWireValue("listing_table")); + assertEquals(2, DatafusionSettings.queryStrategyToWireValue("indexed")); + expectThrows(IllegalArgumentException.class, () -> DatafusionSettings.queryStrategyToWireValue("invalid")); + } + + public void testInvalidQueryStrategyIsRejected() { + Settings settings = Settings.builder().put("datafusion.indexed.query_strategy", "bogus").build(); + expectThrows(IllegalArgumentException.class, () -> DatafusionSettings.INDEXED_QUERY_STRATEGY.get(settings)); + } + public void testBatchSizeZeroIsRejected() { Settings settings = Settings.builder().put("datafusion.indexed.batch_size", 0).build(); expectThrows(IllegalArgumentException.class, () -> DatafusionSettings.INDEXED_BATCH_SIZE.get(settings)); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/FilterDelegationForIndexFullConversionTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/FilterDelegationForIndexFullConversionTests.java index 9d1a47f61b3ba..670f1dc2551c6 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/FilterDelegationForIndexFullConversionTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/FilterDelegationForIndexFullConversionTests.java @@ -460,8 +460,8 @@ public FragmentConvertor getFragmentConvertor() { public FragmentInstructionHandlerFactory getInstructionHandlerFactory() { return new FragmentInstructionHandlerFactory() { @Override - public Optional createShardScanNode() { - return Optional.of(new ShardScanInstructionNode()); + public Optional createShardScanNode(boolean requestsRowIds) { + return Optional.of(new ShardScanInstructionNode(requestsRowIds)); } @Override @@ -474,8 +474,12 @@ public Optional createFilterDelegationNode( } @Override - public Optional createShardScanWithDelegationNode(FilterTreeShape treeShape, int delegatedPredicateCount) { - return Optional.of(new ShardScanWithDelegationInstructionNode(treeShape, delegatedPredicateCount)); + public Optional createShardScanWithDelegationNode( + FilterTreeShape treeShape, + int delegatedPredicateCount, + boolean requestsRowIds + ) { + return Optional.of(new ShardScanWithDelegationInstructionNode(treeShape, delegatedPredicateCount, requestsRowIds)); } @Override diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/WireConfigSnapshotTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/WireConfigSnapshotTests.java index 45cd4bc71f900..a7967145c93d3 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/WireConfigSnapshotTests.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/WireConfigSnapshotTests.java @@ -16,8 +16,8 @@ public class WireConfigSnapshotTests extends OpenSearchTestCase { - public void testByteSizeEquals68() { - assertEquals(68L, WireConfigSnapshot.BYTE_SIZE); + public void testByteSize() { + assertEquals(72L, WireConfigSnapshot.BYTE_SIZE); } public void testWriteToWritesCorrectValuesAtCorrectOffsets() { @@ -30,6 +30,7 @@ public void testWriteToWritesCorrectValuesAtCorrectOffsets() { .maxCollectorParallelism(4) .singleCollectorStrategy(2) .treeCollectorStrategy(1) + .queryStrategy(2) .build(); try (Arena arena = Arena.ofConfined()) { @@ -44,6 +45,7 @@ public void testWriteToWritesCorrectValuesAtCorrectOffsets() { assertEquals(4, segment.get(ValueLayout.JAVA_INT, 56)); // max_collector_parallelism assertEquals(2, segment.get(ValueLayout.JAVA_INT, 60)); // single_collector_strategy assertEquals(1, segment.get(ValueLayout.JAVA_INT, 64)); // tree_collector_strategy + assertEquals(2, segment.get(ValueLayout.JAVA_INT, 68)); // query_strategy = IndexedPredicateOnly } } @@ -84,6 +86,7 @@ public void testBuilderDefaultsMatchExpected() { assertEquals(1, snapshot.maxCollectorParallelism()); assertEquals(2, snapshot.singleCollectorStrategy()); // page_range_split assertEquals(1, snapshot.treeCollectorStrategy()); // tighten_outer_bounds + assertEquals(2, snapshot.queryStrategy()); // IndexedPredicateOnly } public void testBuilderCopyPreservesAllFields() { @@ -96,6 +99,7 @@ public void testBuilderCopyPreservesAllFields() { .maxCollectorParallelism(8) .singleCollectorStrategy(0) .treeCollectorStrategy(2) + .queryStrategy(1) .build(); WireConfigSnapshot copy = WireConfigSnapshot.builder(original).build(); @@ -108,5 +112,6 @@ public void testBuilderCopyPreservesAllFields() { assertEquals(original.maxCollectorParallelism(), copy.maxCollectorParallelism()); assertEquals(original.singleCollectorStrategy(), copy.singleCollectorStrategy()); assertEquals(original.treeCollectorStrategy(), copy.treeCollectorStrategy()); + assertEquals(original.queryStrategy(), copy.queryStrategy()); } } diff --git a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java index 70801932e0941..9acaf730ce826 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java +++ b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPluginTests.java @@ -301,8 +301,8 @@ public byte[] attachPartialAggOnTop(RelNode partialAggFragment, byte[] innerByte public FragmentInstructionHandlerFactory getInstructionHandlerFactory() { return new FragmentInstructionHandlerFactory() { @Override - public Optional createShardScanNode() { - return Optional.of(new ShardScanInstructionNode()); + public Optional createShardScanNode(boolean requestsRowIds) { + return Optional.of(new ShardScanInstructionNode(requestsRowIds)); } @Override @@ -315,8 +315,12 @@ public Optional createFilterDelegationNode( } @Override - public Optional createShardScanWithDelegationNode(FilterTreeShape treeShape, int delegatedPredicateCount) { - return Optional.of(new ShardScanWithDelegationInstructionNode(treeShape, delegatedPredicateCount)); + public Optional createShardScanWithDelegationNode( + FilterTreeShape treeShape, + int delegatedPredicateCount, + boolean requestsRowIds + ) { + return Optional.of(new ShardScanWithDelegationInstructionNode(treeShape, delegatedPredicateCount, requestsRowIds)); } @Override diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java index e1ddccffdd056..b6995aaae0a53 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/AnalyticsPlugin.java @@ -20,6 +20,7 @@ import org.opensearch.analytics.exec.DefaultPlanExecutor; import org.opensearch.analytics.exec.QueryPlanExecutor; import org.opensearch.analytics.exec.QueryScheduler; +import org.opensearch.analytics.exec.ReaderContextStore; import org.opensearch.analytics.exec.Scheduler; import org.opensearch.analytics.exec.action.AnalyticsQueryAction; import org.opensearch.analytics.planner.CapabilityRegistry; @@ -99,6 +100,7 @@ public AnalyticsPlugin() {} private SqlOperatorTable operatorTable; private AnalyticsSearchService searchService; private CoordinatorAllocatorHandle coordinatorAllocatorHandle; + private ReaderContextStore readerContextStore; @SuppressWarnings("rawtypes") @Override @@ -131,7 +133,10 @@ public Collection createComponents( for (AnalyticsSearchBackendPlugin be : backEnds) { backEndsByName.put(be.name(), be); } - searchService = new AnalyticsSearchService(backEndsByName, nativeAllocator, namedWriteableRegistry); + readerContextStore = new ReaderContextStore(threadPool); + clusterService.getClusterSettings() + .addSettingsUpdateConsumer(ReaderContextStore.READER_CONTEXT_KEEP_ALIVE, readerContextStore::setKeepAlive); + searchService = new AnalyticsSearchService(backEndsByName, nativeAllocator, namedWriteableRegistry, readerContextStore); DefaultEngineContext ctx = new DefaultEngineContext(clusterService, indexNameExpressionResolver, operatorTable, backEndsByName); // Build the coordinator allocator under POOL_QUERY here, in the plugin, so that the // plugin's lifecycle owns its lifetime. The Guice-bound DefaultPlanExecutor consumes @@ -167,7 +172,7 @@ public Collection createGuiceModules() { @Override public List> getSettings() { - return List.of(COORDINATOR_BUFFER_LIMIT); + return List.of(COORDINATOR_BUFFER_LIMIT, ReaderContextStore.READER_CONTEXT_KEEP_ALIVE); } @Override 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 f48a0bf8b629f..d62652df0c3d5 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 @@ -9,6 +9,7 @@ package org.opensearch.analytics.exec; import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.BigIntVector; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.analytics.backend.AnalyticsOperationListener; @@ -16,6 +17,7 @@ import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.backend.SearchExecEngine; import org.opensearch.analytics.backend.ShardScanExecutionContext; +import org.opensearch.analytics.exec.action.FetchByRowIdsRequest; import org.opensearch.analytics.exec.action.FragmentExecutionRequest; import org.opensearch.analytics.exec.task.AnalyticsShardTask; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; @@ -31,6 +33,7 @@ import org.opensearch.common.concurrent.GatedCloseable; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.tasks.TaskCancelledException; +import org.opensearch.index.engine.dataformat.DocumentInput; import org.opensearch.index.engine.exec.IndexReaderProvider; import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; import org.opensearch.index.shard.IndexShard; @@ -68,27 +71,31 @@ public class AnalyticsSearchService implements AutoCloseable { private final Map backends; private final AnalyticsOperationListener listener; private final NamedWriteableRegistry namedWriteableRegistry; + /** Cross-phase reader cache for QTF — query phase stores, fetch phase acquires. */ + private final ReaderContextStore readerContextStore; private TaskResourceTrackingService taskResourceTrackingService; private final BufferAllocator allocator; private final ArrowNativeAllocator nativeAllocator; public AnalyticsSearchService(Map backends, ArrowNativeAllocator nativeAllocator) { - this(backends, List.of(), nativeAllocator, null); + this(backends, List.of(), nativeAllocator, null, null); } public AnalyticsSearchService( Map backends, ArrowNativeAllocator nativeAllocator, - NamedWriteableRegistry namedWriteableRegistry + NamedWriteableRegistry namedWriteableRegistry, + ReaderContextStore readerContextStore ) { - this(backends, List.of(), nativeAllocator, namedWriteableRegistry); + this(backends, List.of(), nativeAllocator, namedWriteableRegistry, readerContextStore); } public AnalyticsSearchService( Map backends, List listeners, ArrowNativeAllocator nativeAllocator, - NamedWriteableRegistry namedWriteableRegistry + NamedWriteableRegistry namedWriteableRegistry, + ReaderContextStore readerContextStore ) { this.backends = backends; this.listener = new AnalyticsOperationListener.CompositeListener(listeners); @@ -103,6 +110,7 @@ public AnalyticsSearchService( BufferAllocator queryPool = nativeAllocator.getPoolAllocator(NativeAllocatorPoolConfig.POOL_QUERY); this.allocator = queryPool.newChildAllocator("analytics-search-service", 0, Long.MAX_VALUE); this.namedWriteableRegistry = namedWriteableRegistry; + this.readerContextStore = readerContextStore; } @Override @@ -155,6 +163,125 @@ public void executeFragmentStreamingAsync( } } + /** + * QTF fetch phase: retrieves specific rows by global row ID via the backend SPI and + * streams batches via {@link StreamingFragmentResponseHandler}. Forks onto + * {@code executor} so the iterator drain doesn't pin the transport thread — mirrors + * {@link #executeFragmentStreamingAsync}. + * + *

Reuses the {@link ReaderContext} opened during the query phase. If the context + * is missing (expired before fetch arrived, or query-phase reader-store invariant + * broken), the call fails — there is no cold-start fallback because shard-global + * {@code __row_id__} values produced by one reader cannot be reinterpreted by + * another (segment topology may differ across reopens). + */ + public void executeFetchByRowIdsAsync( + FetchByRowIdsRequest request, + IndexShard shard, + AnalyticsShardTask task, + StreamingFragmentResponseHandler responseHandler, + Executor executor + ) { + try { + executor.execute(() -> drainFetchByRowIds(request, shard, task, responseHandler)); + } catch (Exception e) { + responseHandler.onFailure(e); + } + } + + /** + * Acquires the per-shard {@link ReaderContext}, materialises the rowId vector, invokes + * the backend, drains the stream into {@code responseHandler}, and releases all resources + * in a single try-with-resources scope. Runs on the caller's executor — exhausting the + * iterator here lets the native engine apply backpressure when the channel is slow. + */ + private void drainFetchByRowIds( + FetchByRowIdsRequest request, + IndexShard shard, + AnalyticsShardTask task, + StreamingFragmentResponseHandler responseHandler + ) { + if (task != null && task.isCancelled()) { + responseHandler.onFailure(new TaskCancelledException("Fetch task cancelled before execution: " + task.getReasonCancelled())); + return; + } + long[] rowIds = request.getRowIds(); + String[] columns = request.getColumns(); + if (rowIds == null || rowIds.length == 0 || columns == null || columns.length == 0) { + responseHandler.onFailure( + new IllegalArgumentException( + "fetch on " + + shard.shardId() + + " requires non-empty rowIds and columns; got rowIds=" + + (rowIds == null ? "null" : rowIds.length) + + ", columns=" + + (columns == null ? "null" : columns.length) + ) + ); + return; + } + ReaderContext readerContext = readerContextStore.acquireContext(request.getQueryId(), shard.shardId()); + if (readerContext == null) { + responseHandler.onFailure( + new IllegalStateException( + "No ReaderContext for queryId=" + + request.getQueryId() + + " on " + + shard.shardId() + + " — query phase missing or context expired" + ) + ); + return; + } + assert assertFetchInvariants(readerContext, request.getQueryId()); + AnalyticsSearchBackendPlugin backend = backends.get(request.getBackendId()); + if (backend == null) { + readerContextStore.releaseContext(request.getQueryId(), shard.shardId()); + responseHandler.onFailure( + new IllegalStateException( + "No backend registered for backendId=" + + request.getBackendId() + + " on " + + shard.shardId() + + "; available: " + + backends.keySet() + ) + ); + return; + } + // Caller contract: rowIds must already be sorted ascending (RowSelection invariant on + // native side). Asserted here so violations are caught in dev builds before the FFM call. + assert assertAscending(rowIds); + BigIntVector rowIdVector = null; + FragmentResources resources = null; + try { + rowIdVector = new BigIntVector(DocumentInput.ROW_ID_FIELD, allocator); + rowIdVector.allocateNew(rowIds.length); + for (int i = 0; i < rowIds.length; i++) { + rowIdVector.set(i, rowIds[i]); + } + rowIdVector.setValueCount(rowIds.length); + EngineResultStream stream = backend.fetchByRowIds(readerContext.getReader(), rowIdVector, columns, allocator); + // FragmentResources keeps the rowIdVector alive until the stream drains — closing + // it earlier would pull off-heap memory out from under the native FFM call. + resources = new FragmentResources(readerContextStore, readerContext, null, stream, null, rowIdVector); + } catch (Exception e) { + if (rowIdVector != null) rowIdVector.close(); + readerContextStore.releaseContext(request.getQueryId(), shard.shardId()); + responseHandler.onFailure(new RuntimeException("Failed to execute fetch-by-row-ids on " + shard.shardId(), e)); + return; + } + try (FragmentResources ctx = resources) { + Iterator it = ctx.stream().iterator(); + while (it.hasNext()) { + responseHandler.onBatch(it.next()); + } + responseHandler.onComplete(); + } catch (Exception e) { + responseHandler.onFailure(e); + } + } + /** * Callback interface for async fragment streaming results. */ @@ -169,12 +296,17 @@ public interface StreamingFragmentResponseHandler { private FragmentResources startFragment(FragmentExecutionRequest request, ResolvedFragment resolved, IndexShard shard, Task task) throws IOException { GatedCloseable gatedReader = resolved.readerProvider.acquireReader(); + // QTF: hand the reader to the store so the fetch phase can reuse it without re-opening. + // FragmentResources holds a reference to the ReaderContext; close() releases it back + // to the store, the reaper closes after keepAlive. + ReaderContext readerContext = readerContextStore.createContext(request.getQueryId(), shard.shardId(), gatedReader); + assert assertReaderInvariants(gatedReader, readerContext, request.getQueryId(), shard); SearchExecEngine engine = null; EngineResultStream stream = null; BackendExecutionContext backendContext = null; Runnable trackerCleanup = null; try { - ShardScanExecutionContext ctx = buildContext(request, gatedReader.get(), resolved.plan, shard, task); + ShardScanExecutionContext ctx = buildContext(request, readerContext.getReader(), resolved.plan, shard, task); AnalyticsSearchBackendPlugin backend = backends.get(resolved.plan.getBackendId()); // Apply instruction handlers in order — each builds upon the previous handler's backend context @@ -219,7 +351,7 @@ public void trackEnd(long threadId) { engine = backend.getSearchExecEngineProvider().createSearchExecEngine(ctx, backendContext); stream = engine.execute(ctx); - return new FragmentResources(gatedReader, engine, stream, trackerCleanup); + return new FragmentResources(readerContextStore, readerContext, engine, stream, trackerCleanup); } catch (Exception e) { LOGGER.error( () -> new org.apache.logging.log4j.message.ParameterizedMessage( @@ -231,7 +363,7 @@ public void trackEnd(long threadId) { e ); try { - new FragmentResources(gatedReader, engine, stream, trackerCleanup).close(); + new FragmentResources(readerContextStore, readerContext, engine, stream, trackerCleanup).close(); } catch (Exception suppressed) { e.addSuppressed(suppressed); } @@ -305,4 +437,43 @@ private ShardScanExecutionContext buildContext( return ctx; } + // ── Assertion helpers (invoked only when -ea is enabled; bodies are dead in production) ── + + private static boolean assertReaderInvariants( + GatedCloseable gatedReader, + ReaderContext readerContext, + String queryId, + IndexShard shard + ) { + if (gatedReader == null) { + throw new AssertionError("acquireReader returned null for shard " + shard.shardId()); + } + if (readerContext == null) { + throw new AssertionError("createContext returned null for queryId=" + queryId); + } + if (readerContext.getReader() == null) { + throw new AssertionError("ReaderContext returned null reader for queryId=" + queryId); + } + return true; + } + + private boolean assertFetchInvariants(ReaderContext readerContext, String queryId) { + if (readerContext.getReader() == null) { + throw new AssertionError("acquired ReaderContext has null reader for queryId=" + queryId); + } + if (backends.isEmpty()) { + throw new AssertionError("no backends registered — service constructor invariant violated"); + } + return true; + } + + private static boolean assertAscending(long[] values) { + for (int i = 1; i < values.length; i++) { + if (values[i] < values[i - 1]) { + throw new AssertionError("rowIds not ascending at index " + i + ": " + values[i - 1] + " > " + values[i]); + } + } + return true; + } + } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java index 51837e66843bf..491555ac898a6 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/AnalyticsSearchTransportService.java @@ -9,6 +9,8 @@ package org.opensearch.analytics.exec; import org.opensearch.analytics.backend.EngineResultBatch; +import org.opensearch.analytics.exec.action.FetchByRowIdsAction; +import org.opensearch.analytics.exec.action.FetchByRowIdsRequest; import org.opensearch.analytics.exec.action.FragmentExecutionAction; import org.opensearch.analytics.exec.action.FragmentExecutionArrowResponse; import org.opensearch.analytics.exec.action.FragmentExecutionRequest; @@ -27,11 +29,11 @@ import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.StreamTransportService; import org.opensearch.transport.Transport; +import org.opensearch.transport.TransportChannel; import org.opensearch.transport.TransportException; +import org.opensearch.transport.TransportRequest; import org.opensearch.transport.TransportRequestOptions; import org.opensearch.transport.TransportResponseHandler; -import org.opensearch.transport.stream.StreamErrorCode; -import org.opensearch.transport.stream.StreamException; import org.opensearch.transport.stream.StreamTransportResponse; import java.io.IOException; @@ -71,6 +73,7 @@ public AnalyticsSearchTransportService( this.transportService = streamTransportService; this.clusterService = clusterService; registerStreamingFragmentHandler(this.transportService, searchService, indicesService); + registerFetchByRowIdsHandler(this.transportService, searchService, indicesService); } private static void registerStreamingFragmentHandler( @@ -91,33 +94,75 @@ private static void registerStreamingFragmentHandler( request, shard, (AnalyticsShardTask) task, - new AnalyticsSearchService.StreamingFragmentResponseHandler() { - @Override - public void onBatch(EngineResultBatch batch) throws Exception { - channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); - } - - @Override - public void onComplete() { - channel.completeStream(); - } + channelResponseHandler(channel), + transportService.getThreadPool().executor(ThreadPool.Names.SEARCH) + ); + } + ); + } - @Override - public void onFailure(Exception e) { - if (e instanceof StreamException se && se.getErrorCode() == StreamErrorCode.CANCELLED) { - return; - } - try { - channel.sendResponse(e); - } catch (Exception ignored) {} - } - }, + /** + * Mirrors {@link #registerStreamingFragmentHandler} for the QTF fetch-by-rowids path. + * Forks the iterator drain onto the SEARCH executor via + * {@link AnalyticsSearchService#executeFetchByRowIdsAsync} so a slow coordinator parks + * a search thread, not the transport thread, and the engine sees natural backpressure + * through the blocking {@code channel.sendResponseBatch} call. + */ + private static void registerFetchByRowIdsHandler( + StreamTransportService transportService, + AnalyticsSearchService searchService, + IndicesService indicesService + ) { + transportService.registerRequestHandler( + FetchByRowIdsAction.NAME, + ThreadPool.Names.SAME, + false, + true, + AdmissionControlActionType.SEARCH, + FetchByRowIdsRequest::new, + (request, channel, task) -> { + IndexShard shard = indicesService.indexServiceSafe(request.getShardId().getIndex()).getShard(request.getShardId().id()); + searchService.executeFetchByRowIdsAsync( + request, + shard, + (AnalyticsShardTask) task, + channelResponseHandler(channel), transportService.getThreadPool().executor(ThreadPool.Names.SEARCH) ); } ); } + /** + * Adapter from {@link AnalyticsSearchService.StreamingFragmentResponseHandler} to the + * channel streaming API. Each batch is sent on the channel; onComplete completes the + * stream; onFailure ignores cancellation and forwards everything else as an exception + * response. Shared by the streaming-fragment and fetch-by-rowids handlers — the dispatch + * shape is identical, only the request type differs. + */ + private static AnalyticsSearchService.StreamingFragmentResponseHandler channelResponseHandler(TransportChannel channel) { + return new AnalyticsSearchService.StreamingFragmentResponseHandler() { + @Override + public void onBatch(EngineResultBatch batch) throws Exception { + channel.sendResponseBatch(new FragmentExecutionArrowResponse(batch.getArrowRoot())); + } + + @Override + public void onComplete() { + channel.completeStream(); + } + + @Override + public void onFailure(Exception e) { + try { + channel.sendResponse(e); + } catch (Exception sendException) { + throw new RuntimeException(sendException); + } + } + }; + } + Transport.Connection getConnection(String clusterAlias, String nodeId) { DiscoveryNode node = clusterService.state().nodes().get(nodeId); return transportService.getConnection(node); @@ -129,6 +174,40 @@ public void dispatchFragmentStreaming( StreamingResponseListener listener, Task parentTask, PendingExecutions pending + ) { + dispatchStreaming(FragmentExecutionAction.NAME, request, targetNode, listener, parentTask, pending); + } + + /** + * Dispatches a QTF fetch-by-rowids RPC to {@code targetNode}. Mirrors + * {@link #dispatchFragmentStreaming} — same streaming-response handler shape, same + * {@link PendingExecutions} gating, same cancellation propagation. Different action + * name routes the request to {@link FetchByRowIdsAction} on the data node. + */ + public void dispatchFetchByRowIds( + FetchByRowIdsRequest request, + DiscoveryNode targetNode, + StreamingResponseListener listener, + Task parentTask, + PendingExecutions pending + ) { + dispatchStreaming(FetchByRowIdsAction.NAME, request, targetNode, listener, parentTask, pending); + } + + /** + * Shared streaming dispatch path for {@link FragmentExecutionAction} and + * {@link FetchByRowIdsAction}. Drains the response stream inline — backpressure flows + * because {@code listener.onStreamResponse} blocks until the downstream sink accepts + * the batch, which gates the next {@code stream.nextResponse} call and propagates + * gRPC flow control back to the data node. + */ + private void dispatchStreaming( + String actionName, + TransportRequest request, + DiscoveryNode targetNode, + StreamingResponseListener listener, + Task parentTask, + PendingExecutions pending ) { TransportResponseHandler handler = new TransportResponseHandler<>() { @Override @@ -193,7 +272,7 @@ public void handleException(TransportException e) { pending.tryRun(() -> { try { Transport.Connection connection = getConnection(null, targetNode.getId()); - transportService.sendChildRequest(connection, FragmentExecutionAction.NAME, request, parentTask, options, handler); + transportService.sendChildRequest(connection, actionName, request, parentTask, options, handler); } catch (Exception e) { try { listener.onFailure(e); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java index 92b8e5e1041be..69af5e863c740 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/FragmentResources.java @@ -8,43 +8,67 @@ package org.opensearch.analytics.exec; +import org.apache.arrow.vector.BigIntVector; import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.backend.SearchExecEngine; import org.opensearch.analytics.backend.ShardScanExecutionContext; -import org.opensearch.common.concurrent.GatedCloseable; -import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; /** - * Holds the per-fragment resources (reader, engine, result stream) kept alive for the - * duration of a streaming fragment execution, and releases them in reverse order on close. + * Holds the per-fragment resources (reader context, engine, result stream) kept alive for + * the duration of a streaming fragment execution, and releases them in reverse order on close. + * + *

The reader is owned by {@link ReaderContextStore}, not by this class — close releases + * (does not free) the context, so the reader stays alive across the QTF query→fetch + * boundary. The store's reaper closes the underlying reader after keepAlive elapses. * * @opensearch.internal */ public final class FragmentResources implements AutoCloseable { - private final GatedCloseable gatedReader; + private final ReaderContextStore readerContextStore; + private final ReaderContext readerContext; private final SearchExecEngine engine; private final EngineResultStream stream; private final Runnable onClose; + /** + * Off-heap rowId buffer kept alive across the fetch stream's lifetime. Non-null only + * for the QTF fetch path, where the native side reads rowIds directly via the + * BigIntVector's data-buffer address — closing it before the stream drains would + * pull memory out from under the FFM call. + */ + private final BigIntVector rowIdVector; public FragmentResources( - GatedCloseable gatedReader, + ReaderContextStore readerContextStore, + ReaderContext readerContext, SearchExecEngine engine, - EngineResultStream stream + EngineResultStream stream, + Runnable onClose ) { - this(gatedReader, engine, stream, null); + this(readerContextStore, readerContext, engine, stream, onClose, null); } public FragmentResources( - GatedCloseable gatedReader, + ReaderContextStore readerContextStore, + ReaderContext readerContext, SearchExecEngine engine, EngineResultStream stream, - Runnable onClose + Runnable onClose, + BigIntVector rowIdVector ) { - this.gatedReader = gatedReader; + assert assertCtorInvariants(readerContextStore, readerContext); + this.readerContextStore = readerContextStore; + this.readerContext = readerContext; this.engine = engine; this.stream = stream; this.onClose = onClose; + this.rowIdVector = rowIdVector; + } + + private static boolean assertCtorInvariants(ReaderContextStore store, ReaderContext ctx) { + if (store == null) throw new AssertionError("readerContextStore is required for FragmentResources"); + if (ctx == null) throw new AssertionError("readerContext is required for FragmentResources"); + return true; } public EngineResultStream stream() { @@ -63,7 +87,17 @@ public void close() throws Exception { } first = closeQuietly(stream, first); first = closeQuietly(engine, first); - first = closeQuietly(gatedReader, first); + first = closeQuietly(rowIdVector, first); + // Release (not close) — the store's reaper closes after keepAlive, and the QTF + // fetch phase may still need this reader before then. + if (readerContext != null) { + try { + readerContextStore.releaseContext(readerContext.getQueryId(), readerContext.getShardId()); + } catch (Exception e) { + if (first == null) first = e; + else first.addSuppressed(e); + } + } if (first != null) throw first; } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/OrdinalAppendingSink.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/OrdinalAppendingSink.java new file mode 100644 index 0000000000000..f227171577067 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/OrdinalAppendingSink.java @@ -0,0 +1,58 @@ +/* + * 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.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.opensearch.analytics.spi.ExchangeSink; + +/** + * Pure decorator over an {@link ExchangeSink} that stamps a per-source ordinal + * onto every batch as a constant Int32 column. Used by Late Materialization + * (Query-Then-Fetch) to mark each shard's batches with their {@code ___ugsi} + * before the rows hit the reduce sink, so post-reduce code can group by source + * shard for fan-out fetches. + * + *

Lifecycle is delegated entirely to the wrapped sink — this decorator only + * adds a column on the {@code feed(VSR, int)} path. + * + *

Append is zero-copy: the new VSR shares existing {@link org.apache.arrow.vector.FieldVector}s + * with the input by reference; only the constant ordinal column is freshly + * allocated. See {@link VectorUtils#appendConstantInt}. + * + * @opensearch.internal + */ +public final class OrdinalAppendingSink implements ExchangeSink { + + private final ExchangeSink delegate; + private final BufferAllocator allocator; + private final String columnName; + + public OrdinalAppendingSink(ExchangeSink delegate, BufferAllocator allocator, String columnName) { + this.delegate = delegate; + this.allocator = allocator; + this.columnName = columnName; + } + + @Override + public void feed(VectorSchemaRoot batch, int sourceOrdinal) { + VectorSchemaRoot withOrdinal = VectorUtils.appendConstantInt(batch, columnName, sourceOrdinal, allocator); + delegate.feed(withOrdinal); + } + + @Override + public void feed(VectorSchemaRoot batch) { + delegate.feed(batch); + } + + @Override + public void close() { + delegate.close(); + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java index ab7ba7f0c12a4..90d721adc29f2 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/QueryContext.java @@ -14,9 +14,12 @@ import org.opensearch.analytics.backend.AnalyticsOperationListener; import org.opensearch.analytics.exec.task.AnalyticsQueryTask; import org.opensearch.analytics.planner.dag.QueryDAG; +import org.opensearch.analytics.planner.dag.ShardExecutionTarget; import org.opensearch.threadpool.ThreadPool; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.Executor; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -41,6 +44,22 @@ public class QueryContext { private final boolean ownsAllocator; private volatile ExecutorService localTaskExecutor; private boolean closed; // guarded by `this` + /** + * HACK: side-table for cross-stage routing of resolved {@link ShardExecutionTarget}s. + * Today's only consumer is the QTF (late-materialization) Phase C, which needs to map + * an incoming row's {@code ___ugsi} ordinal back to the {@code (DiscoveryNode, ShardId)} + * to dispatch a fetch. Stage 1 (SHARD_FRAGMENT) populates this once after resolve; + * Stage 3 (LM) reads it. + * + *

TODO: this is a placeholder seam. {@code QueryContext} should not be a generic + * "things stages leave for other stages to find" map. Cleaner shapes: cache on + * {@code Stage} alongside {@code targetResolver}, or reify a typed cross-stage routing + * table. Revisit when a second consumer appears or when extending QTF to UNION/JOIN. + * + *

Single-threaded write inside one stage's {@code materializeTasks}; reads happen + * only after that stage SUCCEEDED → plain {@link HashMap} suffices. + */ + private final Map> resolvedTargetsByStage = new HashMap<>(); public QueryContext( QueryDAG dag, @@ -104,6 +123,29 @@ public List operationListeners() { return operationListeners; } + /** + * Records the {@link ShardExecutionTarget}s resolved for a stage. Called once by the + * stage execution after {@code TargetResolver.resolve(...)} runs. See the field-level + * Javadoc on {@code resolvedTargetsByStage} for context on why this lives on + * {@code QueryContext}. + */ + public void recordResolvedTargets(int stageId, List targets) { + Map byOrdinal = new HashMap<>(targets.size()); + for (ShardExecutionTarget t : targets) { + byOrdinal.put(t.ordinal(), t); + } + resolvedTargetsByStage.put(stageId, byOrdinal); + } + + /** + * Returns the resolved targets for a stage keyed by per-shard ordinal (UGSI), or + * {@code null} if that stage hasn't resolved yet (or doesn't have a resolver). The + * Map is built once at record time so callers can do O(1) ordinal-to-target lookup. + */ + public Map getResolvedTargets(int stageId) { + return resolvedTargetsByStage.get(stageId); + } + public BufferAllocator bufferAllocator() { return allocator; } 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..a643ab62a262a --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContext.java @@ -0,0 +1,105 @@ +/* + * 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.core.index.shard.ShardId; +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 ShardId shardId; + 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, ShardId shardId, GatedCloseable gatedReader, long keepAliveMillis) { + this.queryId = queryId; + this.shardId = shardId; + this.gatedReader = gatedReader; + this.keepAliveMillis = keepAliveMillis; + this.lastAccessTime = new AtomicLong(System.currentTimeMillis()); + } + + public String getQueryId() { + return queryId; + } + + public ShardId getShardId() { + return shardId; + } + + 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..0a72a5ae539d3 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/ReaderContextStore.java @@ -0,0 +1,135 @@ +/* + * 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.core.index.shard.ShardId; +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 + ); + + /** + * Multi-shard QTF queries call {@link #createContext} once per shard with the same + * {@code queryId}. Keying by {@code queryId} alone would collide and overwrite. + */ + public record Key(String queryId, ShardId shardId) { + } + + 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/shard. + * Acquires the reader and marks it in-use. + */ + public ReaderContext createContext(String queryId, ShardId shardId, GatedCloseable gatedReader) { + ReaderContext ctx = new ReaderContext(queryId, shardId, gatedReader, defaultKeepAliveMillis); + ctx.markInUse(); + activeContexts.put(new Key(queryId, shardId), ctx); + return ctx; + } + + /** + * Get an existing context by queryId/shardId. Returns null if not found or expired. + */ + public ReaderContext getContext(String queryId, ShardId shardId) { + return activeContexts.get(new Key(queryId, shardId)); + } + + /** + * Acquire a context for use (fetch phase). Marks it in-use. + * Returns null if not found or already closed. + */ + public ReaderContext acquireContext(String queryId, ShardId shardId) { + ReaderContext ctx = activeContexts.get(new Key(queryId, shardId)); + 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, ShardId shardId) { + ReaderContext ctx = activeContexts.get(new Key(queryId, shardId)); + if (ctx != null) { + ctx.markDone(); + } + } + + /** + * Remove and close a context (fetch complete, no longer needed). + */ + public void freeContext(String queryId, ShardId shardId) { + ReaderContext ctx = activeContexts.remove(new Key(queryId, shardId)); + if (ctx != null) { + try { + ctx.close(); + } catch (Exception e) { + logger.warn("[ReaderContextStore] Failed to close context for query={} shard={}: {}", queryId, shardId, 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={} shard={}", ctx.getQueryId(), ctx.getShardId()); + freeContext(ctx.getQueryId(), ctx.getShardId()); + } + } + } + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/VectorUtils.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/VectorUtils.java new file mode 100644 index 0000000000000..b4d91bed8b989 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/VectorUtils.java @@ -0,0 +1,78 @@ +/* + * 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.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; + +import java.util.ArrayList; +import java.util.List; + +/** + * Utilities for Arrow {@link VectorSchemaRoot} manipulation in the analytics engine. + * + * @opensearch.internal + */ +public final class VectorUtils { + + private VectorUtils() {} + + /** + * Builds a new {@link VectorSchemaRoot} whose schema is the input's schema with + * an additional non-nullable Int32 column appended at the end. The returned VSR's + * existing columns are the same {@link org.apache.arrow.vector.FieldVector} + * instances as the input — only the new constant column is a fresh allocation. + * No data is copied. + * + *

Position contract: the new column is appended at the end of the schema. This + * must align with the planner-side declaration that adds helper columns (e.g. + * {@code ___ugsi}) at the end of the row type — see + * {@code RelNodeUtils.appendField}. + * + *

TODO: position is hard-coded to "end" today because the planner happens to + * declare helper columns at the end. If a future helper is declared in the middle + * of the row type, this helper needs to take a column index parameter (or a name + * + reference schema to look up the index) instead of assuming the tail. Coupling + * the runtime append position to the planner's declared position is the design + * goal; we just don't pay the index-lookup cost today. + * + *

Ownership semantics: the returned VSR shares its leading FieldVectors with + * the input by reference (Arrow ref-counting). Callers must not close the input + * VSR independently — closing the returned VSR releases both the new column and + * the shared columns. + * + *

Implementation: builds a new VSR from the existing FieldVector references plus the + * constant column. We avoid {@code VectorSchemaRoot.addVector} because Arrow 18.x's + * precondition rejects appending at index == size (only insert before an existing column + * is allowed); the Iterable constructor accepts the appended layout directly. + * + * @param input source batch; not closed by this method + * @param name name of the new column + * @param value constant value written to every row of the new column + * @param allocator allocator for the new {@link IntVector} + * @return new VSR with the appended column; row count and existing column data unchanged + */ + public static VectorSchemaRoot appendConstantInt(VectorSchemaRoot input, String name, int value, BufferAllocator allocator) { + int rowCount = input.getRowCount(); + IntVector constantVector = new IntVector(name, allocator); + constantVector.allocateNew(rowCount); + for (int i = 0; i < rowCount; i++) { + constantVector.set(i, value); + } + constantVector.setValueCount(rowCount); + + // Zero-copy: existing FieldVectors are reused by reference; only the new column allocates. + List combined = new ArrayList<>(input.getFieldVectors().size() + 1); + combined.addAll(input.getFieldVectors()); + combined.add(constantVector); + return new VectorSchemaRoot(combined); + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FetchByRowIdsAction.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FetchByRowIdsAction.java new file mode 100644 index 0000000000000..23bd9cb3b4220 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FetchByRowIdsAction.java @@ -0,0 +1,30 @@ +/* + * 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.action; + +import org.opensearch.action.ActionType; + +/** + * {@link ActionType} singleton for the QTF (Query-Then-Fetch / late-materialization) + * fetch-by-rowids action. Sibling of {@link FragmentExecutionAction} — different semantics + * (no Substrait fragment, no plan conversion) so the dispatch table needs a distinct name, + * but it shares the {@code indices:data/read/analytics/*} prefix so the existing security + * permission grant covers both. Reuses {@link FragmentExecutionArrowResponse} as the + * per-batch streamed response type. + */ +public class FetchByRowIdsAction extends ActionType { + + public static final String NAME = "indices:data/read/analytics/fetch_by_row_ids"; + + public static final FetchByRowIdsAction INSTANCE = new FetchByRowIdsAction(); + + private FetchByRowIdsAction() { + super(NAME, FragmentExecutionArrowResponse::new); + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FetchByRowIdsRequest.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FetchByRowIdsRequest.java new file mode 100644 index 0000000000000..def456113c0ea --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FetchByRowIdsRequest.java @@ -0,0 +1,109 @@ +/* + * 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.action; + +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.analytics.exec.task.AnalyticsShardTask; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.index.shard.ShardId; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.tasks.Task; + +import java.io.IOException; +import java.util.Map; + +/** + * Per-shard QTF (Query-Then-Fetch / late-materialization) fetch-by-rowids request. Carries + * the ascending {@code rowIds} (data-node contract) plus the column projection. Routed to + * the data node by {@link FetchByRowIdsAction}; resolved against the same backend that ran + * the query phase via {@code backendId} — replaces the prior 5-arg + * {@code AnalyticsSearchService.executeFetchByRowIds} that picked a backend with + * {@code backends.values().iterator().next()}. + * + * @opensearch.internal + */ +public class FetchByRowIdsRequest extends ActionRequest implements ShardInvocationRequest { + + private final String queryId; + private final int stageId; + private final ShardId shardId; + private final String backendId; + private final long[] rowIds; + private final String[] columns; + + public FetchByRowIdsRequest(String queryId, int stageId, ShardId shardId, String backendId, long[] rowIds, String[] columns) { + this.queryId = queryId; + this.stageId = stageId; + this.shardId = shardId; + this.backendId = backendId; + this.rowIds = rowIds; + this.columns = columns; + } + + public FetchByRowIdsRequest(StreamInput in) throws IOException { + super(in); + this.queryId = in.readString(); + this.stageId = in.readInt(); + this.shardId = new ShardId(in); + this.backendId = in.readString(); + this.rowIds = in.readLongArray(); + this.columns = in.readStringArray(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(queryId); + out.writeInt(stageId); + shardId.writeTo(out); + out.writeString(backendId); + out.writeLongArray(rowIds); + out.writeStringArray(columns); + } + + @Override + public String getQueryId() { + return queryId; + } + + @Override + public int getStageId() { + return stageId; + } + + @Override + public ShardId getShardId() { + return shardId; + } + + public String getBackendId() { + return backendId; + } + + public long[] getRowIds() { + return rowIds; + } + + public String[] getColumns() { + return columns; + } + + @Override + public Task createTask(long id, String type, String action, TaskId parentTaskId, Map headers) { + String desc = "queryId[" + queryId + "] stageId[" + stageId + "] shardId[" + shardId + "] rowIds[" + rowIds.length + "]"; + return new AnalyticsShardTask(id, type, action, desc, parentTaskId, headers); + } + + @Override + public ActionRequestValidationException validate() { + return null; + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java index fd137abb95c50..1e1513b3aa189 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/FragmentExecutionRequest.java @@ -34,7 +34,7 @@ * * @opensearch.internal */ -public class FragmentExecutionRequest extends ActionRequest { +public class FragmentExecutionRequest extends ActionRequest implements ShardInvocationRequest { private final String queryId; private final int stageId; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/ShardInvocationRequest.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/ShardInvocationRequest.java new file mode 100644 index 0000000000000..6fbacf1f0a264 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/action/ShardInvocationRequest.java @@ -0,0 +1,27 @@ +/* + * 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.action; + +import org.opensearch.core.index.shard.ShardId; + +/** + * Cross-cutting shape for any shard-targeted analytics action: identifies the originating + * query, the stage, and the target shard. Implemented by {@link FragmentExecutionRequest} + * and {@link FetchByRowIdsRequest}; lets shared infrastructure (failure-listener wrapping, + * resource accounting, logging) operate on either request type without an instanceof check. + * + * @opensearch.internal + */ +public interface ShardInvocationRequest { + String getQueryId(); + + int getStageId(); + + ShardId getShardId(); +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java new file mode 100644 index 0000000000000..ab64bf3cbc86b --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecution.java @@ -0,0 +1,562 @@ +/* + * 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.stage; + +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.analytics.backend.ExchangeSource; +import org.opensearch.analytics.exec.AnalyticsSearchTransportService; +import org.opensearch.analytics.exec.PendingExecutions; +import org.opensearch.analytics.exec.QueryContext; +import org.opensearch.analytics.exec.RowProducingSink; +import org.opensearch.analytics.exec.StreamingResponseListener; +import org.opensearch.analytics.exec.action.FetchByRowIdsRequest; +import org.opensearch.analytics.exec.action.FragmentExecutionArrowResponse; +import org.opensearch.analytics.exec.stage.coordinator.LocalStageTask; +import org.opensearch.analytics.exec.stage.coordinator.LocalTaskRunner; +import org.opensearch.analytics.exec.stage.coordinator.ReduceStageExecution; +import org.opensearch.analytics.exec.stage.shard.ShardFragmentStageExecution; +import org.opensearch.analytics.planner.ArrowCalciteTypes; +import org.opensearch.analytics.planner.RelNodeUtils; +import org.opensearch.analytics.planner.dag.ShardExecutionTarget; +import org.opensearch.analytics.planner.dag.Stage; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; +import org.opensearch.analytics.spi.DataConsumer; +import org.opensearch.analytics.spi.ExchangeSink; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.core.action.ActionListener; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; + +/** + * QTF (Query-Then-Fetch / late-materialization) Scatter-Gather stage execution. + * + *

This stage orchestrates the four QTF phases internally — there is no Substrait + * fragment and no DataFusion compute. The stage sits above a Sort+Limit reduce stage + * (its only child) and below the parent stage that consumes the fully-materialized + * top-K rows. + * + * 

+ *   parent stage (Post-Sort)
+ *      ▲                        ◄── feeds parent's input sink
+ *      │   stitched K rows
+ *   LateMaterializationStageExecution
+ *      │  Phase D: stitch → emit                 (Java)
+ *      │  Phase C: gather per-node fetch results (Java)
+ *      │  Phase B: scatter fetch by ___ugsi      (Java + transport)
+ *      │  Phase A: drain reduce output           (Java)
+ *      ▲                        ◄── filled by child Sort+Limit reduce
+ *      │   K rows of [sort_cols, ___row_id, ___ugsi]
+ *   child stage (Sort+Limit COORDINATOR_REDUCE)
+ *
+ * + *

Lifecycle

+ * + *
    + *
  1. {@link StageExecution.State#CREATED} — built by {@code LateMaterializationStageExecutionFactory}; + * the parent (Post-Sort) stage's start is gated on this stage's SUCCEEDED.
    2. + *
  2. {@link StageExecution.State#RUNNING} — entered when the child Sort+Limit stage + * SUCCEEDED. The cascade in {@code PlanWalker} fires {@link #start()}; we drain the + * child's output, fan out fetches, stitch, feed the parent's input sink.
    3. + *
  3. {@link StageExecution.State#SUCCEEDED} — every fetch returned, every stitched batch + * fed to the parent sink, parent stage can start.
    4. + *
+ * + *

Implementation status

+ * + *

SKELETON ONLY. Today {@link #start()} throws + * {@link UnsupportedOperationException}. The four phases below are the work to land: + * + *

Phase A — drain reduce output (Java)

+ * + *

The child stage's reduced K rows are buffered in its + * {@link RowProducingSink}-style output (arriving via {@code feed(VSR)} on whatever + * sink the LM stage provides via {@link #inputSink(int)}). At {@link #start()} time the + * full input is available — block-read it. + * + * 

+ *   ExchangeSource src = (ExchangeSource) downstreamSink;   // see inputSink() below
+ *   for (VectorSchemaRoot batch : src.readResult()) {
+ *       // each batch has schema [sort_cols, ___row_id, ___ugsi]
+ *       buffer.add(batch);  // takes ownership; close on terminal
+ *   }
+ *
+ * + *

Phase B — scatter fetch by {@code ___ugsi} (Java + transport)

+ * + *

Group rows by {@code ___ugsi} ordinal. Map each ordinal → (nodeId, shardId, + * indexUUID) using a coord-side side table built when the child stage's per-shard tasks + * dispatched. For each {@code (nodeId, list of (___row_id, rowPos))}, send a + * {@code FetchByRowIdRequest} via {@link #transport}. Use + * {@code stage.getStageId()} for correlation, {@code config.parentTask()} for cancel + * propagation. + * + *

New transport action required. Mirror + * {@link org.opensearch.analytics.exec.action.FragmentExecutionAction} — define + * {@code FetchByRowIdAction}, request and response types, register on + * {@link org.opensearch.analytics.exec.AnalyticsSearchTransportService}. + * + *

Data-node handler required. The handler reads the columns named in the + * wrapper's {@link OpenSearchLateMaterialization#getAboveAnchorPhysicalFields()} from each shard's + * doc-id mapped row, returning {@code (rowPos, fetched_col_values)} batches. + * Backend-specific: DataFusion's parquet reader can produce these from a Lucene docId + * lookup (see {@code ShardScanInstructionHandler} for the registration pattern; the + * fetch-by-rowid handler is a sibling, not a fragment-instruction). + * + *

Phase C — gather (Java)

+ * + *

Per-node responses arrive out of order. Maintain a {@code ConcurrentHashMap}. As each response arrives, decode batches, write per-row + * cells keyed by {@code rowPos}. Use an {@link java.util.concurrent.atomic.AtomicInteger} + * for in-flight count (mirrors {@link ShardFragmentStageExecution}'s pattern). When + * count hits zero, transition to Phase D. + * + *

Phase D — stitch and feed parent (Java)

+ * + *

Walk {@code rowPos = 0..K-1}, building output batches with schema + * {@code [sort_cols, fetch_cols...]} (helper columns {@code ___row_id} + {@code ___ugsi} + * stripped per {@link OpenSearchLateMaterialization} contract). Each row's sort cols + * come from the buffered Phase A batches; fetched cols from the Phase C map. Feed + * batches into {@link #parentSink}. On EOF call {@link ExchangeSink#close()} on + * parentSink and transition to {@link StageExecution.State#SUCCEEDED}. + * + *

Cancellation

+ * + *

{@link #cancel(String)} cancels the parent task — propagates to in-flight fetch + * transports and the (already-completed) child stage. Mirrors + * {@link ShardFragmentStageExecution#cancel(String)}. + * + *

Cross-references

+ * + *
    + *
  • {@link OpenSearchLateMaterialization} — wrapper RelNode in the stage's fragment; + * carries {@code getAboveAnchorPhysicalFields()} (List<RelDataTypeField>) and + * {@code getAboveAnchorPhysicalFieldStorage()} (List<FieldStorageInfo>).
    • + *
  • {@link ShardFragmentStageExecution} — closest existing analogue for transport + * fan-out + per-shard response handling.
    • + *
  • {@link ReduceStageExecution} — closest existing analogue for "consumes child's + * reduce output, feeds parent."
    • + *
  • {@link AnalyticsSearchTransportService#dispatchFragmentStreaming} — pattern for + * Arrow-streaming RPC; the fetch transport is similar but with a different request + * type.
    • + *
+ * + * @opensearch.internal + */ +public final class LateMaterializationStageExecution extends AbstractStageExecution implements DataConsumer, DataProducer { + + private static final Logger logger = LogManager.getLogger(LateMaterializationStageExecution.class); + + private final QueryContext config; + private final ExchangeSink parentSink; + private final ClusterService clusterService; + private final AnalyticsSearchTransportService transport; + + /** + * Sink the child Sort+Limit stage feeds into (Phase A input). + * + *

Returned by {@link #inputSink(int)}; the child stage writes K rows here. The + * stage execution drains this sink at {@link #start()} time. Today a + * {@link RowProducingSink} works as a buffer (it implements both {@link ExchangeSink} + * and {@link ExchangeSource}); a custom sink may be needed if Phase A wants to + * stream-decode rather than block-buffer. + */ + private final RowProducingSink childInputBuffer = new RowProducingSink(); + + /** Wrapper RelNode pulled from the stage's fragment — carries aboveAnchorPhysicalFields + aboveAnchorPhysicalFieldStorage. */ + private final OpenSearchLateMaterialization wrapper; + + /** + * Stage id of the SHARD_FRAGMENT descendant whose resolved targets feed our fetch + * dispatch. Walked once at construction; used in Phase C to look up + * {@code config.getResolvedTargets(shardStageId)}. + * + *

HACK: this id is paired with the side-table on {@link QueryContext} to map + * {@code ___ugsi → (DiscoveryNode, ShardId)}. See {@code QueryContext.resolvedTargetsByStage} + * for the rationale and revisit conditions — short version: stages leaving + * lookup state for other stages is a placeholder seam, not the long-term shape. + */ + private final int shardStageId; + + /** + * BackendId honored by the data node when serving fetches. Resolved by the factory + * from the SHARD_FRAGMENT descendant's first plan alternative. + * + *

FIXME assumes single-alternative shard plan, or all alternatives agree on backendId. + * When multi-backend per-shard plans become real the data node will need to feed back + * which backend it picked during the query phase so the fetch goes to the same one. + */ + private final String fetchBackendId; + + /** + * Per-node PendingExecutions, mirroring {@code ShardTaskRunner}'s pattern. Each node + * gets its own concurrency budget so a slow shard on one node can't block dispatches + * to others. {@link ConcurrentHashMap} because PendingExecutions callbacks run on + * arbitrary transport threads. + */ + private final Map pendingByNodeId = new ConcurrentHashMap<>(); + + public LateMaterializationStageExecution( + Stage stage, + QueryContext config, + ExchangeSink parentSink, + ClusterService clusterService, + AnalyticsSearchTransportService transport, + int shardStageId, + String fetchBackendId + ) { + super(stage, config.queryId(), config.operationListeners(), config.parentTask()); + this.config = config; + this.parentSink = parentSink; + this.clusterService = clusterService; + this.transport = transport; + this.shardStageId = shardStageId; + this.fetchBackendId = fetchBackendId; + this.runner = new LocalTaskRunner(config.localTaskExecutor()); + this.wrapper = RelNodeUtils.findNode(stage.getFragment(), OpenSearchLateMaterialization.class); + if (this.wrapper == null) { + throw new IllegalStateException( + "LATE_MATERIALIZATION stage " + stage.getStageId() + " missing OpenSearchLateMaterialization marker in fragment" + ); + } + logger.info( + "[LMStage] CREATED stageId={} aboveAnchorPhysicalFieldCount={} backendId={}", + stage.getStageId(), + wrapper.getAboveAnchorPhysicalFields().size(), + fetchBackendId + ); + } + + /** + * Phase A input. Child stage (Sort+Limit reduce) writes its K rows here. Returned + * sink must be the same one we drain in {@link #start()}. + */ + @Override + public ExchangeSink inputSink(int childStageId) { + return childInputBuffer; + } + + /** + * Phase D output for the parent stage. Parent stage (Post-Sort) reads from this + * source after we transition to SUCCEEDED. + */ + @Override + public ExchangeSource outputSource() { + if (parentSink instanceof ExchangeSource source) { + return source; + } + throw new UnsupportedOperationException("parentSink does not implement ExchangeSource: " + parentSink.getClass().getSimpleName()); + } + + /** + * Phase A drain. Walks the child's buffered VSRs in arrival (sort) order, extracts + * {@code ___row_id} and {@code ___ugsi} per row, and groups into per-shard fetch plans + * keyed by {@code ___ugsi}. Each plan carries parallel {@code long[] rowIds} + + * {@code int[] positions} arrays where {@code position} is the row's sort-order index + * (0..K-1) — used by the stitcher to place fetched cells back into sort order. + * + *

Order contract with the data node: {@code AnalyticsSearchService.executeFetchByRowIds} + * requires ascending rowIds in the request and the native side returns rows in the same + * ascending order. Phase B will sort each plan's {@code (rowIds, positions)} as a unit + * before dispatch; Phase C then walks request and response in lockstep — no per-row + * lookup map needed. + * + *

Phases B/C/D will follow in subsequent commits; for now Phase A populates + * {@link #drainedRowIds} / {@link #drainedUgsis} / {@link #drainedRowCount} and closes + * the parent sink without producing output rows so the surrounding wiring stays + * exercised. + */ + @Override + protected List materializeTasks() { + return List.of(new LocalStageTask(new StageTaskId(getStageId(), 0), this::drainAndClose)); + } + + /** Drained {@code ___row_id} per row, indexed by sort-order position. Populated by Phase A. */ + private long[] drainedRowIds; + + /** Drained {@code ___ugsi} per row, indexed by sort-order position. Parallel to {@link #drainedRowIds}. */ + private int[] drainedUgsis; + + /** Total rows drained from the child (= K, ≤ anchor's LIMIT). */ + private int drainedRowCount; + + /** Per-shard fetch plans keyed by {@code ___ugsi}. Built by Phase B. */ + private Map plansByUgsi; + + private void drainAndClose(ActionListener listener) { + try { + try { + drainAndGroupByUgsi(); + } finally { + // Drain copied helper columns into primitive arrays; the buffered VSRs are + // no longer needed and their query-allocator buffers must be released here. + childInputBuffer.close(); + } + plansByUgsi = groupByUgsi(); + scatterFetchAndStitch(listener); + // drainAndClose returns immediately after dispatching shards; the outer + // ActionListener fires asynchronously via Stitcher's onComplete callback when + // the last shard's response stream terminates. + } catch (Exception e) { + try { + parentSink.close(); + } catch (Exception ignore) {} + listener.onFailure(e); + } + } + + /** + * Phase C + D combined. Builds a {@link Stitcher} sized to {@link #drainedRowCount}, + * fires one async {@link FetchByRowIdsRequest} per shard via + * {@link AnalyticsSearchTransportService#dispatchFetchByRowIds}, and lets the per-shard + * {@link GatherListener}s feed batches into the stitcher as they arrive. The stitcher + * triggers {@code outerListener.onResponse} (or {@code onFailure}) when the last shard + * completes — no blocking on this thread. + */ + private void scatterFetchAndStitch(ActionListener outerListener) throws Exception { + if (drainedRowCount == 0) { + // K=0: nothing to dispatch. Parent stage owns parentSink.close() via + // onTerminalTransition; closing here races its reduce() task ("sink closed before reduce"). + outerListener.onResponse(null); + return; + } + + // Output schema: aboveAnchorPhysicalFields converted to Arrow Fields. ___row_id is + // NOT included — it's a helper consumed during stitch, never surfaced. + List aboveFields = wrapper.getAboveAnchorPhysicalFields(); + List outputFields = new ArrayList<>(aboveFields.size()); + for (RelDataTypeField f : aboveFields) { + outputFields.add(Field.nullable(f.getName(), ArrowCalciteTypes.toArrow(f.getType()))); + } + + // Fetch projection sent to the data node: ___row_id (required by the data-node + // contract — see AnalyticsSearchService.executeFetchByRowIds) + the physical fields + // the user wants displayed. + String[] columns = new String[aboveFields.size() + 1]; + columns[0] = OpenSearchLateMaterialization.ROW_ID_FIELD; + for (int i = 0; i < aboveFields.size(); i++) { + columns[i + 1] = aboveFields.get(i).getName(); + } + + Map targetsByUgsi = config.getResolvedTargets(shardStageId); + if (targetsByUgsi == null) { + throw new IllegalStateException( + "No resolved targets for shardStageId=" + shardStageId + " — shard fragment stage didn't record targets" + ); + } + + Stitcher stitcher = new Stitcher(config.bufferAllocator(), outputFields, drainedRowCount, plansByUgsi.size(), parentSink, () -> { + Exception failure = this.stitcher.surfaceableFailure(); + if (failure == null) { + outerListener.onResponse(null); + } else { + outerListener.onFailure(failure); + } + }); + this.stitcher = stitcher; + + for (Map.Entry entry : plansByUgsi.entrySet()) { + int ugsi = entry.getKey(); + ShardFetchPlan plan = entry.getValue(); + ShardExecutionTarget target = targetsByUgsi.get(ugsi); + if (target == null) { + stitcher.shardFailed(new IllegalStateException("No resolved target for ugsi=" + ugsi)); + continue; + } + FetchByRowIdsRequest request = new FetchByRowIdsRequest( + config.queryId(), + stage.getStageId(), + target.shardId(), + fetchBackendId, + plan.rowIds(), + columns + ); + // Per-node PendingExecutions: mirrors ShardTaskRunner — keeps a slow node from + // blocking dispatches to other nodes. + PendingExecutions pending = pendingByNodeId.computeIfAbsent( + target.node().getId(), + n -> new PendingExecutions(config.maxConcurrentShardRequests()) + ); + transport.dispatchFetchByRowIds(request, target.node(), new GatherListener(stitcher, plan), config.parentTask(), pending); + } + } + + /** Stashed for the {@code onComplete} closure to read {@link Stitcher#surfaceableFailure()}. */ + private Stitcher stitcher; + + /** + * Per-shard streaming response listener: forwards each Arrow batch to the + * {@link Stitcher} (which copies cells into the output VSR at the row's sort-order + * position), tracks {@code rowsCopiedSoFar} across batches, and signals shard + * completion / failure to the stitcher's countdown. + */ + private static final class GatherListener implements StreamingResponseListener { + private final Stitcher stitcher; + private final ShardFetchPlan plan; + private int rowsCopiedSoFar; + + GatherListener(Stitcher stitcher, ShardFetchPlan plan) { + this.stitcher = stitcher; + this.plan = plan; + } + + @Override + public void onStreamResponse(FragmentExecutionArrowResponse response, boolean isLast) { + VectorSchemaRoot batch = response.getRoot(); + try { + stitcher.acceptBatch(batch, plan.positions(), rowsCopiedSoFar); + rowsCopiedSoFar += batch.getRowCount(); + } catch (Exception e) { + stitcher.shardFailed(e); + return; + } finally { + // Stitcher.acceptBatch only reads + copyFromSafe; ownership of the response + // batch's query-allocator buffers stays with this listener. + if (batch != null) batch.close(); + } + if (isLast) stitcher.shardComplete(); + } + + @Override + public void onFailure(Exception e) { + stitcher.shardFailed(e); + } + } + + /** + * Reads K rows of {@code [reduce-set, ___row_id, ___ugsi]} from the child input + * batches (already in sort order) and copies the helper columns into two parallel + * primitive arrays indexed by sort-order position. Reduce-set columns are discarded. + * + *

Pre-sizes the arrays via {@link RowProducingSink#getRowCount()} — the sink + * maintains {@code totalRows} incrementally on each {@code feed()}, so this is O(1) + * and exact. No need to walk the batches twice. + * + *

Despite the name, the actual partition-by-{@code ___ugsi} into per-shard fetch + * plans is Phase B's job; here we only flatten the helper columns into the two + * parallel arrays Phase B will then walk to build its per-shard plans. + */ + private void drainAndGroupByUgsi() { + long total = childInputBuffer.getRowCount(); + if (total > Integer.MAX_VALUE) { + throw new IllegalStateException("Drained row count exceeds Integer.MAX_VALUE: " + total); + } + int K = (int) total; + drainedRowIds = new long[K]; + drainedUgsis = new int[K]; + + int position = 0; + for (VectorSchemaRoot vsr : childInputBuffer.readResult()) { + int rowIdIdx = vsr.getSchema().getFields().indexOf(vsr.getSchema().findField(OpenSearchLateMaterialization.ROW_ID_FIELD)); + int ugsiIdx = vsr.getSchema().getFields().indexOf(vsr.getSchema().findField(OpenSearchLateMaterialization.UGSI_FIELD)); + if (rowIdIdx < 0 || ugsiIdx < 0) { + throw new IllegalStateException( + "LM stage drain expected [" + + OpenSearchLateMaterialization.ROW_ID_FIELD + + ", " + + OpenSearchLateMaterialization.UGSI_FIELD + + "] in batch schema; got " + + vsr.getSchema() + ); + } + BigIntVector rowIds = (BigIntVector) vsr.getVector(rowIdIdx); + IntVector ugsis = (IntVector) vsr.getVector(ugsiIdx); + int rows = vsr.getRowCount(); + for (int i = 0; i < rows; i++) { + drainedRowIds[position] = rowIds.get(i); + drainedUgsis[position] = ugsis.get(i); + position++; + } + } + drainedRowCount = position; + logger.debug("[LMStage] phase-A stageId={} drainedRows={}", getStageId(), drainedRowCount); + } + + /** + * Phase B. Partitions the drained flat arrays by {@code ___ugsi} into one + * {@link ShardFetchPlan} per shard. Each plan's parallel {@code rowIds} and + * {@code positions} arrays are sorted by {@code rowId} ascending — required by the + * data node's fetch contract (ascending input rowIds, ascending response order). + * + *

Two passes over the flat arrays: one to count rows per shard so we can + * allocate exactly-sized per-shard arrays, one to fill them. No per-row boxing. + */ + private Map groupByUgsi() { + Map sizesByUgsi = new HashMap<>(); + for (int i = 0; i < drainedRowCount; i++) { + sizesByUgsi.merge(drainedUgsis[i], 1, Integer::sum); + } + Map rowIdsByUgsi = new HashMap<>(sizesByUgsi.size()); + Map positionsByUgsi = new HashMap<>(sizesByUgsi.size()); + Map cursorByUgsi = new HashMap<>(sizesByUgsi.size()); + for (Map.Entry e : sizesByUgsi.entrySet()) { + rowIdsByUgsi.put(e.getKey(), new long[e.getValue()]); + positionsByUgsi.put(e.getKey(), new int[e.getValue()]); + cursorByUgsi.put(e.getKey(), 0); + } + for (int i = 0; i < drainedRowCount; i++) { + int ugsi = drainedUgsis[i]; + int cur = cursorByUgsi.get(ugsi); + rowIdsByUgsi.get(ugsi)[cur] = drainedRowIds[i]; + positionsByUgsi.get(ugsi)[cur] = i; + cursorByUgsi.put(ugsi, cur + 1); + } + Map plans = new HashMap<>(sizesByUgsi.size()); + for (Integer ugsi : sizesByUgsi.keySet()) { + long[] rowIds = rowIdsByUgsi.get(ugsi); + int[] positions = positionsByUgsi.get(ugsi); + sortParallelByRowIdAscending(rowIds, positions); + plans.put(ugsi, new ShardFetchPlan(ugsi, rowIds, positions)); + } + logger.debug("[LMStage] phase-B stageId={} shards={}", getStageId(), plans.size()); + return plans; + } + + /** + * Sorts {@code rowIds} ascending in place, permuting {@code positions} along. + * Indirect sort via an index array — ~O(N log N) with bounded extra allocation; fine + * for typical per-shard N (= K/numShards ≤ a few thousand). + */ + static void sortParallelByRowIdAscending(long[] rowIds, int[] positions) { + int n = rowIds.length; + if (n <= 1) return; + Integer[] order = new Integer[n]; + for (int i = 0; i < n; i++) + order[i] = i; + Arrays.sort(order, (a, b) -> Long.compare(rowIds[a], rowIds[b])); + long[] sortedRowIds = new long[n]; + int[] sortedPositions = new int[n]; + for (int i = 0; i < n; i++) { + sortedRowIds[i] = rowIds[order[i]]; + sortedPositions[i] = positions[order[i]]; + } + System.arraycopy(sortedRowIds, 0, rowIds, 0, n); + System.arraycopy(sortedPositions, 0, positions, 0, n); + } + + /** + * Per-shard fetch plan: {@code rowIds} sorted ascending, {@code positions[j]} the + * sort-order position in the LM stage's drained-row index for {@code rowIds[j]}. + * Phase C dispatches {@code rowIds} to shard {@code ugsi}'s data node; the response + * arrives in the same ascending order, and Phase D writes each response cell into + * {@code stitched[positions[j]]}. + */ + record ShardFetchPlan(int ugsi, long[] rowIds, int[] positions) { + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageExecutionBuilder.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageExecutionBuilder.java index 389a93fb24f35..c0428221e665e 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageExecutionBuilder.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/StageExecutionBuilder.java @@ -13,6 +13,7 @@ import org.opensearch.analytics.exec.AnalyticsSearchTransportService; import org.opensearch.analytics.exec.QueryContext; import org.opensearch.analytics.exec.RowProducingSink; +import org.opensearch.analytics.exec.stage.coordinator.LateMaterializationStageExecutionFactory; import org.opensearch.analytics.exec.stage.coordinator.LocalComputeStageExecutionFactory; import org.opensearch.analytics.exec.stage.coordinator.PassThroughStageExecution; import org.opensearch.analytics.exec.stage.coordinator.ReduceStageExecutionFactory; @@ -62,6 +63,12 @@ public StageExecutionBuilder(ClusterService clusterService, AnalyticsSearchTrans registerFactory(StageExecutionType.COORDINATOR_REDUCE, new ReduceStageExecutionFactory()); registerFactory(StageExecutionType.LOCAL_PASSTHROUGH, (stage, sink, config) -> new PassThroughStageExecution(stage, config, sink)); registerFactory(StageExecutionType.LOCAL_COMPUTE, new LocalComputeStageExecutionFactory()); + // QTF (late-materialization) Scatter-Gather. Skeleton today — + // LateMaterializationStageExecution.start() throws UnsupportedOperationException. + // The DAG, FragmentConversion, and stage wiring are all in place; the four phases + // (drain → scatter fetch → gather → stitch) are documented inside the execution + // class and the new transport action / data-node handler are the remaining work. + registerFactory(StageExecutionType.LATE_MATERIALIZATION, new LateMaterializationStageExecutionFactory(clusterService, dispatcher)); } /** @@ -106,7 +113,10 @@ public StageExecution buildRootExecution(Stage rootStage, QueryContext config) { */ public StageExecution buildExecution(Stage stage, StageExecution parentExec, QueryContext config) { ExchangeSink sink = switch (stage.getExecutionType()) { - case SHARD_FRAGMENT, COORDINATOR_REDUCE, LOCAL_PASSTHROUGH, LOCAL_COMPUTE -> resolveRowSink(stage, parentExec); + case SHARD_FRAGMENT, COORDINATOR_REDUCE, LOCAL_PASSTHROUGH, LOCAL_COMPUTE, LATE_MATERIALIZATION -> resolveRowSink( + stage, + parentExec + ); }; return buildStageExecution(stage, sink, config); } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/Stitcher.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/Stitcher.java new file mode 100644 index 0000000000000..9a191c7b06e3e --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/Stitcher.java @@ -0,0 +1,200 @@ +/* + * 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.stage; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; +import org.opensearch.analytics.spi.ExchangeSink; + +import java.util.List; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.atomic.AtomicInteger; + +/** + * Per-LM-stage stitcher. Owns one pre-allocated output {@link VectorSchemaRoot} sized to K + * (the survivor row count) and accumulates fetch-by-rowid response batches from N shards + * into it. Per-shard listeners pass the position-array slice corresponding to each batch; + * the stitcher copies each response row to the output via {@link FieldVector#copyFromSafe}. + * + *

TODO incremental emission. Today the entire stitched VSR is emitted in a single + * {@code parentSink.feed} call after every shard's stream terminates — the post-LM stage + * (Stage 3) cannot start until LM's terminal SUCCEEDED transition. This forces Stage 3 to + * use the buffered MemTable sink. To enable streaming Stage 3 (which would let Camp-A + * post-LM ops — Filter / Project / hash Aggregate — process rows as they arrive), the + * stitcher would need to emit per-shard sub-batches at each shard's completion rather + * than holding everything in {@code output} until the end. The natural emission unit is + * "rows whose entire shard response has arrived" — those land at non-contiguous positions + * across the K-row range, so the wire-side position-sortedness invariant would have to be + * dropped (downstream sees rows in shard-completion order, not sort order). Camp-B post-LM + * ops (Sort / TopN / global-frame Aggregate) need the full set anyway and would still buffer. + * Defer until a real workload demands it. + * + *

Position is derived from the order contract with the data node: + * {@code AnalyticsSearchService.executeFetchByRowIds} requires ascending {@code rowIds} in + * the request and the native side returns rows in the same ascending order, so the + * shard's per-shard {@code positions[]} array (built during Phase B in the same order as + * the request's {@code rowIds[]}) lines up positionally with the response. The listener + * tracks where its previous batch left off via {@code rowsCopiedSoFar}; this stitcher reads + * {@code positions[rowsCopiedSoFar .. rowsCopiedSoFar + batchRows)}. + * + *

Lifecycle: + *

    + *
  1. {@link #acceptBatch} — listener calls per response batch from a shard. Synchronized + * so concurrent shards can't interleave Arrow buffer writes.
    2. + *
  2. {@link #shardComplete} — listener calls once when its shard's stream terminates. + * When all shards complete, the stitcher feeds the output to {@code parentSink} and + * runs {@code onComplete}.
    3. + *
  3. {@link #shardFailed} — listener calls on per-shard failure. Failures are collected + * (first surfaced as primary, others added as suppressed); completion path runs once + * and emits failure to the LM stage's outer listener.
    4. + *
+ * + * @opensearch.internal + */ +public final class Stitcher { + + private static final Logger logger = LogManager.getLogger(Stitcher.class); + + private final VectorSchemaRoot output; + private final int totalRows; + private final int aboveColCount; + private final ExchangeSink parentSink; + private final AtomicInteger pendingShards; + private final ConcurrentLinkedQueue failures = new ConcurrentLinkedQueue<>(); + private final Runnable onComplete; + private final Object outputLock = new Object(); + + public Stitcher( + BufferAllocator allocator, + List outputFields, + int totalRows, + int shardCount, + ExchangeSink parentSink, + Runnable onComplete + ) { + this.output = VectorSchemaRoot.create(new Schema(outputFields), allocator); + // TODO pre-size vectors to totalRows via setInitialCapacity to avoid copyFromSafe + // re-allocs during stitch. Today we let Arrow grow buffers dynamically — fine for + // small K but wasteful for large K (varchar buffers in particular grow geometrically). + this.output.allocateNew(); + this.totalRows = totalRows; + this.aboveColCount = outputFields.size(); + this.parentSink = parentSink; + this.pendingShards = new AtomicInteger(shardCount); + this.onComplete = onComplete; + } + + /** + * Copies one shard's response batch into the output VSR. Reads positions starting at + * {@code rowsCopiedSoFar} (the listener's per-shard "rows copied so far" counter) and + * places each response row at {@code positions[rowsCopiedSoFar + srcRow]}. + * + *

The response batch's column layout is {@code [___row_id, fetch-col-0, fetch-col-1, ...]} — + * fetch cols ordered as the request's {@code columns[]} (excluding the helper). The + * helper column is identified by name and skipped during copy. + */ + public void acceptBatch(VectorSchemaRoot batch, int[] positions, int rowsCopiedSoFar) { + synchronized (outputLock) { + int rowIdIdx = batch.getSchema().getFields().indexOf(batch.getSchema().findField(OpenSearchLateMaterialization.ROW_ID_FIELD)); + if (rowIdIdx < 0) { + throw new IllegalStateException( + "Fetch response missing " + OpenSearchLateMaterialization.ROW_ID_FIELD + " column; got " + batch.getSchema() + ); + } + int batchRows = batch.getRowCount(); + if (rowsCopiedSoFar + batchRows > positions.length) { + throw new IllegalStateException( + "Shard returned more rows than requested: positions.length=" + + positions.length + + " rowsCopiedSoFar=" + + rowsCopiedSoFar + + " batchRows=" + + batchRows + ); + } + int responseColCount = batch.getSchema().getFields().size(); + for (int srcRow = 0; srcRow < batchRows; srcRow++) { + int dstRow = positions[rowsCopiedSoFar + srcRow]; + int outCol = 0; + for (int srcCol = 0; srcCol < responseColCount; srcCol++) { + if (srcCol == rowIdIdx) continue; + if (outCol >= aboveColCount) { + throw new IllegalStateException( + "Response column count exceeds output schema: outCol=" + outCol + " aboveColCount=" + aboveColCount + ); + } + FieldVector srcVec = batch.getVector(srcCol); + FieldVector dstVec = output.getVector(outCol); + dstVec.copyFromSafe(srcRow, dstRow, srcVec); + outCol++; + } + } + } + } + + /** Signals that one shard's stream has terminated successfully. Triggers emit when last. */ + public void shardComplete() { + if (pendingShards.decrementAndGet() == 0) { + finish(); + } + } + + /** + * Records a per-shard failure and proceeds via the completion path. All shard failures + * are retained — the first becomes the primary surfaced exception, the rest land as + * suppressed exceptions so diagnostics aren't lost. + * + *

FIXME: today other in-flight shards continue running until they too complete or fail + * — we don't fast-cancel. When the LM stage gains a {@code cancel(reason)} path that + * propagates to in-flight transports, hook it here. + */ + public void shardFailed(Exception e) { + failures.offer(e); + if (pendingShards.decrementAndGet() == 0) { + finish(); + } + } + + private void finish() { + try { + if (failures.isEmpty()) { + output.setRowCount(totalRows); + parentSink.feed(output); + parentSink.close(); + logger.debug("[Stitcher] emitted rows={}", totalRows); + } else { + output.close(); + try { + parentSink.close(); + } catch (Exception ignore) {} + } + } finally { + onComplete.run(); + } + } + + /** + * Returns the surfaceable failure: the first one collected, with subsequent failures + * attached as {@code addSuppressed}. {@code null} if no shard failed. + */ + public Exception surfaceableFailure() { + Exception primary = failures.poll(); + if (primary == null) return null; + for (Exception next; (next = failures.poll()) != null;) { + primary.addSuppressed(next); + } + return primary; + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LateMaterializationStageExecutionFactory.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LateMaterializationStageExecutionFactory.java new file mode 100644 index 0000000000000..2afa474dd1ae9 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/LateMaterializationStageExecutionFactory.java @@ -0,0 +1,70 @@ +/* + * 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.stage.coordinator; + +import org.opensearch.analytics.exec.AnalyticsSearchTransportService; +import org.opensearch.analytics.exec.QueryContext; +import org.opensearch.analytics.exec.stage.LateMaterializationStageExecution; +import org.opensearch.analytics.exec.stage.StageExecution; +import org.opensearch.analytics.exec.stage.StageExecutionFactory; +import org.opensearch.analytics.planner.dag.Stage; +import org.opensearch.analytics.planner.dag.StageExecutionType; +import org.opensearch.analytics.spi.ExchangeSink; +import org.opensearch.cluster.service.ClusterService; + +/** + * Builds executions for {@link StageExecutionType#LATE_MATERIALIZATION} (QTF Scatter-Gather) + * stages. Pulls the {@code OpenSearchLateMaterialization} marker out of the stage's fragment + * for fetch-list metadata and hands the resulting context to {@link LateMaterializationStageExecution}. + * + *

This stage has no Substrait fragment ({@link + * org.opensearch.analytics.planner.dag.FragmentConversionDriver} skips it), so unlike + * {@link ReduceStageExecutionFactory} we do not pull plan bytes / sink providers / instructions + * from the stage. The wrapper RelNode itself carries everything we need: the fetch list + * (which columns to fetch) and per-column storage info (how the data-node should read them). + * + * @opensearch.internal + */ +public final class LateMaterializationStageExecutionFactory implements StageExecutionFactory { + + private final ClusterService clusterService; + private final AnalyticsSearchTransportService transport; + + public LateMaterializationStageExecutionFactory(ClusterService clusterService, AnalyticsSearchTransportService transport) { + this.clusterService = clusterService; + this.transport = transport; + } + + @Override + public StageExecution createExecution(Stage stage, ExchangeSink sink, QueryContext config) { + Stage shardStage = findShardFragmentDescendant(stage); + if (shardStage == null) { + throw new IllegalStateException("LATE_MATERIALIZATION stage " + stage.getStageId() + " has no SHARD_FRAGMENT descendant"); + } + return new LateMaterializationStageExecution( + stage, + config, + sink, + clusterService, + transport, + shardStage.getStageId(), + shardStage.getPlanAlternatives().get(0).backendId() + ); + } + + /** DFS for the SHARD_FRAGMENT descendant; null if none. */ + private static Stage findShardFragmentDescendant(Stage stage) { + for (Stage child : stage.getChildStages()) { + if (child.getExecutionType() == StageExecutionType.SHARD_FRAGMENT) return child; + Stage deeper = findShardFragmentDescendant(child); + if (deeper != null) return deeper; + } + return null; + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecution.java index 769a925392a11..2f237566d9b97 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecution.java @@ -8,12 +8,14 @@ package org.opensearch.analytics.exec.stage.coordinator; +import org.apache.arrow.memory.BufferAllocator; import org.opensearch.analytics.backend.ExchangeSource; import org.opensearch.analytics.exec.QueryContext; import org.opensearch.analytics.exec.stage.AbstractStageExecution; import org.opensearch.analytics.exec.stage.SinkProvidingStageExecution; import org.opensearch.analytics.exec.stage.StageTask; import org.opensearch.analytics.exec.stage.StageTaskId; +import org.opensearch.analytics.planner.dag.InputSinkDecorator; import org.opensearch.analytics.planner.dag.Stage; import org.opensearch.analytics.spi.ExchangeSink; import org.opensearch.analytics.spi.MultiInputExchangeSink; @@ -39,12 +41,14 @@ public final class ReduceStageExecution extends AbstractStageExecution implement private final ReducingExchangeSink backendSink; private final ExchangeSink downstream; private final Executor reduceExecutor; + private final BufferAllocator allocator; public ReduceStageExecution(Stage stage, QueryContext config, ReducingExchangeSink backendSink, ExchangeSink downstream) { super(stage, config.queryId(), config.operationListeners(), config.parentTask()); this.backendSink = backendSink; this.downstream = downstream; this.reduceExecutor = config.reduceExecutor(); + this.allocator = config.bufferAllocator(); this.runner = new LocalTaskRunner(config.schedulerExecutor()); } @@ -62,10 +66,17 @@ public void closeChildInput(int childStageId) { @Override public ExchangeSink inputSink(int childStageId) { - if (backendSink instanceof MultiInputExchangeSink multi) { - return multi.sinkForChild(childStageId); + InputSinkDecorator decorator = stage.getInputSinkDecorator(); + // sinkForChild routing only applies for Union/Join shapes with multiple child stages. + if (stage.getChildStages().size() > 1) { + if (decorator != null) { + throw new IllegalStateException( + "InputSinkDecorator on a multi-input reducer (stageId=" + getStageId() + ") is not supported" + ); + } + return ((MultiInputExchangeSink) backendSink).sinkForChild(childStageId); } - return backendSink; + return decorator != null ? decorator.decorate(backendSink, allocator) : backendSink; } @Override diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java index 242526af895af..8b426545db8c1 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/coordinator/ReduceStageExecutionFactory.java @@ -13,6 +13,7 @@ import org.opensearch.analytics.exec.stage.StageExecutionFactory; import org.opensearch.analytics.planner.dag.Stage; import org.opensearch.analytics.planner.dag.StageExecutionType; +import org.opensearch.analytics.planner.dag.StagePlan; import org.opensearch.analytics.spi.BackendExecutionContext; import org.opensearch.analytics.spi.ExchangeSink; import org.opensearch.analytics.spi.ExchangeSinkContext; @@ -141,7 +142,10 @@ private static List buildChildInputs(Stage stage } List inputs = new ArrayList<>(children.size()); for (Stage child : children) { - byte[] producerPlanBytes = child.getPlanAlternatives().getFirst().convertedBytes(); + // postDecorationSchemaBytes wins over the producer's natural schema when an input + // decorator widens the wire schema (e.g. QTF's OrdinalAppendingSink). + StagePlan plan = child.getPlanAlternatives().getFirst(); + byte[] producerPlanBytes = plan.postDecorationSchemaBytes() != null ? plan.postDecorationSchemaBytes() : plan.convertedBytes(); inputs.add(new ExchangeSinkContext.ChildInput(child.getStageId(), producerPlanBytes)); } return inputs; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java index bf11446b26203..f787cec290b1e 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardFragmentStageExecution.java @@ -63,9 +63,15 @@ protected List materializeTasks() { List resolved = stage.getTargetResolver().resolve(clusterService.state(), null); // Empty list → base short-circuits to SUCCEEDED (nothing to dispatch). List tasks = new ArrayList<>(resolved.size()); + List shardTargets = new ArrayList<>(resolved.size()); for (int i = 0; i < resolved.size(); i++) { - tasks.add(new ShardStageTask(new StageTaskId(getStageId(), i), resolved.get(i))); + ExecutionTarget target = resolved.get(i); + tasks.add(new ShardStageTask(new StageTaskId(getStageId(), i), target)); + shardTargets.add((ShardExecutionTarget) target); } + // Side-table for cross-stage routing (e.g. QTF Phase C maps ___ugsi → target). + // See QueryContext.resolvedTargetsByStage Javadoc for HACK rationale. + config.recordResolvedTargets(getStageId(), shardTargets); return tasks; } @@ -93,7 +99,7 @@ public ExchangeSource outputSource() { * offload: reordering would let isLast race ahead and drop earlier batches via the * stage-terminal short-circuit. Inline also preserves end-to-end backpressure. */ - StreamingResponseListener responseListenerFor(ActionListener listener) { + StreamingResponseListener responseListenerFor(int sourceOrdinal, ActionListener listener) { return new StreamingResponseListener<>() { @Override public void onStreamResponse(FragmentExecutionArrowResponse response, boolean isLast) { @@ -107,7 +113,7 @@ public void onStreamResponse(FragmentExecutionArrowResponse response, boolean is return; } try { - outputSink.feed(vsr); + outputSink.feed(vsr, sourceOrdinal); } catch (Exception e) { // Sink didn't take ownership — close the VSR before surfacing. RuntimeException wrapped = new RuntimeException("Stage " + getStageId() + " sink feed failed", e); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardTaskRunner.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardTaskRunner.java index 8e17e40205d27..cec5b718487c2 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardTaskRunner.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/exec/stage/shard/ShardTaskRunner.java @@ -52,7 +52,13 @@ public void run(ShardStageTask task, ActionListener listener) { ShardExecutionTarget target = (ShardExecutionTarget) task.target(); FragmentExecutionRequest request = requestBuilder.apply(target); PendingExecutions pending = pendingFor(target); - transport.dispatchFragmentStreaming(request, target.node(), stage.responseListenerFor(listener), config.parentTask(), pending); + transport.dispatchFragmentStreaming( + request, + target.node(), + stage.responseListenerFor(target.ordinal(), listener), + config.parentTask(), + pending + ); } private PendingExecutions pendingFor(ShardExecutionTarget target) { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/ArrowCalciteTypes.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/ArrowCalciteTypes.java new file mode 100644 index 0000000000000..307c1abf53dcb --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/ArrowCalciteTypes.java @@ -0,0 +1,60 @@ +/* + * 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.planner; + +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.TimeUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.calcite.rel.type.RelDataType; + +/** + * Bidirectional Arrow ↔ Calcite type converter for single types. + * + *

Used by the QTF (late-materialization) Phase C in + * {@code LateMaterializationStageExecution} to translate the above-anchor physical fields' + * Calcite {@link RelDataType}s into Arrow {@link ArrowType}s for the fetch-stage output + * schema. The {@code AggregateFunction.ArrowToCalciteTypeMapper} (in the SPI module) handles + * the inverse direction for {@code IntermediateField} resolution; this class is kept as the + * single authority for the Calcite→Arrow direction needed outside that resolver. + * + *

FIXME [FixBeforeMainMerge] coverage gaps: TIMESTAMP currently hardcodes MILLISECOND + * (Calcite precision is ignored — see toArrow note), and several Calcite/Arrow types are + * still unmapped (TIMESTAMP_WITH_LOCAL_TIME_ZONE, DATE, TIME, SMALLINT, TINYINT, DECIMAL, + * Arrow Date/Time/Decimal/...). Audit and broaden before merge so the QTF path tolerates + * non-keyword/non-date columns end-to-end. + */ +public final class ArrowCalciteTypes { + + private ArrowCalciteTypes() {} + + /** + * Convert a Calcite {@link RelDataType} to the corresponding Arrow type. + */ + public static ArrowType toArrow(RelDataType t) { + return switch (t.getSqlTypeName()) { + case BIGINT -> new ArrowType.Int(64, true); + case INTEGER -> new ArrowType.Int(32, true); + case DOUBLE -> new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); + case REAL, FLOAT -> new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); + // Utf8View matches what the DataFusion/parquet path on the data node emits for + // string columns; switching here keeps the coordinator-side Stitcher's pre-allocated + // output type aligned so copyFromSafe doesn't trip on a VARCHAR/VIEWVARCHAR mismatch. + case VARCHAR, CHAR -> ArrowType.Utf8View.INSTANCE; + case VARBINARY, BINARY -> ArrowType.Binary.INSTANCE; + case BOOLEAN -> ArrowType.Bool.INSTANCE; + // TODO: TIMESTAMP_WITH_LOCAL_TIME_ZONE, DATE, TIME, SMALLINT, TINYINT, DECIMAL still missing. + // TODO: hardcoded MILLISECOND to match what DateParquetField emits on the data node; + // Calcite's reported precision doesn't track the wire-level Arrow precision today, so + // honouring t.getPrecision() here would break Stitcher copyFromSafe. Revisit when + // Calcite types carry the data-node-side Arrow precision faithfully. + case TIMESTAMP -> new ArrowType.Timestamp(TimeUnit.MILLISECOND, null); + default -> throw new IllegalArgumentException("Unsupported Calcite type: " + t.getSqlTypeName()); + }; + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java index 908797ae62489..394dbeeda9076 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerImpl.java @@ -36,6 +36,7 @@ import org.opensearch.analytics.planner.rules.OpenSearchFilterRule; import org.opensearch.analytics.planner.rules.OpenSearchJoinRule; import org.opensearch.analytics.planner.rules.OpenSearchJoinSplitRule; +import org.opensearch.analytics.planner.rules.OpenSearchLateMaterializationRewriter; import org.opensearch.analytics.planner.rules.OpenSearchProjectRule; import org.opensearch.analytics.planner.rules.OpenSearchSortRule; import org.opensearch.analytics.planner.rules.OpenSearchSortSplitRule; @@ -45,6 +46,7 @@ import org.opensearch.analytics.planner.rules.OpenSearchValuesRule; import java.util.List; +import java.util.Optional; /** * Central planner for the Analytics Plugin. @@ -62,10 +64,6 @@ * {@link RuleProfilingListener} (when profiling is enabled on * {@link PlannerContext}) can be threaded through every planner created here. * - *

TODO: eliminate copyToCluster — have frontends create RelNodes with Volcano cluster. - *

TODO: DAG construction (cut at exchange boundaries, build stage tree) - *

TODO: Per-stage plan forking (multiple plan generation) - *

TODO: Fragment conversion (backend.getFragmentConvertor()) *

TODO: Join strategy selection, sort removal via CBO * * @opensearch.internal @@ -105,11 +103,16 @@ public static RelNode runAllOptimizations(RelNode rawRelNode, PlannerContext con // AnnotatedPredicates under OR/NOT (Lucene call buys nothing in those positions). modifiedRelNode = cbo(modifiedRelNode, rawRelNode, context, listener); LOGGER.info("After CBO:\n{}", RelOptUtil.toString(modifiedRelNode)); + Optional lateMat = OpenSearchLateMaterializationRewriter.rewrite(modifiedRelNode); + if (lateMat.isPresent()) { + modifiedRelNode = lateMat.get(); + LOGGER.info("After late-materialization:\n{}", RelOptUtil.toString(modifiedRelNode)); + } if (listener != null) { RuleProfilingListener.PlannerProfile profile = listener.snapshot(); context.recordProfilingResults(profile); - LOGGER.info("Planner profile:\n{}", profile.format()); + LOGGER.info("Planner profile for raw RelNode is :\n{}", profile.format()); } return modifiedRelNode; } @@ -189,9 +192,24 @@ private static RelNode reduceExpressions(RelNode input, RuleProfilingListener li private static RelNode pushdownRules(RelNode input, RuleProfilingListener listener) { HepProgramBuilder builder = new HepProgramBuilder(); builder.addMatchOrder(HepMatchOrder.BOTTOM_UP); - // Push Filters below Project/Aggregate/Join. + // SORT_PROJECT_TRANSPOSE + PROJECT_MERGE assist QTF (late-materialization) detection. + // SqlToRelConverter shapes `SELECT ... ORDER BY UPPER(URL) LIMIT N` as + // Sort($1) ← Project(URL, UPPER(URL)) ← Scan + // (the order-by expression is materialized into the Project so the Sort can reference + // it as a slot). SORT_PROJECT_TRANSPOSE flips this to + // Project(URL, UPPER(URL)) ← Sort($1) ← Scan + // putting the Project above the Sort, which is the shape the QTF rewriter recognizes + // as "topmost above-anchor operator." Calcite's RelRoot.project() then trims the + // helper sort-key column from the user-visible output. PROJECT_MERGE collapses any + // adjacent Projects so the rewriter sees at most one Project layer above the anchor. builder.addRuleCollection( - List.of(CoreRules.FILTER_PROJECT_TRANSPOSE, CoreRules.FILTER_AGGREGATE_TRANSPOSE, CoreRules.FILTER_INTO_JOIN) + List.of( + CoreRules.FILTER_PROJECT_TRANSPOSE, + CoreRules.FILTER_AGGREGATE_TRANSPOSE, + CoreRules.FILTER_INTO_JOIN, + CoreRules.SORT_PROJECT_TRANSPOSE, + CoreRules.PROJECT_MERGE + ) ); // Merge adjacent Filters into one — must run after transposes so any // auto-injected NOT NULL collapses with the user's WHERE. diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java index f3af635bc5a7f..bc7cd5d43ba73 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/RelNodeUtils.java @@ -13,6 +13,12 @@ import org.apache.calcite.plan.hep.HepRelVertex; import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.rex.RexShuttle; import org.opensearch.analytics.planner.rel.OpenSearchAggregate; import org.opensearch.analytics.planner.rel.OpenSearchConvention; import org.opensearch.analytics.planner.rel.OpenSearchDistribution; @@ -25,8 +31,12 @@ import org.opensearch.analytics.planner.rel.OpenSearchTableScan; import org.opensearch.analytics.planner.rel.OpenSearchUnion; import org.opensearch.analytics.planner.rel.OpenSearchValues; +import org.opensearch.analytics.spi.FieldStorageInfo; +import java.util.HashSet; +import java.util.LinkedHashSet; import java.util.List; +import java.util.Set; /** * Copies an OpenSearch RelNode tree to a new cluster so all nodes register @@ -203,4 +213,95 @@ private static boolean collectIndices(RelNode node, java.util.Set indice return true; } + /** Collects every {@link RexInputRef} index appearing inside a {@link RexNode} tree. */ + public static Set collectInputRefs(RexNode node) { + Set out = new HashSet<>(); + node.accept(new RexShuttle() { + @Override + public RexNode visitInputRef(RexInputRef ref) { + out.add(ref.getIndex()); + return ref; + } + }); + return out; + } + + /** + * Resolves a derived expression to the ordered list of physical-field names it depends on, + * deduped by first-appearance. Used by {@link OpenSearchProject#getOutputFieldStorage} and + * {@link OpenSearchAggregate#getOutputFieldStorage} to populate + * {@link FieldStorageInfo#getDependsOnPhysicalCols} per Invariant 1 of the QTF v2 algorithm. + * + *

For each {@code RexInputRef} encountered (depth-first order): + *

    + *
  • If the input FSI at that index is non-derived, add its field name.
    • + *
  • If the input FSI at that index is derived, recurse into its + * {@code dependsOnPhysicalCols} (already resolved by the upstream operator).
    • + *
+ */ + public static LinkedHashSet resolvePhysicalDeps(RexNode node, List inputStorage) { + LinkedHashSet deps = new LinkedHashSet<>(); + node.accept(new RexShuttle() { + @Override + public RexNode visitInputRef(RexInputRef ref) { + int idx = ref.getIndex(); + if (idx >= inputStorage.size()) { + throw new IllegalStateException( + "RexInputRef[" + + idx + + "] has no matching FieldStorageInfo entry " + + "(input only declares " + + inputStorage.size() + + " columns) — " + + "the upstream operator did not record storage for every output column" + ); + } + FieldStorageInfo src = inputStorage.get(idx); + if (src.isDerived()) { + deps.addAll(src.getDependsOnPhysicalCols()); + } else { + deps.add(src.getFieldName()); + } + return ref; + } + }); + return deps; + } + + /** + * Returns a copy of {@code base} with one extra field {@code (name, type)} appended. + * Used by rewrites that augment a rowType with synthetic helper columns. + */ + public static RelDataType appendField(RelDataTypeFactory typeFactory, RelDataType base, String name, RelDataType type) { + RelDataTypeFactory.Builder builder = typeFactory.builder(); + for (RelDataTypeField f : base.getFieldList()) { + builder.add(f.getName(), f.getType()); + } + builder.add(name, type); + return builder.build(); + } + + /** + * {@link RexShuttle} that rewrites every {@link RexInputRef} via {@code remap[oldIdx]}. + * Throws when {@code remap[oldIdx] < 0} (referenced column was dropped). Output ref's + * type is sourced from {@code newRowType}. + */ + public static final class IndexRemapShuttle extends RexShuttle { + private final int[] remap; + private final RelDataType newRowType; + + public IndexRemapShuttle(int[] remap, RelDataType newRowType) { + this.remap = remap; + this.newRowType = newRowType; + } + + @Override + public RexNode visitInputRef(RexInputRef ref) { + int newIdx = remap[ref.getIndex()]; + if (newIdx < 0) { + throw new IllegalStateException("RexInputRef references dropped column at original idx " + ref.getIndex()); + } + return new RexInputRef(newIdx, newRowType.getFieldList().get(newIdx).getType()); + } + } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java index 33b3e4780542a..42e914b74b9f5 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DAGBuilder.java @@ -9,10 +9,12 @@ package org.opensearch.analytics.planner.dag; import org.apache.calcite.rel.RelNode; +import org.opensearch.analytics.exec.OrdinalAppendingSink; import org.opensearch.analytics.planner.CapabilityRegistry; import org.opensearch.analytics.planner.CapabilityResolutionUtils; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan; import org.opensearch.analytics.planner.rel.OpenSearchTableScan; @@ -31,6 +33,17 @@ * below becomes a child stage and the reducer's own {@link ExchangeInfo} drives * the parent stage's input wiring. Stage IDs are assigned bottom-up. * + *

TODO move DAGBuilder AFTER PlanForker. Today this runs pre-fork while the + * CBO output may carry multiple viable backends per operator, so the cut helpers + * (e.g. {@code cutAtLateMaterialization}) have to pick a backend from the viable + * list to fetch an {@code ExchangeSinkProvider} — at this point the sink-provider + * choice is technically ambiguous. Running DAGBuilder after PlanForker would + * collapse each operator's viable list to a single resolved backend per + * alternative; the cut helpers could then assert exactly one viable backend and + * throw an exception if not, instead of silently picking the first. Cleaner + * separation of concerns: PlanForker does backend resolution, DAGBuilder does + * stage cuts. + * * @opensearch.internal */ public class DAGBuilder { @@ -52,6 +65,13 @@ public static QueryDAG build( // Cut directly: child stage is the subtree below, root fragment is // ExchangeReducer → StageInputScan. rootFragment = cutAtExchange(reducer, counter, childStages, registry, clusterService, indexNameExpressionResolver); + } else if (cboOutput instanceof OpenSearchLateMaterialization lm) { + // LM at root, no above-ops (e.g. `source = t | where ... | sort col | head N`): + // promote the LM stage to rootStage and skip the synthetic post-LM stage that would + // wrap a bare StageInputScan placeholder. + cutAtLateMaterialization(lm, counter, childStages, registry, clusterService, indexNameExpressionResolver); + assert childStages.size() == 1 : "cutAtLateMaterialization must add exactly one child (the LM stage)"; + return new QueryDAG(newQueryId(), childStages.getFirst()); } else { rootFragment = sever(cboOutput, counter, childStages, registry, clusterService, indexNameExpressionResolver); } @@ -80,7 +100,16 @@ public static QueryDAG build( : null; Stage rootStage = new Stage(counter[0]++, rootFragment, childStages, null, sinkProvider, rootTargetResolver); - return new QueryDAG(UUID.randomUUID().toString(), rootStage); + return new QueryDAG(newQueryId(), rootStage); + } + + /** + * Mints a per-DAG queryId. TODO revisit if uniqueness is relied upon for correctness — + * random UUID v4 is statistically safe but doesn't coordinate with task / context-id + * allocation elsewhere in the engine. + */ + private static String newQueryId() { + return UUID.randomUUID().toString(); } private static RelNode sever( @@ -95,6 +124,8 @@ private static RelNode sever( for (RelNode input : node.getInputs()) { if (input instanceof OpenSearchExchangeReducer reducer) { newInputs.add(cutAtExchange(reducer, counter, childStages, registry, clusterService, indexNameExpressionResolver)); + } else if (input instanceof OpenSearchLateMaterialization lm) { + newInputs.add(cutAtLateMaterialization(lm, counter, childStages, registry, clusterService, indexNameExpressionResolver)); } else { newInputs.add(sever(input, counter, childStages, registry, clusterService, indexNameExpressionResolver)); } @@ -110,6 +141,116 @@ private static RelNode sever( return changed ? node.copy(node.getTraitSet(), newInputs) : node; } + /** + * Cuts at an {@link OpenSearchLateMaterialization} wrapper. Two cuts happen: + * + *

    + *
  1. Reduce child: the wrapper's input subtree (Sort+Limit + ER + scans + * below) becomes the LM stage's child stage — a {@code COORDINATOR_REDUCE} + * gathering shard scans. QTF only fires multi-shard (single-shard collapse + * short-circuits in the rewriter), so the reduce child always has + * grandchildren and always gets a sink provider.
    2. + *
  2. LM stage itself: a fresh {@link Stage} with fragment + * {@code Wrapper ← StageInputScan(reduce-child)}. Returned to the caller + * as a {@link OpenSearchStageInputScan} so the caller's parent fragment + * slots in a schema-bearing placeholder.
    3. + *
+ * + *

The caller (the {@link #sever} walk for the wrapper's parent) attaches whatever + * post-LM ops sit above the wrapper on top of the returned StageInputScan. Those ops + * end up in a vanilla {@code COORDINATOR_REDUCE} stage that runs them via Substrait + * over the LM stage's stitched output. The LM stage itself runs Java-only + * scatter/gather/stitch. + * + *

If the wrapper has no parent ops (the no-above-ops case), the caller's parent + * fragment will just BE this StageInputScan. {@link #build} promotes the LM stage + * to root in that case to avoid a degenerate empty COORDINATOR_REDUCE wrapper. + */ + private static RelNode cutAtLateMaterialization( + OpenSearchLateMaterialization lm, + int[] counter, + List parentChildStages, + CapabilityRegistry registry, + ClusterService clusterService, + IndexNameExpressionResolver indexNameExpressionResolver + ) { + // 1. Reduce child — Sort+Limit reduce above shard scans. Multi-shard QTF only. + List reduceChildren = new ArrayList<>(); + RelNode reduceFragment = sever(lm.getInput(), counter, reduceChildren, registry, clusterService, indexNameExpressionResolver); + if (reduceChildren.isEmpty()) { + throw new IllegalStateException( + "QTF rewriter fired but the wrapper's input has no ExchangeReducer below it — " + + "single-shard collapse should have short-circuited the rewriter." + ); + } + int reduceStageId = counter[0]++; + List reduceViable = CapabilityResolutionUtils.filterByReduceCapability(registry, lm.getViableBackends()); + ExchangeSinkProvider reduceSinkProvider = registry.getBackend(reduceViable.getFirst()).getExchangeSinkProvider(); + Stage reduceStage = new Stage( + reduceStageId, + reduceFragment, + reduceChildren, + /*exchangeInfo=*/ null, + reduceSinkProvider, + /*targetResolver=*/ null + ); + // Reducer feeds the LM stage. Stamp every shard's batches with their target.ordinal() + // as ___ugsi BEFORE the backend's reduce sees them so the LM stage can group rows by + // source shard for fan-out fetches. + reduceStage.setInputSinkDecorator( + (sink, allocator) -> new OrdinalAppendingSink(sink, allocator, OpenSearchLateMaterialization.UGSI_FIELD) + ); + + // 2. LM stage itself — fragment is the wrapper rooted at StageInputScan(reduceStage). + OpenSearchRelNode lmInput = (OpenSearchRelNode) lm.getInput(); + OpenSearchStageInputScan reduceStageInput = new OpenSearchStageInputScan( + lm.getCluster(), + lm.getTraitSet(), + reduceStageId, + lm.getInput().getRowType(), + lm.getViableBackends(), + lmInput.getOutputFieldStorage() + ); + OpenSearchLateMaterialization lmFragment = new OpenSearchLateMaterialization( + lm.getCluster(), + lm.getTraitSet(), + reduceStageInput, + lm.getAboveAnchorPhysicalFields(), + lm.getAboveAnchorPhysicalFieldStorage(), + lm.getViableBackends() + ); + int lmStageId = counter[0]++; + Stage lmStage = new Stage( + lmStageId, + lmFragment, + List.of(reduceStage), + /*exchangeInfo=*/ null, + /*sinkProvider=*/ null, + /*targetResolver=*/ null + ); + parentChildStages.add(lmStage); + + // 3. Hand back StageInputScan(LM) so post-LM ops end up in their own COORDINATOR_REDUCE. + // Schema is the wrapper's output rowType (= aboveAnchorPhysicalFields). + // + // TODO Stage 3 sink mode. The COORDINATOR_REDUCE that wraps the post-LM ops currently + // inherits whatever sink the cluster setting selects (streaming vs memtable). For QTF + // it should be memtable: LM emits a single VSR after full stitch (see Stitcher TODO), + // so streaming buys nothing and the eager-scheduling deadlock-avoidance the streaming + // sink is designed for doesn't apply. Once Stitcher supports incremental emission, we + // can pick streaming for Camp-A post-LM ops (Filter/Project/hash Aggregate) and keep + // memtable for Camp-B (Sort/TopN/global Aggregate). Detect post-LM op shape here and + // hand a per-stage hint to the sink-provider selection. + return new OpenSearchStageInputScan( + lm.getCluster(), + lm.getTraitSet(), + lmStageId, + lmFragment.getRowType(), + lm.getViableBackends(), + lm.getAboveAnchorPhysicalFieldStorage() + ); + } + private static RelNode cutAtExchange( OpenSearchExchangeReducer reducer, int[] counter, @@ -142,15 +283,16 @@ private static RelNode cutAtExchange( new Stage(childStageId, childFragment, grandchildren, reducer.getExchangeInfo(), childSinkProvider, targetResolver) ); - // Replace the reducer's input with a StageInputScan placeholder. - // The root fragment ends at the reducer; the child stage fragment starts below it. - // StageInputScan signals where the Scheduler feeds Arrow batches from the child stage. + // Use the reducer's OUTPUT rowType so QTF's appended ___ugsi (set on erRowType by the + // rewriter) flows into the parent stage's partition schema. No-op for non-QTF reducers. + OpenSearchRelNode reducerInput = (OpenSearchRelNode) reducer.getInput(); OpenSearchStageInputScan stageInput = new OpenSearchStageInputScan( reducer.getCluster(), reducer.getTraitSet(), childStageId, - reducer.getInput().getRowType(), - reducer.getViableBackends() + reducer.getRowType(), + reducer.getViableBackends(), + reducerInput.getOutputFieldStorage() ); return new OpenSearchExchangeReducer( reducer.getCluster(), diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java index 54f930322de95..59aa34bf35e29 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/DistributedAggregateRewriter.java @@ -89,7 +89,8 @@ static RelNode rewrite(OpenSearchAggregate finalAgg) { stageInput.getTraitSet(), stageInput.getChildStageId(), overriddenExchangeType, - stageInput.getViableBackends() + stageInput.getViableBackends(), + stageInput.getOutputFieldStorage() ); newFinalInput = exchange.copy(exchange.getTraitSet(), List.of(newStageInput)); } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java index bddaf04b1262e..a7da71ba0c395 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/FragmentConversionDriver.java @@ -10,6 +10,7 @@ import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.logical.LogicalFilter; +import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rex.RexBuilder; import org.apache.calcite.rex.RexCall; import org.apache.calcite.rex.RexNode; @@ -25,6 +26,7 @@ import org.opensearch.analytics.planner.rel.OpenSearchAggregate; import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer; import org.opensearch.analytics.planner.rel.OpenSearchFilter; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan; import org.opensearch.analytics.planner.rel.OpenSearchTableScan; @@ -92,6 +94,17 @@ private static void convertStage(Stage stage, CapabilityRegistry registry) { for (Stage child : stage.getChildStages()) { convertStage(child, registry); } + // After children are converted, surface any decorator-induced schema delta as + // postDecorationSchemaBytes on the child plans. The reduce sink consults this when + // registering the partition so the catalog binding matches what the decorator delivers. + populatePostDecorationSchemas(stage, registry); + // LM stage runs Java-only scatter/gather/stitch — no Substrait compute. Emit a + // stub Read carrying the wrapper's output schema so Stage 3's parent reduce sink + // can derive the partition schema via the standard producerPlanBytes path. + if (stage.getExecutionType() == StageExecutionType.LATE_MATERIALIZATION) { + convertLateMaterializationStage(stage, registry); + return; + } List converted = new ArrayList<>(stage.getPlanAlternatives().size()); for (StagePlan plan : stage.getPlanAlternatives()) { AnalyticsSearchBackendPlugin backend = registry.getBackend(plan.backendId()); @@ -128,6 +141,59 @@ private static void convertStage(Stage stage, CapabilityRegistry registry) { } } + /** + * Detect a decorator-induced schema delta between a child stage's produced rowType and + * what the parent declares it expects, and emit a schema-only Read for partition registration. + * + *

The expected rowType lives on the parent's {@code OpenSearchStageInputScan(childStageId)} + * placeholder, which {@code DAGBuilder.cutAtExchange} sets to the reducer's output rowType + * (widened by the rewriter when a decorator like {@code OrdinalAppendingSink} runs). The + * produced rowType is the child's fragment top. When the two differ, the producer's natural + * schema undersells what arrives at the partition boundary post-decorator — the reduce sink + * needs the wider one. + * + *

TODO: Uses {@link RelNodeUtils#findNode} which only walks the first-input chain. Fine for + * QTF today (linear fragments). When QTF extends to Joins/Unions, multi-input fragments will + * have multiple {@code StageInputScan} leaves and this needs a multi-leaf walker. + */ + private static void populatePostDecorationSchemas(Stage stage, CapabilityRegistry registry) { + for (Stage child : stage.getChildStages()) { + OpenSearchStageInputScan inputScan = RelNodeUtils.findNode(stage.getFragment(), OpenSearchStageInputScan.class); + if (inputScan == null || inputScan.getChildStageId() != child.getStageId()) continue; + RelDataType produced = child.getFragment().getRowType(); + RelDataType expected = inputScan.getRowType(); + // Cheap int compare first, then digest string compare via equals. + if (produced.getFieldCount() == expected.getFieldCount() && produced.equals(expected)) continue; + + List updated = new ArrayList<>(child.getPlanAlternatives().size()); + for (StagePlan plan : child.getPlanAlternatives()) { + FragmentConvertor convertor = registry.getBackend(plan.backendId()).getFragmentConvertor(); + byte[] postDecorationBytes = convertor.convertSchemaOnlyRead(child.getStageId(), expected); + updated.add(plan.withPostDecorationSchemaBytes(postDecorationBytes)); + } + child.setPlanAlternatives(updated); + } + } + + /** + * Stub Substrait for the LM stage: a {@code Read { named_table: "input-"; + * base_schema: wrapperOutput }} the parent reduce sink can register against. The plan's + * {@code resolvedFragment} IS the wrapper (DAGBuilder builds it that way), but we don't + * convert it — the LM stage runs Java-only scatter/gather/stitch and emits no Substrait + * compute. We only need the schema-bearing Read so Stage 3's reduce sink derives a + * partition schema via the standard producerPlanBytes path. + */ + private static void convertLateMaterializationStage(Stage stage, CapabilityRegistry registry) { + List converted = new ArrayList<>(stage.getPlanAlternatives().size()); + for (StagePlan plan : stage.getPlanAlternatives()) { + OpenSearchLateMaterialization wrapper = (OpenSearchLateMaterialization) plan.resolvedFragment(); + FragmentConvertor convertor = registry.getBackend(plan.backendId()).getFragmentConvertor(); + byte[] bytes = convertor.convertSchemaOnlyRead(stage.getStageId(), wrapper.getRowType()); + converted.add(plan.withConvertedBytes(bytes, List.of()).withInstructions(List.of())); + } + stage.setPlanAlternatives(converted); + } + private static List assembleInstructions( AnalyticsSearchBackendPlugin backend, StagePlan plan, @@ -138,12 +204,15 @@ private static List assembleInstructions( LinkedList instructions = new LinkedList<>(); RelNode leaf = findLeaf(plan.resolvedFragment()); - if (leaf instanceof OpenSearchTableScan) { + if (leaf instanceof OpenSearchTableScan tableScan) { + // QTF narrows the Scan to [belowAnchorPhysicalFields..., __row_id__]; signal that to the + // backend so it picks the row-id-aware table provider regardless of delegation. + boolean requestsRowIds = tableScan.getRowType().getFieldNames().contains(OpenSearchLateMaterialization.ROW_ID_FIELD); List delegated = delegationBytes.getResult(); if (!delegated.isEmpty()) { - factory.createShardScanWithDelegationNode(treeShape, delegated.size()).ifPresent(instructions::add); + factory.createShardScanWithDelegationNode(treeShape, delegated.size(), requestsRowIds).ifPresent(instructions::add); } else { - factory.createShardScanNode().ifPresent(instructions::add); + factory.createShardScanNode(requestsRowIds).ifPresent(instructions::add); } } return instructions; @@ -371,6 +440,17 @@ private static byte[] convertReduceNode( RelNode finalAggFragment = openSearchNode.stripAnnotations(finalAggInputs, resolver); return convertor.convertFragment(finalAggFragment); } + + // LM-fed reduce stage (post-LM Stage 3): the LM stage emits a stitched VSR straight + // into this stage's input partition with no ExchangeReducer between, so the + // StageInputScan leaf sits directly under this op. Convert the whole node as one + // fragment; the convertor's rewriteStageInputScans turns the leaf into a NamedScan + // so isthmus can serialize it. + boolean allChildrenAreStageInputScan = !node.getInputs().isEmpty() + && node.getInputs().stream().allMatch(input -> input instanceof OpenSearchStageInputScan); + if (allChildrenAreStageInputScan) { + return convertor.convertFragment(strip(node, delegationBytes)); + } } // Multi-input node (Join, Union, Intersect, Minus): isthmus handles all of them diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/InputSinkDecorator.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/InputSinkDecorator.java new file mode 100644 index 0000000000000..7e90b250d0c38 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/InputSinkDecorator.java @@ -0,0 +1,43 @@ +/* + * 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.planner.dag; + +import org.apache.arrow.memory.BufferAllocator; +import org.opensearch.analytics.spi.ExchangeSink; + +/** + * Wraps a stage's incoming child sink with feature-specific decoration before + * the producer's batches reach it. Set on the {@link Stage} at DAG-build time + * (today only by {@code DAGBuilder.cutAtLateMaterialization}); applied at + * sink-resolution time inside the parent execution's {@code inputSink(...)}. + * + *

Implementations are stateless factories: each invocation builds a fresh + * wrapper around the supplied sink. The wrapper does not own the sink's + * lifecycle — close semantics delegate to the wrapped sink. + * + *

Today's only producer of this hook is the QTF (late-materialization) DAG + * cut, which installs an {@code OrdinalAppendingSink} so each shard's batches + * carry a {@code ___ugsi} ordinal before reduce. The interface is generic so + * future cross-cutting concerns can reuse it without touching reducer code. + * + * @opensearch.internal + */ +@FunctionalInterface +public interface InputSinkDecorator { + + /** + * Build a new {@link ExchangeSink} that wraps {@code sink} with this decorator's + * behavior. The returned sink delegates lifecycle (close) to the wrapped sink. + * + * @param sink the producer-facing sink to wrap (already resolved per-child if + * the underlying sink is a {@link org.opensearch.analytics.spi.MultiInputExchangeSink}) + * @param allocator allocator the decorator may use for any buffers it allocates + */ + ExchangeSink decorate(ExchangeSink sink, BufferAllocator allocator); +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ShardExecutionTarget.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ShardExecutionTarget.java index c13559d3fbf16..0af19eb424604 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ShardExecutionTarget.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ShardExecutionTarget.java @@ -15,18 +15,31 @@ * Execution target for a data node shard scan. The data node registers the shard * as a named table source before executing the fragment. * + *

{@link #ordinal()} is a per-shard sequence number assigned by the + * {@link TargetResolver} that produced this target. It is the index of this + * target within the resolver's output list and is stable for the lifetime of + * the resolved list. Generic per-shard property — consumed today by QTF + * (Late Materialization stamps it as the {@code ___ugsi} column on every batch + * the shard produces) but not LM-specific. + * * @opensearch.internal */ public final class ShardExecutionTarget extends ExecutionTarget { private final ShardId shardId; + private final int ordinal; - public ShardExecutionTarget(DiscoveryNode node, ShardId shardId) { + public ShardExecutionTarget(DiscoveryNode node, ShardId shardId, int ordinal) { super(node); this.shardId = shardId; + this.ordinal = ordinal; } public ShardId shardId() { return shardId; } + + public int ordinal() { + return ordinal; + } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ShardTargetResolver.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ShardTargetResolver.java index 9348489fdbf64..c230ccbe54418 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ShardTargetResolver.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/ShardTargetResolver.java @@ -59,12 +59,13 @@ public List resolve(ClusterState clusterState, @Nullable Object GroupShardsIterator shardIterators = clusterService.operationRouting() .searchShards(clusterState, concreteNames, null, null); List targets = new ArrayList<>(); + int ordinal = 0; for (ShardIterator shardIt : shardIterators) { ShardRouting shard = shardIt.nextOrNull(); if (shard != null) { DiscoveryNode node = clusterState.nodes().get(shard.currentNodeId()); if (node != null) { - targets.add(new ShardExecutionTarget(node, shard.shardId())); + targets.add(new ShardExecutionTarget(node, shard.shardId(), ordinal++)); } } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java index 6d4a942236c21..bd8b7b8976dad 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/Stage.java @@ -9,6 +9,8 @@ package org.opensearch.analytics.planner.dag; import org.apache.calcite.rel.RelNode; +import org.opensearch.analytics.planner.RelNodeUtils; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; import org.opensearch.analytics.spi.ExchangeSinkProvider; import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory; import org.opensearch.common.Nullable; @@ -47,6 +49,14 @@ public class Stage { private final StageExecutionType executionType; private List planAlternatives; private FragmentInstructionHandlerFactory instructionHandlerFactory; + /** + * Optional decorator wrapping this stage's incoming child sink. Set at DAG-build + * time (today only by {@code DAGBuilder.cutAtLateMaterialization}); applied at + * sink-resolution time inside the parent execution's {@code inputSink(...)}. + * Null when this stage doesn't need any decoration. + */ + @Nullable + private InputSinkDecorator inputSinkDecorator; public Stage( int stageId, @@ -128,11 +138,26 @@ public void setInstructionHandlerFactory(FragmentInstructionHandlerFactory instr this.instructionHandlerFactory = instructionHandlerFactory; } + @Nullable + public InputSinkDecorator getInputSinkDecorator() { + return inputSinkDecorator; + } + + public void setInputSinkDecorator(InputSinkDecorator inputSinkDecorator) { + this.inputSinkDecorator = inputSinkDecorator; + } + private StageExecutionType setStageExecutionType( ExchangeSinkProvider exchangeSinkProvider, TargetResolver targetResolver, RelNode fragment ) { + // QTF Scatter-Gather marker — orchestrates fetch-by-rowid + stitch internally, + // no targetResolver / no sinkProvider. Checked first so other categories don't + // accidentally claim the wrapper-stage. + if (RelNodeUtils.findNode(fragment, OpenSearchLateMaterialization.class) != null) { + return StageExecutionType.LATE_MATERIALIZATION; + } if (targetResolver != null) { return StageExecutionType.SHARD_FRAGMENT; } else if (hasComputeLeaf(fragment)) { diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StageExecutionType.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StageExecutionType.java index cafdf731b03b7..7d36e358c18dc 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StageExecutionType.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StageExecutionType.java @@ -42,5 +42,14 @@ public enum StageExecutionType { * literal-row sources ({@code LogicalValues}) and any future leaf operator * whose data lives on the coordinator rather than on a shard. */ - LOCAL_COMPUTE + LOCAL_COMPUTE, + /** + * QTF (late-materialization) Scatter-Gather stage. Drains the upstream Sort+Limit + * output, fans out fetch-by-rowid requests to data nodes (one per UGSI), stitches + * fetched columns back by row position, and emits the wrapper's output schema + * upstream. No Substrait fragment — the stage execution drives fetch transport + * directly. Marked by an {@code OpenSearchLateMaterialization} node in the + * stage's fragment. + */ + LATE_MATERIALIZATION } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StagePlan.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StagePlan.java index afa941ccaa5c3..c3084ebda370a 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StagePlan.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/StagePlan.java @@ -21,28 +21,36 @@ * are narrowed to exactly one backend, plus the converted bytes produced by * the backend's {@link FragmentConvertor}. * - * @param resolvedFragment fragment with all viableBackends narrowed to single choices - * @param backendId the primary backend for this plan - * @param convertedBytes backend-specific serialized plan bytes (null before conversion) - * @param delegatedExpressions serialized delegated expressions (empty if no delegation) - * @param instructions ordered instruction nodes for data-node execution (empty before resolution) + * @param resolvedFragment fragment with all viableBackends narrowed to single choices + * @param backendId the primary backend for this plan + * @param convertedBytes backend-specific serialized plan bytes (null before conversion) + * @param delegatedExpressions serialized delegated expressions (empty if no delegation) + * @param instructions ordered instruction nodes for data-node execution (empty before resolution) + * @param postDecorationSchemaBytes schema-only Read taking precedence over {@code convertedBytes} when + * the parent stage's input decorator widens this stage's wire schema; + * {@code null} when the producer's schema is already authoritative. * @opensearch.internal */ public record StagePlan(RelNode resolvedFragment, String backendId, byte[] convertedBytes, List delegatedExpressions, - List instructions) { + List instructions, byte[] postDecorationSchemaBytes) { /** Creates a StagePlan before conversion (bytes not yet available). */ public StagePlan(RelNode resolvedFragment, String backendId) { - this(resolvedFragment, backendId, null, List.of(), List.of()); + this(resolvedFragment, backendId, null, List.of(), List.of(), null); } /** Returns a copy with converted bytes and delegated expressions populated. */ public StagePlan withConvertedBytes(byte[] bytes, List delegatedExpressions) { - return new StagePlan(resolvedFragment, backendId, bytes, delegatedExpressions, List.of()); + return new StagePlan(resolvedFragment, backendId, bytes, delegatedExpressions, List.of(), postDecorationSchemaBytes); } /** Returns a copy with instructions populated. */ public StagePlan withInstructions(List instructions) { - return new StagePlan(resolvedFragment, backendId, convertedBytes, delegatedExpressions, instructions); + return new StagePlan(resolvedFragment, backendId, convertedBytes, delegatedExpressions, instructions, postDecorationSchemaBytes); + } + + /** Returns a copy with post-decoration schema bytes populated. */ + public StagePlan withPostDecorationSchemaBytes(byte[] bytes) { + return new StagePlan(resolvedFragment, backendId, convertedBytes, delegatedExpressions, instructions, bytes); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedPredicate.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedPredicate.java index b9955148506f2..da0eb5972bce0 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedPredicate.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/AnnotatedPredicate.java @@ -129,6 +129,26 @@ public RexNode makePlaceholder(RexBuilder rexBuilder) { return DelegatedPredicateFunction.makeCall(rexBuilder, annotationId); } + /** + * Override {@link RexCall#clone(RelDataType, List)} so that {@link org.apache.calcite.rex.RexShuttle}-based + * walks (e.g. {@code IndexRemapShuttle} during the QTF rewriter's narrowed-Scan rebuild) + * preserve the {@code AnnotatedPredicate} subclass when an operand is remapped. Without + * this override, {@code RexCall.clone} returns a plain {@code RexCall} carrying only the + * {@code ANNOTATED_PREDICATE} operator name — the {@code annotationId} / {@code viableBackends} + * / {@code performanceDelegationBackends} fields are lost, and {@code FragmentConversionDriver.strip}'s + * {@code instanceof AnnotatedPredicate} check fails to unwrap it, leaving the operator in + * the plan when it reaches the Substrait visitor. + */ + @Override + public RexCall clone(RelDataType type, List operands) { + if (operands.size() != 1) { + throw new IllegalArgumentException( + "AnnotatedPredicate must wrap exactly one operand (the original predicate); got " + operands.size() + ); + } + return new AnnotatedPredicate(type, operands.get(0), viableBackends, annotationId, performanceDelegationBackends); + } + @Override protected String computeDigest(boolean withType) { return "ANNOTATED_PREDICATE(id=" diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java index 619b0009a52bc..a3f8b0902da9e 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchAggregate.java @@ -21,12 +21,14 @@ import org.apache.calcite.rel.logical.LogicalAggregate; import org.apache.calcite.rel.metadata.RelMetadataQuery; import org.apache.calcite.rex.RexLiteral; +import org.apache.calcite.rex.RexNode; import org.apache.calcite.util.ImmutableBitSet; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.spi.FieldStorageInfo; import java.util.ArrayList; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; @@ -131,9 +133,32 @@ public List getOutputFieldStorage() { } } - // Agg results: derived columns with no physical storage + // Agg results: derived columns whose physical-deps are the union of arg refs' deps + // (preserving first-seen order across argList, then rexList). for (AggregateCall aggCall : getAggCallList()) { - outputStorage.add(FieldStorageInfo.derivedColumn(aggCall.getName(), aggCall.getType().getSqlTypeName())); + LinkedHashSet deps = new LinkedHashSet<>(); + for (int argIdx : aggCall.getArgList()) { + if (argIdx >= inputStorage.size()) { + throw new IllegalStateException( + "AggregateCall arg[" + + argIdx + + "] has no matching FieldStorageInfo entry " + + "(input only declares " + + inputStorage.size() + + " columns)" + ); + } + FieldStorageInfo src = inputStorage.get(argIdx); + if (src.isDerived()) { + deps.addAll(src.getDependsOnPhysicalCols()); + } else { + deps.add(src.getFieldName()); + } + } + for (RexNode rex : aggCall.rexList) { + deps.addAll(RelNodeUtils.resolvePhysicalDeps(rex, inputStorage)); + } + outputStorage.add(FieldStorageInfo.derivedColumn(aggCall.getName(), aggCall.getType().getSqlTypeName(), deps)); } return outputStorage; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java index 22b16d4f2ecb7..9b7a33a2f4bc3 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchExchangeReducer.java @@ -16,6 +16,7 @@ import org.apache.calcite.rel.RelWriter; import org.apache.calcite.rel.convert.ConverterImpl; import org.apache.calcite.rel.metadata.RelMetadataQuery; +import org.apache.calcite.rel.type.RelDataType; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.dag.ExchangeInfo; import org.opensearch.analytics.spi.FieldStorageInfo; @@ -34,10 +35,17 @@ public class OpenSearchExchangeReducer extends ConverterImpl implements OpenSear private final List viableBackends; private final ExchangeInfo exchangeInfo; + /** + * Non-null only when QTF (or a future rule) declares additional coord-side columns + * on the ER's output (e.g. {@code ___ugsi} appended at runtime by + * {@code ShardFragmentStageExecution.responseListenerFor}). Null in the default case + * so {@link ConverterImpl#deriveRowType()} drives. + */ + private final RelDataType overrideRowType; /** Convenience constructor — defaults to {@link ExchangeInfo#singleton()}. */ public OpenSearchExchangeReducer(RelOptCluster cluster, RelTraitSet traitSet, RelNode input, List viableBackends) { - this(cluster, traitSet, input, viableBackends, ExchangeInfo.singleton()); + this(cluster, traitSet, input, viableBackends, ExchangeInfo.singleton(), null); } public OpenSearchExchangeReducer( @@ -46,6 +54,23 @@ public OpenSearchExchangeReducer( RelNode input, List viableBackends, ExchangeInfo exchangeInfo + ) { + this(cluster, traitSet, input, viableBackends, exchangeInfo, null); + } + + /** + * Overload taking an explicit {@code overrideRowType}. Used by the QTF rule to declare + * {@code ___ugsi} on the ER's output schema — the column is appended coord-side at + * runtime in {@code ShardFragmentStageExecution.responseListenerFor} per task. Schema + * declaration here lets the reduce sink's schema-validation pass. + */ + public OpenSearchExchangeReducer( + RelOptCluster cluster, + RelTraitSet traitSet, + RelNode input, + List viableBackends, + ExchangeInfo exchangeInfo, + RelDataType overrideRowType ) { // ConverterImpl makes this a Calcite-recognized trait converter — inserted by // Volcano via OpenSearchDistributionTraitDef.convert when a downstream operator @@ -53,6 +78,12 @@ public OpenSearchExchangeReducer( super(cluster, null, traitSet, input); this.viableBackends = viableBackends; this.exchangeInfo = exchangeInfo; + this.overrideRowType = overrideRowType; + } + + @Override + public RelDataType deriveRowType() { + return overrideRowType != null ? overrideRowType : super.deriveRowType(); } @Override @@ -76,7 +107,7 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { - return new OpenSearchExchangeReducer(getCluster(), traitSet, sole(inputs), viableBackends, exchangeInfo); + return new OpenSearchExchangeReducer(getCluster(), traitSet, sole(inputs), viableBackends, exchangeInfo, overrideRowType); } /** @@ -101,12 +132,26 @@ public RelWriter explainTerms(RelWriter pw) { @Override public RelNode copyResolved(String backend, List children, List resolvedAnnotations) { - return new OpenSearchExchangeReducer(getCluster(), getTraitSet(), children.getFirst(), List.of(backend), exchangeInfo); + return new OpenSearchExchangeReducer( + getCluster(), + getTraitSet(), + children.getFirst(), + List.of(backend), + exchangeInfo, + overrideRowType + ); } @Override public RelNode stripAnnotations(List strippedChildren) { // ExchangeReducer is an infrastructure node — strip children but keep the node itself. - return new OpenSearchExchangeReducer(getCluster(), getTraitSet(), strippedChildren.getFirst(), viableBackends, exchangeInfo); + return new OpenSearchExchangeReducer( + getCluster(), + getTraitSet(), + strippedChildren.getFirst(), + viableBackends, + exchangeInfo, + overrideRowType + ); } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchLateMaterialization.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchLateMaterialization.java new file mode 100644 index 0000000000000..9df7455d0b467 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchLateMaterialization.java @@ -0,0 +1,181 @@ +/* + * 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.planner.rel; + +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelOptCost; +import org.apache.calcite.plan.RelOptPlanner; +import org.apache.calcite.plan.RelTraitSet; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelWriter; +import org.apache.calcite.rel.SingleRel; +import org.apache.calcite.rel.metadata.RelMetadataQuery; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.opensearch.analytics.spi.FieldStorageInfo; +import org.opensearch.index.engine.dataformat.DocumentInput; + +import java.util.ArrayList; +import java.util.List; +import java.util.Set; + +/** + * Stage marker for QTF (Query-Then-Fetch / late materialization). Sits above the anchor Sort. + * Input rowType is whatever the anchor produces ({@code [reduce-set, ___row_id, ___ugsi]} in + * the typical shape) — the wrapper does not derive its output rowType from input. + * Output rowType is {@code aboveAnchorPhysicalFields} only (in topmost-op order, named with + * physical names). The Scatter-Gather stage drains {@code (___row_id, ___ugsi)} pairs from + * input, fetches the {@code aboveAnchorPhysicalFields} per survivor, and emits batches + * matching the wrapper's output rowType. + * + *

{@code DAGBuilder} pattern-matches this node and emits a {@code LATE_MATERIALIZATION} + * stage with custom execution. No backend {@link org.opensearch.analytics.spi.FragmentConvertor} + * is invoked — the stage drives fetch transport directly. + * + *

TODO: revisit when extending QTF to Joins / Unions — multi-source row-id semantics + * (which side's {@code ___row_id} / {@code ___ugsi} survives, fetch fan-out across + * multiple input branches) are not handled here and the single-input {@code SingleRel} + * shape will need to grow. + * + *

TODO: this wrapper is effectively an exchange (data crosses node boundaries via + * scatter-gather fetch), but it isn't introduced by CBO trait propagation like + * {@code OpenSearchExchangeReducer} — the QTF rewriter inserts it as a post-CBO HEP + * pass, so its child stage receives a stitched VSR with no Substrait {@code ExchangeRel} + * representing the boundary. As a consequence, post-LM stages need a special + * {@code allChildrenAreStageInputScan} branch in {@link + * org.opensearch.analytics.planner.dag.FragmentConversionDriver}. Revisit whether QTF + * detection can be modeled as a custom distribution trait so CBO inserts the LM node + * the same way it inserts gather Reducers, removing the special-case in conversion. + * + * @opensearch.internal + */ +public class OpenSearchLateMaterialization extends SingleRel implements OpenSearchRelNode { + + /** Shard-produced row id, last column on Scan output, propagates up through Sort. */ + public static final String ROW_ID_FIELD = DocumentInput.ROW_ID_FIELD; + + /** Coord-appended UGSI (shardOrd + indexUUID + nodeId), declared on ER output. */ + public static final String UGSI_FIELD = "___ugsi"; + + /** Helper columns consumed internally by the Scatter-Gather stage (not in wrapper output). */ + public static final Set RESERVED_LATE_MATERIALIZATION_FIELDS = Set.of(ROW_ID_FIELD, UGSI_FIELD); + + private final List aboveAnchorPhysicalFields; + private final List aboveAnchorPhysicalFieldStorage; + private final List viableBackends; + + public OpenSearchLateMaterialization( + RelOptCluster cluster, + RelTraitSet traitSet, + RelNode input, + List aboveAnchorPhysicalFields, + List aboveAnchorPhysicalFieldStorage, + List viableBackends + ) { + super(cluster, traitSet, input); + if (aboveAnchorPhysicalFields.size() != aboveAnchorPhysicalFieldStorage.size()) { + throw new IllegalArgumentException( + "aboveAnchorPhysicalFields size " + + aboveAnchorPhysicalFields.size() + + " != aboveAnchorPhysicalFieldStorage size " + + aboveAnchorPhysicalFieldStorage.size() + ); + } + this.aboveAnchorPhysicalFields = List.copyOf(aboveAnchorPhysicalFields); + this.aboveAnchorPhysicalFieldStorage = List.copyOf(aboveAnchorPhysicalFieldStorage); + this.viableBackends = viableBackends; + this.rowType = computeRowType(cluster.getTypeFactory(), this.aboveAnchorPhysicalFields); + } + + /** Output rowType = aboveAnchorPhysicalFields, in iteration order. Input rowType is irrelevant. */ + private static RelDataType computeRowType(RelDataTypeFactory typeFactory, List aboveAnchorPhysicalFields) { + RelDataTypeFactory.Builder builder = typeFactory.builder(); + for (RelDataTypeField f : aboveAnchorPhysicalFields) { + builder.add(f.getName(), f.getType()); + } + return builder.build(); + } + + public List getAboveAnchorPhysicalFields() { + return aboveAnchorPhysicalFields; + } + + /** Per-column storage info for {@link #aboveAnchorPhysicalFields}, in the same order. */ + public List getAboveAnchorPhysicalFieldStorage() { + return aboveAnchorPhysicalFieldStorage; + } + + @Override + public List getViableBackends() { + return viableBackends; + } + + /** + * Output storage = {@code aboveAnchorPhysicalFieldStorage} only — same length and order + * as the wrapper's output rowType. Input storage is irrelevant; the wrapper exposes + * fetched physical fields, nothing else. + */ + @Override + public List getOutputFieldStorage() { + return aboveAnchorPhysicalFieldStorage; + } + + @Override + public RelNode copy(RelTraitSet traitSet, List inputs) { + return new OpenSearchLateMaterialization( + getCluster(), + traitSet, + sole(inputs), + aboveAnchorPhysicalFields, + aboveAnchorPhysicalFieldStorage, + viableBackends + ); + } + + @Override + public RelOptCost computeSelfCost(RelOptPlanner planner, RelMetadataQuery mq) { + // Reduces wire bytes by deferring aboveAnchorPhysicalFields until after Sort+Limit; cheap relative to ER. + return planner.getCostFactory().makeTinyCost(); + } + + @Override + public RelWriter explainTerms(RelWriter pw) { + List aboveAnchorPhysicalFieldNames = new ArrayList<>(aboveAnchorPhysicalFields.size()); + for (RelDataTypeField f : aboveAnchorPhysicalFields) { + aboveAnchorPhysicalFieldNames.add(f.getName()); + } + return super.explainTerms(pw).item("aboveAnchorPhysicalFields", aboveAnchorPhysicalFieldNames) + .item("viableBackends", viableBackends); + } + + @Override + public RelNode copyResolved(String backend, List children, List resolvedAnnotations) { + return new OpenSearchLateMaterialization( + getCluster(), + getTraitSet(), + children.getFirst(), + aboveAnchorPhysicalFields, + aboveAnchorPhysicalFieldStorage, + List.of(backend) + ); + } + + @Override + public RelNode stripAnnotations(List strippedChildren) { + // Stage marker — must be cut by DAGBuilder.cutAtLateMaterialization so it never + // reaches FragmentConversionDriver. If we hit this path, DAGBuilder didn't cut + // at the wrapper and the backend's FragmentConvertor would choke on an unknown + // RelNode. Fail loud instead. + throw new IllegalStateException( + "OpenSearchLateMaterialization reached FragmentConversionDriver — DAGBuilder " + + "must cut at the wrapper. This is a planner / DAGBuilder bug." + ); + } +} diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java index 2cebc689b0766..6220b16e62909 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchProject.java @@ -31,6 +31,7 @@ import java.util.ArrayList; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.function.Function; @@ -76,7 +77,8 @@ public List getOutputFieldStorage() { result.add(inputStorage.get(ref.getIndex())); } else { String fieldName = getRowType().getFieldList().get(i).getName(); - result.add(FieldStorageInfo.derivedColumn(fieldName, getRowType().getFieldList().get(i).getType().getSqlTypeName())); + LinkedHashSet deps = RelNodeUtils.resolvePhysicalDeps(expr, inputStorage); + result.add(FieldStorageInfo.derivedColumn(fieldName, getRowType().getFieldList().get(i).getType().getSqlTypeName(), deps)); } } return result; diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchStageInputScan.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchStageInputScan.java index d8c5e68df0a6f..66439a818f573 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchStageInputScan.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchStageInputScan.java @@ -34,17 +34,20 @@ public class OpenSearchStageInputScan extends AbstractRelNode implements OpenSea private final int childStageId; private final List viableBackends; + private final List outputFieldStorage; public OpenSearchStageInputScan( RelOptCluster cluster, RelTraitSet traitSet, int childStageId, RelDataType rowType, - List viableBackends + List viableBackends, + List outputFieldStorage ) { super(cluster, traitSet); this.childStageId = childStageId; this.viableBackends = viableBackends; + this.outputFieldStorage = outputFieldStorage; this.rowType = rowType; } @@ -59,12 +62,12 @@ public List getViableBackends() { @Override public List getOutputFieldStorage() { - return List.of(); + return outputFieldStorage; } @Override public OpenSearchStageInputScan copy(RelTraitSet traitSet, java.util.List inputs) { - return new OpenSearchStageInputScan(getCluster(), traitSet, childStageId, rowType, viableBackends); + return new OpenSearchStageInputScan(getCluster(), traitSet, childStageId, rowType, viableBackends, outputFieldStorage); } @Override @@ -79,7 +82,7 @@ public RelWriter explainTerms(RelWriter pw) { @Override public RelNode copyResolved(String backend, List children, List resolvedAnnotations) { - return new OpenSearchStageInputScan(getCluster(), getTraitSet(), childStageId, rowType, List.of(backend)); + return new OpenSearchStageInputScan(getCluster(), getTraitSet(), childStageId, rowType, List.of(backend), outputFieldStorage); } @Override diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java index 8909f637ccf1d..b1c6e258b915e 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rel/OpenSearchTableScan.java @@ -15,8 +15,8 @@ import org.apache.calcite.rel.RelNode; import org.apache.calcite.rel.RelWriter; import org.apache.calcite.rel.core.TableScan; -import org.apache.calcite.rel.logical.LogicalTableScan; import org.apache.calcite.rel.metadata.RelMetadataQuery; +import org.apache.calcite.rel.type.RelDataType; import org.opensearch.analytics.spi.FieldStorageInfo; import java.util.List; @@ -30,6 +30,12 @@ public class OpenSearchTableScan extends TableScan implements OpenSearchRelNode private final List viableBackends; private final List outputFieldStorage; + /** + * Non-null only when QTF (or a future rule) needs a rowType different from + * {@code getTable().getRowType()} — e.g. fetch cols dropped, {@code ___row_id} + * appended. Null in the default case so {@link TableScan#deriveRowType()} drives. + */ + private final RelDataType overrideRowType; public OpenSearchTableScan( RelOptCluster cluster, @@ -37,10 +43,33 @@ public OpenSearchTableScan( RelOptTable table, List viableBackends, List outputFieldStorage + ) { + this(cluster, traitSet, table, viableBackends, outputFieldStorage, null); + } + + /** + * Overload taking an explicit {@code overrideRowType}. Used by the QTF rule to + * narrow the scan to {@code [sort/filter cols, ___row_id]} after dropping the fetch + * list. {@code outputFieldStorage} must align 1:1 with {@code overrideRowType}'s + * fields (helper columns get synthetic {@link FieldStorageInfo} entries). + */ + public OpenSearchTableScan( + RelOptCluster cluster, + RelTraitSet traitSet, + RelOptTable table, + List viableBackends, + List outputFieldStorage, + RelDataType overrideRowType ) { super(cluster, traitSet, List.of(), table); this.viableBackends = viableBackends; this.outputFieldStorage = outputFieldStorage; + this.overrideRowType = overrideRowType; + } + + @Override + public RelDataType deriveRowType() { + return overrideRowType != null ? overrideRowType : super.deriveRowType(); } /** @@ -84,7 +113,7 @@ public List getOutputFieldStorage() { @Override public RelNode copy(RelTraitSet traitSet, List inputs) { - return new OpenSearchTableScan(getCluster(), traitSet, getTable(), viableBackends, outputFieldStorage); + return new OpenSearchTableScan(getCluster(), traitSet, getTable(), viableBackends, outputFieldStorage, overrideRowType); } @Override @@ -99,11 +128,16 @@ public RelWriter explainTerms(RelWriter pw) { @Override public RelNode copyResolved(String backend, List children, List resolvedAnnotations) { - return new OpenSearchTableScan(getCluster(), getTraitSet(), getTable(), List.of(backend), outputFieldStorage); + return new OpenSearchTableScan(getCluster(), getTraitSet(), getTable(), List.of(backend), outputFieldStorage, overrideRowType); } @Override public RelNode stripAnnotations(List strippedChildren) { - return LogicalTableScan.create(getCluster(), getTable(), List.of()); + // OpenSearchTableScan carries no operator annotations to strip and already exposes the + // correct schema via deriveRowType() (which honours overrideRowType when QTF narrows it). + // Returning this directly keeps the override visible to isthmus's Substrait conversion; + // converting to LogicalTableScan would defer to the underlying RelOptTable's wide rowType + // and silently drop helper columns like __row_id__. + return this; } } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java new file mode 100644 index 0000000000000..a839bc2fd6f90 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchLateMaterializationRewriter.java @@ -0,0 +1,641 @@ +/* + * 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.planner.rules; + +import org.apache.calcite.rel.RelCollations; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.rex.RexInputRef; +import org.apache.calcite.rex.RexNode; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.analytics.planner.RelNodeUtils; +import org.opensearch.analytics.planner.RelNodeUtils.IndexRemapShuttle; +import org.opensearch.analytics.planner.rel.OpenSearchAggregate; +import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer; +import org.opensearch.analytics.planner.rel.OpenSearchFilter; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; +import org.opensearch.analytics.planner.rel.OpenSearchProject; +import org.opensearch.analytics.planner.rel.OpenSearchRelNode; +import org.opensearch.analytics.planner.rel.OpenSearchSort; +import org.opensearch.analytics.planner.rel.OpenSearchTableScan; +import org.opensearch.analytics.spi.FieldStorageInfo; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** + * QTF (Query-Then-Fetch / late materialization) post-CBO rewrite. Two phases: + * + *

    + *
  1. Detect ({@link #detect}) — read-only walk that runs the allow-list checks, + * computes the {@link Detection} bundle ({@code aboveAnchorPhysicalFields}, + * {@code belowAnchorPhysicalFields}, {@code anchorSlotToPhysicalField}), and applies + * the skip predicate. Returns {@code null} when QTF doesn't apply.
    2. + *
  2. Rewrite ({@link #applyRewrite}) — pure transform consuming a {@link Detection}. + * Builds the narrowed Scan, walks the below chain, declares {@code ___ugsi} on the + * ExchangeReducer, swaps the wrapper in for the anchor, and remaps RexNodes in the + * above chain.
    3. + *
+ * + *

Skip predicate

+ * Fire QTF iff {@code aboveAnchorPhysicalFields - belowAnchorPhysicalFields} is non-empty. + * If every above-anchor reference is already in the reduce-set, the non-QTF path reads + * the same physical fields and skips the fetch round-trip — strictly cheaper. + * + *

Allow-lists

+ * Above-anchor: {@link OpenSearchProject} (no {@code RexOver}), {@link OpenSearchFilter}, + * {@link OpenSearchSort}. {@link OpenSearchAggregate} is rejected for now — group / agg-call + * remapping under a moving wrapper schema is a follow-up. + *

+ * Below-anchor intermediate ops: {@link OpenSearchFilter}, {@link OpenSearchExchangeReducer}, + * passthrough {@link OpenSearchProject}. + * + *

Plan shape after rewrite

+ * 
+ *   AboveOps (RexInputRefs remapped)
+ *   └── OpenSearchLateMaterialization        (output rowType = aboveAnchorPhysicalFields)
+ *        └── OpenSearchSort                  (anchor; collation remapped to narrowed-Scan space)
+ *             └── ...
+ *                  └── OpenSearchExchangeReducer  (rowType: [reduce-set, ___row_id, ___ugsi])
+ *                       └── ...
+ *                            └── OpenSearchTableScan   (narrowed: [reduce-set, ___row_id])
+ *
+ * + * @opensearch.internal + */ +public final class OpenSearchLateMaterializationRewriter { + + private static final Logger LOGGER = LogManager.getLogger(OpenSearchLateMaterializationRewriter.class); + + // TODO : Add support for Aggregate, Joins and Union here. + private static final Set> ABOVE_ALLOWED = Set.of( + OpenSearchFilter.class, + OpenSearchProject.class, + OpenSearchSort.class + ); + + private static final Set> BELOW_INTERMEDIATE_ALLOWED = Set.of( + OpenSearchFilter.class, + OpenSearchExchangeReducer.class, + OpenSearchProject.class + ); + + // TODO : [Human Generated] Don't Delete until fixed. + // TODO [Design] : We need to create Rewriting for Distributed Query Execution as a separate PlannerPhase. + // TODO : One categorization that applies is Correctness v/s Performance. So, a RewritePhase -> LateMatRewriter. + // TODO : Late Materialization is a rewrite done for performance. + // TODO : Rewrite for TopK Approximation is done for Correctness as a response to an User ExecutionHint in request. + + private OpenSearchLateMaterializationRewriter() {} + + /** Returns the rewritten root iff QTF matched and fired; {@link Optional#empty()} otherwise. */ + public static Optional rewrite(RelNode root) { + Detection detection = detect(root); + if (detection == null) return Optional.empty(); + LOGGER.debug( + "[QTF] fired: aboveAnchorPhysicalFields={}, belowAnchorPhysicalFields={}", + detection.aboveAnchorPhysicalFields(), + detection.belowAnchorPhysicalFields() + ); + return Optional.of(applyRewrite(root, detection)); + } + + // ── Phase 1 — Detect ─────────────────────────────────────────────── + + /** + * Walks the plan, validates allow-lists, computes the data structures rewrite needs, + * and applies the skip predicate. {@code null} return means "decline QTF." + */ + private static Detection detect(RelNode root) { + AnchorContext anchorCtx = findAnchor(root); + if (anchorCtx == null) return null; + if (!isAboveAllowed(anchorCtx.aboveAnchorOperators)) { + LOGGER.debug("[QTF] above-anchor allow-list rejected; skipping rewrite"); + return null; + } + + BelowChain belowChain = analyzeBelow(anchorCtx.anchor.getInput()); + if (belowChain == null) { + LOGGER.debug("[QTF] below-anchor allow-list rejected; skipping rewrite"); + return null; + } + + // Single-shard plans don't trigger late materialization. + if (!belowChain.hasExchangeReducer()) { + LOGGER.debug("[QTF] single-shard plan (no ExchangeReducer below anchor); skipping rewrite"); + return null; + } + + Set belowAnchorPhysicalFields = computeBelowAnchorPhysicalFields(anchorCtx.anchor, belowChain); + LinkedHashSet aboveAnchorPhysicalFields = computeAboveAnchorPhysicalFields( + anchorCtx.aboveAnchorOperators, + anchorCtx.anchor + ); + + // Skip predicate: aboveAnchorPhysicalFields - belowAnchorPhysicalFields must be non-empty. + boolean hasFetchOnly = aboveAnchorPhysicalFields.stream().anyMatch(name -> !belowAnchorPhysicalFields.contains(name)); + if (!hasFetchOnly) { + LOGGER.debug("[QTF] aboveAnchorPhysicalFields ⊆ belowAnchorPhysicalFields; QTF would not save any I/O — skipping"); + return null; + } + + List anchorSlotToPhysicalField = buildAnchorSlotToPhysicalField(belowChain); + + return new Detection( + anchorCtx.anchor, + anchorCtx.aboveAnchorOperators, + belowChain, + belowAnchorPhysicalFields, + aboveAnchorPhysicalFields, + anchorSlotToPhysicalField + ); + } + + /** + * Walks {@code root} downward along the single-input spine, picking the deepest + * {@link OpenSearchSort} with non-empty collation and non-null fetch as the anchor. + * Two-Sort shapes ({@code Sort(fetch) ← Sort(collation+fetch) ← ...}) put the lower + * Sort as anchor and the upper Sort in the above chain (allowed there). + * + *

Scope: single-input plans only. The walk terminates at any operator whose + * {@code getInputs().size() != 1} — multi-input ops (Join, Union, Intersect, Minus) and + * leaves (TableScan, Values). QTF across a Join or Union is not handled here and will + * need dedicated rewriters with multi-input semantics (per-branch detection, merged + * rebuild, and richer {@code ___row_id}/{@code ___ugsi} encoding to disambiguate which + * branch a survivor row came from). Tracked in {@link OpenSearchLateMaterialization}'s + * class-level TODO. + */ + private static AnchorContext findAnchor(RelNode root) { + List chainTopDown = new ArrayList<>(); + int deepestAnchorDepth = -1; + OpenSearchSort deepestAnchor = null; + + RelNode cur = RelNodeUtils.unwrapHep(root); + while (cur.getInputs().size() == 1) { + if (cur instanceof OpenSearchSort sort && !sort.getCollation().getFieldCollations().isEmpty() && sort.fetch != null) { + deepestAnchor = sort; + deepestAnchorDepth = chainTopDown.size(); + } + chainTopDown.add(cur); + cur = RelNodeUtils.unwrapHep(cur.getInput(0)); + } + if (deepestAnchor == null) return null; + return new AnchorContext(deepestAnchor, chainTopDown.subList(0, deepestAnchorDepth)); + } + + private static boolean isAboveAllowed(List chain) { + for (RelNode n : chain) { + if (!ABOVE_ALLOWED.contains(n.getClass())) return false; + if (n instanceof OpenSearchProject project && project.containsOver()) return false; + } + return true; + } + + /** + * Below-chain analysis. Walks {@code anchor.input} down to the Scan. Returns null + * when an op falls outside the below-allow-list, when a below-Project is non-passthrough, + * or when there's no Scan at the bottom. + */ + private static BelowChain analyzeBelow(RelNode subtree) { + List chain = new ArrayList<>(); + int[] belowProjOutToScan = null; + OpenSearchTableScan scan = null; + boolean hasExchangeReducer = false; + RelNode n = RelNodeUtils.unwrapHep(subtree); + while (n != null) { + if (n instanceof OpenSearchTableScan s) { + scan = s; + break; + } + if (!BELOW_INTERMEDIATE_ALLOWED.contains(n.getClass())) return null; + if (n instanceof OpenSearchExchangeReducer) hasExchangeReducer = true; + if (n instanceof OpenSearchProject p) { + if (belowProjOutToScan != null) { + LOGGER.debug("[QTF] multiple Projects below anchor — skipping"); + return null; + } + int[] outToScan = passthroughMap(p); + if (outToScan == null) { + // TODO: derived below-Project lost-opportunity case. Today's algorithm + // declines plans like: + // SELECT description FROM hits ORDER BY UPPER(URL) LIMIT 10 + // → Project(description) ← Sort($1 ASC) ← Project(description, UPPER(URL)) ← Scan + // The derived col (UPPER(URL)) is consumed only by the anchor's collation; we + // *could* push the derived expression above the wrapper, narrow the Scan to + // {URL}, sort below, and fetch {description} for survivors. Skipping for now — + // adding requires threading the derived RexNode into Detection and emitting + // a synthesized above-Project during rewrite. Separate slice. The non-QTF + // path remains correct in the meantime. + LOGGER.debug("[QTF] expression below-Project — skipping (derived-pushup not yet implemented)"); + return null; + } + belowProjOutToScan = outToScan; + } + chain.add(n); + if (n.getInputs().isEmpty()) return null; + n = RelNodeUtils.unwrapHep(n.getInput(0)); + } + if (scan == null) return null; + return new BelowChain(chain, belowProjOutToScan, scan, hasExchangeReducer); + } + + /** Returns output→scanIdx map iff every project is a {@link RexInputRef}; else null. */ + private static int[] passthroughMap(OpenSearchProject p) { + int[] out = new int[p.getProjects().size()]; + for (int i = 0; i < p.getProjects().size(); i++) { + if (!(p.getProjects().get(i) instanceof RexInputRef ref)) return null; + out[i] = ref.getIndex(); + } + return out; + } + + /** + * {@code BelowAnchorPhysicalFields} = anchor sort cols ∪ below-filter cols, expressed + * as physical (Scan-level) field names. The narrowed Scan's rowType is built from this + * set in scan-original order. + */ + private static Set computeBelowAnchorPhysicalFields(OpenSearchSort anchor, BelowChain belowChain) { + Set fields = new HashSet<>(); + List scanFields = belowChain.scan.getRowType().getFieldList(); + // Anchor sort cols (anchor space → scan space → physical name) + for (RelFieldCollation fc : anchor.getCollation().getFieldCollations()) { + int scanIdx = (belowChain.belowProjOutToScan == null) ? fc.getFieldIndex() : belowChain.belowProjOutToScan[fc.getFieldIndex()]; + fields.add(scanFields.get(scanIdx).getName()); + } + // Below-filter cols (already in scan space) + for (RelNode op : belowChain.chain) { + if (op instanceof OpenSearchFilter f) { + for (int refIdx : RelNodeUtils.collectInputRefs(f.getCondition())) { + fields.add(scanFields.get(refIdx).getName()); + } + } + } + return fields; + } + + /** + * {@code AboveAnchorPhysicalFields} = the union of physical-field deps across the + * topmost-above-anchor operator's output FSIs, in first-appearance order. Order is + * the wrapper-output and fetch-RPC order. + * + *

When {@code aboveAnchorOperators} is empty the anchor itself is the topmost op — + * its FSIs walk up from the below-chain via inheritance. + */ + private static LinkedHashSet computeAboveAnchorPhysicalFields(List aboveAnchorOperators, OpenSearchSort anchor) { + OpenSearchRelNode topmost = aboveAnchorOperators.isEmpty() ? anchor : (OpenSearchRelNode) aboveAnchorOperators.getFirst(); + LinkedHashSet fields = new LinkedHashSet<>(); + for (FieldStorageInfo fsi : topmost.getOutputFieldStorage()) { + if (!fsi.isDerived()) { + fields.add(fsi.getFieldName()); + } else { + fields.addAll(fsi.getDependsOnPhysicalCols()); + } + } + return fields; + } + + /** + * Builds {@code anchorSlotToPhysicalField}: for each slot in {@code anchor.rowType}, + * the physical field name it ultimately reads from. Used during rewrite for above-op + * RexInputRef remapping (anchor.slot → physicalName → wrapperOut.idx). + */ + private static List buildAnchorSlotToPhysicalField(BelowChain belowChain) { + // anchor.rowType inherits from belowChain.chain[0]'s rowType (which inherits ER → Filter → Scan). + RelNode topBelowOp = belowChain.chain.isEmpty() ? belowChain.scan : belowChain.chain.get(0); + int slotCount = topBelowOp.getRowType().getFieldCount(); + List scanFields = belowChain.scan.getRowType().getFieldList(); + List out = new ArrayList<>(slotCount); + for (int slot = 0; slot < slotCount; slot++) { + int scanIdx = (belowChain.belowProjOutToScan == null) ? slot : belowChain.belowProjOutToScan[slot]; + out.add(scanFields.get(scanIdx).getName()); + } + return out; + } + + // ── Phase 2 — Rewrite ────────────────────────────────────────────── + + /** + * Pure transform. Builds the narrowed Scan, walks the below chain (declaring + * {@code ___ugsi} on the ER), rebuilds the anchor's collation, instantiates the wrapper, + * and remaps RexNodes in the above chain. + */ + private static RelNode applyRewrite(RelNode root, Detection detection) { + OpenSearchTableScan origScan = detection.belowChain.scan; + + // 2a. Narrowed Scan: reduce-set in scan-original order + ___row_id. + NarrowedScan narrowed = buildNarrowedScan(origScan, detection.belowAnchorPhysicalFields); + + // 2b. Walk below chain bottom-up; rebuild operators with narrowed input. ER appends ___ugsi. + BelowRebuild belowRebuild = rebuildBelowChain(detection.belowChain, narrowed.newScan, narrowed.scanIdxRemap); + + // 2c. Anchor Sort: collation remapped from anchor space to narrowed-Scan space. + OpenSearchSort newAnchor = rebuildAnchor(detection.anchor, belowRebuild.rebuilt, belowRebuild.anchorSlotRemap); + + // 2d. Wrapper. Output rowType = aboveAnchorPhysicalFields in iteration order. + OpenSearchLateMaterialization wrapper = buildWrapper(newAnchor, detection.aboveAnchorPhysicalFields, origScan, detection.anchor); + + // 2e. Above chain: every op's RexInputRefs remapped by column name from its origChild's + // rowType to its newChild's rowType. Pass-through ops (Filter, Sort) leak the narrowed + // rowType upward, so a single immediate-parent remap is insufficient — every above op + // needs the same treatment, recursively. + return rebuildAboveChain(RelNodeUtils.unwrapHep(root), detection.anchor, wrapper); + } + + // ── 2a. Narrowed Scan ───────────────────────────────────────────── + + private static NarrowedScan buildNarrowedScan(OpenSearchTableScan origScan, Set belowAnchorPhysicalFields) { + RelDataTypeFactory typeFactory = origScan.getCluster().getTypeFactory(); + List origFields = origScan.getRowType().getFieldList(); + List origStorage = origScan.getOutputFieldStorage(); + + RelDataTypeFactory.Builder rowTypeBuilder = typeFactory.builder(); + List newStorage = new ArrayList<>(belowAnchorPhysicalFields.size() + 1); + int[] scanIdxRemap = new int[origFields.size()]; + Arrays.fill(scanIdxRemap, -1); + + int nextNewIdx = 0; + for (int origIdx = 0; origIdx < origFields.size(); origIdx++) { + RelDataTypeField field = origFields.get(origIdx); + if (belowAnchorPhysicalFields.contains(field.getName())) { + scanIdxRemap[origIdx] = nextNewIdx++; + rowTypeBuilder.add(field.getName(), field.getType()); + newStorage.add(origStorage.get(origIdx)); + } + } + rowTypeBuilder.add(OpenSearchLateMaterialization.ROW_ID_FIELD, typeFactory.createSqlType(SqlTypeName.BIGINT)); + newStorage.add(FieldStorageInfo.derivedColumn(OpenSearchLateMaterialization.ROW_ID_FIELD, SqlTypeName.BIGINT)); + + OpenSearchTableScan newScan = new OpenSearchTableScan( + origScan.getCluster(), + origScan.getTraitSet(), + origScan.getTable(), + origScan.getViableBackends(), + newStorage, + rowTypeBuilder.build() + ); + return new NarrowedScan(newScan, scanIdxRemap); + } + + // ── 2b. Below chain rebuild ─────────────────────────────────────── + + private static BelowRebuild rebuildBelowChain(BelowChain belowChain, RelNode newScan, int[] scanIdxRemap) { + RelDataTypeFactory typeFactory = newScan.getCluster().getTypeFactory(); + RelNode rebuilt = newScan; + int[] anchorSlotRemap = scanIdxRemap; + + for (int i = belowChain.chain.size() - 1; i >= 0; i--) { + RelNode orig = belowChain.chain.get(i); + switch (orig) { + case OpenSearchFilter f -> { + RexNode remapped = f.getCondition().accept(new IndexRemapShuttle(scanIdxRemap, rebuilt.getRowType())); + rebuilt = new OpenSearchFilter(f.getCluster(), f.getTraitSet(), rebuilt, remapped, f.getViableBackends()); + } + case OpenSearchProject p -> { + BelowProjectRebuild rebuild = rebuildBelowProject(p, rebuilt, scanIdxRemap, typeFactory); + rebuilt = rebuild.rebuilt; + anchorSlotRemap = rebuild.outputRemap; + } + case OpenSearchExchangeReducer er -> { + // Invariant 4: declare ___ugsi on the ER's output rowType (materialized at runtime + // by OrdinalAppendingSink before DataFusion's reduce sees the batch). + RelDataType erRowType = RelNodeUtils.appendField( + typeFactory, + rebuilt.getRowType(), + OpenSearchLateMaterialization.UGSI_FIELD, + typeFactory.createSqlType(SqlTypeName.INTEGER) + ); + rebuilt = new OpenSearchExchangeReducer( + er.getCluster(), + er.getTraitSet(), + rebuilt, + er.getViableBackends(), + er.getExchangeInfo(), + erRowType + ); + } + default -> throw new IllegalStateException("Unexpected below-anchor operator: " + orig.getClass().getSimpleName()); + } + } + return new BelowRebuild(rebuilt, anchorSlotRemap); + } + + private static BelowProjectRebuild rebuildBelowProject( + OpenSearchProject p, + RelNode newChild, + int[] scanIdxRemap, + RelDataTypeFactory typeFactory + ) { + // Below-Project is passthrough; its output→scan map was captured during analyzeBelow, + // but we recompute here from p's projects since each project is a RexInputRef. + int[] origOutToScan = new int[p.getProjects().size()]; + for (int i = 0; i < p.getProjects().size(); i++) { + origOutToScan[i] = ((RexInputRef) p.getProjects().get(i)).getIndex(); + } + + List newProjects = new ArrayList<>(); + List newNames = new ArrayList<>(); + int[] outputRemap = new int[origOutToScan.length]; + Arrays.fill(outputRemap, -1); + for (int origOut = 0; origOut < origOutToScan.length; origOut++) { + int newScanIdx = scanIdxRemap[origOutToScan[origOut]]; + if (newScanIdx < 0) continue; // source dropped (now fetched, not in narrowed Scan) + RelDataTypeField field = newChild.getRowType().getFieldList().get(newScanIdx); + outputRemap[origOut] = newProjects.size(); + newProjects.add(new RexInputRef(newScanIdx, field.getType())); + newNames.add(p.getRowType().getFieldList().get(origOut).getName()); + } + // Pass through ___row_id (always last in narrowed Scan rowType). + int rowIdIdx = newChild.getRowType().getFieldCount() - 1; + RelDataTypeField rowIdField = newChild.getRowType().getFieldList().get(rowIdIdx); + newProjects.add(new RexInputRef(rowIdIdx, rowIdField.getType())); + newNames.add(OpenSearchLateMaterialization.ROW_ID_FIELD); + + RelDataTypeFactory.Builder pb = typeFactory.builder(); + for (int j = 0; j < newProjects.size(); j++) { + pb.add(newNames.get(j), newProjects.get(j).getType()); + } + OpenSearchProject rebuilt = new OpenSearchProject( + p.getCluster(), + p.getTraitSet(), + newChild, + newProjects, + pb.build(), + p.getViableBackends() + ); + return new BelowProjectRebuild(rebuilt, outputRemap); + } + + // ── 2c. Anchor rebuild ──────────────────────────────────────────── + + private static OpenSearchSort rebuildAnchor(OpenSearchSort anchor, RelNode newInput, int[] anchorSlotRemap) { + List newCollations = new ArrayList<>(anchor.getCollation().getFieldCollations().size()); + for (RelFieldCollation fc : anchor.getCollation().getFieldCollations()) { + int newIdx = anchorSlotRemap[fc.getFieldIndex()]; + if (newIdx < 0) { + throw new IllegalStateException( + "Sort collation references slot " + fc.getFieldIndex() + " whose physical field was dropped from the narrowed Scan" + ); + } + newCollations.add(fc.withFieldIndex(newIdx)); + } + return new OpenSearchSort( + anchor.getCluster(), + anchor.getTraitSet(), + newInput, + RelCollations.of(newCollations), + anchor.offset, + anchor.fetch, + anchor.getViableBackends() + ); + } + + // ── 2d. Wrapper ─────────────────────────────────────────────────── + + /** + * Builds the {@link OpenSearchLateMaterialization} wrapper. Output rowType has one + * field per element of {@code aboveAnchorPhysicalFields} (in iteration order), each + * named with the physical field name and typed from the original Scan. + */ + private static OpenSearchLateMaterialization buildWrapper( + OpenSearchSort newAnchor, + LinkedHashSet aboveAnchorPhysicalFields, + OpenSearchTableScan origScan, + OpenSearchSort origAnchor + ) { + Map origScanIdxByName = new HashMap<>(); + List origFields = origScan.getRowType().getFieldList(); + for (int i = 0; i < origFields.size(); i++) { + origScanIdxByName.put(origFields.get(i).getName(), i); + } + List origStorage = origScan.getOutputFieldStorage(); + + List wrapperFields = new ArrayList<>(aboveAnchorPhysicalFields.size()); + List wrapperStorage = new ArrayList<>(aboveAnchorPhysicalFields.size()); + for (String name : aboveAnchorPhysicalFields) { + Integer scanIdx = origScanIdxByName.get(name); + if (scanIdx == null) { + throw new IllegalStateException( + "aboveAnchorPhysicalFields references [" + name + "] which is not present in the original Scan rowType" + ); + } + wrapperFields.add(origFields.get(scanIdx)); + wrapperStorage.add(origStorage.get(scanIdx)); + } + + return new OpenSearchLateMaterialization( + newAnchor.getCluster(), + newAnchor.getTraitSet(), + newAnchor, + wrapperFields, + wrapperStorage, + origAnchor.getViableBackends() + ); + } + + // ── 2e. Above chain rebuild ─────────────────────────────────────── + + /** + * Walks down the above chain, swapping the anchor's slot with {@code wrapper}, then on + * the way up rewrites every op's RexInputRefs via a by-name remap from {@code origChild}'s + * rowType to {@code newChild}'s rowType. Names are the stable identity that survives + * narrowing — they exist in both rowTypes verbatim for kept columns, and resolve to -1 + * (rejected by {@link IndexRemapShuttle}) for dropped ones. + */ + private static RelNode rebuildAboveChain(RelNode current, OpenSearchSort origAnchor, OpenSearchLateMaterialization wrapper) { + if (current == origAnchor) return wrapper; + if (current.getInputs().size() != 1) { + throw new IllegalStateException("Multi-input parent in QTF chain: " + current.getClass().getSimpleName()); + } + RelNode origChild = RelNodeUtils.unwrapHep(current.getInput(0)); + RelNode newChild = rebuildAboveChain(origChild, origAnchor, wrapper); + + int[] remap = buildByNameRemap(origChild.getRowType(), newChild.getRowType()); + IndexRemapShuttle shuttle = new IndexRemapShuttle(remap, newChild.getRowType()); + + switch (current) { + case OpenSearchProject project -> { + List newExprs = new ArrayList<>(project.getProjects().size()); + for (RexNode expr : project.getProjects()) { + newExprs.add(expr.accept(shuttle)); + } + return project.copy(project.getTraitSet(), newChild, newExprs, project.getRowType()); + } + case OpenSearchFilter filter -> { + return filter.copy(filter.getTraitSet(), newChild, filter.getCondition().accept(shuttle)); + } + case OpenSearchSort sort -> { + List remapped = new ArrayList<>(sort.getCollation().getFieldCollations().size()); + for (RelFieldCollation fc : sort.getCollation().getFieldCollations()) { + int newIdx = remap[fc.getFieldIndex()]; + if (newIdx < 0) { + throw new IllegalStateException( + "Above-anchor Sort references column at slot " + fc.getFieldIndex() + " not present in narrowed rowType" + ); + } + remapped.add(fc.withFieldIndex(newIdx)); + } + return sort.copy(sort.getTraitSet(), newChild, RelCollations.of(remapped), sort.offset, sort.fetch); + } + default -> throw new IllegalStateException("Unexpected above-anchor operator: " + current.getClass().getSimpleName()); + } + } + + /** Maps each index of {@code origType} to the index of the same-named field in {@code newType}, or -1 if dropped. */ + private static int[] buildByNameRemap(RelDataType origType, RelDataType newType) { + Map newIdxByName = new HashMap<>(newType.getFieldCount()); + List newFields = newType.getFieldList(); + for (int i = 0; i < newFields.size(); i++) { + newIdxByName.put(newFields.get(i).getName(), i); + } + List origFields = origType.getFieldList(); + int[] remap = new int[origFields.size()]; + for (int i = 0; i < origFields.size(); i++) { + Integer mapped = newIdxByName.get(origFields.get(i).getName()); + remap[i] = mapped == null ? -1 : mapped; + } + return remap; + } + + // ── Records ──────────────────────────────────────────────────────── + + /** + * Output of phase 1. Contains everything phase 2 needs — phase 2 must not re-walk + * the original plan to recover any of these fields. + */ + private record Detection(OpenSearchSort anchor, List aboveAnchorOperators, BelowChain belowChain, Set< + String> belowAnchorPhysicalFields, LinkedHashSet aboveAnchorPhysicalFields, List anchorSlotToPhysicalField) { + } + + private record AnchorContext(OpenSearchSort anchor, List aboveAnchorOperators) { + } + + private record BelowChain(List chain, int[] belowProjOutToScan, OpenSearchTableScan scan, boolean hasExchangeReducer) { + } + + private record NarrowedScan(OpenSearchTableScan newScan, int[] scanIdxRemap) { + } + + private record BelowRebuild(RelNode rebuilt, int[] anchorSlotRemap) { + } + + private record BelowProjectRebuild(OpenSearchProject rebuilt, int[] outputRemap) { + } +} 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..c9cff0ef76a40 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/QueryThenFetchDataNodeTests.java @@ -0,0 +1,226 @@ +/* + * 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.core.index.Index; +import org.opensearch.core.index.shard.ShardId; +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 static final ShardId SHARD_0 = new ShardId(new Index("idx", "uuid"), 0); + + 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", SHARD_0, gatedReader); + assertNotNull(ctx); + assertSame(originalReader, ctx.getReader()); + + // Query phase completes, release context (reader stays alive for fetch) + store.releaseContext("query-1", SHARD_0); + assertFalse("Reader must not be closed between phases", closed.get()); + + // Simulate fetch phase: acquire the same context + ReaderContext fetchCtx = store.acquireContext("query-1", SHARD_0); + 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", SHARD_0); + 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", SHARD_0, mockGatedReader(closed)); + store.releaseContext("query-expired", SHARD_0); + + // 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", SHARD_0); + 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", SHARD_0); + 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", SHARD_0, mockGatedReader(closed)); + // Query phase done + store.releaseContext("query-2", SHARD_0); + + // 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", SHARD_0); + } + + /** + * 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", SHARD_0, mockGatedReader(closed)); + store.releaseContext("query-3", SHARD_0); + + ReaderContext fetchCtx = store.acquireContext("query-3", SHARD_0); + assertNotNull(fetchCtx); + assertFalse("Reader must not be closed during fetch", closed.get()); + + // Simulate completeFetch + fetchCtx.markDone(); + store.freeContext("query-3", SHARD_0); + + assertTrue("Reader must be closed after freeContext", closed.get()); + assertEquals(0, store.activeCount()); + assertNull("Context must be removed from store", store.getContext("query-3", SHARD_0)); + } + + /** + * 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", SHARD_0, gated1); + store.createContext("q2", SHARD_0, gated2); + assertEquals(2, store.activeCount()); + + // Release both query phases + store.releaseContext("q1", SHARD_0); + store.releaseContext("q2", SHARD_0); + + // Fetch q1 + ReaderContext fetch1 = store.acquireContext("q1", SHARD_0); + 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", SHARD_0); + assertTrue("q1 reader closed", closed1.get()); + assertFalse("q2 reader still alive", closed2.get()); + assertEquals(1, store.activeCount()); + + // Fetch q2 + ReaderContext fetch2 = store.acquireContext("q2", SHARD_0); + assertNotNull(fetch2); + assertSame(reader2, fetch2.getReader()); + fetch2.markDone(); + store.freeContext("q2", SHARD_0); + 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", SHARD_0, mockGatedReader(closed)); + + // Simulate "query failure": release the context without explicit cleanup + store.releaseContext("query-failed", SHARD_0); + + // 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", SHARD_0); + 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..46fde61a46545 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ReaderContextStoreTests.java @@ -0,0 +1,180 @@ +/* + * 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.core.index.Index; +import org.opensearch.core.index.shard.ShardId; +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 static final ShardId SHARD_0 = new ShardId(new Index("idx", "uuid"), 0); + private static final ShardId SHARD_1 = new ShardId(new Index("idx", "uuid"), 1); + + 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", SHARD_0, gated); + + assertNotNull(ctx); + assertEquals(1, store.activeCount()); + + // Release from query phase + store.releaseContext("q1", SHARD_0); + + // Acquire for fetch phase + ReaderContext fetched = store.acquireContext("q1", SHARD_0); + 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", SHARD_0)); + } + + public void testAcquireWhileInUseReturnsNull() { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + store.createContext("q1", SHARD_0, mockGatedReader(closed)); + // Context is already in-use from createContext + + assertNull("Should not acquire while in-use", store.acquireContext("q1", SHARD_0)); + } + + public void testFreeContextClosesReader() { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + store.createContext("q1", SHARD_0, mockGatedReader(closed)); + store.releaseContext("q1", SHARD_0); + store.freeContext("q1", SHARD_0); + + 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", SHARD_0, mockGatedReader(closed)); + store.releaseContext("q1", SHARD_0); + store.freeContext("q1", SHARD_0); + + assertNull("Should not find after free", store.getContext("q1", SHARD_0)); + } + + public void testReaperCleansExpiredContexts() throws Exception { + AtomicBoolean closed = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool, 50); // 50ms keepAlive + + store.createContext("q1", SHARD_0, mockGatedReader(closed)); + store.releaseContext("q1", SHARD_0); + + // 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", SHARD_0, 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", SHARD_0); + store.freeContext("q1", SHARD_0); + } + + public void testMultipleContexts() { + AtomicBoolean closed1 = new AtomicBoolean(false); + AtomicBoolean closed2 = new AtomicBoolean(false); + ReaderContextStore store = new ReaderContextStore(threadPool); + + store.createContext("q1", SHARD_0, mockGatedReader(closed1)); + store.createContext("q2", SHARD_1, mockGatedReader(closed2)); + assertEquals(2, store.activeCount()); + + store.releaseContext("q1", SHARD_0); + store.freeContext("q1", SHARD_0); + assertEquals(1, store.activeCount()); + assertTrue(closed1.get()); + assertFalse(closed2.get()); + + store.releaseContext("q2", SHARD_1); + store.freeContext("q2", SHARD_1); + 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", SHARD_0, mockGatedReader(closed)); + ReaderContext ctx = store.getContext("q1", SHARD_0); + assertNotNull(ctx); + assertNotNull(ctx.getReader()); + + // Query done + store.releaseContext("q1", SHARD_0); + assertFalse("Reader stays open between phases", closed.get()); + + // Fetch phase: acquire + use + ReaderContext fetchCtx = store.acquireContext("q1", SHARD_0); + assertNotNull(fetchCtx); + assertSame(ctx.getReader(), fetchCtx.getReader()); + + // Fetch done: free + fetchCtx.markDone(); + store.freeContext("q1", SHARD_0); + 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..513c2ac86dd45 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/ReaderContextTests.java @@ -0,0 +1,121 @@ +/* + * 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.core.index.Index; +import org.opensearch.core.index.shard.ShardId; +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 static final ShardId SHARD_0 = new ShardId(new Index("idx", "uuid"), 0); + + 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", SHARD_0, 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", SHARD_0, 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", SHARD_0, 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", SHARD_0, 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", SHARD_0, 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", SHARD_0, 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", SHARD_0, gated, 30_000); + + assertSame(reader, ctx.getReader()); + } + + public void testGetQueryId() { + GatedCloseable gated = mockGatedReader(); + ReaderContext ctx = new ReaderContext("test-query-123", SHARD_0, gated, 30_000); + + assertEquals("test-query-123", ctx.getQueryId()); + } + + public void testLastAccessTimeUpdatedOnMarkInUse() throws Exception { + GatedCloseable gated = mockGatedReader(); + ReaderContext ctx = new ReaderContext("q1", SHARD_0, 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/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecutionTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecutionTests.java new file mode 100644 index 0000000000000..f0a15d75e8305 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/LateMaterializationStageExecutionTests.java @@ -0,0 +1,181 @@ +/* + * 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.stage; + +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelTraitSet; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.hep.HepProgramBuilder; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.Values; +import org.apache.calcite.rel.logical.LogicalValues; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.sql.type.SqlTypeName; +import org.opensearch.analytics.exec.AnalyticsSearchTransportService; +import org.opensearch.analytics.exec.QueryContext; +import org.opensearch.analytics.exec.task.AnalyticsQueryTask; +import org.opensearch.analytics.planner.dag.Stage; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; +import org.opensearch.analytics.spi.ExchangeSink; +import org.opensearch.analytics.spi.FieldStorageInfo; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.List; +import java.util.concurrent.ExecutorService; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link LateMaterializationStageExecution}. + * + *

Focused on the K=0 short-circuit path: when the child Sort+Limit reduce produces zero + * rows, the LM stage must NOT close its parent's input sink. Closing it pre-emptively races + * the parent {@code ReduceStageExecution}'s reduce() task — once the JVM thread pools warm + * up, the reduce task runs after the close and the backend sink throws + * {@code IllegalStateException("sink closed before reduce")}, which the IT surfaces as + * HTTP 500 / {@code NoSuchElementException} on the wire. + */ +public class LateMaterializationStageExecutionTests extends OpenSearchTestCase { + + private RelOptCluster cluster; + + @Override + public void setUp() throws Exception { + super.setUp(); + RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + HepPlanner planner = new HepPlanner(new HepProgramBuilder().build()); + cluster = RelOptCluster.create(planner, rexBuilder); + } + + /** + * Empty-K (zero rows from child) must NOT close {@code parentSink}. The parent stage + * (e.g. {@code ReduceStageExecution}) owns the sink's lifecycle via its own + * {@code onTerminalTransition}; closing it from the LM stage races the parent's + * {@code reduce()} task and produces "sink closed before reduce" on warm runtimes. + */ + public void testKZero_doesNotCloseParentSink() { + CapturingSink parentSink = new CapturingSink(); + LateMaterializationStageExecution exec = newLmStage(parentSink); + + scheduleAndDispatch(exec); + + assertEquals(StageExecution.State.SUCCEEDED, exec.getState()); + assertFalse("LM stage must not close parentSink on K=0; parent stage owns close via onTerminalTransition", parentSink.closed); + } + + // ── Fixture builders ───────────────────────────────────────────────────── + + /** + * Construct a Stage whose fragment is exactly an OpenSearchLateMaterialization wrapping + * an empty LogicalValues source. Simulates the LM stage's own fragment with no child + * input batches yet supplied — drainAndGroupByUgsi will see total=0 and the K=0 path + * fires. + */ + private LateMaterializationStageExecution newLmStage(ExchangeSink parentSink) { + RelDataType rowType = cluster.getTypeFactory() + .builder() + .add(OpenSearchLateMaterialization.ROW_ID_FIELD, cluster.getTypeFactory().createSqlType(SqlTypeName.BIGINT)) + .add(OpenSearchLateMaterialization.UGSI_FIELD, cluster.getTypeFactory().createSqlType(SqlTypeName.INTEGER)) + .build(); + Values emptyInput = (Values) LogicalValues.createEmpty(cluster, rowType); + + // Wrapper above-anchor fields = a single string column. The K=0 branch never + // dereferences them, so a minimal list is enough. + List aboveFields = List.of(rowType.getFieldList().get(0)); + List aboveStorage = List.of( + FieldStorageInfo.derivedColumn(OpenSearchLateMaterialization.ROW_ID_FIELD, SqlTypeName.BIGINT) + ); + + OpenSearchLateMaterialization wrapper = new OpenSearchLateMaterialization( + cluster, + RelTraitSet.createEmpty(), + emptyInput, + aboveFields, + aboveStorage, + List.of("datafusion") + ); + + Stage stage = mock(Stage.class); + when(stage.getStageId()).thenReturn(2); + when(stage.getFragment()).thenReturn((RelNode) wrapper); + + QueryContext config = mock(QueryContext.class); + when(config.queryId()).thenReturn("test-query"); + when(config.operationListeners()).thenReturn(List.of()); + when(config.parentTask()).thenReturn(mock(AnalyticsQueryTask.class)); + when(config.localTaskExecutor()).thenReturn(inlineExecutor()); + + return new LateMaterializationStageExecution( + stage, + config, + parentSink, + mock(ClusterService.class), + mock(AnalyticsSearchTransportService.class), + /* shardStageId */ 0, + /* fetchBackendId */ "datafusion" + ); + } + + private static ExecutorService inlineExecutor() { + return mock(ExecutorService.class, invocation -> { + if ("execute".equals(invocation.getMethod().getName())) { + ((Runnable) invocation.getArgument(0)).run(); + return null; + } + return null; + }); + } + + private static void scheduleAndDispatch(LateMaterializationStageExecution exec) { + exec.start(); + @SuppressWarnings("unchecked") + org.opensearch.analytics.exec.task.TaskRunner dispatcher = (org.opensearch.analytics.exec.task.TaskRunner< + StageTask>) exec.taskRunner(); + if (dispatcher == null) return; + for (StageTask task : exec.tasks()) { + task.transitionTo(StageTaskState.RUNNING); + dispatcher.run(task, new org.opensearch.core.action.ActionListener<>() { + @Override + public void onResponse(Void unused) { + task.transitionTo(StageTaskState.FINISHED); + exec.onTaskTerminal(task, null); + } + + @Override + public void onFailure(Exception cause) { + task.transitionTo(StageTaskState.FAILED); + exec.onTaskTerminal(task, cause); + } + }); + } + } + + /** Bare ExchangeSink that records whether it has been closed. */ + private static final class CapturingSink implements ExchangeSink { + boolean closed = false; + + @Override + public void feed(VectorSchemaRoot batch) { + batch.close(); + } + + @Override + public void close() { + closed = true; + } + } +} diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java index c50dc8184cb05..100b5533b4906 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardFragmentStageExecutionTests.java @@ -224,7 +224,7 @@ private Stage mockStageWithTargets(int n) { for (int i = 0; i < n; i++) { DiscoveryNode node = mock(DiscoveryNode.class); when(node.getId()).thenReturn("test-node-" + i); - targets.add(new ShardExecutionTarget(node, new ShardId("idx", "_na_", i))); + targets.add(new ShardExecutionTarget(node, new ShardId("idx", "_na_", i), i)); } when(resolver.resolve(any(ClusterState.class), any())).thenReturn(targets); when(stage.getTargetResolver()).thenReturn(resolver); diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardTaskRunnerTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardTaskRunnerTests.java index 266edcfd1d357..fd4ac6c82d5ed 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardTaskRunnerTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/exec/stage/ShardTaskRunnerTests.java @@ -101,7 +101,7 @@ private ShardTaskRunner newRunner(List capturedQueues) { private static ShardStageTask shardTask(int partitionId, String nodeId) { DiscoveryNode node = mock(DiscoveryNode.class); when(node.getId()).thenReturn(nodeId); - ShardExecutionTarget target = new ShardExecutionTarget(node, new ShardId("idx", "_na_", partitionId)); + ShardExecutionTarget target = new ShardExecutionTarget(node, new ShardId("idx", "_na_", partitionId), partitionId); return new ShardStageTask(new StageTaskId(0, partitionId), target); } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/BasePlannerRulesTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/BasePlannerRulesTests.java index 43305f8c8eaa0..741147e278caf 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/BasePlannerRulesTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/BasePlannerRulesTests.java @@ -32,6 +32,7 @@ import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.util.ImmutableBitSet; import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; import org.opensearch.analytics.spi.AggregateCapability; import org.opensearch.analytics.spi.AggregateFunction; @@ -319,7 +320,7 @@ protected static void assertPipelineViableBackends( private static RelNode skipExchangeReducers(RelNode rel) { RelNode current = rel; - while (current instanceof OpenSearchExchangeReducer) { + while (current instanceof OpenSearchExchangeReducer || current instanceof OpenSearchLateMaterialization) { current = RelNodeUtils.unwrapHep(current.getInputs().get(0)); } return current; diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ClickBench.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ClickBench.java index 378a98f6215b9..710f3bc737fa4 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ClickBench.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/ClickBench.java @@ -8,6 +8,8 @@ package org.opensearch.analytics.planner; +import java.util.Collections; +import java.util.LinkedHashMap; import java.util.Map; /** @@ -16,27 +18,29 @@ *

A subset of the ClickBench {@code hits} index covering the basic SQL types * exercised by pushdown / filter-delegation tests: integer, long, short, keyword, date. * Add fields when a test needs them. + * + *

Backed by {@link LinkedHashMap} so column order is deterministic across JVM + * launches. {@code Map.of}'s iteration order is JVM-seed-salted, which surfaces as + * flaky tests when assertions depend on column ordering (e.g. plan-shape tests + * comparing scan rowType field order). */ public final class ClickBench { public static final String INDEX = "hits"; - public static final Map> BASIC_FIELDS = Map.of( - "CounterID", - Map.of("type", "integer"), - "UserID", - Map.of("type", "long"), - "URL", - Map.of("type", "keyword"), - "Title", - Map.of("type", "keyword"), - "EventDate", - Map.of("type", "date"), - "AdvEngineID", - Map.of("type", "short"), - "ParamPrice", - Map.of("type", "long") - ); + public static final Map> BASIC_FIELDS; + + static { + LinkedHashMap> fields = new LinkedHashMap<>(); + fields.put("CounterID", Map.of("type", "integer")); + fields.put("UserID", Map.of("type", "long")); + fields.put("URL", Map.of("type", "keyword")); + fields.put("Title", Map.of("type", "keyword")); + fields.put("EventDate", Map.of("type", "date")); + fields.put("AdvEngineID", Map.of("type", "short")); + fields.put("ParamPrice", Map.of("type", "long")); + BASIC_FIELDS = Collections.unmodifiableMap(fields); + } private ClickBench() {} } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/LateMaterializationPlanShapeTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/LateMaterializationPlanShapeTests.java new file mode 100644 index 0000000000000..9b7c1a31e07fd --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/LateMaterializationPlanShapeTests.java @@ -0,0 +1,559 @@ +/* + * 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.planner; + +import org.apache.calcite.plan.RelOptUtil; +import org.apache.calcite.rel.RelFieldCollation; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.RelVisitor; +import org.apache.calcite.rel.type.RelDataTypeField; +import org.apache.calcite.rex.RexInputRef; +import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; +import org.opensearch.analytics.planner.rel.OpenSearchProject; +import org.opensearch.analytics.planner.rel.OpenSearchSort; +import org.opensearch.analytics.planner.rel.OpenSearchTableScan; +import org.opensearch.cluster.ClusterState; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +/** + * Plan-shape tests for the QTF (late-materialization) post-CBO rewriter, v2. + * + *

Each test drives a SQL string (Calcite dialect) through the full planner and asserts + * the structural post-conditions QTF must produce: + *

    + *
  • Scan rowType narrowed to {@code BelowAnchorPhysicalFields + ___row_id}.
    • + *
  • ER output declares {@code ___ugsi} as the last field. Single-shard plans do not trigger + * QTF (no gather to skip).
    • + *
  • Wrapper output rowType = {@code AboveAnchorPhysicalFields}, in + * {@code TopmostOperatorAboveAnchor} order. Helpers stripped.
    • + *
  • Outer Project's RexInputRefs are remapped to wrapper-output indices.
    • + *
+ * + *

Skip predicate: QTF fires iff + * {@code AboveAnchorPhysicalFields - BelowAnchorPhysicalFields} is non-empty. + */ +public class LateMaterializationPlanShapeTests extends BasePlannerRulesTests { + + // ── Tests that fire QTF ──────────────────────────────────────────── + + public void testQtfFires_simpleSortProject() { + // SELECT URL, EventDate FROM hits ORDER BY EventDate LIMIT 10 + // AboveAnchorPhysicalFields = [URL, EventDate] (topmost Project, in SELECT order) + // BelowAnchorPhysicalFields = {EventDate} + // FetchOnly = {URL} → fire + assertQtfFired( + "SELECT URL, EventDate FROM hits ORDER BY EventDate LIMIT 10", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "EventDate"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("URL", "EventDate"), + Expect.outerProjectExprIndices(0, 1) + ); + } + + public void testQtfFires_withWhere() { + // SELECT URL, EventDate FROM hits WHERE CounterID = 5 ORDER BY EventDate LIMIT 10 + // AboveAnchorPhysicalFields = [URL, EventDate] + // BelowAnchorPhysicalFields = {CounterID, EventDate} + // FetchOnly = {URL} → fire + assertQtfFired( + "SELECT URL, EventDate FROM hits WHERE CounterID = 5 ORDER BY EventDate LIMIT 10", + 2, + Expect.scanCols("CounterID", "EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "EventDate"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("URL", "EventDate"), + Expect.outerProjectExprIndices(0, 1) + ); + } + + public void testQtfFires_sortColAlsoProjected() { + // SELECT EventDate, URL FROM hits ORDER BY EventDate LIMIT 10 + // Wrapper output is in topmost-op (SELECT) order, NOT scan order. + // EventDate is in BOTH BelowAnchor and AboveAnchor — refetched per v2 (no passthrough trick). + assertQtfFired( + "SELECT EventDate, URL FROM hits ORDER BY EventDate LIMIT 10", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("EventDate", "URL"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("EventDate", "URL"), + Expect.outerProjectExprIndices(0, 1) + ); + } + + public void testQtfDeclined_singleShard() { + // Single shard: CBO inserts no ExchangeReducer below the anchor (the scan's + // SOURCE(SINGLETON) already satisfies the parent Sort's demand). QTF's win comes from + // avoiding cross-node materialization of fetch-only columns through the gather; with + // no gather there's nothing to save, so the rewriter declines. + assertQtfDeclined("SELECT URL, EventDate FROM hits ORDER BY EventDate LIMIT 10", 1); + } + + public void testQtfFires_descendingSort() { + // DESC collation preserved through anchor rebuild. + assertQtfFired( + "SELECT URL, EventDate FROM hits ORDER BY EventDate DESC LIMIT 10", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "EventDate"), + Expect.erHasUgsi(true), + Expect.collationDirections(RelFieldCollation.Direction.DESCENDING) + ); + } + + public void testQtfFires_multiKeySort() { + // SELECT URL, EventDate, CounterID FROM hits ORDER BY EventDate, CounterID LIMIT 10 + // BelowAnchor = {EventDate, CounterID}, AboveAnchor = [URL, EventDate, CounterID] + // FetchOnly = {URL} → fire. All three referenced fields end up in fetch list (v2 refetches). + assertQtfFired( + "SELECT URL, EventDate, CounterID FROM hits ORDER BY EventDate, CounterID LIMIT 10", + 2, + Expect.scanCols("CounterID", "EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "EventDate", "CounterID"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("URL", "EventDate", "CounterID") + ); + } + + public void testQtfFires_filterColAlsoProjected() { + // SELECT URL, CounterID FROM hits WHERE CounterID = 5 ORDER BY EventDate LIMIT 10 + // BelowAnchor = {CounterID, EventDate}, AboveAnchor = [URL, CounterID] + // FetchOnly = {URL} → fire. CounterID refetched. + assertQtfFired( + "SELECT URL, CounterID FROM hits WHERE CounterID = 5 ORDER BY EventDate LIMIT 10", + 2, + Expect.scanCols("CounterID", "EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "CounterID"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("URL", "CounterID") + ); + } + + public void testQtfFires_offset() { + // OFFSET preserved on the rebuilt anchor Sort. + assertQtfFired( + "SELECT URL, EventDate FROM hits ORDER BY EventDate LIMIT 10 OFFSET 5", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "EventDate"), + Expect.erHasUgsi(true) + ); + } + + public void testQtfFires_compositeExpressionWithDedup() { + // SELECT URL || '-' || URL AS combined, UPPER(URL) AS upper_url FROM hits ORDER BY EventDate LIMIT 10 + // Both above-Project outputs derive from URL. dependsOn walk + LinkedHashSet dedup + // collapses to AboveAnchorPhysicalFields = [URL]. + // BelowAnchor = {EventDate}; FetchOnly = {URL} → fire. + // Outer Project carries RexCall expressions, not bare InputRefs — outerProjectExprIndices skipped. + assertQtfFired( + "SELECT URL || '-' || URL AS combined, UPPER(URL) AS upper_url FROM hits ORDER BY EventDate LIMIT 10", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("URL"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("URL") + ); + } + + public void testQtfFires_compositeExpressionMultiCol() { + // SELECT URL || '-' || Title AS combined FROM hits ORDER BY EventDate LIMIT 10 + // dependsOn walk yields {URL, Title} in evaluation order. + // BelowAnchor = {EventDate}; FetchOnly = {URL, Title} → fire. + assertQtfFired( + "SELECT URL || '-' || Title AS combined FROM hits ORDER BY EventDate LIMIT 10", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "Title"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("URL", "Title") + ); + } + + public void testQtfFires_starProjection() { + // SELECT * FROM hits ORDER BY EventDate LIMIT 10 + // Topmost Project's outputs are passthrough refs to every physical field. + // AboveAnchor = [CounterID, UserID, URL, Title, EventDate, AdvEngineID, ParamPrice]. + // BelowAnchor = {EventDate}. FetchOnly = everything except EventDate → fire. + assertQtfFired( + "SELECT * FROM hits ORDER BY EventDate LIMIT 10", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("CounterID", "UserID", "URL", "Title", "EventDate", "AdvEngineID", "ParamPrice"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("CounterID", "UserID", "URL", "Title", "EventDate", "AdvEngineID", "ParamPrice") + ); + } + + public void testQtfFires_outerFilter() { + // SELECT URL, EventDate FROM hits WHERE EventDate > '2020-01-01' ORDER BY EventDate LIMIT 10 + // Filter on EventDate is below-Filter (predicate pushdown), so EventDate stays in BelowAnchor. + // AboveAnchor = [URL, EventDate]; FetchOnly = {URL} → fire. + assertQtfFired( + "SELECT URL, EventDate FROM hits WHERE EventDate > DATE '2020-01-01' ORDER BY EventDate LIMIT 10", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "EventDate"), + Expect.erHasUgsi(true), + Expect.wrapperOutput("URL", "EventDate") + ); + } + + /** + * Above-anchor chain has a Sort that doesn't reshape the rowType (passthrough on type, even + * if collation differs). The PPL/SQL frontends emit this when a SystemLimit Sort wraps the + * user's Sort+Limit. The rewriter must remap RexInputRefs in EVERY above op — not just the + * immediate parent — because each passthrough Sort/Filter leaks the narrowed rowType upward. + * + *

The implicit guarantee under test: the rewrite completes without Calcite's RexChecker + * throwing "RexInputRef out of range" — that's exactly what would fire if the remap stopped + * at the anchor's immediate parent and missed the outer Project sitting above the + * passthrough Sort. {@code assertQtfFired} catches such an exception on the planner call. + */ + public void testQtfFires_aboveChainHasPassthroughSort() { + assertQtfFired( + "SELECT URL, EventDate FROM (" + + " SELECT URL, EventDate FROM hits ORDER BY EventDate LIMIT 10" + + ") AS sub ORDER BY EventDate LIMIT 100", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "EventDate"), + Expect.wrapperOutput("URL", "EventDate") + ); + } + + /** + * Above-anchor Sort's collation index must be remapped from anchor-input space to + * wrapper-output space, even when there's no Project sitting between Sort and anchor. + */ + public void testQtfFires_aboveChainSortRemapsCollation() { + // Outer Sort by EventDate (= sort key, present in BelowAnchor); anchor Sort by EventDate. + // After narrowing, scan = [EventDate, ___row_id] = 2 cols; wrapper output = [URL, EventDate]. + // Outer Sort's collation index must remap from anchor-space to wrapper-space (where + // EventDate is at position 1). + assertQtfFired( + "SELECT * FROM (" + " SELECT URL, EventDate FROM hits ORDER BY EventDate LIMIT 10" + ") AS sub ORDER BY EventDate", + 2, + Expect.scanCols("EventDate"), + Expect.aboveAnchorPhysicalFields("URL", "EventDate"), + Expect.wrapperOutput("URL", "EventDate") + ); + } + + // ── Tests that decline QTF ───────────────────────────────────────── + + public void testQtfDeclined_pureLimit() { + // SELECT URL FROM hits LIMIT 10 + // No ORDER BY → no anchor → no rewrite. + assertQtfDeclined("SELECT URL FROM hits LIMIT 10", 2); + } + + public void testQtfDeclined_skipPredicate_sortColIsOnlyProjection() { + // SELECT EventDate FROM hits ORDER BY EventDate LIMIT 10 + // AboveAnchor = {EventDate}; BelowAnchor = {EventDate}; FetchOnly = {} → skip. + // Refetching EventDate by row id when it already rode through the reduce would be a wash. + assertQtfDeclined("SELECT EventDate FROM hits ORDER BY EventDate LIMIT 10", 2); + } + + public void testQtfDeclined_skipPredicate_filterColIsOnlyProjection() { + // SELECT CounterID FROM hits WHERE CounterID = 5 ORDER BY EventDate LIMIT 10 + // AboveAnchor = {CounterID}; BelowAnchor = {CounterID, EventDate}; FetchOnly = {} → skip. + assertQtfDeclined("SELECT CounterID FROM hits WHERE CounterID = 5 ORDER BY EventDate LIMIT 10", 2); + } + + public void testQtfDeclined_aggregateBelowSort() { + // SELECT CounterID, COUNT(*) FROM hits GROUP BY CounterID ORDER BY CounterID LIMIT 10 + // Aggregate sits below the anchor — not in BELOW_INTERMEDIATE_ALLOWED, declined. + assertQtfDeclined("SELECT CounterID, COUNT(*) AS c FROM hits GROUP BY CounterID ORDER BY CounterID LIMIT 10", 2); + } + + public void testQtfDeclined_aggregateAboveAnchor() { + // Aggregate above the anchor — explicitly excluded from ABOVE_ALLOWED in v2. + // (Above-Aggregate group/aggCall remap under a moving wrapper rowType is a follow-up.) + // SELECT inner.CounterID, COUNT(*) FROM (SELECT CounterID FROM hits ORDER BY CounterID LIMIT 100) AS inner GROUP BY inner.CounterID + assertQtfDeclined( + "SELECT inner_q.CounterID, COUNT(*) AS c " + + "FROM (SELECT CounterID FROM hits ORDER BY CounterID LIMIT 100) AS inner_q " + + "GROUP BY inner_q.CounterID", + 2 + ); + } + + public void testQtfDeclined_windowInOuterProject() { + // SELECT URL, SUM(ParamPrice) OVER () FROM hits ORDER BY EventDate LIMIT 10 + // RexOver in the above-Project — declined (window's global frame needs SINGLETON input; + // the wrapper would break that). + assertQtfDeclined("SELECT URL, SUM(ParamPrice) OVER () AS sp FROM hits ORDER BY EventDate LIMIT 10", 2); + } + + public void testQtfDeclined_expressionProjectBelowAnchor() { + // SELECT URL FROM hits ORDER BY (CounterID + 1) LIMIT 10 + // The sort key (CounterID + 1) gets materialized via a derived below-Project. + // Today's algorithm declines (TODO in passthroughMap — derived-pushup not yet supported). + assertQtfDeclined("SELECT URL FROM hits ORDER BY (CounterID + 1) LIMIT 10", 2); + } + + // ── Composable assert API ────────────────────────────────────────── + + private void assertQtfFired(String sql, int shardCount, Expect... expectations) { + RelNode optimized = optimize(sql, shardCount); + String planText = RelOptUtil.toString(optimized); + Inspector ctx = new Inspector(optimized); + if (ctx.wrapper == null) { + fail("Expected QTF wrapper in plan for SQL: " + sql + "\nPlan:\n" + planText); + } + for (Expect e : expectations) { + e.check(ctx, sql, planText); + } + } + + private void assertQtfDeclined(String sql, int shardCount) { + RelNode optimized = optimize(sql, shardCount); + Inspector ctx = new Inspector(optimized); + if (ctx.wrapper != null) { + fail("QTF should NOT have fired for SQL: " + sql + "\nPlan:\n" + RelOptUtil.toString(optimized)); + } + } + + private abstract static class Expect { + abstract void check(Inspector ctx, String sql, String planText); + + /** Scan rowType (post-narrowing) carries exactly these original cols + {@code ___row_id}. */ + static Expect scanCols(String... expectedNamesInOrder) { + return new Expect() { + @Override + void check(Inspector ctx, String sql, String plan) { + List actual = fieldNames(ctx.scan.getRowType().getFieldList()); + List expected = new ArrayList<>(Arrays.asList(expectedNamesInOrder)); + expected.add(OpenSearchLateMaterialization.ROW_ID_FIELD); + if (!expected.equals(actual)) { + fail( + "Scan rowType mismatch.\n expected: " + + expected + + "\n actual: " + + actual + + "\nSQL: " + + sql + + "\nPlan:\n" + + plan + ); + } + } + }; + } + + /** + * Wrapper carries exactly these field names as its {@code AboveAnchorPhysicalFields} + * (i.e. fetch list, in order). + */ + static Expect aboveAnchorPhysicalFields(String... expectedNamesInOrder) { + return new Expect() { + @Override + void check(Inspector ctx, String sql, String plan) { + List actual = fieldNames(ctx.wrapper.getAboveAnchorPhysicalFields()); + List expected = Arrays.asList(expectedNamesInOrder); + if (!expected.equals(actual)) { + fail( + "Wrapper aboveAnchorPhysicalFields mismatch.\n expected: " + + expected + + "\n actual: " + + actual + + "\nSQL: " + + sql + + "\nPlan:\n" + + plan + ); + } + } + }; + } + + /** + * Wrapper output rowType is exactly these names (= {@code AboveAnchorPhysicalFields} in + * topmost-op order). Helpers must be absent. + */ + static Expect wrapperOutput(String... expectedNamesInOrder) { + return new Expect() { + @Override + void check(Inspector ctx, String sql, String plan) { + List actual = fieldNames(ctx.wrapper.getRowType().getFieldList()); + List expected = Arrays.asList(expectedNamesInOrder); + if (!expected.equals(actual)) { + fail( + "Wrapper output rowType mismatch.\n expected: " + + expected + + "\n actual: " + + actual + + "\nSQL: " + + sql + + "\nPlan:\n" + + plan + ); + } + if (actual.contains(OpenSearchLateMaterialization.ROW_ID_FIELD) + || actual.contains(OpenSearchLateMaterialization.UGSI_FIELD)) { + fail("Wrapper output leaked helper col(s) " + actual + "\nSQL: " + sql + "\nPlan:\n" + plan); + } + } + }; + } + + /** ER carries (or does not carry) {@code ___ugsi} as its last column. */ + static Expect erHasUgsi(boolean expected) { + return new Expect() { + @Override + void check(Inspector ctx, String sql, String plan) { + if (expected) { + if (ctx.er == null) fail("Expected ER in plan but none found.\nSQL: " + sql + "\nPlan:\n" + plan); + List erFields = fieldNames(ctx.er.getRowType().getFieldList()); + if (!OpenSearchLateMaterialization.UGSI_FIELD.equals(erFields.get(erFields.size() - 1))) { + fail("ER output missing ___ugsi as last col. Fields: " + erFields + "\nSQL: " + sql + "\nPlan:\n" + plan); + } + } else { + if (ctx.er != null) fail("Did not expect ER in plan (single-shard).\nSQL: " + sql + "\nPlan:\n" + plan); + } + } + }; + } + + /** + * Outer Project (immediately above wrapper) RexInputRef indices, in expression order. + * Skip this assertion for projects whose expressions are not bare InputRefs (e.g. + * derived expressions like {@code UPPER(URL)}). + */ + static Expect outerProjectExprIndices(int... expectedIndices) { + return new Expect() { + @Override + void check(Inspector ctx, String sql, String plan) { + if (ctx.outerProject == null) { + fail("No outer Project above wrapper.\nSQL: " + sql + "\nPlan:\n" + plan); + return; + } + int[] actual = new int[ctx.outerProject.getProjects().size()]; + for (int i = 0; i < actual.length; i++) { + if (ctx.outerProject.getProjects().get(i) instanceof RexInputRef ref) { + actual[i] = ref.getIndex(); + } else { + fail( + "Outer Project expr [" + + i + + "] is not a RexInputRef: " + + ctx.outerProject.getProjects().get(i) + + "\nSQL: " + + sql + + "\nPlan:\n" + + plan + ); + return; + } + } + if (!Arrays.equals(expectedIndices, actual)) { + fail( + "Outer Project RexInputRef indices mismatch.\n expected: " + + Arrays.toString(expectedIndices) + + "\n actual: " + + Arrays.toString(actual) + + "\nSQL: " + + sql + + "\nPlan:\n" + + plan + ); + } + } + }; + } + + /** Anchor Sort collation directions in order. */ + static Expect collationDirections(RelFieldCollation.Direction... expected) { + return new Expect() { + @Override + void check(Inspector ctx, String sql, String plan) { + List fc = ctx.anchor.getCollation().getFieldCollations(); + if (fc.size() != expected.length) { + fail( + "Anchor collation size mismatch. expected=" + + expected.length + + " actual=" + + fc.size() + + "\nSQL: " + + sql + + "\nPlan:\n" + + plan + ); + } + for (int i = 0; i < expected.length; i++) { + if (fc.get(i).getDirection() != expected[i]) { + fail( + "Anchor collation[" + + i + + "] direction mismatch. expected=" + + expected[i] + + " actual=" + + fc.get(i).getDirection() + + "\nSQL: " + + sql + + "\nPlan:\n" + + plan + ); + } + } + } + }; + } + } + + /** Walks an optimized plan once, capturing the QTF-relevant nodes. */ + private static final class Inspector { + OpenSearchLateMaterialization wrapper; + OpenSearchSort anchor; + OpenSearchTableScan scan; + OpenSearchExchangeReducer er; + OpenSearchProject outerProject; + + Inspector(RelNode root) { + new RelVisitor() { + @Override + public void visit(RelNode node, int ordinal, RelNode parent) { + if (node instanceof OpenSearchLateMaterialization w && wrapper == null) { + wrapper = w; + if (parent instanceof OpenSearchProject p) outerProject = p; + if (w.getInput() instanceof OpenSearchSort s) anchor = s; + } + if (node instanceof OpenSearchExchangeReducer r && er == null) er = r; + if (node instanceof OpenSearchTableScan t && scan == null) scan = t; + super.visit(node, ordinal, parent); + } + }.go(root); + } + } + + private static List fieldNames(List fields) { + List out = new ArrayList<>(fields.size()); + for (RelDataTypeField f : fields) + out.add(f.getName()); + return out; + } + + private RelNode optimize(String sql, int shardCount) { + ClusterState state = SqlPlannerTestFixture.clusterStateWith(ClickBench.INDEX, ClickBench.BASIC_FIELDS, "parquet", shardCount); + PlannerContext context = new PlannerContext( + new CapabilityRegistry(List.of(DATAFUSION, LUCENE), FieldStorageResolver::new), + state, + false + ); + RelNode parsed = SqlPlannerTestFixture.parseSql(sql, state); + return PlannerImpl.runAllOptimizations(parsed, context); + } +} diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockBackend.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockBackend.java index d79ace5f85b78..6ea860b120d81 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockBackend.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockBackend.java @@ -157,8 +157,8 @@ public Map delegatedPredicateSeria public FragmentInstructionHandlerFactory getInstructionHandlerFactory() { return new FragmentInstructionHandlerFactory() { @Override - public Optional createShardScanNode() { - return Optional.of(new ShardScanInstructionNode()); + public Optional createShardScanNode(boolean requestsRowIds) { + return Optional.of(new ShardScanInstructionNode(requestsRowIds)); } @Override @@ -171,8 +171,12 @@ public Optional createFilterDelegationNode( } @Override - public Optional createShardScanWithDelegationNode(FilterTreeShape treeShape, int delegatedPredicateCount) { - return Optional.of(new ShardScanWithDelegationInstructionNode(treeShape, delegatedPredicateCount)); + public Optional createShardScanWithDelegationNode( + FilterTreeShape treeShape, + int delegatedPredicateCount, + boolean requestsRowIds + ) { + return Optional.of(new ShardScanWithDelegationInstructionNode(treeShape, delegatedPredicateCount, requestsRowIds)); } @Override diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java index 24a56f3a546bd..2b4ec2a66bbba 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockDataFusionBackend.java @@ -184,7 +184,10 @@ protected Set scanCapabilities() { // Logical connectives (projection-side composition: `case(a and b, …)`) ScalarFunction.AND, ScalarFunction.OR, - ScalarFunction.NOT + ScalarFunction.NOT, + // String — used by QTF plan-shape tests covering composite expressions / dedup. + ScalarFunction.CONCAT, + ScalarFunction.UPPER ); private static final Set PROJECT_CAPS; diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/PlanShapeTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/PlanShapeTests.java index 89e50c63ae346..5b3a6c50601f9 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/PlanShapeTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/PlanShapeTests.java @@ -55,15 +55,16 @@ public class PlanShapeTests extends PlanShapeTestBase { public void testSortHeadAfterStats_dropsRedundantOuterSort() { RelNode input = topKAfterStats(/* withRedundantOuterSort */ true); RelNode result = runPlanner(input, multiShardContext()); + // SORT_PROJECT_TRANSPOSE + PROJECT_MERGE in pushdown collapse the column-swap Projects + // around the Sort: the inner+outer Project pair merges into identity and drops, the Sort + // remaps its collation from $0 (post-swap) to $1 (Aggregate's cnt column). assertPlanShape( """ - OpenSearchProject(k=[$1], cnt=[$0], viableBackends=[[mock-parquet]]) - OpenSearchSort(sort0=[$0], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) - OpenSearchProject(cnt=[$1], k=[$0], viableBackends=[[mock-parquet]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[FINAL], viableBackends=[[mock-parquet]]) - OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) - OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[FINAL], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, result ); @@ -76,15 +77,15 @@ public void testSortHeadAfterStats_dropsRedundantOuterSort() { public void testSortHeadAfterStats_singleSortFetchPreserved() { RelNode input = topKAfterStats(/* withRedundantOuterSort */ false); RelNode result = runPlanner(input, multiShardContext()); + // Same Project-merge / Sort-transpose as above; the swap-Project pair collapses, + // Sort collation remaps from $0 to $1 (cnt) over the Aggregate output. assertPlanShape( """ - OpenSearchProject(k=[$1], cnt=[$0], viableBackends=[[mock-parquet]]) - OpenSearchSort(sort0=[$0], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) - OpenSearchProject(cnt=[$1], k=[$0], viableBackends=[[mock-parquet]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[FINAL], viableBackends=[[mock-parquet]]) - OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) - OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[FINAL], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, result ); @@ -99,16 +100,16 @@ public void testSortHeadAfterStats_outerSortWithDifferentKeyKept() { // Inner sort by cnt ($0 below swap), outer sort by k ($0 above swap which maps to k). RelNode input = topKAfterStats(/* withRedundantOuterSort */ true, /* outerSortField */ 0); RelNode result = runPlanner(input, multiShardContext()); + // Project-merge collapses the swap pair; outer Sort sees Aggregate output (k=$0, cnt=$1) + // directly. Outer sort on k stays $0; inner sort on cnt remaps to $1. assertPlanShape( """ OpenSearchSort(sort0=[$0], dir0=[ASC], viableBackends=[[mock-parquet]]) - OpenSearchProject(k=[$1], cnt=[$0], viableBackends=[[mock-parquet]]) - OpenSearchSort(sort0=[$0], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) - OpenSearchProject(cnt=[$1], k=[$0], viableBackends=[[mock-parquet]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[FINAL], viableBackends=[[mock-parquet]]) - OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) - OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[FINAL], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, result ); @@ -291,10 +292,12 @@ public void testSortThenProjectThenLimit_multiShard() { ); RelNode result = runPlanner(limit, buildContext("parquet", 3, fields)); + // SORT_PROJECT_TRANSPOSE pushes the outer pure-LIMIT Sort below the identity Project, + // producing the QTF-friendly two-Sort shape Project(identity) ← Sort(fetch) ← Sort(coll) ← ER. assertPlanShape( """ - OpenSearchSort(fetch=[3], viableBackends=[[mock-parquet]]) - OpenSearchProject(name=[$0], score=[$1], viableBackends=[[mock-parquet]]) + OpenSearchProject(name=[$0], score=[$1], viableBackends=[[mock-parquet]]) + OpenSearchSort(fetch=[3], viableBackends=[[mock-parquet]]) OpenSearchSort(sort0=[$1], dir0=[ASC], viableBackends=[[mock-parquet]]) OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SortRuleTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SortRuleTests.java index 6282eadfea2ff..34b083093674b 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SortRuleTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SortRuleTests.java @@ -18,6 +18,7 @@ import org.apache.calcite.sql.type.SqlTypeName; import org.opensearch.analytics.planner.rel.OpenSearchAggregate; import org.opensearch.analytics.planner.rel.OpenSearchFilter; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; import org.opensearch.analytics.planner.rel.OpenSearchSort; import org.opensearch.analytics.planner.rel.OpenSearchTableScan; import org.opensearch.analytics.planner.rules.OpenSearchFilterRule; @@ -44,10 +45,12 @@ private void assertSortPipeline( ) { logger.info("Plan:\n{}", RelOptUtil.toString(result)); assertPipelineViableBackends(result, types, Set.of(MockDataFusionBackend.NAME)); + // After QTF rewrite the wrapper sits at the root; the anchor Sort is its input. + OpenSearchSort sort = (OpenSearchSort) (result instanceof OpenSearchLateMaterialization wrap ? wrap.getInput() : result); if (fetch < 0) { - assertNull("Sort without limit must have null fetch", ((OpenSearchSort) result).fetch); + assertNull("Sort without limit must have null fetch", sort.fetch); } else { - assertNotNull("Sort with limit must have non-null fetch", ((OpenSearchSort) result).fetch); + assertNotNull("Sort with limit must have non-null fetch", sort.fetch); } } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SqlPlannerTestFixture.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SqlPlannerTestFixture.java index 240f9d82f8efb..5aaf239786c40 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SqlPlannerTestFixture.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/SqlPlannerTestFixture.java @@ -90,7 +90,12 @@ public static RelNode parseSql(String sql, ClusterState clusterState) { throw new UnsupportedOperationException("View expansion not used in tests"); }, validator, catalogReader, cluster, StandardConvertletTable.INSTANCE, SqlToRelConverter.config()); - return converter.convertQuery(parsed, true, true).rel; + // Use RelRoot.project() rather than .rel: SqlToRelConverter at top level appends + // ORDER BY columns to the projection without trimming them (intentional — Calcite's + // contract is that the caller applies the trim via RelRoot's `fields` mask). Without + // this, queries like `SELECT URL ORDER BY EventDate` return a 2-column result. + // See org.apache.calcite.rel.RelRoot's class doc for the canonical example. + return converter.convertQuery(parsed, true, true).project(); } /** diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGBuilderTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGBuilderTests.java index 1ada34f75cc6e..e0a4140bee379 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGBuilderTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGBuilderTests.java @@ -10,8 +10,12 @@ import com.google.common.collect.ImmutableList; import org.apache.calcite.plan.RelOptUtil; +import org.apache.calcite.rel.RelCollation; +import org.apache.calcite.rel.RelCollations; import org.apache.calcite.rel.RelDistribution; +import org.apache.calcite.rel.RelFieldCollation; import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.logical.LogicalSort; import org.apache.calcite.rel.logical.LogicalValues; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rex.RexLiteral; @@ -19,10 +23,21 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.analytics.planner.BasePlannerRulesTests; +import org.opensearch.analytics.planner.CapabilityRegistry; +import org.opensearch.analytics.planner.ClickBench; +import org.opensearch.analytics.planner.FieldStorageResolver; +import org.opensearch.analytics.planner.MockDataFusionBackend; +import org.opensearch.analytics.planner.PlannerContext; +import org.opensearch.analytics.planner.PlannerImpl; +import org.opensearch.analytics.planner.RelNodeUtils; +import org.opensearch.analytics.planner.SqlPlannerTestFixture; import org.opensearch.analytics.planner.rel.OpenSearchExchangeReducer; +import org.opensearch.analytics.planner.rel.OpenSearchLateMaterialization; import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan; import org.opensearch.analytics.planner.rel.OpenSearchTableScan; import org.opensearch.analytics.planner.rel.OpenSearchValues; +import org.opensearch.analytics.spi.FragmentConvertor; +import org.opensearch.cluster.ClusterState; import java.util.List; @@ -146,4 +161,258 @@ public void testValuesRootHasNoTargetResolverButHasSink() { assertNull("no TableScan → no ShardTargetResolver", dag.rootStage().getTargetResolver()); assertNotNull("compute leaf needs a sink provider for its local backend output", dag.rootStage().getExchangeSinkProvider()); } + + // ── QTF (late-materialization) DAG shapes ────────────────────────────── + + /** + * Drives SQL through the full planner so the QTF rewriter fires, then runs the result + * through {@link DAGBuilder}. Returns the four-stage DAG documented at + * {@link #testQtfDag_multiShardFourStages}. + */ + private QueryDAG buildQtfDag(String sql, int shardCount) { + ClusterState state = SqlPlannerTestFixture.clusterStateWith(ClickBench.INDEX, ClickBench.BASIC_FIELDS, "parquet", shardCount); + PlannerContext context = new PlannerContext( + new CapabilityRegistry(List.of(DATAFUSION, LUCENE), FieldStorageResolver::new), + state, + false + ); + RelNode parsed = SqlPlannerTestFixture.parseSql(sql, state); + RelNode cbo = PlannerImpl.runAllOptimizations(parsed, context); + QueryDAG dag = DAGBuilder.build(cbo, context.getCapabilityRegistry(), mockClusterService(), TEST_RESOLVER); + LOGGER.info("QTF QueryDAG:\n{}", dag); + return dag; + } + + /** + * Multi-shard QTF produces four stages, dataflow runs bottom-up from shards to the post-LM + * coordinator reduce: + * 

+     *   Stage 0 SHARD_FRAGMENT       (Filter? + Scan)
+     *     ↓
+     *   Stage 1 COORDINATOR_REDUCE   (Sort+Limit over reduce-set, ___row_id, ___ugsi)
+     *     ↓   inputSinkDecorator = OrdinalAppendingSink (stamps shard ordinal as ___ugsi)
+     *   Stage 2 LATE_MATERIALIZATION (wrapper rooted at StageInputScan(stage 1))
+     *     ↓
+     *   Stage 3 COORDINATOR_REDUCE   (post-LM ops: outer Project, etc.)  ← root
+     *
+ * + *

Stage 3 is the post-LM compute stage, separated by {@link DAGBuilder} so the LM stage + * itself stays a pure scatter/gather/stitch and the post-LM Project (or any future + * Filter/Aggregate/Sort) runs on Substrait via the standard reduce path. + */ + public void testQtfDag_multiShardFourStages() { + QueryDAG dag = buildQtfDag("SELECT URL, EventDate FROM hits ORDER BY EventDate LIMIT 10", 2); + assertBottomUpIds(dag.rootStage()); + + Stage postLm = dag.rootStage(); + assertEquals(StageExecutionType.COORDINATOR_REDUCE, postLm.getExecutionType()); + assertNull("post-LM reduce carries no input decorator", postLm.getInputSinkDecorator()); + assertEquals(1, postLm.getChildStages().size()); + + Stage lm = postLm.getChildStages().get(0); + assertEquals(StageExecutionType.LATE_MATERIALIZATION, lm.getExecutionType()); + assertNotNull( + "LM fragment must contain OpenSearchLateMaterialization wrapper", + RelNodeUtils.findNode(lm.getFragment(), OpenSearchLateMaterialization.class) + ); + assertEquals(1, lm.getChildStages().size()); + + Stage reduce = lm.getChildStages().get(0); + assertEquals(StageExecutionType.COORDINATOR_REDUCE, reduce.getExecutionType()); + assertNotNull("LM cut must install OrdinalAppendingSink decorator on reducer", reduce.getInputSinkDecorator()); + assertEquals(1, reduce.getChildStages().size()); + + // StageInputScan must carry ___ugsi — cutAtExchange uses the reducer's output rowType. + OpenSearchStageInputScan reduceInputScan = RelNodeUtils.findNode(reduce.getFragment(), OpenSearchStageInputScan.class); + assertNotNull(reduceInputScan); + assertEquals(0, reduceInputScan.getChildStageId()); + assertTrue( + "StageInputScan rowType must include " + OpenSearchLateMaterialization.UGSI_FIELD, + reduceInputScan.getRowType().getFieldNames().contains(OpenSearchLateMaterialization.UGSI_FIELD) + ); + + Stage scan = reduce.getChildStages().get(0); + assertEquals(StageExecutionType.SHARD_FRAGMENT, scan.getExecutionType()); + assertNull("scan stage carries no input decorator", scan.getInputSinkDecorator()); + assertNotNull("scan stage must have a target resolver", scan.getTargetResolver()); + assertEquals(0, scan.getChildStages().size()); + } + + /** + * Stage 0's shard fragment must propagate {@code __row_id__} on its Scan output. The + * rewriter narrows the Scan to {@code [belowAnchorPhysicalFields..., __row_id__]} via an + * override rowType; this asserts that override survives DAG cuts so the converted + * Substrait declares the helper column. Without it, Stage 1's reduce plan references + * {@code input-0.__row_id__} but the partition exposes only the physical cols and + * DataFusion fails the registration with "No field named __row_id__". + */ + public void testQtfDag_stage0ScanCarriesRowIdHelper() { + QueryDAG dag = buildQtfDag("SELECT URL, EventDate FROM hits ORDER BY EventDate LIMIT 10", 2); + Stage scan = dag.rootStage().getChildStages().getFirst().getChildStages().getFirst().getChildStages().getFirst(); + assertEquals(StageExecutionType.SHARD_FRAGMENT, scan.getExecutionType()); + + OpenSearchTableScan tableScan = RelNodeUtils.findNode(scan.getFragment(), OpenSearchTableScan.class); + assertNotNull("Stage 0 fragment must contain an OpenSearchTableScan", tableScan); + + List scanFieldNames = tableScan.getRowType().getFieldNames(); + assertTrue( + "Stage 0 Scan rowType must carry " + OpenSearchLateMaterialization.ROW_ID_FIELD + ", got " + scanFieldNames, + scanFieldNames.contains(OpenSearchLateMaterialization.ROW_ID_FIELD) + ); + } + + /** + * Reducer stages outside QTF carry no {@code inputSinkDecorator} — confirms the decorator + * is a QTF-only attachment, not a regression on every reduce. + */ + public void testReducerHasNoInputSinkDecoratorForNonQtf() { + QueryDAG dag = buildDAG(3, makeAggregate(sumCall())); + assertEquals(StageExecutionType.COORDINATOR_REDUCE, dag.rootStage().getExecutionType()); + assertNull(dag.rootStage().getInputSinkDecorator()); + } + + /** + * QTF wrapper ends up at the post-CBO root with NO operators above it (e.g. PPL frontends + * that don't add an outer Project around Sort+Limit). DAGBuilder must NOT route the resulting + * post-LM root stage to COORDINATOR_REDUCE — the fragment is a bare {@link OpenSearchStageInputScan} + * with zero inputs, which {@code FragmentConversionDriver.convertReduceNode} cannot handle + * (falls through to {@code getInputs().getFirst()} → NoSuchElementException). Route to + * LOCAL_PASSTHROUGH instead so the stage just relays Stitcher output upward. + * + *

Repro: build {@code Sort(Filter(Scan))} directly (bypasses SQL parse which would add an outer + * Project), drive through the planner + DAGBuilder + FragmentConversionDriver, assert + * {@code convertAll} succeeds. + */ + public void testQtfDag_lmAtRoot_noOuterProject_convertsCleanly() { + // Build LogicalSort(LogicalFilter(stubScan)) — no Project anywhere above the anchor. + // mockTable's columns mirror ClickBench's BASIC_FIELDS so QTF detects fetch-only fields. + RelNode scan = stubScan( + mockTable( + "test_index", + new String[] { "CounterID", "UserID", "URL", "Title", "EventDate", "AdvEngineID", "ParamPrice" }, + new SqlTypeName[] { + SqlTypeName.INTEGER, + SqlTypeName.BIGINT, + SqlTypeName.VARCHAR, + SqlTypeName.VARCHAR, + SqlTypeName.DATE, + SqlTypeName.SMALLINT, + SqlTypeName.BIGINT } + ) + ); + // ILIKE on URL ($2) — fetch-only column not in the anchor's sort key. + RelNode filter = makeFilter(scan, makeEquals(2, SqlTypeName.VARCHAR, "x")); + // Sort on EventDate ($4) with fetch=10 — anchor. + RelCollation collation = RelCollations.of(new RelFieldCollation(4, RelFieldCollation.Direction.ASCENDING)); + LogicalSort sort = LogicalSort.create( + filter, + collation, + null, + rexBuilder.makeLiteral(10, typeFactory.createSqlType(SqlTypeName.INTEGER), true) + ); + + // Drive through marking + CBO + LM rewrite. RecordingConvertor swapped onto the + // datafusion mock backend so FragmentConversionDriver.convertAll has something to call. + RecordingConvertor convertor = new RecordingConvertor(); + MockDataFusionBackend df = new MockDataFusionBackend() { + @Override + public FragmentConvertor getFragmentConvertor() { + return convertor; + } + }; + PlannerContext context = buildContext("parquet", 2, ClickBench.BASIC_FIELDS, List.of(df, LUCENE)); + RelNode cbo = runPlanner(sort, context); + LOGGER.info("Post-CBO/LM RelNode:\n{}", RelOptUtil.toString(cbo)); + + // QTF must have fired (LM wrapper present) and there must be no Project above it. + OpenSearchLateMaterialization wrapper = RelNodeUtils.findNode(cbo, OpenSearchLateMaterialization.class); + assertNotNull("QTF rewriter should fire (URL fetch-only)", wrapper); + assertSame("LM wrapper should be at the root (no Project above)", wrapper, RelNodeUtils.unwrapHep(cbo)); + + // Build DAG + run conversion. The bug is that Stage 3 (root) gets a bare StageInputScan + // and convertReduceNode trips on getInputs().getFirst(). + QueryDAG dag = DAGBuilder.build(cbo, context.getCapabilityRegistry(), mockClusterService(), TEST_RESOLVER); + LOGGER.info("QueryDAG:\n{}", dag); + PlanForker.forkAll(dag, context.getCapabilityRegistry()); + FragmentConversionDriver.convertAll(dag, context.getCapabilityRegistry()); + + // LM at root with no above-ops: the LM stage IS the root — no synthetic post-LM stage. + assertEquals( + "LM at CBO root must be promoted to rootStage; no synthetic post-LM stage", + StageExecutionType.LATE_MATERIALIZATION, + dag.rootStage().getExecutionType() + ); + assertNotNull( + "rootStage must contain the LM wrapper", + RelNodeUtils.findNode(dag.rootStage().getFragment(), OpenSearchLateMaterialization.class) + ); + } + + /** + * Contrast to {@link #testQtfDag_lmAtRoot_noOuterProject_convertsCleanly}: SAME query shape + * but WITH an outer Project above Sort+Limit. Post-LM, Stage 3's fragment is + * {@code Project(StageInputScan)} — real coordinator-side compute. Must be + * COORDINATOR_REDUCE so the convertor serializes the Project for execution. + */ + public void testQtfDag_lmAtRoot_withOuterProject_isCoordReduce() { + RelNode scan = stubScan( + mockTable( + "test_index", + new String[] { "CounterID", "UserID", "URL", "Title", "EventDate", "AdvEngineID", "ParamPrice" }, + new SqlTypeName[] { + SqlTypeName.INTEGER, + SqlTypeName.BIGINT, + SqlTypeName.VARCHAR, + SqlTypeName.VARCHAR, + SqlTypeName.DATE, + SqlTypeName.SMALLINT, + SqlTypeName.BIGINT } + ) + ); + // Filter on CounterID ($0) so filter/sort cols don't subsume the project — keeps URL fetch-only. + RelNode filter = makeFilter(scan, makeEquals(0, SqlTypeName.INTEGER, 5)); + RelCollation collation = RelCollations.of(new RelFieldCollation(4, RelFieldCollation.Direction.ASCENDING)); + LogicalSort sort = LogicalSort.create( + filter, + collation, + null, + rexBuilder.makeLiteral(10, typeFactory.createSqlType(SqlTypeName.INTEGER), true) + ); + // Outer Project: pick URL ($2) and EventDate ($4) — mirrors PPL `... | fields URL, EventDate`. + RelDataType varcharType = typeFactory.createSqlType(SqlTypeName.VARCHAR); + RelDataType dateType = typeFactory.createSqlType(SqlTypeName.DATE); + org.apache.calcite.rel.logical.LogicalProject project = org.apache.calcite.rel.logical.LogicalProject.create( + sort, + List.of(), + List.of(rexBuilder.makeInputRef(varcharType, 2), rexBuilder.makeInputRef(dateType, 4)), + List.of("URL", "EventDate") + ); + + RecordingConvertor convertor = new RecordingConvertor(); + MockDataFusionBackend df = new MockDataFusionBackend() { + @Override + public FragmentConvertor getFragmentConvertor() { + return convertor; + } + }; + PlannerContext context = buildContext("parquet", 2, ClickBench.BASIC_FIELDS, List.of(df, LUCENE)); + RelNode cbo = runPlanner(project, context); + LOGGER.info("Post-CBO/LM RelNode:\n{}", RelOptUtil.toString(cbo)); + + OpenSearchLateMaterialization wrapper = RelNodeUtils.findNode(cbo, OpenSearchLateMaterialization.class); + assertNotNull("QTF should fire", wrapper); + assertNotSame("Outer Project should sit above the LM wrapper", wrapper, RelNodeUtils.unwrapHep(cbo)); + + QueryDAG dag = DAGBuilder.build(cbo, context.getCapabilityRegistry(), mockClusterService(), TEST_RESOLVER); + LOGGER.info("QueryDAG:\n{}", dag); + PlanForker.forkAll(dag, context.getCapabilityRegistry()); + FragmentConversionDriver.convertAll(dag, context.getCapabilityRegistry()); + + // Stage 3 has real compute (the Project) on top of the StageInputScan → COORDINATOR_REDUCE. + assertEquals( + "post-LM root with above-ops must be COORDINATOR_REDUCE", + StageExecutionType.COORDINATOR_REDUCE, + dag.rootStage().getExecutionType() + ); + } } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGShapeTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGShapeTests.java index b2bcb50117b2f..3cece8c78a2e3 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGShapeTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/DAGShapeTests.java @@ -109,8 +109,8 @@ public void testJoinDag_case1_singleShardSameTable() { OpenSearchJoin(condition=[=($0, $2)], joinType=[left], viableBackends=[[mock-parquet]]) OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) - OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) - OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) + OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, dag @@ -129,8 +129,8 @@ public void testJoinDag_case2_multiShardSameTable() { OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) OpenSearchStageInputScan(childStageId=[0], viableBackends=[[mock-parquet]]) - OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) - OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) + OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) OpenSearchStageInputScan(childStageId=[1], viableBackends=[[mock-parquet]]) Stage 0 exchange=SINGLETON @@ -159,8 +159,8 @@ public void testJoinDag_case3_singleShardDifferentTables() { OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchTableScan(table=[[left_idx]], viableBackends=[[mock-parquet]]) Stage 1 exchange=SINGLETON - OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) - OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) + OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) OpenSearchTableScan(table=[[right_idx]], viableBackends=[[mock-parquet]]) """, dag @@ -179,8 +179,8 @@ public void testJoinDag_case4_multiShardDifferentTables() { OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) OpenSearchStageInputScan(childStageId=[0], viableBackends=[[mock-parquet]]) - OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) - OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) + OpenSearchProject(status=[$0], size=[$1], viableBackends=[[mock-parquet]]) + OpenSearchSort(fetch=[50000], viableBackends=[[mock-parquet]]) OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) OpenSearchStageInputScan(childStageId=[1], viableBackends=[[mock-parquet]]) Stage 0 exchange=SINGLETON @@ -201,12 +201,10 @@ public void testTopKAfterStatsDag_multiShard() { """ QueryDAG(queryId=) Stage 1 - OpenSearchProject(k=[$1], cnt=[$0], viableBackends=[[mock-parquet]]) - OpenSearchSort(sort0=[$0], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) - OpenSearchProject(cnt=[$1], k=[$0], viableBackends=[[mock-parquet]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[FINAL], viableBackends=[[mock-parquet]]) - OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) - OpenSearchStageInputScan(childStageId=[0], viableBackends=[[mock-parquet]]) + OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[FINAL], viableBackends=[[mock-parquet]]) + OpenSearchExchangeReducer(viableBackends=[[mock-parquet]], exchange=[ExchangeInfo[distributionType=SINGLETON, partitionKeyIndices=[]]]) + OpenSearchStageInputScan(childStageId=[0], viableBackends=[[mock-parquet]]) Stage 0 exchange=SINGLETON OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[PARTIAL], viableBackends=[[mock-parquet]]) OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) @@ -223,11 +221,9 @@ public void testTopKAfterStatsDag_singleShard() { assertDagShape(""" QueryDAG(queryId=) Stage 0 - OpenSearchProject(k=[$1], cnt=[$0], viableBackends=[[mock-parquet]]) - OpenSearchSort(sort0=[$0], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) - OpenSearchProject(cnt=[$1], k=[$0], viableBackends=[[mock-parquet]]) - OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[SINGLE], viableBackends=[[mock-parquet]]) - OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) + OpenSearchSort(sort0=[$1], dir0=[ASC], fetch=[2], viableBackends=[[mock-parquet]]) + OpenSearchAggregate(group=[{0}], cnt=[COUNT()], mode=[SINGLE], viableBackends=[[mock-parquet]]) + OpenSearchTableScan(table=[[test_index]], viableBackends=[[mock-parquet]]) """, dag); } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java index 25ae2acb6e3d9..1af860ecc9bd3 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/FragmentConversionDriverTests.java @@ -42,7 +42,6 @@ import org.opensearch.analytics.planner.rel.OpenSearchFilter; import org.opensearch.analytics.planner.rel.OpenSearchProject; import org.opensearch.analytics.planner.rel.OpenSearchSort; -import org.opensearch.analytics.planner.rel.OpenSearchTableScan; import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; import org.opensearch.analytics.spi.DelegatedPredicateFunction; import org.opensearch.analytics.spi.DelegatedPredicateSerializer; @@ -73,10 +72,14 @@ public class FragmentConversionDriverTests extends BasePlannerRulesTests { private static final Logger LOGGER = LogManager.getLogger(FragmentConversionDriverTests.class); + // Operator markers that carry annotations and must be stripped before isthmus conversion. + // OpenSearchTableScan is intentionally NOT in this set: stripAnnotations leaves it in place + // because it carries no annotations and isthmus only reads its rowType + qualified name. + // Stripping it to LogicalTableScan would silently drop QTF's overrideRowType (helper cols + // like __row_id__) — see OpenSearchTableScan#stripAnnotations. private static final Set OPENSEARCH_OPERATORS = Set.of( OpenSearchFilter.class.getSimpleName(), OpenSearchAggregate.class.getSimpleName(), - OpenSearchTableScan.class.getSimpleName(), OpenSearchSort.class.getSimpleName(), OpenSearchProject.class.getSimpleName() ); @@ -1718,44 +1721,4 @@ public void testNotCorrectnessAndPerfDelegation() { assertEquals(FilterTreeShape.CONJUNCTIVE, treeShapeOf(plan)); } - // ---- RecordingConvertor ---- - - /** Records which convertor method was called and what was passed. */ - private static class RecordingConvertor implements FragmentConvertor { - boolean shardScanCalled; - boolean finalAggCalled; - String shardScanTableName; - RelNode shardScanFragment; - RelNode reduceFragment; - - @Override - public byte[] convertFragment(RelNode fragment) { - // Distinguish shard-scan vs reduce/final by walking down the leftmost spine - // to find a TableScan-shaped leaf (annotations are stripped before this is - // called, so OpenSearchTableScan has been rewritten to LogicalTableScan). - org.apache.calcite.rel.core.TableScan scan = org.opensearch.analytics.planner.RelNodeUtils.findNode( - fragment, - org.apache.calcite.rel.core.TableScan.class - ); - if (scan != null) { - this.shardScanCalled = true; - this.shardScanTableName = scan.getTable().getQualifiedName().getLast(); - this.shardScanFragment = fragment; - return ("shard:" + this.shardScanTableName).getBytes(StandardCharsets.UTF_8); - } - this.finalAggCalled = true; - this.reduceFragment = fragment; - return "reduce".getBytes(StandardCharsets.UTF_8); - } - - @Override - public byte[] attachFragmentOnTop(RelNode fragment, byte[] innerBytes) { - return ("attach:" + new String(innerBytes, StandardCharsets.UTF_8)).getBytes(StandardCharsets.UTF_8); - } - - @Override - public byte[] attachPartialAggOnTop(RelNode partialAggFragment, byte[] innerBytes) { - return ("partialAgg:" + new String(innerBytes, StandardCharsets.UTF_8)).getBytes(StandardCharsets.UTF_8); - } - } } diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/PlanForkerTests.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/PlanForkerTests.java index ec120a356c3db..64dd6bd233c9d 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/PlanForkerTests.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/PlanForkerTests.java @@ -138,8 +138,11 @@ public void testSortQueryShapes() { 3, makeSort(makeFilter(stubScan(mockTable("test_index", "status", "size")), makeEquals(0, SqlTypeName.INTEGER, 200)), 10) ); - assertEquals(1, sortFilterDag.rootStage().getPlanAlternatives().size()); - for (StagePlan plan : sortFilterDag.rootStage().getPlanAlternatives()) { + // QTF rewriter inserts a LATE_MATERIALIZATION root stage above the original Sort+Limit + // reduce stage; this test was written pre-QTF to assert on the Sort stage. Skip past LM. + Stage sortFilterRoot = effectiveSortRoot(sortFilterDag.rootStage()); + assertEquals(1, sortFilterRoot.getPlanAlternatives().size()); + for (StagePlan plan : sortFilterRoot.getPlanAlternatives()) { assertTrue(plan.resolvedFragment() instanceof OpenSearchSort); } @@ -155,12 +158,28 @@ public void testSortQueryShapes() { 10 ) ); - assertEquals(1, sortAggDag.rootStage().getPlanAlternatives().size()); - for (StagePlan plan : sortAggDag.rootStage().getPlanAlternatives()) { + Stage sortAggRoot = effectiveSortRoot(sortAggDag.rootStage()); + assertEquals(1, sortAggRoot.getPlanAlternatives().size()); + for (StagePlan plan : sortAggRoot.getPlanAlternatives()) { assertTrue(plan.resolvedFragment() instanceof OpenSearchSort); } } + /** + * QTF wraps the original sort+limit reduce inside a 4-stage spine: post-LM COORDINATOR_REDUCE + * (root) → LATE_MATERIALIZATION → sort+limit COORDINATOR_REDUCE → SHARD_FRAGMENT. Descend + * down the first-child chain until we hit the stage whose fragment is the OpenSearchSort the + * pre-QTF assertions expect. + */ + private static Stage effectiveSortRoot(Stage root) { + Stage stage = root; + while (!(stage.getFragment() instanceof OpenSearchSort)) { + if (stage.getChildStages().isEmpty()) return stage; + stage = stage.getChildStages().getFirst(); + } + return stage; + } + /** * Aggregate(Filter(Scan)) — most common OLAP shape. Verifies that forking narrows * annotations consistently through the entire tree: both the aggregate root and the diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/RecordingConvertor.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/RecordingConvertor.java new file mode 100644 index 0000000000000..a33fc74f47234 --- /dev/null +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/dag/RecordingConvertor.java @@ -0,0 +1,62 @@ +/* + * 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.planner.dag; + +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.TableScan; +import org.apache.calcite.rel.type.RelDataType; +import org.opensearch.analytics.planner.RelNodeUtils; +import org.opensearch.analytics.spi.FragmentConvertor; + +import java.nio.charset.StandardCharsets; + +/** + * Records which {@link FragmentConvertor} method was called and what was passed. + * Shared across DAG / FragmentConversionDriver tests so any test can assert on + * the convertor surface that {@link FragmentConversionDriver#convertAll} hits. + */ +public class RecordingConvertor implements FragmentConvertor { + public boolean shardScanCalled; + public boolean finalAggCalled; + public String shardScanTableName; + public RelNode shardScanFragment; + public RelNode reduceFragment; + + @Override + public byte[] convertFragment(RelNode fragment) { + // Distinguish shard-scan vs reduce/final by walking down the leftmost spine to find a + // TableScan-shaped leaf (annotations are stripped before this is called, so + // OpenSearchTableScan has been rewritten to LogicalTableScan). + TableScan scan = RelNodeUtils.findNode(fragment, TableScan.class); + if (scan != null) { + this.shardScanCalled = true; + this.shardScanTableName = scan.getTable().getQualifiedName().getLast(); + this.shardScanFragment = fragment; + return ("shard:" + this.shardScanTableName).getBytes(StandardCharsets.UTF_8); + } + this.finalAggCalled = true; + this.reduceFragment = fragment; + return "reduce".getBytes(StandardCharsets.UTF_8); + } + + @Override + public byte[] attachFragmentOnTop(RelNode fragment, byte[] innerBytes) { + return ("attach:" + new String(innerBytes, StandardCharsets.UTF_8)).getBytes(StandardCharsets.UTF_8); + } + + @Override + public byte[] attachPartialAggOnTop(RelNode partialAggFragment, byte[] innerBytes) { + return ("partialAgg:" + new String(innerBytes, StandardCharsets.UTF_8)).getBytes(StandardCharsets.UTF_8); + } + + @Override + public byte[] convertSchemaOnlyRead(int stageId, RelDataType schema) { + return ("schemaOnlyRead:" + stageId).getBytes(StandardCharsets.UTF_8); + } +} diff --git a/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/be/datafusion/QtfDerivedAboveProjectIT.java b/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/be/datafusion/QtfDerivedAboveProjectIT.java new file mode 100644 index 0000000000000..0452d713f4529 --- /dev/null +++ b/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/be/datafusion/QtfDerivedAboveProjectIT.java @@ -0,0 +1,225 @@ +/* + * 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.be.datafusion; + +import org.apache.logging.log4j.Level; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.Version; +import org.opensearch.action.admin.indices.create.CreateIndexResponse; +import org.opensearch.analytics.AnalyticsPlugin; +import org.opensearch.analytics.exec.DefaultPlanExecutor; +import org.opensearch.analytics.sql.SqlPlanRunner; +import org.opensearch.arrow.flight.transport.FlightStreamPlugin; +import org.opensearch.arrow.allocator.ArrowBasePlugin; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.FeatureFlags; +import org.opensearch.composite.CompositeDataFormatPlugin; +import org.opensearch.index.engine.dataformat.stub.MockCommitterEnginePlugin; +import org.opensearch.parquet.ParquetDataFormatPlugin; +import org.opensearch.plugins.Plugin; +import org.opensearch.plugins.PluginInfo; +import org.opensearch.test.MockLogAppender; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.time.LocalDate; +import java.time.LocalDateTime; +import java.util.Collection; +import java.util.Collections; +import java.util.List; +import java.util.Locale; + +/** + * End-to-end IT for QTF (late materialization) Stage 3 — verifies the post-LM + * COORDINATOR_REDUCE actually runs derived expressions over the LM-stitched output. + * + *

Two assertions per test: + *

    + *
  1. Engagement — the rewriter's "[QTF] fired" debug line is captured by a + * {@link MockLogAppender}, proving QTF actually triggered. The rewriter logger is + * flipped to DEBUG via cluster settings only for the duration of the test (it stays + * off by default in production).
    2. + *
  2. Stage 3 compute — the SELECT carries a derived expression + * ({@code UPPER(URL)}). If Stage 3 weren't separated and run as Substrait by + * DAGBuilder's wrapper-cut, the engine would either error (LM stage doesn't run + * Substrait) or return raw URLs. Asserting on uppercased values proves the + * LM-stitched VSR was fed into Stage 3 and the post-LM Project executed.
    3. + *
+ * + * @opensearch.internal + */ +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.SUITE, numDataNodes = 2, numClientNodes = 0) +public class QtfDerivedAboveProjectIT extends OpenSearchIntegTestCase { + + private static final String INDEX = "qtf_derived_idx"; + private static final String REWRITER_LOGGER_NAME = "org.opensearch.analytics.planner.rules.OpenSearchLateMaterializationRewriter"; + + @Override + protected Collection> nodePlugins() { + return List.of(ArrowBasePlugin.class, CompositeDataFormatPlugin.class, MockCommitterEnginePlugin.class); + } + + @Override + protected Collection additionalNodePlugins() { + return List.of( + classpathPlugin(FlightStreamPlugin.class, List.of(ArrowBasePlugin.class.getName())), + classpathPlugin(AnalyticsPlugin.class, Collections.emptyList()), + classpathPlugin(ParquetDataFormatPlugin.class, Collections.emptyList()), + classpathPlugin(DataFusionPlugin.class, List.of(AnalyticsPlugin.class.getName())) + ); + } + + private static PluginInfo classpathPlugin(Class pluginClass, List extendedPlugins) { + return new PluginInfo( + pluginClass.getName(), + "classpath plugin", + "NA", + Version.CURRENT, + "1.8", + pluginClass.getName(), + null, + extendedPlugins, + false + ); + } + + @Override + protected Settings nodeSettings(int nodeOrdinal) { + return Settings.builder() + .put(super.nodeSettings(nodeOrdinal)) + .put(FeatureFlags.PLUGGABLE_DATAFORMAT_EXPERIMENTAL_FLAG, true) + .put(FeatureFlags.STREAM_TRANSPORT, true) + .build(); + } + + /** + * SELECT UPPER(URL), EventDate FROM qtf_derived_idx WHERE CounterID > 0 ORDER BY EventDate LIMIT 5 + * — multi-shard so QTF fires, outer Project carries a non-passthrough RexCall (UPPER) that + * must run on the post-LM Stage 3. Verifies engagement via log capture and result via + * uppercased URL. + * + *

The {@code WHERE CounterID > 0} predicate is here to give the indexed query strategy + * an index_filter() to attach against — the default {@code datafusion.indexed.query_strategy=indexed} + * rejects plans with no filter. + */ + public void testQtfFires_outerProjectRunsDerivedUpperOnStage3() throws Exception { + createAndSeedIndex(2); + + Logger rewriterLogger = LogManager.getLogger(REWRITER_LOGGER_NAME); + try (MockLogAppender appender = MockLogAppender.createForLoggers(rewriterLogger)) { + appender.addExpectation( + new MockLogAppender.SeenEventExpectation( + "QTF rewriter must fire for multi-shard sort+limit with above-anchor fetch-only cols", + REWRITER_LOGGER_NAME, + Level.DEBUG, + "*[QTF] fired*" + ) + ); + + withRewriterLogLevel("DEBUG", () -> { + SqlPlanRunner runner = sqlPlanRunner(); + List rows = runner.executeSql( + "SELECT UPPER(URL) AS u, EventDate FROM " + INDEX + " WHERE CounterID > 0 ORDER BY EventDate LIMIT 5" + ); + + // 10 seeded docs with EventDate 2026-05-01..2026-05-10; ASC LIMIT 5 → first 5. + assertEquals("LIMIT 5 must yield 5 rows", 5, rows.size()); + for (int i = 0; i < 5; i++) { + Object[] row = rows.get(i); + String expectedUrl = "HTTPS://EXAMPLE.COM/PAGE" + i; + assertEquals("row " + i + " UPPER(URL) mismatch", expectedUrl, row[0]); + + // EventDate is returned as a LocalDateTime by the executor (TIMESTAMP at + // midnight on the seeded date). + LocalDateTime expectedDate = LocalDate.of(2026, 5, i + 1).atStartOfDay(); + assertEquals("row " + i + " EventDate mismatch", expectedDate, row[1]); + } + }); + + appender.assertAllExpectationsMatched(); + } + } + + // ── Infrastructure ────────────────────────────────────────────────────── + + /** + * Flips the QTF rewriter's logger to {@code level} via cluster settings (the same path + * production operators would use), runs the body, and resets the override afterward. + * The rewriter's "fired" log line is intentionally DEBUG so prod stays log-quiet on the + * hot path; tests that need to observe engagement raise it just for their duration. + */ + private void withRewriterLogLevel(String level, ThrowingRunnable body) throws Exception { + client().admin() + .cluster() + .prepareUpdateSettings() + .setTransientSettings(Settings.builder().put("logger." + REWRITER_LOGGER_NAME, level).build()) + .get(); + try { + body.run(); + } finally { + client().admin() + .cluster() + .prepareUpdateSettings() + .setTransientSettings(Settings.builder().putNull("logger." + REWRITER_LOGGER_NAME).build()) + .get(); + } + } + + @FunctionalInterface + private interface ThrowingRunnable { + void run() throws Exception; + } + + private SqlPlanRunner sqlPlanRunner() { + String node = internalCluster().getNodeNames()[0]; + ClusterService clusterService = internalCluster().getInstance(ClusterService.class, node); + DefaultPlanExecutor executor = internalCluster().getInstance(DefaultPlanExecutor.class, node); + return new SqlPlanRunner(clusterService, executor); + } + + private void createAndSeedIndex(int shardCount) { + Settings indexSettings = Settings.builder() + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, shardCount) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) + .put("index.pluggable.dataformat.enabled", true) + .put("index.pluggable.dataformat", "composite") + .put("index.composite.primary_data_format", "parquet") + .putList("index.composite.secondary_data_formats") + .build(); + + CreateIndexResponse response = client().admin() + .indices() + .prepareCreate(INDEX) + .setSettings(indexSettings) + .setMapping("URL", "type=keyword", "EventDate", "type=date", "CounterID", "type=integer") + .get(); + assertTrue("index creation must be acknowledged", response.isAcknowledged()); + ensureGreen(INDEX); + + // 10 docs spread across two shards. URLs are lowercase so UPPER on Stage 3 is observable. + // CounterID is monotonically increasing so `WHERE CounterID > 0` matches all rows. + for (int i = 0; i < 10; i++) { + client().prepareIndex(INDEX) + .setId(String.valueOf(i)) + .setSource( + "URL", + "https://example.com/page" + i, + "EventDate", + "2026-05-" + String.format(Locale.ROOT, "%02d", i + 1), + "CounterID", + i + 1 + ) + .get(); + } + client().admin().indices().prepareRefresh(INDEX).get(); + client().admin().indices().prepareFlush(INDEX).get(); + } +} diff --git a/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/be/datafusion/QtfSubstraitDumpIT.java b/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/be/datafusion/QtfSubstraitDumpIT.java new file mode 100644 index 0000000000000..56ee88619ed7a --- /dev/null +++ b/sandbox/qa/analytics-engine-coordinator/src/internalClusterTest/java/org/opensearch/be/datafusion/QtfSubstraitDumpIT.java @@ -0,0 +1,316 @@ +/* + * 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.be.datafusion; + +import org.apache.calcite.avatica.util.Casing; +import org.apache.calcite.config.CalciteConnectionConfigImpl; +import org.apache.calcite.jdbc.CalciteSchema; +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.plan.RelOptCluster; +import org.apache.calcite.plan.RelOptUtil; +import org.apache.calcite.plan.hep.HepPlanner; +import org.apache.calcite.plan.hep.HepProgramBuilder; +import org.apache.calcite.prepare.CalciteCatalogReader; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.schema.SchemaPlus; +import org.apache.calcite.sql.SqlNode; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.parser.SqlParseException; +import org.apache.calcite.sql.parser.SqlParser; +import org.apache.calcite.sql.validate.SqlValidator; +import org.apache.calcite.sql.validate.SqlValidatorUtil; +import org.apache.calcite.sql2rel.SqlToRelConverter; +import org.apache.calcite.sql2rel.StandardConvertletTable; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.Version; +import org.opensearch.analytics.planner.CapabilityRegistry; +import org.opensearch.analytics.planner.FieldStorageResolver; +import org.opensearch.analytics.planner.PlannerContext; +import org.opensearch.analytics.planner.PlannerImpl; +import org.opensearch.analytics.planner.dag.DAGBuilder; +import org.opensearch.analytics.planner.dag.FragmentConversionDriver; +import org.opensearch.analytics.planner.dag.PlanForker; +import org.opensearch.analytics.planner.dag.QueryDAG; +import org.opensearch.analytics.planner.dag.Stage; +import org.opensearch.analytics.planner.dag.StagePlan; +import org.opensearch.analytics.schema.OpenSearchSchemaBuilder; +import org.opensearch.cluster.ClusterName; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.routing.GroupShardsIterator; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.routing.OperationRouting; +import org.opensearch.cluster.routing.ShardIterator; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.xcontent.MediaTypeRegistry; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Properties; + +import io.substrait.extension.DefaultExtensionCatalog; +import io.substrait.extension.SimpleExtension; +import io.substrait.proto.Plan; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Diagnostic dump: drives a QTF SQL query through the full planner stack + * (parse → CBO → QTF rewrite → DAG cut → fragment conversion) and prints both + * the per-stage RelNode trees and the per-stage Substrait Plan bytes (parsed back + * into proto for human-readable output). + * + *

Not a regression test — it has no asserts beyond "things didn't crash." + * Run with {@code --info} to see the dumped plans. + * + *

Lives in the QA module because the analytics-engine plugin's test classpath + * doesn't have substrait, but this module does. + * + * @opensearch.internal + */ +public class QtfSubstraitDumpIT extends OpenSearchTestCase { + + private static final Logger LOGGER = LogManager.getLogger(QtfSubstraitDumpIT.class); + + private static final String INDEX = "hits"; + private static final IndexNameExpressionResolver TEST_RESOLVER = new IndexNameExpressionResolver(new ThreadContext(Settings.EMPTY)); + + public void testDumpQtfPipeline() throws Exception { + String sql = "SELECT URL, EventDate FROM hits WHERE CounterID = 5 ORDER BY EventDate LIMIT 10"; + + Map> fields = new LinkedHashMap<>(); + fields.put("CounterID", Map.of("type", "integer")); + fields.put("UserID", Map.of("type", "long")); + fields.put("URL", Map.of("type", "keyword")); + fields.put("Title", Map.of("type", "keyword")); + fields.put("EventDate", Map.of("type", "date")); + + ClusterState clusterState = clusterStateWith(INDEX, fields, "parquet", 2); + + // Real Substrait extensions, loaded via the DataFusionPlugin's exact path. + SimpleExtension.ExtensionCollection extensions = loadExtensions(); + + // Stub DataFusionPlugin: only the bits getCapabilityProvider() and + // getFragmentConvertor() touch (no native runtime, no plugin lifecycle). + DataFusionPlugin dfPlugin = mock(DataFusionPlugin.class); + when(dfPlugin.name()).thenReturn("datafusion"); + when(dfPlugin.getSupportedFormats()).thenReturn(List.of("parquet")); + when(dfPlugin.getSubstraitExtensions()).thenReturn(extensions); + + DataFusionAnalyticsBackendPlugin dfBackend = new DataFusionAnalyticsBackendPlugin(dfPlugin); + + PlannerContext context = new PlannerContext( + new CapabilityRegistry(List.of(dfBackend), FieldStorageResolver::new), + clusterState, + false + ); + + // Parse → CBO → DAG → forker → fragment conversion. + RelNode parsed = parseSql(sql, clusterState); + LOGGER.info("[QTF-DUMP] sql:\n{}", sql); + LOGGER.info("[QTF-DUMP] parsed RelNode:\n{}", RelOptUtil.toString(parsed)); + + RelNode cbo = PlannerImpl.runAllOptimizations(parsed, context); + LOGGER.info("[QTF-DUMP] post-CBO+QTF RelNode:\n{}", RelOptUtil.toString(cbo)); + + QueryDAG dag = DAGBuilder.build(cbo, context.getCapabilityRegistry(), mockClusterService(), TEST_RESOLVER); + LOGGER.info("[QTF-DUMP] QueryDAG (pre-conversion):\n{}", dag); + + PlanForker.forkAll(dag, context.getCapabilityRegistry()); + FragmentConversionDriver.convertAll(dag, context.getCapabilityRegistry()); + LOGGER.info("[QTF-DUMP] QueryDAG (post-conversion, with backend-resolved fragments):\n{}", dag); + + // Walk every stage and dump its substrait Plan(s). + dumpSubstraitPerStage(dag.rootStage()); + } + + /** + * Stage 0's converted Substrait base_schema must include {@code __row_id__}. The rewriter + * adds the helper to the narrowed Scan rowType; if the override is dropped during DAG + * cuts or fragment conversion, Stage 1's reduce plan references {@code input-0.__row_id__} + * but the partition exposes only the physical cols and DataFusion fails registration with + * "No field named __row_id__". + */ + public void testStage0ScanCarriesRowIdInConvertedSubstrait() throws Exception { + String sql = "SELECT URL, EventDate FROM hits ORDER BY EventDate LIMIT 10"; + QueryDAG dag = buildAndConvertQtfDag(sql); + + // Stage 0 sits two levels below the LM stage (root post-LM reduce → LM → reduce → scan). + Stage scan = dag.rootStage().getChildStages().getFirst().getChildStages().getFirst().getChildStages().getFirst(); + byte[] bytes = scan.getPlanAlternatives().getFirst().convertedBytes(); + Plan plan = Plan.parseFrom(bytes); + + List baseSchemaNames = plan.getRelations(0).getRoot().getInput().getRead().getBaseSchema().getNamesList(); + assertTrue( + "Stage 0 base_schema must include __row_id__; got " + baseSchemaNames, + baseSchemaNames.contains("__row_id__") + ); + } + + /** + * Reusable harness: parses SQL, runs the planner, builds the DAG, forks plans, and runs + * fragment conversion. Returns the fully-converted DAG ready for per-stage assertions. + */ + private QueryDAG buildAndConvertQtfDag(String sql) { + Map> fields = new LinkedHashMap<>(); + fields.put("CounterID", Map.of("type", "integer")); + fields.put("UserID", Map.of("type", "long")); + fields.put("URL", Map.of("type", "keyword")); + fields.put("Title", Map.of("type", "keyword")); + fields.put("EventDate", Map.of("type", "date")); + ClusterState clusterState = clusterStateWith(INDEX, fields, "parquet", 2); + + SimpleExtension.ExtensionCollection extensions = loadExtensions(); + DataFusionPlugin dfPlugin = mock(DataFusionPlugin.class); + when(dfPlugin.name()).thenReturn("datafusion"); + when(dfPlugin.getSupportedFormats()).thenReturn(List.of("parquet")); + when(dfPlugin.getSubstraitExtensions()).thenReturn(extensions); + DataFusionAnalyticsBackendPlugin dfBackend = new DataFusionAnalyticsBackendPlugin(dfPlugin); + PlannerContext context = new PlannerContext( + new CapabilityRegistry(List.of(dfBackend), FieldStorageResolver::new), + clusterState, + false + ); + + RelNode parsed = parseSql(sql, clusterState); + RelNode cbo = PlannerImpl.runAllOptimizations(parsed, context); + QueryDAG dag = DAGBuilder.build(cbo, context.getCapabilityRegistry(), mockClusterService(), TEST_RESOLVER); + PlanForker.forkAll(dag, context.getCapabilityRegistry()); + FragmentConversionDriver.convertAll(dag, context.getCapabilityRegistry()); + return dag; + } + + private void dumpSubstraitPerStage(Stage stage) throws Exception { + LOGGER.info("[QTF-DUMP] === Stage {} ({}) — alternatives: {} ===", + stage.getStageId(), + stage.getExecutionType(), + stage.getPlanAlternatives().size() + ); + for (int i = 0; i < stage.getPlanAlternatives().size(); i++) { + StagePlan alt = stage.getPlanAlternatives().get(i); + byte[] bytes = alt.convertedBytes(); + if (bytes == null || bytes.length == 0) { + LOGGER.info("[QTF-DUMP] stage {} alt[{}] backend={} — no Substrait bytes (stage type doesn't convert)", + stage.getStageId(), i, alt.backendId() + ); + continue; + } + Plan plan = Plan.parseFrom(bytes); + LOGGER.info("[QTF-DUMP] stage {} alt[{}] backend={} — Substrait Plan ({} bytes):\n{}", + stage.getStageId(), i, alt.backendId(), bytes.length, plan + ); + } + for (Stage child : stage.getChildStages()) { + dumpSubstraitPerStage(child); + } + } + + // ── parse helpers (inlined from analytics-engine's SqlPlannerTestFixture) ───── + + private static RelNode parseSql(String sql, ClusterState clusterState) { + SchemaPlus schema = OpenSearchSchemaBuilder.buildSchema(clusterState); + RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(); + CalciteCatalogReader catalogReader = new CalciteCatalogReader( + CalciteSchema.from(schema), + Collections.singletonList(""), + typeFactory, + new CalciteConnectionConfigImpl(new Properties()) + ); + SqlValidator validator = SqlValidatorUtil.newValidator( + SqlStdOperatorTable.instance(), + catalogReader, + typeFactory, + SqlValidator.Config.DEFAULT + ); + HepPlanner hepPlanner = new HepPlanner(new HepProgramBuilder().build()); + RelOptCluster cluster = RelOptCluster.create(hepPlanner, new RexBuilder(typeFactory)); + + SqlParser.Config parserConfig = SqlParser.config().withUnquotedCasing(Casing.UNCHANGED); + SqlNode parsedNode; + try { + parsedNode = SqlParser.create(sql, parserConfig).parseQuery(); + } catch (SqlParseException e) { + throw new AssertionError("Failed to parse SQL: " + sql, e); + } + SqlToRelConverter converter = new SqlToRelConverter( + (rowType, queryString, schemaPath, viewPath) -> { throw new UnsupportedOperationException("View expansion not used"); }, + validator, + catalogReader, + cluster, + StandardConvertletTable.INSTANCE, + SqlToRelConverter.config() + ); + return converter.convertQuery(parsedNode, true, true).project(); + } + + private static ClusterState clusterStateWith(String indexName, Map> fields, String primaryDataFormat, int shardCount) { + try (XContentBuilder mapping = XContentBuilder.builder(MediaTypeRegistry.JSON.xContent())) { + mapping.startObject().field("properties", fields).endObject(); + IndexMetadata indexMetadata = IndexMetadata.builder(indexName) + .settings( + Settings.builder() + .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT.id) + .put("index.composite.primary_data_format", primaryDataFormat) + .putList("index.composite.secondary_data_formats", "lucene") + ) + .numberOfShards(shardCount) + .numberOfReplicas(0) + .putMapping(mapping.toString()) + .build(); + Metadata metadata = Metadata.builder().put(indexMetadata, false).build(); + return ClusterState.builder(new ClusterName("test")).metadata(metadata).build(); + } catch (Exception e) { + throw new AssertionError("Failed to build ClusterState", e); + } + } + + private static SimpleExtension.ExtensionCollection loadExtensions() { + Thread t = Thread.currentThread(); + ClassLoader prev = t.getContextClassLoader(); + try { + t.setContextClassLoader(QtfSubstraitDumpIT.class.getClassLoader()); + SimpleExtension.ExtensionCollection delegationExtensions = SimpleExtension.load(List.of("/delegation_functions.yaml")); + SimpleExtension.ExtensionCollection scalarExtensions = SimpleExtension.load(List.of("/opensearch_scalar_functions.yaml")); + SimpleExtension.ExtensionCollection arrayExtensions = SimpleExtension.load(List.of("/opensearch_array_functions.yaml")); + SimpleExtension.ExtensionCollection aggregateExtensions = SimpleExtension.load(List.of("/opensearch_aggregate_functions.yaml")); + SimpleExtension.ExtensionCollection roundingOverloads = SimpleExtension.load(List.of("/opensearch_rounding_overloads.yaml")); + return DefaultExtensionCatalog.DEFAULT_COLLECTION + .merge(delegationExtensions) + .merge(scalarExtensions) + .merge(arrayExtensions) + .merge(aggregateExtensions) + .merge(roundingOverloads); + } finally { + t.setContextClassLoader(prev); + } + } + + @SuppressWarnings("unchecked") + private static ClusterService mockClusterService() { + ClusterService clusterService = mock(ClusterService.class); + ClusterState state = mock(ClusterState.class); + OperationRouting routing = mock(OperationRouting.class); + when(clusterService.state()).thenReturn(state); + when(clusterService.operationRouting()).thenReturn(routing); + when(routing.searchShards(any(), any(), any(), any())).thenReturn(new GroupShardsIterator(List.of())); + return clusterService; + } +}