Instances are not thread-safe and must be closed. Both {@link #collect} and - * {@link #executeStream} consume the DataFrame: a successfully consumed DataFrame cannot be - * consumed again by either method (or by other executors such as {@link #count}), and {@link - * #close()} on an already-consumed instance is a no-op. + *
Instances must be closed. Both {@link #collect} and {@link #executeStream} consume the + * DataFrame: a successfully consumed DataFrame cannot be consumed again by either method (or by + * other executors such as {@link #count}), and {@link #close()} on an already-consumed instance is + * a no-op. + * + *
Instances are safe to share between threads. This class pins its native handle for the + * duration of every call, so a {@link #close} or a consuming operation racing with work on another + * thread cannot free the plan out from under it; the call that loses such a race throws {@link + * IllegalStateException}. When several threads race to consume the same DataFrame, exactly one + * succeeds. {@link #close} and the consuming operations block until calls already in flight have + * returned. */ public final class DataFrame implements AutoCloseable { static { NativeLibraryLoader.loadLibrary(); } - private long nativeHandle; + private static final String CLOSED = "DataFrame is closed or already collected"; + + private final NativeHandle handle; DataFrame(long nativeHandle) { if (nativeHandle == 0) { throw new IllegalArgumentException("DataFrame native handle is null"); } - this.nativeHandle = nativeHandle; + this.handle = new NativeHandle(nativeHandle, CLOSED); } /** @@ -68,14 +78,9 @@ public final class DataFrame implements AutoCloseable { * {@link #executeStream(BufferAllocator)} for analytics-scale queries. */ public ArrowReader collect(BufferAllocator allocator) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator); - long handle = nativeHandle; - nativeHandle = 0; try { - collectDataFrame(handle, stream.memoryAddress()); + collectDataFrame(handle.claim(), stream.memoryAddress()); return Data.importArrayStream(allocator, stream); } catch (Throwable e) { stream.close(); @@ -98,14 +103,9 @@ public ArrowReader collect(BufferAllocator allocator) { * use this method. */ public ArrowReader executeStream(BufferAllocator allocator) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } ArrowArrayStream stream = ArrowArrayStream.allocateNew(allocator); - long handle = nativeHandle; - nativeHandle = 0; try { - executeStreamDataFrame(handle, stream.memoryAddress()); + executeStreamDataFrame(handle.claim(), stream.memoryAddress()); return Data.importArrayStream(allocator, stream); } catch (Throwable e) { stream.close(); @@ -121,10 +121,13 @@ public ArrowReader executeStream(BufferAllocator allocator) { * schema carries no buffer data. */ public Schema schema() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + byte[] ipcBytes; + long h = handle.acquire(); + try { + ipcBytes = schemaIpc(h); + } finally { + handle.release(); } - byte[] ipcBytes = schemaIpc(nativeHandle); try { return MessageSerializer.deserializeSchema( new ReadChannel(Channels.newChannel(new ByteArrayInputStream(ipcBytes)))); @@ -143,10 +146,12 @@ public Schema schema() { * #show()} or {@link #collect(BufferAllocator)}. */ public DataFrame explain(boolean verbose, boolean analyze) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(explainPlan(h, verbose, analyze)); + } finally { + handle.release(); } - return new DataFrame(explainPlan(nativeHandle, verbose, analyze)); } /** @@ -160,10 +165,12 @@ public DataFrame explain(boolean verbose, boolean analyze) { * @throws RuntimeException if execution fails. */ public DataFrame cache() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(cachePlan(h)); + } finally { + handle.release(); } - return new DataFrame(cachePlan(nativeHandle)); } /** @@ -178,34 +185,42 @@ public DataFrame cache() { * @throws RuntimeException if execution fails. */ public DataFrame describe() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(describePlan(h)); + } finally { + handle.release(); } - return new DataFrame(describePlan(nativeHandle)); } /** Execute the plan and return the number of rows. */ public long count() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return countRows(h); + } finally { + handle.release(); } - return countRows(nativeHandle); } /** Execute the plan and print formatted batches to native stdout. */ public void show() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + showDataFrame(h); + } finally { + handle.release(); } - showDataFrame(nativeHandle); } /** Execute the plan and print the first {@code limit} rows to native stdout. */ public void show(int limit) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + showDataFrameWithLimit(h, limit); + } finally { + handle.release(); } - showDataFrameWithLimit(nativeHandle, limit); } /** @@ -213,10 +228,12 @@ public void show(int limit) { * closed independently. */ public DataFrame select(String... columnNames) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(selectColumns(h, columnNames)); + } finally { + handle.release(); } - return new DataFrame(selectColumns(nativeHandle, columnNames)); } /** @@ -224,10 +241,12 @@ public DataFrame select(String... columnNames) { * DataFrame's own schema. The receiver remains usable and must still be closed independently. */ public DataFrame filter(String predicate) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(filterRows(h, predicate)); + } finally { + handle.release(); } - return new DataFrame(filterRows(nativeHandle, predicate)); } /** @@ -249,10 +268,12 @@ public DataFrame limit(int skip, int fetch) { if (fetch < 0) { throw new IllegalArgumentException("fetch must be non-negative, was " + fetch); } - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(limitRows(h, skip, fetch)); + } finally { + handle.release(); } - return new DataFrame(limitRows(nativeHandle, skip, fetch)); } /** @@ -260,10 +281,12 @@ public DataFrame limit(int skip, int fetch) { * independently. */ public DataFrame distinct() { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(distinctRows(h)); + } finally { + handle.release(); } - return new DataFrame(distinctRows(nativeHandle)); } /** @@ -271,18 +294,22 @@ public DataFrame distinct() { * and must still be closed independently. */ public DataFrame dropColumns(String... columnNames) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(dropColumns(h, columnNames)); + } finally { + handle.release(); } - return new DataFrame(dropColumns(nativeHandle, columnNames)); } /** Rename a column. The receiver remains usable and must still be closed independently. */ public DataFrame withColumnRenamed(String oldName, String newName) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + return new DataFrame(renameColumn(h, oldName, newName)); + } finally { + handle.release(); } - return new DataFrame(renameColumn(nativeHandle, oldName, newName)); } /** @@ -294,16 +321,18 @@ public DataFrame withColumnRenamed(String oldName, String newName) { * @throws IllegalArgumentException if {@code name} or {@code expr} is {@code null}. */ public DataFrame withColumn(String name, String expr) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (name == null) { - throw new IllegalArgumentException("withColumn name must be non-null"); - } - if (expr == null) { - throw new IllegalArgumentException("withColumn expr must be non-null"); + long h = handle.acquire(); + try { + if (name == null) { + throw new IllegalArgumentException("withColumn name must be non-null"); + } + if (expr == null) { + throw new IllegalArgumentException("withColumn expr must be non-null"); + } + return new DataFrame(withColumnExpr(h, name, expr)); + } finally { + handle.release(); } - return new DataFrame(withColumnExpr(nativeHandle, name, expr)); } /** @@ -322,16 +351,18 @@ public DataFrame unnestColumns(String... columns) { * @throws IllegalArgumentException if {@code options} or {@code columns} is {@code null}. */ public DataFrame unnestColumns(UnnestOptions options, String... columns) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (options == null) { - throw new IllegalArgumentException("unnestColumns options must be non-null"); - } - if (columns == null) { - throw new IllegalArgumentException("unnestColumns columns must be non-null"); + long h = handle.acquire(); + try { + if (options == null) { + throw new IllegalArgumentException("unnestColumns options must be non-null"); + } + if (columns == null) { + throw new IllegalArgumentException("unnestColumns columns must be non-null"); + } + return new DataFrame(unnestColumns(h, columns, options.preserveNulls())); + } finally { + handle.release(); } - return new DataFrame(unnestColumns(nativeHandle, columns, options.preserveNulls())); } // -- Set operations ------------------------------------------------------ @@ -365,7 +396,7 @@ public DataFrame unnestColumns(UnnestOptions options, String... columns) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame union(DataFrame other) { - return new DataFrame(unionRows(nativeHandle, otherHandle("union", other))); + return combine("union", other, DataFrame::unionRows); } /** @@ -377,7 +408,7 @@ public DataFrame union(DataFrame other) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame unionDistinct(DataFrame other) { - return new DataFrame(unionDistinctRows(nativeHandle, otherHandle("unionDistinct", other))); + return combine("unionDistinct", other, DataFrame::unionDistinctRows); } /** @@ -388,7 +419,7 @@ public DataFrame unionDistinct(DataFrame other) { * @throws RuntimeException if column types disagree on a shared name. */ public DataFrame unionByName(DataFrame other) { - return new DataFrame(unionByNameRows(nativeHandle, otherHandle("unionByName", other))); + return combine("unionByName", other, DataFrame::unionByNameRows); } /** @@ -399,8 +430,7 @@ public DataFrame unionByName(DataFrame other) { * @throws RuntimeException if column types disagree on a shared name. */ public DataFrame unionByNameDistinct(DataFrame other) { - return new DataFrame( - unionByNameDistinctRows(nativeHandle, otherHandle("unionByNameDistinct", other))); + return combine("unionByNameDistinct", other, DataFrame::unionByNameDistinctRows); } /** @@ -420,7 +450,7 @@ public DataFrame unionByNameDistinct(DataFrame other) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame intersect(DataFrame other) { - return new DataFrame(intersectRows(nativeHandle, otherHandle("intersect", other))); + return combine("intersect", other, DataFrame::intersectRows); } /** @@ -431,8 +461,7 @@ public DataFrame intersect(DataFrame other) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame intersectDistinct(DataFrame other) { - return new DataFrame( - intersectDistinctRows(nativeHandle, otherHandle("intersectDistinct", other))); + return combine("intersectDistinct", other, DataFrame::intersectDistinctRows); } /** @@ -452,7 +481,7 @@ public DataFrame intersectDistinct(DataFrame other) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame except(DataFrame other) { - return new DataFrame(exceptRows(nativeHandle, otherHandle("except", other))); + return combine("except", other, DataFrame::exceptRows); } /** @@ -463,24 +492,44 @@ public DataFrame except(DataFrame other) { * @throws RuntimeException if the schemas are incompatible. */ public DataFrame exceptDistinct(DataFrame other) { - return new DataFrame(exceptDistinctRows(nativeHandle, otherHandle("exceptDistinct", other))); + return combine("exceptDistinct", other, DataFrame::exceptDistinctRows); } /** - * Validate the receiver and the other DataFrame and return {@code other.nativeHandle}. Common - * preamble for the eight set-operation methods so the validation logic stays in one place. + * Validate the receiver and {@code other}, pin both handles, and apply a two-handle native set + * operation. Common body for the eight set-operation methods so the validation and pinning stay + * in one place. + * + *
Pinning two handles cannot deadlock: {@link NativeHandle#acquire()} never blocks, so a + * thread holding the receiver's pin is never waiting on the other DataFrame's lifetime. */ - private long otherHandle(String op, DataFrame other) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (other == null) { - throw new IllegalArgumentException(op + " other must be non-null"); + private DataFrame combine(String op, DataFrame other, LongBinaryOperator nativeOp) { + long left = handle.acquire(); + try { + if (other == null) { + throw new IllegalArgumentException(op + " other must be non-null"); + } + long right = other.acquireAsOperand(op + " other DataFrame is closed or already collected"); + try { + return new DataFrame(nativeOp.applyAsLong(left, right)); + } finally { + other.handle.release(); + } + } finally { + handle.release(); } - if (other.nativeHandle == 0) { - throw new IllegalStateException(op + " other DataFrame is closed or already collected"); + } + + /** + * Pin this DataFrame as the non-receiver operand of a binary operation, reporting a closed handle + * with {@code message} so the caller can distinguish which side was already gone. + */ + private long acquireAsOperand(String message) { + try { + return handle.acquire(); + } catch (IllegalStateException e) { + throw new IllegalStateException(message); } - return other.nativeHandle; } /** @@ -495,25 +544,27 @@ private long otherHandle(String op, DataFrame other) { * @throws RuntimeException if a sort column does not exist in this DataFrame's schema. */ public DataFrame sort(SortExpr... exprs) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (exprs == null) { - throw new IllegalArgumentException("sort exprs must be non-null"); - } - String[] columns = new String[exprs.length]; - boolean[] ascending = new boolean[exprs.length]; - boolean[] nullsFirst = new boolean[exprs.length]; - for (int i = 0; i < exprs.length; i++) { - SortExpr e = exprs[i]; - if (e == null) { - throw new IllegalArgumentException("sort exprs[" + i + "] must be non-null"); + long h = handle.acquire(); + try { + if (exprs == null) { + throw new IllegalArgumentException("sort exprs must be non-null"); } - columns[i] = e.column(); - ascending[i] = e.ascending(); - nullsFirst[i] = e.nullsFirst(); + String[] columns = new String[exprs.length]; + boolean[] ascending = new boolean[exprs.length]; + boolean[] nullsFirst = new boolean[exprs.length]; + for (int i = 0; i < exprs.length; i++) { + SortExpr e = exprs[i]; + if (e == null) { + throw new IllegalArgumentException("sort exprs[" + i + "] must be non-null"); + } + columns[i] = e.column(); + ascending[i] = e.ascending(); + nullsFirst[i] = e.nullsFirst(); + } + return new DataFrame(sortRows(h, columns, ascending, nullsFirst)); + } finally { + handle.release(); } - return new DataFrame(sortRows(nativeHandle, columns, ascending, nullsFirst)); } /** @@ -524,13 +575,15 @@ public DataFrame sort(SortExpr... exprs) { * @throws RuntimeException if the underlying repartition plan rejects the request. */ public DataFrame repartitionRoundRobin(int numPartitions) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (numPartitions <= 0) { - throw new IllegalArgumentException("numPartitions must be positive, was " + numPartitions); + long h = handle.acquire(); + try { + if (numPartitions <= 0) { + throw new IllegalArgumentException("numPartitions must be positive, was " + numPartitions); + } + return new DataFrame(repartitionRoundRobinRows(h, numPartitions)); + } finally { + handle.release(); } - return new DataFrame(repartitionRoundRobinRows(nativeHandle, numPartitions)); } /** @@ -544,24 +597,26 @@ public DataFrame repartitionRoundRobin(int numPartitions) { * @throws RuntimeException if a partition column does not exist in this DataFrame's schema. */ public DataFrame repartitionHash(int numPartitions, String... columns) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (numPartitions <= 0) { - throw new IllegalArgumentException("numPartitions must be positive, was " + numPartitions); - } - if (columns == null) { - throw new IllegalArgumentException("repartitionHash columns must be non-null"); - } - if (columns.length == 0) { - throw new IllegalArgumentException("repartitionHash requires at least one column"); - } - for (int i = 0; i < columns.length; i++) { - if (columns[i] == null) { - throw new IllegalArgumentException("repartitionHash columns[" + i + "] must be non-null"); + long h = handle.acquire(); + try { + if (numPartitions <= 0) { + throw new IllegalArgumentException("numPartitions must be positive, was " + numPartitions); + } + if (columns == null) { + throw new IllegalArgumentException("repartitionHash columns must be non-null"); } + if (columns.length == 0) { + throw new IllegalArgumentException("repartitionHash requires at least one column"); + } + for (int i = 0; i < columns.length; i++) { + if (columns[i] == null) { + throw new IllegalArgumentException("repartitionHash columns[" + i + "] must be non-null"); + } + } + return new DataFrame(repartitionHashRows(h, numPartitions, columns)); + } finally { + handle.release(); } - return new DataFrame(repartitionHashRows(nativeHandle, numPartitions, columns)); } /** @@ -580,8 +635,17 @@ public DataFrame repartitionHash(int numPartitions, String... columns) { */ public DataFrame join(DataFrame right, JoinType type, String[] leftCols, String[] rightCols) { checkJoinArgs(right, type, leftCols, rightCols); - return new DataFrame( - joinDataFrame(nativeHandle, right.nativeHandle, type.code(), leftCols, rightCols, null)); + long l = handle.acquire(); + try { + long r = right.acquireAsOperand("right DataFrame is closed or already collected"); + try { + return new DataFrame(joinDataFrame(l, r, type.code(), leftCols, rightCols, null)); + } finally { + right.handle.release(); + } + } finally { + handle.release(); + } } /** @@ -601,11 +665,20 @@ public DataFrame join(DataFrame right, JoinType type, String[] leftCols, String[ public DataFrame join( DataFrame right, JoinType type, String[] leftCols, String[] rightCols, String filter) { checkJoinArgs(right, type, leftCols, rightCols); - if (filter == null) { - throw new IllegalArgumentException("join filter must be non-null"); + long l = handle.acquire(); + try { + long r = right.acquireAsOperand("right DataFrame is closed or already collected"); + try { + if (filter == null) { + throw new IllegalArgumentException("join filter must be non-null"); + } + return new DataFrame(joinDataFrame(l, r, type.code(), leftCols, rightCols, filter)); + } finally { + right.handle.release(); + } + } finally { + handle.release(); } - return new DataFrame( - joinDataFrame(nativeHandle, right.nativeHandle, type.code(), leftCols, rightCols, filter)); } /** @@ -640,17 +713,20 @@ public DataFrame joinOn(DataFrame right, JoinType type, String... predicates) { throw new IllegalArgumentException("joinOn predicates must not contain null"); } } - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (right.nativeHandle == 0) { - throw new IllegalStateException("right DataFrame is closed or already collected"); + long l = handle.acquire(); + try { + long r = right.acquireAsOperand("right DataFrame is closed or already collected"); + try { + return new DataFrame(joinOnDataFrame(l, r, type.code(), predicates)); + } finally { + right.handle.release(); + } + } finally { + handle.release(); } - return new DataFrame( - joinOnDataFrame(nativeHandle, right.nativeHandle, type.code(), predicates)); } - private void checkJoinArgs( + private static void checkJoinArgs( DataFrame right, JoinType type, String[] leftCols, String[] rightCols) { if (right == null) { throw new IllegalArgumentException("join right must be non-null"); @@ -671,12 +747,6 @@ private void checkJoinArgs( + " and " + rightCols.length); } - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (right.nativeHandle == 0) { - throw new IllegalStateException("right DataFrame is closed or already collected"); - } } /** @@ -698,15 +768,17 @@ public void writeParquet(String path) { * etc.). */ public void writeParquet(String path, ParquetWriteOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); + long h = handle.acquire(); + try { + writeParquetWithOptions( + h, + path, + options.compression(), + options.singleFileOutput() != null, + options.singleFileOutput() != null && options.singleFileOutput()); + } finally { + handle.release(); } - writeParquetWithOptions( - nativeHandle, - path, - options.compression(), - options.singleFileOutput() != null, - options.singleFileOutput() != null && options.singleFileOutput()); } /** @@ -729,16 +801,18 @@ public void writeCsv(String path) { * etc.). */ public void writeCsv(String path, CsvWriteOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (path == null) { - throw new IllegalArgumentException("writeCsv path must be non-null"); - } - if (options == null) { - throw new IllegalArgumentException("writeCsv options must be non-null"); + long h = handle.acquire(); + try { + if (path == null) { + throw new IllegalArgumentException("writeCsv path must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("writeCsv options must be non-null"); + } + writeCsvWithOptions(h, path, options.toBytes()); + } finally { + handle.release(); } - writeCsvWithOptions(nativeHandle, path, options.toBytes()); } /** @@ -761,23 +835,30 @@ public void writeJson(String path) { * etc.). */ public void writeJson(String path, JsonWriteOptions options) { - if (nativeHandle == 0) { - throw new IllegalStateException("DataFrame is closed or already collected"); - } - if (path == null) { - throw new IllegalArgumentException("writeJson path must be non-null"); - } - if (options == null) { - throw new IllegalArgumentException("writeJson options must be non-null"); + long h = handle.acquire(); + try { + if (path == null) { + throw new IllegalArgumentException("writeJson path must be non-null"); + } + if (options == null) { + throw new IllegalArgumentException("writeJson options must be non-null"); + } + writeJsonWithOptions(h, path, options.toBytes()); + } finally { + handle.release(); } - writeJsonWithOptions(nativeHandle, path, options.toBytes()); } + /** + * Release the native plan. Blocks until calls already in flight on this DataFrame return, so that + * no native call can be left dereferencing a freed plan. Idempotent, and a no-op if the DataFrame + * was already consumed by {@link #collect} or {@link #executeStream}. + */ @Override public void close() { - if (nativeHandle != 0) { - closeDataFrame(nativeHandle); - nativeHandle = 0; + long h = handle.claimQuietly(); + if (h != 0) { + closeDataFrame(h); } } diff --git a/core/src/main/java/org/apache/datafusion/NativeHandle.java b/core/src/main/java/org/apache/datafusion/NativeHandle.java new file mode 100644 index 0000000..a5f2731 --- /dev/null +++ b/core/src/main/java/org/apache/datafusion/NativeHandle.java @@ -0,0 +1,172 @@ +/* + * 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; + +/** + * Owns a pointer to a native allocation and serialises its lifetime against concurrent use. + * + *
A bare {@code long} field guarded by {@code if (handle == 0) throw} is a time-of-check / + * time-of-use bug: a thread can read a live handle, and a {@code close()} on another thread can + * free the allocation before the first thread's JNI call dereferences it. This class closes that + * window by pinning the handle for the duration of each native call and deferring the free until + * every in-flight call has drained. + * + *
Usage is a pin around each native call: + * + *
{@code
+ * long h = handle.acquire();
+ * try {
+ * someNativeCall(h, ...);
+ * } finally {
+ * handle.release();
+ * }
+ * }
+ *
+ * and a claim for the operation that hands the pointer back to Rust to be freed or consumed: + * + *
{@code
+ * long h = handle.claimQuietly();
+ * if (h != 0) {
+ * closeNative(h);
+ * }
+ * }
+ *
+ * Two properties matter for callers: + * + *
The monitor is held only for bookkeeping, never across a native call, so independent + * operations on the same object still run concurrently. + */ +final class NativeHandle { + + /** Message for the {@link IllegalStateException} raised once the handle is gone. */ + private final String closedMessage; + + /** The native pointer, or 0 once claimed. */ + private long handle; + + /** Number of threads currently inside a native call with this handle pinned. */ + private int inFlight; + + /** + * Set once a claim has begun. New pins are refused from this point on, even though {@link + * #handle} stays readable until the drain completes. + */ + private boolean claimed; + + /** + * @param handle a non-zero native pointer + * @param closedMessage the {@link IllegalStateException} message to use once the handle is gone; + * lets each owner keep its own wording + */ + NativeHandle(long handle, String closedMessage) { + if (handle == 0) { + throw new IllegalArgumentException(closedMessage + ": native handle is null"); + } + this.handle = handle; + this.closedMessage = closedMessage; + } + + /** + * Pin the handle for the duration of one native call and return it. Never blocks. + * + * @throws IllegalStateException if the handle has been claimed + */ + synchronized long acquire() { + if (claimed) { + throw new IllegalStateException(closedMessage); + } + inFlight++; + return handle; + } + + /** + * Release a pin taken by {@link #acquire()}. Must be called from a {@code finally} block so that + * a native call which throws still drains. + */ + synchronized void release() { + if (inFlight == 0) { + throw new IllegalStateException("release() without a matching acquire()"); + } + inFlight--; + if (inFlight == 0 && claimed) { + notifyAll(); + } + } + + /** + * Take exclusive ownership of the handle, blocking until in-flight calls drain, and return it for + * freeing or consumption. Subsequent {@link #acquire()} calls throw. + * + * @throws IllegalStateException if the handle has already been claimed + */ + synchronized long claim() { + if (claimed) { + throw new IllegalStateException(closedMessage); + } + return drainAndTake(); + } + + /** + * As {@link #claim()}, but returns 0 rather than throwing when the handle has already been + * claimed. Used by {@code close()}, which is specified to be idempotent. + */ + synchronized long claimQuietly() { + if (claimed) { + return 0; + } + return drainAndTake(); + } + + /** + * Refuse further pins, wait for the outstanding ones to drain, then surrender the pointer. + * + *
Interruption is deferred rather than obeyed: returning early would either leak the native + * allocation or free it under a call that is still using it. The flag is restored so the caller + * can act on it after the handle is safely accounted for. + */ + private long drainAndTake() { + claimed = true; + boolean interrupted = false; + while (inFlight > 0) { + try { + wait(); + } catch (InterruptedException e) { + interrupted = true; + } + } + long claimedHandle = handle; + handle = 0; + if (interrupted) { + Thread.currentThread().interrupt(); + } + return claimedHandle; + } +} diff --git a/core/src/main/java/org/apache/datafusion/SessionContext.java b/core/src/main/java/org/apache/datafusion/SessionContext.java index b68cda5..df5a745 100644 --- a/core/src/main/java/org/apache/datafusion/SessionContext.java +++ b/core/src/main/java/org/apache/datafusion/SessionContext.java @@ -38,30 +38,41 @@ /** * A DataFusion session context. * - *
Instances are not thread-safe. Concurrent calls to any of {@link #sql}, - * {@link #registerParquet}, or {@link #close} from different threads can produce a use-after-free - * on the native side. Callers must externally synchronize, or confine each context to a single - * thread. + *
Instances are safe to share between threads. The underlying native session is itself + * thread-safe, and this class pins its native handle for the duration of every call, so a {@link + * #close} racing with work on another thread cannot free the session out from under it. A call that + * loses such a race throws {@link IllegalStateException}. + * + *
{@link #close} blocks until calls already in flight on this context have returned, then + * releases the native session. Calling it a second time, concurrently or otherwise, is a no-op. + * + *
Thread safety here covers the handle's lifetime, not the semantics of overlapping operations:
+ * registering a table concurrently with a query that reads it still races in the ordinary way, and
+ * whether the query observes the registration is undefined.
*/
public final class SessionContext implements AutoCloseable {
static {
NativeLibraryLoader.loadLibrary();
}
- private long nativeHandle;
+ private static final String CLOSED = "SessionContext is closed";
+
+ private final NativeHandle handle;
public SessionContext() {
- this.nativeHandle = createSessionContext();
- if (this.nativeHandle == 0) {
+ long nativeHandle = createSessionContext();
+ if (nativeHandle == 0) {
throw new RuntimeException("Failed to create native SessionContext");
}
+ this.handle = new NativeHandle(nativeHandle, CLOSED);
}
SessionContext(byte[] optionsBytes) {
- this.nativeHandle = createSessionContextWithOptions(optionsBytes);
- if (this.nativeHandle == 0) {
+ long nativeHandle = createSessionContextWithOptions(optionsBytes);
+ if (nativeHandle == 0) {
throw new RuntimeException("Failed to create native SessionContext");
}
+ this.handle = new NativeHandle(nativeHandle, CLOSED);
}
/** Start configuring a {@link SessionContext}. */
@@ -74,11 +85,12 @@ public static SessionContextBuilder builder() {
* until {@link DataFrame#collect} is called.
*/
public DataFrame sql(String query) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
+ long h = handle.acquire();
+ try {
+ return new DataFrame(createDataFrame(h, query));
+ } finally {
+ handle.release();
}
- long dfHandle = createDataFrame(nativeHandle, query);
- return new DataFrame(dfHandle);
}
/**
@@ -92,11 +104,12 @@ public DataFrame sql(String query) {
* planning fails.
*/
public DataFrame fromProto(byte[] planBytes) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
+ long h = handle.acquire();
+ try {
+ return new DataFrame(createDataFrameFromProto(h, planBytes));
+ } finally {
+ handle.release();
}
- long dfHandle = createDataFrameFromProto(nativeHandle, planBytes);
- return new DataFrame(dfHandle);
}
/**
@@ -126,14 +139,15 @@ public DataFrame fromProto(byte[] planBytes) {
* if the native crate was built without the {@code substrait} feature.
*/
public DataFrame fromSubstrait(byte[] planBytes) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
- }
- if (planBytes == null) {
- throw new IllegalArgumentException("fromSubstrait planBytes must be non-null");
+ long h = handle.acquire();
+ try {
+ if (planBytes == null) {
+ throw new IllegalArgumentException("fromSubstrait planBytes must be non-null");
+ }
+ return new DataFrame(createDataFrameFromSubstrait(h, planBytes));
+ } finally {
+ handle.release();
}
- long dfHandle = createDataFrameFromSubstrait(nativeHandle, planBytes);
- return new DataFrame(dfHandle);
}
/**
@@ -165,11 +179,13 @@ public DataFrame fromSubstrait(byte[] planBytes) {
* (should not happen in practice -- tracker registration is done by the constructor).
*/
public MemoryUsage memoryUsage() {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
+ long h = handle.acquire();
+ try {
+ long[] values = memoryUsageNative(h);
+ return new MemoryUsage(values[0], values[1]);
+ } finally {
+ handle.release();
}
- long[] values = memoryUsageNative(nativeHandle);
- return new MemoryUsage(values[0], values[1]);
}
/**
@@ -195,12 +211,14 @@ public MemoryUsage memoryUsage() {
* feature.
*/
public RuntimeStats runtimeStats() {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
+ long h = handle.acquire();
+ try {
+ long[] s = runtimeStatsNative(h);
+ return new RuntimeStats(
+ (int) s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10]);
+ } finally {
+ handle.release();
}
- long[] s = runtimeStatsNative(nativeHandle);
- return new RuntimeStats(
- (int) s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7], s[8], s[9], s[10]);
}
/**
@@ -210,10 +228,13 @@ public RuntimeStats runtimeStats() {
* @throws RuntimeException if {@code tableName} is not registered in this context.
*/
public Schema tableSchema(String tableName) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
+ byte[] ipcBytes;
+ long h = handle.acquire();
+ try {
+ ipcBytes = tableSchemaIpc(h, tableName);
+ } finally {
+ handle.release();
}
- byte[] ipcBytes = tableSchemaIpc(nativeHandle, tableName);
try {
return MessageSerializer.deserializeSchema(
new ReadChannel(Channels.newChannel(new ByteArrayInputStream(ipcBytes))));
@@ -241,13 +262,15 @@ public Schema tableSchema(String tableName) {
* @throws IllegalStateException if this context is closed.
*/
public String getOption(String key) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
- }
- if (key == null) {
- throw new IllegalArgumentException("getOption key must be non-null");
+ long h = handle.acquire();
+ try {
+ if (key == null) {
+ throw new IllegalArgumentException("getOption key must be non-null");
+ }
+ return getOptionNative(h, key);
+ } finally {
+ handle.release();
}
- return getOptionNative(nativeHandle, key);
}
public void registerCsv(String name, String path) {
@@ -261,15 +284,17 @@ public void registerCsv(String name, String path) {
* @throws RuntimeException if registration fails (path not found, schema inference error, etc.).
*/
public void registerCsv(String name, String path, CsvReadOptions options) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
+ long h = handle.acquire();
+ try {
+ registerCsvWithOptions(
+ h,
+ name,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
+ } finally {
+ handle.release();
}
- registerCsvWithOptions(
- nativeHandle,
- name,
- path,
- options.toBytes(),
- options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
}
/** Read a CSV file as a {@link DataFrame} without registering it. */
@@ -283,16 +308,17 @@ public DataFrame readCsv(String path) {
* @throws RuntimeException if the read fails.
*/
public DataFrame readCsv(String path, CsvReadOptions options) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
+ long h = handle.acquire();
+ try {
+ return new DataFrame(
+ readCsvWithOptions(
+ h,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) : null));
+ } finally {
+ handle.release();
}
- long dfHandle =
- readCsvWithOptions(
- nativeHandle,
- path,
- options.toBytes(),
- options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
- return new DataFrame(dfHandle);
}
public void registerJson(String name, String path) {
@@ -306,24 +332,26 @@ public void registerJson(String name, String path) {
* @throws RuntimeException if registration fails (path not found, schema inference error, etc.).
*/
public void registerJson(String name, String path, NdJsonReadOptions options) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
- }
- if (name == null) {
- throw new IllegalArgumentException("registerJson name must be non-null");
- }
- if (path == null) {
- throw new IllegalArgumentException("registerJson path must be non-null");
- }
- if (options == null) {
- throw new IllegalArgumentException("registerJson options must be non-null");
+ long h = handle.acquire();
+ try {
+ if (name == null) {
+ throw new IllegalArgumentException("registerJson name must be non-null");
+ }
+ if (path == null) {
+ throw new IllegalArgumentException("registerJson path must be non-null");
+ }
+ if (options == null) {
+ throw new IllegalArgumentException("registerJson options must be non-null");
+ }
+ registerJsonWithOptions(
+ h,
+ name,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
+ } finally {
+ handle.release();
}
- registerJsonWithOptions(
- nativeHandle,
- name,
- path,
- options.toBytes(),
- options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
}
/** Read a newline-delimited JSON file as a {@link DataFrame} without registering it. */
@@ -338,22 +366,23 @@ public DataFrame readJson(String path) {
* @throws RuntimeException if the read fails.
*/
public DataFrame readJson(String path, NdJsonReadOptions options) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
- }
- if (path == null) {
- throw new IllegalArgumentException("readJson path must be non-null");
- }
- if (options == null) {
- throw new IllegalArgumentException("readJson options must be non-null");
+ long h = handle.acquire();
+ try {
+ if (path == null) {
+ throw new IllegalArgumentException("readJson path must be non-null");
+ }
+ if (options == null) {
+ throw new IllegalArgumentException("readJson options must be non-null");
+ }
+ return new DataFrame(
+ readJsonWithOptions(
+ h,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) : null));
+ } finally {
+ handle.release();
}
- long dfHandle =
- readJsonWithOptions(
- nativeHandle,
- path,
- options.toBytes(),
- options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
- return new DataFrame(dfHandle);
}
public void registerParquet(String name, String path) {
@@ -366,15 +395,17 @@ public void registerParquet(String name, String path) {
* @throws RuntimeException if registration fails (path not found, schema mismatch, etc.).
*/
public void registerParquet(String name, String path, ParquetReadOptions options) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
+ long h = handle.acquire();
+ try {
+ registerParquetWithOptions(
+ h,
+ name,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
+ } finally {
+ handle.release();
}
- registerParquetWithOptions(
- nativeHandle,
- name,
- path,
- options.toBytes(),
- options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
}
/** Read a parquet file as a {@link DataFrame} without registering it. */
@@ -388,16 +419,17 @@ public DataFrame readParquet(String path) {
* @throws RuntimeException if the read fails.
*/
public DataFrame readParquet(String path, ParquetReadOptions options) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
+ long h = handle.acquire();
+ try {
+ return new DataFrame(
+ readParquetWithOptions(
+ h,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) : null));
+ } finally {
+ handle.release();
}
- long dfHandle =
- readParquetWithOptions(
- nativeHandle,
- path,
- options.toBytes(),
- options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
- return new DataFrame(dfHandle);
}
/** Register an Arrow IPC file (or directory of Arrow IPC files) as a table. */
@@ -414,24 +446,26 @@ public void registerArrow(String name, String path) {
* @throws RuntimeException if registration fails (path not found, schema mismatch, etc.).
*/
public void registerArrow(String name, String path, ArrowReadOptions options) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
- }
- if (name == null) {
- throw new IllegalArgumentException("registerArrow name must be non-null");
- }
- if (path == null) {
- throw new IllegalArgumentException("registerArrow path must be non-null");
- }
- if (options == null) {
- throw new IllegalArgumentException("registerArrow options must be non-null");
+ long h = handle.acquire();
+ try {
+ if (name == null) {
+ throw new IllegalArgumentException("registerArrow name must be non-null");
+ }
+ if (path == null) {
+ throw new IllegalArgumentException("registerArrow path must be non-null");
+ }
+ if (options == null) {
+ throw new IllegalArgumentException("registerArrow options must be non-null");
+ }
+ registerArrowWithOptions(
+ h,
+ name,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
+ } finally {
+ handle.release();
}
- registerArrowWithOptions(
- nativeHandle,
- name,
- path,
- options.toBytes(),
- options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
}
/** Read an Arrow IPC file as a {@link DataFrame} without registering it. */
@@ -446,22 +480,23 @@ public DataFrame readArrow(String path) {
* @throws RuntimeException if the read fails.
*/
public DataFrame readArrow(String path, ArrowReadOptions options) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
- }
- if (path == null) {
- throw new IllegalArgumentException("readArrow path must be non-null");
- }
- if (options == null) {
- throw new IllegalArgumentException("readArrow options must be non-null");
+ long h = handle.acquire();
+ try {
+ if (path == null) {
+ throw new IllegalArgumentException("readArrow path must be non-null");
+ }
+ if (options == null) {
+ throw new IllegalArgumentException("readArrow options must be non-null");
+ }
+ return new DataFrame(
+ readArrowWithOptions(
+ h,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) : null));
+ } finally {
+ handle.release();
}
- long dfHandle =
- readArrowWithOptions(
- nativeHandle,
- path,
- options.toBytes(),
- options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
- return new DataFrame(dfHandle);
}
/** Register an Avro file (or directory of Avro files) as a table. */
@@ -478,24 +513,26 @@ public void registerAvro(String name, String path) {
* @throws RuntimeException if registration fails (path not found, schema mismatch, etc.).
*/
public void registerAvro(String name, String path, AvroReadOptions options) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
- }
- if (name == null) {
- throw new IllegalArgumentException("registerAvro name must be non-null");
- }
- if (path == null) {
- throw new IllegalArgumentException("registerAvro path must be non-null");
- }
- if (options == null) {
- throw new IllegalArgumentException("registerAvro options must be non-null");
+ long h = handle.acquire();
+ try {
+ if (name == null) {
+ throw new IllegalArgumentException("registerAvro name must be non-null");
+ }
+ if (path == null) {
+ throw new IllegalArgumentException("registerAvro path must be non-null");
+ }
+ if (options == null) {
+ throw new IllegalArgumentException("registerAvro options must be non-null");
+ }
+ registerAvroWithOptions(
+ h,
+ name,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
+ } finally {
+ handle.release();
}
- registerAvroWithOptions(
- nativeHandle,
- name,
- path,
- options.toBytes(),
- options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
}
/** Read an Avro file as a {@link DataFrame} without registering it. */
@@ -510,22 +547,23 @@ public DataFrame readAvro(String path) {
* @throws RuntimeException if the read fails.
*/
public DataFrame readAvro(String path, AvroReadOptions options) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
- }
- if (path == null) {
- throw new IllegalArgumentException("readAvro path must be non-null");
- }
- if (options == null) {
- throw new IllegalArgumentException("readAvro options must be non-null");
+ long h = handle.acquire();
+ try {
+ if (path == null) {
+ throw new IllegalArgumentException("readAvro path must be non-null");
+ }
+ if (options == null) {
+ throw new IllegalArgumentException("readAvro options must be non-null");
+ }
+ return new DataFrame(
+ readAvroWithOptions(
+ h,
+ path,
+ options.toBytes(),
+ options.schema() != null ? serializeSchemaIpc(options.schema()) : null));
+ } finally {
+ handle.release();
}
- long dfHandle =
- readAvroWithOptions(
- nativeHandle,
- path,
- options.toBytes(),
- options.schema() != null ? serializeSchemaIpc(options.schema()) : null);
- return new DataFrame(dfHandle);
}
/**
@@ -539,19 +577,21 @@ public DataFrame readAvro(String path, AvroReadOptions options) {
* incompatible signature, schema serialisation failure).
*/
public void registerUdf(ScalarUdf udf) {
- if (nativeHandle == 0) {
- throw new IllegalStateException("SessionContext is closed");
+ long h = handle.acquire();
+ try {
+ java.util.Objects.requireNonNull(udf, "udf");
+ ScalarFunction impl = udf.impl();
+ String name = udf.name();
+ Volatility volatility = udf.volatility();
+ List