Add SPI for doc_values based collection

sandbox/libs/analytics-framework/licenses/jackson-core-2.22.2.jar.sha1

@@ -0,0 +1 @@
+13a748ea3e329fa220076e021b45c8391b32420c

sandbox/libs/analytics-framework/licenses/jackson-databind-2.22.2.jar.sha1

@@ -0,0 +1 @@
+921bd2092b0c539b2876de7063d55c72edcd05d3

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

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.
+ *
+ * <p>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.
+     *
+     * <p>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();
+}

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.
+ *
+ * <p>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();
+}

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<InputColumn> 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<InputColumn> 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());
+    }
+}

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<InputColumn> 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());
+            }
+        }
+    }
+}

sandbox/plugins/analytics-backend-datafusion/licenses/jackson-datatype-jdk8-2.22.2.jar.sha1

@@ -0,0 +1 @@
+0b900dd7125fa16cfdf46135d3ffb3243d0f8b88

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