feat(datasource): add DataSource interface and SessionContext.registe…

…rDataSource Java API

core/src/main/java/org/apache/datafusion/DataSource.java

@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.datafusion;
+
+import org.apache.arrow.vector.ipc.ArrowReader;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+/**
+ * A Java-implemented table that can be registered with a {@link SessionContext}.
+ *
+ * <p>Each call to {@link #scan()} must return a fresh, independent {@link ArrowReader} so that
+ * queries which touch the table more than once (self-joins, {@code UNION ALL}, repeated reads) work
+ * correctly. The returned reader is closed by the framework when the stream ends.
+ *
+ * <p>The schema returned by {@link #schema()} is captured once at registration time. Every batch
+ * produced by every {@code ArrowReader} returned from {@link #scan()} must conform to it; a
+ * mismatch fails the query.
+ */
+public interface DataSource {
+  /** The fixed schema of this table. Called once, at registration time. */
+  Schema schema();
+
+  /**
+   * Open a fresh batch stream for this table. Called once per query that scans the table.
+   *
+   * <p>Each invocation MUST return an independent {@link ArrowReader}. The reader's schema MUST
+   * equal {@link #schema()}.
+   */
+  ArrowReader scan();
+}

core/src/main/java/org/apache/datafusion/SessionContext.java

@@ -392,6 +392,38 @@ public void registerUdf(ScalarUdf udf) {
     registerScalarUdf(nativeHandle, name, signatureBytes, volatility.code(), impl);
   }
 
+  /**
+   * Register a Java-implemented data source as a table. SQL queries that reference {@code name}
+   * call back into {@code source} to fetch batches.
+   *
+   * <p>{@link DataSource#schema()} is called once here, on the calling thread, and cached on the
+   * native side. {@link DataSource#scan()} is called once per query that touches the table, on a
+   * Tokio worker thread; it must return a fresh, independent {@link
+   * org.apache.arrow.vector.ipc.ArrowReader} on every call.
+   *
+   * @throws IllegalArgumentException if {@code name} or {@code source} is {@code null}.
+   * @throws IllegalStateException if {@code source.schema()} returns {@code null}, or this context
+   *     is closed.
+   * @throws RuntimeException if native registration fails.
+   */
+  public void registerDataSource(String name, DataSource source) {
+    if (nativeHandle == 0) {
+      throw new IllegalStateException("SessionContext is closed");
+    }
+    if (name == null) {
+      throw new IllegalArgumentException("registerDataSource name must be non-null");
+    }
+    if (source == null) {
+      throw new IllegalArgumentException("registerDataSource source must be non-null");
+    }
+    Schema schema = source.schema();
+    if (schema == null) {
+      throw new IllegalStateException("DataSource.schema returned null");
+    }
+    byte[] schemaIpc = serializeSchemaIpc(schema);
+    registerDataSourceNative(nativeHandle, name, schemaIpc, source);
+  }
+
   private static byte[] serializeSchemaIpc(Schema schema) {
     ByteArrayOutputStream baos = new ByteArrayOutputStream();
     try (BufferAllocator allocator = new RootAllocator();
@@ -453,4 +485,7 @@ private static native long readJsonWithOptions(
 
   private static native void registerScalarUdf(
       long handle, String name, byte[] signatureSchemaBytes, byte volatility, ScalarFunction impl);
+
+  private static native void registerDataSourceNative(
+      long handle, String name, byte[] schemaIpcBytes, DataSource source);
 }

core/src/main/java/org/apache/datafusion/internal/JniBridge.java

@@ -23,12 +23,15 @@
 import java.util.List;
 
 import org.apache.arrow.c.ArrowArray;
+import org.apache.arrow.c.ArrowArrayStream;
 import org.apache.arrow.c.ArrowSchema;
 import org.apache.arrow.c.Data;
 import org.apache.arrow.memory.RootAllocator;
 import org.apache.arrow.vector.FieldVector;
 import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.ipc.ArrowReader;
 import org.apache.datafusion.ColumnarValue;
+import org.apache.datafusion.DataSource;
 import org.apache.datafusion.ScalarFunction;
 import org.apache.datafusion.ScalarFunctionArgs;
 
@@ -139,4 +142,30 @@ public static byte invokeScalarUdf(
       return resultKind;
     }
   }
+
+  /**
+   * Open a fresh batch stream from a Java {@link DataSource} and export it through the supplied
+   * Arrow C Data Interface address. Called from native code; not for application use.
+   *
+   * <p>On success, ownership of the returned reader transfers to the FFI stream's release callback,
+   * so the native side closing the stream also closes the reader. On any failure during export, the
+   * reader is closed here before the exception propagates.
+   */
+  public static void invokeDataSourceScan(DataSource source, long ffiStreamAddr) {
+    ArrowReader reader = source.scan();
+    if (reader == null) {
+      throw new IllegalStateException("DataSource.scan returned null");
+    }
+    ArrowArrayStream stream = ArrowArrayStream.wrap(ffiStreamAddr);
+    try {
+      Data.exportArrayStream(ALLOCATOR, reader, stream);
+    } catch (Throwable t) {
+      try {
+        reader.close();
+      } catch (Exception ignored) {
+        // best-effort cleanup; original error wins
+      }
+      throw t;
+    }
+  }
 }