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..51a4cc84e --- /dev/null +++ b/docs/src/operations/dml/replace.md @@ -0,0 +1,60 @@ +# 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. +- **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-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..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 @@ -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,23 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) Optimize(table, args) } + override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) + : ReplaceWhere = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + // 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 = { + 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..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 @@ -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,23 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) Optimize(table, args) } + override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) + : ReplaceWhere = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + // 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 = { + 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..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 @@ -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,23 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) Optimize(table, args) } + override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) + : ReplaceWhere = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + // 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 = { + 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..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 @@ -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,23 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) Optimize(table, args) } + override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) + : ReplaceWhere = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + // 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 = { + 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..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 @@ -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,23 @@ class LanceSqlExtensionsAstBuilder(delegate: ParserInterface) Optimize(table, args) } + override def visitReplaceWhere(ctx: LanceSqlExtensionsParser.ReplaceWhereContext) + : ReplaceWhere = { + val table = UnresolvedIdentifier(visitMultipartIdentifier(ctx.multipartIdentifier())) + // 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 = { + 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..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,6 +43,17 @@ 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 body=replaceBody #replaceWhere + ; + +// 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 + : .+ ; multipartIdentifier @@ -113,6 +124,7 @@ NOT: 'NOT'; OF: 'OF'; OPTIMIZE: 'OPTIMIZE'; PRIMARY: 'PRIMARY'; +REPLACE: 'REPLACE'; SET: 'SET'; SHOW: 'SHOW'; TABLE: 'TABLE'; @@ -123,6 +135,7 @@ UPDATE: 'UPDATE'; USING: 'USING'; VACUUM: 'VACUUM'; VERSION: 'VERSION'; +WHERE: 'WHERE'; WITH: 'WITH'; TRUE: 'TRUE'; @@ -173,3 +186,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..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 @@ -35,6 +35,23 @@ 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 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 7bab7bef5..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 @@ -99,6 +99,21 @@ 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; + + /** + * 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; @@ -117,6 +132,8 @@ private LanceSparkWriteOptions(Builder builder) { this.namespace = builder.namespace; this.tableId = builder.tableId; this.version = builder.version; + this.replaceWhere = builder.replaceWhere; + this.replaceWhereEqualities = builder.replaceWhereEqualities; } /** Creates a new builder for LanceSparkWriteOptions. */ @@ -217,6 +234,18 @@ public Long getVersion() { return version; } + /** Returns the {@code REPLACE ... WHERE ...} row filter, or null for an ordinary write. */ + 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() @@ -236,7 +265,9 @@ public Builder toBuilder() { .storageOptions(storageOptions) .namespace(namespace) .tableId(tableId) - .version(version); + .version(version) + .replaceWhere(replaceWhere) + .replaceWhereEqualities(replaceWhereEqualities); } /** Returns a copy of these options with version set to the given version. */ @@ -328,7 +359,9 @@ public boolean equals(Object o) { && Objects.equals(blobPackFileSizeThreshold, that.blobPackFileSizeThreshold) && Objects.equals(storageOptions, that.storageOptions) && Objects.equals(tableId, that.tableId) - && Objects.equals(version, that.version); + && Objects.equals(version, that.version) + && Objects.equals(replaceWhere, that.replaceWhere) + && Objects.equals(replaceWhereEqualities, that.replaceWhereEqualities); } @Override @@ -349,7 +382,9 @@ public int hashCode() { blobPackFileSizeThreshold, storageOptions, tableId, - version); + version, + replaceWhere, + replaceWhereEqualities); } /** Builder for creating LanceSparkWriteOptions instances. */ @@ -371,6 +406,8 @@ public static class Builder { private LanceNamespace namespace; private List tableId; private Long version; + private String replaceWhere; + private String replaceWhereEqualities; private Builder() {} @@ -464,6 +501,18 @@ 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; + } + + /** 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. * @@ -472,6 +521,16 @@ 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(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/utils/ParserUtils.java b/lance-spark-base_2.12/src/main/java/org/lance/spark/utils/ParserUtils.java index 4b90860ed..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 @@ -56,4 +56,126 @@ 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) { + // 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' && s.charAt(i) != '\r') { + i++; + } + return i; + } + + 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(); + 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 depth == 0 ? i : 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 5cba05f2e..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 @@ -15,18 +15,26 @@ 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; +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; @@ -34,15 +42,21 @@ 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; +import java.util.ArrayList; import java.util.Arrays; 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; +import static org.lance.spark.join.FragmentAwareJoinUtils.extractRowIndex; + public class LanceBatchWrite implements BatchWrite { private static final Logger logger = LoggerFactory.getLogger(LanceBatchWrite.class); @@ -192,7 +206,14 @@ 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(), + writeOptions.getReplaceWhereEqualities(), + fragments); + } else if (isOverwrite) { operation = Overwrite.builder().fragments(fragments).schema(arrowSchema).build(); } else { operation = Append.builder().fragments(fragments).build(); @@ -226,6 +247,113 @@ 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, 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); + 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.getLongCardinality() == 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) { + 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 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 matchingDeletionsByFragment( + 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()) { + 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(); + 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 deletionsByFragment; + } + @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/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/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..789da6692 --- /dev/null +++ b/lance-spark-base_2.12/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ReplaceWhereExec.scala @@ -0,0 +1,151 @@ +/* + * 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.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 `. + * + * 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 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 new file mode 100644 index 000000000..fc1845a39 --- /dev/null +++ b/lance-spark-base_2.12/src/test/java/org/lance/spark/write/BaseReplaceWhereTest.java @@ -0,0 +1,430 @@ +/* + * 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))); + } + + /** + * 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. + */ + @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))); + } + + /** + * 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))); + } + + /** 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))); + } + + /** + * 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() { + 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)); + } + + 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. + 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); + } + } +}