FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

feat(substrait): add SessionContext.fromSubstrait gated behind opt-in… · apache/datafusion-java@d078c12 · GitHub

Repository navigation

Commit d078c12

Browse files
authored
feat(substrait): add SessionContext.fromSubstrait gated behind opt-in Cargo feature (#80)
1 parent ff969cb commit d078c12

7 files changed

Lines changed: 595 additions & 24 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.
@@ -581,6 +617,8 @@ public void close() {
581617

582618
private static native long createDataFrameFromProto(long handle, byte[] planBytes);
583619

620+
private static native long createDataFrameFromSubstrait(long handle, byte[] planBytes);
621+
584622
private static native byte[] tableSchemaIpc(long handle, String tableName);
585623

586624
private static native String getOptionNative(long handle, String key);
Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,202 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.datafusion;
21+
22+
import static org.junit.jupiter.api.Assertions.assertEquals;
23+
import static org.junit.jupiter.api.Assertions.assertThrows;
24+
import static org.junit.jupiter.api.Assertions.assertTrue;
25+
26+
import java.io.IOException;
27+
import java.io.UncheckedIOException;
28+
import java.nio.charset.StandardCharsets;
29+
import java.nio.file.Files;
30+
import java.nio.file.Path;
31+
32+
import org.junit.jupiter.api.Assumptions;
33+
import org.junit.jupiter.api.BeforeAll;
34+
import org.junit.jupiter.api.Test;
35+
import org.junit.jupiter.api.io.TempDir;
36+
37+
import io.substrait.proto.NamedStruct;
38+
import io.substrait.proto.Plan;
39+
import io.substrait.proto.PlanRel;
40+
import io.substrait.proto.ReadRel;
41+
import io.substrait.proto.Rel;
42+
import io.substrait.proto.RelRoot;
43+
import io.substrait.proto.Type;
44+
import io.substrait.proto.Version;
45+
46+
/**
47+
* Round-trip tests for {@link SessionContext#fromSubstrait(byte[])}. Plans are constructed by hand
48+
* via the protobuf-Java classes shipped by {@code io.substrait:core} (test scope) — no external
49+
* SQL→Substrait compiler is required.
50+
*
51+
* <p>The {@code substrait} Cargo feature is off by default in {@code native/Cargo.toml}; if the
52+
* native crate was built without it, every test here is skipped (see {@link #checkFeatureEnabled}).
53+
* Run {@code (cd native && cargo build --features substrait)} before {@code ./mvnw test} to
54+
* exercise this class.
55+
*/
56+
class SessionContextSubstraitTest {
57+
58+
@BeforeAll
59+
static void checkFeatureEnabled() {
60+
// The native crate's feature-off stub returns this exact error; everything else (including a
61+
// prost decode failure on garbage bytes) means the feature is compiled in.
62+
try (SessionContext ctx = new SessionContext()) {
63+
ctx.fromSubstrait(new byte[] {0});
64+
// fromSubstrait should never succeed on a single zero byte — but if it ever does, the
65+
// feature is clearly compiled in.
66+
} catch (RuntimeException e) {
67+
String msg = e.getMessage() == null ? "" : e.getMessage();
68+
Assumptions.assumeFalse(
69+
msg.contains("substrait` Cargo feature"),
70+
"datafusion-jni built without the `substrait` feature; skipping Substrait tests");
71+
}
72+
}
73+
74+
/**
75+
* Build a minimal Substrait {@code Plan} that scans a registered named table {@code tableName}
76+
* with two columns {@code (id int32, v int32)} and projects them through unchanged.
77+
*/
78+
private static Plan namedTableScanPlan(String tableName) {
79+
NamedStruct schema =
80+
NamedStruct.newBuilder()
81+
.addNames("id")
82+
.addNames("v")
83+
.setStruct(
84+
Type.Struct.newBuilder()
85+
.addTypes(
86+
Type.newBuilder()
87+
.setI64(
88+
Type.I64
89+
.newBuilder()
90+
.setNullability(Type.Nullability.NULLABILITY_REQUIRED)))
91+
.addTypes(
92+
Type.newBuilder()
93+
.setI64(
94+
Type.I64
95+
.newBuilder()
96+
.setNullability(Type.Nullability.NULLABILITY_REQUIRED)))
97+
.setNullability(Type.Nullability.NULLABILITY_REQUIRED))
98+
.build();
99+
ReadRel read =
100+
ReadRel.newBuilder()
101+
.setBaseSchema(schema)
102+
.setNamedTable(ReadRel.NamedTable.newBuilder().addNames(tableName))
103+
.build();
104+
Rel rel = Rel.newBuilder().setRead(read).build();
105+
RelRoot root = RelRoot.newBuilder().setInput(rel).addNames("id").addNames("v").build();
106+
return Plan.newBuilder()
107+
.setVersion(Version.newBuilder().setMinorNumber(62))
108+
.addRelations(PlanRel.newBuilder().setRoot(root))
109+
.build();
110+
}
111+
112+
/** Write {@code id,v} CSV to a temp file with the given int rows and return the path. */
113+
private static Path writeCsv(Path dir, String name, int[][] rows) {
114+
try {
115+
Path csv = dir.resolve(name);
116+
StringBuilder sb = new StringBuilder("id,v\n");
117+
for (int[] r : rows) {
118+
sb.append(r[0]).append(',').append(r[1]).append('\n');
119+
}
120+
Files.writeString(csv, sb.toString(), StandardCharsets.UTF_8);
121+
return csv;
122+
} catch (IOException e) {
123+
throw new UncheckedIOException(e);
124+
}
125+
}
126+
127+
@Test
128+
void executesScanOfRegisteredTable(@TempDir Path tempDir) {
129+
Path csv = writeCsv(tempDir, "t.csv", new int[][] {{1, 10}, {2, 20}, {3, 30}});
130+
try (SessionContext ctx = new SessionContext()) {
131+
ctx.registerCsv("t", csv.toAbsolutePath().toString());
132+
Plan plan = namedTableScanPlan("t");
133+
try (DataFrame df = ctx.fromSubstrait(plan.toByteArray())) {
134+
assertEquals(3L, df.count());
135+
}
136+
}
137+
}
138+
139+
@Test
140+
void scanResultIsLazyDataFrame(@TempDir Path tempDir) {
141+
// The DataFrame returned by fromSubstrait is a lazy plan, just like ctx.sql(...) — we can
142+
// chain transformations against it before executing.
143+
Path csv = writeCsv(tempDir, "t.csv", new int[][] {{1, 10}, {2, 20}, {3, 30}, {4, 40}});
144+
try (SessionContext ctx = new SessionContext()) {
145+
ctx.registerCsv("t", csv.toAbsolutePath().toString());
146+
byte[] bytes = namedTableScanPlan("t").toByteArray();
147+
try (DataFrame scanned = ctx.fromSubstrait(bytes);
148+
DataFrame filtered = scanned.filter("v >= 20")) {
149+
assertEquals(3L, filtered.count());
150+
}
151+
}
152+
}
153+
154+
@Test
155+
void unknownTableThrows() {
156+
try (SessionContext ctx = new SessionContext()) {
157+
// Not registered.
158+
Plan plan = namedTableScanPlan("missing_table");
159+
RuntimeException e =
160+
assertThrows(RuntimeException.class, () -> ctx.fromSubstrait(plan.toByteArray()));
161+
// The DataFusion side surfaces a "table not found" / equivalent error message; match
162+
// loosely so the test isn't brittle to upstream wording changes.
163+
assertTrue(
164+
e.getMessage().toLowerCase().contains("missing_table")
165+
|| e.getMessage().toLowerCase().contains("not found")
166+
|| e.getMessage().toLowerCase().contains("table"),
167+
() -> "unexpected error message: " + e.getMessage());
168+
}
169+
}
170+
171+
@Test
172+
void invalidBytesThrow() {
173+
try (SessionContext ctx = new SessionContext()) {
174+
// A short non-protobuf byte sequence — prost's decode will reject this.
175+
byte[] garbage = new byte[] {(byte) 0xFF, (byte) 0xFE, 0x01, 0x02, 0x03};
176+
assertThrows(RuntimeException.class, () -> ctx.fromSubstrait(garbage));
177+
}
178+
}
179+
180+
@Test
181+
void emptyBytesThrow() {
182+
try (SessionContext ctx = new SessionContext()) {
183+
// An empty Plan has no relations; the consumer rejects it.
184+
byte[] emptyPlan = Plan.newBuilder().build().toByteArray();
185+
assertThrows(RuntimeException.class, () -> ctx.fromSubstrait(emptyPlan));
186+
}
187+
}
188+
189+
@Test
190+
void nullBytesThrowIllegalArgument() {
191+
try (SessionContext ctx = new SessionContext()) {
192+
assertThrows(IllegalArgumentException.class, () -> ctx.fromSubstrait(null));
193+
}
194+
}
195+
196+
@Test
197+
void closedContextThrows() {
198+
SessionContext ctx = new SessionContext();
199+
ctx.close();
200+
assertThrows(IllegalStateException.class, () -> ctx.fromSubstrait(new byte[] {0}));
201+
}
202+
}

0 commit comments

Comments
 (0)

Back | FazBrowse Home | New Git URL