From 8aeca3f8cfa91ca5cddd181985fb4f243459aad1 Mon Sep 17 00:00:00 2001 From: Pucheng Yang Date: Tue, 11 Aug 2026 00:10:41 +0000 Subject: [PATCH 1/6] feat(sql): add REPLACE ... WHERE for atomic predicate-scoped overwrite Adds a `REPLACE WHERE AS ` SQL command that atomically replaces the rows matching the predicate with the result of the query, in a single Lance `Update` commit (one table version). This is the predicate-scoped analogue of Iceberg's `INSERT OVERWRITE ... PARTITION(...)`: Lance has no partition spec, so the region to replace is chosen by a row filter rather than a declared partition. The delete of matching rows and the append of new rows land in one `Operation.Update{removedFragmentIds, updatedFragments, newFragments}`, so readers never see a deleted-but-not-reinserted state and a crash cannot leave the region half-written. Fragments fully covered by the predicate are dropped; fragments that only partially match keep their non-matching rows via a deletion vector. Implementation reuses the existing distributed write pipeline: the predicate rides through as an internal write option and `LanceBatchWrite.commit()` branches to build the atomic `Update`. Adds the grammar rule + per-version AST builders, the `ReplaceWhere` logical plan, strategy mapping, and `ReplaceWhereExec`, plus unit tests and docs. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/src/operations/dml/.pages | 2 + docs/src/operations/dml/replace.md | 53 ++++ .../LanceSqlExtensionsAstBuilder.scala | 21 +- .../LanceSqlExtensionsAstBuilder.scala | 21 +- .../lance/spark/write/ReplaceWhereTest.java | 16 + .../LanceSqlExtensionsAstBuilder.scala | 21 +- .../LanceSqlExtensionsAstBuilder.scala | 21 +- .../LanceSqlExtensionsAstBuilder.scala | 21 +- .../parser/extensions/LanceSqlExtensions.g4 | 29 ++ .../java/org/lance/spark/LanceConstant.java | 8 + .../lance/spark/LanceSparkWriteOptions.java | 36 ++- .../lance/spark/write/LanceBatchWrite.java | 82 +++++- .../catalyst/plans/logical/ReplaceWhere.scala | 49 ++++ .../v2/LanceDataSourceV2Strategy.scala | 4 + .../datasources/v2/ReplaceWhereExec.scala | 70 +++++ .../spark/write/BaseReplaceWhereTest.java | 274 ++++++++++++++++++ 16 files changed, 719 insertions(+), 9 deletions(-) create mode 100644 docs/src/operations/dml/replace.md create mode 100644 lance-spark-3.5_2.12/src/test/java/org/lance/spark/write/ReplaceWhereTest.java create mode 100644 lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ReplaceWhere.scala create mode 100644 lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ReplaceWhereExec.scala create mode 100644 lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java diff --git a/docs/src/operations/dml/.pages b/docs/src/operations/dml/.pages index 057a3360b..58de64ac6 100644 --- a/docs/src/operations/dml/.pages +++ b/docs/src/operations/dml/.pages @@ -1,6 +1,8 @@ title: DML nav: - insert-into.md + - insert-overwrite.md + - replace.md - update.md - delete.md - add-columns.md diff --git a/docs/src/operations/dml/replace.md b/docs/src/operations/dml/replace.md new file mode 100644 index 000000000..53a28c006 --- /dev/null +++ b/docs/src/operations/dml/replace.md @@ -0,0 +1,53 @@ +# REPLACE ... WHERE + +`REPLACE` atomically replaces the rows of a table that match a predicate with the result of a +query, in a single table version. It is the predicate-scoped analogue of `INSERT OVERWRITE`: rather +than replacing the whole table, it replaces only the rows selected by `WHERE`. + +The delete of the matching rows and the append of the new rows are committed as one atomic Lance +`Update` operation, so readers never observe a state where the old rows are gone but the new rows +are not yet present, and a failure cannot leave the region half-written. + +## Syntax + +```sql +REPLACE
WHERE AS +``` + +- `` is any SQL boolean expression over the table's columns. It selects the existing rows + to delete. +- `` is any `SELECT` producing rows with the table's schema. Its result becomes the new rows. + +## Examples + +=== "Spark SQL" + ```sql + -- Replace one day's data with freshly computed rows + REPLACE lance.db.events + WHERE dt = '2026-08-01' + AS SELECT id, dt, value FROM staging_events WHERE dt = '2026-08-01'; + + -- Predicates may span a range or combine conditions + REPLACE lance.db.events + WHERE dt >= '2026-08-01' AND dt < '2026-08-08' + AS SELECT id, dt, value FROM staging_events; + ``` + +=== "PySpark" + ```python + spark.sql( + "REPLACE lance.db.events " + "WHERE dt = '2026-08-01' " + "AS SELECT id, dt, value FROM staging_events WHERE dt = '2026-08-01'" + ) + ``` + +## Notes + +- The `spark.sql.extensions` entry `org.lance.spark.extensions.LanceSparkSessionExtensions` must be + configured, as `REPLACE` is a Lance SQL extension. +- A row filter that matches no existing rows makes `REPLACE` behave as a plain append of the query + result. +- When a fragment contains both matching and non-matching rows, only the matching rows are removed + (via a deletion vector); the rest are preserved. A fragment whose rows all match is dropped + outright. diff --git a/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 1ca2fb957..919767777 100644 --- a/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -14,9 +14,10 @@ package org.apache.spark.sql.catalyst.parser.extensions import org.antlr.v4.runtime.ParserRuleContext +import org.antlr.v4.runtime.misc.Interval import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} +import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, ReplaceWhere, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} import org.lance.spark.utils.{FieldPathUtils, ParserUtils} import scala.collection.JavaConverters._ @@ -80,6 +81,24 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) Optimize(table, args) } + override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) + : ReplaceWhere = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + // Recover the predicate and query as raw source text by character interval. The Lance grammar + // captures them as opaque token runs, so slicing the original stream is what preserves + // operators and spacing. The query is then parsed by Spark's own parser (delegate); the + // predicate is left as text and handed to Lance at execution time. + val predicate = originalText(ctx.predicate) + val queryText = originalText(ctx.query) + val query = delegate.parsePlan(queryText) + ReplaceWhere(table, predicate, query) + } + + private def originalText(ctx: ParserRuleContext): String = { + val stream = ctx.getStart.getInputStream + stream.getText(Interval.of(ctx.getStart.getStartIndex, ctx.getStop.getStopIndex)).trim + } + override def visitVacuum(ctx: LanceSqlExtensionsParser.VacuumContext): Vacuum = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => diff --git a/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 40a228e97..1c4a64ca8 100644 --- a/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -14,9 +14,10 @@ package org.apache.spark.sql.catalyst.parser.extensions import org.antlr.v4.runtime.ParserRuleContext +import org.antlr.v4.runtime.misc.Interval import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} +import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, ReplaceWhere, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} import org.lance.spark.utils.{FieldPathUtils, ParserUtils} import scala.collection.JavaConverters._ @@ -80,6 +81,24 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) Optimize(table, args) } + override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) + : ReplaceWhere = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + // Recover the predicate and query as raw source text by character interval. The Lance grammar + // captures them as opaque token runs, so slicing the original stream is what preserves + // operators and spacing. The query is then parsed by Spark's own parser (delegate); the + // predicate is left as text and handed to Lance at execution time. + val predicate = originalText(ctx.predicate) + val queryText = originalText(ctx.query) + val query = delegate.parsePlan(queryText) + ReplaceWhere(table, predicate, query) + } + + private def originalText(ctx: ParserRuleContext): String = { + val stream = ctx.getStart.getInputStream + stream.getText(Interval.of(ctx.getStart.getStartIndex, ctx.getStop.getStopIndex)).trim + } + override def visitVacuum(ctx: LanceSqlExtensionsParser.VacuumContext): Vacuum = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => diff --git a/lance-spark-3.5_2.12/src/test/java/org/lance/spark/write/ReplaceWhereTest.java b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/write/ReplaceWhereTest.java new file mode 100644 index 000000000..0d06ff6d7 --- /dev/null +++ b/lance-spark-3.5_2.12/src/test/java/org/lance/spark/write/ReplaceWhereTest.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.write; + +public class ReplaceWhereTest extends BaseReplaceWhereTest {} diff --git a/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 86068fd67..1f3c12875 100644 --- a/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -14,9 +14,10 @@ package org.apache.spark.sql.catalyst.parser.extensions import org.antlr.v4.runtime.ParserRuleContext +import org.antlr.v4.runtime.misc.Interval import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} +import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, ReplaceWhere, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} import org.lance.spark.utils.{FieldPathUtils, ParserUtils} import scala.jdk.CollectionConverters._ @@ -80,6 +81,24 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) Optimize(table, args) } + override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) + : ReplaceWhere = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + // Recover the predicate and query as raw source text by character interval. The Lance grammar + // captures them as opaque token runs, so slicing the original stream is what preserves + // operators and spacing. The query is then parsed by Spark's own parser (delegate); the + // predicate is left as text and handed to Lance at execution time. + val predicate = originalText(ctx.predicate) + val queryText = originalText(ctx.query) + val query = delegate.parsePlan(queryText) + ReplaceWhere(table, predicate, query) + } + + private def originalText(ctx: ParserRuleContext): String = { + val stream = ctx.getStart.getInputStream + stream.getText(Interval.of(ctx.getStart.getStartIndex, ctx.getStop.getStopIndex)).trim + } + override def visitVacuum(ctx: LanceSqlExtensionsParser.VacuumContext): Vacuum = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => diff --git a/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 86068fd67..1f3c12875 100644 --- a/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -14,9 +14,10 @@ package org.apache.spark.sql.catalyst.parser.extensions import org.antlr.v4.runtime.ParserRuleContext +import org.antlr.v4.runtime.misc.Interval import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} +import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, ReplaceWhere, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} import org.lance.spark.utils.{FieldPathUtils, ParserUtils} import scala.jdk.CollectionConverters._ @@ -80,6 +81,24 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) Optimize(table, args) } + override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) + : ReplaceWhere = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + // Recover the predicate and query as raw source text by character interval. The Lance grammar + // captures them as opaque token runs, so slicing the original stream is what preserves + // operators and spacing. The query is then parsed by Spark's own parser (delegate); the + // predicate is left as text and handed to Lance at execution time. + val predicate = originalText(ctx.predicate) + val queryText = originalText(ctx.query) + val query = delegate.parsePlan(queryText) + ReplaceWhere(table, predicate, query) + } + + private def originalText(ctx: ParserRuleContext): String = { + val stream = ctx.getStart.getInputStream + stream.getText(Interval.of(ctx.getStart.getStartIndex, ctx.getStop.getStopIndex)).trim + } + override def visitVacuum(ctx: LanceSqlExtensionsParser.VacuumContext): Vacuum = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => diff --git a/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 86068fd67..1f3c12875 100644 --- a/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -14,9 +14,10 @@ package org.apache.spark.sql.catalyst.parser.extensions import org.antlr.v4.runtime.ParserRuleContext +import org.antlr.v4.runtime.misc.Interval import org.apache.spark.sql.catalyst.analysis.{UnresolvedIdentifier, UnresolvedRelation} import org.apache.spark.sql.catalyst.parser.{ParseException, ParserInterface} -import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} +import org.apache.spark.sql.catalyst.plans.logical.{AddColumnsBackfill, AddIndex, LanceCreateBranch, LanceCreateTag, LanceDropBranch, LanceDropIndex, LanceDropTag, LanceNamedArgument, LanceShowBranches, LanceShowTags, LogicalPlan, Optimize, ReplaceWhere, SetUnenforcedPrimaryKey, ShowIndexes, UpdateColumnsBackfill, Vacuum} import org.lance.spark.utils.{FieldPathUtils, ParserUtils} import scala.jdk.CollectionConverters._ @@ -80,6 +81,24 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) Optimize(table, args) } + override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) + : ReplaceWhere = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + // Recover the predicate and query as raw source text by character interval. The Lance grammar + // captures them as opaque token runs, so slicing the original stream is what preserves + // operators and spacing. The query is then parsed by Spark's own parser (delegate); the + // predicate is left as text and handed to Lance at execution time. + val predicate = originalText(ctx.predicate) + val queryText = originalText(ctx.query) + val query = delegate.parsePlan(queryText) + ReplaceWhere(table, predicate, query) + } + + private def originalText(ctx: ParserRuleContext): String = { + val stream = ctx.getStart.getInputStream + stream.getText(Interval.of(ctx.getStart.getStartIndex, ctx.getStop.getStopIndex)).trim + } + override def visitVacuum(ctx: LanceSqlExtensionsParser.VacuumContext): Vacuum = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) val args = ctx.namedArgument().asScala.map(a => diff --git a/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 b/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 index 788525249..d38fe05ec 100644 --- a/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 +++ b/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 @@ -43,6 +43,21 @@ statement | OPTIMIZE multipartIdentifier (WITH '(' (namedArgument (',' namedArgument)*)? ')')? #optimize | VACUUM multipartIdentifier (WITH '(' (namedArgument (',' namedArgument)*)? ')')? #vacuum | ALTER TABLE multipartIdentifier SET UNENFORCED PRIMARY KEY '(' columnList ')' #setUnenforcedPrimaryKey + | REPLACE multipartIdentifier WHERE predicate=predicateText AS query=queryText #replaceWhere + ; + +// The predicate (between WHERE and AS) and the query (after AS) are captured verbatim from the +// original SQL text and re-parsed with Spark's own parser, so the Lance grammar itself does not +// need to model SQL expressions or SELECT statements. The predicate runs up to the first AS +// keyword (its separator); the query is everything after AS, so it may itself contain AS (e.g. +// column aliases). Both are recovered by character interval in the AST builder, which preserves +// operators and whitespace regardless of how individual characters tokenized. +predicateText + : (~AS)+ + ; + +queryText + : .+ ; multipartIdentifier @@ -113,6 +128,7 @@ NOT: 'NOT'; OF: 'OF'; OPTIMIZE: 'OPTIMIZE'; PRIMARY: 'PRIMARY'; +REPLACE: 'REPLACE'; SET: 'SET'; SHOW: 'SHOW'; TABLE: 'TABLE'; @@ -123,6 +139,7 @@ UPDATE: 'UPDATE'; USING: 'USING'; VACUUM: 'VACUUM'; VERSION: 'VERSION'; +WHERE: 'WHERE'; WITH: 'WITH'; TRUE: 'TRUE'; @@ -173,3 +190,15 @@ fragment LETTER : [A-Z] ; +WS + : [ \t\r\n ]+ -> skip + ; + +// Catch-all so any character not matched by a specific token above (SQL operators like '>', '<', +// '*', etc.) still produces a token instead of a silently-dropped lexer error. This lets the raw +// predicate/query regions of REPLACE tokenize completely; their text is recovered by character +// interval, so the specific token kind does not matter. +ANY + : . + ; + diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceConstant.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceConstant.java index ebb5012be..7e65e2a4e 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceConstant.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceConstant.java @@ -35,6 +35,14 @@ public class LanceConstant { public static final String BACKFILL_COLUMNS_KEY = "backfill_columns"; public static final String UPDATE_COLUMNS_KEY = "update_columns"; + /** + * Internal write option carrying the row filter for a {@code REPLACE ... WHERE ... AS ...} + * command. When present, the batch write commits a single atomic {@code Update} that deletes the + * existing rows matching this predicate and appends the newly written fragments, instead of a + * plain append. Set on the driver by {@code ReplaceWhereExec}; not a user-facing option. + */ + public static final String REPLACE_WHERE_KEY = "__lance_replace_where"; + /** * Internal write option carrying the encoded blob source credential/open contexts for an INSERT * whose query reads blob columns. Set on the driver by {@code LanceBlobSourceContextRule} and 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..d31299395 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 @@ -99,6 +99,14 @@ public class LanceSparkWriteOptions implements Serializable { /** Use this version to open the dataset and apply write if set. */ private final Long version; + /** + * Row filter for a {@code REPLACE ... WHERE ...} command, or null for ordinary writes. When set, + * the batch commit deletes existing rows matching this predicate and appends the new fragments in + * a single atomic {@code Update}. Deliberately excluded from {@link #toWriteParams} so it is + * never forwarded to the native library as a storage option. + */ + private final String replaceWhere; + private LanceSparkWriteOptions(Builder builder) { this.datasetUri = builder.datasetUri; this.writeMode = builder.writeMode; @@ -117,6 +125,7 @@ private LanceSparkWriteOptions(Builder builder) { this.namespace = builder.namespace; this.tableId = builder.tableId; this.version = builder.version; + this.replaceWhere = builder.replaceWhere; } /** Creates a new builder for LanceSparkWriteOptions. */ @@ -217,6 +226,11 @@ public Long getVersion() { return version; } + /** Returns the {@code REPLACE ... WHERE ...} row filter, or null for an ordinary write. */ + public String getReplaceWhere() { + return replaceWhere; + } + /** Returns a builder pre-populated with all fields from this instance. */ public Builder toBuilder() { return builder() @@ -236,7 +250,8 @@ public Builder toBuilder() { .storageOptions(storageOptions) .namespace(namespace) .tableId(tableId) - .version(version); + .version(version) + .replaceWhere(replaceWhere); } /** Returns a copy of these options with version set to the given version. */ @@ -328,7 +343,8 @@ 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(replaceWhere, that.replaceWhere); } @Override @@ -349,7 +365,8 @@ public int hashCode() { blobPackFileSizeThreshold, storageOptions, tableId, - version); + version, + replaceWhere); } /** Builder for creating LanceSparkWriteOptions instances. */ @@ -371,6 +388,7 @@ public static class Builder { private LanceNamespace namespace; private List tableId; private Long version; + private String replaceWhere; private Builder() {} @@ -464,6 +482,12 @@ public Builder version(Long version) { return this; } + /** Sets the {@code REPLACE ... WHERE ...} row filter; null for an ordinary write. */ + public Builder replaceWhere(String replaceWhere) { + this.replaceWhere = replaceWhere; + return this; + } + /** * Parses options from a map, extracting write-specific settings. * @@ -472,6 +496,12 @@ public Builder version(Long version) { */ public Builder fromOptions(Map options) { this.storageOptions = new HashMap<>(options); + // Internal REPLACE ... WHERE filter: promote to the typed field and strip it from the + // storage options so it is never forwarded to the native library. + if (options.containsKey(LanceConstant.REPLACE_WHERE_KEY)) { + this.replaceWhere = options.get(LanceConstant.REPLACE_WHERE_KEY); + this.storageOptions.remove(LanceConstant.REPLACE_WHERE_KEY); + } if (options.containsKey(CONFIG_WRITE_MODE)) { this.writeMode = WriteMode.valueOf(options.get(CONFIG_WRITE_MODE).toUpperCase()); } diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java index 5cba05f2e..6c7f0448e 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java @@ -17,16 +17,23 @@ import org.lance.Dataset; import org.lance.FragmentMetadata; import org.lance.Transaction; +import org.lance.ipc.LanceScanner; +import org.lance.ipc.ScanOptions; import org.lance.memwal.ShardingSpec; import org.lance.namespace.LanceNamespace; import org.lance.operation.Append; import org.lance.operation.Operation; import org.lance.operation.Overwrite; +import org.lance.operation.Update; +import org.lance.spark.LanceConstant; import org.lance.spark.LanceRuntime; import org.lance.spark.LanceSparkWriteOptions; import org.lance.spark.utils.BlobSourceContext; import org.lance.spark.utils.Utils; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ArrowReader; import org.apache.arrow.vector.types.pojo.Schema; import org.apache.spark.sql.connector.write.BatchWrite; import org.apache.spark.sql.connector.write.DataWriterFactory; @@ -37,12 +44,16 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.stream.Collectors; +import static org.lance.spark.join.FragmentAwareJoinUtils.extractFragmentId; +import static org.lance.spark.join.FragmentAwareJoinUtils.extractRowIndex; + public class LanceBatchWrite implements BatchWrite { private static final Logger logger = LoggerFactory.getLogger(LanceBatchWrite.class); @@ -192,7 +203,9 @@ public void commit(WriterCommitMessage[] messages) { "version must be set (resolved in LanceBatchWrite constructor)"); try (Dataset ds = Utils.openDatasetBuilder(writeOptions).build()) { Operation operation; - if (isOverwrite) { + if (writeOptions.getReplaceWhere() != null) { + operation = buildReplaceOperation(ds, writeOptions.getReplaceWhere(), fragments); + } else if (isOverwrite) { operation = Overwrite.builder().fragments(fragments).schema(arrowSchema).build(); } else { operation = Append.builder().fragments(fragments).build(); @@ -226,6 +239,73 @@ public void commit(WriterCommitMessage[] messages) { } } + /** + * Builds an atomic {@link Update} that replaces the rows matching {@code predicate} with the + * newly written {@code newFragments}. The existing rows are found by scanning the open dataset + * for their physical row addresses; each affected fragment is rewritten with those rows deleted + * (added to {@code updatedFragments}), or dropped entirely when all of its rows match (added to + * {@code removedFragmentIds}). Deletes and the append land in a single table version. + * + *

This is correct regardless of physical layout: a fragment that only partially matches the + * predicate keeps its non-matching rows via a deletion vector, while a fragment fully covered by + * the predicate is removed outright. + */ + private static Operation buildReplaceOperation( + Dataset ds, String predicate, List newFragments) { + Map> rowIndexesByFragment = matchingRowIndexesByFragment(ds, predicate); + + List removedFragmentIds = new ArrayList<>(); + List updatedFragments = new ArrayList<>(); + for (Map.Entry> entry : rowIndexesByFragment.entrySet()) { + int fragmentId = entry.getKey(); + FragmentMetadata updated = ds.getFragment(fragmentId).deleteRows(entry.getValue()); + if (updated == null) { + // All rows in the fragment matched the predicate; drop the whole fragment. + removedFragmentIds.add((long) fragmentId); + } else { + updatedFragments.add(updated); + } + } + + return Update.builder() + .removedFragmentIds(removedFragmentIds) + .updatedFragments(updatedFragments) + .newFragments(newFragments) + .build(); + } + + /** + * Scans the dataset for rows matching {@code predicate} and groups their physical row indexes by + * fragment id, decoding the 64-bit {@code _rowaddr} (fragment id in the high 32 bits, row index + * in the low 32 bits). Returns an empty map when no existing row matches. + */ + private static Map> matchingRowIndexesByFragment( + Dataset ds, String predicate) { + Map> rowIndexesByFragment = new java.util.HashMap<>(); + ScanOptions scanOptions = + new ScanOptions.Builder() + .columns(java.util.Collections.emptyList()) + .withRowAddress(true) + .filter(predicate) + .build(); + try (LanceScanner scanner = ds.newScan(scanOptions); + ArrowReader reader = scanner.scanBatches()) { + while (reader.loadNextBatch()) { + VectorSchemaRoot batch = reader.getVectorSchemaRoot(); + FieldVector rowAddrVector = batch.getVector(LanceConstant.ROW_ADDRESS); + for (int i = 0; i < batch.getRowCount(); i++) { + long rowAddress = ((Number) rowAddrVector.getObject(i)).longValue(); + rowIndexesByFragment + .computeIfAbsent(extractFragmentId(rowAddress), k -> new ArrayList<>()) + .add(extractRowIndex(rowAddress)); + } + } + } catch (Exception e) { + throw new RuntimeException("Failed to scan rows for REPLACE ... WHERE " + predicate, e); + } + return rowIndexesByFragment; + } + @Override public void abort(WriterCommitMessage[] messages) { // For staged tables, the dataset is managed by StagedCommit (via abortStagedChanges) diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ReplaceWhere.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ReplaceWhere.scala new file mode 100644 index 000000000..2f1f30f80 --- /dev/null +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/ReplaceWhere.scala @@ -0,0 +1,49 @@ +/* + * 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.apache.spark.sql.catalyst.plans.logical + +import org.apache.spark.sql.catalyst.expressions.Attribute + +/** + * Logical plan node for the Lance `REPLACE

WHERE AS ` command. + * + * The command atomically replaces the rows of the target table matching {@code predicate} with the + * result of {@code query}, in a single table version (an atomic delete + append). It is the + * partition-overwrite analogue of Iceberg's `INSERT OVERWRITE ... PARTITION(...)`: the rows to drop + * are chosen by the predicate rather than by a declared partition spec, which Lance does not have. + * + * @param table The target Lance table whose matching rows are replaced. + * @param predicate The row filter, as raw SQL text captured verbatim from the original statement. + * It is handed to Lance to select the rows to delete, so its semantics match Lance's own filter + * evaluation rather than being resolved as a Catalyst expression here. + * @param query The source query whose result becomes the new rows for the matched region. + */ +case class ReplaceWhere( + table: LogicalPlan, + predicate: String, + query: LogicalPlan) extends Command { + + override def children: Seq[LogicalPlan] = Seq(table, query) + + override def output: Seq[Attribute] = Seq.empty + + override protected def withNewChildrenInternal( + newChildren: IndexedSeq[LogicalPlan]): ReplaceWhere = { + copy(table = newChildren(0), predicate = predicate, query = newChildren(1)) + } + + override def simpleString(maxFields: Int): String = { + s"ReplaceWhere predicate=[$predicate]" + } +} diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDataSourceV2Strategy.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDataSourceV2Strategy.scala index f17e0d52c..9c8a226f5 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDataSourceV2Strategy.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/LanceDataSourceV2Strategy.scala @@ -72,6 +72,10 @@ case class LanceDataSourceV2Strategy(session: SparkSession) extends SparkStrateg case SetUnenforcedPrimaryKey(ResolvedIdentifier(catalog, ident), columns) => SetUnenforcedPrimaryKeyExec(asTableCatalog(catalog), ident, columns) :: Nil + case ReplaceWhere(ResolvedIdentifier(catalog, ident), predicate, query) + if query.resolved => + ReplaceWhereExec(asTableCatalog(catalog), ident, predicate, query) :: Nil + case _ => Nil } diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ReplaceWhereExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ReplaceWhereExec.scala new file mode 100644 index 000000000..f932942af --- /dev/null +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ReplaceWhereExec.scala @@ -0,0 +1,70 @@ +/* + * 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.apache.spark.sql.execution.datasources.v2 + +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.plans.logical.{AppendData, LogicalPlan} +import org.apache.spark.sql.connector.catalog._ +import org.lance.spark.{LanceConstant, LanceDataset} + +/** + * Physical plan for `REPLACE
WHERE AS `. + * + * The command reuses the ordinary distributed write pipeline to materialize the query result into + * new Lance fragments, and carries the row filter through as an internal write option + * ([[LanceConstant.REPLACE_WHERE_KEY]]). The batch-write commit then turns the append into a single + * atomic `Update` that deletes the existing rows matching the predicate and adds the new fragments, + * so the replacement is one table version (an atomic delete + append) rather than two commits. + */ +case class ReplaceWhereExec( + catalog: TableCatalog, + ident: Identifier, + predicate: String, + query: LogicalPlan) + extends LeafV2CommandExec { + + override def output: Seq[Attribute] = Seq.empty + + override protected def run(): Seq[InternalRow] = { + val originalTable = catalog.loadTable(ident) match { + case lanceTable: LanceDataset => lanceTable + case other => + throw new UnsupportedOperationException( + s"REPLACE ... WHERE is only supported for Lance tables, but got: ${other.getClass}") + } + + // Write through a relation built on the target table's schema so the query is validated and + // written exactly like a normal INSERT. The predicate rides along as an internal write option; + // it is consumed at commit time to compute the rows to delete. + val relation = DataSourceV2Relation.create( + new LanceDataset( + originalTable.readOptions(), + originalTable.schema(), + originalTable.getInitialStorageOptions, + originalTable.getNamespaceImpl, + originalTable.getNamespaceProperties, + originalTable.getManagedVersioning, + originalTable.getFileFormatVersion), + Some(catalog), + Some(ident)) + + val append = + AppendData.byPosition(relation, query, Map(LanceConstant.REPLACE_WHERE_KEY -> predicate)) + val qe = session.sessionState.executePlan(append) + qe.assertCommandExecuted() + + Nil + } +} diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java new file mode 100644 index 000000000..f56deb906 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java @@ -0,0 +1,274 @@ +/* + * 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.Dataset; +import org.lance.spark.LanceDataset; + +import org.apache.arrow.memory.RootAllocator; +import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.connector.catalog.Identifier; +import org.apache.spark.sql.connector.catalog.TableCatalog; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +import java.nio.file.Path; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.UUID; +import java.util.stream.Collectors; + +/** + * Tests for the {@code REPLACE
WHERE AS } command, which atomically + * replaces the rows matching the predicate with the result of the query in a single table version. + */ +public abstract class BaseReplaceWhereTest { + protected SparkSession spark; + protected TableCatalog catalog; + protected String catalogName = "lance_ns"; + + @TempDir protected Path tempDir; + + @BeforeEach + void setup() { + spark = + SparkSession.builder() + .appName("lance-replace-where-test") + .master("local") + .config( + "spark.sql.catalog." + catalogName, "org.lance.spark.LanceNamespaceSparkCatalog") + .config("spark.sql.catalog." + catalogName + ".impl", getNsImpl()) + .config( + "spark.sql.extensions", "org.lance.spark.extensions.LanceSparkSessionExtensions") + .getOrCreate(); + + Map additionalConfigs = getAdditionalNsConfigs(); + for (Map.Entry entry : additionalConfigs.entrySet()) { + spark.conf().set("spark.sql.catalog." + catalogName + "." + entry.getKey(), entry.getValue()); + } + + catalog = (TableCatalog) spark.sessionState().catalogManager().catalog(catalogName); + spark.sql("CREATE NAMESPACE IF NOT EXISTS " + catalogName + ".default"); + } + + @AfterEach + void tearDown() { + if (spark != null) { + spark.stop(); + } + } + + protected String getNsImpl() { + return "dir"; + } + + protected Map getAdditionalNsConfigs() { + Map configs = new HashMap<>(); + configs.put("root", tempDir.toString()); + return configs; + } + + /** Replacing a partition that lives in its own fragment removes it and appends the new rows. */ + @Test + public void testReplaceSinglePartition() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + + // One INSERT per dt → one fragment per dt. + op.insert(Arrays.asList(Row.of(1, "2026-08-01", 100), Row.of(2, "2026-08-01", 200))); + op.insert(Arrays.asList(Row.of(3, "2026-08-02", 300))); + + op.replace("dt = '2026-08-01'", "SELECT 10 AS id, '2026-08-01' AS dt, 999 AS value"); + + op.check(Arrays.asList(Row.of(3, "2026-08-02", 300), Row.of(10, "2026-08-01", 999))); + } + + /** REPLACE must not touch partitions outside the predicate. */ + @Test + public void testReplaceLeavesOtherPartitionsUntouched() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + + op.insert(Arrays.asList(Row.of(1, "2026-08-01", 100))); + op.insert(Arrays.asList(Row.of(2, "2026-08-02", 200))); + op.insert(Arrays.asList(Row.of(3, "2026-08-03", 300))); + + op.replace("dt = '2026-08-02'", "SELECT 20 AS id, '2026-08-02' AS dt, 222 AS value"); + + op.check( + Arrays.asList( + Row.of(1, "2026-08-01", 100), + Row.of(3, "2026-08-03", 300), + Row.of(20, "2026-08-02", 222))); + } + + /** + * When a single fragment straddles the predicate boundary (holds both matching and non-matching + * rows), only the matching rows are removed; the rest survive via a deletion vector. + */ + @Test + public void testReplacePartiallyMatchingFragment() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + + // A single INSERT → single fragment containing two different dt values. + op.insert( + Arrays.asList( + Row.of(1, "2026-08-01", 100), + Row.of(2, "2026-08-02", 200), + Row.of(3, "2026-08-01", 300))); + + op.replace("dt = '2026-08-01'", "SELECT 9 AS id, '2026-08-01' AS dt, 900 AS value"); + + op.check(Arrays.asList(Row.of(2, "2026-08-02", 200), Row.of(9, "2026-08-01", 900))); + } + + /** Replacing a partition that has no existing rows is a plain append. */ + @Test + public void testReplaceNonExistingPartitionAppends() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + + op.insert(Arrays.asList(Row.of(1, "2026-08-01", 100))); + + op.replace("dt = '2026-08-09'", "SELECT 5 AS id, '2026-08-09' AS dt, 500 AS value"); + + op.check(Arrays.asList(Row.of(1, "2026-08-01", 100), Row.of(5, "2026-08-09", 500))); + } + + /** The replacement is a single atomic commit: exactly one new table version is produced. */ + @Test + public void testReplaceIsSingleAtomicCommit() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + + op.insert(Arrays.asList(Row.of(1, "2026-08-01", 100))); + long versionBefore = op.latestVersion(); + + op.replace("dt = '2026-08-01'", "SELECT 2 AS id, '2026-08-01' AS dt, 200 AS value"); + + Assertions.assertEquals( + versionBefore + 1, + op.latestVersion(), + "REPLACE ... WHERE must bump the table version exactly once (atomic delete + append)"); + op.check(Arrays.asList(Row.of(2, "2026-08-01", 200))); + } + + private class TableOperator { + private final SparkSession spark; + private final String catalogName; + private final String tableName; + + TableOperator(SparkSession spark, String catalogName) { + this.spark = spark; + this.catalogName = catalogName; + this.tableName = "replace_test_" + UUID.randomUUID().toString().replace("-", ""); + } + + String fullName() { + return catalogName + ".default." + tableName; + } + + void create() { + spark.sql("CREATE TABLE " + fullName() + " (id INT NOT NULL, dt STRING, value INT)"); + } + + void insert(List rows) { + spark.sql( + String.format( + "INSERT INTO %s VALUES %s", + fullName(), rows.stream().map(Row::insertSql).collect(Collectors.joining(", ")))); + } + + void replace(String predicate, String query) { + spark.sql(String.format("REPLACE %s WHERE %s AS %s", fullName(), predicate, query)); + } + + long latestVersion() { + // Resolve the table's dataset URI through the catalog, then open it to read the current + // manifest version. This mirrors how other connector tests read a table's version. + try { + String datasetUri = + ((LanceDataset) + ((TableCatalog) spark.sessionState().catalogManager().catalog(catalogName)) + .loadTable(Identifier.of(new String[] {"default"}, tableName))) + .readOptions() + .getDatasetUri(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE); + Dataset dataset = Dataset.open(datasetUri, allocator)) { + return dataset.version(); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + void check(List expected) { + List actual = + spark + .sql("SELECT id, dt, value FROM " + fullName() + " ORDER BY id") + .collectAsList() + .stream() + .map(row -> Row.of(row.getInt(0), row.getString(1), row.getInt(2))) + .collect(Collectors.toList()); + Assertions.assertEquals(expected, actual); + } + } + + private static class Row { + int id; + String dt; + int value; + + static Row of(int id, String dt, int value) { + Row row = new Row(); + row.id = id; + row.dt = dt; + row.value = value; + return row; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + Row row = (Row) o; + return id == row.id && value == row.value && Objects.equals(dt, row.dt); + } + + @Override + public int hashCode() { + return Objects.hash(id, dt, value); + } + + @Override + public String toString() { + return String.format("Row(id=%s, dt=%s, value=%s)", id, dt, value); + } + + private String insertSql() { + return String.format("(%d, '%s', %d)", id, dt, value); + } + } +} From 498fdef8c6c0f615c5db8a49a9fc10ac93aaa63b Mon Sep 17 00:00:00 2001 From: Pucheng Yang Date: Tue, 11 Aug 2026 00:33:19 +0000 Subject: [PATCH 2/6] fix(sql): split REPLACE on top-level AS and bound deletion memory Addresses lance-gatekeeper review feedback on the REPLACE ... WHERE command: 1. Predicate/query split: the grammar previously delimited on the first AS token, which broke predicates containing AS (e.g. CAST(dt AS STRING)). The grammar now captures everything after WHERE as one raw region, and the AST builder splits it at the first *top-level* AS via ParserUtils.splitReplaceBody (honoring parentheses, string/backtick literals, and line/block comments). Adds a CAST-in-predicate regression test. 2. Deletion planning memory: LanceBatchWrite collected a boxed Integer per matched row across all fragments, so a partition-scale replace could OOM the driver. Deletions are now accumulated as a compressed RoaringBitmap per fragment and materialized to a row-index list one fragment at a time, keeping driver memory bounded regardless of total matched row count. This mirrors the existing SparkPositionDeltaWrite pattern. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../LanceSqlExtensionsAstBuilder.scala | 15 ++- .../LanceSqlExtensionsAstBuilder.scala | 15 ++- .../LanceSqlExtensionsAstBuilder.scala | 15 ++- .../LanceSqlExtensionsAstBuilder.scala | 15 ++- .../LanceSqlExtensionsAstBuilder.scala | 15 ++- .../parser/extensions/LanceSqlExtensions.g4 | 20 ++-- .../org/lance/spark/utils/ParserUtils.java | 110 ++++++++++++++++++ .../lance/spark/write/LanceBatchWrite.java | 28 +++-- .../spark/write/BaseReplaceWhereTest.java | 17 +++ 9 files changed, 187 insertions(+), 63 deletions(-) diff --git a/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 919767777..79d1a9be0 100644 --- a/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-3.4_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -84,14 +84,13 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) : ReplaceWhere = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) - // Recover the predicate and query as raw source text by character interval. The Lance grammar - // captures them as opaque token runs, so slicing the original stream is what preserves - // operators and spacing. The query is then parsed by Spark's own parser (delegate); the - // predicate is left as text and handed to Lance at execution time. - val predicate = originalText(ctx.predicate) - val queryText = originalText(ctx.query) - val query = delegate.parsePlan(queryText) - ReplaceWhere(table, predicate, query) + // Recover the raw text following WHERE by character interval (the grammar captures it as an + // opaque token run, so slicing the original stream preserves operators, quoting, and spacing), + // then split it at the first top-level AS into predicate and query. The query is parsed by + // Spark's own parser (delegate); the predicate is handed to Lance at execution time. + val parts = ParserUtils.splitReplaceBody(originalText(ctx.body)) + val query = delegate.parsePlan(parts(1)) + ReplaceWhere(table, parts(0), query) } private def originalText(ctx: ParserRuleContext): String = { diff --git a/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 1c4a64ca8..3443505b4 100644 --- a/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-3.5_2.12/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -84,14 +84,13 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) : ReplaceWhere = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) - // Recover the predicate and query as raw source text by character interval. The Lance grammar - // captures them as opaque token runs, so slicing the original stream is what preserves - // operators and spacing. The query is then parsed by Spark's own parser (delegate); the - // predicate is left as text and handed to Lance at execution time. - val predicate = originalText(ctx.predicate) - val queryText = originalText(ctx.query) - val query = delegate.parsePlan(queryText) - ReplaceWhere(table, predicate, query) + // Recover the raw text following WHERE by character interval (the grammar captures it as an + // opaque token run, so slicing the original stream preserves operators, quoting, and spacing), + // then split it at the first top-level AS into predicate and query. The query is parsed by + // Spark's own parser (delegate); the predicate is handed to Lance at execution time. + val parts = ParserUtils.splitReplaceBody(originalText(ctx.body)) + val query = delegate.parsePlan(parts(1)) + ReplaceWhere(table, parts(0), query) } private def originalText(ctx: ParserRuleContext): String = { diff --git a/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 1f3c12875..cfdd809c3 100644 --- a/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-4.0_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -84,14 +84,13 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) : ReplaceWhere = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) - // Recover the predicate and query as raw source text by character interval. The Lance grammar - // captures them as opaque token runs, so slicing the original stream is what preserves - // operators and spacing. The query is then parsed by Spark's own parser (delegate); the - // predicate is left as text and handed to Lance at execution time. - val predicate = originalText(ctx.predicate) - val queryText = originalText(ctx.query) - val query = delegate.parsePlan(queryText) - ReplaceWhere(table, predicate, query) + // Recover the raw text following WHERE by character interval (the grammar captures it as an + // opaque token run, so slicing the original stream preserves operators, quoting, and spacing), + // then split it at the first top-level AS into predicate and query. The query is parsed by + // Spark's own parser (delegate); the predicate is handed to Lance at execution time. + val parts = ParserUtils.splitReplaceBody(originalText(ctx.body)) + val query = delegate.parsePlan(parts(1)) + ReplaceWhere(table, parts(0), query) } private def originalText(ctx: ParserRuleContext): String = { diff --git a/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 1f3c12875..cfdd809c3 100644 --- a/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-4.1_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -84,14 +84,13 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) : ReplaceWhere = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) - // Recover the predicate and query as raw source text by character interval. The Lance grammar - // captures them as opaque token runs, so slicing the original stream is what preserves - // operators and spacing. The query is then parsed by Spark's own parser (delegate); the - // predicate is left as text and handed to Lance at execution time. - val predicate = originalText(ctx.predicate) - val queryText = originalText(ctx.query) - val query = delegate.parsePlan(queryText) - ReplaceWhere(table, predicate, query) + // Recover the raw text following WHERE by character interval (the grammar captures it as an + // opaque token run, so slicing the original stream preserves operators, quoting, and spacing), + // then split it at the first top-level AS into predicate and query. The query is parsed by + // Spark's own parser (delegate); the predicate is handed to Lance at execution time. + val parts = ParserUtils.splitReplaceBody(originalText(ctx.body)) + val query = delegate.parsePlan(parts(1)) + ReplaceWhere(table, parts(0), query) } private def originalText(ctx: ParserRuleContext): String = { diff --git a/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala b/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala index 1f3c12875..cfdd809c3 100644 --- a/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala +++ b/lance-spark-4.2_2.13/src/main/scala/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensionsAstBuilder.scala @@ -84,14 +84,13 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) : ReplaceWhere = { val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) - // Recover the predicate and query as raw source text by character interval. The Lance grammar - // captures them as opaque token runs, so slicing the original stream is what preserves - // operators and spacing. The query is then parsed by Spark's own parser (delegate); the - // predicate is left as text and handed to Lance at execution time. - val predicate = originalText(ctx.predicate) - val queryText = originalText(ctx.query) - val query = delegate.parsePlan(queryText) - ReplaceWhere(table, predicate, query) + // Recover the raw text following WHERE by character interval (the grammar captures it as an + // opaque token run, so slicing the original stream preserves operators, quoting, and spacing), + // then split it at the first top-level AS into predicate and query. The query is parsed by + // Spark's own parser (delegate); the predicate is handed to Lance at execution time. + val parts = ParserUtils.splitReplaceBody(originalText(ctx.body)) + val query = delegate.parsePlan(parts(1)) + ReplaceWhere(table, parts(0), query) } private def originalText(ctx: ParserRuleContext): String = { diff --git a/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 b/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 index d38fe05ec..d39c151bb 100644 --- a/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 +++ b/lance-spark-base_2.12/src/main/antlr4/org/apache/spark/sql/catalyst/parser/extensions/LanceSqlExtensions.g4 @@ -43,20 +43,16 @@ statement | OPTIMIZE multipartIdentifier (WITH '(' (namedArgument (',' namedArgument)*)? ')')? #optimize | VACUUM multipartIdentifier (WITH '(' (namedArgument (',' namedArgument)*)? ')')? #vacuum | ALTER TABLE multipartIdentifier SET UNENFORCED PRIMARY KEY '(' columnList ')' #setUnenforcedPrimaryKey - | REPLACE multipartIdentifier WHERE predicate=predicateText AS query=queryText #replaceWhere + | REPLACE multipartIdentifier WHERE body=replaceBody #replaceWhere ; -// The predicate (between WHERE and AS) and the query (after AS) are captured verbatim from the -// original SQL text and re-parsed with Spark's own parser, so the Lance grammar itself does not -// need to model SQL expressions or SELECT statements. The predicate runs up to the first AS -// keyword (its separator); the query is everything after AS, so it may itself contain AS (e.g. -// column aliases). Both are recovered by character interval in the AST builder, which preserves -// operators and whitespace regardless of how individual characters tokenized. -predicateText - : (~AS)+ - ; - -queryText +// Everything after WHERE (the predicate, the AS separator, and the query) is captured verbatim as +// one raw region and re-parsed with Spark's own parser, so the Lance grammar itself does not need +// to model SQL expressions or SELECT statements. The AST builder splits this text at the first +// top-level `AS` (honoring parentheses, string literals, and comments), which lets predicates +// contain `AS` themselves (e.g. `CAST(x AS STRING)`) and lets the query contain column aliases. +// Recovering by character interval preserves operators and whitespace regardless of tokenization. +replaceBody : .+ ; diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java index 4b90860ed..a0c4609ee 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java @@ -56,4 +56,114 @@ public static String quoteIdentifier(String identifier) { } return "`" + identifier + "`"; } + + /** + * Splits the raw body of a {@code REPLACE ... WHERE AS } command into its + * predicate and query at the first top-level {@code AS} keyword — one that is not nested + * inside parentheses, a string/backtick literal, or a comment. This lets the predicate itself + * contain {@code AS} (e.g. {@code CAST(dt AS STRING)}) and lets the query contain column aliases. + * + * @param body the source text following {@code WHERE}, e.g. {@code "dt = '1' AS SELECT ..."} + * @return a two-element array {@code [predicate, query]}, both trimmed + * @throws IllegalArgumentException if no top-level {@code AS} separator is found + */ + public static String[] splitReplaceBody(String body) { + int depth = 0; + int i = 0; + int n = body.length(); + while (i < n) { + char c = body.charAt(i); + if (c == '\'' || c == '"' || c == '`') { + i = skipQuoted(body, i, c); + continue; + } + if (c == '-' && i + 1 < n && body.charAt(i + 1) == '-') { + i = skipLineComment(body, i); + continue; + } + if (c == '/' && i + 1 < n && body.charAt(i + 1) == '*') { + i = skipBlockComment(body, i); + continue; + } + if (c == '(') { + depth++; + } else if (c == ')') { + depth--; + } else if (depth == 0 && isAsKeywordAt(body, i)) { + String predicate = body.substring(0, i).trim(); + String query = body.substring(i + 2).trim(); + if (predicate.isEmpty() || query.isEmpty()) { + throw new IllegalArgumentException( + "REPLACE ... WHERE requires a non-empty predicate and query around AS: " + body); + } + return new String[] {predicate, query}; + } + i++; + } + throw new IllegalArgumentException( + "REPLACE ... WHERE requires an AS separator between the predicate and query: " + body); + } + + /** Returns the index just past the closing quote for the literal starting at {@code start}. */ + private static int skipQuoted(String s, int start, char quote) { + int i = start + 1; + int n = s.length(); + while (i < n) { + char c = s.charAt(i); + if (c == '\\' && quote != '`') { + i += 2; // escaped char in a '...'/"..." literal + continue; + } + if (c == quote) { + // A doubled quote is an escaped quote, not a terminator. + if (i + 1 < n && s.charAt(i + 1) == quote) { + i += 2; + continue; + } + return i + 1; + } + i++; + } + return n; + } + + private static int skipLineComment(String s, int start) { + int i = start + 2; + int n = s.length(); + while (i < n && s.charAt(i) != '\n') { + i++; + } + return i; + } + + private static int skipBlockComment(String s, int start) { + int i = start + 2; + int n = s.length(); + while (i + 1 < n && !(s.charAt(i) == '*' && s.charAt(i + 1) == '/')) { + i++; + } + return Math.min(i + 2, n); + } + + /** + * Whether a standalone {@code AS} keyword (word-bounded, case-insensitive) begins at {@code i}. + */ + private static boolean isAsKeywordAt(String s, int i) { + int n = s.length(); + if (i + 2 > n) { + return false; + } + char a = s.charAt(i); + char b = s.charAt(i + 1); + if (!((a == 'a' || a == 'A') && (b == 's' || b == 'S'))) { + return false; + } + boolean leftBoundary = i == 0 || !isWordChar(s.charAt(i - 1)); + boolean rightBoundary = i + 2 == n || !isWordChar(s.charAt(i + 2)); + return leftBoundary && rightBoundary; + } + + private static boolean isWordChar(char c) { + return Character.isLetterOrDigit(c) || c == '_'; + } } diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java index 6c7f0448e..e26365d40 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java @@ -41,6 +41,7 @@ import org.apache.spark.sql.connector.write.WriterCommitMessage; import org.apache.spark.sql.types.StructType; import org.apache.spark.sql.util.LanceArrowUtils; +import org.roaringbitmap.RoaringBitmap; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -252,13 +253,17 @@ public void commit(WriterCommitMessage[] messages) { */ private static Operation buildReplaceOperation( Dataset ds, String predicate, List newFragments) { - Map> rowIndexesByFragment = matchingRowIndexesByFragment(ds, predicate); + Map deletionsByFragment = matchingDeletionsByFragment(ds, predicate); List removedFragmentIds = new ArrayList<>(); List updatedFragments = new ArrayList<>(); - for (Map.Entry> entry : rowIndexesByFragment.entrySet()) { + for (Map.Entry entry : deletionsByFragment.entrySet()) { int fragmentId = entry.getKey(); - FragmentMetadata updated = ds.getFragment(fragmentId).deleteRows(entry.getValue()); + // Materialize the row indexes for a single fragment at a time; the aggregate deletion state + // stays compressed as RoaringBitmaps so driver memory does not grow with total matched rows. + List rowIndexes = new ArrayList<>(entry.getValue().getCardinality()); + entry.getValue().forEach((org.roaringbitmap.IntConsumer) rowIndexes::add); + FragmentMetadata updated = ds.getFragment(fragmentId).deleteRows(rowIndexes); if (updated == null) { // All rows in the fragment matched the predicate; drop the whole fragment. removedFragmentIds.add((long) fragmentId); @@ -275,13 +280,14 @@ private static Operation buildReplaceOperation( } /** - * Scans the dataset for rows matching {@code predicate} and groups their physical row indexes by - * fragment id, decoding the 64-bit {@code _rowaddr} (fragment id in the high 32 bits, row index - * in the low 32 bits). Returns an empty map when no existing row matches. + * Scans the dataset for rows matching {@code predicate} and collects their physical row indexes + * per fragment as {@link RoaringBitmap}s, decoding the 64-bit {@code _rowaddr} (fragment id in + * the high 32 bits, row index in the low 32 bits). A compressed bitmap per fragment keeps driver + * memory bounded regardless of how many rows match. Returns an empty map when no row matches. */ - private static Map> matchingRowIndexesByFragment( + private static Map matchingDeletionsByFragment( Dataset ds, String predicate) { - Map> rowIndexesByFragment = new java.util.HashMap<>(); + Map deletionsByFragment = new java.util.HashMap<>(); ScanOptions scanOptions = new ScanOptions.Builder() .columns(java.util.Collections.emptyList()) @@ -295,15 +301,15 @@ private static Map> matchingRowIndexesByFragment( FieldVector rowAddrVector = batch.getVector(LanceConstant.ROW_ADDRESS); for (int i = 0; i < batch.getRowCount(); i++) { long rowAddress = ((Number) rowAddrVector.getObject(i)).longValue(); - rowIndexesByFragment - .computeIfAbsent(extractFragmentId(rowAddress), k -> new ArrayList<>()) + deletionsByFragment + .computeIfAbsent(extractFragmentId(rowAddress), k -> new RoaringBitmap()) .add(extractRowIndex(rowAddress)); } } } catch (Exception e) { throw new RuntimeException("Failed to scan rows for REPLACE ... WHERE " + predicate, e); } - return rowIndexesByFragment; + return deletionsByFragment; } @Override diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java index f56deb906..436aee446 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java @@ -153,6 +153,23 @@ public void testReplaceNonExistingPartitionAppends() { op.check(Arrays.asList(Row.of(1, "2026-08-01", 100), Row.of(5, "2026-08-09", 500))); } + /** + * A predicate may itself contain {@code AS} (e.g. inside a {@code CAST}); the command must split + * on the top-level {@code AS} separator, not the first {@code AS} token. + */ + @Test + public void testReplacePredicateWithCast() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + + op.insert(Arrays.asList(Row.of(1, "2026-08-01", 100), Row.of(2, "2026-08-02", 200))); + + op.replace( + "CAST(dt AS STRING) = '2026-08-01'", "SELECT 3 AS id, '2026-08-01' AS dt, 300 AS value"); + + op.check(Arrays.asList(Row.of(2, "2026-08-02", 200), Row.of(3, "2026-08-01", 300))); + } + /** The replacement is a single atomic commit: exactly one new table version is produced. */ @Test public void testReplaceIsSingleAtomicCommit() { From a79e7b6367684866e323a83ed4bdd1abe443d00d Mon Sep 17 00:00:00 2001 From: Pucheng Yang Date: Tue, 11 Aug 2026 00:50:30 +0000 Subject: [PATCH 3/6] fix(sql): depth-aware comment skipping and full-match fragment fast path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second round of lance-gatekeeper feedback on REPLACE ... WHERE: 1. Nested block comments: splitReplaceBody's block-comment skip stopped at the first `*/`, so an `AS` inside an outer comment (e.g. `/* outer /* inner */ AS ... */`) was mistaken for the separator. Comment skipping is now depth-aware, matching Spark's nested-comment behavior. Adds a nested-block-comment regression test. 2. Per-fragment deletion memory: a fully-matched fragment is now dropped without materializing any per-row list — matched cardinality is compared against the fragment's live row count and, on equality, the fragment id goes straight to removedFragmentIds. This covers the common whole-partition-per-fragment case with memory independent of fragment size. Only genuinely partial fragments materialize a list, bounded by a single fragment's row count (the native deleteRows(List) API's inherent limit). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../org/lance/spark/utils/ParserUtils.java | 17 +++++++++++--- .../lance/spark/write/LanceBatchWrite.java | 22 ++++++++++++++----- .../spark/write/BaseReplaceWhereTest.java | 18 +++++++++++++++ 3 files changed, 48 insertions(+), 9 deletions(-) diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java index a0c4609ee..2ea4cfd8d 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java @@ -137,12 +137,23 @@ private static int skipLineComment(String s, int start) { } private static int skipBlockComment(String s, int start) { + // Spark supports nested block comments, so track depth: an inner `*/` closes only the inner + // comment, and an `AS` remains commented out until the outermost comment closes. int i = start + 2; int n = s.length(); - while (i + 1 < n && !(s.charAt(i) == '*' && s.charAt(i + 1) == '/')) { - i++; + int depth = 1; + while (i + 1 < n && depth > 0) { + if (s.charAt(i) == '/' && s.charAt(i + 1) == '*') { + depth++; + i += 2; + } else if (s.charAt(i) == '*' && s.charAt(i + 1) == '/') { + depth--; + i += 2; + } else { + i++; + } } - return Math.min(i + 2, n); + return depth == 0 ? i : n; } /** diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java index e26365d40..b04fed452 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java @@ -15,6 +15,7 @@ import org.lance.CommitBuilder; import org.lance.Dataset; +import org.lance.Fragment; import org.lance.FragmentMetadata; import org.lance.Transaction; import org.lance.ipc.LanceScanner; @@ -259,13 +260,22 @@ private static Operation buildReplaceOperation( List updatedFragments = new ArrayList<>(); for (Map.Entry entry : deletionsByFragment.entrySet()) { int fragmentId = entry.getKey(); - // Materialize the row indexes for a single fragment at a time; the aggregate deletion state - // stays compressed as RoaringBitmaps so driver memory does not grow with total matched rows. - List rowIndexes = new ArrayList<>(entry.getValue().getCardinality()); - entry.getValue().forEach((org.roaringbitmap.IntConsumer) rowIndexes::add); - FragmentMetadata updated = ds.getFragment(fragmentId).deleteRows(rowIndexes); + Fragment fragment = ds.getFragment(fragmentId); + RoaringBitmap matched = entry.getValue(); + // Fast path: when every live row in the fragment matches, drop the fragment outright without + // materializing any per-row list. This covers the common partition-overwrite case (a whole + // partition living in its own fragment) and keeps driver memory independent of fragment size. + if (matched.getCardinality() == fragment.metadata().getNumRows()) { + removedFragmentIds.add((long) fragmentId); + continue; + } + // Partial match: the native deleteRows takes a List, so the surviving indexes for + // this one fragment are materialized here. This list is bounded by a single fragment's row + // count, not by the total matched row count across the table. + List rowIndexes = new ArrayList<>(matched.getCardinality()); + matched.forEach((org.roaringbitmap.IntConsumer) rowIndexes::add); + FragmentMetadata updated = fragment.deleteRows(rowIndexes); if (updated == null) { - // All rows in the fragment matched the predicate; drop the whole fragment. removedFragmentIds.add((long) fragmentId); } else { updatedFragments.add(updated); diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java index 436aee446..03c7888bd 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java @@ -170,6 +170,24 @@ public void testReplacePredicateWithCast() { op.check(Arrays.asList(Row.of(2, "2026-08-02", 200), Row.of(3, "2026-08-01", 300))); } + /** + * A nested block comment containing {@code AS} must not be treated as the command separator; the + * split honors nested {@code /* ... */} the way Spark's parser does. + */ + @Test + public void testReplacePredicateWithNestedBlockComment() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + + op.insert(Arrays.asList(Row.of(1, "2026-08-01", 100), Row.of(2, "2026-08-02", 200))); + + op.replace( + "dt = '2026-08-01' /* outer /* inner */ AS still-comment */", + "SELECT 3 AS id, '2026-08-01' AS dt, 300 AS value"); + + op.check(Arrays.asList(Row.of(2, "2026-08-02", 200), Row.of(3, "2026-08-01", 300))); + } + /** The replacement is a single atomic commit: exactly one new table version is produced. */ @Test public void testReplaceIsSingleAtomicCommit() { From 4ce419bc2857c04622a89b68a5e0989d54ec0841 Mon Sep 17 00:00:00 2001 From: Pucheng Yang Date: Tue, 11 Aug 2026 03:32:13 +0000 Subject: [PATCH 4/6] fix(sql): end predicate line comments at carriage return Third round of lance-gatekeeper feedback on REPLACE ... WHERE: - splitReplaceBody's line-comment skip only stopped at \n, but Spark ends a line comment at \r too, so a predicate ending in `-- ...\r` swallowed the AS separator. Line comments now end at either \r or \n. Adds a regression test. - Use RoaringBitmap.getLongCardinality() when comparing against a fragment's live row count in the full-match fast path, avoiding int truncation for very large fragments. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../java/org/lance/spark/utils/ParserUtils.java | 3 ++- .../org/lance/spark/write/LanceBatchWrite.java | 2 +- .../lance/spark/write/BaseReplaceWhereTest.java | 14 ++++++++++++++ 3 files changed, 17 insertions(+), 2 deletions(-) diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java index 2ea4cfd8d..72c50e05f 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java @@ -128,9 +128,10 @@ private static int skipQuoted(String s, int start, char quote) { } private static int skipLineComment(String s, int start) { + // A line comment ends at the next line terminator; Spark treats both \n and \r as terminators. int i = start + 2; int n = s.length(); - while (i < n && s.charAt(i) != '\n') { + while (i < n && s.charAt(i) != '\n' && s.charAt(i) != '\r') { i++; } return i; diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java index b04fed452..a282215af 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java @@ -265,7 +265,7 @@ private static Operation buildReplaceOperation( // Fast path: when every live row in the fragment matches, drop the fragment outright without // materializing any per-row list. This covers the common partition-overwrite case (a whole // partition living in its own fragment) and keeps driver memory independent of fragment size. - if (matched.getCardinality() == fragment.metadata().getNumRows()) { + if (matched.getLongCardinality() == fragment.metadata().getNumRows()) { removedFragmentIds.add((long) fragmentId); continue; } diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java index 03c7888bd..c70a62fd1 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java @@ -188,6 +188,20 @@ public void testReplacePredicateWithNestedBlockComment() { op.check(Arrays.asList(Row.of(2, "2026-08-02", 200), Row.of(3, "2026-08-01", 300))); } + /** A line comment in the predicate ends at a carriage return, matching Spark's parser. */ + @Test + public void testReplacePredicateWithCarriageReturnLineComment() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + + op.insert(Arrays.asList(Row.of(1, "2026-08-01", 100), Row.of(2, "2026-08-02", 200))); + + op.replace( + "dt = '2026-08-01' -- AS ignored\r", "SELECT 3 AS id, '2026-08-01' AS dt, 300 AS value"); + + op.check(Arrays.asList(Row.of(2, "2026-08-02", 200), Row.of(3, "2026-08-01", 300))); + } + /** The replacement is a single atomic commit: exactly one new table version is produced. */ @Test public void testReplaceIsSingleAtomicCommit() { From 473e6ba45cfd69e17bb3cba4f2f3f1d61582843b Mon Sep 17 00:00:00 2001 From: Pucheng Yang Date: Fri, 14 Aug 2026 23:31:47 +0000 Subject: [PATCH 5/6] test(sql): cover REPLACE on a fragment with pre-existing deletions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a regression test verifying REPLACE ... WHERE is correct when the target fragment already carries a deletion vector: the full-fragment drop compares matched live rows against getNumRows() (physical − deletions) and a filtered scan enumerates only live rows, so matching all remaining live rows drops the fragment cleanly rather than mis-counting against the physical total. Verified against lance-core semantics (getNumRows is logical; deleteRows takes physical offsets = _rowaddr low 32 bits; filtered scans exclude deleted rows). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../spark/write/BaseReplaceWhereTest.java | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java index c70a62fd1..c42e16412 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java @@ -153,6 +153,35 @@ public void testReplaceNonExistingPartitionAppends() { op.check(Arrays.asList(Row.of(1, "2026-08-01", 100), Row.of(5, "2026-08-09", 500))); } + /** + * REPLACE must be correct on a fragment that already carries a deletion vector. The full-fragment + * drop compares matched live rows against the fragment's live row count ({@code getNumRows()} = + * physical − deletions), and a filtered scan enumerates only live rows, so a REPLACE that matches + * all remaining live rows of a partially-deleted fragment drops it cleanly rather than + * mis-counting against the physical total. + */ + @Test + public void testReplaceFragmentWithPreExistingDeletions() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + + // Single fragment for dt=2026-08-01 with three rows, plus another partition. + op.insert( + Arrays.asList( + Row.of(1, "2026-08-01", 100), + Row.of(2, "2026-08-01", 200), + Row.of(3, "2026-08-01", 300))); + op.insert(Arrays.asList(Row.of(9, "2026-08-02", 900))); + + // Delete one row from the first fragment, leaving a pre-existing deletion vector. + op.delete("id = 2"); + + // REPLACE now matches all *remaining live* rows of that fragment (ids 1 and 3). + op.replace("dt = '2026-08-01'", "SELECT 5 AS id, '2026-08-01' AS dt, 500 AS value"); + + op.check(Arrays.asList(Row.of(5, "2026-08-01", 500), Row.of(9, "2026-08-02", 900))); + } + /** * A predicate may itself contain {@code AS} (e.g. inside a {@code CAST}); the command must split * on the top-level {@code AS} separator, not the first {@code AS} token. @@ -250,6 +279,10 @@ void replace(String predicate, String query) { spark.sql(String.format("REPLACE %s WHERE %s AS %s", fullName(), predicate, query)); } + void delete(String predicate) { + spark.sql(String.format("DELETE FROM %s WHERE %s", fullName(), predicate)); + } + long latestVersion() { // Resolve the table's dataset URI through the catalog, then open it to read the current // manifest version. This mirrors how other connector tests read a table's version. From 497011c938895b8e47e16bc35e87095e141fe389 Mon Sep 17 00:00:00 2001 From: Pucheng Yang Date: Sat, 15 Aug 2026 00:05:17 +0000 Subject: [PATCH 6/6] feat(sql): metadata-only fragment removal for REPLACE on zonemapped columns MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For a REPLACE ... WHERE whose predicate is a pure conjunction of `column = literal` equalities on columns that have a zonemap index, fragments that the zonemap proves are fully covered (every zone pinned to the required value, no nulls) are now dropped by id without scanning their rows. The remaining (unproven) fragments are scanned exactly as before, restricted to their ids. When every matching fragment is proven covered, the delete-planning scan is skipped entirely — replacing a partition that occupies its own fragments becomes metadata-only (O(#fragments) instead of O(rows)). Mechanics: - ReplaceWhereExec parses the predicate with Spark's parser and, only when it is a conjunction of equalities on string/integral columns, encodes the terms as JSON in a new internal write option (else the option is omitted). - ReplaceCoverage proves per-fragment coverage from getZonemapStats: a fragment qualifies only if, for every equality column, all its zones have min == max == value with zero nulls. Multiple single-column zonemaps combine by intersection for multi-column predicates. - Correctness is fail-safe: any missing zonemap, unproven zone, non-equality predicate, or value/format mismatch excludes the fragment and defers to the exact scan, so the optimization never changes which rows are replaced. Tests: zonemap-covered single-partition drop, multi-column zonemap, and a range-predicate fallback, each asserting results identical to the scan path. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/src/operations/dml/replace.md | 7 + .../java/org/lance/spark/LanceConstant.java | 9 + .../lance/spark/LanceSparkWriteOptions.java | 35 +++- .../lance/spark/write/LanceBatchWrite.java | 42 ++++- .../lance/spark/write/ReplaceCoverage.java | 154 ++++++++++++++++++ .../datasources/v2/ReplaceWhereExec.scala | 87 +++++++++- .../spark/write/BaseReplaceWhereTest.java | 74 +++++++++ 7 files changed, 397 insertions(+), 11 deletions(-) create mode 100644 lance-spark-base_2.12/src/main/java/org/lance/spark/write/ReplaceCoverage.java diff --git a/docs/src/operations/dml/replace.md b/docs/src/operations/dml/replace.md index 53a28c006..51a4cc84e 100644 --- a/docs/src/operations/dml/replace.md +++ b/docs/src/operations/dml/replace.md @@ -51,3 +51,10 @@ REPLACE
WHERE AS - When a fragment contains both matching and non-matching rows, only the matching rows are removed (via a deletion vector); the rest are preserved. A fragment whose rows all match is dropped outright. +- **Performance:** for an equality predicate (a conjunction of `column = value`) on columns that + have a zonemap index, `REPLACE` uses the zonemap to drop fully-covered fragments by id without + scanning their rows — so replacing a partition that occupies its own fragments is metadata-only. + A predicate with no zonemap, or any non-equality shape (ranges, `OR`, functions), still works but + falls back to scanning the affected fragments to find the rows to delete. Multiple single-column + zonemaps combine for a multi-column predicate (e.g. a zonemap on `dt` and one on `hr` for + `WHERE dt = ... AND hr = ...`). diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceConstant.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceConstant.java index 7e65e2a4e..ec05cfc76 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceConstant.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/LanceConstant.java @@ -43,6 +43,15 @@ public class LanceConstant { */ public static final String REPLACE_WHERE_KEY = "__lance_replace_where"; + /** + * Internal write option carrying the JSON-encoded equality terms of a {@code REPLACE ... WHERE} + * predicate (a pure conjunction of {@code column = literal}), when the predicate has that shape. + * The commit uses these terms plus zonemap statistics to drop fully-covered fragments by id + * without scanning their rows; when absent, commit falls back to the exact scan-based deletion. + * Set on the driver by {@code ReplaceWhereExec}; not a user-facing option. + */ + public static final String REPLACE_WHERE_EQUALITY_KEY = "__lance_replace_where_equality"; + /** * Internal write option carrying the encoded blob source credential/open contexts for an INSERT * whose query reads blob columns. Set on the driver by {@code LanceBlobSourceContextRule} and 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 d31299395..87baf0411 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 @@ -107,6 +107,13 @@ public class LanceSparkWriteOptions implements Serializable { */ private final String replaceWhere; + /** + * JSON-encoded equality terms of a {@code REPLACE ... WHERE} predicate, or null. When present, + * the commit may drop fully-covered fragments via zonemap statistics without scanning them. Like + * {@link #replaceWhere}, excluded from {@link #toWriteParams}. + */ + private final String replaceWhereEqualities; + private LanceSparkWriteOptions(Builder builder) { this.datasetUri = builder.datasetUri; this.writeMode = builder.writeMode; @@ -126,6 +133,7 @@ private LanceSparkWriteOptions(Builder builder) { this.tableId = builder.tableId; this.version = builder.version; this.replaceWhere = builder.replaceWhere; + this.replaceWhereEqualities = builder.replaceWhereEqualities; } /** Creates a new builder for LanceSparkWriteOptions. */ @@ -231,6 +239,13 @@ public String getReplaceWhere() { return replaceWhere; } + /** + * Returns the JSON-encoded equality terms of the REPLACE predicate, or null if not applicable. + */ + public String getReplaceWhereEqualities() { + return replaceWhereEqualities; + } + /** Returns a builder pre-populated with all fields from this instance. */ public Builder toBuilder() { return builder() @@ -251,7 +266,8 @@ public Builder toBuilder() { .namespace(namespace) .tableId(tableId) .version(version) - .replaceWhere(replaceWhere); + .replaceWhere(replaceWhere) + .replaceWhereEqualities(replaceWhereEqualities); } /** Returns a copy of these options with version set to the given version. */ @@ -344,7 +360,8 @@ public boolean equals(Object o) { && Objects.equals(storageOptions, that.storageOptions) && Objects.equals(tableId, that.tableId) && Objects.equals(version, that.version) - && Objects.equals(replaceWhere, that.replaceWhere); + && Objects.equals(replaceWhere, that.replaceWhere) + && Objects.equals(replaceWhereEqualities, that.replaceWhereEqualities); } @Override @@ -366,7 +383,8 @@ public int hashCode() { storageOptions, tableId, version, - replaceWhere); + replaceWhere, + replaceWhereEqualities); } /** Builder for creating LanceSparkWriteOptions instances. */ @@ -389,6 +407,7 @@ public static class Builder { private List tableId; private Long version; private String replaceWhere; + private String replaceWhereEqualities; private Builder() {} @@ -488,6 +507,12 @@ public Builder replaceWhere(String replaceWhere) { return this; } + /** Sets the JSON-encoded equality terms of the REPLACE predicate; null when not applicable. */ + public Builder replaceWhereEqualities(String replaceWhereEqualities) { + this.replaceWhereEqualities = replaceWhereEqualities; + return this; + } + /** * Parses options from a map, extracting write-specific settings. * @@ -502,6 +527,10 @@ public Builder fromOptions(Map options) { this.replaceWhere = options.get(LanceConstant.REPLACE_WHERE_KEY); this.storageOptions.remove(LanceConstant.REPLACE_WHERE_KEY); } + if (options.containsKey(LanceConstant.REPLACE_WHERE_EQUALITY_KEY)) { + this.replaceWhereEqualities = options.get(LanceConstant.REPLACE_WHERE_EQUALITY_KEY); + this.storageOptions.remove(LanceConstant.REPLACE_WHERE_EQUALITY_KEY); + } if (options.containsKey(CONFIG_WRITE_MODE)) { this.writeMode = WriteMode.valueOf(options.get(CONFIG_WRITE_MODE).toUpperCase()); } diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java index a282215af..bae2b9f40 100644 --- a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/LanceBatchWrite.java @@ -51,6 +51,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.stream.Collectors; import static org.lance.spark.join.FragmentAwareJoinUtils.extractFragmentId; @@ -206,7 +207,12 @@ public void commit(WriterCommitMessage[] messages) { try (Dataset ds = Utils.openDatasetBuilder(writeOptions).build()) { Operation operation; if (writeOptions.getReplaceWhere() != null) { - operation = buildReplaceOperation(ds, writeOptions.getReplaceWhere(), fragments); + operation = + buildReplaceOperation( + ds, + writeOptions.getReplaceWhere(), + writeOptions.getReplaceWhereEqualities(), + fragments); } else if (isOverwrite) { operation = Overwrite.builder().fragments(fragments).schema(arrowSchema).build(); } else { @@ -253,11 +259,36 @@ public void commit(WriterCommitMessage[] messages) { * the predicate is removed outright. */ private static Operation buildReplaceOperation( - Dataset ds, String predicate, List newFragments) { - Map deletionsByFragment = matchingDeletionsByFragment(ds, predicate); - + Dataset ds, String predicate, String equalitiesJson, List newFragments) { List removedFragmentIds = new ArrayList<>(); List updatedFragments = new ArrayList<>(); + + // Metadata-only fast path: for an equality predicate, use zonemap statistics to identify + // fragments whose every live row provably matches (each indexed equality column has all zones + // pinned to the required value), and drop them by id without reading any rows. Fragments that + // cannot be proven fully covered are handled by the exact scan below, so this only ever avoids + // work — it never changes which rows are replaced. + Set fullyCoveredFragmentIds = + ReplaceCoverage.fullyCoveredFragmentIds(ds, equalitiesJson); + for (int fragmentId : fullyCoveredFragmentIds) { + removedFragmentIds.add((long) fragmentId); + } + + // Scan only the fragments not already dropped by metadata, restricting the scan to their ids. + // When every fragment is proven covered, this list is empty and the scan is skipped entirely — + // the whole-partition case reads no data rows at all (O(#fragments), not O(rows)). + List scanFragmentIds = new ArrayList<>(); + for (Fragment fragment : ds.getFragments()) { + if (!fullyCoveredFragmentIds.contains(fragment.getId())) { + scanFragmentIds.add(fragment.getId()); + } + } + + Map deletionsByFragment = + scanFragmentIds.isEmpty() + ? java.util.Collections.emptyMap() + : matchingDeletionsByFragment(ds, predicate, scanFragmentIds); + for (Map.Entry entry : deletionsByFragment.entrySet()) { int fragmentId = entry.getKey(); Fragment fragment = ds.getFragment(fragmentId); @@ -296,13 +327,14 @@ private static Operation buildReplaceOperation( * memory bounded regardless of how many rows match. Returns an empty map when no row matches. */ private static Map matchingDeletionsByFragment( - Dataset ds, String predicate) { + Dataset ds, String predicate, List scanFragmentIds) { Map deletionsByFragment = new java.util.HashMap<>(); ScanOptions scanOptions = new ScanOptions.Builder() .columns(java.util.Collections.emptyList()) .withRowAddress(true) .filter(predicate) + .fragmentIds(scanFragmentIds) .build(); try (LanceScanner scanner = ds.newScan(scanOptions); ArrowReader reader = scanner.scanBatches()) { diff --git a/lance-spark-base_2.12/src/main/java/org/lance/spark/write/ReplaceCoverage.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/ReplaceCoverage.java new file mode 100644 index 000000000..b5953eb36 --- /dev/null +++ b/lance-spark-base_2.12/src/main/java/org/lance/spark/write/ReplaceCoverage.java @@ -0,0 +1,154 @@ +/* + * 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.Dataset; +import org.lance.index.scalar.ZoneStats; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; + +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Set; + +/** + * Identifies, from zonemap statistics alone, which fragments are provably covered in full + * by a {@code REPLACE ... WHERE} equality predicate — so they can be dropped by id without scanning + * their rows. + * + *

The predicate must be a pure conjunction of {@code column = literal} terms (encoded as JSON by + * {@code ReplaceWhereExec}). A fragment is considered fully covered only when, for every + * equality column, the fragment has at least one zonemap zone and all of its zones are + * pinned to the required value (zone {@code min == max == value}) with no nulls. This is + * deliberately conservative: any column without a zonemap, any zone whose bounds are not pinned to + * the value, or any value/format mismatch simply excludes the fragment, which then falls back to + * the exact scan-based deletion. The method therefore only ever avoids work; it never changes which + * rows are replaced. + */ +final class ReplaceCoverage { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private ReplaceCoverage() {} + + /** + * Returns the ids of fragments provably covered in full by the equality predicate, or an empty + * set when the fast path does not apply (no equality terms, missing zonemaps, or nothing proven). + * + * @param ds the dataset being replaced into + * @param equalitiesJson JSON array of {@code {"column":..,"value":..}} terms, or null + */ + static Set fullyCoveredFragmentIds(Dataset ds, String equalitiesJson) { + if (equalitiesJson == null || equalitiesJson.isEmpty()) { + return Collections.emptySet(); + } + + List equalities = parseEqualities(equalitiesJson); + if (equalities.isEmpty()) { + return Collections.emptySet(); + } + + Set covered = null; + for (String[] equality : equalities) { + String column = equality[0]; + String value = equality[1]; + Set pinned = fragmentsPinnedToValue(ds, column, value); + if (pinned.isEmpty()) { + // No fragment can be proven covered for this column (e.g. no zonemap on it), so the + // conjunction cannot cover any fragment. + return Collections.emptySet(); + } + if (covered == null) { + covered = pinned; + } else { + covered.retainAll(pinned); + } + if (covered.isEmpty()) { + return Collections.emptySet(); + } + } + return covered == null ? Collections.emptySet() : covered; + } + + /** + * Returns the ids of fragments whose zonemap on {@code column} proves every live row equals + * {@code value}: the fragment has at least one zone, and all of its zones have {@code min == max + * == value} (by canonical string form) with zero nulls. Returns an empty set if the column has no + * zonemap index. + */ + private static Set fragmentsPinnedToValue(Dataset ds, String column, String value) { + List zones = ds.getZonemapStats(column); + if (zones == null || zones.isEmpty()) { + return Collections.emptySet(); + } + + // Group zones per fragment and track, per fragment, whether every zone is pinned to the value. + Set candidate = new HashSet<>(); + Set disqualified = new HashSet<>(); + for (ZoneStats zone : zones) { + int fragmentId = zone.getFragmentId(); + if (disqualified.contains(fragmentId)) { + continue; + } + if (isZonePinnedToValue(zone, value)) { + candidate.add(fragmentId); + } else { + candidate.remove(fragmentId); + disqualified.add(fragmentId); + } + } + candidate.removeAll(disqualified); + return candidate; + } + + /** + * A zone is pinned to {@code value} when it contains only that value: min == max == value, no + * nulls. + */ + private static boolean isZonePinnedToValue(ZoneStats zone, String value) { + if (zone.getNullCount() != 0) { + return false; + } + Comparable min = zone.getMin(); + Comparable max = zone.getMax(); + if (min == null || max == null) { + return false; + } + // Compare by canonical string form: zonemap min/max box as Long/Double/String via JNI, and the + // required value is the literal's toString() from ReplaceWhereExec. A mismatch (including any + // formatting difference) conservatively fails the proof and defers to the exact scan. + return value.equals(min.toString()) && value.equals(max.toString()); + } + + private static List parseEqualities(String json) { + try { + JsonNode array = MAPPER.readTree(json); + java.util.List result = new java.util.ArrayList<>(); + for (JsonNode node : array) { + JsonNode column = node.get("column"); + JsonNode value = node.get("value"); + if (column == null || value == null) { + return Collections.emptyList(); + } + result.add(new String[] {column.asText(), value.asText()}); + } + return result; + } catch (Exception e) { + // Malformed encoding: skip the fast path entirely and let the exact scan handle the delete. + return Collections.emptyList(); + } + } +} diff --git a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ReplaceWhereExec.scala b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ReplaceWhereExec.scala index f932942af..789da6692 100644 --- a/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ReplaceWhereExec.scala +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ReplaceWhereExec.scala @@ -14,11 +14,15 @@ package org.apache.spark.sql.execution.datasources.v2 import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.catalyst.analysis.UnresolvedAttribute +import org.apache.spark.sql.catalyst.expressions.{And, Attribute, EqualTo, Expression, Literal} import org.apache.spark.sql.catalyst.plans.logical.{AppendData, LogicalPlan} import org.apache.spark.sql.connector.catalog._ +import org.apache.spark.sql.types.{ByteType, IntegerType, LongType, ShortType, StringType} import org.lance.spark.{LanceConstant, LanceDataset} +import scala.collection.mutable.ArrayBuffer + /** * Physical plan for `REPLACE

WHERE AS `. * @@ -60,11 +64,88 @@ case class ReplaceWhereExec( Some(catalog), Some(ident)) - val append = - AppendData.byPosition(relation, query, Map(LanceConstant.REPLACE_WHERE_KEY -> predicate)) + val options = Map(LanceConstant.REPLACE_WHERE_KEY -> predicate) ++ + equalityTermsJson(predicate).map(LanceConstant.REPLACE_WHERE_EQUALITY_KEY -> _) + + val append = AppendData.byPosition(relation, query, options) val qe = session.sessionState.executePlan(append) qe.assertCommandExecuted() Nil } + + /** + * If the predicate is a pure conjunction of `column = literal` equality terms on string or + * integral columns, returns their JSON encoding for the metadata-only fragment-drop fast path at + * commit time. Returns `None` for any other predicate shape (ranges, OR, functions, other types), + * in which case commit falls back to the exact scan-based deletion — so this only ever enables an + * optimization, never changes which rows are replaced. + */ + private def equalityTermsJson(predicate: String): Option[String] = { + val parsed = + try { + session.sessionState.sqlParser.parseExpression(predicate) + } catch { + case _: Throwable => return None + } + + val terms = ArrayBuffer.empty[(String, String)] + if (!collectEqualities(parsed, terms)) { + return None + } + // A column appearing twice with different required values can never match; let the scan path + // handle that (it will simply find no rows). Only emit when each column maps to one value. + val byColumn = terms.groupBy(_._1) + if (terms.isEmpty || byColumn.exists(_._2.map(_._2).distinct.size > 1)) { + return None + } + val json = + byColumn + .map { case (col, pairs) => (col, pairs.head._2) } + .map { case (col, value) => s"""{"column":${quote(col)},"value":${quote(value)}}""" } + .mkString("[", ",", "]") + Some(json) + } + + /** + * Walks a conjunction, collecting `column = literal` pairs into `out`. Returns false (disabling + * the fast path) as soon as any node is not an AND or a supported equality on a simple column + * reference and string/integral literal. + */ + private def collectEqualities(expr: Expression, out: ArrayBuffer[(String, String)]): Boolean = + expr match { + case And(left, right) => collectEqualities(left, out) && collectEqualities(right, out) + case EqualTo(col, lit: Literal) if columnName(col).isDefined && lit.value != null => + supportedLiteral(lit) match { + case Some(value) => + out += ((columnName(col).get, value)) + true + case None => false + } + case EqualTo(lit: Literal, col) if columnName(col).isDefined && lit.value != null => + supportedLiteral(lit) match { + case Some(value) => + out += ((columnName(col).get, value)) + true + case None => false + } + case _ => false + } + + private def columnName(expr: Expression): Option[String] = expr match { + case u: UnresolvedAttribute if u.nameParts.size == 1 => Some(u.nameParts.head) + case _ => None + } + + /** Canonical string form for a literal the zonemap comparison can reproduce, else None. */ + private def supportedLiteral(lit: Literal): Option[String] = lit.dataType match { + case StringType => Some(lit.value.toString) + case ByteType | ShortType | IntegerType | LongType => Some(lit.value.toString) + case _ => None + } + + private def quote(s: String): String = { + val escaped = s.replace("\\", "\\\\").replace("\"", "\\\"") + "\"" + escaped + "\"" + } } diff --git a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java index c42e16412..fc1845a39 100644 --- a/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java @@ -231,6 +231,74 @@ public void testReplacePredicateWithCarriageReturnLineComment() { op.check(Arrays.asList(Row.of(2, "2026-08-02", 200), Row.of(3, "2026-08-01", 300))); } + /** + * With a zonemap index on the predicate column, a partition that occupies its own fragment is + * dropped via the metadata-only fast path (no row scan). This asserts the result is identical to + * the scan-based path — the optimization must not change which rows are replaced. + */ + @Test + public void testReplaceWithZonemapCoveredPartition() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + op.createZonemap("dt"); + + // One INSERT per dt → one fragment per dt, so dt=2026-08-01's fragment is fully covered. + op.insert(Arrays.asList(Row.of(1, "2026-08-01", 100), Row.of(2, "2026-08-01", 200))); + op.insert(Arrays.asList(Row.of(3, "2026-08-02", 300))); + + op.replace("dt = '2026-08-01'", "SELECT 10 AS id, '2026-08-01' AS dt, 999 AS value"); + + op.check(Arrays.asList(Row.of(3, "2026-08-02", 300), Row.of(10, "2026-08-01", 999))); + } + + /** + * Multi-column equality with zonemaps on both columns: only fragments pinned to BOTH values are + * dropped by metadata, and the result matches the exact semantics. + */ + @Test + public void testReplaceWithZonemapMultiColumn() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + op.createZonemap("dt"); + op.createZonemap("value"); + + // Each INSERT is its own fragment; only the first is (dt=2026-08-01, value=100). + op.insert(Arrays.asList(Row.of(1, "2026-08-01", 100))); + op.insert(Arrays.asList(Row.of(2, "2026-08-01", 200))); + op.insert(Arrays.asList(Row.of(3, "2026-08-02", 100))); + + op.replace( + "dt = '2026-08-01' AND value = 100", "SELECT 9 AS id, '2026-08-01' AS dt, 100 AS value"); + + op.check( + Arrays.asList( + Row.of(2, "2026-08-01", 200), + Row.of(3, "2026-08-02", 100), + Row.of(9, "2026-08-01", 100))); + } + + /** + * A non-equality predicate (range) is not eligible for the metadata fast path and must fall back + * to the exact scan, still producing the correct result even with a zonemap present. + */ + @Test + public void testReplaceRangePredicateFallsBack() { + TableOperator op = new TableOperator(spark, catalogName); + op.create(); + op.createZonemap("value"); + + op.insert(Arrays.asList(Row.of(1, "2026-08-01", 100), Row.of(2, "2026-08-01", 200))); + op.insert(Arrays.asList(Row.of(3, "2026-08-02", 300))); + + op.replace("value >= 300", "SELECT 7 AS id, '2026-08-02' AS dt, 700 AS value"); + + op.check( + Arrays.asList( + Row.of(1, "2026-08-01", 100), + Row.of(2, "2026-08-01", 200), + Row.of(7, "2026-08-02", 700))); + } + /** The replacement is a single atomic commit: exactly one new table version is produced. */ @Test public void testReplaceIsSingleAtomicCommit() { @@ -283,6 +351,12 @@ void delete(String predicate) { spark.sql(String.format("DELETE FROM %s WHERE %s", fullName(), predicate)); } + void createZonemap(String column) { + spark.sql( + String.format( + "ALTER TABLE %s CREATE INDEX %s_zm USING zonemap (%s)", fullName(), column, column)); + } + long latestVersion() { // Resolve the table's dataset URI through the catalog, then open it to read the current // manifest version. This mirrors how other connector tests read a table's version.