| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -68,6 +68,11 @@ under the License. | |||
| 68 | 68 | <artifactId>avro</artifactId> | |
| 69 | 69 | <scope>test</scope> | |
| 70 | 70 | </dependency> | |
| 71 | + <dependency> | ||
| 72 | + <groupId>io.substrait</groupId> | ||
| 73 | + <artifactId>core</artifactId> | ||
| 74 | + <scope>test</scope> | ||
| 75 | + </dependency> | ||
| 71 | 76 | </dependencies> | |
| 72 | 77 | ||
| 73 | 78 | <build> | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -99,6 +99,42 @@ public DataFrame fromProto(byte[] planBytes) { | |||
| 99 | 99 | return new DataFrame(dfHandle); | |
| 100 | 100 | } | |
| 101 | 101 | ||
| 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 | + | ||
| 102 | 138 | /** | |
| 103 | 139 | * Snapshot the session's memory pool: bytes currently held and the peak observed since this | |
| 104 | 140 | * session was created. Thread-safe; can be polled while queries run. | |
@@ -581,6 +617,8 @@ public void close() { | |||
| 581 | 617 | ||
| 582 | 618 | private static native long createDataFrameFromProto(long handle, byte[] planBytes); | |
| 583 | 619 | ||
| 620 | + private static native long createDataFrameFromSubstrait(long handle, byte[] planBytes); | ||
| 621 | + | ||
| 584 | 622 | private static native byte[] tableSchemaIpc(long handle, String tableName); | |
| 585 | 623 | ||
| 586 | 624 | private static native String getOptionNative(long handle, String key); | |
| Original file line number | Diff line number | Diff 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 | + } | ||
| Back | FazBrowse Home | New Git URL |
0 commit comments