diff --git a/docs/src/.pages b/docs/src/.pages index cb47b6fc6..9b1a3e756 100644 --- a/docs/src/.pages +++ b/docs/src/.pages @@ -3,5 +3,6 @@ nav: - Install: install.md - Config: config.md - Performance: performance.md + - Streaming: streaming.md - Operations: operations - Vendors: vendors \ No newline at end of file diff --git a/docs/src/streaming.md b/docs/src/streaming.md new file mode 100644 index 000000000..a34ab4912 --- /dev/null +++ b/docs/src/streaming.md @@ -0,0 +1,139 @@ +# Spark Structured Streaming + +The Lance Spark connector supports writing data to Lance tables from a Spark Structured Streaming +query. Each non-empty micro-batch lands as one Lance append (or overwrite, in Complete output mode), +so the table is queryable between batches and time-travel reads see at most one snapshot per epoch +(empty append epochs issue no transaction — see the Exactly-once semantics section below). + +> Streaming writes are supported on all built versions (Spark **3.4+**). On Spark 3.4, +> `StreamingWrite.useCommitCoordinator` is absent from the interface, so it is inert there; Lance +> commits each epoch atomically via `CommitBuilder`, so the commit coordinator is immaterial on +> every version. + +## Quick start + +```python +(spark.readStream + .schema(schema) + .parquet("/path/to/source") + .writeStream + .format("lance") + .option("streamingQueryId", "my-query-v1") # required, see below + .option("checkpointLocation", "/path/to/checkpoint") + .trigger(availableNow=True) + .toTable("lance_ns.default.events")) +``` + +The target table must already exist; the streaming sink does not auto-create. Create it first with +plain SQL DDL: + +```sql +CREATE TABLE lance_ns.default.events (id BIGINT, name STRING); +``` + +## Required option: `streamingQueryId` + +`streamingQueryId` is the **idempotency key** for the query. It must be: + +- **Globally unique** across all concurrent streaming queries that write to the same Lance + table. Two queries sharing a `streamingQueryId` will dedupe each other's epochs and produce + incorrect results. +- **Stable across restarts** — picking a different value on restart loses replay protection for + the prior epoch. + +The connector fails fast at query start if the option is absent. + +## Output modes + +| Spark mode | Lance behavior | +|---|---| +| `append` (default) | Each epoch issues a Lance `Append` operation. | +| `complete` | Each epoch issues a Lance `Overwrite` (replacing all existing rows). Requires the upstream query to be a streaming aggregation. | +| `update` | Update-mode rows are routed through the append path — **delta rows are appended, not merged**. This matches the `SupportsStreamingUpdateAsAppend` contract; native MERGE-style upsert is not implemented. | + +## Exactly-once semantics + +Each non-empty micro-batch performs **one** Lance transaction stamped with `lance.streaming.queryId` +and `lance.streaming.epochId` in the transaction properties. On replay, the connector scans recent +transaction history (`DatasetDelta.listTransactions`) for an existing `(queryId, epochId)` pair — +finding one means the prior attempt already committed, so the current call is a no-op. + +**Empty epochs** in `append` mode are no-ops: no Lance transaction is issued. Spark's checkpoint +advances independently, and replays of empty epochs find no prior commit, see no fragments, and skip +again. In `complete` mode an empty epoch still issues **one** `Overwrite` to truncate the table to +the new (empty) result — skipping it would leave the previous epoch's rows on disk. + +### Bounded at-least-once fallback + +The scan window covers the last `lance.streaming.dedupe.lookback.versions` versions, so a duplicate +can slip through once **`lance.streaming.dedupe.lookback.versions` or more unrelated commits land +between a crash and the restart** — that many intervening commits push the prior commit out of the +window. The default lookback is 100 versions; users with very high commit churn can raise it up to +10 000. + +A duplicate can also slip through if a version inside the lookback window is removed by `VACUUM` +(or other version cleanup) before a replay: the dedupe scan can no longer read that version's +transaction, so it skips dedupe for the epoch and the replay may re-commit. Rather than failing the +query, the sink logs a `WARN` and proceeds (the window self-heals as new versions land). To avoid +this, run `VACUUM`/cleanup with a retention window larger than `lance.streaming.dedupe.lookback.versions` +worth of commits on a table that has an active streaming writer. + +## Configuration + +| Option | Default | Purpose | +|---|---|---| +| `streamingQueryId` | — | Required. Globally unique idempotency key. | +| `lance.streaming.dedupe.lookback.versions` | 100 | Number of recent versions the dedupe scan reads for an already-committed `(queryId, epochId)` pair. The scan runs on **every** epoch (not only at restart), so this is a per-commit cost. Max 10 000. Raise for high commit churn; lower to reduce per-commit dedupe-scan latency. | + +The two transaction-property keys (`lance.streaming.queryId`, `lance.streaming.epochId`) are part +of the stability contract — external tooling can rely on them when inspecting Lance transaction +history. + +## Limitations + +- **CTAS / `REPLACE TABLE`** is rejected. Stage commits commit exactly once, which is structurally + incompatible with the per-epoch streaming cadence. +- **Streaming reads** (`MicroBatchStream`) are not yet supported — only writes. Tracked + separately. +- **Row-level UPDATE/DELETE** via Spark's position-delta API is not supported in streaming. + Append-style writes only. +- **Blob v2 source-context / copy-through** is not applied on the streaming path. Streaming writes + store blob data directly (which works), but the batch-only optimization that copies blob + references from a source Lance table does not fire for streaming writes. +- **High-frequency writes (`< 30s` trigger) and Spark 4.1 Real-Time Mode** are out of scope for + this sink — the trigger→fragments/day math below shows why. A MemWAL-backed sink is the intended + path for those workloads (follow-up). +- The Lance table must exist before the streaming query starts; the sink does not auto-create. + +## Trigger cadence, fragment count & compaction + +Each non-empty micro-batch advances the Lance manifest by one version and writes at least one +fragment. Manifest size grows linearly with fragment count, and read-path performance degrades +for very large manifests. Choose a trigger cadence with the following fragment-per-day cost in +mind (Spark actually schedules the next batch at `max(triggerInterval, batchCompletionTime)`, so +busy pipelines produce fewer than the nominal count): + +| Trigger | Fragments/run | Fragments/day (typical) | Notes | +|---|---|---|---| +| `RealTime(...)` (Spark 4.1+) | — | — | ❌ Rejected by Spark's `RealTimeModeAllowlist` — only Kafka / ForeachWriter / Console / Memory sinks are allowlisted. Use one of those, or wait for the MemWAL sink. | +| `Continuous(...)`, source is Lance | — | — | ❌ Rejected at planning: Lance does not declare `TableCapability.CONTINUOUS_READ`. | +| `Continuous(...)`, non-Lance source (e.g. Kafka) | — | — | ⚠️ Passes source-side validation and reaches `Write.toStreaming()`, but per-epoch commit cadence is far too high for a Lance sink — expect catastrophic manifest growth. **Not recommended.** | +| `ProcessingTime("10s")` | 1 per epoch | ~8,640 | ❌ Not recommended for this sink; wait for the MemWAL sink. | +| `ProcessingTime("30s")` | 1 per epoch | ~2,880 | ⚠️ Borderline; requires frequent `OPTIMIZE`. | +| `ProcessingTime("1min")` | 1 per epoch | ~1,440 | ✅ Comparable to a daily batch job. | +| `ProcessingTime("5min")` | 1 per epoch | ~288 | ✅ Comfortable. | +| `AvailableNow` | 1 per run | Depends on run cadence | ✅ Ideal (batch-style backfill). | + +### `foreachBatch` is not an escape hatch + +`foreachBatch` invokes a batch write per epoch and produces the same manifest growth this sink +does. A user willing to buffer records across epochs and hand-roll `(queryId, epochId)` +idempotency in `foreachBatch` could reduce the per-epoch fragment cost, but the logic is +error-prone. The MemWAL follow-up is the intended path for high-frequency writes. + +### OPTIMIZE cadence + +For any long-running streaming query, schedule `OPTIMIZE` on the target table at a cadence +appropriate to your commit volume — for example, every 1000 epochs at a 5-minute trigger, or +daily at a 1-minute trigger. + diff --git a/integration-tests/test_lance_spark.py b/integration-tests/test_lance_spark.py index dbb9b3abd..bcbf3549a 100644 --- a/integration-tests/test_lance_spark.py +++ b/integration-tests/test_lance_spark.py @@ -13,7 +13,10 @@ """ import os +import shutil +import tempfile import time +import uuid import pytest from packaging.version import Version from pyspark.sql.types import StructType, StructField, IntegerType, StringType, DoubleType, BinaryType @@ -3112,5 +3115,94 @@ def test_update_preserves_row_ids(self, spark, test_table): assert values == {1: 100, 2: 201, 3: 301, 4: 401, 5: 501} +class TestStreamingWrite: + """Structured Streaming write sink, end-to-end through the bundled jar and + each storage backend. + + The streaming-specific logic (dedupe, empty epochs, overwrite, version + strip, lookback window) is unit-covered by ``BaseStreamingWriteTest``, which + runs against a local ``dir`` catalog. These two cases pin the real + ``writeStream`` -> object-store path that the JUnit suite cannot exercise: + the bundled jar in a live Spark cluster, the PySpark ``writeStream`` API, + and writes against a credential-bearing backend (MinIO / Azurite / Glue). + """ + + _SCHEMA = StructType([ + StructField("id", IntegerType(), False), + StructField("name", StringType(), True), + ]) + + def _write_source_batch(self, spark, batch_dir, ids): + """Write one source micro-batch as a single Parquet file.""" + rows = [(i, f"row-{i}") for i in ids] + ( + spark.createDataFrame(rows, self._SCHEMA) + .coalesce(1) + .write.parquet(batch_dir) + ) + + def _run_stream(self, spark, src_dir, checkpoint, table, query_id): + """Drain all currently-available source files into ``table`` once.""" + query = ( + spark.readStream + .schema(self._SCHEMA) + # Spark's Parquet writer creates per-batch sub-directories — let the + # streaming file source descend into them. + .option("recursiveFileLookup", "true") + .parquet(src_dir) + .writeStream + .format("lance") + .option("streamingQueryId", query_id) + .option("checkpointLocation", checkpoint) + .trigger(availableNow=True) + .toTable(table) + ) + query.processAllAvailable() + query.stop() + + def test_streaming_write_append_roundtrip(self, spark, test_table): + """Two micro-batches on one checkpoint accumulate in the target table.""" + spark.sql(f"CREATE TABLE {test_table} (id INT NOT NULL, name STRING)") + tmp = tempfile.mkdtemp(prefix="lance-stream-") + try: + src = os.path.join(tmp, "src") + ckpt = os.path.join(tmp, "ckpt") + query_id = f"qid-{uuid.uuid4().hex}" + + self._write_source_batch(spark, os.path.join(src, "b1"), [1, 2, 3]) + self._run_stream(spark, src, ckpt, test_table, query_id) + assert spark.table(test_table).count() == 3 + + # A second batch on the SAME checkpoint is a new epoch and appends. + self._write_source_batch(spark, os.path.join(src, "b2"), [4, 5]) + self._run_stream(spark, src, ckpt, test_table, query_id) + assert spark.table(test_table).count() == 5 + finally: + shutil.rmtree(tmp, ignore_errors=True) + + def test_streaming_write_dedupe_on_replay(self, spark, test_table): + """Re-running with the same streamingQueryId but a fresh checkpoint + re-delivers epoch 0; the dedupe scan finds the prior (queryId, epochId) + commit and skips it, so rows are not double-inserted.""" + spark.sql(f"CREATE TABLE {test_table} (id INT NOT NULL, name STRING)") + tmp = tempfile.mkdtemp(prefix="lance-stream-") + try: + src = os.path.join(tmp, "src") + query_id = f"qid-{uuid.uuid4().hex}" + + # A single source file => exactly one micro-batch (epoch 0) per run. + self._write_source_batch(spark, os.path.join(src, "b1"), [10, 20, 30]) + + self._run_stream(spark, src, os.path.join(tmp, "ckptA"), test_table, query_id) + assert spark.table(test_table).count() == 3 + + # Simulate a restart that lost its checkpoint but kept the queryId: + # epoch 0 is re-delivered against the same data and must dedupe. + self._run_stream(spark, src, os.path.join(tmp, "ckptB"), test_table, query_id) + assert spark.table(test_table).count() == 3 + finally: + shutil.rmtree(tmp, ignore_errors=True) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/lance-spark-3.4_2.12/src/test/java/org/lance/spark/streaming/TestStreamingWrite.java b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/streaming/TestStreamingWrite.java new file mode 100644 index 000000000..c3d1a0ce9 --- /dev/null +++ b/lance-spark-3.4_2.12/src/test/java/org/lance/spark/streaming/TestStreamingWrite.java @@ -0,0 +1,16 @@ +/* + * Licensed 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.lance.spark.streaming; + +public class TestStreamingWrite extends BaseStreamingWriteTest {} diff --git a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/streaming/TestStreamingWrite.java b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/streaming/TestStreamingWrite.java new file mode 100644 index 000000000..c3d1a0ce9 --- /dev/null +++ b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/streaming/TestStreamingWrite.java @@ -0,0 +1,16 @@ +/* + * Licensed 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.lance.spark.streaming; + +public class TestStreamingWrite extends BaseStreamingWriteTest {} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java index 14b5eeaee..76134458a 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceDataset.java @@ -58,13 +58,17 @@ public class LanceDataset private static final Set CAPABILITIES = ImmutableSet.of( - TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, TableCapability.TRUNCATE); + TableCapability.BATCH_READ, + TableCapability.BATCH_WRITE, + TableCapability.STREAMING_WRITE, + TableCapability.TRUNCATE); // Blob v2 is read as descriptor structs, but written as BINARY from sparkSchema. private static final Set CAPABILITIES_WITH_BLOB_V2 = ImmutableSet.of( TableCapability.BATCH_READ, TableCapability.BATCH_WRITE, + TableCapability.STREAMING_WRITE, TableCapability.TRUNCATE, TableCapability.ACCEPT_ANY_SCHEMA); diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java index 7bab7bef5..f51425a6f 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceSparkWriteOptions.java @@ -62,6 +62,19 @@ public class LanceSparkWriteOptions implements Serializable { public static final String CONFIG_MAX_BATCH_BYTES = "max_batch_bytes"; public static final String CONFIG_BLOB_PACK_FILE_SIZE_THRESHOLD = "blob_pack_file_size_threshold"; + // Structured Streaming options (see docs/src/streaming.md). + public static final String CONFIG_STREAMING_QUERY_ID = "streamingQueryId"; + public static final String CONFIG_STREAMING_DEDUPE_LOOKBACK_VERSIONS = + "lance.streaming.dedupe.lookback.versions"; + + // Transaction-property keys stamped into the Lance commit. The streaming dedupe scan looks for + // these keys to detect a replay of an already-committed epoch. + public static final String STREAMING_QUERY_ID_PROP = "lance.streaming.queryId"; + public static final String STREAMING_EPOCH_ID_PROP = "lance.streaming.epochId"; + + public static final int DEFAULT_STREAMING_DEDUPE_LOOKBACK_VERSIONS = 100; + public static final int MAX_STREAMING_DEDUPE_LOOKBACK_VERSIONS = 10_000; + private static final WriteMode DEFAULT_WRITE_MODE = WriteMode.APPEND; private static final boolean DEFAULT_USE_QUEUED_WRITE_BUFFER = false; private static final int DEFAULT_QUEUE_DEPTH = 8; @@ -99,6 +112,18 @@ public class LanceSparkWriteOptions implements Serializable { /** Use this version to open the dataset and apply write if set. */ private final Long version; + /** + * Idempotency key for Structured Streaming. Required for the streaming write path; null on the + * batch path. Persisted as {@link #STREAMING_QUERY_ID_PROP} in each epoch's Lance transaction. + */ + private final String streamingQueryId; + + /** + * Number of versions to scan backwards when looking for a previously-committed (queryId, epochId) + * pair during replay. Bounded above by {@link #MAX_STREAMING_DEDUPE_LOOKBACK_VERSIONS}. + */ + private final int streamingDedupeLookbackVersions; + private LanceSparkWriteOptions(Builder builder) { this.datasetUri = builder.datasetUri; this.writeMode = builder.writeMode; @@ -117,6 +142,8 @@ private LanceSparkWriteOptions(Builder builder) { this.namespace = builder.namespace; this.tableId = builder.tableId; this.version = builder.version; + this.streamingQueryId = builder.streamingQueryId; + this.streamingDedupeLookbackVersions = builder.streamingDedupeLookbackVersions; } /** Creates a new builder for LanceSparkWriteOptions. */ @@ -217,6 +244,15 @@ public Long getVersion() { return version; } + /** Nullable: non-null only on the streaming write path. */ + public String getStreamingQueryId() { + return streamingQueryId; + } + + public int getStreamingDedupeLookbackVersions() { + return streamingDedupeLookbackVersions; + } + /** Returns a builder pre-populated with all fields from this instance. */ public Builder toBuilder() { return builder() @@ -236,7 +272,9 @@ public Builder toBuilder() { .storageOptions(storageOptions) .namespace(namespace) .tableId(tableId) - .version(version); + .version(version) + .streamingQueryId(streamingQueryId) + .streamingDedupeLookbackVersions(streamingDedupeLookbackVersions); } /** Returns a copy of these options with version set to the given version. */ @@ -328,7 +366,9 @@ public boolean equals(Object o) { && Objects.equals(blobPackFileSizeThreshold, that.blobPackFileSizeThreshold) && Objects.equals(storageOptions, that.storageOptions) && Objects.equals(tableId, that.tableId) - && Objects.equals(version, that.version); + && Objects.equals(version, that.version) + && Objects.equals(streamingQueryId, that.streamingQueryId) + && streamingDedupeLookbackVersions == that.streamingDedupeLookbackVersions; } @Override @@ -349,7 +389,9 @@ public int hashCode() { blobPackFileSizeThreshold, storageOptions, tableId, - version); + version, + streamingQueryId, + streamingDedupeLookbackVersions); } /** Builder for creating LanceSparkWriteOptions instances. */ @@ -371,6 +413,8 @@ public static class Builder { private LanceNamespace namespace; private List tableId; private Long version; + private String streamingQueryId; + private int streamingDedupeLookbackVersions = DEFAULT_STREAMING_DEDUPE_LOOKBACK_VERSIONS; private Builder() {} @@ -464,6 +508,27 @@ public Builder version(Long version) { return this; } + public Builder streamingQueryId(String streamingQueryId) { + this.streamingQueryId = streamingQueryId; + return this; + } + + public Builder streamingDedupeLookbackVersions(int versions) { + // Note: avoid Preconditions.checkArgument with the (boolean, String, int, int) overload — + // that one was added in Guava 20.0 and Spark 3.x runtimes still ship Guava 14.0.1, which + // causes NoSuchMethodError at runtime even though compile + unit tests resolve a newer + // Guava from the Maven classpath. + if (versions <= 0 || versions > MAX_STREAMING_DEDUPE_LOOKBACK_VERSIONS) { + throw new IllegalArgumentException( + "streamingDedupeLookbackVersions must be in (0, " + + MAX_STREAMING_DEDUPE_LOOKBACK_VERSIONS + + "], got " + + versions); + } + this.streamingDedupeLookbackVersions = versions; + return this; + } + /** * Parses options from a map, extracting write-specific settings. * @@ -515,6 +580,25 @@ public Builder fromOptions(Map options) { Preconditions.checkArgument(parsed > 0, "blob_pack_file_size_threshold must be positive"); this.blobPackFileSizeThreshold = parsed; } + if (options.containsKey(CONFIG_STREAMING_QUERY_ID)) { + this.streamingQueryId = options.get(CONFIG_STREAMING_QUERY_ID); + } + if (options.containsKey(CONFIG_STREAMING_DEDUPE_LOOKBACK_VERSIONS)) { + String raw = options.get(CONFIG_STREAMING_DEDUPE_LOOKBACK_VERSIONS); + int parsed; + try { + parsed = Integer.parseInt(raw); + } catch (NumberFormatException e) { + throw new IllegalArgumentException( + CONFIG_STREAMING_DEDUPE_LOOKBACK_VERSIONS + + " must be an integer in (0, " + + MAX_STREAMING_DEDUPE_LOOKBACK_VERSIONS + + "], got: " + + raw, + e); + } + streamingDedupeLookbackVersions(parsed); + } return this; } diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceStreamingWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceStreamingWrite.java new file mode 100644 index 000000000..d46002368 --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceStreamingWrite.java @@ -0,0 +1,196 @@ +/* + * Licensed 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.lance.spark.write; + +import org.lance.FragmentMetadata; +import org.lance.memwal.ShardingSpec; +import org.lance.spark.LanceSparkWriteOptions; +import org.lance.spark.utils.BlobSourceContext; + +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.connector.write.DataWriter; +import org.apache.spark.sql.connector.write.PhysicalWriteInfo; +import org.apache.spark.sql.connector.write.WriterCommitMessage; +import org.apache.spark.sql.connector.write.streaming.StreamingDataWriterFactory; +import org.apache.spark.sql.connector.write.streaming.StreamingWrite; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.util.LanceArrowUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +/** + * Spark Structured Streaming sink for Lance. + * + *

Each micro-batch produces a single Lance transaction stamped with the streaming queryId and + * epochId. Replay dedupe scans recent transaction history for the same (queryId, epochId) pair; see + * {@link StreamingCommitProtocol}. + */ +public class LanceStreamingWrite implements StreamingWrite { + private static final Logger LOG = LoggerFactory.getLogger(LanceStreamingWrite.class); + + private final StructType schema; + private final LanceSparkWriteOptions writeOptions; + private final boolean overwrite; + private final Map initialStorageOptions; + private final String namespaceImpl; + private final Map namespaceProperties; + private final List tableId; + private final boolean managedVersioning; + private final ShardingSpec shardingSpec; + private final Map blobSourceContexts; + private final String queryId; + + public LanceStreamingWrite( + StructType schema, + LanceSparkWriteOptions writeOptions, + boolean overwrite, + Map initialStorageOptions, + String namespaceImpl, + Map namespaceProperties, + List tableId, + boolean managedVersioning, + ShardingSpec shardingSpec, + Map blobSourceContexts) { + this.queryId = requireStreamingQueryId(writeOptions.getStreamingQueryId()); + this.schema = schema; + // Streaming intentionally does NOT pin the dataset version. A long-running query produces + // many epochs; each commit must open the dataset at the current latest version so the + // dedupe scan window and the transaction's readVersion both reflect on-disk reality. + this.writeOptions = writeOptions; + this.overwrite = overwrite; + this.initialStorageOptions = initialStorageOptions; + this.namespaceImpl = namespaceImpl; + this.namespaceProperties = namespaceProperties; + this.tableId = tableId; + this.managedVersioning = managedVersioning; + this.shardingSpec = shardingSpec; + this.blobSourceContexts = + blobSourceContexts == null ? Collections.emptyMap() : blobSourceContexts; + } + + private static String requireStreamingQueryId(String value) { + if (value == null || value.trim().isEmpty()) { + throw new IllegalArgumentException( + "Structured Streaming writes to Lance require the '" + + LanceSparkWriteOptions.CONFIG_STREAMING_QUERY_ID + + "' option (a stable, globally unique query id used as the idempotency key)."); + } + // Trim once and use the trimmed value as the durable idempotency key, so the value validated + // here is byte-for-byte the value stamped into and compared against the transaction + // properties. Otherwise a queryId differing only by surrounding whitespace would silently + // fail to dedupe against itself. + return value.trim(); + } + + @Override + public StreamingDataWriterFactory createStreamingWriterFactory(PhysicalWriteInfo info) { + LanceDataWriter.WriterFactory inner = + new LanceDataWriter.WriterFactory( + schema, + writeOptions, + initialStorageOptions, + namespaceImpl, + namespaceProperties, + tableId, + shardingSpec, + blobSourceContexts); + return new EpochAwareWriterFactory(inner); + } + + @Override + public void commit(long epochId, WriterCommitMessage[] messages) { + List fragments = + Arrays.stream(messages) + .filter(m -> m != null) + .map(m -> (LanceBatchWrite.TaskCommit) m) + .map(LanceBatchWrite.TaskCommit::getFragments) + .map(LanceDataWriter::stripRowIdMeta) + .flatMap(List::stream) + .collect(Collectors.toList()); + + Schema arrowSchema = + LanceArrowUtils.toArrowSchema(schema, "UTC", true, writeOptions.isUseLargeVarTypes()); + // On the streaming path the output mode is the ONLY overwrite signal: Complete mode drives + // Spark's truncate() which sets `overwrite`. We deliberately ignore writeOptions.isOverwrite() + // (a batch write_mode option / catalog default) — otherwise a stray `write_mode=overwrite` + // would silently turn an Append stream into a per-epoch truncate, wiping the table on every + // (even empty) epoch. + boolean isOverwrite = overwrite; + + StreamingCommitProtocol.commitEpoch( + writeOptions, + initialStorageOptions, + fragments, + arrowSchema, + isOverwrite, + queryId, + epochId, + namespaceImpl, + namespaceProperties, + tableId, + managedVersioning); + } + + @Override + public void abort(long epochId, WriterCommitMessage[] messages) { + // No driver-side resources to release. Per-task writers release their own buffers/resolvers on + // close(); any fragment data files already written by a task whose epoch then aborts are left + // unreferenced (no transaction committed) and reclaimed by Lance cleanup/VACUUM — identical to + // the batch abort model (LanceBatchWrite#abort is likewise a no-op). + LOG.debug("streaming epoch aborted queryId={} epochId={}", queryId, epochId); + } + + /** + * Lance commits are atomic via {@link org.lance.CommitBuilder}; we do not need Spark's commit + * coordinator. Returning {@code false} here keeps the commit path identical to {@link + * LanceBatchWrite#useCommitCoordinator()}. + * + *

Annotation intentionally omitted: this method is a {@code default} on {@link StreamingWrite} + * in Spark 3.5+ but not declared at all in 3.4. Omitting {@code @Override} keeps this class + * compilable across the supported version matrix. On 3.4 the method is absent from the interface, + * so this {@code false} is inert there; it only takes effect on 3.5+. Lance's atomic commit makes + * the coordinator immaterial either way. + */ + public boolean useCommitCoordinator() { + return false; + } + + /** + * Adapter that wraps the existing batch {@link LanceDataWriter.WriterFactory} so it can fulfill + * the streaming factory contract. The {@code epochId} is ignored on the writer side — Lance + * fragments produced by a task are written under fresh UUIDs each epoch and are surfaced back to + * the driver via the {@link WriterCommitMessage}; the streaming identity is applied at commit + * time, not per-row. + */ + static final class EpochAwareWriterFactory implements StreamingDataWriterFactory { + private static final long serialVersionUID = 1L; + private final LanceDataWriter.WriterFactory delegate; + + EpochAwareWriterFactory(LanceDataWriter.WriterFactory delegate) { + this.delegate = delegate; + } + + @Override + public DataWriter createWriter(int partitionId, long taskId, long epochId) { + return delegate.createWriter(partitionId, taskId); + } + } +} diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/SparkWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/SparkWrite.java index 5639b85d2..676a54f49 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/SparkWrite.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/SparkWrite.java @@ -32,6 +32,7 @@ import org.apache.spark.sql.connector.write.Write; import org.apache.spark.sql.connector.write.WriteBuilder; import org.apache.spark.sql.connector.write.streaming.StreamingWrite; +import org.apache.spark.sql.internal.connector.SupportsStreamingUpdateAsAppend; import org.apache.spark.sql.types.StructType; import java.util.Collections; @@ -166,11 +167,35 @@ public BatchWrite toBatch() { @Override public StreamingWrite toStreaming() { - throw new UnsupportedOperationException(); + if (stagedCommit != null) { + throw new UnsupportedOperationException( + "Structured Streaming writes are incompatible with staged commits (CTAS / REPLACE " + + "TABLE). Create the target Lance table first, then start the streaming query."); + } + // Streaming always opens the dataset at the current latest version on every epoch (see + // LanceStreamingWrite's constructor note). Strip any pinned `version` so a stale snapshot + // carried in via a per-write option or a catalog default cannot break the dedupe scan window + // or the per-epoch readVersion. + LanceSparkWriteOptions streamingOptions = + writeOptions.getVersion() == null + ? writeOptions + : writeOptions.toBuilder().version(null).build(); + return new LanceStreamingWrite( + schema, + streamingOptions, + overwrite, + initialStorageOptions, + namespaceImpl, + namespaceProperties, + tableId, + managedVersioning, + shardingSpec(), + blobSourceContexts); } /** Spark write builder. */ - public static class SparkWriteBuilder implements SupportsTruncate, WriteBuilder { + public static class SparkWriteBuilder + implements SupportsTruncate, SupportsStreamingUpdateAsAppend, WriteBuilder { private final LanceSparkWriteOptions writeOptions; private final StructType schema; private boolean overwrite = false; diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/StreamingCommitProtocol.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/StreamingCommitProtocol.java new file mode 100644 index 000000000..1f0f21646 --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/StreamingCommitProtocol.java @@ -0,0 +1,249 @@ +/* + * Licensed 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.lance.spark.write; + +import org.lance.CommitBuilder; +import org.lance.Dataset; +import org.lance.FragmentMetadata; +import org.lance.Transaction; +import org.lance.delta.DatasetDelta; +import org.lance.delta.DatasetDeltaBuilder; +import org.lance.namespace.LanceNamespace; +import org.lance.operation.Append; +import org.lance.operation.Operation; +import org.lance.operation.Overwrite; +import org.lance.spark.LanceRuntime; +import org.lance.spark.LanceSparkWriteOptions; +import org.lance.spark.utils.Utils; + +import org.apache.arrow.vector.types.pojo.Schema; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +/** + * Driver-side commit helper for Structured Streaming. + * + *

Each non-empty micro-batch is committed in a single Lance {@link Transaction} stamped + * with {@link LanceSparkWriteOptions#STREAMING_QUERY_ID_PROP} and {@link + * LanceSparkWriteOptions#STREAMING_EPOCH_ID_PROP} in the transaction properties. Replay dedupe + * scans recent transaction history via {@link DatasetDelta#listTransactions()} for an existing + * (queryId, epochId) pair — finding one means the previous attempt already committed, so the + * current call is a no-op. + * + *

Empty append epochs are skipped entirely (no transaction): Spark's checkpoint advances + * regardless, and replays find no prior commit, see no fragments, and skip again. Empty + * overwrite (Complete-mode) epochs still commit one empty Overwrite to truncate the table to + * the new empty result — skipping would leave the prior epoch's rows on disk. + * + *

Failure window: the scan window covers the last {@code streamingDedupeLookbackVersions} + * versions, so a duplicate can slip through once {@code streamingDedupeLookbackVersions} or more + * unrelated commits land between a crash and the restart — that many intervening commits push the + * prior commit out of the window. Users with very high commit churn can raise the lookback up to + * {@link LanceSparkWriteOptions#MAX_STREAMING_DEDUPE_LOOKBACK_VERSIONS}. + */ +final class StreamingCommitProtocol { + private static final Logger LOG = LoggerFactory.getLogger(StreamingCommitProtocol.class); + + private StreamingCommitProtocol() {} + + /** + * Commits a streaming micro-batch. Empty {@code fragments} are skipped in append mode, but in + * overwrite (Complete) mode still commit an empty Overwrite to truncate the table. + * + * @return {@code true} if a transaction was issued, {@code false} on dedupe or empty-append skip. + */ + static boolean commitEpoch( + LanceSparkWriteOptions writeOptions, + Map initialStorageOptions, + List fragments, + Schema arrowSchema, + boolean isOverwrite, + String queryId, + long epochId, + String namespaceImpl, + Map namespaceProperties, + List tableId, + boolean managedVersioning) { + Boolean enableStableRowIds = writeOptions.getEnableStableRowIds(); + + try (Dataset ds = Utils.openDatasetBuilder(writeOptions).build()) { + long currentVersion = ds.version(); + + Optional replayVersion = findReplay(ds, currentVersion, queryId, epochId, writeOptions); + if (replayVersion.isPresent()) { + LOG.info( + "streaming epoch already committed (replay); queryId={} epochId={} foundInVersion={}", + queryId, + epochId, + replayVersion.get()); + return false; + } + + // Append: an empty epoch issues no transaction — the version does not advance, and a replay + // of an empty epoch finds no prior commit and skips again. Overwrite (Complete output mode): + // an empty epoch must STILL commit an empty Overwrite to truncate the table to the new + // (empty) result; skipping would leave the previous epoch's rows on disk, silently violating + // Complete-mode semantics. Both branches run only after the replay-dedupe check above, so a + // replayed empty overwrite remains a no-op. + if (fragments.isEmpty() && !isOverwrite) { + LOG.info( + "streaming empty append epoch skipped (no fragments); queryId={} epochId={}", + queryId, + epochId); + return false; + } + + Operation operation; + if (isOverwrite) { + operation = Overwrite.builder().fragments(fragments).schema(arrowSchema).build(); + } else { + operation = Append.builder().fragments(fragments).build(); + } + + Map txnProps = new HashMap<>(2); + txnProps.put(LanceSparkWriteOptions.STREAMING_QUERY_ID_PROP, queryId); + txnProps.put(LanceSparkWriteOptions.STREAMING_EPOCH_ID_PROP, Long.toString(epochId)); + + // Merge the namespace-vended credentials (from describeTable, carried in + // initialStorageOptions) on top of the static write storage options — mirroring + // LanceBatchWrite.commit. Without this, streaming commits to a credential-vending + // namespace/object-store table would authenticate with the static options only. + CommitBuilder commitBuilder = + new CommitBuilder(ds) + .writeParams( + LanceRuntime.mergeStorageOptions( + writeOptions.getStorageOptions(), initialStorageOptions)); + if (enableStableRowIds != null) { + commitBuilder.useStableRowIds(enableStableRowIds); + } + if (managedVersioning) { + LanceNamespace namespace = + LanceRuntime.getOrCreateNamespace(namespaceImpl, namespaceProperties); + commitBuilder + .namespaceClient(namespace) + .tableId(tableId) + .namespaceClientManagedVersioning(true); + } + + try (Transaction txn = + new Transaction.Builder() + .readVersion(currentVersion) + .operation(operation) + .transactionProperties(txnProps) + .build(); + Dataset committed = commitBuilder.execute(txn)) { + LOG.info( + "streaming epoch committed queryId={} epochId={} fragments={} version={}", + queryId, + epochId, + fragments.size(), + committed.version()); + return true; + } + } + } + + /** + * Scans transaction history backwards from {@code currentVersion} for at most {@code lookback} + * versions, returning the version where {@code (queryId, epochId)} was committed if found. + * + *

The Lance Transaction API does not expose the exact committed version, and {@code + * listTransactions()} may be sparse (lance-core omits versions that have no transaction file), so + * the index-derived version is a best-effort estimate used only for logging — the skip decision + * keys solely on the {@code (queryId, epochId)} property match, which is unaffected by + * sparseness. + * + *

If the scan fails (e.g. a version in the window was removed by VACUUM/cleanup, which makes + * {@code listTransactions()} throw), this returns {@link Optional#empty()} and logs a WARN — + * degrading to the documented bounded at-least-once fallback rather than failing the epoch and + * wedging the query on retry. + */ + private static Optional findReplay( + Dataset ds, + long currentVersion, + String queryId, + long epochId, + LanceSparkWriteOptions writeOptions) { + int lookback = writeOptions.getStreamingDedupeLookbackVersions(); + if (lookback <= 0 || currentVersion <= 0) { + return Optional.empty(); + } + long beginExclusive = Math.max(0L, currentVersion - lookback); + if (beginExclusive >= currentVersion) { + return Optional.empty(); + } + String epochIdStr = Long.toString(epochId); + try (DatasetDelta delta = + new DatasetDeltaBuilder(ds) + .withBeginVersion(beginExclusive) + .withEndVersion(currentVersion) + .build()) { + List transactions = delta.listTransactions(); + try { + // Walk newest-first. listTransactions() returns ascending-version order, but the list may + // be sparse, so the index is not a reliable version number (see findReplay's javadoc). + for (int i = transactions.size() - 1; i >= 0; i--) { + Transaction txn = transactions.get(i); + Optional> props = txn.transactionProperties(); + if (!props.isPresent()) { + continue; + } + Map p = props.get(); + if (queryId.equals(p.get(LanceSparkWriteOptions.STREAMING_QUERY_ID_PROP)) + && epochIdStr.equals(p.get(LanceSparkWriteOptions.STREAMING_EPOCH_ID_PROP))) { + return Optional.of(beginExclusive + (long) i + 1L); + } + } + } finally { + // Each history Transaction is AutoCloseable and may hold a native Arrow schema handle; + // close them all so a long-running query does not leak one handle per micro-batch. + closeQuietly(transactions); + } + } catch (RuntimeException e) { + // A version inside the lookback window may have been removed by VACUUM/cleanup since it was + // written; lance-core's listTransactions() checks out every version in the range and throws + // when one is gone. Failing here would fail the epoch — and because the version stays gone, + // Spark would retry the same micro-batch forever, permanently wedging the query. Instead we + // degrade to the documented bounded at-least-once fallback: skip dedupe for this epoch (a + // replay may then produce a duplicate) and let the commit proceed. The window self-heals as + // currentVersion advances past the removed versions. + LOG.warn( + "streaming dedupe scan failed for versions ({}, {}] (a version may have been removed " + + "by VACUUM/cleanup); skipping dedupe for this epoch (at-least-once fallback, a " + + "replay may duplicate). queryId={} epochId={}", + beginExclusive, + currentVersion, + queryId, + epochId, + e); + return Optional.empty(); + } + return Optional.empty(); + } + + private static void closeQuietly(List transactions) { + for (Transaction txn : transactions) { + try { + txn.close(); + } catch (RuntimeException e) { + LOG.debug("failed to close history transaction during dedupe scan", e); + } + } + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java index 4b34b6424..7a4a8243d 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/LanceSparkWriteOptionsTest.java @@ -29,6 +29,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; /** Tests for {@link LanceSparkWriteOptions}. */ @@ -351,4 +352,53 @@ public void testRebuiltOverwriteOptionsSurviveJavaSerialization() throws Excepti copy.getNamespace(), "namespace is transient: the non-null stub set above must not survive serialization"); } + + @Test + public void streamingDedupeLookbackDefaultsTo100() { + LanceSparkWriteOptions opts = LanceSparkWriteOptions.from(TEMP_URL); + assertEquals( + LanceSparkWriteOptions.DEFAULT_STREAMING_DEDUPE_LOOKBACK_VERSIONS, + opts.getStreamingDedupeLookbackVersions()); + } + + @Test + public void streamingOptionsRoundTripThroughOptionMap() { + Map options = new HashMap<>(); + options.put(LanceSparkWriteOptions.CONFIG_STREAMING_QUERY_ID, "q-42"); + options.put(LanceSparkWriteOptions.CONFIG_STREAMING_DEDUPE_LOOKBACK_VERSIONS, "250"); + LanceSparkWriteOptions opts = + LanceSparkWriteOptions.builder().datasetUri(TEMP_URL).fromOptions(options).build(); + assertEquals("q-42", opts.getStreamingQueryId()); + assertEquals(250, opts.getStreamingDedupeLookbackVersions()); + } + + @Test + public void streamingDedupeLookbackRejectsOutOfRangeValues() { + assertThrows( + IllegalArgumentException.class, + () -> LanceSparkWriteOptions.builder().streamingDedupeLookbackVersions(0)); + assertThrows( + IllegalArgumentException.class, + () -> LanceSparkWriteOptions.builder().streamingDedupeLookbackVersions(-1)); + assertThrows( + IllegalArgumentException.class, + () -> + LanceSparkWriteOptions.builder() + .streamingDedupeLookbackVersions( + LanceSparkWriteOptions.MAX_STREAMING_DEDUPE_LOOKBACK_VERSIONS + 1)); + } + + @Test + public void streamingDedupeLookbackNonNumericGivesActionableError() { + Map options = new HashMap<>(); + options.put(LanceSparkWriteOptions.CONFIG_STREAMING_DEDUPE_LOOKBACK_VERSIONS, "not-a-number"); + IllegalArgumentException ex = + assertThrows( + IllegalArgumentException.class, + () -> + LanceSparkWriteOptions.builder().datasetUri(TEMP_URL).fromOptions(options).build()); + assertTrue( + ex.getMessage().contains(LanceSparkWriteOptions.CONFIG_STREAMING_DEDUPE_LOOKBACK_VERSIONS), + "error should name the option key, got: " + ex.getMessage()); + } } diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/streaming/BaseStreamingWriteTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/streaming/BaseStreamingWriteTest.java new file mode 100644 index 000000000..1832e4803 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/streaming/BaseStreamingWriteTest.java @@ -0,0 +1,519 @@ +/* + * Licensed 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.lance.spark.streaming; + +import org.lance.Dataset; +import org.lance.WriteParams; +import org.lance.spark.LanceDataset; +import org.lance.spark.LanceRuntime; +import org.lance.spark.LanceSparkReadOptions; +import org.lance.spark.LanceSparkWriteOptions; +import org.lance.spark.write.LanceStreamingWrite; + +import org.apache.spark.sql.Row; +import org.apache.spark.sql.RowFactory; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.catalyst.InternalRow; +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.TableCatalog; +import org.apache.spark.sql.connector.write.DataWriter; +import org.apache.spark.sql.connector.write.WriterCommitMessage; +import org.apache.spark.sql.connector.write.streaming.StreamingDataWriterFactory; +import org.apache.spark.sql.streaming.StreamingQuery; +import org.apache.spark.sql.streaming.StreamingQueryException; +import org.apache.spark.sql.streaming.Trigger; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Metadata; +import org.apache.spark.sql.types.StructField; +import org.apache.spark.sql.types.StructType; +import org.apache.spark.unsafe.types.UTF8String; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.UUID; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * End-to-end tests for {@link LanceStreamingWrite}. Concrete extensions live in each Spark + * version-specific test module so the same suite runs on Spark 3.4, 3.5, 4.0, and 4.1. + * + *

The streaming sink is exercised both through Spark's streaming engine ({@code writeStream ... + * toTable}) and through direct {@code commit} calls — the latter lets us deterministically verify + * dedupe and empty-epoch behavior without depending on Spark's checkpoint advancement. + */ +public abstract class BaseStreamingWriteTest { + + protected SparkSession spark; + protected TableCatalog catalog; + protected final String catalogName = "lance_ns"; + + @TempDir protected Path tempDir; + + @BeforeEach + void setup() throws IOException { + spark = + SparkSession.builder() + .appName("lance-streaming-write-test") + .master("local") + .config( + "spark.sql.catalog." + catalogName, "org.lance.spark.LanceNamespaceSparkCatalog") + .config("spark.sql.catalog." + catalogName + ".impl", "dir") + .config("spark.sql.catalog." + catalogName + ".root", tempDir.toString()) + .config("spark.sql.session.timeZone", "UTC") + .getOrCreate(); + catalog = (TableCatalog) spark.sessionState().catalogManager().catalog(catalogName); + // Create the "default" namespace in manifest mode so deregisterTable works. + spark.sql("CREATE NAMESPACE IF NOT EXISTS " + catalogName + ".default"); + } + + @AfterEach + void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + // ---------- end-to-end tests ---------- + + @Test + public void testAppendHappyPath() throws Exception { + String tableName = uniqueTable("stream_append"); + String fullName = qualified(tableName); + spark.sql("CREATE TABLE " + fullName + " (id INT NOT NULL, name STRING)"); + + Path src = tempDir.resolve("src-" + tableName); + Path checkpoint = tempDir.resolve("ckpt-" + tableName); + String queryId = "qid-" + tableName; + + // Batch 1: 3 rows. + writeBatchParquet(src, new int[] {1, 2, 3}); + runStreamingBatch(src, checkpoint, fullName, queryId); + assertEquals(3L, countRows(fullName)); + + // Batch 2: 2 more rows on the same checkpoint. + writeBatchParquet(src, new int[] {4, 5}); + runStreamingBatch(src, checkpoint, fullName, queryId); + assertEquals(5L, countRows(fullName)); + } + + @Test + public void testMissingStreamingQueryIdFails() throws Exception { + String tableName = uniqueTable("stream_no_qid"); + String fullName = qualified(tableName); + spark.sql("CREATE TABLE " + fullName + " (id INT NOT NULL, name STRING)"); + + Path src = tempDir.resolve("src-" + tableName); + Path checkpoint = tempDir.resolve("ckpt-" + tableName); + writeBatchParquet(src, new int[] {1}); + + StreamingQueryException ex = + assertThrows( + StreamingQueryException.class, + () -> { + StreamingQuery q = + spark + .readStream() + .schema(rowSchema()) + .option("recursiveFileLookup", "true") + .parquet(src.toString()) + .writeStream() + .format("lance") + .option("checkpointLocation", checkpoint.toString()) + .trigger(Trigger.AvailableNow()) + .toTable(fullName); + q.processAllAvailable(); + q.stop(); + }); + assertTrue( + ex.getMessage().contains(LanceSparkWriteOptions.CONFIG_STREAMING_QUERY_ID), + "expected error to mention streamingQueryId, got: " + ex.getMessage()); + } + + @Test + public void testDedupeOnReplayedEpoch() throws Exception { + String tableName = uniqueTable("stream_dedupe"); + String fullName = qualified(tableName); + spark.sql("CREATE TABLE " + fullName + " (id INT NOT NULL, name STRING)"); + + String queryId = "qid-" + tableName; + Path src = tempDir.resolve("src-" + tableName); + Path checkpoint = tempDir.resolve("ckpt-" + tableName); + writeBatchParquet(src, new int[] {10, 20}); + runStreamingBatch(src, checkpoint, fullName, queryId); + assertEquals(2L, countRows(fullName)); + long versionAfterFirst = currentVersion(tableName); + + // Re-issue the SAME (queryId, epochId=0) directly — the very first epoch in a fresh + // checkpoint is 0, so this matches what the previous Spark query committed. NOTE: this replays + // an EMPTY epoch, so it only proves the no-op path is reached; it cannot distinguish a dedupe + // hit from the empty-epoch skip. testDedupeReplayWithNonEmptyEpoch is the real dedupe guard. + LanceStreamingWrite sink = directSink(tableName, queryId, 100); + sink.commit(0L, new WriterCommitMessage[0]); + assertEquals(2L, countRows(fullName)); + assertEquals(versionAfterFirst, currentVersion(tableName)); + } + + @Test + public void testMultipleEpochsOnSameSinkAdvanceVersionMonotonically() throws Exception { + // End-to-end regression guard for the stale-pinned-version bug: each epoch must open the + // dataset at the current latest version and advance the manifest. maxFilesPerTrigger=1 forces + // three separate micro-batches. NOTE: Spark may rebuild the StreamingWrite per micro-batch, so + // this does not guarantee one sink instance spans all epochs — the + // single-instance-across-epochs + // path is pinned directly by testDedupeReplayWithNonEmptyEpoch and + // testOverwriteEpochReplacesAndEmptyEpochTruncates, which reuse one directSink across commits. + String tableName = uniqueTable("stream_multi_epoch"); + String fullName = qualified(tableName); + spark.sql("CREATE TABLE " + fullName + " (id INT NOT NULL, name STRING)"); + + Path src = tempDir.resolve("src-" + tableName); + Path checkpoint = tempDir.resolve("ckpt-" + tableName); + String queryId = "qid-" + tableName; + + // Three separate parquet files → three epochs at maxFilesPerTrigger=1. + writeBatchParquet(src, new int[] {1}); + writeBatchParquet(src, new int[] {2}); + writeBatchParquet(src, new int[] {3}); + + long versionBefore = currentVersion(tableName); + StreamingQuery q = + spark + .readStream() + .schema(rowSchema()) + .option("recursiveFileLookup", "true") + .option("maxFilesPerTrigger", "1") + .parquet(src.toString()) + .writeStream() + .format("lance") + .option(LanceSparkWriteOptions.CONFIG_STREAMING_QUERY_ID, queryId) + .option("checkpointLocation", checkpoint.toString()) + .trigger(Trigger.AvailableNow()) + .toTable(fullName); + q.processAllAvailable(); + q.stop(); + + assertEquals(3L, countRows(fullName), "all three rows should be visible after the query"); + long versionAfter = currentVersion(tableName); + assertTrue( + versionAfter >= versionBefore + 3, + "each epoch must advance the dataset version; before=" + + versionBefore + + " after=" + + versionAfter); + } + + @Test + public void testEmptyEpochIsNoOp() throws Exception { + String tableName = uniqueTable("stream_empty"); + String fullName = qualified(tableName); + spark.sql("CREATE TABLE " + fullName + " (id INT NOT NULL, name STRING)"); + long versionBefore = currentVersion(tableName); + + String queryId = "qid-" + tableName; + LanceStreamingWrite sink = directSink(tableName, queryId, 100); + sink.commit(42L, new WriterCommitMessage[0]); + + // Empty epochs are skipped entirely — no Lance transaction is issued, so the dataset + // version does not advance. Replays of empty epochs are also no-ops. + assertEquals(0L, countRows(fullName)); + assertEquals(versionBefore, currentVersion(tableName)); + + sink.commit(42L, new WriterCommitMessage[0]); + assertEquals(versionBefore, currentVersion(tableName)); + } + + @Test + public void testDedupeReplayWithNonEmptyEpoch() throws Exception { + // Unlike testDedupeOnReplayedEpoch (which replays an EMPTY epoch and would pass even if dedupe + // were broken, since empty epochs are no-ops regardless), this replays a NON-EMPTY epoch with + // real fragments. If the (queryId, epochId) dedupe scan failed to fire, the rows would be + // inserted twice and the version would advance — so this actually pins the idempotency + // contract. + String tableName = uniqueTable("stream_dedupe_real"); + String fullName = qualified(tableName); + spark.sql("CREATE TABLE " + fullName + " (id INT NOT NULL, name STRING)"); + + String queryId = "qid-" + tableName; + LanceStreamingWrite sink = directSink(tableName, queryId, 100); + + sink.commit(0L, writeEpoch(sink, 0L, new int[] {10, 20})); + assertEquals(2L, countRows(fullName)); + long versionAfterFirst = currentVersion(tableName); + + // Replay epoch 0 with freshly written fragments (a real restart re-runs the writers and + // produces new fragment files); the dedupe scan must still skip on the (queryId, 0) match. + sink.commit(0L, writeEpoch(sink, 0L, new int[] {10, 20})); + assertEquals(2L, countRows(fullName), "replayed non-empty epoch must not double-insert"); + assertEquals( + versionAfterFirst, + currentVersion(tableName), + "replayed epoch must not issue a new transaction"); + } + + @Test + public void testDedupeExpiresOnceEpochFallsOutOfLookbackWindow() throws Exception { + // The dedupe scan is BOUNDED to the last `lookback` versions (the documented at-least-once + // fallback). With lookback=1, a single unrelated commit after epoch 0 pushes epoch 0's commit + // out of the window, so a later replay of epoch 0 is NO LONGER deduped and re-inserts. This + // pins the bounded-window contract — a regression making the window unbounded would otherwise + // ship green (every other dedupe test replays the immediately preceding epoch, always + // in-window). + String tableName = uniqueTable("stream_dedupe_expire"); + String fullName = qualified(tableName); + spark.sql("CREATE TABLE " + fullName + " (id INT NOT NULL, name STRING)"); + + String queryId = "qid-" + tableName; + LanceStreamingWrite sink = directSink(tableName, queryId, 1); + + sink.commit(0L, writeEpoch(sink, 0L, new int[] {10, 20})); + assertEquals(2L, countRows(fullName)); + + // One unrelated (different epochId) commit advances the version past epoch 0's commit, so + // epoch 0 now sits outside the 1-version lookback window. + sink.commit(1L, writeEpoch(sink, 1L, new int[] {30})); + assertEquals(3L, countRows(fullName)); + + // Replay epoch 0: its prior commit is outside the window, so the scan cannot find it and the + // rows are re-inserted — the documented bounded at-least-once fallback. + sink.commit(0L, writeEpoch(sink, 0L, new int[] {10, 20})); + assertEquals( + 5L, + countRows(fullName), + "epoch 0 fell out of the lookback window, so its replay must re-insert (at-least-once)"); + } + + @Test + public void testOverwriteEpochReplacesAndEmptyEpochTruncates() throws Exception { + // Complete output mode drives overwrite=true: each epoch fully replaces the table, and an + // empty epoch must truncate to empty (NOT be skipped, which would leave the prior epoch's rows + // on disk). + String tableName = uniqueTable("stream_overwrite"); + String fullName = qualified(tableName); + spark.sql("CREATE TABLE " + fullName + " (id INT NOT NULL, name STRING)"); + + String queryId = "qid-" + tableName; + LanceStreamingWrite sink = directSink(tableName, queryId, 100, true); + + sink.commit(0L, writeEpoch(sink, 0L, new int[] {1, 2, 3})); + assertEquals(3L, countRows(fullName)); + + // Epoch 1 overwrites: the table holds ONLY the new rows, not the union. + sink.commit(1L, writeEpoch(sink, 1L, new int[] {4, 5})); + assertEquals(2L, countRows(fullName), "overwrite epoch must replace, not append"); + + // Epoch 2 is empty: Complete mode must truncate the table to empty. + sink.commit(2L, new WriterCommitMessage[0]); + assertEquals(0L, countRows(fullName), "empty overwrite epoch must truncate to empty"); + + // The empty overwrite still stamped (queryId, 2), so replaying it is a no-op. + long versionAfterEmpty = currentVersion(tableName); + sink.commit(2L, new WriterCommitMessage[0]); + assertEquals(0L, countRows(fullName)); + assertEquals( + versionAfterEmpty, currentVersion(tableName), "replayed empty overwrite is a no-op"); + } + + @Test + public void testWriteModeOverwriteOptionDoesNotTruncateAppendStream() throws Exception { + // A stray write_mode=overwrite (e.g. a catalog default meant for the batch path) must NOT turn + // an Append-mode streaming query into a per-epoch truncate. The output mode — Spark's + // truncate()-driven `overwrite` flag — is the sole overwrite signal on the streaming path. + String tableName = uniqueTable("stream_wm_overwrite"); + String fullName = qualified(tableName); + spark.sql("CREATE TABLE " + fullName + " (id INT NOT NULL, name STRING)"); + + LanceDataset table = + (LanceDataset) catalog.loadTable(Identifier.of(new String[] {"default"}, tableName)); + LanceSparkReadOptions read = table.readOptions(); + LanceSparkWriteOptions writeOptions = + LanceSparkWriteOptions.builder() + .datasetUri(read.getDatasetUri()) + .storageOptions(read.getStorageOptions()) + .writeMode(WriteParams.WriteMode.OVERWRITE) + .streamingQueryId("qid-" + tableName) + .streamingDedupeLookbackVersions(100) + .build(); + // overwrite=false → Append output mode, despite write_mode=overwrite in the options. + LanceStreamingWrite sink = + new LanceStreamingWrite( + table.schema(), + writeOptions, + false, + Collections.emptyMap(), + null, + Collections.emptyMap(), + null, + false, + null, + Collections.emptyMap()); + + sink.commit(0L, writeEpoch(sink, 0L, new int[] {1, 2, 3})); + assertEquals(3L, countRows(fullName)); + sink.commit(1L, writeEpoch(sink, 1L, new int[] {4, 5})); + assertEquals( + 5L, countRows(fullName), "append stream must accumulate despite write_mode=overwrite"); + } + + @Test + public void testWhitespaceQueryIdRejected() throws Exception { + String tableName = uniqueTable("stream_ws_qid"); + spark.sql("CREATE TABLE " + qualified(tableName) + " (id INT NOT NULL, name STRING)"); + IllegalArgumentException ex = + assertThrows(IllegalArgumentException.class, () -> directSink(tableName, " ", 100)); + assertTrue( + ex.getMessage().contains(LanceSparkWriteOptions.CONFIG_STREAMING_QUERY_ID), + "expected error to mention streamingQueryId, got: " + ex.getMessage()); + } + + // ---------- helpers ---------- + + /** + * Drives the streaming writer factory exactly as Spark's task layer does — write rows, {@code + * commit()} to obtain the {@link WriterCommitMessage}, then {@code close()} — producing real + * Lance fragments for the given epoch. + */ + private WriterCommitMessage[] writeEpoch(LanceStreamingWrite sink, long epochId, int[] ids) + throws IOException { + StreamingDataWriterFactory factory = sink.createStreamingWriterFactory(null); + DataWriter writer = factory.createWriter(0, 0L, epochId); + try { + for (int id : ids) { + writer.write(new GenericInternalRow(new Object[] {id, UTF8String.fromString("row-" + id)})); + } + return new WriterCommitMessage[] {writer.commit()}; + } finally { + writer.close(); + } + } + + private String uniqueTable(String base) { + return base + "_" + UUID.randomUUID().toString().replace("-", ""); + } + + private String qualified(String tableName) { + return catalogName + ".default." + tableName; + } + + private StructType rowSchema() { + return new StructType( + new StructField[] { + new StructField("id", DataTypes.IntegerType, false, Metadata.empty()), + new StructField("name", DataTypes.StringType, true, Metadata.empty()) + }); + } + + private void writeBatchParquet(Path dir, int[] ids) throws IOException { + Files.createDirectories(dir); + List rows = new ArrayList<>(ids.length); + for (int id : ids) { + rows.add(RowFactory.create(id, "row-" + id)); + } + Path target = dir.resolve("batch-" + UUID.randomUUID() + ".parquet"); + spark.createDataFrame(rows, rowSchema()).coalesce(1).write().parquet(target.toString()); + } + + private void runStreamingBatch(Path srcDir, Path checkpoint, String targetTable, String queryId) + throws Exception { + StreamingQuery q = + spark + .readStream() + .schema(rowSchema()) + // Spark's parquet writer creates per-batch sub-directories — let the streaming + // file source descend into them. + .option("recursiveFileLookup", "true") + .parquet(srcDir.toString()) + .writeStream() + .format("lance") + .option(LanceSparkWriteOptions.CONFIG_STREAMING_QUERY_ID, queryId) + .option("checkpointLocation", checkpoint.toString()) + .trigger(Trigger.AvailableNow()) + .toTable(targetTable); + q.processAllAvailable(); + q.stop(); + } + + private long countRows(String fullName) { + return spark.sql("SELECT * FROM " + fullName).count(); + } + + private long currentVersion(String tableName) throws Exception { + LanceDataset table = + (LanceDataset) catalog.loadTable(Identifier.of(new String[] {"default"}, tableName)); + LanceSparkReadOptions read = table.readOptions(); + try (Dataset ds = openDataset(read)) { + return ds.version(); + } + } + + private Dataset openDataset(LanceSparkReadOptions read) { + if (read.hasNamespace()) { + return Dataset.open() + .allocator(LanceRuntime.allocator()) + .namespaceClient(read.getNamespace()) + .tableId(read.getTableId()) + .readOptions(read.toReadOptions()) + .build(); + } + return Dataset.open() + .allocator(LanceRuntime.allocator()) + .uri(read.getDatasetUri()) + .readOptions(read.toReadOptions()) + .build(); + } + + private LanceStreamingWrite directSink(String tableName, String queryId, int lookback) + throws Exception { + return directSink(tableName, queryId, lookback, false); + } + + private LanceStreamingWrite directSink( + String tableName, String queryId, int lookback, boolean overwrite) throws Exception { + LanceDataset table = + (LanceDataset) catalog.loadTable(Identifier.of(new String[] {"default"}, tableName)); + LanceSparkReadOptions read = table.readOptions(); + StructType schema = table.schema(); + + LanceSparkWriteOptions writeOptions = + LanceSparkWriteOptions.builder() + .datasetUri(read.getDatasetUri()) + .storageOptions(read.getStorageOptions()) + .streamingQueryId(queryId) + .streamingDedupeLookbackVersions(lookback) + .build(); + + return new LanceStreamingWrite( + schema, + writeOptions, + overwrite, + Collections.emptyMap(), + null, + Collections.emptyMap(), + null, + false, + null, + Collections.emptyMap()); + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/SparkWriteTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/SparkWriteTest.java index 2e84bb76c..fa7ac5a9d 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/SparkWriteTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/SparkWriteTest.java @@ -104,10 +104,71 @@ public void testToBatchReturnsLanceBatchWrite(TestInfo testInfo) { } @Test - public void testToStreamingThrowsUnsupportedOperationException(TestInfo testInfo) { + public void testToStreamingRequiresStreamingQueryId(TestInfo testInfo) { String datasetUri = createDataset(testInfo.getTestMethod().get().getName()); Write write = createBuilder(datasetUri).build(); - assertThrows(UnsupportedOperationException.class, write::toStreaming); + // Without the streamingQueryId option the constructor fails fast — exercises the option + // validation in LanceStreamingWrite#requireStreamingQueryId. + IllegalArgumentException ex = assertThrows(IllegalArgumentException.class, write::toStreaming); + assertTrue( + ex.getMessage().contains(LanceSparkWriteOptions.CONFIG_STREAMING_QUERY_ID), + "Expected error message to mention streamingQueryId, got: " + ex.getMessage()); + } + + @Test + public void testToStreamingReturnsLanceStreamingWrite(TestInfo testInfo) { + String datasetUri = createDataset(testInfo.getTestMethod().get().getName()); + LanceSparkWriteOptions writeOptions = + LanceSparkWriteOptions.builder() + .datasetUri(datasetUri) + .streamingQueryId("test-query-id") + .build(); + SparkWrite.SparkWriteBuilder builder = + new SparkWrite.SparkWriteBuilder( + SPARK_SCHEMA, + writeOptions, + Collections.emptyMap(), + null, + Collections.emptyMap(), + Arrays.asList("default", "test_table"), + false, + null, + Collections.emptyMap()); + Write write = builder.build(); + assertInstanceOf(LanceStreamingWrite.class, write.toStreaming()); + } + + @Test + public void testToStreamingRejectsStagedCommit(TestInfo testInfo) { + String datasetUri = createDataset(testInfo.getTestMethod().get().getName()); + LanceSparkWriteOptions writeOptions = + LanceSparkWriteOptions.builder() + .datasetUri(datasetUri) + .streamingQueryId("test-query-id") + .build(); + SparkWrite.SparkWriteBuilder builder = + new SparkWrite.SparkWriteBuilder( + SPARK_SCHEMA, + writeOptions, + Collections.emptyMap(), + null, + Collections.emptyMap(), + Arrays.asList("default", "test_table"), + false, + null, + Collections.emptyMap()); + // Simulate a CTAS / REPLACE TABLE flow that staged the commit before toStreaming. + builder.setStagedCommit( + StagedCommit.forNewTable( + ARROW_SCHEMA, + datasetUri, + StagedCommitOptions.pathBased(Collections.emptyMap(), false))); + Write write = builder.build(); + UnsupportedOperationException ex = + assertThrows(UnsupportedOperationException.class, write::toStreaming); + assertTrue( + ex.getMessage().contains("staged"), + "expected error to mention staged commits, got: " + ex.getMessage()); } @Test