Skip to content

Commit eaa74c6

Browse files
LantaoJinLantaoJin
committed
Merge remote-tracking branch 'upstream/main' into feat/dataframe-cancellation
2 parents 1d3011d + f9cdc15 commit eaa74c6

11 files changed

Lines changed: 799 additions & 25 deletions

File tree

‎core/pom.xml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,11 @@ under the License.
6868
<artifactId>avro</artifactId>
6969
<scope>test</scope>
7070
</dependency>
71+
<dependency>
72+
<groupId>io.substrait</groupId>
73+
<artifactId>core</artifactId>
74+
<scope>test</scope>
75+
</dependency>
7176
</dependencies>
7277

7378
<build>

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

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,42 @@ public DataFrame fromProto(byte[] planBytes) {
9999
return new DataFrame(dfHandle);
100100
}
101101

102+
/**
103+
* Decode a <a href="https://substrait.io/">Substrait</a> {@code Plan} message and return a lazy
104+
* {@link DataFrame}. The plan is not executed until {@link DataFrame#collect} or {@link
105+
* DataFrame#executeStream} is called.
106+
*
107+
* <p>{@code planBytes} must be a serialised {@code substrait.proto.Plan}. The plan is translated
108+
* to a DataFusion {@link DataFrame} against this context's catalog: any tables referenced by the
109+
* plan must already be registered (see {@link #registerCsv}, {@link #registerParquet}, etc.).
110+
*
111+
* <p>This entry point lets Java callers compile plans elsewhere — Calcite via <a
112+
* href="https://github.com/substrait-io/substrait-java">Isthmus</a>, custom planners, or any
113+
* other Substrait-emitting tool — and hand them to DataFusion without round-tripping through SQL.
114+
*
115+
* <p>Substrait support is gated behind the {@code substrait} Cargo feature on the native crate
116+
* and is <strong>off by default</strong>. Rebuild the native crate with {@code cargo build
117+
* --features substrait} (or {@code cargo build --features substrait,protoc} for hermetic builds
118+
* that vendor {@code protoc} via {@code cmake}) to enable it. If invoked against a native binary
119+
* built without the feature, this method throws {@link RuntimeException} pointing at the flag.
120+
*
121+
* @throws IllegalArgumentException if {@code planBytes} is {@code null}.
122+
* @throws IllegalStateException if this context is closed.
123+
* @throws RuntimeException if the bytes are not a valid {@code substrait.proto.Plan}, if
124+
* Substrait→DataFusion translation fails (e.g. the plan references an unregistered table), or
125+
* if the native crate was built without the {@code substrait} feature.
126+
*/
127+
public DataFrame fromSubstrait(byte[] planBytes) {
128+
if (nativeHandle == 0) {
129+
throw new IllegalStateException("SessionContext is closed");
130+
}
131+
if (planBytes == null) {
132+
throw new IllegalArgumentException("fromSubstrait planBytes must be non-null");
133+
}
134+
long dfHandle = createDataFrameFromSubstrait(nativeHandle, planBytes);
135+
return new DataFrame(dfHandle);
136+
}
137+
102138
/**
103139
* Snapshot the session's memory pool: bytes currently held and the peak observed since this
104140
* session was created. Thread-safe; can be polled while queries run.
@@ -599,6 +635,8 @@ public void close() {
599635

600636
private static native long createDataFrameFromProto(long handle, byte[] planBytes);
601637

638+
private static native long createDataFrameFromSubstrait(long handle, byte[] planBytes);
639+
602640
private static native byte[] tableSchemaIpc(long handle, String tableName);
603641

604642
private static native String getOptionNative(long handle, String key);

‎core/src/main/java/org/apache/datafusion/SessionContextBuilder.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import java.util.Map;
2626

2727
import org.apache.datafusion.protobuf.ConfigOption;
28+
import org.apache.datafusion.protobuf.DiskManagerOptions;
2829
import org.apache.datafusion.protobuf.MemoryLimit;
2930
import org.apache.datafusion.protobuf.SessionOptions;
3031

@@ -40,6 +41,8 @@ public final class SessionContextBuilder {
4041
private Long memoryLimitBytes;
4142
private Double memoryLimitFraction;
4243
private String tempDirectory;
44+
private boolean spillDisabled;
45+
private Long maxTempDirectorySize;
4346
private CacheManagerOptions cacheManager;
4447
private final LinkedHashMap<String, String> options = new LinkedHashMap<>();
4548
private final List<ObjectStoreOptions> objectStores = new ArrayList<>();
@@ -94,6 +97,39 @@ public SessionContextBuilder tempDirectory(String path) {
9497
return this;
9598
}
9699

100+
/**
101+
* Disable on-disk spill entirely. Queries that need spill fail with a {@link
102+
* ResourcesExhaustedException} rather than going to disk; useful for memory-only execution
103+
* profiles or environments without writable disk.
104+
*
105+
* <p>Mutually exclusive with {@link #tempDirectory(String)} — the combination throws at {@link
106+
* #build()} time. {@link #maxTempDirectorySize(long)} is allowed alongside this setter but is a
107+
* no-op (no directory to cap).
108+
*/
109+
public SessionContextBuilder disableSpill() {
110+
this.spillDisabled = true;
111+
return this;
112+
}
113+
114+
/**
115+
* Cap the cumulative bytes used by spill files under {@link #tempDirectory(String)}. Mirrors
116+
* upstream {@code RuntimeEnvBuilder::with_max_temp_directory_size} 1:1. Once exceeded, queries
117+
* that need more spill space fail with a {@link ResourcesExhaustedException}. Combinable with
118+
* {@link #disableSpill()} but a no-op there.
119+
*
120+
* <p>{@code 0} is accepted — upstream documents zero as legal and equivalent to "no spill
121+
* allowed". Negative values are rejected.
122+
*
123+
* @throws IllegalArgumentException if {@code bytes} is negative.
124+
*/
125+
public SessionContextBuilder maxTempDirectorySize(long bytes) {
126+
if (bytes < 0) {
127+
throw new IllegalArgumentException("maxTempDirectorySize must be non-negative, got " + bytes);
128+
}
129+
this.maxTempDirectorySize = bytes;
130+
return this;
131+
}
132+
97133
/**
98134
* Set an arbitrary {@code datafusion.*} config option by string key. Mirrors DataFusion's {@code
99135
* ConfigOptions::set(key, value)} API — see the DataFusion configuration reference for the full
@@ -219,9 +255,15 @@ public SessionContextBuilder registerObjectStore(ObjectStoreOptions options) {
219255
/**
220256
* Construct a {@link SessionContext} with the configured options.
221257
*
258+
* @throws IllegalStateException if {@link #disableSpill()} was called alongside {@link
259+
* #tempDirectory(String)} — the combination is contradictory.
222260
* @throws RuntimeException if the native side fails to construct the context.
223261
*/
224262
public SessionContext build() {
263+
if (spillDisabled && tempDirectory != null) {
264+
throw new IllegalStateException(
265+
"disableSpill() is mutually exclusive with tempDirectory(...)");
266+
}
225267
return new SessionContext(toBytes());
226268
}
227269

@@ -249,6 +291,18 @@ byte[] toBytes() {
249291
if (tempDirectory != null) {
250292
b.setTempDirectory(tempDirectory);
251293
}
294+
DiskManagerOptions.Builder dm = null;
295+
if (spillDisabled) {
296+
dm = DiskManagerOptions.newBuilder().setDisabled(true);
297+
}
298+
if (maxTempDirectorySize != null) {
299+
dm =
300+
(dm != null ? dm : DiskManagerOptions.newBuilder())
301+
.setMaxTempDirectorySize(maxTempDirectorySize);
302+
}
303+
if (dm != null) {
304+
b.setDiskManager(dm.build());
305+
}
252306
if (cacheManager != null) {
253307
b.setCacheManager(cacheManager.toProto());
254308
}

‎core/src/test/java/org/apache/datafusion/SessionContextBuilderTest.java‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -347,4 +347,119 @@ void getOptionRejectsRuntimeKeys() {
347347
"expected runtime-key error, got: " + thrown.getMessage());
348348
}
349349
}
350+
351+
// -------------------------------------------------------------------------
352+
// disk_manager surface (disable, size cap)
353+
// -------------------------------------------------------------------------
354+
355+
@Test
356+
void disableSpillRoundTripsThroughProto() throws Exception {
357+
byte[] bytes = SessionContext.builder().disableSpill().toBytes();
358+
SessionOptions parsed = SessionOptions.parseFrom(bytes);
359+
assertTrue(parsed.hasDiskManager());
360+
assertTrue(parsed.getDiskManager().getDisabled());
361+
}
362+
363+
@Test
364+
void disableSpillProducesUsableContext() throws Exception {
365+
// SELECT 1 doesn't spill, so disabling spill must not break it.
366+
try (BufferAllocator allocator = new RootAllocator();
367+
SessionContext ctx = SessionContext.builder().disableSpill().build();
368+
DataFrame df = ctx.sql("SELECT 1");
369+
ArrowReader reader = df.collect(allocator)) {
370+
assertTrue(reader.loadNextBatch());
371+
}
372+
}
373+
374+
@Test
375+
void disableSpillAndTempDirectoryConflictThrowsAtBuild() {
376+
SessionContextBuilder b = SessionContext.builder().disableSpill().tempDirectory("/tmp/x");
377+
assertThrows(IllegalStateException.class, b::build);
378+
}
379+
380+
@Test
381+
void maxTempDirectorySizeRoundTripsThroughProto() throws Exception {
382+
byte[] bytes = SessionContext.builder().maxTempDirectorySize(20L << 30).toBytes();
383+
SessionOptions parsed = SessionOptions.parseFrom(bytes);
384+
assertTrue(parsed.hasDiskManager());
385+
assertTrue(parsed.getDiskManager().hasMaxTempDirectorySize());
386+
assertEquals(20L << 30, parsed.getDiskManager().getMaxTempDirectorySize());
387+
}
388+
389+
@Test
390+
void maxTempDirectorySizeRejectsNegative() {
391+
SessionContextBuilder b = SessionContext.builder();
392+
assertThrows(IllegalArgumentException.class, () -> b.maxTempDirectorySize(-1));
393+
}
394+
395+
@Test
396+
void maxTempDirectorySizeZeroBuildsCleanly() throws Exception {
397+
// Upstream allows 0 -- "no spill allowed" -- so the Java setter mirrors
398+
// that. Sanity-check that the context constructs and SELECT 1 (which
399+
// doesn't spill) still works.
400+
try (BufferAllocator allocator = new RootAllocator();
401+
SessionContext ctx = SessionContext.builder().maxTempDirectorySize(0).build();
402+
DataFrame df = ctx.sql("SELECT 1");
403+
ArrowReader reader = df.collect(allocator)) {
404+
assertTrue(reader.loadNextBatch());
405+
}
406+
}
407+
408+
@Test
409+
void disableSpillCombinesWithMaxTempDirectorySize() throws Exception {
410+
// The cap is a no-op when spill is disabled (no directory to cap), but
411+
// setting both must not throw -- callers may have code that always
412+
// configures the cap and conditionally disables spill.
413+
byte[] bytes = SessionContext.builder().disableSpill().maxTempDirectorySize(1L << 30).toBytes();
414+
SessionOptions parsed = SessionOptions.parseFrom(bytes);
415+
assertTrue(parsed.getDiskManager().getDisabled());
416+
assertTrue(parsed.getDiskManager().hasMaxTempDirectorySize());
417+
}
418+
419+
@Test
420+
void diskManagerFieldIsAbsentWhenNothingSet() throws Exception {
421+
// Existing-callers-see-no-change contract: a builder that doesn't touch
422+
// any of the new disk_manager setters produces no disk_manager field on
423+
// the wire.
424+
byte[] bytes = SessionContext.builder().batchSize(8192).toBytes();
425+
SessionOptions parsed = SessionOptions.parseFrom(bytes);
426+
assertFalse(parsed.hasDiskManager());
427+
}
428+
429+
@Test
430+
void tempDirectoryStaysOnLegacyField() throws Exception {
431+
// tempDirectory(String) writes the existing SessionOptions.temp_directory
432+
// field, not disk_manager -- bytes identical to pre-PR behaviour for
433+
// builders that touch only that setter.
434+
byte[] bytes = SessionContext.builder().tempDirectory("/tmp/df").toBytes();
435+
SessionOptions parsed = SessionOptions.parseFrom(bytes);
436+
assertTrue(parsed.hasTempDirectory());
437+
assertEquals("/tmp/df", parsed.getTempDirectory());
438+
assertFalse(parsed.hasDiskManager());
439+
}
440+
441+
@Test
442+
void disableSpillUnderMemoryPressureThrowsResourcesExhausted() throws java.io.IOException {
443+
// Pin the Javadoc claim on disableSpill() / maxTempDirectorySize that
444+
// out-of-memory queries surface as ResourcesExhaustedException -- not
445+
// the parent DataFusionException, not a generic RuntimeException. A 1 MiB
446+
// pool plus a 200k-row sort plus no spill forces the OOM path
447+
// deterministically.
448+
try (BufferAllocator allocator = new RootAllocator();
449+
SessionContext ctx =
450+
SessionContext.builder().memoryLimit(1L << 20, 1.0).disableSpill().build()) {
451+
DataFusionException e =
452+
assertThrows(
453+
DataFusionException.class,
454+
() -> {
455+
try (DataFrame df =
456+
ctx.sql("SELECT i FROM generate_series(1, 200000) AS t(i) ORDER BY i DESC")) {
457+
df.collect(allocator).close();
458+
}
459+
});
460+
assertTrue(
461+
e instanceof ResourcesExhaustedException,
462+
"expected ResourcesExhaustedException, got " + e.getClass().getName());
463+
}
464+
}
350465
}

0 commit comments

Comments
 (0)