diff --git a/sandbox/libs/analytics-framework/licenses/jackson-core-2.22.1.jar.sha1 b/sandbox/libs/analytics-framework/licenses/jackson-core-2.22.1.jar.sha1 deleted file mode 100644 index e2c41405d169c..0000000000000 --- a/sandbox/libs/analytics-framework/licenses/jackson-core-2.22.1.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -da7ffb60088d7e8f37ecdd3b617520971cc7b9bf \ No newline at end of file diff --git a/sandbox/libs/analytics-framework/licenses/jackson-core-2.22.2.jar.sha1 b/sandbox/libs/analytics-framework/licenses/jackson-core-2.22.2.jar.sha1 new file mode 100644 index 0000000000000..67dfe67ab4b45 --- /dev/null +++ b/sandbox/libs/analytics-framework/licenses/jackson-core-2.22.2.jar.sha1 @@ -0,0 +1 @@ +13a748ea3e329fa220076e021b45c8391b32420c \ No newline at end of file diff --git a/sandbox/libs/analytics-framework/licenses/jackson-databind-2.22.1.jar.sha1 b/sandbox/libs/analytics-framework/licenses/jackson-databind-2.22.1.jar.sha1 deleted file mode 100644 index 7b38d1b99d9d0..0000000000000 --- a/sandbox/libs/analytics-framework/licenses/jackson-databind-2.22.1.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -9e2fb91831cce9cb9262909cd76647508949f232 \ No newline at end of file diff --git a/sandbox/libs/analytics-framework/licenses/jackson-databind-2.22.2.jar.sha1 b/sandbox/libs/analytics-framework/licenses/jackson-databind-2.22.2.jar.sha1 new file mode 100644 index 0000000000000..de250cf9085be --- /dev/null +++ b/sandbox/libs/analytics-framework/licenses/jackson-databind-2.22.2.jar.sha1 @@ -0,0 +1 @@ +921bd2092b0c539b2876de7063d55c72edcd05d3 \ No newline at end of file diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java index 6681cd65a0cbc..21889ecd10f54 100644 --- a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/backend/ShardScanExecutionContext.java @@ -12,6 +12,7 @@ import org.apache.lucene.search.QueryCache; import org.apache.lucene.search.QueryCachingPolicy; import org.opensearch.analytics.spi.CommonExecutionContext; +import org.opensearch.analytics.spi.DelegationThreadTracker; import org.opensearch.analytics.spi.ShuffleBufferRegistry; import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.index.shard.ShardId; @@ -41,6 +42,7 @@ public class ShardScanExecutionContext implements CommonExecutionContext { private QueryCachingPolicy queryCachingPolicy; private ShardId shardId; private boolean hasPartialAggregate; + private DelegationThreadTracker delegationThreadTracker; /** * Constructs an execution context. @@ -172,4 +174,14 @@ public boolean hasPartialAggregate() { public void setHasPartialAggregate(boolean hasPartialAggregate) { this.hasPartialAggregate = hasPartialAggregate; } + + /** Returns the tracker used to attribute callback work to the owning task. */ + public DelegationThreadTracker getDelegationThreadTracker() { + return delegationThreadTracker; + } + + /** Sets the tracker used to attribute callback work to the owning task. */ + public void setDelegationThreadTracker(DelegationThreadTracker delegationThreadTracker) { + this.delegationThreadTracker = delegationThreadTracker; + } } 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 071d1578561f1..0ed660ff7793f 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 @@ -10,11 +10,15 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.BigIntVector; +import org.apache.calcite.rel.RelNode; import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.cluster.ClusterState; +import org.opensearch.common.concurrent.GatedCloseable; import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; import org.opensearch.index.shard.IndexShard; +import org.opensearch.tasks.Task; +import java.io.IOException; import java.util.Collections; import java.util.List; import java.util.Map; @@ -43,6 +47,38 @@ public interface AnalyticsSearchBackendPlugin { /** Unique backend name (e.g., "datafusion", "lucene"). */ String name(); + /** Supplies the complete backend registry after all extensions have loaded. */ + default void bindBackends(Map backends) {} + + /** Whether this backend can execute plans over a caller-provided Arrow batch source. */ + default boolean supportsArrowBatchSourceExecution() { + return false; + } + + /** Compiles a fragment over a named Arrow input. */ + default byte[] compileArrowBatchSourcePlan(RelNode fragment, boolean partialAggregate) { + throw new UnsupportedOperationException("Arrow batch source compilation not implemented for [" + name() + "]"); + } + + /** Attaches a coordinator-local fragment to an already compiled Arrow source plan. */ + default byte[] attachArrowBatchSourcePlan(RelNode fragment, byte[] innerPlanBytes) { + throw new UnsupportedOperationException("Arrow batch source composition not implemented for [" + name() + "]"); + } + + /** + * Executes a backend-native plan over a caller-provided Arrow batch source. + * Ownership of {@code sourceFactory} transfers to this method, including on failure. + */ + default EngineResultStream executeArrowBatchSource( + BufferAllocator resultAllocator, + ArrowBatchSourcePlan plan, + ArrowBatchSourceFactory sourceFactory, + Task task, + DelegationThreadTracker threadTracker + ) { + throw new UnsupportedOperationException("Arrow batch source execution not implemented for [" + name() + "]"); + } + /** * Returns the capability provider for this backend. * Used by the coordinator-side planner ({@code CapabilityRegistry}) to determine @@ -93,6 +129,23 @@ default FragmentInstructionHandlerFactory getInstructionHandlerFactory() { throw new UnsupportedOperationException("getInstructionHandlerFactory not implemented for [" + name() + "]"); } + /** + * Acquires the point-in-time shard reader used by this backend. Pluggable-format backends use + * the shard's {@link org.opensearch.index.engine.exec.IndexReaderProvider} directly. A backend + * that also supports a standard engine can override this method to adapt that engine's native + * searcher into the shared {@link Reader} contract. + */ + default GatedCloseable acquireReader(IndexShard shard) throws IOException { + try { + return shard.getReaderProvider().acquireReader(); + } catch (UnsupportedOperationException e) { + throw new UnsupportedOperationException( + "Backend [" + name() + "] cannot acquire a reader from " + shard.getReaderProvider().getClass().getName(), + e + ); + } + } + /** * Prepare a filter delegation handle for the given delegated expressions. * Called by Core after all instruction handlers have run, when the plan has delegation. diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ArrowBatchSource.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ArrowBatchSource.java new file mode 100644 index 0000000000000..246e4da9c2348 --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ArrowBatchSource.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.spi; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; + +/** + * Source for one sequential Arrow input stream. + * + *

Implementations return ownership of every non-null batch to the caller. A null + * batch signals EOF. Sources are single-consumer and close must be idempotent. + * Implementations that can block should override {@link #cancel()} so a concurrent + * cancellation request can make {@link #nextBatch()} return promptly. + * + * @opensearch.internal + */ +public interface ArrowBatchSource extends AutoCloseable { + + /** Allocator that owns returned batches and exported Arrow C Data buffers. */ + BufferAllocator allocator(); + + /** Returns the next owned batch, or {@code null} at EOF. */ + VectorSchemaRoot nextBatch() throws Exception; + + /** + * Requests cooperative cancellation. + * + *

This method can run concurrently with {@link #nextBatch()}. Implementations must + * return promptly and must not close resources still in use by that call. The default + * is a no-op for compatibility; such sources can delay release until a pending {@link #nextBatch()} call returns. + */ + default void cancel() {} + + @Override + void close(); +} diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ArrowBatchSourceFactory.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ArrowBatchSourceFactory.java new file mode 100644 index 0000000000000..8499ad5dde773 --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ArrowBatchSourceFactory.java @@ -0,0 +1,53 @@ +/* + * 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.spi; + +/** + * Reusable factory for DataFusion-driven Arrow scans. + * + *

The projection contains zero-based indices into the factory's declared input + * schema. Every open call returns an independent sequential cursor. Closing the factory + * must be idempotent. + * + * @opensearch.internal + */ +public interface ArrowBatchSourceFactory extends AutoCloseable { + + /** Logical type of one source column. */ + enum ColumnKind { + /** Signed 64-bit integer. */ + LONG, + /** UTF-8 keyword. */ + KEYWORD, + /** Millisecond timestamp. */ + TIMESTAMP, + /** Boolean value encoded as zero or one. */ + BOOLEAN, + /** IEEE 754 single-precision floating-point value. */ + FLOAT, + /** IEEE 754 double-precision floating-point value. */ + DOUBLE, + /** Opaque bytes. */ + BINARY, + /** Encoded IP address bytes. */ + IP + } + + /** One source column. Multi-valued columns use an Arrow list of the declared kind. */ + record InputColumn(String name, ColumnKind kind, boolean multiValued) { + public InputColumn(String name, ColumnKind kind) { + this(name, kind, false); + } + } + + ArrowBatchSource open(int[] projection) throws Exception; + + @Override + void close(); +} diff --git a/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ArrowBatchSourcePlan.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ArrowBatchSourcePlan.java new file mode 100644 index 0000000000000..f830331aa3fd9 --- /dev/null +++ b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/ArrowBatchSourcePlan.java @@ -0,0 +1,96 @@ +/* + * 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.spi; + +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.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.InputColumn; +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; +import java.util.List; +import java.util.Objects; + +/** + * Engine plan that consumes one named {@link ArrowBatchSourceFactory} input. + * + * @opensearch.internal + */ +public record ArrowBatchSourcePlan(String inputId, byte[] planBytes, List inputColumns) implements Writeable { + + public ArrowBatchSourcePlan { + inputId = Objects.requireNonNull(inputId, "inputId"); + if (inputId.isBlank()) { + throw new IllegalArgumentException("inputId must not be blank"); + } + planBytes = Objects.requireNonNull(planBytes, "planBytes").clone(); + if (planBytes.length == 0) { + throw new IllegalArgumentException("planBytes must not be empty"); + } + inputColumns = List.copyOf(Objects.requireNonNull(inputColumns, "inputColumns")); + } + + public ArrowBatchSourcePlan(StreamInput input) throws IOException { + this(input.readString(), input.readByteArray(), input.readList(ArrowBatchSourcePlan::readInputColumn)); + } + + @Override + public byte[] planBytes() { + return planBytes.clone(); + } + + /** Arrow schema declared by this plan's input columns. */ + public Schema inputSchema() { + return schemaFor(inputColumns); + } + + /** Builds the source schema for an input-column list. */ + public static Schema schemaFor(List columns) { + return new Schema(columns.stream().map(ArrowBatchSourcePlan::toField).toList()); + } + + @Override + public void writeTo(StreamOutput output) throws IOException { + output.writeString(inputId); + output.writeByteArray(planBytes); + output.writeCollection(inputColumns, ArrowBatchSourcePlan::writeInputColumn); + } + + private static Field toField(InputColumn column) { + ArrowType type = switch (column.kind()) { + case LONG -> new ArrowType.Int(64, true); + case KEYWORD -> new ArrowType.Utf8View(); + case TIMESTAMP -> new ArrowType.Timestamp(TimeUnit.MILLISECOND, null); + case BOOLEAN -> ArrowType.Bool.INSTANCE; + case FLOAT -> new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE); + case DOUBLE -> new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE); + case BINARY, IP -> new ArrowType.Binary(); + }; + if (column.multiValued()) { + Field item = new Field("item", FieldType.nullable(type), null); + return new Field(column.name(), FieldType.nullable(new ArrowType.List()), List.of(item)); + } + return new Field(column.name(), FieldType.nullable(type), null); + } + + private static InputColumn readInputColumn(StreamInput input) throws IOException { + return new InputColumn(input.readString(), input.readEnum(ArrowBatchSourceFactory.ColumnKind.class), input.readBoolean()); + } + + private static void writeInputColumn(StreamOutput output, InputColumn column) throws IOException { + output.writeString(column.name()); + output.writeEnum(column.kind()); + output.writeBoolean(column.multiValued()); + } +} diff --git a/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/ArrowBatchSourcePlanTests.java b/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/ArrowBatchSourcePlanTests.java new file mode 100644 index 0000000000000..55cfe58f8832b --- /dev/null +++ b/sandbox/libs/analytics-framework/src/test/java/org/opensearch/analytics/spi/ArrowBatchSourcePlanTests.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.spi; + +import org.apache.arrow.vector.types.FloatingPointPrecision; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.ColumnKind; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.InputColumn; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.List; + +public class ArrowBatchSourcePlanTests extends OpenSearchTestCase { + + public void testSchemaAndWireRoundTrip() throws Exception { + List columns = List.of( + new InputColumn("long", ColumnKind.LONG), + new InputColumn("keyword", ColumnKind.KEYWORD), + new InputColumn("timestamp", ColumnKind.TIMESTAMP), + new InputColumn("boolean", ColumnKind.BOOLEAN), + new InputColumn("float", ColumnKind.FLOAT), + new InputColumn("double", ColumnKind.DOUBLE), + new InputColumn("binary", ColumnKind.BINARY), + new InputColumn("ip", ColumnKind.IP), + new InputColumn("many_ip", ColumnKind.IP, true) + ); + ArrowBatchSourcePlan original = new ArrowBatchSourcePlan("input-0", new byte[] { 1, 2 }, columns); + Schema schema = original.inputSchema(); + + assertTrue(schema.findField("long").getType() instanceof ArrowType.Int); + assertTrue(schema.findField("keyword").getType() instanceof ArrowType.Utf8View); + assertTrue(schema.findField("timestamp").getType() instanceof ArrowType.Timestamp); + assertTrue(schema.findField("boolean").getType() instanceof ArrowType.Bool); + assertEquals(FloatingPointPrecision.SINGLE, ((ArrowType.FloatingPoint) schema.findField("float").getType()).getPrecision()); + assertEquals(FloatingPointPrecision.DOUBLE, ((ArrowType.FloatingPoint) schema.findField("double").getType()).getPrecision()); + assertTrue(schema.findField("binary").getType() instanceof ArrowType.Binary); + assertTrue(schema.findField("ip").getType() instanceof ArrowType.Binary); + assertTrue(schema.findField("many_ip").getType() instanceof ArrowType.List); + assertTrue(schema.findField("many_ip").getChildren().get(0).getType() instanceof ArrowType.Binary); + + try (BytesStreamOutput out = new BytesStreamOutput()) { + original.writeTo(out); + try (StreamInput in = out.bytes().streamInput()) { + ArrowBatchSourcePlan decoded = new ArrowBatchSourcePlan(in); + assertEquals(original.inputId(), decoded.inputId()); + assertArrayEquals(original.planBytes(), decoded.planBytes()); + assertEquals(columns, decoded.inputColumns()); + } + } + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/licenses/jackson-datatype-jdk8-2.22.1.jar.sha1 b/sandbox/plugins/analytics-backend-datafusion/licenses/jackson-datatype-jdk8-2.22.1.jar.sha1 deleted file mode 100644 index 8aec5b1d3886e..0000000000000 --- a/sandbox/plugins/analytics-backend-datafusion/licenses/jackson-datatype-jdk8-2.22.1.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -ef21a94e7b5603f57ca1efe9215513d210d96320 \ No newline at end of file diff --git a/sandbox/plugins/analytics-backend-datafusion/licenses/jackson-datatype-jdk8-2.22.2.jar.sha1 b/sandbox/plugins/analytics-backend-datafusion/licenses/jackson-datatype-jdk8-2.22.2.jar.sha1 new file mode 100644 index 0000000000000..ef40d77269bbb --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/licenses/jackson-datatype-jdk8-2.22.2.jar.sha1 @@ -0,0 +1 @@ +0b900dd7125fa16cfdf46135d3ffb3243d0f8b88 \ No newline at end of file diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs index 4a38104684f9b..c598a09d0411b 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs @@ -1924,6 +1924,22 @@ pub unsafe fn close_local_session(ptr: i64) { } } +/// Registers a Java-backed Arrow batch source on a local session. +/// +/// # Safety +/// `session_ptr` must be a valid local-session pointer. +pub unsafe fn register_arrow_batch_source_provider( + session_ptr: i64, + input_id: &str, + schema_ipc: &[u8], + binding_id: i64, + task_id: i64, +) -> Result<(), DataFusionError> { + let session = &mut *(session_ptr as *mut LocalSession); + let schema = schema_from_ipc_bytes(schema_ipc)?; + session.register_arrow_batch_source_provider(input_id, schema, binding_id, task_id) +} + /// Registers a streaming input on the session under `input_id`. The schema is /// derived by lowering `partial_plan_bytes` (the producer side's substrait) to /// a physical plan and reading its output schema — that is the schema the diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/callbacks.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/callbacks.rs new file mode 100644 index 0000000000000..38167845adc73 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/callbacks.rs @@ -0,0 +1,192 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +use std::sync::atomic::{AtomicBool, AtomicPtr, Ordering}; + +use arrow_array::ffi::{from_ffi, FFI_ArrowArray}; +use arrow_array::{RecordBatch, StructArray}; +use arrow_schema::ffi::FFI_ArrowSchema; +use datafusion::common::DataFusionError; + +const CANCELLED: i64 = -1; +const ERROR: i64 = -2; +const EMPTY_BATCH: i64 = -3; +const ERROR_CAPACITY: usize = 2048; + +type CreateSourceFn = unsafe extern "C" fn(i64, *const i32, i64, *mut u8, i64) -> i32; +type NextBatchFn = + unsafe extern "C" fn(i64, i32, *mut FFI_ArrowArray, *mut FFI_ArrowSchema, *mut u8, i64) -> i64; +type CancelSourceFn = unsafe extern "C" fn(i64, i32); +type ReleaseSourceFn = unsafe extern "C" fn(i64, i32); + +static CREATE_SOURCE: AtomicPtr<()> = AtomicPtr::new(std::ptr::null_mut()); +static NEXT_BATCH: AtomicPtr<()> = AtomicPtr::new(std::ptr::null_mut()); +static CANCEL_SOURCE: AtomicPtr<()> = AtomicPtr::new(std::ptr::null_mut()); +static RELEASE_SOURCE: AtomicPtr<()> = AtomicPtr::new(std::ptr::null_mut()); + +#[no_mangle] +pub unsafe extern "C" fn df_register_arrow_batch_source_callbacks( + create_source: CreateSourceFn, + next_batch: NextBatchFn, + cancel_source: CancelSourceFn, + release_source: ReleaseSourceFn, +) { + CREATE_SOURCE.store(create_source as *mut (), Ordering::Release); + NEXT_BATCH.store(next_batch as *mut (), Ordering::Release); + CANCEL_SOURCE.store(cancel_source as *mut (), Ordering::Release); + RELEASE_SOURCE.store(release_source as *mut (), Ordering::Release); +} + +fn create_callback() -> Result { + let pointer = CREATE_SOURCE.load(Ordering::Acquire); + if pointer.is_null() { + return Err(DataFusionError::Execution( + "Arrow batch source callbacks are not registered".into(), + )); + } + Ok(unsafe { std::mem::transmute::<*mut (), CreateSourceFn>(pointer) }) +} + +fn next_callback() -> Result { + let pointer = NEXT_BATCH.load(Ordering::Acquire); + if pointer.is_null() { + return Err(DataFusionError::Execution( + "Arrow batch source callbacks are not registered".into(), + )); + } + Ok(unsafe { std::mem::transmute::<*mut (), NextBatchFn>(pointer) }) +} + +fn cancel_callback() -> Option { + let pointer = CANCEL_SOURCE.load(Ordering::Acquire); + (!pointer.is_null()).then(|| unsafe { std::mem::transmute::<*mut (), CancelSourceFn>(pointer) }) +} + +fn release_callback() -> Option { + let pointer = RELEASE_SOURCE.load(Ordering::Acquire); + (!pointer.is_null()) + .then(|| unsafe { std::mem::transmute::<*mut (), ReleaseSourceFn>(pointer) }) +} + +#[derive(Debug)] +pub struct ArrowBatchSourceHandle { + binding_id: i64, + source_key: i32, + cancelled: AtomicBool, +} + +impl ArrowBatchSourceHandle { + pub fn create(binding_id: i64, projection: &[usize]) -> Result { + let projected: Result, _> = + projection.iter().copied().map(i32::try_from).collect(); + let projected = projected.map_err(|_| { + DataFusionError::Execution("Arrow batch source projection exceeds i32".into()) + })?; + let callback = create_callback()?; + let mut error = [0u8; ERROR_CAPACITY]; + let key = unsafe { + callback( + binding_id, + projected.as_ptr(), + projected.len() as i64, + error.as_mut_ptr(), + error.len() as i64, + ) + }; + if key < 0 { + return Err(callback_error( + &error, + format!("Java failed to create Arrow batch source binding {binding_id}"), + )); + } + Ok(Self { + binding_id, + source_key: key, + cancelled: AtomicBool::new(false), + }) + } + + pub fn next_batch(&self) -> Result, DataFusionError> { + let callback = next_callback()?; + let mut array = FFI_ArrowArray::empty(); + let mut schema = FFI_ArrowSchema::empty(); + let mut error = [0u8; ERROR_CAPACITY]; + let status = unsafe { + callback( + self.binding_id, + self.source_key, + &mut array, + &mut schema, + error.as_mut_ptr(), + error.len() as i64, + ) + }; + match status { + 0 => Ok(None), + CANCELLED => Err(DataFusionError::Execution("query cancelled".into())), + ERROR => Err(callback_error( + &error, + "Java Arrow batch source failed".into(), + )), + EMPTY_BATCH => import_batch(array, schema, 0), + rows if rows > 0 => import_batch(array, schema, rows as usize), + other => Err(DataFusionError::Execution(format!( + "unexpected Arrow batch callback status {other}" + ))), + } + } + + /// Requests cooperative cancellation without releasing the active source. + pub fn cancel(&self) { + if self.cancelled.swap(true, Ordering::AcqRel) == false { + if let Some(callback) = cancel_callback() { + unsafe { callback(self.binding_id, self.source_key) }; + } + } + } +} + +impl Drop for ArrowBatchSourceHandle { + fn drop(&mut self) { + self.cancel(); + if let Some(callback) = release_callback() { + unsafe { callback(self.binding_id, self.source_key) }; + } + } +} + +fn import_batch( + array: FFI_ArrowArray, + schema: FFI_ArrowSchema, + expected_rows: usize, +) -> Result, DataFusionError> { + let mut data = unsafe { from_ffi(array, &schema) }.map_err(|error| { + DataFusionError::Execution(format!("failed to import Java Arrow batch: {error}")) + })?; + data.align_buffers(); + let batch = RecordBatch::from(StructArray::from(data)); + if batch.num_rows() != expected_rows { + return Err(DataFusionError::Execution(format!( + "Arrow batch callback returned row count {expected_rows}, imported {}", + batch.num_rows() + ))); + } + Ok(Some(batch)) +} + +fn callback_error(buffer: &[u8], fallback: String) -> DataFusionError { + let length = buffer + .iter() + .position(|value| *value == 0) + .unwrap_or(buffer.len()); + if length == 0 { + DataFusionError::Execution(fallback) + } else { + DataFusionError::Execution(String::from_utf8_lossy(&buffer[..length]).into_owned()) + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/mod.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/mod.rs new file mode 100644 index 0000000000000..a4a81daba4e5e --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/mod.rs @@ -0,0 +1,12 @@ +/* + * 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. + */ + +pub mod callbacks; +pub mod table_provider; + +pub use table_provider::ArrowBatchTableProvider; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/table_provider.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/table_provider.rs new file mode 100644 index 0000000000000..a5a0f728158f2 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/arrow_batch_source/table_provider.rs @@ -0,0 +1,321 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +use std::fmt; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; +use async_trait::async_trait; +use datafusion::catalog::{Session, TableProvider}; +use datafusion::common::{DataFusionError, Result}; +use datafusion::datasource::TableType; +use datafusion::execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties}; +use futures::{Future, Stream}; +use tokio::task::JoinHandle; +use tokio_util::sync::CancellationToken; + +use super::callbacks::ArrowBatchSourceHandle; + +/// DataFusion table backed by externally allocated Arrow batches. +/// +/// The Java owner must use a breaker-accounted allocator. Imported buffers remain charged +/// to that allocator for their lifetime and do not use coordinator reduce-budget admission. +#[derive(Debug)] +pub struct ArrowBatchTableProvider { + schema: SchemaRef, + binding_id: i64, + task_id: i64, +} + +impl ArrowBatchTableProvider { + pub fn new(schema: SchemaRef, binding_id: i64, task_id: i64) -> Self { + Self { + schema, + binding_id, + task_id, + } + } +} + +#[async_trait] +impl TableProvider for ArrowBatchTableProvider { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> Result> { + Ok(vec![ + TableProviderFilterPushDown::Unsupported; + filters.len() + ]) + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + let projection = projection + .cloned() + .unwrap_or_else(|| (0..self.schema.fields().len()).collect()); + let projected_schema = Arc::new(self.schema.project(&projection)?); + Ok(Arc::new(ArrowBatchSourceExec::new( + projected_schema, + projection, + self.binding_id, + crate::query_tracker::get_cancellation_token(self.task_id), + ))) + } +} + +pub struct ArrowBatchSourceExec { + schema: SchemaRef, + projection: Vec, + binding_id: i64, + cancellation_token: Option, + properties: Arc, +} + +impl ArrowBatchSourceExec { + fn new( + schema: SchemaRef, + projection: Vec, + binding_id: i64, + cancellation_token: Option, + ) -> Self { + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + Partitioning::UnknownPartitioning(1), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { + schema, + projection, + binding_id, + cancellation_token, + properties, + } + } +} + +impl fmt::Debug for ArrowBatchSourceExec { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ArrowBatchSourceExec") + .field("projection", &self.projection) + .field("binding_id", &self.binding_id) + .finish() + } +} + +impl DisplayAs for ArrowBatchSourceExec { + fn fmt_as( + &self, + _display_type: DisplayFormatType, + formatter: &mut fmt::Formatter<'_>, + ) -> fmt::Result { + write!( + formatter, + "ArrowBatchSourceExec: projection={:?}", + self.projection + ) + } +} + +impl ExecutionPlan for ArrowBatchSourceExec { + fn name(&self) -> &str { + "ArrowBatchSourceExec" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn execute( + &self, + partition: usize, + _context: Arc, + ) -> Result { + if partition != 0 { + return Err(DataFusionError::Execution(format!( + "ArrowBatchSourceExec partition {partition} out of range" + ))); + } + let source = Arc::new(ArrowBatchSourceHandle::create( + self.binding_id, + &self.projection, + )?); + Ok(Box::pin(ArrowBatchRecordBatchStream::new( + Arc::clone(&self.schema), + source, + self.cancellation_token.clone(), + ))) + } +} + +pub(super) struct ArrowBatchRecordBatchStream { + schema: SchemaRef, + source: Arc, + cancellation: Option + Send>>>, + cancellation_requested: bool, + pending: Option>>>, + finished: bool, +} + +impl ArrowBatchRecordBatchStream { + pub(super) fn new( + schema: SchemaRef, + source: Arc, + cancellation_token: Option, + ) -> Self { + let cancellation = cancellation_token.map(|token| { + Box::pin(async move { token.cancelled().await }) + as Pin + Send>> + }); + Self { + schema, + source, + cancellation, + cancellation_requested: false, + pending: None, + finished: false, + } + } + + fn poll_cancellation(&mut self, context: &mut Context<'_>) { + if self.cancellation_requested { + return; + } + if let Some(cancellation) = self.cancellation.as_mut() { + if cancellation.as_mut().poll(context).is_ready() { + self.cancellation_requested = true; + self.source.cancel(); + } + } + } + + fn cancelled(&mut self) -> Poll>> { + self.finished = true; + Poll::Ready(Some(Err(DataFusionError::Execution( + "query cancelled".into(), + )))) + } +} + +impl fmt::Debug for ArrowBatchRecordBatchStream { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ArrowBatchRecordBatchStream") + .field("finished", &self.finished) + .finish() + } +} + +impl Stream for ArrowBatchRecordBatchStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + if self.finished { + return Poll::Ready(None); + } + self.poll_cancellation(context); + if self.cancellation_requested && self.pending.is_none() { + return self.cancelled(); + } + if self.pending.is_none() { + let source = Arc::clone(&self.source); + self.pending = Some(tokio::task::spawn_blocking(move || source.next_batch())); + } + let handle = self.pending.as_mut().expect("pending source request"); + match Pin::new(handle).poll(context) { + Poll::Pending => Poll::Pending, + Poll::Ready(result) => { + self.pending = None; + if self.cancellation_requested { + return self.cancelled(); + } + match result { + Err(error) => { + self.finished = true; + Poll::Ready(Some(Err(DataFusionError::Execution(format!( + "Arrow batch source task failed: {error}" + ))))) + } + Ok(Err(error)) => { + self.finished = true; + Poll::Ready(Some(Err(error))) + } + Ok(Ok(None)) => { + self.finished = true; + Poll::Ready(None) + } + Ok(Ok(Some(batch))) => { + if batch.schema().fields() != self.schema.fields() { + self.finished = true; + return Poll::Ready(Some(Err(DataFusionError::Execution(format!( + "Arrow batch schema mismatch: expected {:?}, got {:?}", + self.schema, + batch.schema() + ))))); + } + Poll::Ready(Some(Ok(batch))) + } + } + } + } + } +} + +impl Drop for ArrowBatchRecordBatchStream { + fn drop(&mut self) { + // Cancellation is cooperative. Sources that keep the default no-op cancel method + // remain retained by an active blocking task until that pull returns safely. + self.source.cancel(); + if let Some(handle) = self.pending.take() { + handle.abort(); + } + } +} + +impl RecordBatchStream for ArrowBatchRecordBatchStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs index a3c0f2e3397e1..fe3ab459336f1 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs @@ -664,6 +664,26 @@ pub unsafe extern "C" fn df_destroy_custom_cache_manager(ptr: i64) { } } +/// Registers a Java-backed Arrow batch source provider on a local session. +#[ffm_safe] +#[no_mangle] +pub unsafe extern "C" fn df_register_arrow_batch_source_provider( + session_ptr: i64, + input_id_ptr: *const u8, + input_id_len: i64, + schema_ptr: *const u8, + schema_len: i64, + binding_id: i64, + task_id: i64, +) -> i64 { + let input_id = str_from_raw(input_id_ptr, input_id_len) + .map_err(|error| format!("df_register_arrow_batch_source_provider: input_id: {error}"))?; + let schema = slice::from_raw_parts(schema_ptr, schema_len as usize); + api::register_arrow_batch_source_provider(session_ptr, input_id, schema, binding_id, task_id) + .map_err(|error| error.to_string())?; + Ok(0) +} + /// Registers a streaming partition input on the session. Schema is derived by /// lowering the producer-side substrait `partial_plan_bytes`; the resulting /// IPC-encoded schema is written into the caller-allocated `out_ptr/out_cap` diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs index c5fddc00610bd..4604effd37a8a 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/lib.rs @@ -22,6 +22,7 @@ pub const ROW_ID_COLUMN_NAME: &str = "__row_id__"; pub(crate) mod agg_mode; pub mod api; +pub mod arrow_batch_source; pub mod cache; pub mod can_match; pub mod cancellation; diff --git a/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs b/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs index 42bf70d811c4a..cbdeeabcfea13 100644 --- a/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs +++ b/sandbox/plugins/analytics-backend-datafusion/rust/src/local_executor.rs @@ -160,6 +160,26 @@ impl LocalSession { Ok(sender) } + /// Registers a pull-based Java Arrow batch source. + pub fn register_arrow_batch_source_provider( + &mut self, + name: &str, + schema: SchemaRef, + binding_id: i64, + task_id: i64, + ) -> Result<(), DataFusionError> { + let table = + crate::arrow_batch_source::ArrowBatchTableProvider::new(schema, binding_id, task_id); + self.ctx + .register_table(name, Arc::new(table)) + .map_err(|error| { + DataFusionError::Execution(format!( + "Failed to register Arrow batch source table '{name}': {error}" + )) + })?; + Ok(()) + } + /// Registers an in-memory input on the session under `name`, holding all /// `batches` in a single [`MemTable`] partition. /// 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 1e3f75b5d9800..6a29100cfc08d 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 @@ -11,6 +11,8 @@ import org.apache.arrow.memory.BufferAllocator; import org.apache.arrow.vector.BigIntVector; import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.Aggregate; import org.apache.calcite.sql.SqlOperator; import org.apache.calcite.sql.fun.SqlLibraryOperators; import org.apache.calcite.sql.fun.SqlStdOperatorTable; @@ -21,10 +23,13 @@ import org.opensearch.analytics.exec.shuffle.ShuffleCompression; import org.opensearch.analytics.exec.task.AnalyticsShardTask; import org.opensearch.analytics.planner.CalciteToArrowSchema; +import org.opensearch.analytics.planner.dag.BackendPlanAdapter; 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.ArrowBatchSourceFactory; +import org.opensearch.analytics.spi.ArrowBatchSourcePlan; import org.opensearch.analytics.spi.BackendCapabilityProvider; import org.opensearch.analytics.spi.BackendExecutionContext; import org.opensearch.analytics.spi.DataTransferCapability; @@ -61,6 +66,7 @@ import org.opensearch.index.engine.dataformat.DataFormatRegistry; import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; import org.opensearch.index.shard.IndexShard; +import org.opensearch.tasks.Task; import java.util.HashSet; import java.util.List; @@ -909,6 +915,47 @@ public FragmentConvertor getFragmentConvertor() { return new DataFusionFragmentConvertor(plugin.getSubstraitExtensions()); } + @Override + public boolean supportsArrowBatchSourceExecution() { + return true; + } + + @Override + public byte[] compileArrowBatchSourcePlan(RelNode fragment, boolean partialAggregate) { + RelNode adapted = BackendPlanAdapter.adaptFragment(fragment, getCapabilityProvider()); + DataFusionFragmentConvertor convertor = new DataFusionFragmentConvertor(plugin.getSubstraitExtensions()); + if (partialAggregate == false) { + return convertor.convertFragment(adapted); + } + if (adapted instanceof Aggregate == false) { + throw new IllegalArgumentException("partial Arrow source plan must be rooted at an Aggregate"); + } + Aggregate aggregate = (Aggregate) adapted; + return convertor.attachPartialAggOnTop(aggregate, convertor.convertFragment(aggregate.getInput())); + } + + @Override + public byte[] attachArrowBatchSourcePlan(RelNode fragment, byte[] innerPlanBytes) { + RelNode adapted = BackendPlanAdapter.adaptFragment(fragment, getCapabilityProvider()); + return new DataFusionFragmentConvertor(plugin.getSubstraitExtensions()).attachFragmentOnTop(adapted, innerPlanBytes); + } + + @Override + public EngineResultStream executeArrowBatchSource( + BufferAllocator resultAllocator, + ArrowBatchSourcePlan plan, + ArrowBatchSourceFactory sourceFactory, + Task task, + DelegationThreadTracker threadTracker + ) { + DataFusionService service = plugin.getDataFusionService(); + if (service == null) { + sourceFactory.close(); + throw new IllegalStateException("DataFusionService not initialized"); + } + return new DatafusionArrowBatchSourceExecutor(service).execute(resultAllocator, plan, sourceFactory, task, threadTracker); + } + @Override public SearchExecEngineProvider getSearchExecEngineProvider() { return (ctx, backendContext) -> { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutor.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutor.java new file mode 100644 index 0000000000000..627eb83f51eca --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutor.java @@ -0,0 +1,179 @@ +/* + * 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.arrow.memory.BufferAllocator; +import org.opensearch.analytics.backend.EngineResultBatch; +import org.opensearch.analytics.backend.EngineResultStream; +import org.opensearch.analytics.exec.FragmentResources; +import org.opensearch.analytics.exec.task.AnalyticsShardTask; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory; +import org.opensearch.analytics.spi.ArrowBatchSourcePlan; +import org.opensearch.analytics.spi.DelegationThreadTracker; +import org.opensearch.be.datafusion.arrow.ArrowBatchSourceCallbacks; +import org.opensearch.be.datafusion.nativelib.NativeBridge; +import org.opensearch.be.datafusion.nativelib.StreamHandle; +import org.opensearch.core.tasks.TaskCancelledException; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.Task; + +import java.util.Iterator; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicBoolean; + +/** Executes one DataFusion plan over a caller-provided Arrow batch source. */ +final class DatafusionArrowBatchSourceExecutor { + + private final DataFusionService service; + + DatafusionArrowBatchSourceExecutor(DataFusionService service) { + this.service = Objects.requireNonNull(service, "service"); + } + + EngineResultStream execute( + BufferAllocator resultAllocator, + ArrowBatchSourcePlan plan, + ArrowBatchSourceFactory sourceFactory, + Task task, + DelegationThreadTracker threadTracker + ) { + Objects.requireNonNull(resultAllocator, "resultAllocator"); + Objects.requireNonNull(plan, "plan"); + Objects.requireNonNull(sourceFactory, "sourceFactory"); + if (task instanceof CancellableTask cancellableTask && cancellableTask.isCancelled()) { + sourceFactory.close(); + throw new TaskCancelledException("Arrow batch source execution cancelled before setup"); + } + + ArrowBatchSourceCallbacks.Registration registration; + try { + registration = ArrowBatchSourceCallbacks.register(sourceFactory, threadTracker); + } catch (RuntimeException | Error throwable) { + closeAfterFailure(throwable, sourceFactory::close); + throw throwable; + } + + DatafusionLocalSession session = null; + StreamHandle output = null; + AnalyticsShardTask shardTask = task instanceof AnalyticsShardTask analyticsShardTask ? analyticsShardTask : null; + long taskId = task == null ? 0L : task.getId(); + try { + session = new DatafusionLocalSession(service.getNativeRuntime().get()); + NativeBridge.registerArrowBatchSourceProvider( + session.getPointer(), + plan.inputId(), + ArrowSchemaIpc.toBytes(plan.inputSchema()), + registration.bindingId(), + taskId + ); + if (shardTask != null) { + shardTask.setCancellationListener(() -> NativeBridge.cancelQuery(taskId)); + } + long streamPointer = NativeBridge.executeLocalPlan(session.getPointer(), plan.planBytes(), taskId); + output = new StreamHandle(streamPointer, service.getNativeRuntime()); + DatafusionResultStream delegate = new DatafusionResultStream(output, resultAllocator); + return new OwnedResultStream(delegate, session, registration, shardTask, taskId); + } catch (RuntimeException | Error throwable) { + if (shardTask != null) { + shardTask.clearCancellationListener(); + } + if (taskId != 0L) { + closeAfterFailure(throwable, () -> NativeBridge.cancelQuery(taskId)); + } + if (output != null) { + closeAfterFailure(throwable, output::close); + } + if (session != null) { + closeAfterFailure(throwable, session::close); + } + closeAfterFailure(throwable, registration::close); + throw throwable; + } + } + + private static void closeAfterFailure(Throwable failure, Runnable closeAction) { + try { + closeAction.run(); + } catch (RuntimeException | Error closeFailure) { + failure.addSuppressed(closeFailure); + } + } + + /** Owns the native output, local session, callback binding, and source factory. */ + private static final class OwnedResultStream implements EngineResultStream, FragmentResources.MetricsCapable { + private final DatafusionResultStream delegate; + private final DatafusionLocalSession session; + private final ArrowBatchSourceCallbacks.Registration registration; + private final AnalyticsShardTask shardTask; + private final long taskId; + private final AtomicBoolean closed = new AtomicBoolean(); + + private OwnedResultStream( + DatafusionResultStream delegate, + DatafusionLocalSession session, + ArrowBatchSourceCallbacks.Registration registration, + AnalyticsShardTask shardTask, + long taskId + ) { + this.delegate = delegate; + this.session = session; + this.registration = registration; + this.shardTask = shardTask; + this.taskId = taskId; + } + + @Override + public Iterator iterator() { + return delegate.iterator(); + } + + @Override + public byte[] getMetricsJson() { + return delegate.getMetricsJson(); + } + + @Override + public void close() { + if (closed.compareAndSet(false, true) == false) { + return; + } + if (shardTask != null) { + shardTask.clearCancellationListener(); + } + Throwable failure = null; + if (taskId != 0L) { + try { + NativeBridge.cancelQuery(taskId); + } catch (RuntimeException | Error throwable) { + failure = throwable; + } + } + failure = close(delegate::close, failure); + failure = close(session::close, failure); + failure = close(registration::close, failure); + if (failure instanceof RuntimeException runtimeException) { + throw runtimeException; + } + if (failure instanceof Error error) { + throw error; + } + } + + private static Throwable close(Runnable closeAction, Throwable failure) { + try { + closeAction.run(); + } catch (RuntimeException | Error closeFailure) { + if (failure == null) { + return closeFailure; + } + failure.addSuppressed(closeFailure); + } + return failure; + } + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionLocalSession.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionLocalSession.java index 9457220da9ea1..e47ff63c56e63 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionLocalSession.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/DatafusionLocalSession.java @@ -16,9 +16,8 @@ * coordinator-reduce path ({@link DatafusionReduceSink}). * *

The session holds a DataFusion {@code SessionContext} bound to the node-global runtime's - * memory pool and disk manager. It owns any input partition streams registered via - * {@link NativeBridge#registerPartitionStream(long, String, byte[])} and drops them when the - * session itself is closed. + * memory pool and disk manager. It owns table providers and input partition streams registered + * through {@link NativeBridge}; closing the session drops those native registrations. */ public final class DatafusionLocalSession extends NativeHandle { diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java index bc71c033508b3..79028681a6017 100644 --- a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/PplAggregateCallRewriter.java @@ -23,6 +23,7 @@ import org.apache.calcite.rex.RexLiteral; import org.apache.calcite.rex.RexNode; import org.apache.calcite.sql.SqlAggFunction; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; import org.apache.calcite.sql.type.SqlTypeName; import java.util.ArrayList; @@ -30,8 +31,8 @@ import java.util.Set; /** - * Rewrites PPL state-expanding aggregates (TAKE / FIRST / LAST / LIST / VALUES / PATTERN / - * PERCENTILE_APPROX) onto local stubs the substrait emitter binds via + * Rewrites PPL custom aggregates (CHECKED_LONG_SUM / TAKE / FIRST / LAST / LIST / VALUES / + * PATTERN / PERCENTILE_APPROX) onto standard operators or local stubs the substrait emitter binds via * {@link DataFusionFragmentConvertor}'s ADDITIONAL_AGGREGATE_SIGS. Also normalises any * RexLiteral{SqlTypeName.SYMBOL} in upstream Projects to VARCHAR — isthmus's * LiteralConverter rejects unregistered Enum classes, and PPL's percentile_approx / @@ -121,6 +122,7 @@ private static AggregateCall rewriteCall(Aggregate agg, AggregateCall call) { boolean targetDistinct = call.isDistinct(); RelDataType explicitReturnType = call.getType(); switch (aggregation.getName().toUpperCase(java.util.Locale.ROOT)) { + case "CHECKED_LONG_SUM" -> targetOp = SqlStdOperatorTable.SUM; case "TAKE" -> targetOp = DataFusionFragmentConvertor.LOCAL_TAKE_OP; case "FIRST" -> targetOp = DataFusionFragmentConvertor.LOCAL_FIRST_OP; case "LAST" -> targetOp = DataFusionFragmentConvertor.LOCAL_LAST_OP; diff --git a/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java new file mode 100644 index 0000000000000..b0a181e897cb4 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/main/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacks.java @@ -0,0 +1,426 @@ +/* + * 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.arrow; + +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.Data; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.analytics.spi.ArrowBatchSource; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory; +import org.opensearch.analytics.spi.DelegationThreadTracker; +import org.opensearch.core.tasks.TaskCancelledException; + +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import java.nio.charset.StandardCharsets; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; + +/** Java upcall targets for a native DataFusion Arrow batch source. */ +public final class ArrowBatchSourceCallbacks { + + public static final long CANCELLED = -1L; + public static final long ERROR = -2L; + public static final long EMPTY_BATCH = -3L; + + private static final Logger LOGGER = LogManager.getLogger(ArrowBatchSourceCallbacks.class); + private static final AtomicLong NEXT_BINDING_ID = new AtomicLong(1L); + private static final ConcurrentHashMap BINDINGS = new ConcurrentHashMap<>(); + + private ArrowBatchSourceCallbacks() {} + + /** + * Registers one query-scoped factory and returns its independently allocated binding. + * The factory must allocate batches through a caller-owned, breaker-accounted allocator; + * native consumers import those externally accounted buffers without taking ownership of + * their memory admission. + */ + public static Registration register(ArrowBatchSourceFactory factory, DelegationThreadTracker tracker) { + if (factory == null) { + throw new NullPointerException("factory"); + } + while (true) { + long bindingId = NEXT_BINDING_ID.getAndIncrement(); + if (bindingId <= 0L) { + throw new IllegalStateException("Arrow batch source binding IDs exhausted"); + } + Binding binding = new Binding(bindingId, factory, tracker); + if (BINDINGS.putIfAbsent(bindingId, binding) == null) { + return new Registration(bindingId, binding); + } + } + } + + /** Owns one callback binding. Close requests cleanup and is idempotent. */ + public static final class Registration implements AutoCloseable { + private final long bindingId; + private final Binding binding; + private final AtomicBoolean closed = new AtomicBoolean(); + + private Registration(long bindingId, Binding binding) { + this.bindingId = bindingId; + this.binding = binding; + } + + public long bindingId() { + return bindingId; + } + + @Override + public void close() { + if (closed.compareAndSet(false, true)) { + binding.requestClose(); + } + } + } + + /** FFM upcall: opens a projected source and returns its positive key. */ + public static int createSource( + long bindingId, + MemorySegment projectionPointer, + long projectionLength, + MemorySegment errorPointer, + long errorCapacity + ) { + CallbackLease lease = acquire(bindingId, false); + if (lease == null) { + writeError(errorPointer, errorCapacity, "Arrow batch source binding is closed: " + bindingId); + return -1; + } + try (lease) { + int length = Math.toIntExact(projectionLength); + if (length < 0) { + throw new IllegalArgumentException("projection length must be non-negative"); + } + int[] projection; + if (length == 0) { + projection = new int[0]; + } else { + if (projectionPointer == null || projectionPointer.equals(MemorySegment.NULL)) { + throw new IllegalArgumentException("projection pointer is null"); + } + long byteLength = (long) length * Integer.BYTES; + MemorySegment projectionView = projectionPointer.byteSize() >= byteLength + ? projectionPointer.asSlice(0L, byteLength) + : projectionPointer.reinterpret(byteLength); + projection = projectionView.toArray(ValueLayout.JAVA_INT); + } + return lease.binding.open(projection); + } catch (Throwable throwable) { + writeError(errorPointer, errorCapacity, throwable.toString()); + LOGGER.debug("Failed to create Arrow batch source for binding {}", bindingId, throwable); + return -1; + } + } + + /** FFM upcall: exports the next owned batch through Arrow C Data. */ + public static long nextBatch( + long bindingId, + int sourceKey, + MemorySegment arrayPointer, + MemorySegment schemaPointer, + MemorySegment errorPointer, + long errorCapacity + ) { + CallbackLease lease = acquire(bindingId, false); + if (lease == null) { + writeError(errorPointer, errorCapacity, "Arrow batch source binding is closed: " + bindingId); + return ERROR; + } + try (lease) { + SourceEntry source = lease.binding.source(sourceKey); + if (source == null) { + writeError(errorPointer, errorCapacity, "Unknown Arrow batch source key: " + sourceKey); + return ERROR; + } + return source.exportNextBatch(arrayPointer, schemaPointer); + } catch (TaskCancelledException cancelled) { + return CANCELLED; + } catch (Throwable throwable) { + writeError(errorPointer, errorCapacity, throwable.toString()); + LOGGER.debug("Failed to read Arrow batch source binding {} key {}", bindingId, sourceKey, throwable); + return ERROR; + } + } + + /** FFM upcall: cooperatively cancels one source without waiting for an active pull. */ + public static void cancelSource(long bindingId, int sourceKey) { + CallbackLease lease = acquire(bindingId, true); + if (lease == null) { + return; + } + try (lease) { + lease.binding.cancel(sourceKey); + } catch (Throwable throwable) { + LOGGER.warn("Failed to cancel Arrow batch source binding {} key {}", bindingId, sourceKey, throwable); + } + } + + /** FFM upcall: releases one source. Unknown or already-released keys are ignored. */ + public static void releaseSource(long bindingId, int sourceKey) { + CallbackLease lease = acquire(bindingId, true); + if (lease == null) { + return; + } + try (lease) { + lease.binding.release(sourceKey); + } catch (Throwable throwable) { + LOGGER.warn("Failed to release Arrow batch source binding {} key {}", bindingId, sourceKey, throwable); + } + } + + private static CallbackLease acquire(long bindingId, boolean allowClosing) { + Binding binding = BINDINGS.get(bindingId); + if (binding == null || binding.acquire(allowClosing) == false) { + return null; + } + return new CallbackLease(binding); + } + + private static void writeError(MemorySegment output, long capacity, String message) { + if (output == null || output.equals(MemorySegment.NULL) || capacity <= 0L) { + return; + } + byte[] bytes = message.getBytes(StandardCharsets.UTF_8); + int length = Math.toIntExact(Math.min(bytes.length, capacity - 1L)); + MemorySegment view = output.reinterpret(capacity); + MemorySegment.copy(bytes, 0, view, ValueLayout.JAVA_BYTE, 0L, length); + view.set(ValueLayout.JAVA_BYTE, length, (byte) 0); + } + + private static final class CallbackLease implements AutoCloseable { + private final Binding binding; + private final long trackedThreadId; + private boolean closed; + + private CallbackLease(Binding binding) { + this.binding = binding; + this.trackedThreadId = binding.trackStart(); + } + + @Override + public void close() { + if (closed == false) { + closed = true; + try { + binding.trackEnd(trackedThreadId); + } finally { + binding.releaseCallback(); + } + } + } + } + + private static final class Binding { + private final long bindingId; + private final ArrowBatchSourceFactory factory; + private final DelegationThreadTracker tracker; + private final Object openLock = new Object(); + private final Map sources = new HashMap<>(); + private int nextSourceKey = 1; + private int activeCallbacks; + private boolean closing; + private boolean factoryClosed; + + private Binding(long bindingId, ArrowBatchSourceFactory factory, DelegationThreadTracker tracker) { + this.bindingId = bindingId; + this.factory = factory; + this.tracker = tracker; + } + + private synchronized boolean acquire(boolean allowClosing) { + if (factoryClosed || (closing && allowClosing == false)) { + return false; + } + activeCallbacks++; + return true; + } + + private int open(int[] projection) throws Exception { + synchronized (this) { + if (closing) { + throw new IllegalStateException("Arrow batch source binding is closing: " + bindingId); + } + } + ArrowBatchSource source; + synchronized (openLock) { + synchronized (this) { + if (closing) { + throw new IllegalStateException("Arrow batch source binding is closing: " + bindingId); + } + } + source = factory.open(projection); + } + if (source == null) { + throw new IllegalStateException("Arrow batch source factory returned null"); + } + SourceEntry entry = new SourceEntry(source); + String rejection = null; + int sourceKey = -1; + synchronized (this) { + if (closing) { + rejection = "Arrow batch source binding is closing: " + bindingId; + } else { + sourceKey = nextSourceKey++; + if (sourceKey <= 0) { + rejection = "Arrow batch source keys exhausted for binding " + bindingId; + } else { + sources.put(sourceKey, entry); + } + } + } + if (rejection != null) { + entry.close(); + throw new IllegalStateException(rejection); + } + return sourceKey; + } + + private synchronized SourceEntry source(int sourceKey) { + return sources.get(sourceKey); + } + + private synchronized void cancel(int sourceKey) { + SourceEntry source = sources.get(sourceKey); + if (source != null) { + source.cancel(); + } + } + + private synchronized void release(int sourceKey) { + SourceEntry source = sources.remove(sourceKey); + if (source != null) { + source.close(); + } + tryFinishClose(); + } + + private void requestClose() { + List openSources; + synchronized (this) { + closing = true; + openSources = List.copyOf(sources.values()); + } + for (SourceEntry source : openSources) { + try { + source.cancel(); + } catch (Throwable throwable) { + LOGGER.warn("Failed to cancel Arrow batch source while closing binding {}", bindingId, throwable); + } + } + synchronized (this) { + tryFinishClose(); + } + } + + private synchronized void releaseCallback() { + activeCallbacks--; + tryFinishClose(); + } + + private synchronized void tryFinishClose() { + if (closing == false || factoryClosed || activeCallbacks != 0) { + return; + } + // Native stream teardown normally releases each source, but an output stream can be + // closed before its provider is polled. Registration owns final cleanup and must not + // wait for a native release callback that may never arrive. + List remainingSources = List.copyOf(sources.values()); + sources.clear(); + for (SourceEntry source : remainingSources) { + try { + source.close(); + } catch (Throwable throwable) { + LOGGER.warn("Failed to close Arrow batch source while closing binding {}", bindingId, throwable); + } + } + factoryClosed = true; + try { + factory.close(); + } finally { + BINDINGS.remove(bindingId, this); + } + } + + private long trackStart() { + if (tracker == null) { + return -1L; + } + try { + return tracker.trackStart(); + } catch (Throwable throwable) { + LOGGER.warn("Failed to start Arrow batch source callback tracking", throwable); + return -1L; + } + } + + private void trackEnd(long threadId) { + if (tracker == null || threadId < 0L) { + return; + } + try { + tracker.trackEnd(threadId); + } catch (Throwable throwable) { + LOGGER.warn("Failed to finish Arrow batch source callback tracking", throwable); + } + } + } + + private static final class SourceEntry { + private final ArrowBatchSource source; + private final Object cancellationLock = new Object(); + private final AtomicBoolean cancellationRequested = new AtomicBoolean(); + private boolean closed; + + private SourceEntry(ArrowBatchSource source) { + this.source = source; + } + + private synchronized long exportNextBatch(MemorySegment arrayPointer, MemorySegment schemaPointer) throws Exception { + if (closed) { + throw new IllegalStateException("Arrow batch source is closed"); + } + VectorSchemaRoot root = source.nextBatch(); + if (root == null) { + return 0L; + } + try { + ArrowArray array = ArrowArray.wrap(arrayPointer.address()); + ArrowSchema schema = ArrowSchema.wrap(schemaPointer.address()); + Data.exportVectorSchemaRoot(source.allocator(), root, null, array, schema); + return root.getRowCount() == 0 ? EMPTY_BATCH : root.getRowCount(); + } finally { + root.close(); + } + } + + private void cancel() { + synchronized (cancellationLock) { + if (cancellationRequested.compareAndSet(false, true)) { + source.cancel(); + } + } + } + + private synchronized void close() { + cancel(); + if (closed == false) { + closed = true; + source.close(); + } + } + } +} 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 347a5a481d61d..9303988186dfc 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 @@ -14,6 +14,7 @@ import org.opensearch.analytics.spi.QueryExecutionMetrics; import org.opensearch.analytics.spi.ShardSortBounds; import org.opensearch.be.datafusion.NativeErrorConverter; +import org.opensearch.be.datafusion.arrow.ArrowBatchSourceCallbacks; import org.opensearch.be.datafusion.stats.DataFusionStats; import org.opensearch.be.datafusion.stats.NativeExecutorsStats; import org.opensearch.be.datafusion.stats.TaskMonitorStats; @@ -30,6 +31,8 @@ import java.lang.foreign.SymbolLookup; import java.lang.foreign.ValueLayout; import java.lang.invoke.MethodHandle; +import java.lang.invoke.MethodHandles; +import java.lang.invoke.MethodType; import java.util.Collections; import java.util.HashMap; import java.util.LinkedHashMap; @@ -115,9 +118,11 @@ private static RuntimeException rethrowConverted(RuntimeException e) { private static final MethodHandle FREE_METRICS_BUF; private static final MethodHandle SQL_TO_SUBSTRAIT; private static final MethodHandle REGISTER_FILTER_TREE_CALLBACKS; + private static final MethodHandle REGISTER_ARROW_BATCH_SOURCE_CALLBACKS; private static final MethodHandle CREATE_LOCAL_SESSION; private static final MethodHandle CLOSE_LOCAL_SESSION; private static final MethodHandle REGISTER_PARTITION_STREAM; + private static final MethodHandle REGISTER_ARROW_BATCH_SOURCE_PROVIDER; private static final MethodHandle EXECUTE_LOCAL_PLAN; private static final MethodHandle SENDER_SEND; private static final MethodHandle SENDER_TERMINATE_EARLY; @@ -335,6 +340,22 @@ private static RuntimeException rethrowConverted(RuntimeException e) { FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG) ); + // i64 df_register_arrow_batch_source_provider(session_ptr, input_id_ptr, input_id_len, + // schema_ptr, schema_len, binding_id, task_id) + REGISTER_ARROW_BATCH_SOURCE_PROVIDER = linker.downcallHandle( + lib.find("df_register_arrow_batch_source_provider").orElseThrow(), + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG + ) + ); + // i64 df_register_partition_stream(session_ptr, input_id_ptr, input_id_len, // partial_plan_ptr, partial_plan_len, // out_ptr, out_cap, out_len) @@ -408,6 +429,11 @@ private static RuntimeException rethrowConverted(RuntimeException e) { ) ); + REGISTER_ARROW_BATCH_SOURCE_CALLBACKS = linker.downcallHandle( + lib.find("df_register_arrow_batch_source_callbacks").orElseThrow(), + FunctionDescriptor.ofVoid(ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.ADDRESS, ValueLayout.ADDRESS) + ); + // Same signature as df_register_memtable but for SessionContextHandle (shard-scan path). REGISTER_MEMTABLE_ON_SESSION_CONTEXT = linker.downcallHandle( lib.find("df_register_memtable_on_session_context").orElseThrow(), @@ -657,6 +683,7 @@ private static RuntimeException rethrowConverted(RuntimeException e) { // caller step required — as soon as this class is loaded, callbacks // are installed and `df_execute_indexed_query` can dispatch into Java. installFilterTreeCallbacks(linker); + installArrowBatchSourceCallbacks(linker); CLOSE_SESSION_CONTEXT = linker.downcallHandle( lib.find("df_close_session_context").orElseThrow(), @@ -758,6 +785,85 @@ private static RuntimeException rethrowConverted(RuntimeException e) { private NativeBridge() {} + private static void installArrowBatchSourceCallbacks(Linker linker) { + try { + Arena arena = Arena.global(); + MethodHandles.Lookup lookup = MethodHandles.lookup(); + MethodHandle createSource = lookup.findStatic( + ArrowBatchSourceCallbacks.class, + "createSource", + MethodType.methodType(int.class, long.class, MemorySegment.class, long.class, MemorySegment.class, long.class) + ); + MethodHandle nextBatch = lookup.findStatic( + ArrowBatchSourceCallbacks.class, + "nextBatch", + MethodType.methodType( + long.class, + long.class, + int.class, + MemorySegment.class, + MemorySegment.class, + MemorySegment.class, + long.class + ) + ); + MethodHandle cancelSource = lookup.findStatic( + ArrowBatchSourceCallbacks.class, + "cancelSource", + MethodType.methodType(void.class, long.class, int.class) + ); + MethodHandle releaseSource = lookup.findStatic( + ArrowBatchSourceCallbacks.class, + "releaseSource", + MethodType.methodType(void.class, long.class, int.class) + ); + MemorySegment createSourceStub = linker.upcallStub( + createSource, + FunctionDescriptor.of( + ValueLayout.JAVA_INT, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG + ), + arena + ); + MemorySegment nextBatchStub = linker.upcallStub( + nextBatch, + FunctionDescriptor.of( + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_LONG, + ValueLayout.JAVA_INT, + ValueLayout.ADDRESS, + ValueLayout.ADDRESS, + ValueLayout.ADDRESS, + ValueLayout.JAVA_LONG + ), + arena + ); + MemorySegment cancelSourceStub = linker.upcallStub( + cancelSource, + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG, ValueLayout.JAVA_INT), + arena + ); + MemorySegment releaseSourceStub = linker.upcallStub( + releaseSource, + FunctionDescriptor.ofVoid(ValueLayout.JAVA_LONG, ValueLayout.JAVA_INT), + arena + ); + NativeCall.invokeVoid( + REGISTER_ARROW_BATCH_SOURCE_CALLBACKS, + createSourceStub, + nextBatchStub, + cancelSourceStub, + releaseSourceStub + ); + } catch (Throwable throwable) { + throw new ExceptionInInitializerError(throwable); + } + } + private static void installFilterTreeCallbacks(Linker linker) { try { java.lang.foreign.Arena arena = java.lang.foreign.Arena.global(); @@ -1359,8 +1465,8 @@ public static byte[] sqlToSubstrait(long readerPtr, String tableName, String sql // ---- Coordinator-reduce exports ---- /** - * Pair returned from {@link #registerPartitionStream} / {@link #registerMemtable}: the - * native sender pointer (or 0 for memtable) plus the Arrow IPC-encoded schema the native + * Pair returned from {@link #registerPartitionStream} or {@link #registerMemtable}: the + * native sender pointer (or 0 for non-stream inputs) plus the Arrow IPC-encoded schema the native * session derived by lowering the producer-side substrait. The Java tripwire * ({@code typesMatch} in {@code DatafusionReduceSink}) validates fed batches against this * schema, and downstream callers decode it once into an Arrow {@link org.apache.arrow.vector.types.pojo.Schema}. @@ -1384,6 +1490,35 @@ public static void closeLocalSession(long sessionPtr) { NativeCall.invokeVoid(CLOSE_LOCAL_SESSION, sessionPtr); } + /** + * Registers a Java Arrow batch source provider. The binding ID selects the Java factory; + * the task ID is used only for native cancellation and accounting. The production owner + * must back source batches with a breaker-accounted Arrow allocator. Native imports retain + * those externally accounted buffers but do not perform separate coordinator admission. + */ + public static void registerArrowBatchSourceProvider(long sessionPtr, String inputId, byte[] schemaIpc, long bindingId, long taskId) { + NativeHandle.validatePointer(sessionPtr, "session"); + if (schemaIpc.length == 0) { + throw new IllegalArgumentException("schemaIpc must not be empty"); + } + if (bindingId <= 0L) { + throw new IllegalArgumentException("bindingId must be positive"); + } + try (var call = new NativeCall()) { + var id = call.str(inputId); + call.invoke( + REGISTER_ARROW_BATCH_SOURCE_PROVIDER, + sessionPtr, + id.segment(), + id.len(), + call.bytes(schemaIpc), + (long) schemaIpc.length, + bindingId, + taskId + ); + } + } + /** * Registers an input partition stream on the session under {@code inputId}, deriving the * input schema by lowering the producer-side {@code partialPlanBytes}. Returns the native 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 e4fb5ba6b5720..3009d77b30b6b 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 @@ -742,11 +742,56 @@ public void testApproxCountDistinctRenamed() throws Exception { assertTrue("must find approx_distinct in extension declarations", foundApproxDistinct); } + /** PPL's overflow-checking long sum must bind to DataFusion's standard {@code sum} aggregate. */ + public void testCheckedLongSumRewrittenToStandardSum() throws Exception { + RelDataType longRow = typeFactory.builder() + .add("A", typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true)) + .build(); + RelNode scan = new DataFusionFragmentConvertor.StageInputTableScan(cluster, cluster.traitSet(), "test_index", longRow); + SqlAggFunction checkedLongSum = new SqlAggFunction( + "CHECKED_LONG_SUM", + null, + SqlKind.SUM, + ReturnTypes.BIGINT_FORCE_NULLABLE, + null, + OperandTypes.NUMERIC, + SqlFunctionCategory.USER_DEFINED_FUNCTION, + false, + false, + Optionality.FORBIDDEN + ) { + }; + AggregateCall sumCall = AggregateCall.create( + checkedLongSum, + false, + List.of(0), + -1, + typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true), + "sum_col" + ); + LogicalAggregate agg = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(sumCall)); + + Plan plan = decodeSubstrait(newConvertor().convertFragment(agg)); + List aggregateNames = plan.getExtensionsList() + .stream() + .filter(SimpleExtensionDeclaration::hasExtensionFunction) + .map(declaration -> declaration.getExtensionFunction().getName()) + .toList(); + + assertTrue( + "CHECKED_LONG_SUM must emit as standard sum: " + aggregateNames, + aggregateNames.stream().anyMatch(name -> name.equals("sum") || name.startsWith("sum:")) + ); + assertFalse( + "CHECKED_LONG_SUM must not leak into Substrait: " + aggregateNames, + aggregateNames.stream().anyMatch(name -> name.contains("CHECKED_LONG_SUM")) + ); + } + /** * SUM aggregate is not affected by the rename map — its extension function * name remains unchanged. */ - public void testOtherFunctionsNotRenamed() throws Exception { RelNode scan = buildTableScan("test_index", "A"); LogicalAggregate agg = buildSumAggregate(scan, 0); diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutorTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutorTests.java new file mode 100644 index 0000000000000..74da66d6d3c29 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/DatafusionArrowBatchSourceExecutorTests.java @@ -0,0 +1,539 @@ +/* + * 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.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ViewVarCharVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.plan.RelOptCluster; +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.AggregateCall; +import org.apache.calcite.rel.logical.LogicalAggregate; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.ImmutableBitSet; +import org.apache.lucene.analysis.standard.StandardAnalyzer; +import org.apache.lucene.document.BinaryDocValuesField; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.InetAddressPoint; +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.document.SortedDocValuesField; +import org.apache.lucene.document.SortedNumericDocValuesField; +import org.apache.lucene.document.SortedSetDocValuesField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.MatchAllDocsQuery; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.NumericUtils; +import org.opensearch.analytics.backend.EngineResultBatch; +import org.opensearch.analytics.backend.EngineResultStream; +import org.opensearch.analytics.exec.task.AnalyticsShardTask; +import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.ColumnKind; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.InputColumn; +import org.opensearch.analytics.spi.ArrowBatchSourcePlan; +import org.opensearch.be.lucene.DocValuesBatchSourceFactory; +import org.opensearch.core.tasks.TaskCancelledException; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.test.OpenSearchTestCase; + +import java.net.InetAddress; +import java.nio.file.Path; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.Map; + +import io.substrait.extension.DefaultExtensionCatalog; + +/** End-to-end Lucene doc-values to DataFusion Arrow source execution test. */ +public class DatafusionArrowBatchSourceExecutorTests extends OpenSearchTestCase { + + public void testCompilesNamedArrowSourcePlan() throws Exception { + byte[] planBytes = convert(inputScan("input-0")); + + io.substrait.proto.Plan plan = io.substrait.proto.Plan.parseFrom(planBytes); + assertEquals("input-0", plan.getRelations(0).getRoot().getInput().getRead().getNamedTable().getNames(0)); + } + + public void testSetupFailureClosesTransferredFactory() throws Exception { + try ( + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer())) + ) { + writer.commit(); + try (DirectoryReader reader = DirectoryReader.open(writer)) { + int initialRefCount = reader.getRefCount(); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("x", ColumnKind.LONG)), + allocator, + null + ); + ArrowBatchSourcePlan plan = new ArrowBatchSourcePlan( + "input-0", + new byte[] { 1 }, + List.of(new InputColumn("x", ColumnKind.LONG)) + ); + DatafusionArrowBatchSourceExecutor executor = new DatafusionArrowBatchSourceExecutor(DataFusionService.builder().build()); + + expectThrows(IllegalStateException.class, () -> executor.execute(allocator, plan, factory, null, null)); + assertEquals(initialRefCount, reader.getRefCount()); + expectThrows(IllegalStateException.class, () -> factory.open(new int[] { 0 })); + } + } + } + + public void testCancelledBeforeSetupClosesTransferredFactory() throws Exception { + try ( + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer())) + ) { + writer.commit(); + try (DirectoryReader reader = DirectoryReader.open(writer)) { + int initialRefCount = reader.getRefCount(); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("x", ColumnKind.LONG)), + allocator, + null + ); + AnalyticsShardTask task = new AnalyticsShardTask( + 98_002L, + "test", + "arrow-source", + "cancelled", + TaskId.EMPTY_TASK_ID, + Map.of() + ); + task.cancel("test cancellation"); + ArrowBatchSourcePlan plan = new ArrowBatchSourcePlan( + "input-0", + new byte[] { 1 }, + List.of(new InputColumn("x", ColumnKind.LONG)) + ); + DatafusionArrowBatchSourceExecutor executor = new DatafusionArrowBatchSourceExecutor(DataFusionService.builder().build()); + + expectThrows(TaskCancelledException.class, () -> executor.execute(allocator, plan, factory, task, null)); + assertEquals(initialRefCount, reader.getRefCount()); + expectThrows(IllegalStateException.class, () -> factory.open(new int[] { 0 })); + } + } + } + + public void testEarlyResultStreamCloseReleasesLuceneFactory() throws Exception { + Path spillDirectory = createTempDir("arrow-source-early-close-spill"); + DataFusionService service = DataFusionService.builder() + .memoryPoolLimit(64L * 1024L * 1024L) + .spillMemoryLimit(32L * 1024L * 1024L) + .spillDirectory(spillDirectory.toString()) + .cpuThreads(2) + .build(); + service.start(); + DatafusionArrowBatchSourceExecutor executor = new DatafusionArrowBatchSourceExecutor(service); + + try ( + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer())) + ) { + addDocument(writer, 11L, "a"); + writer.commit(); + try (DirectoryReader reader = DirectoryReader.open(writer)) { + int initialRefCount = reader.getRefCount(); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("x", ColumnKind.LONG)), + allocator, + null + ); + ArrowBatchSourcePlan plan = new ArrowBatchSourcePlan( + "input-0", + convert(inputScan("input-0")), + List.of(new InputColumn("x", ColumnKind.LONG)) + ); + EngineResultStream stream = executor.execute(allocator, plan, factory, null, null); + assertTrue("factory and native source hold reader references", reader.getRefCount() > initialRefCount); + + stream.close(); + stream.close(); + + assertBusy(() -> assertEquals(initialRefCount, reader.getRefCount())); + expectThrows(IllegalStateException.class, () -> factory.open(new int[] { 0 })); + } + } finally { + service.close(); + } + } + + public void testExecutesLuceneDocValuesThroughDataFusion() throws Exception { + Path spillDirectory = createTempDir("arrow-source-spill"); + DataFusionService service = DataFusionService.builder() + .memoryPoolLimit(64L * 1024L * 1024L) + .spillMemoryLimit(32L * 1024L * 1024L) + .spillDirectory(spillDirectory.toString()) + .cpuThreads(2) + .build(); + service.start(); + DatafusionArrowBatchSourceExecutor executor = new DatafusionArrowBatchSourceExecutor(service); + + try ( + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer())) + ) { + addDocument(writer, 11L, "a"); + addDocument(writer, 22L, "b"); + addDocument(writer, null, "missing"); + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(writer)) { + int initialRefCount = reader.getRefCount(); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("x", ColumnKind.LONG)), + allocator, + null + ); + ArrowBatchSourcePlan plan = new ArrowBatchSourcePlan( + "input-0", + convert(filteredInput("input-0", 15L)), + List.of(new InputColumn("x", ColumnKind.LONG)) + ); + AnalyticsShardTask task = new AnalyticsShardTask( + 98_001L, + "test", + "arrow-source", + "arrow-source", + TaskId.EMPTY_TASK_ID, + Collections.emptyMap() + ); + + try (EngineResultStream stream = executor.execute(allocator, plan, factory, task, null)) { + Iterator batches = stream.iterator(); + assertTrue(batches.hasNext()); + EngineResultBatch batch = batches.next(); + try (VectorSchemaRoot root = batch.getArrowRoot()) { + assertEquals(1, root.getRowCount()); + BigIntVector values = (BigIntVector) root.getVector("x"); + assertEquals(22L, values.get(0)); + } + assertFalse(batches.hasNext()); + } + + expectThrows(IllegalStateException.class, () -> factory.open(new int[] { 0 })); + assertEquals(initialRefCount, reader.getRefCount()); + + DocValuesBatchSourceFactory countFactory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("x", ColumnKind.LONG)), + allocator, + null + ); + ArrowBatchSourcePlan countPlan = new ArrowBatchSourcePlan( + "input-0", + countSubstrait("input-0", true), + List.of(new InputColumn("x", ColumnKind.LONG)) + ); + try (EngineResultStream stream = executor.execute(allocator, countPlan, countFactory, task, null)) { + Iterator batches = stream.iterator(); + assertTrue(batches.hasNext()); + try (VectorSchemaRoot root = batches.next().getArrowRoot()) { + assertEquals(1, root.getRowCount()); + assertEquals(1, root.getFieldVectors().size()); + assertEquals(3L, ((BigIntVector) root.getVector(0)).get(0)); + } + assertFalse(batches.hasNext()); + } + expectThrows(IllegalStateException.class, () -> countFactory.open(new int[0])); + assertEquals(initialRefCount, reader.getRefCount()); + + DocValuesBatchSourceFactory fieldCountFactory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("x", ColumnKind.LONG)), + allocator, + null + ); + ArrowBatchSourcePlan fieldCountPlan = new ArrowBatchSourcePlan( + "input-0", + countSubstrait("input-0", false), + List.of(new InputColumn("x", ColumnKind.LONG)) + ); + try (EngineResultStream stream = executor.execute(allocator, fieldCountPlan, fieldCountFactory, task, null)) { + Iterator batches = stream.iterator(); + assertTrue(batches.hasNext()); + try (VectorSchemaRoot root = batches.next().getArrowRoot()) { + assertEquals(2L, ((BigIntVector) root.getVector(0)).get(0)); + } + assertFalse(batches.hasNext()); + } + expectThrows(IllegalStateException.class, () -> fieldCountFactory.open(new int[] { 0 })); + assertEquals(initialRefCount, reader.getRefCount()); + + DocValuesBatchSourceFactory keywordFactory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("keyword", ColumnKind.KEYWORD)), + allocator, + null + ); + ArrowBatchSourcePlan keywordPlan = new ArrowBatchSourcePlan( + "input-0", + convert(inputScan("input-0", "keyword", SqlTypeName.VARCHAR)), + List.of(new InputColumn("keyword", ColumnKind.KEYWORD)) + ); + try (EngineResultStream stream = executor.execute(allocator, keywordPlan, keywordFactory, task, null)) { + Iterator batches = stream.iterator(); + assertTrue(batches.hasNext()); + try (VectorSchemaRoot root = batches.next().getArrowRoot()) { + assertEquals(3, root.getRowCount()); + assertEquals("a", root.getVector("keyword").getObject(0).toString()); + assertEquals("b", root.getVector("keyword").getObject(1).toString()); + assertEquals("missing", root.getVector("keyword").getObject(2).toString()); + } + assertFalse(batches.hasNext()); + } + expectThrows(IllegalStateException.class, () -> keywordFactory.open(new int[] { 0 })); + assertEquals(initialRefCount, reader.getRefCount()); + } + } finally { + service.close(); + } + } + + public void testExecutesAdditionalDocValueTypesThroughDataFusion() throws Exception { + Path spillDirectory = createTempDir("arrow-source-types-spill"); + DataFusionService service = DataFusionService.builder() + .memoryPoolLimit(64L * 1024L * 1024L) + .spillMemoryLimit(32L * 1024L * 1024L) + .spillDirectory(spillDirectory.toString()) + .cpuThreads(2) + .build(); + service.start(); + DatafusionArrowBatchSourceExecutor executor = new DatafusionArrowBatchSourceExecutor(service); + byte[] binary = new byte[] { 0, 1, (byte) 0xFF }; + byte[] firstIp = InetAddressPoint.encode(InetAddress.getByName("192.0.2.1")); + byte[] secondIp = InetAddressPoint.encode(InetAddress.getByName("2001:db8::1")); + + try ( + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer())) + ) { + Document first = new Document(); + first.add(new NumericDocValuesField("boolean", 1L)); + first.add(new NumericDocValuesField("float", NumericUtils.floatToSortableInt(1.25F))); + first.add(new NumericDocValuesField("double", NumericUtils.doubleToSortableLong(-2.5D))); + first.add(new BinaryDocValuesField("binary", new BytesRef(binary))); + first.add(new SortedSetDocValuesField("ip", new BytesRef(firstIp))); + first.add(new SortedNumericDocValuesField("longs", 1L)); + first.add(new SortedNumericDocValuesField("longs", 2L)); + first.add(new SortedNumericDocValuesField("floats", NumericUtils.floatToSortableInt(-1.5F))); + first.add(new SortedNumericDocValuesField("floats", NumericUtils.floatToSortableInt(2.25F))); + first.add(new SortedNumericDocValuesField("doubles", NumericUtils.doubleToSortableLong(-3.5D))); + first.add(new SortedNumericDocValuesField("doubles", NumericUtils.doubleToSortableLong(4.75D))); + first.add(new SortedSetDocValuesField("keywords", new BytesRef("a"))); + first.add(new SortedSetDocValuesField("keywords", new BytesRef("b"))); + first.add(new SortedSetDocValuesField("ips", new BytesRef(firstIp))); + first.add(new SortedSetDocValuesField("ips", new BytesRef(secondIp))); + writer.addDocument(first); + writer.addDocument(new Document()); + writer.commit(); + + List columns = List.of( + new InputColumn("boolean", ColumnKind.BOOLEAN), + new InputColumn("float", ColumnKind.FLOAT), + new InputColumn("double", ColumnKind.DOUBLE), + new InputColumn("binary", ColumnKind.BINARY), + new InputColumn("ip", ColumnKind.IP), + new InputColumn("longs", ColumnKind.LONG, true), + new InputColumn("floats", ColumnKind.FLOAT, true), + new InputColumn("doubles", ColumnKind.DOUBLE, true), + new InputColumn("keywords", ColumnKind.KEYWORD, true), + new InputColumn("ips", ColumnKind.IP, true) + ); + try (DirectoryReader reader = DirectoryReader.open(writer)) { + int initialRefCount = reader.getRefCount(); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + columns, + allocator, + null + ); + ArrowBatchSourcePlan plan = new ArrowBatchSourcePlan("input-0", convert(inputScan("input-0", columns)), columns); + AnalyticsShardTask task = new AnalyticsShardTask( + 98_003L, + "test", + "arrow-source", + "arrow-source-types", + TaskId.EMPTY_TASK_ID, + Map.of() + ); + + try (EngineResultStream stream = executor.execute(allocator, plan, factory, task, null)) { + Iterator batches = stream.iterator(); + assertTrue(batches.hasNext()); + try (VectorSchemaRoot root = batches.next().getArrowRoot()) { + assertEquals(2, root.getRowCount()); + assertEquals(1, ((BitVector) root.getVector("boolean")).get(0)); + assertEquals(1.25F, ((Float4Vector) root.getVector("float")).get(0), 0F); + assertEquals(-2.5D, ((Float8Vector) root.getVector("double")).get(0), 0D); + assertArrayEquals(binary, ((VarBinaryVector) root.getVector("binary")).get(0)); + assertArrayEquals(firstIp, ((VarBinaryVector) root.getVector("ip")).get(0)); + + ListVector longs = (ListVector) root.getVector("longs"); + int longStart = longs.getElementStartIndex(0); + assertEquals(1L, ((BigIntVector) longs.getDataVector()).get(longStart)); + assertEquals(2L, ((BigIntVector) longs.getDataVector()).get(longStart + 1)); + + ListVector floats = (ListVector) root.getVector("floats"); + int floatStart = floats.getElementStartIndex(0); + assertEquals(-1.5F, ((Float4Vector) floats.getDataVector()).get(floatStart), 0F); + assertEquals(2.25F, ((Float4Vector) floats.getDataVector()).get(floatStart + 1), 0F); + + ListVector doubles = (ListVector) root.getVector("doubles"); + int doubleStart = doubles.getElementStartIndex(0); + assertEquals(-3.5D, ((Float8Vector) doubles.getDataVector()).get(doubleStart), 0D); + assertEquals(4.75D, ((Float8Vector) doubles.getDataVector()).get(doubleStart + 1), 0D); + + ListVector keywords = (ListVector) root.getVector("keywords"); + int keywordStart = keywords.getElementStartIndex(0); + ViewVarCharVector keywordValues = (ViewVarCharVector) keywords.getDataVector(); + assertEquals("a", keywordValues.getObject(keywordStart).toString()); + assertEquals("b", keywordValues.getObject(keywordStart + 1).toString()); + + ListVector ips = (ListVector) root.getVector("ips"); + int ipStart = ips.getElementStartIndex(0); + VarBinaryVector ipValues = (VarBinaryVector) ips.getDataVector(); + assertArrayEquals(firstIp, ipValues.get(ipStart)); + assertArrayEquals(secondIp, ipValues.get(ipStart + 1)); + for (InputColumn column : columns) { + assertTrue(root.getVector(column.name()).isNull(1)); + } + } + assertFalse(batches.hasNext()); + } + assertEquals(initialRefCount, reader.getRefCount()); + } + } finally { + service.close(); + } + } + + private static void addDocument(IndexWriter writer, Long value, String keyword) throws Exception { + Document document = new Document(); + if (value != null) { + document.add(new NumericDocValuesField("x", value)); + } + document.add(new SortedDocValuesField("keyword", new BytesRef(keyword))); + writer.addDocument(document); + } + + private static byte[] countSubstrait(String inputId, boolean countStar) { + RelNode scan = inputScan(inputId); + RelDataTypeFactory typeFactory = scan.getCluster().getTypeFactory(); + AggregateCall count = AggregateCall.create( + SqlStdOperatorTable.COUNT, + false, + countStar ? List.of() : List.of(0), + -1, + typeFactory.createSqlType(SqlTypeName.BIGINT), + "count" + ); + return convert(LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(count))); + } + + private static RelNode filteredInput(String inputId, long lowerBound) { + RelNode scan = inputScan(inputId); + RexBuilder rexBuilder = scan.getCluster().getRexBuilder(); + return org.apache.calcite.rel.logical.LogicalFilter.create( + scan, + rexBuilder.makeCall( + SqlStdOperatorTable.GREATER_THAN, + rexBuilder.makeInputRef(scan, 0), + rexBuilder.makeBigintLiteral(java.math.BigDecimal.valueOf(lowerBound)) + ) + ); + } + + private static RelNode inputScan(String inputId) { + return inputScan(inputId, "x", SqlTypeName.BIGINT); + } + + private static RelNode inputScan(String inputId, List columns) { + RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + HepPlanner planner = new HepPlanner(new HepProgramBuilder().build()); + RelOptCluster cluster = RelOptCluster.create(planner, rexBuilder); + RelDataTypeFactory.Builder fields = typeFactory.builder(); + for (InputColumn column : columns) { + SqlTypeName sqlType = switch (column.kind()) { + case LONG -> SqlTypeName.BIGINT; + case KEYWORD -> SqlTypeName.VARCHAR; + case TIMESTAMP -> SqlTypeName.TIMESTAMP; + case BOOLEAN -> SqlTypeName.BOOLEAN; + case FLOAT -> SqlTypeName.REAL; + case DOUBLE -> SqlTypeName.DOUBLE; + case BINARY, IP -> SqlTypeName.VARBINARY; + }; + RelDataType fieldType = typeFactory.createTypeWithNullability(typeFactory.createSqlType(sqlType), true); + if (column.multiValued()) { + fieldType = typeFactory.createTypeWithNullability(typeFactory.createArrayType(fieldType, -1), true); + } + fields.add(column.name(), fieldType); + } + int childStageId = Integer.parseInt(inputId.substring("input-".length())); + return new OpenSearchStageInputScan(cluster, cluster.traitSet(), childStageId, fields.build(), List.of(), List.of()); + } + + private static RelNode inputScan(String inputId, String fieldName, SqlTypeName type) { + RelDataTypeFactory typeFactory = new JavaTypeFactoryImpl(); + RexBuilder rexBuilder = new RexBuilder(typeFactory); + HepPlanner planner = new HepPlanner(new HepProgramBuilder().build()); + RelOptCluster cluster = RelOptCluster.create(planner, rexBuilder); + RelDataType fieldType = typeFactory.createTypeWithNullability(typeFactory.createSqlType(type), true); + RelDataType rowType = typeFactory.builder().add(fieldName, fieldType).build(); + int childStageId = Integer.parseInt(inputId.substring("input-".length())); + return new OpenSearchStageInputScan(cluster, cluster.traitSet(), childStageId, rowType, List.of(), List.of()); + } + + private static byte[] convert(RelNode node) { + Thread thread = Thread.currentThread(); + ClassLoader previous = thread.getContextClassLoader(); + try { + thread.setContextClassLoader(DatafusionArrowBatchSourceExecutorTests.class.getClassLoader()); + return new DataFusionFragmentConvertor(DefaultExtensionCatalog.DEFAULT_COLLECTION).convertFragment(node); + } finally { + thread.setContextClassLoader(previous); + } + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/LuceneDocValuesBatchSourceTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/LuceneDocValuesBatchSourceTests.java new file mode 100644 index 0000000000000..4ea243bb225be --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/LuceneDocValuesBatchSourceTests.java @@ -0,0 +1,398 @@ +/* + * 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.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.TimeStampMilliVector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ViewVarCharVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.lucene.document.BinaryDocValuesField; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.InetAddressPoint; +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.document.SortedNumericDocValuesField; +import org.apache.lucene.document.SortedSetDocValuesField; +import org.apache.lucene.document.TextField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.MatchAllDocsQuery; +import org.apache.lucene.search.PhraseQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.NumericUtils; +import org.opensearch.analytics.exec.task.AnalyticsShardTask; +import org.opensearch.analytics.spi.ArrowBatchSource; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.ColumnKind; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.InputColumn; +import org.opensearch.be.lucene.DocValuesBatchSourceFactory; +import org.opensearch.core.tasks.TaskCancelledException; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.test.OpenSearchTestCase; + +import java.net.InetAddress; +import java.util.List; +import java.util.Map; + +public class LuceneDocValuesBatchSourceTests extends OpenSearchTestCase { + + public void testProjectionNullsAndEof() throws Exception { + try (Directory directory = newDirectory(); IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + Document first = new Document(); + first.add(new SortedNumericDocValuesField("n", 10)); + first.add(new SortedSetDocValuesField("k", new BytesRef("a"))); + writer.addDocument(first); + + Document second = new Document(); + second.add(new SortedSetDocValuesField("k", new BytesRef("b"))); + writer.addDocument(second); + + Document third = new Document(); + third.add(new SortedNumericDocValuesField("n", 30)); + writer.addDocument(third); + writer.commit(); + + try ( + DirectoryReader reader = DirectoryReader.open(writer); + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("n", ColumnKind.LONG), new InputColumn("k", ColumnKind.KEYWORD)), + allocator, + null + ) + ) { + try (ArrowBatchSource source = factory.open(new int[] { 0, 1 }); VectorSchemaRoot root = source.nextBatch()) { + assertEquals(3, root.getRowCount()); + assertEquals(10L, root.getVector("n").getObject(0)); + assertEquals("a", root.getVector("k").getObject(0).toString()); + assertNull(root.getVector("n").getObject(1)); + assertEquals("b", root.getVector("k").getObject(1).toString()); + assertEquals(30L, root.getVector("n").getObject(2)); + assertNull(root.getVector("k").getObject(2)); + assertNull(source.nextBatch()); + } + try (ArrowBatchSource source = factory.open(new int[] { 1 }); VectorSchemaRoot root = source.nextBatch()) { + assertEquals(1, root.getFieldVectors().size()); + assertEquals("k", root.getVector(0).getName()); + assertEquals(3, root.getRowCount()); + assertNull(source.nextBatch()); + } + IllegalArgumentException error = expectThrows(IllegalArgumentException.class, () -> factory.open(new int[] { 2 })); + assertTrue(error.getMessage(), error.getMessage().contains("outside input schema")); + } + } + } + + public void testScalarFieldTypes() throws Exception { + byte[] binary = new byte[] { 0, 1, (byte) 0xFF }; + byte[] ip = InetAddressPoint.encode(InetAddress.getByName("192.0.2.1")); + try (Directory directory = newDirectory(); IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + Document first = new Document(); + first.add(new NumericDocValuesField("boolean", 1L)); + first.add(new NumericDocValuesField("float", NumericUtils.floatToSortableInt(1.25F))); + first.add(new NumericDocValuesField("double", NumericUtils.doubleToSortableLong(-2.5D))); + first.add(new BinaryDocValuesField("binary", new BytesRef(binary))); + first.add(new SortedSetDocValuesField("ip", new BytesRef(ip))); + writer.addDocument(first); + writer.addDocument(new Document()); + writer.commit(); + + List columns = List.of( + new InputColumn("boolean", ColumnKind.BOOLEAN), + new InputColumn("float", ColumnKind.FLOAT), + new InputColumn("double", ColumnKind.DOUBLE), + new InputColumn("binary", ColumnKind.BINARY), + new InputColumn("ip", ColumnKind.IP) + ); + try ( + DirectoryReader reader = DirectoryReader.open(writer); + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + columns, + allocator, + null + ); + ArrowBatchSource source = factory.open(new int[] { 0, 1, 2, 3, 4 }); + VectorSchemaRoot root = source.nextBatch() + ) { + assertEquals(2, root.getRowCount()); + assertEquals(true, root.getVector("boolean").getObject(0)); + assertEquals(1.25F, (Float) root.getVector("float").getObject(0), 0F); + assertEquals(-2.5D, (Double) root.getVector("double").getObject(0), 0D); + assertArrayEquals(binary, (byte[]) root.getVector("binary").getObject(0)); + assertArrayEquals(ip, (byte[]) root.getVector("ip").getObject(0)); + for (InputColumn column : columns) { + assertNull(root.getVector(column.name()).getObject(1)); + } + assertNull(source.nextBatch()); + } + } + } + + public void testMultiValuedFieldTypes() throws Exception { + byte[] firstIp = InetAddressPoint.encode(InetAddress.getByName("192.0.2.1")); + byte[] secondIp = InetAddressPoint.encode(InetAddress.getByName("2001:db8::1")); + try (Directory directory = newDirectory(); IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + Document first = new Document(); + first.add(new SortedNumericDocValuesField("longs", 1L)); + first.add(new SortedNumericDocValuesField("longs", 2L)); + first.add(new SortedNumericDocValuesField("timestamps", 1_000L)); + first.add(new SortedNumericDocValuesField("timestamps", 2_000L)); + first.add(new SortedNumericDocValuesField("booleans", 0L)); + first.add(new SortedNumericDocValuesField("booleans", 1L)); + first.add(new SortedNumericDocValuesField("floats", NumericUtils.floatToSortableInt(-1.5F))); + first.add(new SortedNumericDocValuesField("floats", NumericUtils.floatToSortableInt(2.25F))); + first.add(new SortedNumericDocValuesField("doubles", NumericUtils.doubleToSortableLong(-3.5D))); + first.add(new SortedNumericDocValuesField("doubles", NumericUtils.doubleToSortableLong(4.75D))); + first.add(new SortedSetDocValuesField("keywords", new BytesRef("a"))); + first.add(new SortedSetDocValuesField("keywords", new BytesRef("b"))); + first.add(new SortedSetDocValuesField("ips", new BytesRef(firstIp))); + first.add(new SortedSetDocValuesField("ips", new BytesRef(secondIp))); + writer.addDocument(first); + writer.addDocument(new Document()); + writer.commit(); + + List columns = List.of( + new InputColumn("longs", ColumnKind.LONG, true), + new InputColumn("timestamps", ColumnKind.TIMESTAMP, true), + new InputColumn("booleans", ColumnKind.BOOLEAN, true), + new InputColumn("floats", ColumnKind.FLOAT, true), + new InputColumn("doubles", ColumnKind.DOUBLE, true), + new InputColumn("keywords", ColumnKind.KEYWORD, true), + new InputColumn("ips", ColumnKind.IP, true) + ); + try ( + DirectoryReader reader = DirectoryReader.open(writer); + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + columns, + allocator, + null + ); + ArrowBatchSource source = factory.open(new int[] { 0, 1, 2, 3, 4, 5, 6 }); + VectorSchemaRoot root = source.nextBatch() + ) { + ListVector longs = (ListVector) root.getVector("longs"); + int longStart = longs.getElementStartIndex(0); + assertEquals(longStart + 2, longs.getElementEndIndex(0)); + assertEquals(1L, ((BigIntVector) longs.getDataVector()).get(longStart)); + assertEquals(2L, ((BigIntVector) longs.getDataVector()).get(longStart + 1)); + + ListVector timestamps = (ListVector) root.getVector("timestamps"); + int timestampStart = timestamps.getElementStartIndex(0); + assertEquals(1_000L, ((TimeStampMilliVector) timestamps.getDataVector()).get(timestampStart)); + assertEquals(2_000L, ((TimeStampMilliVector) timestamps.getDataVector()).get(timestampStart + 1)); + + ListVector booleans = (ListVector) root.getVector("booleans"); + int booleanStart = booleans.getElementStartIndex(0); + assertEquals(0, ((BitVector) booleans.getDataVector()).get(booleanStart)); + assertEquals(1, ((BitVector) booleans.getDataVector()).get(booleanStart + 1)); + + ListVector floats = (ListVector) root.getVector("floats"); + int floatStart = floats.getElementStartIndex(0); + assertEquals(-1.5F, ((Float4Vector) floats.getDataVector()).get(floatStart), 0F); + assertEquals(2.25F, ((Float4Vector) floats.getDataVector()).get(floatStart + 1), 0F); + + ListVector doubles = (ListVector) root.getVector("doubles"); + int doubleStart = doubles.getElementStartIndex(0); + assertEquals(-3.5D, ((Float8Vector) doubles.getDataVector()).get(doubleStart), 0D); + assertEquals(4.75D, ((Float8Vector) doubles.getDataVector()).get(doubleStart + 1), 0D); + + ListVector keywords = (ListVector) root.getVector("keywords"); + int keywordStart = keywords.getElementStartIndex(0); + ViewVarCharVector keywordValues = (ViewVarCharVector) keywords.getDataVector(); + assertEquals("a", keywordValues.getObject(keywordStart).toString()); + assertEquals("b", keywordValues.getObject(keywordStart + 1).toString()); + + ListVector ips = (ListVector) root.getVector("ips"); + int ipStart = ips.getElementStartIndex(0); + VarBinaryVector ipValues = (VarBinaryVector) ips.getDataVector(); + assertArrayEquals(firstIp, ipValues.get(ipStart)); + assertArrayEquals(secondIp, ipValues.get(ipStart + 1)); + + for (InputColumn column : columns) { + assertTrue(((ListVector) root.getVector(column.name())).isNull(1)); + } + assertNull(source.nextBatch()); + } + } + } + + public void testTwoPhaseQueryUsesConfirmedMatches() throws Exception { + try ( + Directory directory = newDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig()); + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE) + ) { + Document match = new Document(); + match.add(new TextField("text", "a b", Field.Store.NO)); + match.add(new NumericDocValuesField("n", 1)); + writer.addDocument(match); + + Document approximationOnly = new Document(); + approximationOnly.add(new TextField("text", "a x b", Field.Store.NO)); + approximationOnly.add(new NumericDocValuesField("n", 2)); + writer.addDocument(approximationOnly); + writer.commit(); + + try ( + DirectoryReader reader = DirectoryReader.open(writer); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new PhraseQuery("text", "a", "b"), + List.of(new InputColumn("n", ColumnKind.LONG)), + allocator, + null + ); + ArrowBatchSource source = factory.open(new int[] { 0 }); + VectorSchemaRoot root = source.nextBatch() + ) { + assertEquals(1, root.getRowCount()); + assertEquals(1L, root.getVector("n").getObject(0)); + assertNull(source.nextBatch()); + } + } + } + + public void testFactoryRetainsReader() throws Exception { + try ( + Directory directory = newDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig()); + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE) + ) { + Document document = new Document(); + document.add(new NumericDocValuesField("n", 42)); + writer.addDocument(document); + writer.commit(); + + DirectoryReader reader = DirectoryReader.open(writer); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("n", ColumnKind.LONG)), + allocator, + null + ); + reader.close(); + try (factory; ArrowBatchSource source = factory.open(new int[] { 0 }); VectorSchemaRoot root = source.nextBatch()) { + assertEquals(42L, root.getVector("n").getObject(0)); + } + } + } + + public void testRejectsMultiValuedColumnsDeclaredScalar() throws Exception { + try ( + Directory directory = newDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig()); + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE) + ) { + Document document = new Document(); + document.add(new SortedNumericDocValuesField("numbers", 1)); + document.add(new SortedNumericDocValuesField("numbers", 2)); + document.add(new SortedSetDocValuesField("keywords", new BytesRef("a"))); + document.add(new SortedSetDocValuesField("keywords", new BytesRef("b"))); + writer.addDocument(document); + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(writer)) { + assertMultiValuedRejected(reader, allocator, new InputColumn("numbers", ColumnKind.LONG), "numeric"); + assertMultiValuedRejected(reader, allocator, new InputColumn("keywords", ColumnKind.KEYWORD), "keyword"); + } + } + } + + public void testCancelledTaskDoesNotAllocateBatch() throws Exception { + try ( + Directory directory = newDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig()); + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE) + ) { + Document document = new Document(); + document.add(new NumericDocValuesField("n", 1)); + writer.addDocument(document); + writer.commit(); + + AnalyticsShardTask task = new AnalyticsShardTask(1, "test", "test", "test", TaskId.EMPTY_TASK_ID, Map.of()); + task.cancel("test cancellation"); + try ( + DirectoryReader reader = DirectoryReader.open(writer); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("n", ColumnKind.LONG)), + allocator, + task + ); + ArrowBatchSource source = factory.open(new int[] { 0 }) + ) { + expectThrows(TaskCancelledException.class, source::nextBatch); + assertEquals(0L, allocator.getAllocatedMemory()); + } + } + } + + public void testCooperativeCancellation() throws Exception { + try ( + Directory directory = newDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig()); + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE) + ) { + Document document = new Document(); + document.add(new NumericDocValuesField("n", 1)); + writer.addDocument(document); + writer.commit(); + + try ( + DirectoryReader reader = DirectoryReader.open(writer); + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(new InputColumn("n", ColumnKind.LONG)), + allocator, + null + ); + ArrowBatchSource source = factory.open(new int[] { 0 }) + ) { + source.cancel(); + expectThrows(TaskCancelledException.class, source::nextBatch); + } + } + } + + private static void assertMultiValuedRejected(DirectoryReader reader, RootAllocator allocator, InputColumn column, String kind) + throws Exception { + try ( + DocValuesBatchSourceFactory factory = new DocValuesBatchSourceFactory( + new IndexSearcher(reader), + new MatchAllDocsQuery(), + List.of(column), + allocator, + null + ); + ArrowBatchSource source = factory.open(new int[] { 0 }) + ) { + IllegalArgumentException error = expectThrows(IllegalArgumentException.class, source::nextBatch); + assertTrue(error.getMessage(), error.getMessage().contains("multi-valued " + kind + " doc values")); + } + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacksTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacksTests.java new file mode 100644 index 0000000000000..a914cfc0c5650 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/datafusion/arrow/ArrowBatchSourceCallbacksTests.java @@ -0,0 +1,301 @@ +/* + * 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.arrow; + +import org.apache.arrow.c.ArrowArray; +import org.apache.arrow.c.ArrowSchema; +import org.apache.arrow.c.CDataDictionaryProvider; +import org.apache.arrow.c.Data; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.document.SortedDocValuesField; +import org.apache.lucene.document.SortedNumericDocValuesField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.MatchAllDocsQuery; +import org.apache.lucene.store.Directory; +import org.apache.lucene.util.BytesRef; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.ColumnKind; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.InputColumn; +import org.opensearch.be.lucene.DocValuesBatchSourceFactory; +import org.opensearch.test.OpenSearchTestCase; + +import java.lang.foreign.Arena; +import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import java.nio.charset.StandardCharsets; +import java.util.Arrays; +import java.util.List; + +public class ArrowBatchSourceCallbacksTests extends OpenSearchTestCase { + + private RootAllocator allocator; + + @Override + public void setUp() throws Exception { + super.setUp(); + allocator = new RootAllocator(Long.MAX_VALUE); + } + + @Override + public void tearDown() throws Exception { + allocator.close(); + super.tearDown(); + } + + public void testProjectionBatchExportEofAndRelease() throws Exception { + try (Directory directory = newDirectory(); IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + Document first = new Document(); + first.add(new NumericDocValuesField("number", 11L)); + first.add(new SortedDocValuesField("keyword", new BytesRef("a"))); + writer.addDocument(first); + Document second = new Document(); + second.add(new NumericDocValuesField("number", 22L)); + second.add(new SortedDocValuesField("keyword", new BytesRef("b"))); + writer.addDocument(second); + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(writer)) { + int initialRefCount = reader.getRefCount(); + DocValuesBatchSourceFactory factory = factory( + reader, + List.of(new InputColumn("number", ColumnKind.LONG), new InputColumn("keyword", ColumnKind.KEYWORD)) + ); + try ( + ArrowBatchSourceCallbacks.Registration registration = ArrowBatchSourceCallbacks.register(factory, null); + Arena arena = Arena.ofConfined(); + ArrowArray array = ArrowArray.allocateNew(allocator); + ArrowSchema cSchema = ArrowSchema.allocateNew(allocator) + ) { + MemorySegment error = arena.allocate(256L); + int sourceKey = ArrowBatchSourceCallbacks.createSource( + registration.bindingId(), + MemorySegment.ofArray(new int[] { 1, 0 }), + 2L, + error, + error.byteSize() + ); + assertTrue(sourceKey > 0); + assertEquals( + 2L, + ArrowBatchSourceCallbacks.nextBatch( + registration.bindingId(), + sourceKey, + MemorySegment.ofAddress(array.memoryAddress()), + MemorySegment.ofAddress(cSchema.memoryAddress()), + error, + error.byteSize() + ) + ); + try (CDataDictionaryProvider dictionaries = new CDataDictionaryProvider()) { + Schema schema = Data.importSchema(allocator, cSchema, dictionaries); + try (VectorSchemaRoot imported = VectorSchemaRoot.create(schema, allocator)) { + Data.importIntoVectorSchemaRoot(allocator, array, imported, dictionaries); + assertEquals( + List.of("keyword", "number"), + imported.getSchema().getFields().stream().map(f -> f.getName()).toList() + ); + assertEquals("a", imported.getVector("keyword").getObject(0).toString()); + assertEquals("b", imported.getVector("keyword").getObject(1).toString()); + assertEquals(11L, imported.getVector("number").getObject(0)); + assertEquals(22L, imported.getVector("number").getObject(1)); + } + } + assertEquals( + 0L, + ArrowBatchSourceCallbacks.nextBatch( + registration.bindingId(), + sourceKey, + MemorySegment.NULL, + MemorySegment.NULL, + error, + error.byteSize() + ) + ); + ArrowBatchSourceCallbacks.releaseSource(registration.bindingId(), sourceKey); + } + assertEquals(initialRefCount, reader.getRefCount()); + expectThrows(IllegalStateException.class, () -> factory.open(new int[] { 0 })); + } + } + } + + public void testInvalidProjectionReturnsTerminatedError() throws Exception { + try (Directory directory = newDirectory(); IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + writer.commit(); + try (DirectoryReader reader = DirectoryReader.open(writer)) { + int initialRefCount = reader.getRefCount(); + DocValuesBatchSourceFactory factory = factory(reader, List.of(new InputColumn("number", ColumnKind.LONG))); + try ( + ArrowBatchSourceCallbacks.Registration registration = ArrowBatchSourceCallbacks.register(factory, null); + Arena arena = Arena.ofConfined() + ) { + MemorySegment error = arena.allocate(32L); + error.fill((byte) 'x'); + assertEquals( + -1, + ArrowBatchSourceCallbacks.createSource( + registration.bindingId(), + MemorySegment.ofArray(new int[] { 1 }), + 1L, + error, + error.byteSize() + ) + ); + assertEquals(0, error.get(ValueLayout.JAVA_BYTE, error.byteSize() - 1L)); + String message = readCString(error, error.byteSize()); + assertTrue(message, message.contains("IllegalArgument")); + } + assertEquals(initialRefCount, reader.getRefCount()); + } + } + } + + public void testMultiValuedSourceErrorIsReturned() throws Exception { + try (Directory directory = newDirectory(); IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + Document document = new Document(); + document.add(new SortedNumericDocValuesField("number", 1L)); + document.add(new SortedNumericDocValuesField("number", 2L)); + writer.addDocument(document); + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(writer)) { + int initialRefCount = reader.getRefCount(); + DocValuesBatchSourceFactory factory = factory(reader, List.of(new InputColumn("number", ColumnKind.LONG))); + try ( + ArrowBatchSourceCallbacks.Registration registration = ArrowBatchSourceCallbacks.register(factory, null); + Arena arena = Arena.ofConfined() + ) { + MemorySegment error = arena.allocate(256L); + int sourceKey = ArrowBatchSourceCallbacks.createSource( + registration.bindingId(), + MemorySegment.ofArray(new int[] { 0 }), + 1L, + error, + error.byteSize() + ); + assertTrue(sourceKey > 0); + assertEquals( + ArrowBatchSourceCallbacks.ERROR, + ArrowBatchSourceCallbacks.nextBatch( + registration.bindingId(), + sourceKey, + MemorySegment.NULL, + MemorySegment.NULL, + error, + error.byteSize() + ) + ); + assertTrue(readCString(error, error.byteSize()).contains("multi-valued numeric doc values")); + ArrowBatchSourceCallbacks.releaseSource(registration.bindingId(), sourceKey); + } + assertEquals(initialRefCount, reader.getRefCount()); + } + } + } + + public void testCancellationReturnsCancelledStatus() throws Exception { + try (Directory directory = newDirectory(); IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + Document document = new Document(); + document.add(new NumericDocValuesField("number", 1L)); + writer.addDocument(document); + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(writer)) { + int initialRefCount = reader.getRefCount(); + DocValuesBatchSourceFactory factory = factory(reader, List.of(new InputColumn("number", ColumnKind.LONG))); + try ( + ArrowBatchSourceCallbacks.Registration registration = ArrowBatchSourceCallbacks.register(factory, null); + Arena arena = Arena.ofConfined() + ) { + MemorySegment error = arena.allocate(128L); + int sourceKey = ArrowBatchSourceCallbacks.createSource( + registration.bindingId(), + MemorySegment.ofArray(new int[] { 0 }), + 1L, + error, + error.byteSize() + ); + assertTrue(sourceKey > 0); + ArrowBatchSourceCallbacks.cancelSource(registration.bindingId(), sourceKey); + assertEquals( + ArrowBatchSourceCallbacks.CANCELLED, + ArrowBatchSourceCallbacks.nextBatch( + registration.bindingId(), + sourceKey, + MemorySegment.NULL, + MemorySegment.NULL, + error, + error.byteSize() + ) + ); + ArrowBatchSourceCallbacks.releaseSource(registration.bindingId(), sourceKey); + } + assertEquals(initialRefCount, reader.getRefCount()); + } + } + } + + public void testRegistrationCloseRejectsNewPullAndIsIdempotent() throws Exception { + try (Directory directory = newDirectory(); IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig())) { + writer.commit(); + try (DirectoryReader reader = DirectoryReader.open(writer); Arena arena = Arena.ofConfined()) { + int initialRefCount = reader.getRefCount(); + DocValuesBatchSourceFactory factory = factory(reader, List.of(new InputColumn("number", ColumnKind.LONG))); + ArrowBatchSourceCallbacks.Registration registration = ArrowBatchSourceCallbacks.register(factory, null); + MemorySegment error = arena.allocate(128L); + int sourceKey = ArrowBatchSourceCallbacks.createSource( + registration.bindingId(), + MemorySegment.ofArray(new int[] { 0 }), + 1L, + error, + error.byteSize() + ); + assertTrue(sourceKey > 0); + + registration.close(); + registration.close(); + + assertEquals( + ArrowBatchSourceCallbacks.ERROR, + ArrowBatchSourceCallbacks.nextBatch( + registration.bindingId(), + sourceKey, + MemorySegment.NULL, + MemorySegment.NULL, + error, + error.byteSize() + ) + ); + assertTrue(readCString(error, error.byteSize()).contains("binding is closed")); + ArrowBatchSourceCallbacks.releaseSource(registration.bindingId(), sourceKey); + assertEquals(initialRefCount, reader.getRefCount()); + expectThrows(IllegalStateException.class, () -> factory.open(new int[] { 0 })); + } + } + } + + private DocValuesBatchSourceFactory factory(DirectoryReader reader, List columns) throws Exception { + return new DocValuesBatchSourceFactory(new IndexSearcher(reader), new MatchAllDocsQuery(), columns, allocator, null); + } + + private static String readCString(MemorySegment segment, long capacity) { + byte[] bytes = segment.reinterpret(capacity).toArray(ValueLayout.JAVA_BYTE); + int length = 0; + while (length < bytes.length && bytes[length] != 0) { + length++; + } + return new String(Arrays.copyOf(bytes, length), StandardCharsets.UTF_8); + } +} diff --git a/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/lucene/LuceneSearchExecEngineTests.java b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/lucene/LuceneSearchExecEngineTests.java new file mode 100644 index 0000000000000..430778000d739 --- /dev/null +++ b/sandbox/plugins/analytics-backend-datafusion/src/test/java/org/opensearch/be/lucene/LuceneSearchExecEngineTests.java @@ -0,0 +1,145 @@ +/* + * 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.lucene; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.lucene.analysis.standard.StandardAnalyzer; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.NumericDocValuesField; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.MatchAllDocsQuery; +import org.apache.lucene.store.ByteBuffersDirectory; +import org.opensearch.analytics.backend.EngineResultStream; +import org.opensearch.analytics.backend.ShardScanExecutionContext; +import org.opensearch.analytics.exec.task.AnalyticsShardTask; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; +import org.opensearch.analytics.spi.ArrowBatchSource; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.ColumnKind; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.InputColumn; +import org.opensearch.analytics.spi.ArrowBatchSourcePlan; +import org.opensearch.analytics.spi.DelegationThreadTracker; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; + +public class LuceneSearchExecEngineTests extends OpenSearchTestCase { + + public void testTransfersPerExecutionDocValuesFactoryToBoundBackend() throws Exception { + try ( + RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer())) + ) { + addDocument(writer, 11L); + addDocument(writer, 22L); + writer.commit(); + + try (DirectoryReader reader = DirectoryReader.open(writer)) { + int initialRefCount = reader.getRefCount(); + IndexSearcher searcher = new IndexSearcher(reader); + ArrowBatchSourcePlan plan = new ArrowBatchSourcePlan( + "input-0", + new byte[] { 1 }, + List.of(new InputColumn("x", ColumnKind.LONG)) + ); + AnalyticsShardTask task = new AnalyticsShardTask( + 91L, + "test", + "arrow-source", + "arrow-source", + TaskId.EMPTY_TASK_ID, + Collections.emptyMap() + ); + TestTracker tracker = new TestTracker(); + ShardScanExecutionContext context = new ShardScanExecutionContext("index", task, null); + context.setAllocator(allocator); + context.setDelegationThreadTracker(tracker); + LuceneSearcherState state = new LuceneSearcherState(searcher, new MatchAllDocsQuery(), List.of(), plan); + AtomicBoolean executorCalled = new AtomicBoolean(); + AnalyticsSearchBackendPlugin backend = new AnalyticsSearchBackendPlugin() { + @Override + public String name() { + return "recording"; + } + + @Override + public EngineResultStream executeArrowBatchSource( + BufferAllocator resultAllocator, + ArrowBatchSourcePlan receivedPlan, + ArrowBatchSourceFactory sourceFactory, + Task receivedTask, + DelegationThreadTracker receivedTracker + ) { + executorCalled.set(true); + assertSame(allocator, resultAllocator); + assertSame(plan, receivedPlan); + assertSame(task, receivedTask); + assertSame(tracker, receivedTracker); + try (sourceFactory; ArrowBatchSource source = sourceFactory.open(new int[] { 0 })) { + try (VectorSchemaRoot root = source.nextBatch()) { + assertEquals(2, root.getRowCount()); + BigIntVector values = (BigIntVector) root.getVector("x"); + assertEquals(11L, values.get(0)); + assertEquals(22L, values.get(1)); + } + assertNull(source.nextBatch()); + } catch (Exception exception) { + throw new RuntimeException(exception); + } + return new EmptyResultStream(); + } + }; + LuceneSearchExecEngine engine = new LuceneSearchExecEngine(state, backend); + + try (EngineResultStream stream = engine.execute(context)) { + assertFalse(stream.iterator().hasNext()); + } + + assertTrue(executorCalled.get()); + assertEquals(initialRefCount, reader.getRefCount()); + } + } + } + + private static void addDocument(IndexWriter writer, long value) throws Exception { + Document document = new Document(); + document.add(new NumericDocValuesField("x", value)); + writer.addDocument(document); + } + + private static final class TestTracker implements DelegationThreadTracker { + @Override + public long trackStart() { + return Thread.currentThread().threadId(); + } + + @Override + public void trackEnd(long threadId) {} + } + + private static final class EmptyResultStream implements EngineResultStream { + @Override + public java.util.Iterator iterator() { + return Collections.emptyIterator(); + } + + @Override + public void close() {} + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSource.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSource.java new file mode 100644 index 0000000000000..8d907aa848333 --- /dev/null +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSource.java @@ -0,0 +1,525 @@ +/* + * 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.lucene; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.BaseFixedWidthVector; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.BitVectorHelper; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float4Vector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.VarBinaryVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ViewVarCharVector; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.impl.UnionListWriter; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.Text; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.index.BinaryDocValues; +import org.apache.lucene.index.DocValues; +import org.apache.lucene.index.DocValuesType; +import org.apache.lucene.index.FieldInfo; +import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.NumericDocValues; +import org.apache.lucene.index.SortedDocValues; +import org.apache.lucene.index.SortedNumericDocValues; +import org.apache.lucene.index.SortedSetDocValues; +import org.apache.lucene.search.DocIdSetIterator; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Scorer; +import org.apache.lucene.search.TwoPhaseIterator; +import org.apache.lucene.search.Weight; +import org.apache.lucene.util.Bits; +import org.apache.lucene.util.BytesRef; +import org.apache.lucene.util.NumericUtils; +import org.opensearch.analytics.spi.ArrowBatchSource; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.ColumnKind; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.InputColumn; +import org.opensearch.analytics.spi.ArrowBatchSourcePlan; +import org.opensearch.common.SuppressForbidden; +import org.opensearch.core.tasks.TaskCancelledException; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.Task; + +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.Locale; +import java.util.concurrent.atomic.AtomicBoolean; + +/** One independent sequential Lucene query/doc-values cursor. */ +@SuppressForbidden(reason = "reference counting keeps the reader alive for the source lifetime") +final class DocValuesBatchSource implements ArrowBatchSource { + + static final int BATCH_SIZE = 65_536; + private static final Logger LOGGER = LogManager.getLogger(DocValuesBatchSource.class); + + private final IndexSearcher searcher; + private final List columns; + private final BufferAllocator allocator; + private final Task task; + private final Weight weight; + private final Schema schema; + private final int[] docs = new int[BATCH_SIZE]; + private final long[] fallbackScratch = new long[BATCH_SIZE]; + private final int[] ordScratch = new int[BATCH_SIZE]; + private final long[] ordRowScratch = new long[BATCH_SIZE]; + private final Text textScratch = new Text(); + private final AtomicBoolean cancelled = new AtomicBoolean(); + private final AtomicBoolean closed = new AtomicBoolean(); + + private int leafIndex = -1; + private LeafReaderContext leaf; + private DocIdSetIterator iterator; + private Bits liveDocs; + private ColumnReader[] readers; + private boolean eof; + + DocValuesBatchSource(IndexSearcher searcher, Weight weight, List columns, BufferAllocator allocator, Task task) + throws IOException { + this.searcher = searcher; + this.columns = List.copyOf(columns); + this.allocator = allocator; + this.task = task; + this.weight = weight; + this.schema = ArrowBatchSourcePlan.schemaFor(columns); + } + + @Override + public BufferAllocator allocator() { + return allocator; + } + + @Override + public synchronized VectorSchemaRoot nextBatch() throws Exception { + ensureOpen(); + checkCancelled(); + while (eof == false) { + if (iterator == null && advanceLeaf() == false) { + eof = true; + return null; + } + int size = 0; + while (size < BATCH_SIZE) { + int doc = iterator.nextDoc(); + if (doc == DocIdSetIterator.NO_MORE_DOCS) { + iterator = null; + break; + } + if (liveDocs == null || liveDocs.get(doc)) { + docs[size++] = doc; + } + if ((size & 0xFFF) == 0) { + checkCancelled(); + } + } + if (size == 0) { + continue; + } + VectorSchemaRoot root = decodeBatch(size); + try { + checkCancelled(); + return root; + } catch (RuntimeException | Error e) { + root.close(); + throw e; + } + } + return null; + } + + private boolean advanceLeaf() throws IOException { + List leaves = searcher.getIndexReader().leaves(); + while (++leafIndex < leaves.size()) { + leaf = leaves.get(leafIndex); + Scorer scorer = weight.scorer(leaf); + if (scorer == null) { + continue; + } + readers = new ColumnReader[columns.size()]; + for (int i = 0; i < readers.length; i++) { + readers[i] = openColumn(leaf, columns.get(i)); + } + liveDocs = leaf.reader().getLiveDocs(); + TwoPhaseIterator twoPhase = scorer.twoPhaseIterator(); + iterator = twoPhase == null ? scorer.iterator() : TwoPhaseIterator.asDocIdSetIterator(twoPhase); + return true; + } + leaf = null; + readers = null; + liveDocs = null; + iterator = null; + return false; + } + + private VectorSchemaRoot decodeBatch(int size) throws IOException { + VectorSchemaRoot root = VectorSchemaRoot.create(schema, allocator); + boolean success = false; + try { + for (int i = 0; i < readers.length; i++) { + FieldVector vector = root.getVector(i); + switch (readers[i]) { + case NumericColumn numeric -> decodeNumeric(numeric, vector, size); + case SortedNumericColumn numeric -> decodeSortedNumeric(numeric, (ListVector) vector, size); + case BinaryColumn binary -> decodeBinary(binary.values(), (VarBinaryVector) vector, size); + case SortedColumn sorted -> decodeSorted(sorted, vector, size); + case SortedSetColumn sortedSet -> decodeSortedSet(sortedSet, (ListVector) vector, size); + case MissingColumn ignored -> decodeMissing(vector, size); + } + } + root.setRowCount(size); + success = true; + return root; + } finally { + if (success == false) { + root.close(); + } + } + } + + private void decodeMissing(FieldVector vector, int size) { + switch (vector) { + case BaseFixedWidthVector fixedWidth -> fixedWidth.allocateNew(size); + case ViewVarCharVector view -> view.allocateNew((long) size * 16, size); + case VarBinaryVector binary -> binary.allocateNew((long) size * 16, size); + case ListVector list -> { + UnionListWriter writer = list.getWriter(); + for (int i = 0; i < size; i++) { + writer.setPosition(i); + writer.writeNull(); + } + writer.setValueCount(size); + return; + } + default -> throw new IllegalStateException("unsupported missing-column vector type " + vector.getClass().getName()); + } + vector.getValidityBuffer().setZero(0, (size + 7) / 8); + vector.setValueCount(size); + } + + private void decodeNumeric(NumericColumn column, FieldVector vector, int size) throws IOException { + column.values().longValues(size, docs, 0, fallbackScratch, 0, 0L); + switch (column.kind()) { + case LONG, TIMESTAMP -> decodeLong(column.validity(), (BaseFixedWidthVector) vector, size); + case BOOLEAN -> { + BitVector booleans = (BitVector) vector; + booleans.allocateNew(size); + for (int i = 0; i < size; i++) { + if (column.validity().advanceExact(docs[i])) { + booleans.set(i, fallbackScratch[i] == 0L ? 0 : 1); + } else { + booleans.setNull(i); + } + } + booleans.setValueCount(size); + } + case FLOAT -> { + Float4Vector floats = (Float4Vector) vector; + floats.allocateNew(size); + for (int i = 0; i < size; i++) { + if (column.validity().advanceExact(docs[i])) { + floats.set(i, NumericUtils.sortableIntToFloat((int) fallbackScratch[i])); + } else { + floats.setNull(i); + } + } + floats.setValueCount(size); + } + case DOUBLE -> { + Float8Vector doubles = (Float8Vector) vector; + doubles.allocateNew(size); + for (int i = 0; i < size; i++) { + if (column.validity().advanceExact(docs[i])) { + doubles.set(i, NumericUtils.sortableLongToDouble(fallbackScratch[i])); + } else { + doubles.setNull(i); + } + } + doubles.setValueCount(size); + } + case KEYWORD, BINARY, IP -> throw new IllegalStateException("non-numeric column kind " + column.kind()); + } + } + + private void decodeLong(NumericDocValues validity, BaseFixedWidthVector vector, int size) throws IOException { + vector.allocateNew(size); + vector.getValidityBuffer().setZero(0, (size + 7) / 8); + for (int i = 0; i < size; i++) { + vector.getDataBuffer().setLong((long) i * Long.BYTES, fallbackScratch[i]); + if (validity.advanceExact(docs[i])) { + BitVectorHelper.setBit(vector.getValidityBuffer(), i); + } + } + vector.setValueCount(size); + } + + private void decodeSorted(SortedColumn column, FieldVector vector, int size) throws IOException { + SortedDocValues values = column.values(); + for (int i = 0; i < size; i++) { + ordScratch[i] = values.advanceExact(docs[i]) ? values.ordValue() : -1; + } + allocateBytes(vector, size); + for (int i = 0; i < size; i++) { + ordRowScratch[i] = ((long) (ordScratch[i] + 1) << 20) | i; + } + Arrays.sort(ordRowScratch, 0, size); + + int lastOrd = Integer.MIN_VALUE; + BytesRef term = null; + for (int i = 0; i < size; i++) { + int row = (int) (ordRowScratch[i] & 0xFFFFF); + int ord = (int) (ordRowScratch[i] >>> 20) - 1; + if (ord < 0) { + vector.setNull(row); + } else { + if (ord != lastOrd) { + term = values.lookupOrd(ord); + lastOrd = ord; + } + setBytes(vector, row, term); + } + } + vector.setValueCount(size); + } + + private void decodeBinary(BinaryDocValues values, VarBinaryVector vector, int size) throws IOException { + vector.allocateNew((long) size * 16, size); + for (int i = 0; i < size; i++) { + if (values.advanceExact(docs[i])) { + BytesRef value = values.binaryValue(); + vector.setSafe(i, value.bytes, value.offset, value.length); + } else { + vector.setNull(i); + } + } + vector.setValueCount(size); + } + + private void decodeSortedNumeric(SortedNumericColumn column, ListVector vector, int size) throws IOException { + SortedNumericDocValues values = column.values(); + UnionListWriter writer = vector.getWriter(); + int valueIndex = 0; + for (int i = 0; i < size; i++) { + writer.setPosition(i); + if (values.advanceExact(docs[i]) == false) { + writer.writeNull(); + continue; + } + writer.startList(); + int count = values.docValueCount(); + for (int j = 0; j < count; j++) { + writeNumeric(writer, column.kind(), values.nextValue()); + if ((++valueIndex & 0xFFF) == 0) { + checkCancelled(); + } + } + writer.endList(); + } + writer.setValueCount(size); + } + + private void decodeSortedSet(SortedSetColumn column, ListVector vector, int size) throws IOException { + SortedSetDocValues values = column.values(); + UnionListWriter writer = vector.getWriter(); + int valueIndex = 0; + for (int i = 0; i < size; i++) { + writer.setPosition(i); + if (values.advanceExact(docs[i]) == false) { + writer.writeNull(); + continue; + } + writer.startList(); + int count = values.docValueCount(); + for (int j = 0; j < count; j++) { + writeBytes(writer, column.kind(), values.lookupOrd(values.nextOrd())); + if ((++valueIndex & 0xFFF) == 0) { + checkCancelled(); + } + } + writer.endList(); + } + writer.setValueCount(size); + } + + private static void writeNumeric(UnionListWriter writer, ColumnKind kind, long value) { + switch (kind) { + case LONG -> writer.bigInt().writeBigInt(value); + case TIMESTAMP -> writer.timeStampMilli().writeTimeStampMilli(value); + case BOOLEAN -> writer.bit().writeBit(value == 0L ? 0 : 1); + case FLOAT -> writer.float4().writeFloat4(NumericUtils.sortableIntToFloat((int) value)); + case DOUBLE -> writer.float8().writeFloat8(NumericUtils.sortableLongToDouble(value)); + case KEYWORD, BINARY, IP -> throw new IllegalStateException("non-numeric column kind " + kind); + } + } + + private void writeBytes(UnionListWriter writer, ColumnKind kind, BytesRef value) { + switch (kind) { + case KEYWORD -> { + textScratch.set(value.bytes, value.offset, value.length); + writer.writeViewVarChar(textScratch); + } + case BINARY, IP -> writer.writeVarBinary(value.bytes, value.offset, value.length); + case LONG, TIMESTAMP, BOOLEAN, FLOAT, DOUBLE -> throw new IllegalStateException("non-binary column kind " + kind); + } + } + + private static void allocateBytes(FieldVector vector, int size) { + switch (vector) { + case ViewVarCharVector text -> text.allocateNew((long) size * 16, size); + case VarBinaryVector binary -> binary.allocateNew((long) size * 16, size); + default -> throw new IllegalStateException("unsupported byte vector type " + vector.getClass().getName()); + } + } + + private static void setBytes(FieldVector vector, int row, BytesRef value) { + switch (vector) { + case ViewVarCharVector text -> text.setSafe(row, value.bytes, value.offset, value.length); + case VarBinaryVector binary -> binary.setSafe(row, value.bytes, value.offset, value.length); + default -> throw new IllegalStateException("unsupported byte vector type " + vector.getClass().getName()); + } + } + + private static ColumnReader openColumn(LeafReaderContext leaf, InputColumn column) throws IOException { + FieldInfo fieldInfo = leaf.reader().getFieldInfos().fieldInfo(column.name()); + if (fieldInfo == null || fieldInfo.getDocValuesType() == DocValuesType.NONE) { + return MissingColumn.INSTANCE; + } + DocValuesType type = fieldInfo.getDocValuesType(); + if (column.multiValued()) { + return switch (column.kind()) { + case KEYWORD, IP -> new SortedSetColumn(sortedSetValues(leaf, column.name(), type), column.kind()); + case BINARY -> throw incompatibleType(column.name(), type); + case LONG, TIMESTAMP, BOOLEAN, FLOAT, DOUBLE -> new SortedNumericColumn( + sortedNumericValues(leaf, column.name(), type), + column.kind() + ); + }; + } + return switch (column.kind()) { + case KEYWORD, IP -> new SortedColumn(sortedValues(leaf, column, type), column.kind()); + case BINARY -> type == DocValuesType.BINARY + ? new BinaryColumn(leaf.reader().getBinaryDocValues(column.name())) + : new SortedColumn(sortedValues(leaf, column, type), column.kind()); + case LONG, TIMESTAMP, BOOLEAN, FLOAT, DOUBLE -> new NumericColumn( + numericValues(leaf, column.name(), type), + numericValues(leaf, column.name(), type), + column.kind() + ); + }; + } + + private static NumericDocValues numericValues(LeafReaderContext leaf, String name, DocValuesType type) throws IOException { + if (type == DocValuesType.NUMERIC) { + return leaf.reader().getNumericDocValues(name); + } + if (type == DocValuesType.SORTED_NUMERIC) { + NumericDocValues singleton = DocValues.unwrapSingleton(leaf.reader().getSortedNumericDocValues(name)); + if (singleton == null) { + throw new IllegalArgumentException("multi-valued numeric doc values are not supported for scalar field [" + name + "]"); + } + return singleton; + } + throw incompatibleType(name, type); + } + + private static SortedNumericDocValues sortedNumericValues(LeafReaderContext leaf, String name, DocValuesType type) throws IOException { + if (type == DocValuesType.NUMERIC) { + return DocValues.singleton(leaf.reader().getNumericDocValues(name)); + } + if (type == DocValuesType.SORTED_NUMERIC) { + return leaf.reader().getSortedNumericDocValues(name); + } + throw incompatibleType(name, type); + } + + private static SortedDocValues sortedValues(LeafReaderContext leaf, InputColumn column, DocValuesType type) throws IOException { + if (type == DocValuesType.SORTED) { + return leaf.reader().getSortedDocValues(column.name()); + } + if (type == DocValuesType.SORTED_SET) { + SortedDocValues singleton = DocValues.unwrapSingleton(leaf.reader().getSortedSetDocValues(column.name())); + if (singleton == null) { + String kind = column.kind().name().toLowerCase(Locale.ROOT); + throw new IllegalArgumentException( + "multi-valued " + kind + " doc values are not supported for scalar field [" + column.name() + "]" + ); + } + return singleton; + } + throw incompatibleType(column.name(), type); + } + + private static SortedSetDocValues sortedSetValues(LeafReaderContext leaf, String name, DocValuesType type) throws IOException { + if (type == DocValuesType.SORTED) { + return DocValues.singleton(leaf.reader().getSortedDocValues(name)); + } + if (type == DocValuesType.SORTED_SET) { + return leaf.reader().getSortedSetDocValues(name); + } + throw incompatibleType(name, type); + } + + private static IllegalArgumentException incompatibleType(String name, DocValuesType type) { + return new IllegalArgumentException("field [" + name + "] has incompatible doc values type [" + type + "]"); + } + + private void checkCancelled() { + if (cancelled.get() || (task instanceof CancellableTask cancellableTask && cancellableTask.isCancelled())) { + throw new TaskCancelledException("doc-values scan cancelled"); + } + } + + private void ensureOpen() { + if (closed.get()) { + throw new IllegalStateException("doc-values source is closed"); + } + } + + @Override + public void cancel() { + cancelled.set(true); + } + + @Override + public void close() { + cancel(); + if (closed.compareAndSet(false, true)) { + try { + searcher.getIndexReader().decRef(); + } catch (IOException e) { + LOGGER.warn("failed to release doc-values source reader", e); + } + } + } + + private sealed interface ColumnReader permits NumericColumn, SortedNumericColumn, BinaryColumn, SortedColumn, SortedSetColumn, + MissingColumn {} + + private enum MissingColumn implements ColumnReader { + INSTANCE + } + + private record NumericColumn(NumericDocValues values, NumericDocValues validity, ColumnKind kind) implements ColumnReader { + } + + private record SortedNumericColumn(SortedNumericDocValues values, ColumnKind kind) implements ColumnReader { + } + + private record BinaryColumn(BinaryDocValues values) implements ColumnReader { + } + + private record SortedColumn(SortedDocValues values, ColumnKind kind) implements ColumnReader { + } + + private record SortedSetColumn(SortedSetDocValues values, ColumnKind kind) implements ColumnReader { + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSourceFactory.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSourceFactory.java new file mode 100644 index 0000000000000..7a3dea8fbb1a5 --- /dev/null +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/DocValuesBatchSourceFactory.java @@ -0,0 +1,97 @@ +/* + * 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.lucene; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.ScoreMode; +import org.apache.lucene.search.Weight; +import org.opensearch.analytics.spi.ArrowBatchSource; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.InputColumn; +import org.opensearch.common.SuppressForbidden; +import org.opensearch.tasks.Task; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; + +/** Reusable factory for independent Lucene doc-values cursors. */ +@SuppressForbidden(reason = "reference counting keeps the reader alive for the factory and source lifetimes") +public final class DocValuesBatchSourceFactory implements ArrowBatchSourceFactory { + + private static final Logger LOGGER = LogManager.getLogger(DocValuesBatchSourceFactory.class); + + private final IndexSearcher searcher; + private final Weight weight; + private final List columns; + private final BufferAllocator allocator; + private final Task task; + private final AtomicBoolean closed = new AtomicBoolean(); + + public DocValuesBatchSourceFactory(IndexSearcher searcher, Query query, List columns, BufferAllocator allocator, Task task) + throws java.io.IOException { + searcher.getIndexReader().incRef(); + boolean success = false; + try { + IndexSearcher uncachedSearcher = new IndexSearcher(searcher.getIndexReader()); + uncachedSearcher.setSimilarity(searcher.getSimilarity()); + uncachedSearcher.setQueryCache(null); + this.searcher = uncachedSearcher; + this.weight = uncachedSearcher.createWeight(uncachedSearcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1f); + this.columns = List.copyOf(columns); + this.allocator = allocator; + this.task = task; + success = true; + } finally { + if (success == false) { + searcher.getIndexReader().decRef(); + } + } + } + + @Override + public synchronized ArrowBatchSource open(int[] projection) throws Exception { + if (closed.get()) { + throw new IllegalStateException("doc-values source factory is closed"); + } + List projected = new ArrayList<>(projection.length); + for (int index : projection) { + if (index < 0 || index >= columns.size()) { + throw new IllegalArgumentException("projection index [" + index + "] outside input schema of size " + columns.size()); + } + projected.add(columns.get(index)); + } + searcher.getIndexReader().incRef(); + boolean success = false; + try { + DocValuesBatchSource source = new DocValuesBatchSource(searcher, weight, projected, allocator, task); + success = true; + return source; + } finally { + if (success == false) { + searcher.getIndexReader().decRef(); + } + } + } + + @Override + public synchronized void close() { + if (closed.compareAndSet(false, true)) { + try { + searcher.getIndexReader().decRef(); + } catch (java.io.IOException e) { + LOGGER.warn("failed to release doc-values source factory reader", e); + } + } + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPlugin.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPlugin.java index 36b26ec95a093..d01d9abd14403 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPlugin.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneAnalyticsBackendPlugin.java @@ -28,15 +28,24 @@ import org.opensearch.analytics.spi.FilterDelegationHandle; import org.opensearch.analytics.spi.FragmentConvertor; import org.opensearch.analytics.spi.FragmentInstructionHandlerFactory; +import org.opensearch.analytics.spi.ProjectCapability; import org.opensearch.analytics.spi.ScalarFunction; import org.opensearch.analytics.spi.ScanCapability; import org.opensearch.analytics.spi.SearchExecEngineProvider; +import org.opensearch.common.concurrent.GatedCloseable; +import org.opensearch.common.util.io.IOUtils; +import org.opensearch.index.engine.DataFormatAwareEngine.DataFormatAwareReader; +import org.opensearch.index.engine.Engine; +import org.opensearch.index.engine.EngineBackedIndexer; import org.opensearch.index.engine.exec.IndexReaderProvider; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryShardContext; +import org.opensearch.index.shard.IndexShard; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; +import java.io.IOException; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -117,6 +126,18 @@ public class LuceneAnalyticsBackendPlugin implements AnalyticsSearchBackendPlugi private static final Set KEYWORD_ONLY = Set.of(FieldType.KEYWORD); + /** Field types whose values the Arrow source reader can expose without type coercion. */ + private static final Set DOC_VALUES_TYPES = Set.of( + FieldType.LONG, + FieldType.DATE, + FieldType.KEYWORD, + FieldType.BOOLEAN, + FieldType.FLOAT, + FieldType.DOUBLE, + FieldType.BINARY, + FieldType.IP + ); + private static final Set FILTER_CAPS; static { Set caps = new HashSet<>(); @@ -125,6 +146,7 @@ public class LuceneAnalyticsBackendPlugin implements AnalyticsSearchBackendPlugi caps.add(new FilterCapability.Standard(op, KEYWORD_ONLY, LUCENE_FORMATS)); } else { caps.add(new FilterCapability.Standard(op, STANDARD_TYPES, LUCENE_FORMATS)); + caps.add(new FilterCapability.Standard(op, DOC_VALUES_TYPES, LUCENE_FORMATS)); } } for (ScalarFunction op : FULL_TEXT_OPS) { @@ -135,28 +157,77 @@ public class LuceneAnalyticsBackendPlugin implements AnalyticsSearchBackendPlugi FILTER_CAPS = caps; } + private static final Set NUMERIC_DOC_VALUES_TYPES = Set.of(FieldType.LONG, FieldType.FLOAT, FieldType.DOUBLE); + + /** Scalar expressions evaluated by DataFusion after Lucene supplies doc-values batches. */ + private static final Set PROJECT_CAPS; + static { + Set returnTypes = new HashSet<>(DOC_VALUES_TYPES); + returnTypes.add(FieldType.DOUBLE); + returnTypes.add(FieldType.FLOAT); + returnTypes.add(FieldType.BOOLEAN); + Set capabilities = new HashSet<>(); + for (ScalarFunction function : List.of( + ScalarFunction.PLUS, + ScalarFunction.MINUS, + ScalarFunction.TIMES, + ScalarFunction.DIVIDE, + ScalarFunction.MOD, + ScalarFunction.CAST, + ScalarFunction.EXTRACT, + ScalarFunction.DATE_FORMAT, + ScalarFunction.REGEXP_REPLACE, + ScalarFunction.CASE, + ScalarFunction.AND, + ScalarFunction.OR, + ScalarFunction.NOT, + ScalarFunction.EQUALS, + ScalarFunction.NOT_EQUALS, + ScalarFunction.GREATER_THAN, + ScalarFunction.GREATER_THAN_OR_EQUAL, + ScalarFunction.LESS_THAN, + ScalarFunction.LESS_THAN_OR_EQUAL + )) { + capabilities.add(new ProjectCapability.Scalar(function, returnTypes, LUCENE_FORMATS, true)); + } + capabilities.add(new ProjectCapability.Scalar(ScalarFunction.CHAR_LENGTH, Set.of(FieldType.LONG), LUCENE_FORMATS, true)); + PROJECT_CAPS = Set.copyOf(capabilities); + } + /** * Lucene-secondary indexes the term dictionary (inverted index) for the same field * types it accepts filters on — keyword / text / match_only_text. The Index * scan capability lets the planner mark Lucene viable as a driver for metadata-only - * operations (count today, group-by-count and top-K terms in future) over scans whose - * fields are listed here. It does NOT imply Lucene can deliver row values; consumers - * needing values (Project, Sort) consult value-producing scan capabilities separately - * and self-restrict, which the chain-agreement filter at PlanForker enforces. + * operations over scans whose fields are listed here. The separate DocValues capability + * covers supported value-producing plans; shape validation rejects unsupported referenced + * columns before selection. */ - private static final Set SCAN_CAPS = Set.of(new ScanCapability.Index(LUCENE_FORMATS, STANDARD_TYPES)); + private static final Set INDEX_SCAN_CAPS = Set.of(new ScanCapability.Index(LUCENE_FORMATS, STANDARD_TYPES)); + private static final Set SCAN_CAPS = Set.of( + new ScanCapability.Index(LUCENE_FORMATS, STANDARD_TYPES), + new ScanCapability.DocValues(LUCENE_FORMATS, DOC_VALUES_TYPES) + ); - /** - * Lucene drives count(*) and (in a follow-up) count(col) over fields it indexes. - * Coupled with the Index scan capability above, this lets PlanForker emit a - * Lucene-driver StagePlan alternative for count-shaped fragments without bypassing - * the existing engine path. - */ - private static final Set AGGREGATE_CAPS = Set.of( + private static final Set COUNT_CAPS = Set.of( AggregateCapability.simple(AggregateFunction.COUNT, STANDARD_TYPES, LUCENE_FORMATS) ); + /** Aggregate shapes supported by either the count fast path or the Arrow source plan. */ + private static final Set AGGREGATE_CAPS; + static { + Set capabilities = new HashSet<>(COUNT_CAPS); + for (AggregateFunction function : List.of(AggregateFunction.SUM, AggregateFunction.SUM0, AggregateFunction.AVG)) { + capabilities.add(AggregateCapability.simple(function, NUMERIC_DOC_VALUES_TYPES, LUCENE_FORMATS)); + } + for (AggregateFunction function : List.of(AggregateFunction.COUNT, AggregateFunction.MIN, AggregateFunction.MAX)) { + capabilities.add(AggregateCapability.simple(function, DOC_VALUES_TYPES, LUCENE_FORMATS)); + } + AGGREGATE_CAPS = Set.copyOf(capabilities); + } + private final LucenePlugin plugin; + private volatile AnalyticsSearchBackendPlugin arrowSourceBackend; + private final BackendShardPreference shardPreference = new LuceneShardPreference(() -> arrowSourceBackend != null); public LuceneAnalyticsBackendPlugin(LucenePlugin plugin) { this.plugin = plugin; @@ -167,6 +238,18 @@ public String name() { return LuceneDataFormat.LUCENE_FORMAT_NAME; } + @Override + public void bindBackends(Map backends) { + List candidates = backends.values() + .stream() + .filter(AnalyticsSearchBackendPlugin::supportsArrowBatchSourceExecution) + .toList(); + if (candidates.size() > 1) { + throw new IllegalStateException("Multiple Arrow batch source execution backends are installed: " + candidates); + } + arrowSourceBackend = candidates.isEmpty() ? null : candidates.getFirst(); + } + @Override public BackendCapabilityProvider getCapabilityProvider() { return new BackendCapabilityProvider() { @@ -182,12 +265,17 @@ public Set filterCapabilities() { @Override public Set scanCapabilities() { - return SCAN_CAPS; + return arrowSourceBackend == null ? INDEX_SCAN_CAPS : SCAN_CAPS; } @Override public Set aggregateCapabilities() { - return AGGREGATE_CAPS; + return arrowSourceBackend == null ? COUNT_CAPS : AGGREGATE_CAPS; + } + + @Override + public Set projectCapabilities() { + return arrowSourceBackend == null ? Set.of() : PROJECT_CAPS; } @Override @@ -202,15 +290,45 @@ public Map delegatedPredicateSeria @Override public BackendShardPreference shardPreference() { - return SHARD_PREFERENCE; + return shardPreference; } }; } - private static final BackendShardPreference SHARD_PREFERENCE = new LuceneShardPreference(); - private static final Logger LOGGER = LogManager.getLogger(LuceneAnalyticsBackendPlugin.class); + /** + * Standard Lucene shards expose an {@link Engine.Searcher}, not the pluggable-format reader + * implemented by composite engines. Adapt that searcher to the shared reader contract so the + * same Lucene execution code can consume both index types. + */ + @Override + public GatedCloseable acquireReader(IndexShard shard) throws IOException { + IndexReaderProvider readerProvider = shard.getReaderProvider(); + if (!(readerProvider instanceof EngineBackedIndexer indexer)) { + return readerProvider.acquireReader(); + } + + Engine.Searcher searcher = shard.acquireSearcher("analytics-lucene"); + GatedCloseable snapshotRef; + try { + snapshotRef = indexer.acquireSnapshot(); + } catch (RuntimeException | Error e) { + searcher.close(); + throw e; + } + try { + DataFormatAwareReader reader = new DataFormatAwareReader( + snapshotRef, + Map.of(plugin.getDataFormat(), new LuceneReader(searcher.getDirectoryReader(), Map.of())) + ); + return new GatedCloseable<>(reader, () -> IOUtils.close(searcher, reader)); + } catch (RuntimeException | Error e) { + IOUtils.closeWhileHandlingException(searcher, snapshotRef); + throw e; + } + } + @Override public FilterDelegationHandle getFilterDelegationHandle(List expressions, CommonExecutionContext ctx) { ShardScanExecutionContext shardCtx = (ShardScanExecutionContext) ctx; @@ -238,7 +356,7 @@ public FilterDelegationHandle getFilterDelegationHandle(ListReuses the same leaf-serializer registry as {@link LuceneSubtreeConvertor} via - * {@link QuerySerializerRegistry} — keyword equality, MATCH, MATCH_PHRASE, etc. all - * round-trip through the same {@link DelegatedPredicateSerializer} → {@link QueryBuilder} - * path. The data-node Lucene driver deserializes the bytes via NamedWriteable and runs - * {@code IndexSearcher.count} on the resulting {@link QueryBuilder#toQuery(QueryShardContext)}. - * - *

Multi-stage / non-shard-scan fragments aren't supported: Lucene drives shard-local - * count fragments only. Reduce or coordinator stages still run on DataFusion, so this - * convertor is invoked only when the planner picked Lucene as the StagePlan's backend — - * which happens exclusively for count-fast-path-eligible shards today. + *

Filters reuse {@link QuerySerializerRegistry}, so delegated and driver execution use + * the same {@link QueryBuilder} conversion. * * @opensearch.internal */ @@ -71,140 +58,97 @@ final class LuceneFragmentConvertor implements FragmentConvertor { private static final Logger LOGGER = LogManager.getLogger(LuceneFragmentConvertor.class); private final Map leafSerializers; + private final AnalyticsSearchBackendPlugin arrowSourceBackend; LuceneFragmentConvertor(Map leafSerializers) { - this.leafSerializers = leafSerializers; + this(leafSerializers, null); } - /** - * True iff the top is an {@link Aggregate} with empty group-set whose every call is - * {@link SqlKind#COUNT} — what {@code IndexSearcher.count} can answer from the term - * dictionary. Read by {@link LuceneShardPreference} to score this fragment. - * - *

Defense-in-depth: PlanForker's chain-agreement filter already narrows aggregate - * alternatives to declared capabilities (prod Lucene declares only COUNT), so this - * guards against capability-declaration drift. - */ - static boolean isCountFastPath(RelNode fragment) { - if (fragment instanceof Aggregate == false) return false; - Aggregate agg = (Aggregate) fragment; - if (agg.getGroupSet().isEmpty() == false) return false; - for (AggregateCall call : agg.getAggCallList()) { - if (call.getAggregation().getKind() != SqlKind.COUNT) return false; - } - return true; + LuceneFragmentConvertor( + Map leafSerializers, + AnalyticsSearchBackendPlugin arrowSourceBackend + ) { + this.leafSerializers = leafSerializers; + this.arrowSourceBackend = arrowSourceBackend; } @Override public byte[] convertFragment(RelNode fragment) { - // Lucene-driver wire format: [columnNames StringCollection] [hasFilter boolean] - // [QueryBuilder NamedWriteable]?. Both ends are controlled (this convertor on the - // coordinator, LuceneScanInstructionHandler on the data node), so a tiny custom - // format is fine — beats threading column names through the InstructionNode. - // columnNames may be empty when the convertor runs against a non-count Lucene - // alternative kept around for delegation (e.g. DF drives, Lucene is the peer); the - // bytes are produced but the data node never invokes them — selector or runtime - // alternative-selection drops this plan before dispatch. - List columnNames = extractAggCallNames(fragment); - QueryBuilder filterQuery = null; - Filter filter = findFilter(fragment); - if (filter != null) { - // strip() in FragmentConversionDriver replaces OpenSearchFilter with a plain - // LogicalFilter, so the field-storage info lives on the OpenSearch ancestor - // below (the TableScan). Walk down past LogicalFilter to find the nearest - // OpenSearchRelNode and use its output field storage. The condition itself was - // already resolved (annotation placeholders unwrapped) by the resolver in strip(). - List fieldStorage = findFieldStorage(filter); - filterQuery = toQueryBuilder(filter.getCondition(), fieldStorage); + LuceneFragmentPlanner.Shape shape = LuceneFragmentPlanner.classify(fragment); + if (shape instanceof LuceneFragmentPlanner.ArrowSourceShape arrowSource) { + return convertArrowSourceShape(arrowSource, false); } - byte[] bytes; - try (BytesStreamOutput out = new BytesStreamOutput()) { - out.writeStringCollection(columnNames); - if (filterQuery == null) { - out.writeBoolean(false); - } else { - out.writeBoolean(true); - out.writeNamedWriteable(filterQuery); - } - bytes = BytesReference.toBytes(out.bytes()); - } catch (IOException e) { - throw new IllegalStateException("Failed to serialize Lucene-driver fragment", e); - } - LOGGER.debug("[lucene-count] convertFragment columnNames={} filterQuery={} bytes={}", columnNames, filterQuery, bytes.length); + + QueryBuilder filterQuery = toQueryBuilder(shape.filter()); + byte[] bytes = LuceneFragmentWirePlan.create(shape.outputNames(), filterQuery, null).toBytes(); + LOGGER.debug( + "[lucene-count] convertFragment outputNames={} filterQuery={} bytes={}", + shape.outputNames(), + filterQuery, + bytes.length + ); return bytes; } - /** - * Walks down to find an Aggregate (Calcite {@link Aggregate} or {@code OpenSearchAggregate}) - * and extracts the user-facing call names. These become the Arrow output column names so - * the coordinator's reduce sink sees the schema it expects. - */ - private static List extractAggCallNames(RelNode root) { - RelNode current = root; - while (current != null) { - if (current instanceof Aggregate agg) { - List names = new ArrayList<>(agg.getAggCallList().size()); - for (AggregateCall call : agg.getAggCallList()) { - names.add(call.getName()); - } - return names; - } - if (current.getInputs().isEmpty()) break; - current = current.getInputs().getFirst(); - } - return List.of(); + private byte[] convertArrowSourceShape(LuceneFragmentPlanner.ArrowSourceShape shape, boolean partialAggregate) { + byte[] planBytes = arrowSourceBackend().compileArrowBatchSourcePlan(shape.rebasedFragment(), partialAggregate); + QueryBuilder filterQuery = toQueryBuilder(shape.filter()); + ArrowBatchSourcePlan sourcePlan = new ArrowBatchSourcePlan(shape.inputId(), planBytes, shape.inputColumns()); + byte[] bytes = LuceneFragmentWirePlan.create(shape.outputNames(), filterQuery, sourcePlan).toBytes(); + LOGGER.debug( + "[lucene-arrow-source] inputColumns={} outputNames={} plan={}B filter={} bytes={}", + shape.inputColumns(), + shape.outputNames(), + planBytes.length, + filterQuery, + bytes.length + ); + return bytes; + } + + private QueryBuilder toQueryBuilder(Filter filter) { + return filter == null ? null : toQueryBuilder(filter.getCondition(), findFieldStorage(filter)); } @Override public byte[] attachPartialAggOnTop(RelNode partialAggFragment, byte[] innerBytes) { - // Lucene-as-driver count fragments DO go through the partial-agg split — the driver's - // FragmentConversionDriver invokes convertFragment on the input subtree (the - // TableScan / Filter, no Aggregate above), then attachPartialAggOnTop on the - // OpenSearchAggregate node. Without this rewrite, innerBytes carries an empty - // columnNames list (extractAggCallNames found no Aggregate in the input) and the - // data-node Lucene exec engine emits a 0-column Arrow batch — the coordinator - // reduce sink then stalls waiting for the count column. - // - // Strategy: re-decode innerBytes' columnNames length-prefix (always present, possibly - // empty), then preserve the remaining tail (hasFilter + optional QueryBuilder) - // verbatim. Re-emit with the partialAggFragment's aggregate-call names as the new - // columnNames. Avoids needing a NamedWriteableRegistry at coordinator-side conversion. - if (!(partialAggFragment instanceof Aggregate agg)) { + if (partialAggFragment instanceof Aggregate == false) { throw new IllegalStateException( "Lucene attachPartialAggOnTop expected an Aggregate fragment, got " + partialAggFragment.getClass().getSimpleName() ); } - List columnNames = new ArrayList<>(agg.getAggCallList().size()); - for (AggregateCall call : agg.getAggCallList()) { - columnNames.add(call.getName()); + LuceneFragmentPlanner.Shape shape = LuceneFragmentPlanner.classify(partialAggFragment); + if (shape instanceof LuceneFragmentPlanner.ArrowSourceShape arrowSource) { + return convertArrowSourceShape(arrowSource, true); } - // Read past the inner columnNames StringCollection to get the byte offset of the - // hasFilter + optional QueryBuilder tail. We then copy the tail verbatim into the new - // bytes prefixed by the aggregate's column names. - int tailOffset; - try (StreamInput in = StreamInput.wrap(innerBytes)) { - in.readStringList(); // discard inner columnNames; we'll write the agg names instead - tailOffset = innerBytes.length - in.available(); - } catch (IOException e) { - throw new IllegalStateException("Failed to decode Lucene innerBytes during partial-agg attach", e); + byte[] bytes = LuceneFragmentWirePlan.fromBytes(innerBytes).withOutputNames(shape.outputNames()).toBytes(); + LOGGER.debug("[lucene-count] attachPartialAggOnTop outputNames={} bytes={}", shape.outputNames(), bytes.length); + return bytes; + } + + @Override + public byte[] attachFragmentOnTop(RelNode fragment, byte[] innerBytes) { + LuceneFragmentWirePlan inner = LuceneFragmentWirePlan.fromBytes(innerBytes); + ArrowBatchSourcePlan innerPlan = inner.arrowSourcePlan(); + if (innerPlan == null) { + throw new UnsupportedOperationException("Cannot attach a fragment to the Lucene count wire format"); } + byte[] planBytes = arrowSourceBackend().attachArrowBatchSourcePlan(fragment, innerPlan.planBytes()); + ArrowBatchSourcePlan plan = new ArrowBatchSourcePlan(innerPlan.inputId(), planBytes, innerPlan.inputColumns()); + return inner.withArrowSourcePlan(plan, LuceneFragmentPlanner.resultNames(fragment)).toBytes(); + } - try (BytesStreamOutput out = new BytesStreamOutput()) { - out.writeStringCollection(columnNames); - out.writeBytes(innerBytes, tailOffset, innerBytes.length - tailOffset); - byte[] bytes = BytesReference.toBytes(out.bytes()); - LOGGER.debug("[lucene-count] attachPartialAggOnTop columnNames={} bytes={}", columnNames, bytes.length); - return bytes; - } catch (IOException e) { - throw new IllegalStateException("Failed to serialize Lucene-driver partial-agg bytes", e); + private AnalyticsSearchBackendPlugin arrowSourceBackend() { + if (arrowSourceBackend == null) { + throw new IllegalStateException("No Arrow batch source execution backend is available"); } + return arrowSourceBackend; } @Override public WireFormat wireFormat() { - // convertFragment emits a custom NamedWriteable wire format ([columnNames][hasFilter] - // [BoolQueryBuilder]?), not self-describing. The orchestrator queries this so it + // convertFragment emits a typed OpenSearch wire payload, not Substrait. The orchestrator queries this so it // knows to emit a separate schema-only stub via convertSchemaOnlyRead for the // coordinator's reduce-sink partition registration. return WireFormat.OPAQUE; @@ -290,28 +234,15 @@ private static Type toSubstraitType(RelDataType type) { case DOUBLE -> Type.newBuilder().setFp64(Type.FP64.newBuilder().setNullability(n)).build(); case FLOAT, REAL -> Type.newBuilder().setFp32(Type.FP32.newBuilder().setNullability(n)).build(); case VARCHAR, CHAR -> Type.newBuilder().setString(Type.String.newBuilder().setNullability(n)).build(); + case DATE, TIMESTAMP, TIMESTAMP_WITH_LOCAL_TIME_ZONE -> Type.newBuilder() + .setPrecisionTimestamp(Type.PrecisionTimestamp.newBuilder().setPrecision(3).setNullability(n)) + .build(); default -> throw new IllegalStateException( "Lucene convertSchemaOnlyRead: unmapped Calcite type " + type.getSqlTypeName() + " for field of type " + type ); }; } - /** - * Walks the linear input chain looking for any Calcite {@link Filter} (covers both - * {@link OpenSearchFilter} and the plain {@code LogicalFilter} that - * {@code FragmentConversionDriver.strip} produces once annotation resolution unwraps the - * filter's condition into native predicate calls). - */ - private static Filter findFilter(RelNode node) { - RelNode current = node; - while (current != null) { - if (current instanceof Filter filter) return filter; - if (current.getInputs().isEmpty()) return null; - current = current.getInputs().getFirst(); - } - return null; - } - /** * Returns the field-storage info for a filter's child operator. When the filter is a * native {@link OpenSearchFilter} this is just its own {@code getOutputFieldStorage()}; diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneFragmentPlanner.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneFragmentPlanner.java new file mode 100644 index 0000000000000..2d76fc9e3ee83 --- /dev/null +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneFragmentPlanner.java @@ -0,0 +1,432 @@ +/* + * 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.lucene; + +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.rel.core.Aggregate; +import org.apache.calcite.rel.core.AggregateCall; +import org.apache.calcite.rel.core.Filter; +import org.apache.calcite.rel.core.Project; +import org.apache.calcite.rel.logical.LogicalProject; +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.apache.calcite.sql.SqlKind; +import org.apache.calcite.util.ImmutableBitSet; +import org.opensearch.analytics.planner.rel.OpenSearchRelNode; +import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.ColumnKind; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory.InputColumn; +import org.opensearch.analytics.spi.FieldStorageInfo; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.TreeSet; + +/** Classifies one Lucene fragment as the count path, Arrow source path, or unsupported. */ +final class LuceneFragmentPlanner { + + private static final String SOURCE_INPUT_ID = "input-0"; + + private LuceneFragmentPlanner() {} + + sealed interface Shape permits CountShape, ArrowSourceShape, UnsupportedShape { + List outputNames(); + + Filter filter(); + } + + record CountShape(List outputNames, Filter filter) implements Shape { + } + + record UnsupportedShape(List outputNames, Filter filter) implements Shape { + } + + static Shape classify(RelNode fragment) { + if (isCountFastPath(fragment)) { + return new CountShape(aggregateOutputNames(fragment), findFilter(fragment)); + } + ArrowSourceShape arrowSource = extractArrowSourceShape(fragment); + if (arrowSource != null) { + return arrowSource; + } + return new UnsupportedShape(aggregateOutputNames(fragment), findFilter(fragment)); + } + + private static List aggregateOutputNames(RelNode root) { + RelNode current = root; + while (current != null) { + if (current instanceof Aggregate aggregate) { + List names = new ArrayList<>(aggregate.getAggCallList().size()); + for (int i = 0; i < aggregate.getAggCallList().size(); i++) { + String name = aggregate.getAggCallList().get(i).getName(); + names.add(name == null ? "EXPR$" + i : name); + } + return names; + } + current = current.getInputs().isEmpty() ? null : current.getInputs().getFirst(); + } + return List.of(); + } + + private static Filter findFilter(RelNode node) { + RelNode current = node; + while (current != null) { + if (current instanceof Filter filter) { + return filter; + } + current = current.getInputs().isEmpty() ? null : current.getInputs().getFirst(); + } + return null; + } + + /** + * True iff the top is an {@link Aggregate} with empty group-set whose every call is + * {@link SqlKind#COUNT} — what {@code IndexSearcher.count} can answer from the term + * dictionary. Read by {@link LuceneShardPreference} to score this fragment. + * + *

Defense-in-depth: PlanForker's chain-agreement filter already narrows aggregate + * alternatives to declared capabilities (prod Lucene declares only COUNT), so this + * guards against capability-declaration drift. + */ + static boolean isCountFastPath(RelNode fragment) { + if (fragment instanceof Aggregate == false) return false; + Aggregate agg = (Aggregate) fragment; + if (agg.getGroupSet().isEmpty() == false || agg.getAggCallList().isEmpty()) return false; + for (AggregateCall call : agg.getAggCallList()) { + if (call.getAggregation().getKind() != SqlKind.COUNT || call.getArgList().isEmpty() == false) return false; + } + return true; + } + + private static InputColumn docValuesColumn(List storage, int ordinal) { + if (ordinal < 0 || ordinal >= storage.size()) { + return null; + } + FieldStorageInfo info = storage.get(ordinal); + if (info.isDerived()) { + return null; + } + List docValueFormats = info.getDocValueFormats(); + if (docValueFormats == null || docValueFormats.contains(LuceneDataFormat.LUCENE_FORMAT_NAME) == false) { + return null; + } + ColumnKind kind = switch (info.getFieldType()) { + case DATE -> ColumnKind.TIMESTAMP; + case LONG -> ColumnKind.LONG; + case KEYWORD -> ColumnKind.KEYWORD; + case BOOLEAN -> ColumnKind.BOOLEAN; + case FLOAT -> ColumnKind.FLOAT; + case DOUBLE -> ColumnKind.DOUBLE; + case BINARY -> ColumnKind.BINARY; + case IP -> ColumnKind.IP; + default -> null; + }; + return kind == null ? null : new InputColumn(info.getFieldName(), kind); + } + + /** A supported fragment rebased onto the named Arrow source table. */ + record ArrowSourceShape(String inputId, RelNode rebasedFragment, List inputColumns, Filter filter, List< + String> outputNames) implements Shape { + } + + /** + * Extracts aggregate and row-returning unary fragments that can read all required input + * values from single-valued Lucene doc values. + */ + private static ArrowSourceShape extractArrowSourceShape(RelNode fragment) { + RelNode originalFragment = fragment; + List wrappers = new ArrayList<>(); + while (fragment instanceof Aggregate == false) { + if (fragment.getInputs().size() != 1 + || (fragment instanceof Project == false + && fragment instanceof Filter == false + && fragment instanceof org.apache.calcite.rel.core.Sort == false)) { + return extractRowArrowSourceShape(originalFragment); + } + wrappers.add(fragment); + fragment = fragment.getInput(0); + } + + Aggregate aggregate = (Aggregate) fragment; + if (aggregate.getGroupSets() != null && aggregate.getGroupSets().size() > 1) { + return null; + } + RelNode below = aggregate.getInput(); + Project project = null; + if (below instanceof Project candidate) { + project = candidate; + below = candidate.getInput(); + } + Filter filter = null; + if (below instanceof Filter candidate) { + filter = candidate; + below = candidate.getInput(); + } + if (below instanceof OpenSearchRelNode == false || below.getInputs().isEmpty() == false) { + return null; + } + List storage = ((OpenSearchRelNode) below).getOutputFieldStorage(); + if (storage == null) { + return null; + } + + TreeSet referenced = new TreeSet<>(); + if (project != null) { + RexShuttle collector = inputReferenceCollector(referenced); + for (RexNode expression : project.getProjects()) { + expression.accept(collector); + } + } else { + for (int ordinal : aggregate.getGroupSet()) { + referenced.add(ordinal); + } + for (AggregateCall call : aggregate.getAggCallList()) { + referenced.addAll(call.getArgList()); + if (call.filterArg >= 0) { + return null; + } + } + } + // Preserve the metadata count fast path for COUNT(*). + if (referenced.isEmpty()) { + return null; + } + + RebasedInput input = rebaseInput(aggregate, below, storage, referenced); + if (input == null) { + return null; + } + RexShuttle remap = inputRemapper(input.oldToNew()); + + RelNode rebasedInput; + Aggregate rebasedAggregate; + if (project != null) { + List remapped = new ArrayList<>(project.getProjects().size()); + for (RexNode expression : project.getProjects()) { + remapped.add(expression.accept(remap)); + } + rebasedInput = LogicalProject.create(input.scan(), project.getHints(), remapped, project.getRowType().getFieldNames()); + rebasedAggregate = aggregate.copy( + aggregate.getTraitSet(), + rebasedInput, + aggregate.getGroupSet(), + aggregate.getGroupSets(), + aggregate.getAggCallList() + ); + } else { + ImmutableBitSet.Builder newGroupSet = ImmutableBitSet.builder(); + for (int ordinal : aggregate.getGroupSet()) { + newGroupSet.set(input.oldToNew()[ordinal]); + } + List remappedCalls = new ArrayList<>(aggregate.getAggCallList().size()); + for (AggregateCall call : aggregate.getAggCallList()) { + List remappedArguments = new ArrayList<>(call.getArgList().size()); + for (int argument : call.getArgList()) { + remappedArguments.add(input.oldToNew()[argument]); + } + remappedCalls.add(call.copy(remappedArguments, -1, call.collation)); + } + rebasedAggregate = aggregate.copy(aggregate.getTraitSet(), input.scan(), newGroupSet.build(), null, remappedCalls); + } + + RelNode rebased = rebasedAggregate; + for (int i = wrappers.size() - 1; i >= 0; i--) { + RelNode wrapper = wrappers.get(i); + rebased = wrapper.copy(wrapper.getTraitSet(), List.of(rebased)); + } + return new ArrowSourceShape(SOURCE_INPUT_ID, rebased, input.columns(), filter, resultNames(originalFragment)); + } + + private static ArrowSourceShape extractRowArrowSourceShape(RelNode fragment) { + Project topProject = null; + RelNode node = fragment; + if (node instanceof Project candidate) { + topProject = candidate; + node = candidate.getInput(); + } + org.apache.calcite.rel.core.Sort sort = null; + if (node instanceof org.apache.calcite.rel.core.Sort candidate) { + if (candidate.offset != null) { + return null; + } + sort = candidate; + node = candidate.getInput(); + } + Project middleProject = null; + if (node instanceof Project candidate) { + middleProject = candidate; + node = candidate.getInput(); + } + Filter filter = null; + if (node instanceof Filter candidate) { + filter = candidate; + node = candidate.getInput(); + } + if (node instanceof OpenSearchRelNode == false || node.getInputs().isEmpty() == false) { + return null; + } + if (topProject == null && sort == null && middleProject == null) { + return null; + } + List storage = ((OpenSearchRelNode) node).getOutputFieldStorage(); + if (storage == null) { + return null; + } + + TreeSet referenced = new TreeSet<>(); + RexShuttle collector = inputReferenceCollector(referenced); + if (middleProject != null) { + for (RexNode expression : middleProject.getProjects()) { + expression.accept(collector); + } + } else if (topProject != null) { + for (RexNode expression : topProject.getProjects()) { + expression.accept(collector); + } + if (sort != null) { + for (org.apache.calcite.rel.RelFieldCollation field : sort.getCollation().getFieldCollations()) { + referenced.add(field.getFieldIndex()); + } + } + } else { + for (int i = 0; i < storage.size(); i++) { + referenced.add(i); + } + } + + // A constant projection still needs one source column to preserve input row count. + if (referenced.isEmpty()) { + TreeSet filterReferences = new TreeSet<>(); + if (filter != null) { + filter.getCondition().accept(inputReferenceCollector(filterReferences)); + } + for (int ordinal : filterReferences) { + if (docValuesColumn(storage, ordinal) != null) { + referenced.add(ordinal); + break; + } + } + if (referenced.isEmpty()) { + for (int i = 0; i < storage.size(); i++) { + if (docValuesColumn(storage, i) != null) { + referenced.add(i); + break; + } + } + } + if (referenced.isEmpty()) { + return null; + } + } + + RebasedInput input = rebaseInput(fragment, node, storage, referenced); + if (input == null) { + return null; + } + RexShuttle remap = inputRemapper(input.oldToNew()); + RelNode rebased = input.scan(); + if (middleProject != null) { + List expressions = new ArrayList<>(middleProject.getProjects().size()); + for (RexNode expression : middleProject.getProjects()) { + expressions.add(expression.accept(remap)); + } + rebased = LogicalProject.create(rebased, middleProject.getHints(), expressions, middleProject.getRowType().getFieldNames()); + } + if (sort != null) { + org.apache.calcite.rel.RelCollation collation = sort.getCollation(); + if (middleProject == null) { + List fields = new ArrayList<>(); + for (org.apache.calcite.rel.RelFieldCollation field : collation.getFieldCollations()) { + fields.add(field.withFieldIndex(input.oldToNew()[field.getFieldIndex()])); + } + collation = org.apache.calcite.rel.RelCollations.of(fields); + } + rebased = org.apache.calcite.rel.logical.LogicalSort.create(rebased, collation, null, sort.fetch); + } + if (topProject != null) { + List expressions = new ArrayList<>(topProject.getProjects().size()); + for (RexNode expression : topProject.getProjects()) { + expressions.add(middleProject == null ? expression.accept(remap) : expression); + } + rebased = LogicalProject.create(rebased, topProject.getHints(), expressions, topProject.getRowType().getFieldNames()); + } + return new ArrowSourceShape(SOURCE_INPUT_ID, rebased, input.columns(), filter, resultNames(fragment)); + } + + private static RebasedInput rebaseInput( + RelNode owner, + RelNode originalScan, + List storage, + TreeSet referenced + ) { + List columns = new ArrayList<>(referenced.size()); + int[] oldToNew = new int[storage.size()]; + Arrays.fill(oldToNew, -1); + RelDataTypeFactory.Builder rowBuilder = owner.getCluster().getTypeFactory().builder(); + List rebasedStorage = new ArrayList<>(referenced.size()); + for (int ordinal : referenced) { + InputColumn column = docValuesColumn(storage, ordinal); + if (column == null || ordinal >= originalScan.getRowType().getFieldCount()) { + return null; + } + oldToNew[ordinal] = columns.size(); + columns.add(column); + RelDataTypeField field = originalScan.getRowType().getFieldList().get(ordinal); + rowBuilder.add(column.name(), field.getType()); + rebasedStorage.add(storage.get(ordinal)); + } + OpenSearchStageInputScan scan = new OpenSearchStageInputScan( + owner.getCluster(), + owner.getTraitSet(), + 0, + rowBuilder.build(), + List.of(), + rebasedStorage + ); + return new RebasedInput(scan, columns, oldToNew); + } + + private static RexShuttle inputReferenceCollector(TreeSet references) { + return new RexShuttle() { + @Override + public RexNode visitInputRef(RexInputRef ref) { + references.add(ref.getIndex()); + return ref; + } + }; + } + + private static RexShuttle inputRemapper(int[] oldToNew) { + return new RexShuttle() { + @Override + public RexNode visitInputRef(RexInputRef ref) { + int mapped = oldToNew[ref.getIndex()]; + if (mapped < 0) { + throw new IllegalStateException("unreferenced input ordinal [" + ref.getIndex() + "] survived source-plan rebasing"); + } + return new RexInputRef(mapped, ref.getType()); + } + }; + } + + static List resultNames(RelNode fragment) { + List names = fragment.getRowType().getFieldNames(); + List result = new ArrayList<>(names.size()); + for (int i = 0; i < names.size(); i++) { + result.add(names.get(i) == null ? "EXPR$" + i : names.get(i)); + } + return result; + } + + private record RebasedInput(OpenSearchStageInputScan scan, List columns, int[] oldToNew) { + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneFragmentWirePlan.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneFragmentWirePlan.java new file mode 100644 index 0000000000000..b1047dd1ba191 --- /dev/null +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneFragmentWirePlan.java @@ -0,0 +1,107 @@ +/* + * 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.lucene; + +import org.opensearch.analytics.spi.ArrowBatchSourcePlan; +import org.opensearch.common.io.stream.BytesStreamOutput; +import org.opensearch.core.common.bytes.BytesReference; +import org.opensearch.core.common.io.stream.NamedWriteableAwareStreamInput; +import org.opensearch.core.common.io.stream.NamedWriteableRegistry; +import org.opensearch.core.common.io.stream.StreamInput; +import org.opensearch.index.query.QueryBuilder; + +import java.io.IOException; +import java.util.List; + +/** Typed wire payload shared by Lucene fragment conversion and shard execution. */ +record LuceneFragmentWirePlan(List outputNames, byte[] filterBytes, ArrowBatchSourcePlan arrowSourcePlan) { + + LuceneFragmentWirePlan { + outputNames = List.copyOf(outputNames); + filterBytes = filterBytes == null ? null : filterBytes.clone(); + } + + @Override + public byte[] filterBytes() { + return filterBytes == null ? null : filterBytes.clone(); + } + + static LuceneFragmentWirePlan create(List outputNames, QueryBuilder filter, ArrowBatchSourcePlan arrowSourcePlan) { + return new LuceneFragmentWirePlan(outputNames, serializeFilter(filter), arrowSourcePlan); + } + + static LuceneFragmentWirePlan fromBytes(byte[] bytes) { + try (StreamInput input = StreamInput.wrap(bytes)) { + LuceneFragmentWirePlan plan = new LuceneFragmentWirePlan( + input.readStringList(), + readOptionalBytes(input), + input.readOptionalWriteable(ArrowBatchSourcePlan::new) + ); + if (input.available() != 0) { + throw new IllegalStateException("Unexpected trailing Lucene fragment bytes"); + } + return plan; + } catch (IOException e) { + throw new IllegalStateException("Failed to deserialize Lucene fragment", e); + } + } + + LuceneFragmentWirePlan withOutputNames(List names) { + return new LuceneFragmentWirePlan(names, filterBytes, arrowSourcePlan); + } + + LuceneFragmentWirePlan withArrowSourcePlan(ArrowBatchSourcePlan plan, List names) { + return new LuceneFragmentWirePlan(names, filterBytes, plan); + } + + QueryBuilder filterQuery(NamedWriteableRegistry registry) { + if (filterBytes == null) { + return null; + } + try (StreamInput rawInput = StreamInput.wrap(filterBytes)) { + StreamInput input = new NamedWriteableAwareStreamInput(rawInput, registry); + QueryBuilder filter = input.readNamedWriteable(QueryBuilder.class); + if (input.available() != 0) { + throw new IllegalStateException("Unexpected trailing Lucene filter bytes"); + } + return filter; + } catch (IOException e) { + throw new IllegalStateException("Failed to deserialize Lucene filter", e); + } + } + + byte[] toBytes() { + try (BytesStreamOutput output = new BytesStreamOutput()) { + output.writeStringCollection(outputNames); + output.writeBoolean(filterBytes != null); + if (filterBytes != null) { + output.writeByteArray(filterBytes); + } + output.writeOptionalWriteable(arrowSourcePlan); + return BytesReference.toBytes(output.bytes()); + } catch (IOException e) { + throw new IllegalStateException("Failed to serialize Lucene fragment", e); + } + } + + private static byte[] readOptionalBytes(StreamInput input) throws IOException { + return input.readBoolean() ? input.readByteArray() : null; + } + + private static byte[] serializeFilter(QueryBuilder filter) { + if (filter == null) { + return null; + } + try (BytesStreamOutput output = new BytesStreamOutput()) { + output.writeNamedWriteable(filter); + return BytesReference.toBytes(output.bytes()); + } catch (IOException e) { + throw new IllegalStateException("Failed to serialize Lucene filter", e); + } + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneQueryConversionUtils.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneQueryConversionUtils.java index a7f826693072a..b3fa4bf52fee5 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneQueryConversionUtils.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneQueryConversionUtils.java @@ -21,6 +21,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.function.Predicate; /** * Helpers for adapting Lucene {@link Query} trees so they can execute against the @@ -62,21 +63,29 @@ private LuceneQueryConversionUtils() {} * when nothing changed */ public static Query rewriteFieldExistsForSecondary(Query query) { + return rewriteFieldExistsForSecondary(query, field -> false); + } + + /** Preserves native existence queries for fields whose reader has doc values. */ + static Query rewriteFieldExistsForSecondary(Query query, Predicate hasDocValues) { // IndexOrDocValuesQuery wraps an index query + a doc-values query; the secondary has no // doc-values, so unwrap to just the index query (TermRangeQuery against the term dictionary). if (query instanceof IndexOrDocValuesQuery idv) { - return rewriteFieldExistsForSecondary(idv.getIndexQuery()); + return rewriteFieldExistsForSecondary(idv.getIndexQuery(), hasDocValues); } if (query instanceof FieldExistsQuery fieldExists) { + if (hasDocValues.test(fieldExists.getField())) { + return fieldExists; + } // null lower/upper bound = unbounded both ends = "any term present for this field". return new TermRangeQuery(fieldExists.getField(), null, null, true, true); } if (query instanceof ConstantScoreQuery constantScore) { - Query inner = rewriteFieldExistsForSecondary(constantScore.getQuery()); + Query inner = rewriteFieldExistsForSecondary(constantScore.getQuery(), hasDocValues); return inner == constantScore.getQuery() ? constantScore : new ConstantScoreQuery(inner); } if (query instanceof BoostQuery boost) { - Query inner = rewriteFieldExistsForSecondary(boost.getQuery()); + Query inner = rewriteFieldExistsForSecondary(boost.getQuery(), hasDocValues); return inner == boost.getQuery() ? boost : new BoostQuery(inner, boost.getBoost()); } if (query instanceof BooleanQuery bool) { @@ -84,7 +93,7 @@ public static Query rewriteFieldExistsForSecondary(Query query) { builder.setMinimumNumberShouldMatch(bool.getMinimumNumberShouldMatch()); boolean changed = false; for (BooleanClause clause : bool.clauses()) { - Query rewritten = rewriteFieldExistsForSecondary(clause.query()); + Query rewritten = rewriteFieldExistsForSecondary(clause.query(), hasDocValues); changed |= rewritten != clause.query(); builder.add(rewritten, clause.occur()); } @@ -94,7 +103,7 @@ public static Query rewriteFieldExistsForSecondary(Query query) { List rewritten = new ArrayList<>(disjunctionMax.getDisjuncts().size()); boolean changed = false; for (Query disjunct : disjunctionMax.getDisjuncts()) { - Query r = rewriteFieldExistsForSecondary(disjunct); + Query r = rewriteFieldExistsForSecondary(disjunct, hasDocValues); changed |= r != disjunct; rewritten.add(r); } diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneScanInstructionHandler.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneScanInstructionHandler.java index 998a5f79637f9..54510d944483c 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneScanInstructionHandler.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneScanInstructionHandler.java @@ -10,27 +10,28 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.apache.lucene.index.DocValuesType; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.MatchAllDocsQuery; import org.apache.lucene.search.Query; import org.opensearch.analytics.backend.ShardScanExecutionContext; +import org.opensearch.analytics.spi.ArrowBatchSourcePlan; import org.opensearch.analytics.spi.BackendExecutionContext; import org.opensearch.analytics.spi.CommonExecutionContext; import org.opensearch.analytics.spi.FragmentInstructionHandler; import org.opensearch.analytics.spi.ShardScanInstructionNode; -import org.opensearch.core.common.io.stream.NamedWriteableAwareStreamInput; -import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.index.engine.exec.IndexReaderProvider; import org.opensearch.index.query.QueryBuilder; import org.opensearch.index.query.QueryShardContext; import java.io.IOException; +import java.util.List; /** * Lucene-side shard-scan instruction handler. Reads a {@link ShardScanInstructionNode} - * produced for a Lucene {@code StagePlan}, acquires the shard's Lucene reader, deserialises - * the filter {@link QueryBuilder} from {@code ShardScanExecutionContext.getFragmentBytes()}, - * compiles it to a Lucene {@link Query}, and returns a {@link LuceneSearcherState} for + * produced for a Lucene {@code StagePlan}, acquires the shard's Lucene reader, decodes its + * {@link LuceneFragmentWirePlan}, compiles the optional filter to a Lucene {@link Query}, + * and returns a {@link LuceneSearcherState} for * {@link LuceneSearchExecEngine} to execute. * *

Empty {@code fragmentBytes} → {@link MatchAllDocsQuery} (count(*) over the whole shard). @@ -63,46 +64,42 @@ public BackendExecutionContext apply( IndexSearcher searcher = luceneReader.searcher(shardCtx.getQueryCache(), shardCtx.getQueryCachingPolicy()); Decoded decoded = decodeFragmentBytes(shardCtx, searcher); LOGGER.debug( - "[lucene-count] shardId={} filterQuery={} columnNames={}", + "[lucene-scan] shardId={} filterQuery={} columnNames={} arrowSourcePlan={}", shardCtx.getShardId(), decoded.filterQuery, - decoded.columnNames + decoded.outputNames, + decoded.arrowSourcePlan != null ); - return new LuceneSearcherState(searcher, decoded.filterQuery, decoded.columnNames); + return new LuceneSearcherState(searcher, decoded.filterQuery, decoded.outputNames, decoded.arrowSourcePlan); } - /** - * Deserializes the wire format produced by {@link LuceneFragmentConvertor#convertFragment}: - * {@code [columnNames String[]] [hasFilter boolean] [QueryBuilder NamedWriteable]?}. - * Empty bytes → no filter, no column names (legacy/defensive fallback that shouldn't - * happen on the Lucene-driver path but stays safe if the wire shape ever drifts). - */ private Decoded decodeFragmentBytes(ShardScanExecutionContext shardCtx, IndexSearcher searcher) { byte[] bytes = shardCtx.getFragmentBytes(); if (bytes == null || bytes.length == 0) { - return new Decoded(new MatchAllDocsQuery(), java.util.List.of()); + return new Decoded(new MatchAllDocsQuery(), List.of(), null); } - try (StreamInput rawInput = StreamInput.wrap(bytes)) { - StreamInput input = new NamedWriteableAwareStreamInput(rawInput, shardCtx.getNamedWriteableRegistry()); - java.util.List columnNames = input.readStringList(); - boolean hasFilter = input.readBoolean(); - Query filterQuery; - if (hasFilter) { - QueryShardContext qsc = LuceneAnalyticsBackendPlugin.buildMinimalQueryShardContext(shardCtx, searcher); - QueryBuilder queryBuilder = input.readNamedWriteable(QueryBuilder.class); - // Rewrite FieldExistsQuery → postings-only equivalent for the doc-values-less - // lucene-secondary segment (same reason as the filter-delegation path). This covers - // the Lucene-driver scan path (count + non-count) executed by LuceneSearchExecEngine. - filterQuery = LuceneQueryConversionUtils.rewriteFieldExistsForSecondary(queryBuilder.toQuery(qsc)); - } else { - filterQuery = new MatchAllDocsQuery(); + LuceneFragmentWirePlan wirePlan = LuceneFragmentWirePlan.fromBytes(bytes); + QueryBuilder queryBuilder = wirePlan.filterQuery(shardCtx.getNamedWriteableRegistry()); + Query filterQuery; + if (queryBuilder == null) { + filterQuery = new MatchAllDocsQuery(); + } else { + QueryShardContext queryContext = LuceneAnalyticsBackendPlugin.buildMinimalQueryShardContext(shardCtx, searcher); + try { + filterQuery = LuceneQueryConversionUtils.rewriteFieldExistsForSecondary( + queryBuilder.toQuery(queryContext), + field -> searcher.getIndexReader().leaves().stream().anyMatch(leaf -> { + var fieldInfo = leaf.reader().getFieldInfos().fieldInfo(field); + return fieldInfo != null && fieldInfo.getDocValuesType() != DocValuesType.NONE; + }) + ); + } catch (IOException e) { + throw new IllegalStateException("Failed to compile Lucene fragment filter", e); } - return new Decoded(filterQuery, columnNames); - } catch (IOException e) { - throw new IllegalStateException("Failed to deserialize Lucene-driver fragment bytes", e); } + return new Decoded(filterQuery, wirePlan.outputNames(), wirePlan.arrowSourcePlan()); } - private record Decoded(Query filterQuery, java.util.List columnNames) { + private record Decoded(Query filterQuery, List outputNames, ArrowBatchSourcePlan arrowSourcePlan) { } } diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java index 701840299c8b1..90371f6a3f829 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneSearchExecEngine.java @@ -24,6 +24,8 @@ import org.opensearch.analytics.backend.EngineResultStream; import org.opensearch.analytics.backend.SearchExecEngine; import org.opensearch.analytics.backend.ShardScanExecutionContext; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; +import org.opensearch.analytics.spi.ArrowBatchSourcePlan; import java.io.IOException; import java.util.ArrayList; @@ -35,12 +37,10 @@ * instruction handler, executes the operation, and returns an {@link EngineResultStream} * the framework drains into the Flight transport. * - *

Today's only operation is the count fast path — - * {@link org.apache.lucene.search.IndexSearcher#count(org.apache.lucene.search.Query)} — - * exported through the Arrow C-Data interface so the result VSR has the same - * foreign-allocation-managed buffer layout DataFusion's result stream produces. Pure-Java - * {@code setSafe}-built VSRs don't survive Flight's {@code VectorTransfer.transferRoot}; - * see {@link LuceneResultStream} for the detailed comparison. + *

Count-only states use the metadata fast path through + * {@link org.apache.lucene.search.IndexSearcher#count(org.apache.lucene.search.Query)}. + * States carrying an {@link ArrowBatchSourcePlan} create one doc-values source factory and + * transfer it to the bound Arrow source execution backend. * *

No deletes gate. {@code IndexSearcher.count} is self-healing: per-leaf * {@code Weight.count(leaf)} returns -1 on dirty leaves and falls back to full iteration — @@ -54,9 +54,15 @@ final class LuceneSearchExecEngine implements SearchExecEngineMirrors the role {@code DataFusionSessionState} plays for the DataFusion backend — * a small immutable state record threaded from instruction handler to search engine. @@ -36,13 +38,24 @@ final class LuceneSearcherState implements BackendExecutionContext { private final Query filterQuery; /** Aggregate-call output names — one Int64 column per name in the result Arrow batch. */ private final List outputColumnNames; + private final ArrowBatchSourcePlan arrowBatchSourcePlan; LuceneSearcherState(IndexSearcher searcher, Query filterQuery, List outputColumnNames) { + this(searcher, filterQuery, outputColumnNames, null); + } + + LuceneSearcherState( + IndexSearcher searcher, + Query filterQuery, + List outputColumnNames, + ArrowBatchSourcePlan arrowBatchSourcePlan + ) { this.searcher = Objects.requireNonNull(searcher, "searcher"); // Never null — see field javadoc. Caller must substitute MatchAllDocsQuery when the // fragment has no filter so the search engine doesn't have to branch. this.filterQuery = Objects.requireNonNull(filterQuery, "filterQuery (use MatchAllDocsQuery for no-filter fragments)"); this.outputColumnNames = List.copyOf(Objects.requireNonNull(outputColumnNames, "outputColumnNames")); + this.arrowBatchSourcePlan = arrowBatchSourcePlan; } IndexSearcher searcher() { @@ -56,4 +69,8 @@ Query filterQuery() { List outputColumnNames() { return outputColumnNames; } + + ArrowBatchSourcePlan arrowBatchSourcePlan() { + return arrowBatchSourcePlan; + } } diff --git a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneShardPreference.java b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneShardPreference.java index cb1c01aa500ee..c7bef5400c466 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneShardPreference.java +++ b/sandbox/plugins/analytics-backend-lucene/src/main/java/org/opensearch/be/lucene/LuceneShardPreference.java @@ -13,6 +13,7 @@ import org.opensearch.analytics.spi.ShardPreferenceContext; import java.util.OptionalInt; +import java.util.function.BooleanSupplier; /** * Lucene's per-shard preference: opt in to drive count-fast-path fragments when the user has @@ -26,6 +27,16 @@ */ final class LuceneShardPreference implements BackendShardPreference { + private final BooleanSupplier arrowSourceAvailable; + + LuceneShardPreference() { + this(() -> true); + } + + LuceneShardPreference(BooleanSupplier arrowSourceAvailable) { + this.arrowSourceAvailable = arrowSourceAvailable; + } + /** Wants-to-drive score — beats generic alternatives (score 0). */ private static final int COUNT_FAST_PATH_SCORE = 100; @@ -34,12 +45,19 @@ final class LuceneShardPreference implements BackendShardPreference { * Lucene alternative just because it appeared first in PlanForker order. */ private static final int NOT_DRIVABLE_SCORE = -1; + /** Prefer supported doc-values plans over generic storage alternatives. */ + private static final int ARROW_SOURCE_SCORE = 50; + @Override public OptionalInt scoreFor(RelNode fragment, ShardPreferenceContext ctx) { if (ctx.preferMetadataDriver() == false) return OptionalInt.empty(); - if (LuceneFragmentConvertor.isCountFastPath(fragment) == false) { - return OptionalInt.of(NOT_DRIVABLE_SCORE); + LuceneFragmentPlanner.Shape shape = LuceneFragmentPlanner.classify(fragment); + if (shape instanceof LuceneFragmentPlanner.CountShape) { + return OptionalInt.of(COUNT_FAST_PATH_SCORE); + } + if (shape instanceof LuceneFragmentPlanner.ArrowSourceShape && arrowSourceAvailable.getAsBoolean()) { + return OptionalInt.of(ARROW_SOURCE_SCORE); } - return OptionalInt.of(COUNT_FAST_PATH_SCORE); + return OptionalInt.of(NOT_DRIVABLE_SCORE); } } 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 0236065560429..18bb714f43e66 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 @@ -25,6 +25,12 @@ import org.apache.calcite.sql.type.OperandTypes; import org.apache.calcite.sql.type.ReturnTypes; import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.lucene.analysis.standard.StandardAnalyzer; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.store.ByteBuffersDirectory; import org.opensearch.action.search.TransportSearchAction; import org.opensearch.analytics.planner.CapabilityRegistry; import org.opensearch.analytics.planner.FieldStorageResolver; @@ -65,6 +71,7 @@ import org.opensearch.cluster.routing.OperationRouting; import org.opensearch.cluster.routing.ShardIterator; import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.concurrent.GatedCloseable; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; @@ -72,8 +79,14 @@ import org.opensearch.core.common.io.stream.NamedWriteableRegistry; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.index.Index; +import org.opensearch.index.engine.Engine; +import org.opensearch.index.engine.EngineBackedIndexer; +import org.opensearch.index.engine.dataformat.DataFormat; +import org.opensearch.index.engine.exec.IndexReaderProvider; +import org.opensearch.index.engine.exec.coord.CatalogSnapshot; import org.opensearch.index.query.MatchQueryBuilder; import org.opensearch.index.query.QueryBuilder; +import org.opensearch.index.shard.IndexShard; import org.opensearch.test.OpenSearchTestCase; import java.io.IOException; @@ -83,6 +96,7 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.atomic.AtomicInteger; import java.util.function.Function; import static org.mockito.ArgumentMatchers.any; @@ -123,6 +137,72 @@ public void setUp() throws Exception { cluster = RelOptCluster.create(new HepPlanner(new HepProgramBuilder().build()), rexBuilder); } + public void testStandardEngineReaderAdapterClosesSearcherAndSnapshot() throws Exception { + try ( + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer())); + DirectoryReader directoryReader = DirectoryReader.open(writer) + ) { + AtomicInteger searcherCloses = new AtomicInteger(); + AtomicInteger snapshotCloses = new AtomicInteger(); + Engine.Searcher searcher = new Engine.Searcher( + "analytics-lucene-test", + directoryReader, + IndexSearcher.getDefaultSimilarity(), + IndexSearcher.getDefaultQueryCache(), + IndexSearcher.getDefaultQueryCachingPolicy(), + searcherCloses::incrementAndGet + ); + CatalogSnapshot snapshot = mock(CatalogSnapshot.class); + GatedCloseable snapshotRef = new GatedCloseable<>(snapshot, snapshotCloses::incrementAndGet); + EngineBackedIndexer indexer = mock(EngineBackedIndexer.class); + when(indexer.acquireSnapshot()).thenReturn(snapshotRef); + IndexShard shard = mock(IndexShard.class); + when(shard.getReaderProvider()).thenReturn(indexer); + when(shard.acquireSearcher("analytics-lucene")).thenReturn(searcher); + DataFormat format = mock(DataFormat.class); + when(format.name()).thenReturn("lucene"); + LucenePlugin plugin = mock(LucenePlugin.class); + when(plugin.getDataFormat()).thenReturn(format); + + GatedCloseable acquired = new LuceneAnalyticsBackendPlugin(plugin).acquireReader(shard); + assertSame(snapshot, acquired.get().catalogSnapshot()); + assertNotNull(acquired.get().getReader(format, LuceneReader.class)); + acquired.close(); + acquired.close(); + + assertEquals(1, searcherCloses.get()); + assertEquals(1, snapshotCloses.get()); + } + } + + public void testStandardEngineReaderAdapterClosesSearcherWhenSnapshotSetupFails() throws Exception { + try ( + ByteBuffersDirectory directory = new ByteBuffersDirectory(); + IndexWriter writer = new IndexWriter(directory, new IndexWriterConfig(new StandardAnalyzer())); + DirectoryReader directoryReader = DirectoryReader.open(writer) + ) { + AtomicInteger searcherCloses = new AtomicInteger(); + Engine.Searcher searcher = new Engine.Searcher( + "analytics-lucene-test", + directoryReader, + IndexSearcher.getDefaultSimilarity(), + IndexSearcher.getDefaultQueryCache(), + IndexSearcher.getDefaultQueryCachingPolicy(), + searcherCloses::incrementAndGet + ); + EngineBackedIndexer indexer = mock(EngineBackedIndexer.class); + when(indexer.acquireSnapshot()).thenThrow(new IllegalStateException("snapshot setup failed")); + IndexShard shard = mock(IndexShard.class); + when(shard.getReaderProvider()).thenReturn(indexer); + when(shard.acquireSearcher("analytics-lucene")).thenReturn(searcher); + LucenePlugin plugin = mock(LucenePlugin.class); + + expectThrows(IllegalStateException.class, () -> new LuceneAnalyticsBackendPlugin(plugin).acquireReader(shard)); + assertEquals(1, searcherCloses.get()); + } + } + /** * MATCH(message, 'hello world') through full pipeline → delegatedQueries contains * valid MatchQueryBuilder bytes with correct field name and query text. diff --git a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneArrowSourcePlanTests.java b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneArrowSourcePlanTests.java new file mode 100644 index 0000000000000..a8d79b170b894 --- /dev/null +++ b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneArrowSourcePlanTests.java @@ -0,0 +1,269 @@ +/* + * 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.lucene; + +import org.apache.arrow.memory.BufferAllocator; +import org.apache.calcite.jdbc.JavaTypeFactoryImpl; +import org.apache.calcite.plan.RelOptCluster; +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.AggregateCall; +import org.apache.calcite.rel.logical.LogicalAggregate; +import org.apache.calcite.rel.logical.LogicalProject; +import org.apache.calcite.rel.type.RelDataType; +import org.apache.calcite.rel.type.RelDataTypeFactory; +import org.apache.calcite.rex.RexBuilder; +import org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.calcite.sql.type.SqlTypeName; +import org.apache.calcite.util.ImmutableBitSet; +import org.opensearch.analytics.backend.EngineResultStream; +import org.opensearch.analytics.planner.rel.OpenSearchStageInputScan; +import org.opensearch.analytics.spi.AnalyticsSearchBackendPlugin; +import org.opensearch.analytics.spi.ArrowBatchSourceFactory; +import org.opensearch.analytics.spi.ArrowBatchSourcePlan; +import org.opensearch.analytics.spi.DelegationThreadTracker; +import org.opensearch.analytics.spi.FieldStorageInfo; +import org.opensearch.analytics.spi.FieldType; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.List; +import java.util.Map; + +public class LuceneArrowSourcePlanTests extends OpenSearchTestCase { + + private RelDataTypeFactory typeFactory; + private RexBuilder rexBuilder; + private RelOptCluster cluster; + + @Override + public void setUp() throws Exception { + super.setUp(); + typeFactory = new JavaTypeFactoryImpl(); + rexBuilder = new RexBuilder(typeFactory); + cluster = RelOptCluster.create(new HepPlanner(new HepProgramBuilder().build()), rexBuilder); + } + + public void testAggregateShapeUsesOnlyReferencedDocValuesColumns() { + RelDataType bigint = nullable(SqlTypeName.BIGINT); + RelDataType rowType = typeFactory.builder().add("key", bigint).add("metric", bigint).add("unused", bigint).build(); + RelNode scan = scan( + rowType, + List.of(storage("key", FieldType.LONG), storage("metric", FieldType.LONG), storage("unused", FieldType.LONG)) + ); + AggregateCall sum = AggregateCall.create(SqlStdOperatorTable.SUM, false, List.of(1), -1, scan, bigint, "sum_metric"); + RelNode aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(0), null, List.of(sum)); + + LuceneFragmentPlanner.ArrowSourceShape shape = arrowSourceShape(aggregate); + + assertNotNull(shape); + assertEquals(List.of("key", "metric"), shape.inputColumns().stream().map(ArrowBatchSourceFactory.InputColumn::name).toList()); + assertEquals(List.of("key", "sum_metric"), shape.outputNames()); + assertTrue(shape.rebasedFragment().getInput(0) instanceof OpenSearchStageInputScan); + assertEquals(List.of("key", "metric"), shape.rebasedFragment().getInput(0).getRowType().getFieldNames()); + } + + public void testCountFieldUsesArrowSourceForNullSemantics() { + RelDataType bigint = nullable(SqlTypeName.BIGINT); + RelNode scan = scan(typeFactory.builder().add("metric", bigint).build(), List.of(storage("metric", FieldType.LONG))); + AggregateCall count = AggregateCall.create( + SqlStdOperatorTable.COUNT, + false, + List.of(0), + -1, + scan, + typeFactory.createSqlType(SqlTypeName.BIGINT), + "count_metric" + ); + RelNode aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(count)); + + assertTrue(LuceneFragmentPlanner.classify(aggregate) instanceof LuceneFragmentPlanner.ArrowSourceShape); + } + + public void testCountStarRetainsMetadataFastPath() { + RelDataType bigint = nullable(SqlTypeName.BIGINT); + RelNode scan = scan(typeFactory.builder().add("metric", bigint).build(), List.of(storage("metric", FieldType.LONG))); + AggregateCall count = AggregateCall.create( + SqlStdOperatorTable.COUNT, + false, + List.of(), + -1, + scan, + typeFactory.createSqlType(SqlTypeName.BIGINT), + "count_star" + ); + RelNode aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(count)); + + assertTrue(LuceneFragmentPlanner.classify(aggregate) instanceof LuceneFragmentPlanner.CountShape); + } + + public void testRowProjectionRebasesKeywordInput() { + RelDataType varchar = nullable(SqlTypeName.VARCHAR); + RelDataType rowType = typeFactory.builder().add("keyword", varchar).add("unused", varchar).build(); + RelNode scan = scan(rowType, List.of(storage("keyword", FieldType.KEYWORD), storage("unused", FieldType.KEYWORD))); + RelNode project = LogicalProject.create(scan, List.of(), List.of(rexBuilder.makeInputRef(scan, 0)), List.of("keyword")); + + LuceneFragmentPlanner.ArrowSourceShape shape = arrowSourceShape(project); + + assertNotNull(shape); + assertEquals(List.of("keyword"), shape.inputColumns().stream().map(ArrowBatchSourceFactory.InputColumn::name).toList()); + assertEquals(ArrowBatchSourceFactory.ColumnKind.KEYWORD, shape.inputColumns().getFirst().kind()); + } + + public void testTimestampProjectionUsesTimestampColumnKind() { + RelDataType timestamp = nullable(SqlTypeName.TIMESTAMP); + RelNode scan = scan(typeFactory.builder().add("event_time", timestamp).build(), List.of(storage("event_time", FieldType.DATE))); + RelNode project = LogicalProject.create(scan, List.of(), List.of(rexBuilder.makeInputRef(scan, 0)), List.of("event_time")); + + LuceneFragmentPlanner.ArrowSourceShape shape = arrowSourceShape(project); + + assertNotNull(shape); + assertEquals(ArrowBatchSourceFactory.ColumnKind.TIMESTAMP, shape.inputColumns().getFirst().kind()); + } + + public void testAdditionalScalarColumnKinds() { + Map kinds = Map.of( + FieldType.BOOLEAN, + ArrowBatchSourceFactory.ColumnKind.BOOLEAN, + FieldType.FLOAT, + ArrowBatchSourceFactory.ColumnKind.FLOAT, + FieldType.DOUBLE, + ArrowBatchSourceFactory.ColumnKind.DOUBLE, + FieldType.BINARY, + ArrowBatchSourceFactory.ColumnKind.BINARY, + FieldType.IP, + ArrowBatchSourceFactory.ColumnKind.IP + ); + for (Map.Entry entry : kinds.entrySet()) { + SqlTypeName sqlType = switch (entry.getKey()) { + case BOOLEAN -> SqlTypeName.BOOLEAN; + case FLOAT -> SqlTypeName.FLOAT; + case DOUBLE -> SqlTypeName.DOUBLE; + case BINARY, IP -> SqlTypeName.VARBINARY; + default -> throw new AssertionError(entry.getKey()); + }; + RelDataType type = nullable(sqlType); + RelNode scan = scan(typeFactory.builder().add("value", type).build(), List.of(storage("value", entry.getKey()))); + RelNode project = LogicalProject.create(scan, List.of(), List.of(rexBuilder.makeInputRef(scan, 0)), List.of("value")); + + ArrowBatchSourceFactory.InputColumn column = arrowSourceShape(project).inputColumns().getFirst(); + assertEquals(entry.getValue(), column.kind()); + assertFalse(column.multiValued()); + } + } + + public void testUnsupportedInputTypeDoesNotCreateSourcePlan() { + RelDataType integer = nullable(SqlTypeName.INTEGER); + RelNode scan = scan(typeFactory.builder().add("value", integer).build(), List.of(storage("value", FieldType.INTEGER))); + RelNode project = LogicalProject.create(scan, List.of(), List.of(rexBuilder.makeInputRef(scan, 0)), List.of("value")); + + assertTrue(LuceneFragmentPlanner.classify(project) instanceof LuceneFragmentPlanner.UnsupportedShape); + } + + public void testAttachedOperatorUpdatesCompiledPlanAndOutputNames() throws Exception { + RelDataType bigint = nullable(SqlTypeName.BIGINT); + RelNode scan = scan(typeFactory.builder().add("metric", bigint).build(), List.of(storage("metric", FieldType.LONG))); + AggregateCall sum = AggregateCall.create(SqlStdOperatorTable.SUM, false, List.of(0), -1, scan, bigint, "sum_metric"); + RelNode aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(sum)); + RelNode project = LogicalProject.create(aggregate, List.of(), List.of(rexBuilder.makeInputRef(aggregate, 0)), List.of("renamed")); + RecordingBackend backend = new RecordingBackend(); + LuceneFragmentConvertor convertor = new LuceneFragmentConvertor(Map.of(), backend); + byte[] inner = convertor.convertFragment(aggregate); + byte[] attached = convertor.attachFragmentOnTop(project, inner); + + assertNotNull(backend.attachedFragment); + LuceneFragmentWirePlan wirePlan = LuceneFragmentWirePlan.fromBytes(attached); + assertArrayEquals(new byte[] { 4, 5, 6 }, wirePlan.arrowSourcePlan().planBytes()); + assertEquals(List.of("renamed"), wirePlan.outputNames()); + assertNull(wirePlan.filterBytes()); + } + + public void testPartialAggregateCompilesAndSerializesArrowSourcePlan() throws Exception { + RelDataType bigint = nullable(SqlTypeName.BIGINT); + RelNode scan = scan(typeFactory.builder().add("metric", bigint).build(), List.of(storage("metric", FieldType.LONG))); + AggregateCall sum = AggregateCall.create(SqlStdOperatorTable.SUM, false, List.of(0), -1, scan, bigint, "sum_metric"); + RelNode aggregate = LogicalAggregate.create(scan, List.of(), ImmutableBitSet.of(), null, List.of(sum)); + RecordingBackend backend = new RecordingBackend(); + LuceneFragmentConvertor convertor = new LuceneFragmentConvertor(Map.of(), backend); + byte[] inner = convertor.convertFragment(scan); + byte[] bytes = convertor.attachPartialAggOnTop(aggregate, inner); + + assertTrue(backend.partialAggregate); + assertNotNull(backend.compiledFragment); + LuceneFragmentWirePlan wirePlan = LuceneFragmentWirePlan.fromBytes(bytes); + assertEquals("input-0", wirePlan.arrowSourcePlan().inputId()); + assertEquals( + List.of(new ArrowBatchSourceFactory.InputColumn("metric", ArrowBatchSourceFactory.ColumnKind.LONG)), + wirePlan.arrowSourcePlan().inputColumns() + ); + assertNull(wirePlan.filterBytes()); + } + + private static LuceneFragmentPlanner.ArrowSourceShape arrowSourceShape(RelNode fragment) { + LuceneFragmentPlanner.Shape shape = LuceneFragmentPlanner.classify(fragment); + assertTrue(shape instanceof LuceneFragmentPlanner.ArrowSourceShape); + return (LuceneFragmentPlanner.ArrowSourceShape) shape; + } + + private RelNode scan(RelDataType rowType, List storage) { + return new OpenSearchStageInputScan(cluster, cluster.traitSet(), 0, rowType, List.of("lucene"), storage); + } + + private FieldStorageInfo storage(String name, FieldType type) { + return new FieldStorageInfo( + name, + type.name().toLowerCase(java.util.Locale.ROOT), + type, + List.of("lucene"), + List.of(), + List.of(), + false + ); + } + + private RelDataType nullable(SqlTypeName type) { + return typeFactory.createTypeWithNullability(typeFactory.createSqlType(type), true); + } + + private static final class RecordingBackend implements AnalyticsSearchBackendPlugin { + private RelNode compiledFragment; + private RelNode attachedFragment; + private boolean partialAggregate; + + @Override + public String name() { + return "recording"; + } + + @Override + public byte[] compileArrowBatchSourcePlan(RelNode fragment, boolean partial) { + compiledFragment = fragment; + partialAggregate = partial; + return new byte[] { 1, 2, 3 }; + } + + @Override + public byte[] attachArrowBatchSourcePlan(RelNode fragment, byte[] innerPlanBytes) { + attachedFragment = fragment; + assertArrayEquals(new byte[] { 1, 2, 3 }, innerPlanBytes); + return new byte[] { 4, 5, 6 }; + } + + @Override + public EngineResultStream executeArrowBatchSource( + BufferAllocator resultAllocator, + ArrowBatchSourcePlan plan, + ArrowBatchSourceFactory sourceFactory, + Task task, + DelegationThreadTracker threadTracker + ) { + throw new UnsupportedOperationException(); + } + } +} diff --git a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneCanDriveFragmentTests.java b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneCanDriveFragmentTests.java index 5e3638b45a7d2..106cf08561dda 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneCanDriveFragmentTests.java +++ b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneCanDriveFragmentTests.java @@ -34,9 +34,9 @@ import static org.mockito.Mockito.when; /** - * Pins {@link LuceneFragmentConvertor#isCountFastPath}: drivable iff top is an Aggregate - * with empty group-set and every call is {@code SqlKind.COUNT}. Read by - * {@link LuceneShardPreference} to score this fragment for the count-fast-path. Guards + * Pins {@link LuceneFragmentPlanner#isCountFastPath}: drivable iff top is an Aggregate + * with empty group-set and every call is {@code COUNT(*)}. Read by + * {@link LuceneShardPreference} to score this fragment for the count fast path. Guards * capability-declaration drift — PlanForker already narrows by declared caps, this is the * second line. */ @@ -55,11 +55,10 @@ public void setUp() throws Exception { public void testCountStarOverEmptyGroupSet_drivable() { TableScan scan = stubScan("status", SqlTypeName.VARCHAR); RelNode agg = aggregate(scan, ImmutableBitSet.of(), countStar(scan)); - assertTrue("COUNT(*) with empty group-set is the canonical Lucene-driver shape", LuceneFragmentConvertor.isCountFastPath(agg)); + assertTrue("COUNT(*) with empty group-set is the canonical Lucene-driver shape", LuceneFragmentPlanner.isCountFastPath(agg)); } - public void testCountFieldOverEmptyGroupSet_drivable() { - // count(field) — same SqlKind.COUNT, just with a field arg. + public void testCountFieldOverEmptyGroupSet_notFastPath() { TableScan scan = stubScan("status", SqlTypeName.VARCHAR); AggregateCall countField = AggregateCall.create( SqlStdOperatorTable.COUNT, @@ -70,9 +69,9 @@ public void testCountFieldOverEmptyGroupSet_drivable() { typeFactory.createSqlType(SqlTypeName.BIGINT), "cnt_status" ); - assertTrue( - "count(field) is also drivable — same SqlKind.COUNT", - LuceneFragmentConvertor.isCountFastPath(aggregate(scan, ImmutableBitSet.of(), countField)) + assertFalse( + "COUNT(field) needs doc values to preserve null semantics", + LuceneFragmentPlanner.isCountFastPath(aggregate(scan, ImmutableBitSet.of(), countField)) ); } @@ -80,7 +79,7 @@ public void testSumOverEmptyGroupSet_notDrivable() { TableScan scan = stubScan("size", SqlTypeName.INTEGER); assertFalse( "SUM needs column values Lucene can't materialise — must be rejected", - LuceneFragmentConvertor.isCountFastPath( + LuceneFragmentPlanner.isCountFastPath( aggregate(scan, ImmutableBitSet.of(), nullableNumeric(SqlStdOperatorTable.SUM, scan, "total_size")) ) ); @@ -90,7 +89,7 @@ public void testMinOverEmptyGroupSet_notDrivable() { TableScan scan = stubScan("size", SqlTypeName.INTEGER); assertFalse( "MIN must be rejected", - LuceneFragmentConvertor.isCountFastPath( + LuceneFragmentPlanner.isCountFastPath( aggregate(scan, ImmutableBitSet.of(), nullableNumeric(SqlStdOperatorTable.MIN, scan, "min_size")) ) ); @@ -100,7 +99,7 @@ public void testMaxOverEmptyGroupSet_notDrivable() { TableScan scan = stubScan("size", SqlTypeName.INTEGER); assertFalse( "MAX must be rejected", - LuceneFragmentConvertor.isCountFastPath( + LuceneFragmentPlanner.isCountFastPath( aggregate(scan, ImmutableBitSet.of(), nullableNumeric(SqlStdOperatorTable.MAX, scan, "max_size")) ) ); @@ -110,7 +109,7 @@ public void testCountPlusSum_mixedAggregate_notDrivable() { // Even one non-COUNT call disqualifies the whole aggregate — every call must be COUNT. TableScan scan = stubScan("size", SqlTypeName.INTEGER); RelNode agg = aggregate(scan, ImmutableBitSet.of(), countStar(scan), nullableNumeric(SqlStdOperatorTable.SUM, scan, "total_size")); - assertFalse("COUNT(*) + SUM mixed must be rejected", LuceneFragmentConvertor.isCountFastPath(agg)); + assertFalse("COUNT(*) + SUM mixed must be rejected", LuceneFragmentPlanner.isCountFastPath(agg)); } public void testCountWithGroupBy_notDrivable() { @@ -118,7 +117,7 @@ public void testCountWithGroupBy_notDrivable() { TableScan scan = stubScan("status", SqlTypeName.VARCHAR); assertFalse( "COUNT(*) GROUP BY status must be rejected — Lucene has no per-group count", - LuceneFragmentConvertor.isCountFastPath(aggregate(scan, ImmutableBitSet.of(0), countStar(scan))) + LuceneFragmentPlanner.isCountFastPath(aggregate(scan, ImmutableBitSet.of(0), countStar(scan))) ); } @@ -130,8 +129,8 @@ public void testNonAggregateTop_notDrivable() { List.of(new RexBuilder(typeFactory).makeInputRef(scan, 0)), List.of("status") ); - assertFalse("Project (no aggregate above) must be rejected", LuceneFragmentConvertor.isCountFastPath(project)); - assertFalse("Bare TableScan must be rejected", LuceneFragmentConvertor.isCountFastPath(scan)); + assertFalse("Project (no aggregate above) must be rejected", LuceneFragmentPlanner.isCountFastPath(project)); + assertFalse("Bare TableScan must be rejected", LuceneFragmentPlanner.isCountFastPath(scan)); } // ---- Helpers ---- diff --git a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneQueryConversionUtilsTests.java b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneQueryConversionUtilsTests.java index a00ab6664970e..8f815b79ac586 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneQueryConversionUtilsTests.java +++ b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/LuceneQueryConversionUtilsTests.java @@ -50,6 +50,11 @@ public void testBareFieldExistsIsRewritten() { assertExistsRewrite(LuceneQueryConversionUtils.rewriteFieldExistsForSecondary(exists("severityText")), "severityText"); } + public void testFieldExistsIsPreservedWhenReaderHasDocValues() { + FieldExistsQuery query = exists("metric"); + assertSame(query, LuceneQueryConversionUtils.rewriteFieldExistsForSecondary(query, "metric"::equals)); + } + public void testFieldExistsUnderConstantScore() { Query in = new ConstantScoreQuery(exists("f")); Query out = LuceneQueryConversionUtils.rewriteFieldExistsForSecondary(in); diff --git a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/PlanAlternativeSelectorTests.java b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/PlanAlternativeSelectorTests.java index abe368d5504c4..adf2b46c4a57b 100644 --- a/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/PlanAlternativeSelectorTests.java +++ b/sandbox/plugins/analytics-backend-lucene/src/test/java/org/opensearch/be/lucene/PlanAlternativeSelectorTests.java @@ -37,6 +37,7 @@ import org.opensearch.analytics.planner.PlannerImpl; import org.opensearch.analytics.planner.dag.BackendPlanAdapter; import org.opensearch.analytics.planner.dag.DAGBuilder; +import org.opensearch.analytics.planner.dag.FragmentConversionDriver; import org.opensearch.analytics.planner.dag.PlanAlternativeSelector; import org.opensearch.analytics.planner.dag.PlanForker; import org.opensearch.analytics.planner.dag.QueryDAG; @@ -91,13 +92,12 @@ /** * Tests for {@link PlanAlternativeSelector} executed against the real {@link LuceneAnalyticsBackendPlugin}. * - *

Lives in the lucene module (rather than analytics-engine) so the production capability surface - * — {@code Index} scan + standard filter + COUNT aggregate, declared only for keyword/text - * types — is consulted directly. If someone widens or narrows {@code STANDARD_TYPES} in the - * production plugin, these tests catch the change without any mock to update. + *

Lives in the Lucene module so the production count and doc-values capability surface is + * consulted directly. If the supported field or operator sets change, these tests catch the + * change without a capability mock to update. * - *

Pipeline executed: {@code PlanForker} → {@code PlanAlternativeSelector} (no convertor — - * selection happens before conversion, mirroring {@code DefaultPlanExecutor.executeInternal}). + *

Pipeline executed: {@code PlanForker} → {@code PlanAlternativeSelector}; the source-plan + * test also runs {@code FragmentConversionDriver} to verify the normal serialized envelope. */ public class PlanAlternativeSelectorTests extends OpenSearchTestCase { @@ -185,6 +185,109 @@ public void testSumOverIntegerField_luceneNeverDrivesSum() { assertEquals("mock-parquet", alternatives.getFirst().backendId()); } + public void testSumOverLongWithoutDocValuesHasNoViableBackend() { + TableScan scan = scanOver("metric", SqlTypeName.BIGINT); + AggregateCall sum = AggregateCall.create( + SqlStdOperatorTable.SUM, + false, + List.of(0), + -1, + scan, + typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true), + "sum_metric" + ); + Map> mappings = Map.of("metric", Map.of("type", "long", "doc_values", false)); + + IllegalStateException failure = expectThrows( + IllegalStateException.class, + () -> forkAndSelect(aggregate(scan, sum), mappings, true, "lucene") + ); + assertTrue(failure.getMessage(), failure.getMessage().contains("No backend can scan all requested fields")); + } + + public void testSumOverLongWithParquetPrimaryDoesNotSelectLuceneDocValues() { + TableScan scan = scanOver("metric", SqlTypeName.BIGINT); + AggregateCall sum = AggregateCall.create( + SqlStdOperatorTable.SUM, + false, + List.of(0), + -1, + scan, + typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true), + "sum_metric" + ); + + QueryDAG dag = forkAndSelect(aggregate(scan, sum), longMappings(), true, "parquet"); + + List alternatives = leafOf(dag).getPlanAlternatives(); + assertEquals(1, alternatives.size()); + assertEquals("mock-parquet", alternatives.getFirst().backendId()); + } + + public void testSumOverLongSelectsLuceneArrowSource() { + TableScan scan = scanOver("metric", SqlTypeName.BIGINT); + AggregateCall sum = AggregateCall.create( + SqlStdOperatorTable.SUM, + false, + List.of(0), + -1, + scan, + typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true), + "sum_metric" + ); + + QueryDAG dag = forkAndSelect(aggregate(scan, sum), longMappings(), true, "lucene", true); + + List alternatives = leafOf(dag).getPlanAlternatives(); + assertEquals(1, alternatives.size()); + assertTrue( + org.apache.calcite.plan.RelOptUtil.toString(alternatives.getFirst().resolvedFragment()), + LuceneFragmentPlanner.classify(alternatives.getFirst().resolvedFragment()) instanceof LuceneFragmentPlanner.ArrowSourceShape + ); + assertEquals("lucene", alternatives.getFirst().backendId()); + assertNotNull(LuceneFragmentWirePlan.fromBytes(alternatives.getFirst().convertedBytes()).arrowSourcePlan()); + } + + public void testSumOverLongIgnoresUnsupportedUnreferencedField() { + TableScan scan = scanOver(List.of("metric", "unsupported"), List.of(SqlTypeName.BIGINT, SqlTypeName.INTEGER)); + AggregateCall sum = AggregateCall.create( + SqlStdOperatorTable.SUM, + false, + List.of(0), + -1, + scan, + typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true), + "sum_metric" + ); + Map> mappings = Map.of("metric", Map.of("type", "long"), "unsupported", Map.of("type", "integer")); + + QueryDAG dag = forkAndSelect(aggregate(scan, sum), mappings, true, "lucene"); + + assertEquals("lucene", leafOf(dag).getPlanAlternatives().getFirst().backendId()); + } + + public void testFilteredSumOverLongSelectsLuceneArrowSource() { + TableScan scan = scanOver("metric", SqlTypeName.BIGINT); + RexNode condition = rexBuilder.makeCall( + SqlStdOperatorTable.GREATER_THAN, + rexBuilder.makeInputRef(scan, 0), + rexBuilder.makeBigintLiteral(java.math.BigDecimal.TEN) + ); + AggregateCall sum = AggregateCall.create( + SqlStdOperatorTable.SUM, + false, + List.of(0), + -1, + scan, + typeFactory.createTypeWithNullability(typeFactory.createSqlType(SqlTypeName.BIGINT), true), + "sum_metric" + ); + + QueryDAG dag = forkAndSelect(aggregate(LogicalFilter.create(scan, condition), sum), longMappings(), true, "lucene"); + + assertEquals("lucene", leafOf(dag).getPlanAlternatives().getFirst().backendId()); + } + /** * Disqualified shape: {@code COUNT(*)} over a TEXT field with {@code index: false}. Even * though the field's type is in Lucene's {@code Index.supportedFieldTypes}, @@ -221,6 +324,38 @@ public void testCountStarWithIndexedKeywordFilter_selectsLucene() { assertEquals("lucene", alternatives.getFirst().backendId()); } + public void testCountStarWithParquetNumericFilterDoesNotSelectLucene() { + TableScan scan = scanOver(List.of("status", "amount"), List.of(SqlTypeName.VARCHAR, SqlTypeName.BIGINT)); + RexNode greaterThan = rexBuilder.makeCall( + SqlStdOperatorTable.GREATER_THAN, + rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.BIGINT), 1), + rexBuilder.makeExactLiteral(java.math.BigDecimal.valueOf(50)) + ); + RelNode plan = aggregate(LogicalFilter.create(scan, greaterThan), countStar(scan)); + Map> mappings = Map.of("status", Map.of("type", "keyword"), "amount", Map.of("type", "long")); + + QueryDAG dag = forkAndSelect(plan, mappings, true); + + List alternatives = leafOf(dag).getPlanAlternatives(); + assertEquals(1, alternatives.size()); + assertEquals("mock-parquet", alternatives.getFirst().backendId()); + } + + public void testHavingOnDerivedCountDoesNotDelegateToLucene() { + TableScan scan = scanOver("status", SqlTypeName.VARCHAR); + RelNode groupedCount = LogicalAggregate.create(scan, ImmutableBitSet.of(0), null, List.of(countStar(scan))); + RexNode greaterThan = rexBuilder.makeCall( + SqlStdOperatorTable.GREATER_THAN, + rexBuilder.makeInputRef(typeFactory.createSqlType(SqlTypeName.BIGINT), 1), + rexBuilder.makeBigintLiteral(java.math.BigDecimal.TEN) + ); + + QueryDAG dag = forkAndSelect(LogicalFilter.create(groupedCount, greaterThan), keywordMappings(), true, "parquet", true); + + assertEquals("mock-parquet", leafOf(dag).getPlanAlternatives().getFirst().backendId()); + assertNoDelegatedExpressions(dag.rootStage()); + } + /** * Depth-3: {@code COUNT(*) WHERE tag='a' OR region='eu' OR message MATCH 'x'}. All three * leaves are Lucene-delegatable (two keyword EQUALS on distinct fields + one MATCH on @@ -284,15 +419,40 @@ public void testCountStarWithFourArmOrAllDelegatable_selectsLucene() { // ---- Plan-execution helpers ---- private QueryDAG forkAndSelect(RelNode plan, Map> fieldMappings, boolean preferMetadataDriver) { + return forkAndSelect(plan, fieldMappings, preferMetadataDriver, "parquet"); + } + + private QueryDAG forkAndSelect( + RelNode plan, + Map> fieldMappings, + boolean preferMetadataDriver, + String primaryFormat + ) { + return forkAndSelect(plan, fieldMappings, preferMetadataDriver, primaryFormat, false); + } + + private QueryDAG forkAndSelect( + RelNode plan, + Map> fieldMappings, + boolean preferMetadataDriver, + String primaryFormat, + boolean convert + ) { AnalyticsSearchBackendPlugin dfBackend = new StubDfBackend(); AnalyticsSearchBackendPlugin luceneBackend = new LuceneAnalyticsBackendPlugin(null); + List backends = List.of(dfBackend, luceneBackend); + Map registry = Map.of(dfBackend.name(), dfBackend, luceneBackend.name(), luceneBackend); + backends.forEach(backend -> backend.bindBackends(registry)); - PlannerContext context = buildContext(fieldMappings, List.of(dfBackend, luceneBackend), preferMetadataDriver); + PlannerContext context = buildContext(fieldMappings, backends, preferMetadataDriver, primaryFormat); RelNode marked = PlannerImpl.runAllOptimizations(plan, context); QueryDAG dag = DAGBuilder.build(marked, context.getCapabilityRegistry(), mockClusterService(), TEST_RESOLVER); PlanForker.forkAll(dag, context.getCapabilityRegistry()); BackendPlanAdapter.adaptAll(dag, context.getCapabilityRegistry()); PlanAlternativeSelector.selectAll(dag, context.getCapabilityRegistry(), preferMetadataDriver); + if (convert) { + FragmentConversionDriver.convertAll(dag, context.getCapabilityRegistry()); + } return dag; } @@ -304,11 +464,24 @@ private static Stage leafOf(QueryDAG dag) { return stage; } + private static void assertNoDelegatedExpressions(Stage stage) { + for (StagePlan plan : stage.getPlanAlternatives()) { + assertTrue(plan.delegatedExpressions().isEmpty()); + } + stage.getChildStages().forEach(PlanAlternativeSelectorTests::assertNoDelegatedExpressions); + } + // ---- Calcite helpers ---- private TableScan scanOver(String fieldName, SqlTypeName type) { + return scanOver(List.of(fieldName), List.of(type)); + } + + private TableScan scanOver(List fieldNames, List types) { RelDataTypeFactory.Builder builder = typeFactory.builder(); - builder.add(fieldName, typeFactory.createSqlType(type)); + for (int i = 0; i < fieldNames.size(); i++) { + builder.add(fieldNames.get(i), typeFactory.createSqlType(types.get(i))); + } RelDataType rowType = builder.build(); RelOptTable table = mock(RelOptTable.class); when(table.getQualifiedName()).thenReturn(List.of("test_index")); @@ -384,6 +557,10 @@ private static Map> integerMappings() { return Map.of("status", Map.of("type", "integer")); } + private static Map> longMappings() { + return Map.of("metric", Map.of("type", "long")); + } + private static Map> nonIndexedTextMappings() { return Map.of("status", Map.of("type", "text", "index", false)); } @@ -407,7 +584,8 @@ private static Map> threeFieldMappings() { private PlannerContext buildContext( Map> fieldMappings, List backends, - boolean preferMetadataDriver + boolean preferMetadataDriver, + String primaryFormat ) { MappingMetadata mappingMetadata = mock(MappingMetadata.class); when(mappingMetadata.sourceAsMap()).thenReturn(Map.of("properties", fieldMappings)); @@ -416,7 +594,7 @@ private PlannerContext buildContext( when(indexMetadata.getIndex()).thenReturn(new Index("test_index", "uuid")); when(indexMetadata.getSettings()).thenReturn( Settings.builder() - .put("index.composite.primary_data_format", "parquet") + .put("index.composite.primary_data_format", primaryFormat) .putList("index.composite.secondary_data_formats", "lucene") .build() ); @@ -463,6 +641,16 @@ public String name() { return "mock-parquet"; } + @Override + public boolean supportsArrowBatchSourceExecution() { + return true; + } + + @Override + public byte[] compileArrowBatchSourcePlan(RelNode fragment, boolean partialAggregate) { + return new byte[] { 1, 2, 3 }; + } + @Override public BackendCapabilityProvider getCapabilityProvider() { return new BackendCapabilityProvider() { diff --git a/sandbox/plugins/analytics-engine/licenses/jackson-dataformat-yaml-2.22.1.jar.sha1 b/sandbox/plugins/analytics-engine/licenses/jackson-dataformat-yaml-2.22.1.jar.sha1 deleted file mode 100644 index 597173dcd7fc3..0000000000000 --- a/sandbox/plugins/analytics-engine/licenses/jackson-dataformat-yaml-2.22.1.jar.sha1 +++ /dev/null @@ -1 +0,0 @@ -af5fde2414e4a8d2617e7890f7068c26b5b66a33 \ No newline at end of file diff --git a/sandbox/plugins/analytics-engine/licenses/jackson-dataformat-yaml-2.22.2.jar.sha1 b/sandbox/plugins/analytics-engine/licenses/jackson-dataformat-yaml-2.22.2.jar.sha1 new file mode 100644 index 0000000000000..78dcb210d0996 --- /dev/null +++ b/sandbox/plugins/analytics-engine/licenses/jackson-dataformat-yaml-2.22.2.jar.sha1 @@ -0,0 +1 @@ +a4f075bf4cc1ee814ab98d69c1612786c2f42bc3 \ No newline at end of file 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 39f6d54c871e2..ba9e0ec0ce24b 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 @@ -120,17 +120,18 @@ public class AnalyticsPlugin extends Plugin implements ExtensiblePlugin, ActionP ); /** - * Controls the metadata-only driver vs. value-producing peer choice when both are viable - * for a stage: + * Controls the Lucene driver vs. value-producing peer choice when both are viable for a + * stage: * *

*/ public static final Setting PREFER_METADATA_DRIVER = Setting.boolSetting( @@ -182,12 +183,16 @@ public Collection createComponents( ArrowNativeAllocator nativeAllocator = pluginComponentRegistry.getComponent(ArrowNativeAllocator.class) .orElseThrow(() -> new IllegalStateException("ArrowNativeAllocator not available; arrow-base plugin must be installed")); - CapabilityRegistry capabilityRegistry = new CapabilityRegistry(backEnds, FieldStorageResolver::new); - Map backEndsByName = new LinkedHashMap<>(); for (AnalyticsSearchBackendPlugin be : backEnds) { backEndsByName.put(be.name(), be); } + Map backendRegistry = Map.copyOf(backEndsByName); + for (AnalyticsSearchBackendPlugin backend : backEnds) { + backend.bindBackends(backendRegistry); + } + CapabilityRegistry capabilityRegistry = new CapabilityRegistry(backEnds, FieldStorageResolver::new); + readerContextStore = new ReaderContextStore(threadPool); clusterService.getClusterSettings() .addSettingsUpdateConsumer(ReaderContextStore.READER_CONTEXT_KEEP_ALIVE, readerContextStore::setKeepAlive); 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 1046cbaa873f1..61471abaf83a8 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 @@ -51,7 +51,6 @@ 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; import org.opensearch.tasks.Task; @@ -763,7 +762,8 @@ default void onCompleteWithMetrics(byte[] metrics) { private FragmentResources startFragment(FragmentExecutionRequest request, ResolvedFragment resolved, IndexShard shard, Task task) throws IOException { - GatedCloseable gatedReader = resolved.readerProvider.acquireReader(); + AnalyticsSearchBackendPlugin backend = backends.get(resolved.plan.getBackendId()); + GatedCloseable gatedReader = backend.acquireReader(shard); // A query that requested top-N docs (row-ids) will be followed by a fetch phase that reuses // this reader. When it does, close() keeps the reader in the store for the fetch; otherwise // close() frees it immediately instead of waiting for the reaper. @@ -783,7 +783,7 @@ private FragmentResources startFragment(FragmentExecutionRequest request, Resolv .stream() .anyMatch(n -> n.type() == org.opensearch.analytics.spi.InstructionType.SETUP_PARTIAL_AGGREGATE) ); - AnalyticsSearchBackendPlugin backend = backends.get(resolved.plan.getBackendId()); + ctx.setDelegationThreadTracker(createThreadTracker(task)); backendContext = applyInstructionHandlers(backend, resolved.plan.getInstructions(), ctx); @@ -804,30 +804,10 @@ private FragmentResources startFragment(FragmentExecutionRequest request, Resolv AnalyticsSearchBackendPlugin acceptingBackend = backends.get(acceptingBackendId); FilterDelegationHandle handle = acceptingBackend.getFilterDelegationHandle(delegation.delegatedExpressions(), ctx); - // Build a thread tracker when task resource tracking is available. - DelegationThreadTracker tracker = null; - if (taskResourceTrackingService != null) { - long taskId = task.getId(); - TaskResourceTrackingService service = taskResourceTrackingService; - tracker = new DelegationThreadTracker() { - @Override - public long trackStart() { - long threadId = Thread.currentThread().threadId(); - service.taskExecutionStartedOnThread(taskId, threadId); - return threadId; - } - - @Override - public void trackEnd(long threadId) { - service.taskExecutionFinishedOnThread(taskId, threadId); - } - }; - } - // Register handle and tracker together under the query's contextId so concurrent // queries have isolated FFM callback bindings. The returned cleanup removes the // binding after query execution completes. - trackerCleanup = backend.configureFilterDelegation(contextId, handle, tracker, backendContext); + trackerCleanup = backend.configureFilterDelegation(contextId, handle, ctx.getDelegationThreadTracker(), backendContext); } // Hash-shuffle producer routing: if the instruction chain produced a @@ -998,8 +978,7 @@ private static BackendExecutionContext applyInstructionHandlers( return backendContext; } - private record ResolvedFragment(IndexReaderProvider readerProvider, FragmentExecutionRequest.PlanAlternative plan, String queryId, - int stageId, String shardIdStr) { + private record ResolvedFragment(FragmentExecutionRequest.PlanAlternative plan, String queryId, int stageId, String shardIdStr) { } /** @@ -1017,11 +996,6 @@ private String selectedBackendId(FragmentExecutionRequest request) { } private ResolvedFragment resolveFragment(FragmentExecutionRequest request, IndexShard shard) { - IndexReaderProvider readerProvider = shard.getReaderProvider(); - if (readerProvider == null) { - throw new IllegalStateException("No ReaderProvider on " + shard.shardId()); - } - // Backend selection happens on the coordinator (PlanAlternativeSelector), so the // request typically carries a single alternative. We still iterate to handle the // case where a stage genuinely has multiple value-producing alternatives — pick the @@ -1044,7 +1018,28 @@ private ResolvedFragment resolveFragment(FragmentExecutionRequest request, Index String shardIdStr = shard.shardId().toString(); listener.onPreFragmentExecution(request.getQueryId(), request.getStageId(), shardIdStr); - return new ResolvedFragment(readerProvider, selectedPlan, request.getQueryId(), request.getStageId(), shardIdStr); + return new ResolvedFragment(selectedPlan, request.getQueryId(), request.getStageId(), shardIdStr); + } + + private DelegationThreadTracker createThreadTracker(Task task) { + if (taskResourceTrackingService == null || task == null) { + return null; + } + long taskId = task.getId(); + TaskResourceTrackingService service = taskResourceTrackingService; + return new DelegationThreadTracker() { + @Override + public long trackStart() { + long threadId = Thread.currentThread().threadId(); + service.taskExecutionStartedOnThread(taskId, threadId); + return threadId; + } + + @Override + public void trackEnd(long threadId) { + service.taskExecutionFinishedOnThread(taskId, threadId); + } + }; } private ShardScanExecutionContext buildContext( diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityRegistry.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityRegistry.java index 092fb968c3fe2..b20ce6bea9b87 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityRegistry.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityRegistry.java @@ -249,7 +249,7 @@ public boolean isOpaqueOperation(String name) { // ---- Field-level lookups (iterates all formats a field has) ---- - /** All backends that can filter on this field across all its storage formats. */ + /** All backends that can filter on this field through a declared scan capability. */ public List filterBackendsForField(ScalarFunction function, FieldStorageInfo field) { FieldType fieldType = field.getFieldType(); List result = new ArrayList<>(); @@ -259,6 +259,9 @@ public List filterBackendsForField(ScalarFunction function, FieldStorage for (String format : field.getIndexFormats()) { result.addAll(filterBackends(function, fieldType, format)); } + List docValueReaders = scanBackendsForField(field); + List indexReaders = indexScanBackendsForField(field); + result.removeIf(backend -> docValueReaders.contains(backend) == false && indexReaders.contains(backend) == false); return result; } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java index cad42db1d6817..c669cae567b7c 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/CapabilityResolutionUtils.java @@ -37,7 +37,16 @@ public static List filterByReduceCapability(CapabilityRegistry registry, } } if (result.isEmpty()) { - throw new IllegalStateException("No viable backend supports coordinator reduce among " + viableBackends); + // A reduce stage consumes exchanged Arrow batches and does not scan the child's + // storage format. Permit any registered sink-capable backend to execute it. + for (AnalyticsSearchBackendPlugin backend : registry.getBackends()) { + if (viableBackends.contains(backend.name()) == false && backend.getExchangeSinkProvider() != null) { + result.add(backend.name()); + } + } + } + if (result.isEmpty()) { + throw new IllegalStateException("No registered backend supports coordinator reduce for " + viableBackends); } return result; } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java index 693f6c97204f9..83e858e91cc7a 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/PlannerContext.java @@ -240,10 +240,9 @@ public OpenSearchDistributionTraitDef getDistributionTraitDef() { /** * Mirrors the {@code analytics.planner.prefer_metadata_driver} cluster setting at planning - * time. When {@code false}, {@code OpenSearchTableScanRule} skips the permissive - * metadata-only-driver gate, so the metadata backend (Lucene today) is never admitted as a - * scan alternative — value-producing peers handle every shape, no late-stage alternative - * pruning needed. + * time. When {@code false}, {@code OpenSearchTableScanRule} does not admit Lucene through + * its permissive index or doc-values gate. Value-producing peers handle every shape, with + * no late-stage Lucene alternative selection. */ public boolean preferMetadataDriver() { return preferMetadataDriver; 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 fa6b73dc7ee2c..30817dd6aefe4 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 @@ -148,7 +148,11 @@ 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); RelNodeUtils.logPlan(LOGGER, "After CBO", modifiedRelNode); - Optional lateMat = OpenSearchLateMaterializationRewriter.rewrite(modifiedRelNode); + // Lucene's Arrow source path does not implement the QTF fetch-by-row-id phase. + Optional lateMat = OpenSearchLateMaterializationRewriter.rewrite( + modifiedRelNode, + scan -> scan.getViableBackends().size() != 1 || scan.getViableBackends().contains("lucene") == false + ); if (lateMat.isPresent()) { modifiedRelNode = lateMat.get(); RelNodeUtils.logPlan(LOGGER, "After late-materialization", modifiedRelNode); diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/BackendPlanAdapter.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/BackendPlanAdapter.java index edd0789909f6a..5df0ce34d07a2 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/BackendPlanAdapter.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/BackendPlanAdapter.java @@ -32,6 +32,7 @@ import org.opensearch.analytics.planner.rel.OpenSearchProject; import org.opensearch.analytics.planner.rel.OpenSearchRelNode; import org.opensearch.analytics.planner.rel.OperatorAnnotation; +import org.opensearch.analytics.spi.BackendCapabilityProvider; import org.opensearch.analytics.spi.FieldStorageInfo; import org.opensearch.analytics.spi.ScalarFunction; import org.opensearch.analytics.spi.ScalarFunctionAdapter; @@ -66,6 +67,12 @@ public static void adaptAll(QueryDAG dag, CapabilityRegistry registry) { adaptStage(dag.rootStage(), registry); } + /** Applies one backend's adapters to a fragment compiled through a sibling execution engine. */ + public static RelNode adaptFragment(RelNode fragment, BackendCapabilityProvider capabilityProvider) { + Adapters adapters = new Adapters(capabilityProvider.scalarFunctionAdapters(), capabilityProvider.windowFunctionAdapters()); + return adaptNode(fragment, adapters); + } + private static void adaptStage(Stage stage, CapabilityRegistry registry) { for (Stage child : stage.getChildStages()) { adaptStage(child, registry); @@ -115,7 +122,19 @@ private static RelNode adaptNode(RelNode node, Adapters adapters) { return DistributedAggregateRewriter.rewrite(withAdaptedChildren); } - return childrenChanged ? node.copy(node.getTraitSet(), adaptedChildren) : node; + RelNode current = childrenChanged ? node.copy(node.getTraitSet(), adaptedChildren) : node; + // Sibling-engine compilation receives annotation-stripped LogicalProject/Filter/Sort + // nodes. Apply the same adapters to their Rex trees as the normal OpenSearch path. + if (current instanceof OpenSearchRelNode == false) { + List fieldStorage = nearestInputStorage(current); + return current.accept(new RexShuttle() { + @Override + public RexNode visitCall(RexCall call) { + return adaptRex(call, adapters, fieldStorage, current.getCluster()); + } + }); + } + return current; } /** @@ -316,6 +335,23 @@ private static ScalarFunction resolveFunction(RexCall call) { return ScalarFunction.fromSqlOperatorWithFallback(call.getOperator()); } + private static List nearestInputStorage(RelNode node) { + if (node.getInputs().isEmpty()) { + return List.of(); + } + RelNode current = RelNodeUtils.unwrapHep(node.getInputs().getFirst()); + while (current != null) { + if (current instanceof OpenSearchRelNode openSearchNode) { + return openSearchNode.getOutputFieldStorage(); + } + if (current.getInputs().isEmpty()) { + break; + } + current = RelNodeUtils.unwrapHep(current.getInputs().getFirst()); + } + return List.of(); + } + /** * Rebind a Project's expressions against a new input whose row type has shifted (typically * in nullability — e.g. FINAL aggregate's rewriter turned a NOT-NULL count into a nullable 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 1ae407ba33813..a0726a2872d31 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 @@ -93,7 +93,12 @@ public static QueryDAG build( registry, ((OpenSearchRelNode) cboOutput).getViableBackends() ); - sinkProvider = registry.getBackend(reduceViable.getFirst()).getExchangeSinkProvider(); + String reduceBackend = reduceViable.getFirst(); + sinkProvider = registry.getBackend(reduceBackend).getExchangeSinkProvider(); + // The sink backend executes this entire coordinator fragment. Bind wrappers above the + // exchange as well as the reduce operator itself; their original viability was derived + // from shard storage and can otherwise conflict with a cross-backend reduce. + rootFragment = PlanForker.bindToBackend(rootFragment, reduceBackend); } // Root needs a shard target only if its fragment actually contains a TableScan. diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/PlanAlternativeSelector.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/PlanAlternativeSelector.java index b9035f5f10569..583711ff1a28d 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/PlanAlternativeSelector.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/PlanAlternativeSelector.java @@ -28,8 +28,9 @@ * preference (or scores empty) are kept as-is — value-producing backends that don't * implement {@link BackendShardPreference} simply pass through. * - *

Today's only consumer is Lucene's count-fast-path. The {@link ShardPreferenceContext} - * surface is intentionally minimal (just the user-facing {@code prefer_metadata_driver} flag); + *

Lucene uses this for its count fast path and supported doc-values Arrow source plans. + * The {@link ShardPreferenceContext} surface is intentionally minimal (just the user-facing + * {@code prefer_metadata_driver} flag); * future inputs (deletes, segment count, query-cache warmth) plug into the same scoring path. * *

TODO: this selection runs on the coordinator using only fragment-shape signals. True @@ -45,7 +46,7 @@ public final class PlanAlternativeSelector { private PlanAlternativeSelector() {} /** - * Collapses each stage's alternatives by score. Stages with ≤1 alternative are untouched. + * Collapses each stage's alternatives by score and rejects a sole negative-scoring plan. * * @param dag plan-forked DAG; modified in place. * @param registry capability registry for backend lookups. @@ -71,8 +72,14 @@ private static void selectStage(Stage stage, CapabilityRegistry registry, ShardP constrainToParentBackends(stage, child, registry, ctx); selectStage(child, registry, ctx); } - if (ctx == null) return; - if (stage.getPlanAlternatives().size() < 2) return; + if (ctx == null || stage.getPlanAlternatives().isEmpty()) return; + if (stage.getPlanAlternatives().size() == 1) { + StagePlan only = stage.getPlanAlternatives().getFirst(); + if (scoreOf(only, registry, ctx) < 0) { + throw new IllegalStateException("Only plan alternative was rejected by backend preference: " + only.backendId()); + } + return; + } // Pick the highest-scoring alternative. Backends without a preference score 0; // a positive score wins. Ties go to the first plan in PlanForker order. @@ -86,6 +93,9 @@ private static void selectStage(Stage stage, CapabilityRegistry registry, ShardP winnerScore = s; } } + if (winnerScore < 0) { + throw new IllegalStateException("All plan alternatives were rejected by backend preference"); + } stage.setPlanAlternatives(List.of(winner)); } diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/PlanForker.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/PlanForker.java index ddee34f007842..daaaad002f506 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/PlanForker.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/dag/PlanForker.java @@ -40,6 +40,20 @@ public static void forkAll(QueryDAG dag, CapabilityRegistry registry) { forkStage(dag.rootStage(), registry); } + /** + * Binds a coordinator-local fragment to the backend selected for its exchange sink. + * Storage capability does not constrain these operators because their leaves consume + * exchanged Arrow batches rather than index segments. + */ + static RelNode bindToBackend(RelNode node, String backend) { + List children = node.getInputs().stream().map(child -> bindToBackend(child, backend)).toList(); + if (node instanceof OpenSearchRelNode openSearchNode) { + List annotations = openSearchNode.getAnnotations().stream().map(a -> a.narrowTo(backend)).toList(); + return openSearchNode.copyResolved(backend, children, annotations); + } + return children.isEmpty() ? node : node.copy(node.getTraitSet(), children); + } + private static void forkStage(Stage stage, CapabilityRegistry registry) { for (Stage child : stage.getChildStages()) { forkStage(child, registry); 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 da0eb5972bce0..b95759b29a0a9 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 @@ -47,6 +47,8 @@ public SqlSyntax getSyntax() { private final RexNode original; private final List viableBackends; private final int annotationId; + /** Derived fields cannot be independently evaluated by a peer storage backend. */ + private final boolean performanceDelegationAllowed; /** * Peer backends that could have evaluated this predicate but lost the narrow. * Empty when the predicate is single-viable (no peer to consult) or hasn't @@ -67,7 +69,17 @@ public SqlSyntax getSyntax() { private final List performanceDelegationBackends; public AnnotatedPredicate(RelDataType type, RexNode original, List viableBackends, int annotationId) { - this(type, original, viableBackends, annotationId, List.of()); + this(type, original, viableBackends, annotationId, true, List.of()); + } + + public AnnotatedPredicate( + RelDataType type, + RexNode original, + List viableBackends, + int annotationId, + boolean performanceDelegationAllowed + ) { + this(type, original, viableBackends, annotationId, performanceDelegationAllowed, List.of()); } private AnnotatedPredicate( @@ -75,12 +87,14 @@ private AnnotatedPredicate( RexNode original, List viableBackends, int annotationId, + boolean performanceDelegationAllowed, List performanceDelegationBackends ) { super(type, ANNOTATED_PREDICATE_OP, List.of(original)); this.original = original; this.viableBackends = viableBackends; this.annotationId = annotationId; + this.performanceDelegationAllowed = performanceDelegationAllowed; this.performanceDelegationBackends = performanceDelegationBackends; } @@ -108,10 +122,10 @@ public List getPerformanceDelegationBackends() { @Override public OperatorAnnotation narrowTo(String backend) { - List peers = (viableBackends.size() > 1 && viableBackends.contains(backend)) + List peers = (performanceDelegationAllowed && viableBackends.size() > 1 && viableBackends.contains(backend)) ? viableBackends.stream().filter(b -> !b.equals(backend)).toList() : List.of(); - return new AnnotatedPredicate(type, original, List.of(backend), annotationId, peers); + return new AnnotatedPredicate(type, original, List.of(backend), annotationId, performanceDelegationAllowed, peers); } @Override @@ -121,7 +135,14 @@ public RexNode unwrap() { @Override public RexNode withAdaptedOriginal(RexNode adaptedOriginal) { - return new AnnotatedPredicate(type, adaptedOriginal, viableBackends, annotationId, performanceDelegationBackends); + return new AnnotatedPredicate( + type, + adaptedOriginal, + viableBackends, + annotationId, + performanceDelegationAllowed, + performanceDelegationBackends + ); } @Override @@ -146,7 +167,14 @@ public RexCall clone(RelDataType type, List operands) { "AnnotatedPredicate must wrap exactly one operand (the original predicate); got " + operands.size() ); } - return new AnnotatedPredicate(type, operands.get(0), viableBackends, annotationId, performanceDelegationBackends); + return new AnnotatedPredicate( + type, + operands.get(0), + viableBackends, + annotationId, + performanceDelegationAllowed, + performanceDelegationBackends + ); } @Override diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java index 6280f236a8433..593b736893774 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchAggregateSplitRule.java @@ -25,6 +25,7 @@ import org.apache.calcite.sql.SqlAggFunction; import org.apache.calcite.sql.type.SqlTypeFamily; import org.apache.calcite.util.ImmutableBitSet; +import org.opensearch.analytics.planner.CapabilityResolutionUtils; import org.opensearch.analytics.planner.PlannerContext; import org.opensearch.analytics.planner.RelNodeUtils; import org.opensearch.analytics.planner.dag.DistributedAggregateRewriter.FinalAggCallBuilder; @@ -215,6 +216,14 @@ public void onMatch(RelOptRuleCall call) { aggregate.getGroupSet().isEmpty() ); + // The PARTIAL reads storage and therefore retains the source backend. FINAL consumes + // exchanged Arrow batches, so it must run on a sink-capable reduce backend. This can be a + // sibling backend (Lucene PARTIAL -> DataFusion FINAL) because no storage scan crosses the + // exchange boundary. + List finalViableBackends = CapabilityResolutionUtils.filterByReduceCapability( + context.getCapabilityRegistry(), + aggregate.getViableBackends() + ); OpenSearchAggregate finalAggregate = new OpenSearchAggregate( aggregate.getCluster(), finalTraits, @@ -223,7 +232,7 @@ public void onMatch(RelOptRuleCall call) { aggregate.getGroupSets(), finalAggCalls, AggregateMode.FINAL, - aggregate.getViableBackends(), + finalViableBackends, aggregate.getCallAnnotations(), finalExtraLiterals, intermediateFields diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java index 8cef3978400e4..2b38e222b8c0b 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchFilterRule.java @@ -136,7 +136,22 @@ private RexNode annotateCondition(RexNode condition, List fiel // every declared FilterCapability has a matching serializer registered, and reject // the plugin otherwise — fail-fast at boot rather than at first dual-viable query. // Needs revisiting. - return new AnnotatedPredicate(rexCall.getType(), rexCall, viableBackends, context.nextAnnotationId()); + return new AnnotatedPredicate( + rexCall.getType(), + rexCall, + viableBackends, + context.nextAnnotationId(), + referencesDerivedField(rexCall, fieldStorageInfos) == false + ); + } + + private boolean referencesDerivedField(RexCall predicate, List fieldStorageInfos) { + PredicateContents contents = new PredicateContents(new HashSet<>(), new ArrayList<>()); + collect(predicate, contents); + return contents.fieldIndices() + .stream() + .map(i -> FieldStorageInfo.resolve(fieldStorageInfos, i)) + .anyMatch(FieldStorageInfo::isDerived); } /** 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 index 46b2cc8c6ae60..5c17fc4fafb1b 100644 --- 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 @@ -109,8 +109,19 @@ private OpenSearchLateMaterializationRewriter() {} /** Returns the rewritten root iff QTF matched and fired; {@link Optional#empty()} otherwise. */ public static Optional rewrite(RelNode root) { + return rewrite(root, scan -> true); + } + + /** + * Applies QTF only when the scan backend can emit row IDs and serve the fetch phase. + */ + public static Optional rewrite(RelNode root, java.util.function.Predicate scanSupportsFetch) { Detection detection = detect(root); if (detection == null) return Optional.empty(); + if (scanSupportsFetch.test(detection.belowChain().scan()) == false) { + LOGGER.debug("[QTF] scan backend cannot serve the fetch phase; skipping rewrite"); + return Optional.empty(); + } LOGGER.debug( "[QTF] fired: aboveAnchorPhysicalFields={}, belowAnchorPhysicalFields={}", detection.aboveAnchorPhysicalFields(), diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java index 3729a0a17adfb..3b007bc50ecaa 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchSortRule.java @@ -59,7 +59,14 @@ public void onMatch(RelOptRuleCall call) { List viableBackends = childViableBackends.stream().filter(sortCapable::contains).toList(); if (viableBackends.isEmpty()) { - throw new IllegalStateException("No backend supports SORT capability among " + childViableBackends); + // A coordinator sort above an exchange is scan-free. It can run on any + // sort-capable sink backend even when Lucene drove the shard scan. + viableBackends = sortCapable.stream() + .filter(backend -> context.getCapabilityRegistry().getBackend(backend).getExchangeSinkProvider() != null) + .toList(); + } + if (viableBackends.isEmpty()) { + throw new IllegalStateException("No registered backend supports SORT for " + childViableBackends); } // plus(): Calcite's Sort constructor asserts the trait set contains the collation. diff --git a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java index 27f3b2974d43b..17f7c7246d0d0 100644 --- a/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java +++ b/sandbox/plugins/analytics-engine/src/main/java/org/opensearch/analytics/planner/rules/OpenSearchTableScanRule.java @@ -83,38 +83,14 @@ public void onMatch(RelOptRuleCall call) { List delegationAcceptors = registry.delegationAcceptors(DelegationType.SCAN); List viableBackends = new ArrayList<>(registry.scanCapableBackends()); - // Two-phase field coverage check: - // 1. Value-producing backends (DocValues / StoredFields) must cover EVERY field — - // downstream ops can need any column's actual value, so a value-driver must be - // able to deliver all of them. Original strict invariant. - // 2. Metadata-only drivers (today: only Lucene via inverted index) stay viable if - // they cover SOME field. Downstream ops that need a column the metadata driver - // can't reach (e.g. Project on a numeric field) self-restrict and PlanForker's - // chain-agreement filter drops the driver from the surviving alternatives. The - // only chain that makes it through end-to-end is the count fast-path shape: - // count(*) / count(col) over filters touching only Lucene-indexable fields. - // - // Without the split, a single non-keyword field in the scan's row type (e.g. - // `amount`) would disqualify Lucene from every query against the index, even - // queries that never reference it. - // - // TODO: today {@code "lucene"} is the only metadata-only driver, identified by - // membership in the per-field {@code FieldStorageInfo.getIndexFormats()}. When a - // second metadata-only backend (e.g. Tantivy) lands — or worse, a backend that - // declares both Index AND DocValues — replace this hardcoded id with a - // first-class identifier on {@code BackendCapabilityProvider} (e.g. a "metadata - // driver" marker) so the planner can tell them apart from value-producing peers - // that happen to also have an inverted index. See - // CapabilityRegistry.metadataOnlyScanBackends history for the prior precomputed - // set; collapsed for now to keep the registry surface small. - final String metadataOnlyDriver = "lucene"; - // When the cluster setting analytics.planner.prefer_metadata_driver is off, skip the - // permissive metadata-only gate entirely — the metadata driver runs the strict - // value-producing check like any other backend, and (since Lucene declares no - // value-producing scan today) gets dropped at the scan level. No alternatives, no - // post-fork pruning needed downstream. + // Value-producing backends must cover every field in the scan row. The metadata + // driver gets one additional permissive route when the preference setting is on: + // it may remain viable if its index or doc-values reader covers any field. Downstream + // operators and the shard preference scorer reject shapes whose referenced values are + // unavailable. + final String metadataDriver = "lucene"; final boolean admitMetadataDriver = context.preferMetadataDriver(); - boolean metadataOnlyCoversAny = false; + boolean metadataDriverCoversAny = false; for (FieldStorageInfo field : fieldStorage) { if (field.isDerived()) { throw new IllegalStateException( @@ -128,9 +104,9 @@ public void onMatch(RelOptRuleCall call) { // a numeric field with the same indexFormats does not — even though its values are // physically in Lucene, no backend declares an Index scan over numerics today. List idxBackends = registry.indexScanBackendsForField(field); - boolean idxCoversMetadataDriver = idxBackends.contains(metadataOnlyDriver); - if (idxCoversMetadataDriver) { - metadataOnlyCoversAny = true; + boolean idxCoversMetadataDriver = idxBackends.contains(metadataDriver); + if (idxCoversMetadataDriver || dvBackends.contains(metadataDriver)) { + metadataDriverCoversAny = true; } LOGGER.debug( "[table-scan] field={} type={} indexFormats={} docValueFormats={} dvBackends={} idxBackends={} idxCoversMetadata={}", @@ -142,19 +118,18 @@ public void onMatch(RelOptRuleCall call) { idxBackends, idxCoversMetadataDriver ); - // Strict: every value-producing candidate must cover this field (or delegate to one - // that does). When admitMetadataDriver=false the metadata driver is held to the same - // strict rule; when true it's exempt here and the permissive check below decides. viableBackends.removeIf(candidate -> { - if (admitMetadataDriver && candidate.equals(metadataOnlyDriver)) return false; // metadata-only handled below if (dvBackends.contains(candidate)) return false; return !delegationSupporters.contains(candidate) || dvBackends.stream().noneMatch(delegationAcceptors::contains); }); } - // Permissive: keep the metadata-only driver viable iff it covers at least one field — - // only consulted when the setting allows it. - if (admitMetadataDriver == false || metadataOnlyCoversAny == false) { - viableBackends.remove(metadataOnlyDriver); + if (admitMetadataDriver) { + if (metadataDriverCoversAny && viableBackends.contains(metadataDriver) == false) { + viableBackends.add(metadataDriver); + } + } else { + // This setting controls both Lucene's metadata count path and its doc-values path. + viableBackends.remove(metadataDriver); } LOGGER.debug("[table-scan] viableBackends={}", viableBackends); diff --git a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java index e7601d060fe6b..aa86652ce57f3 100644 --- a/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java +++ b/sandbox/plugins/analytics-engine/src/test/java/org/opensearch/analytics/planner/MockLuceneBackend.java @@ -18,6 +18,7 @@ import org.opensearch.analytics.spi.FieldType; import org.opensearch.analytics.spi.FilterCapability; import org.opensearch.analytics.spi.ScalarFunction; +import org.opensearch.analytics.spi.ScanCapability; import org.opensearch.common.regex.Regex; import org.opensearch.index.engine.dataformat.ReaderManagerConfig; import org.opensearch.index.engine.exec.EngineReaderManager; @@ -99,6 +100,11 @@ public String name() { return NAME; } + @Override + protected Set scanCapabilities() { + return Set.of(new ScanCapability.Index(LUCENE_FORMATS, STANDARD_TYPES)); + } + @Override protected Set filterCapabilities() { return FILTER_CAPS; diff --git a/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LuceneArrowSourceIT.java b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LuceneArrowSourceIT.java new file mode 100644 index 0000000000000..fc238b990ae74 --- /dev/null +++ b/sandbox/qa/analytics-engine-rest/src/test/java/org/opensearch/analytics/qa/LuceneArrowSourceIT.java @@ -0,0 +1,346 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to this file be licensed under + * the Apache-2.0 license or a compatible open source license. + */ + +package org.opensearch.analytics.qa; + +import org.opensearch.client.Request; +import org.opensearch.client.Response; +import org.opensearch.client.ResponseException; + +import java.io.IOException; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +/** + * End-to-end coverage for the normal Lucene doc-values to DataFusion query path. + * + *

The main fixture is a standard two-shard Lucene index. Its mapping also contains an + * unsupported integer field, so successful long/keyword/date queries + * verify that planning narrows the Arrow source instead of requiring every scan column. Three + * forced flushes exercise multiple Lucene segments, and one document omits {@code metric} to + * exercise null propagation and {@code COUNT(field)} semantics. + */ +public class LuceneArrowSourceIT extends AnalyticsRestTestCase { + + private static final String INDEX = "lucene_arrow_source_e2e"; + private static final String FALLBACK_INDEX = "lucene_arrow_source_fallback"; + private static final String MULTI_VALUE_INDEX = "lucene_arrow_source_multi_value"; + private static final String DOC_VALUES_ONLY_INDEX = "lucene_arrow_source_doc_values_only"; + + private static final List DOCS = List.of( + new Doc(1L, 10L, "alpha", "2024-01-01T00:00:00Z", 101), + new Doc(2L, 20L, "beta", "2024-01-02T00:00:00Z", 102), + new Doc(3L, null, "alpha", "2024-01-03T00:00:00Z", 103), + new Doc(4L, 40L, "beta", "2024-01-04T00:00:00Z", 104), + new Doc(5L, 50L, "gamma", "2024-01-05T00:00:00Z", 105), + new Doc(6L, 60L, "alpha", "2024-01-06T00:00:00Z", 106) + ); + + private static boolean dataProvisioned; + + @Override + protected void onBeforeQuery() throws IOException { + synchronized (LuceneArrowSourceIT.class) { + if (dataProvisioned) { + return; + } + createLucenePrimaryIndex(); + ingestLuceneSegments(); + createFallbackIndex(); + createMultiValueIndex(); + createDocValuesOnlyIndex(); + dataProvisioned = true; + } + } + + public void testFilteredAggregateAcrossShardsAndSegments() throws Exception { + String ppl = "source=" + + INDEX + + " | where id >= 3 | stats sum(metric) as total, count(metric) as non_null, count() as rows"; + Map explain = executeExplain(ppl); + assertStageChoseBackend(explain, "SHARD_FRAGMENT", "lucene"); + assertStageChoseBackend(explain, "COORDINATOR_REDUCE", "datafusion"); + + Map response = executePpl(ppl); + assertEquals(150L, numberCell(response, 0, "total").longValue()); + assertEquals(3L, numberCell(response, 0, "non_null").longValue()); + assertEquals(4L, numberCell(response, 0, "rows").longValue()); + } + + public void testAverageMinAndMaxAggregates() throws Exception { + String ppl = "source=" + INDEX + " | stats avg(metric) as average, min(metric) as minimum, max(metric) as maximum"; + Map response = executePpl(ppl); + + assertEquals(36.0d, numberCell(response, 0, "average").doubleValue(), 0.0d); + assertEquals(10L, numberCell(response, 0, "minimum").longValue()); + assertEquals(60L, numberCell(response, 0, "maximum").longValue()); + Map explain = executeExplain(ppl); + assertStageChoseBackend(explain, "SHARD_FRAGMENT", "lucene"); + assertStageChoseBackend(explain, "COORDINATOR_REDUCE", "datafusion"); + } + + public void testGroupedKeywordAggregate() throws Exception { + String ppl = "source=" + INDEX + " | stats sum(metric) as total, count(metric) as non_null by category | sort category"; + Map response = executePpl(ppl); + Map actual = new HashMap<>(); + List columns = extractColumnNames(response); + for (List row : dataRows(response)) { + String category = row.get(columns.indexOf("category")).toString(); + long total = ((Number) row.get(columns.indexOf("total"))).longValue(); + long nonNull = ((Number) row.get(columns.indexOf("non_null"))).longValue(); + actual.put(category, new long[] { total, nonNull }); + } + + assertEquals(3, actual.size()); + assertArrayEquals(new long[] { 70L, 2L }, actual.get("alpha")); + assertArrayEquals(new long[] { 60L, 2L }, actual.get("beta")); + assertArrayEquals(new long[] { 50L, 1L }, actual.get("gamma")); + assertStageChoseBackend(executeExplain(ppl), "SHARD_FRAGMENT", "lucene"); + } + + public void testProjectionFilterSortAndTimestamp() throws Exception { + String ppl = "source=" + INDEX + " | where id >= 4 | fields id, category, event_time | sort id"; + Map response = executePpl(ppl); + List> rows = dataRows(response); + List columns = extractColumnNames(response); + + assertEquals(3, rows.size()); + assertProjectedRow(rows.get(0), columns, 4L, "beta", "2024-01-04"); + assertProjectedRow(rows.get(1), columns, 5L, "gamma", "2024-01-05"); + assertProjectedRow(rows.get(2), columns, 6L, "alpha", "2024-01-06"); + assertStageChoseBackend(executeExplain(ppl), "SHARD_FRAGMENT", "lucene"); + } + + public void testDocValuesOnlyIndex() throws Exception { + String ppl = "source=" + + DOC_VALUES_ONLY_INDEX + + " | stats sum(metric) as total, count(metric) as rows by category | sort category"; + + Map response = executePpl(ppl); + List> rows = dataRows(response); + List columns = extractColumnNames(response); + assertEquals(2, rows.size()); + assertEquals("alpha", rows.get(0).get(columns.indexOf("category"))); + assertEquals(40L, ((Number) rows.get(0).get(columns.indexOf("total"))).longValue()); + assertEquals(2L, ((Number) rows.get(0).get(columns.indexOf("rows"))).longValue()); + assertEquals("beta", rows.get(1).get(columns.indexOf("category"))); + assertEquals(20L, ((Number) rows.get(1).get(columns.indexOf("total"))).longValue()); + assertEquals(1L, ((Number) rows.get(1).get(columns.indexOf("rows"))).longValue()); + + Map explain = executeExplain(ppl); + assertStageChoseBackend(explain, "SHARD_FRAGMENT", "lucene"); + assertStageChoseBackend(explain, "COORDINATOR_REDUCE", "datafusion"); + } + + public void testCountFastPathAndNumericNullFilter() throws Exception { + String countAll = "source=" + INDEX + " | stats count() as rows"; + String countNull = "source=" + INDEX + " | where isnull(metric) | stats count() as rows"; + + assertEquals(6L, numberCell(executePpl(countAll), 0, "rows").longValue()); + assertEquals(1L, numberCell(executePpl(countNull), 0, "rows").longValue()); + assertStageChoseBackend(executeExplain(countAll), "SHARD_FRAGMENT", "lucene"); + assertStageChoseBackend(executeExplain(countNull), "SHARD_FRAGMENT", "lucene"); + } + + public void testMultiValuedScalarFieldFailsFast() { + String ppl = "source=" + MULTI_VALUE_INDEX + " | stats sum(metric) as total"; + + ResponseException failure = expectThrows(ResponseException.class, () -> executePpl(ppl)); + assertEquals(500, failure.getResponse().getStatusLine().getStatusCode()); + } + + public void testUnsupportedIntegerFallsBackToDataFusion() throws Exception { + String ppl = "source=" + FALLBACK_INDEX + " | stats sum(value) as total"; + + assertEquals(10L, numberCell(executePpl(ppl), 0, "total").longValue()); + assertStageChoseBackend(executeExplain(ppl), "SHARD_FRAGMENT", "datafusion"); + } + + private void createLucenePrimaryIndex() throws IOException { + deleteIfExists(INDEX); + String body = "{" + + "\"settings\": {" + + " \"number_of_shards\": 2," + + " \"number_of_replicas\": 0," + + " \"index.composite.primary_data_format\": \"lucene\"" + + "}," + + "\"mappings\": {\"properties\": {" + + " \"id\": {\"type\": \"long\"}," + + " \"metric\": {\"type\": \"long\"}," + + " \"category\": {\"type\": \"keyword\"}," + + " \"event_time\": {\"type\": \"date\"}," + + " \"unsupported\": {\"type\": \"integer\"}" + + "}}}"; + createIndex(INDEX, body); + } + + private void ingestLuceneSegments() throws IOException { + for (int from = 0; from < DOCS.size(); from += 2) { + StringBuilder bulk = new StringBuilder(); + for (int i = from; i < Math.min(from + 2, DOCS.size()); i++) { + bulk.append("{\"index\": {}}\n"); + bulk.append(DOCS.get(i).toJson()).append('\n'); + } + bulkIndex(INDEX, bulk.toString()); + client().performRequest(new Request("POST", "/" + INDEX + "/_flush?force=true")); + } + } + + private void createMultiValueIndex() throws IOException { + deleteIfExists(MULTI_VALUE_INDEX); + String body = "{" + + "\"settings\": {" + + " \"number_of_shards\": 1," + + " \"number_of_replicas\": 0," + + " \"index.composite.primary_data_format\": \"lucene\"" + + "}," + + "\"mappings\": {\"properties\": {\"metric\": {\"type\": \"long\"}}}" + + "}"; + createIndex(MULTI_VALUE_INDEX, body); + bulkIndex(MULTI_VALUE_INDEX, "{\"index\": {}}\n{\"metric\": [10, 20]}\n"); + } + + private void createDocValuesOnlyIndex() throws IOException { + deleteIfExists(DOC_VALUES_ONLY_INDEX); + String body = "{" + + "\"settings\": {" + + " \"number_of_shards\": 2," + + " \"number_of_replicas\": 0," + + " \"index.composite.primary_data_format\": \"lucene\"" + + "}," + + "\"mappings\": {" + + " \"_source\": {\"enabled\": false}," + + " \"properties\": {" + + " \"metric\": {\"type\": \"long\", \"index\": false, \"doc_values\": true}," + + " \"category\": {\"type\": \"keyword\", \"index\": false, \"doc_values\": true}" + + " }" + + "}}"; + createIndex(DOC_VALUES_ONLY_INDEX, body); + bulkIndex( + DOC_VALUES_ONLY_INDEX, + "{\"index\": {}}\n{\"metric\": 10, \"category\": \"alpha\"}\n" + + "{\"index\": {}}\n{\"metric\": 20, \"category\": \"beta\"}\n" + + "{\"index\": {}}\n{\"metric\": 30, \"category\": \"alpha\"}\n" + ); + } + + private void createFallbackIndex() throws IOException { + deleteIfExists(FALLBACK_INDEX); + String body = "{" + + "\"settings\": {" + + " \"number_of_shards\": 2," + + " \"number_of_replicas\": 0," + + " \"index.pluggable.dataformat.enabled\": true," + + " \"index.pluggable.dataformat\": \"composite\"," + + " \"index.composite.primary_data_format\": \"parquet\"," + + " \"index.composite.secondary_data_formats\": \"lucene\"" + + "}," + + "\"mappings\": {\"properties\": {\"value\": {\"type\": \"integer\"}}}" + + "}"; + createIndex(FALLBACK_INDEX, body); + bulkIndex( + FALLBACK_INDEX, + "{\"index\": {}}\n{\"value\": 2}\n" + + "{\"index\": {}}\n{\"value\": 3}\n" + + "{\"index\": {}}\n{\"value\": 5}\n" + ); + } + + private void createIndex(String index, String body) throws IOException { + Request create = new Request("PUT", "/" + index); + create.setJsonEntity(body); + Map response = assertOkAndParse(client().performRequest(create), "Create " + index); + assertEquals(true, response.get("acknowledged")); + + Request health = new Request("GET", "/_cluster/health/" + index); + health.addParameter("wait_for_status", "green"); + health.addParameter("timeout", "30s"); + client().performRequest(health); + } + + private void bulkIndex(String index, String ndjson) throws IOException { + Request bulk = new Request("POST", "/" + index + "/_bulk"); + bulk.setJsonEntity(ndjson); + bulk.addParameter("refresh", "true"); + bulk.setOptions(bulk.getOptions().toBuilder().addHeader("Content-Type", "application/x-ndjson").build()); + Map response = assertOkAndParse(client().performRequest(bulk), "Bulk index " + index); + assertEquals("Bulk indexing should have no errors: " + response, false, response.get("errors")); + } + + private void deleteIfExists(String index) throws IOException { + Request delete = new Request("DELETE", "/" + index); + delete.addParameter("ignore_unavailable", "true"); + client().performRequest(delete); + } + + private Map executeExplain(String ppl) throws IOException { + Request request = new Request("POST", "/_analytics/ppl/_explain"); + request.setJsonEntity("{\"query\": \"" + escapeJson(ppl) + "\"}"); + Response response = client().performRequest(request); + return assertOkAndParse(response, "EXPLAIN: " + ppl); + } + + @SuppressWarnings("unchecked") + private static void assertStageChoseBackend(Map explain, String executionType, String expectedBackend) { + Map profile = (Map) explain.get("profile"); + assertNotNull("profile present", profile); + List> stages = (List>) profile.get("stages"); + assertNotNull("stages present", stages); + for (Map stage : stages) { + if (executionType.equals(stage.get("execution_type"))) { + assertEquals(executionType + " stage: " + stage, expectedBackend, stage.get("chosen_backend")); + return; + } + } + fail("No " + executionType + " stage in profile: " + stages); + } + + private static Number numberCell(Map response, int rowIndex, String column) { + List columns = extractColumnNames(response); + int columnIndex = columns.indexOf(column); + assertTrue("Missing column " + column + " in " + columns, columnIndex >= 0); + return (Number) dataRows(response).get(rowIndex).get(columnIndex); + } + + @SuppressWarnings("unchecked") + private static List> dataRows(Map response) { + List> rows = (List>) response.get("datarows"); + assertNotNull("Response missing datarows: " + response, rows); + return rows; + } + + private static void assertProjectedRow( + List row, + List columns, + long expectedId, + String expectedCategory, + String expectedDate + ) { + assertEquals(expectedId, ((Number) row.get(columns.indexOf("id"))).longValue()); + assertEquals(expectedCategory, row.get(columns.indexOf("category"))); + Object timestamp = row.get(columns.indexOf("event_time")); + assertNotNull(timestamp); + assertTrue("Unexpected timestamp: " + timestamp, timestamp.toString().contains(expectedDate)); + } + + private record Doc(long id, Long metric, String category, String eventTime, int unsupported) { + String toJson() { + String metricJson = metric == null ? "" : ", \"metric\": " + metric; + return "{\"id\": " + + id + + metricJson + + ", \"category\": \"" + + category + + "\", \"event_time\": \"" + + eventTime + + "\", \"unsupported\": " + + unsupported + + "}"; + } + } +}