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/spi/AnalyticsSearchBackendPlugin.java b/sandbox/libs/analytics-framework/src/main/java/org/opensearch/analytics/spi/AnalyticsSearchBackendPlugin.java index 071d1578561f1..871f5d6ffd296 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 @@ -14,6 +14,7 @@ import org.opensearch.cluster.ClusterState; import org.opensearch.index.engine.exec.IndexReaderProvider.Reader; import org.opensearch.index.shard.IndexShard; +import org.opensearch.tasks.Task; import java.util.Collections; import java.util.List; @@ -43,6 +44,25 @@ public interface AnalyticsSearchBackendPlugin { /** Unique backend name (e.g., "datafusion", "lucene"). */ String name(); + /** Whether this backend can execute plans over a caller-provided Arrow batch source. */ + default boolean supportsArrowBatchSourceExecution() { + return false; + } + + /** + * 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 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..46b109555ec06 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 @@ -25,6 +25,8 @@ 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 +63,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 +912,27 @@ public FragmentConvertor getFragmentConvertor() { return new DataFusionFragmentConvertor(plugin.getSubstraitExtensions()); } + @Override + public boolean supportsArrowBatchSourceExecution() { + return true; + } + + @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/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/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-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-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