From ca97ec4f39d3be6cfbc7a7fab0762318cd583f7f Mon Sep 17 00:00:00 2001 From: Wyatt Alt Date: Fri, 7 Aug 2026 07:44:33 -0700 Subject: [PATCH] fix(dataset)!: conflict a NOT NULL tightening with concurrent writes alter_columns proves a column holds no nulls by scanning, then commits the change as a Project, which conflicts with nothing. A write racing the scan can land nulls in the conflict window, leaving data that fails to read under the tightened schema. Project and Merge now carry preserves_nullability, asserting the operation makes no nullability-affecting schema change. An operation without the assertion conflicts with any concurrent value-write in either commit order, and the loser retries against the new state. False is the no-assertion default, so a transaction from a writer that predates the field can only over-conflict. alter_columns asserts it for a rename and drop_columns for a drop; a tightening withholds it. A merge is included because it can introduce a field that stale data does not supply: add_columns derives new-column nullability from the expression or takes the stream schema verbatim, subcolumn additions merge new children into existing structs, and a cast rewrites a column under a new field id. An append is the one value-write that rebases across a committed merge, and its fragments read as null for fields they predate. The assertion derives from the schema delta: a non-nullable first-new node on any path withholds it, while a nullable new node masks its whole subtree, the same rule the AllNulls transform enforces at the top level. Any new node under a non-nullable top-level column also withholds it, because the reader synthesizes missing subcolumns against the column's declared nullability and could not read the stale fragment at all. The cast path withholds it when any recast field is non-nullable. The python and java operations default to no assertion. Tightening and casting in the same alter_columns call is rejected. The check relies on read_version being the version the data was validated against, the same contract every other conflict rule already depends on. Marked breaking due to new fields on public enum variants. Co-Authored-By: Claude Fable 5 --- java/lance-jni/src/transaction.rs | 20 +- .../main/java/org/lance/operation/Merge.java | 27 +- .../java/org/lance/operation/Project.java | 40 +- .../java/org/lance/operation/MergeTest.java | 25 + .../java/org/lance/operation/ProjectTest.java | 34 ++ protos/transaction.proto | 12 + python/python/lance/dataset.py | 15 + python/python/tests/test_schema_evolution.py | 54 ++ python/src/transaction.rs | 41 +- rust/lance/src/dataset.rs | 3 + rust/lance/src/dataset/fragment.rs | 4 +- rust/lance/src/dataset/schema_evolution.rs | 186 ++++++- .../src/dataset/tests/dataset_merge_update.rs | 509 ++++++++++++++++++ rust/lance/src/dataset/transaction.rs | 138 ++++- rust/lance/src/io/commit/conflict_resolver.rs | 174 ++++++ 15 files changed, 1227 insertions(+), 55 deletions(-) diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index bc5fa4e3476..de81398ee58 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -470,13 +470,19 @@ fn convert_to_java_operation_inner<'local>( ], )?) } - Operation::Project { schema } => { + Operation::Project { + schema, + preserves_nullability, + } => { let java_schema = convert_to_java_schema(env, schema)?; Ok(env.new_object( "org/lance/operation/Project", - "(Lorg/apache/arrow/vector/types/pojo/Schema;)V", - &[JValue::Object(&java_schema)], + "(Lorg/apache/arrow/vector/types/pojo/Schema;Z)V", + &[ + JValue::Object(&java_schema), + JValue::Bool(preserves_nullability as u8), + ], )?) } Operation::Rewrite { @@ -552,16 +558,18 @@ fn convert_to_java_operation_inner<'local>( Operation::Merge { fragments: rust_fragments, schema, + preserves_nullability, } => { let java_fragments = export_vec(env, &rust_fragments)?; let java_schema = convert_to_java_schema(env, schema)?; Ok(env.new_object( "org/lance/operation/Merge", - "(Ljava/util/List;Lorg/apache/arrow/vector/types/pojo/Schema;)V", + "(Ljava/util/List;Lorg/apache/arrow/vector/types/pojo/Schema;Z)V", &[ JValue::Object(&java_fragments), JValue::Object(&java_schema), + JValue::Bool(preserves_nullability as u8), ], )?) } @@ -1044,6 +1052,8 @@ fn convert_to_rust_operation( let op_name = env.get_string_from_method(java_operation, "name")?; let op = match op_name.as_str() { "Project" => Operation::Project { + preserves_nullability: env + .get_boolean_from_method(java_operation, "preservesNullability")?, schema: convert_schema_from_operation( env, java_operation, @@ -1283,6 +1293,8 @@ fn convert_to_rust_operation( })?; Operation::Merge { fragments, + preserves_nullability: env + .get_boolean_from_method(java_operation, "preservesNullability")?, schema: convert_schema_from_operation( env, java_operation, diff --git a/java/src/main/java/org/lance/operation/Merge.java b/java/src/main/java/org/lance/operation/Merge.java index bd83657384b..07ecb880d90 100644 --- a/java/src/main/java/org/lance/operation/Merge.java +++ b/java/src/main/java/org/lance/operation/Merge.java @@ -27,16 +27,26 @@ */ public class Merge extends SchemaOperation { private final List fragments; + // True when this merge makes no nullability-affecting schema change: it + // introduces no field that data staged against an earlier schema could not + // safely omit. Without the assertion the merge conservatively conflicts with + // concurrent appends, whose fragments would omit new columns and read as null. + private final boolean preservesNullability; - protected Merge(List fragments, Schema schema) { + protected Merge(List fragments, Schema schema, boolean preservesNullability) { super(schema); this.fragments = fragments; + this.preservesNullability = preservesNullability; } public List fragments() { return fragments; } + public boolean preservesNullability() { + return preservesNullability; + } + @Override public String name() { return "Merge"; @@ -47,6 +57,7 @@ public String toString() { return MoreObjects.toStringHelper(this) .add("fragments", fragments) .add("schema", schema()) + .add("preservesNullability", preservesNullability) .toString(); } @@ -56,12 +67,13 @@ public boolean equals(Object o) { if (o == null || getClass() != o.getClass()) return false; if (!super.equals(o)) return false; Merge that = (Merge) o; - return Objects.equals(fragments, that.fragments); + return Objects.equals(fragments, that.fragments) + && preservesNullability == that.preservesNullability; } @Override public int hashCode() { - return Objects.hash(super.hashCode(), fragments); + return Objects.hash(super.hashCode(), fragments, preservesNullability); } public static Builder builder() { @@ -71,6 +83,8 @@ public static Builder builder() { public static class Builder { private List fragments; private Schema schema; + // No assertion by default, which conservatively conflicts. + private boolean preservesNullability = false; private Builder() {} @@ -84,8 +98,13 @@ public Builder schema(Schema schema) { return this; } + public Builder preservesNullability(boolean preservesNullability) { + this.preservesNullability = preservesNullability; + return this; + } + public Merge build() { - return new Merge(fragments, schema); + return new Merge(fragments, schema, preservesNullability); } } } diff --git a/java/src/main/java/org/lance/operation/Project.java b/java/src/main/java/org/lance/operation/Project.java index a4c718c4c4b..8c79431a7d1 100644 --- a/java/src/main/java/org/lance/operation/Project.java +++ b/java/src/main/java/org/lance/operation/Project.java @@ -16,14 +16,26 @@ import com.google.common.base.MoreObjects; import org.apache.arrow.vector.types.pojo.Schema; +import java.util.Objects; + /** * Project to a new schema. This Operation only changes the schema, not the data. Note: 1. For * removing columns. The data will be removed after compaction. 2. Project will modify column * positions, not ids(a.k.a. field id) */ public class Project extends SchemaOperation { - private Project(Schema schema) { + // True when this projection makes no nullability-affecting schema change, + // as a rename or a drop does not. Without the assertion the projection + // conservatively conflicts with concurrent value writes. + private final boolean preservesNullability; + + private Project(Schema schema, boolean preservesNullability) { super(schema); + this.preservesNullability = preservesNullability; + } + + public boolean preservesNullability() { + return preservesNullability; } @Override @@ -33,7 +45,23 @@ public String name() { @Override public String toString() { - return MoreObjects.toStringHelper(this).add("schema", schema()).toString(); + return MoreObjects.toStringHelper(this) + .add("schema", schema()) + .add("preservesNullability", preservesNullability) + .toString(); + } + + @Override + public boolean equals(Object o) { + if (!super.equals(o)) { + return false; + } + return preservesNullability == ((Project) o).preservesNullability; + } + + @Override + public int hashCode() { + return Objects.hash(schema(), preservesNullability); } public static Builder builder() { @@ -42,6 +70,7 @@ public static Builder builder() { public static class Builder { private Schema schema; + private boolean preservesNullability; public Builder() {} @@ -50,8 +79,13 @@ public Builder schema(Schema schema) { return this; } + public Builder preservesNullability(boolean preservesNullability) { + this.preservesNullability = preservesNullability; + return this; + } + public Project build() { - return new Project(schema); + return new Project(schema, preservesNullability); } } } diff --git a/java/src/test/java/org/lance/operation/MergeTest.java b/java/src/test/java/org/lance/operation/MergeTest.java index 841ad41d5ea..1e037285217 100644 --- a/java/src/test/java/org/lance/operation/MergeTest.java +++ b/java/src/test/java/org/lance/operation/MergeTest.java @@ -104,10 +104,14 @@ void testMergeNewColumn(@TempDir Path tempDir) throws Exception { Merge.builder() .fragments(Collections.singletonList(evolvedFragment)) .schema(evolvedSchema) + .preservesNullability(true) .build()) .build()) { try (Dataset evolvedDataset = new CommitBuilder(initialDataset).execute(mergeTxn)) { Assertions.assertEquals(3, evolvedDataset.version()); + // The explicit non-default assertion must survive the JNI round trip. + Transaction readBack = evolvedDataset.readTransaction().orElseThrow(); + Assertions.assertTrue(((Merge) readBack.operation()).preservesNullability()); Assertions.assertEquals(rowCount, evolvedDataset.countRows()); Assertions.assertEquals(evolvedSchema, evolvedDataset.getSchema()); Assertions.assertEquals(3, evolvedDataset.getSchema().getFields().size()); @@ -136,6 +140,27 @@ void testMergeNewColumn(@TempDir Path tempDir) throws Exception { } } + @Test + void testPreservesNullabilityEquality() { + Schema schema = + new Schema( + Collections.singletonList(Field.nullable("id", new ArrowType.Int(32, true))), null); + // No assertion by default, and the assertion is part of the operation's identity. + Assertions.assertFalse( + Merge.builder() + .fragments(Collections.emptyList()) + .schema(schema) + .build() + .preservesNullability()); + Assertions.assertNotEquals( + Merge.builder().fragments(Collections.emptyList()).schema(schema).build(), + Merge.builder() + .fragments(Collections.emptyList()) + .schema(schema) + .preservesNullability(true) + .build()); + } + @Test void testMergeNewColumnWithNonContiguousFieldId(@TempDir Path tempDir) throws Exception { String datasetPath = tempDir.resolve("testMergeNewColumnWithNonContiguousFieldId").toString(); diff --git a/java/src/test/java/org/lance/operation/ProjectTest.java b/java/src/test/java/org/lance/operation/ProjectTest.java index fa0c92cc15f..bd3dd7d2960 100644 --- a/java/src/test/java/org/lance/operation/ProjectTest.java +++ b/java/src/test/java/org/lance/operation/ProjectTest.java @@ -30,6 +30,8 @@ import java.util.List; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; public class ProjectTest extends OperationTestBase { @@ -69,4 +71,36 @@ void testProjection(@TempDir Path tempDir) { } } } + + @Test + void testPreservesNullabilityEqualityAndRoundTrip(@TempDir Path tempDir) { + String datasetPath = tempDir.resolve("testAssertsNonNull").toString(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + dataset = testDataset.createEmptyDataset(); + Schema schema = testDataset.getSchema(); + + // The assertion is part of the operation's identity. + assertNotEquals( + Project.builder().schema(schema).preservesNullability(true).build(), + Project.builder().schema(schema).preservesNullability(false).build()); + assertEquals( + Project.builder().schema(schema).preservesNullability(true).build(), + Project.builder().schema(schema).preservesNullability(true).build()); + + // The explicit non-default assertion must survive the JNI round trip. + try (Transaction txn = + new Transaction.Builder() + .readVersion(dataset.version()) + .operation(Project.builder().schema(schema).preservesNullability(true).build()) + .build()) { + try (Dataset committed = new CommitBuilder(dataset).execute(txn)) { + Transaction readBack = committed.readTransaction().orElseThrow(); + Project project = (Project) readBack.operation(); + assertTrue(project.preservesNullability()); + } + } + } + } } diff --git a/protos/transaction.proto b/protos/transaction.proto index c3c5e38190e..dc49aa1304b 100644 --- a/protos/transaction.proto +++ b/protos/transaction.proto @@ -144,12 +144,24 @@ message Transaction { repeated lance.file.Field schema = 2; // Schema metadata. map schema_metadata = 3; + // Set when this merge makes no nullability-affecting schema change: it + // introduces no field that data staged against an earlier schema could + // not safely omit. Without the assertion (including transactions written + // before this field existed) the merge conservatively conflicts with + // concurrent value-writes, which can only cause a retry. + bool preserves_nullability = 4; } // An operation that projects a subset of columns, altering the schema. message Project { // The new schema repeated lance.file.Field schema = 1; + // Set when this projection makes no nullability-affecting schema change, + // as a rename or a drop does not. Without the assertion (including + // transactions written before this field existed) the projection + // conservatively conflicts with concurrent value-writes, which can only + // cause a retry. A nullability tightening must not set this. + bool preserves_nullability = 2; } // An operation that restores a dataset to a previous version. diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index d387880be48..c6c04afcb9d 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -6032,6 +6032,14 @@ class Merge(BaseOperation): schema: LanceSchema or pyarrow.Schema The schema of the new dataset. Passing a LanceSchema is preferred, and passing a pyarrow.Schema is deprecated. + preserves_nullability: bool + True when this merge makes no nullability-affecting schema change: + it introduces no field that data staged against an earlier schema + could not safely omit. Without the assertion (the default) the + merge conservatively conflicts with concurrent appends, whose + fragments would omit new columns and read as null; that can only + cause a retry. Pass True when every column this merge introduces + is nullable to let concurrent appends commit without conflict. Warning ------- @@ -6077,6 +6085,7 @@ class Merge(BaseOperation): fragments: Iterable[FragmentMetadata] schema: LanceSchema | pa.Schema + preserves_nullability: bool = False def __post_init__(self): if isinstance(self.schema, pa.Schema): @@ -6249,6 +6258,11 @@ class Project(BaseOperation): ---------- schema: LanceSchema The lance schema of the new dataset. + preserves_nullability: bool + True when this projection makes no nullability-affecting schema + change, as a rename or a drop does not. Without the assertion + (the default) the projection conservatively conflicts with + concurrent writes, which can only cause a retry. Examples -------- @@ -6277,6 +6291,7 @@ class Project(BaseOperation): """ schema: LanceSchema + preserves_nullability: bool = False @dataclass class UpdateMap: diff --git a/python/python/tests/test_schema_evolution.py b/python/python/tests/test_schema_evolution.py index 7df6962789e..abd89e0cded 100644 --- a/python/python/tests/test_schema_evolution.py +++ b/python/python/tests/test_schema_evolution.py @@ -15,6 +15,7 @@ import pytest from lance import LanceDataset from lance.file import LanceFileReader, LanceFileWriter +from lance.fragment import write_fragments def test_drop_columns(tmp_path: Path): @@ -573,3 +574,56 @@ def test_add_cols_all_null_with_sql(tmp_path: Path): "b": pa.int32(), } ) + + +def test_merge_nullability_assertion(tmp_path: Path): + tbl = pa.table({"value": pa.array([1, 2], pa.int32())}) + lance.write_dataset(tbl, tmp_path) + written_at = lance.dataset(tmp_path).version + + # Stage an append against the original schema, then add a non-nullable + # column. The merge claims non-null, and reading it back must preserve + # the claim, or recommitting it would silently drop the barrier. + fragments = write_fragments( + pa.table({"value": pa.array([7], pa.int32())}), tmp_path, mode="append" + ) + lance.dataset(tmp_path).add_columns({"one": "1"}) + txn = lance.dataset(tmp_path).read_transaction(2) + assert txn is not None + assert txn.operation.preserves_nullability is False + + # The stale append omits the required column, so its rows would read as + # null under the merged schema; the claim refuses it. + op = lance.LanceOperation.Append(fragments) + with pytest.raises(Exception, match="preempted"): + lance.LanceDataset.commit(tmp_path, op, read_version=written_at) + + # A nullable add preserves nullability and skips the barrier. + lance.dataset(tmp_path).add_columns({"copied": "value"}) + txn = lance.dataset(tmp_path).read_transaction(3) + assert txn is not None + assert txn.operation.preserves_nullability is True + + +def test_project_nullability_assertion_round_trips(tmp_path: Path): + tbl = pa.table({"value": pa.array([1, 2], pa.int32())}) + lance.write_dataset(tbl, tmp_path) + lance.dataset(tmp_path).alter_columns({"path": "value", "nullable": False}) + + # Reading the tightening back must preserve the claim, or recommitting it + # would silently drop the concurrency barrier. + txn = lance.dataset(tmp_path).read_transaction(2) + assert txn is not None + assert txn.operation.preserves_nullability is False + + # A Python-built non-assertion must reach the barrier: race an append. + written_at = lance.dataset(tmp_path).version + appended = lance.write_dataset( + pa.table({"value": pa.array([7], pa.int32())}), tmp_path, mode="append" + ) + assert appended.version > written_at + relax = lance.LanceOperation.Project( + schema=txn.operation.schema, preserves_nullability=False + ) + with pytest.raises(Exception, match="preempted"): + lance.LanceDataset.commit(tmp_path, relax, read_version=written_at) diff --git a/python/src/transaction.rs b/python/src/transaction.rs index e710a36df02..f90fb8f330f 100644 --- a/python/src/transaction.rs +++ b/python/src/transaction.rs @@ -434,7 +434,18 @@ impl FromPyObject<'_, '_> for PyLance { .extract::>>()?; let fragments = fragments.into_iter().map(|f| f.0).collect(); - let op = Operation::Merge { schema, fragments }; + // Absent on objects predating the field: no assertion, which + // conservatively conflicts. + let preserves_nullability = ob + .getattr("preserves_nullability") + .and_then(|v| v.extract()) + .unwrap_or(false); + + let op = Operation::Merge { + schema, + fragments, + preserves_nullability, + }; Ok(Self(op)) } "Restore" => { @@ -482,8 +493,17 @@ impl FromPyObject<'_, '_> for PyLance { } "Project" => { let schema = extract_schema(&ob.getattr("schema")?)?; - - let op = Operation::Project { schema }; + // Absent on objects predating the field: no assertion, which + // conservatively conflicts. + let preserves_nullability = ob + .getattr("preserves_nullability") + .and_then(|v| v.extract()) + .unwrap_or(false); + + let op = Operation::Project { + schema, + preserves_nullability, + }; Ok(Self(op)) } "UpdateConfig" => { @@ -636,13 +656,17 @@ impl<'py> IntoPyObject<'py> for PyLance<&Operation> { .expect("Failed to get Delete class"); cls.call1((updated_fragments, deleted_fragment_ids, predicate)) } - Operation::Merge { fragments, schema } => { + Operation::Merge { + fragments, + schema, + preserves_nullability, + } => { let fragments_py = export_vec(py, fragments.as_slice())?; let schema_py = LanceSchema(schema.clone()); let cls = namespace .getattr("Merge") .expect("Failed to get Merge class"); - cls.call1((fragments_py, schema_py)) + cls.call1((fragments_py, schema_py, *preserves_nullability)) } Operation::Restore { version } => { let cls = namespace @@ -674,12 +698,15 @@ impl<'py> IntoPyObject<'py> for PyLance<&Operation> { .expect("Failed to get CreateIndex class"); cls.call1((new_indices_py, removed_indices_py)) } - Operation::Project { schema } => { + Operation::Project { + schema, + preserves_nullability, + } => { let schema_py = LanceSchema(schema.clone()); let cls = namespace .getattr("Project") .expect("Failed to get Project class"); - cls.call1((schema_py,)) + cls.call1((schema_py, *preserves_nullability)) } Operation::ReserveFragments { num_fragments } => { if let Ok(cls) = namespace.getattr("ReserveFragments") { diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index b98a2d1ee11..da46047c523 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -3638,11 +3638,14 @@ impl Dataset { .try_collect::>() .await?; + let preserves_nullability = + !schema_evolution::merge_introduces_required_field(self.schema(), &new_schema); let transaction = Transaction::new( self.manifest.version, Operation::Merge { fragments: updated_fragments, schema: new_schema, + preserves_nullability, }, None, ); diff --git a/rust/lance/src/dataset/fragment.rs b/rust/lance/src/dataset/fragment.rs index 3e53e604b15..621754c55f0 100644 --- a/rust/lance/src/dataset/fragment.rs +++ b/rust/lance/src/dataset/fragment.rs @@ -1999,7 +1999,7 @@ impl FileFragment { read_columns: Option>, batch_size: Option, ) -> Result<(Fragment, Schema)> { - let (fragments, schema, _) = schema_evolution::add_columns_to_fragments( + let (fragments, schema, _, _) = schema_evolution::add_columns_to_fragments( self.dataset.as_ref(), transforms, read_columns, @@ -5727,6 +5727,7 @@ mod tests { let op = Operation::Merge { fragments: merged_fragments, schema: full_schema.clone(), + preserves_nullability: true, }; let dataset = Dataset::commit( @@ -5970,6 +5971,7 @@ mod tests { Operation::Merge { schema, fragments: vec![frag], + preserves_nullability: true, }, Some(dataset.manifest.version), None, diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index 54943594d9e..6a0a47b31ef 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -254,7 +254,7 @@ pub(super) async fn add_columns_to_fragments( read_columns: Option>, fragments: &[FileFragment], batch_size: Option, -) -> Result<(Vec, Schema, Vec)> { +) -> Result<(Vec, Schema, Vec, bool)> { // Check names early (before calling add_columns_impl) to avoid extra work if // the names are wrong. let version = dataset.manifest.data_storage_format.lance_file_version()?; @@ -415,7 +415,58 @@ pub(super) async fn add_columns_to_fragments( }; schema.set_field_id(Some(dataset.manifest.max_field_id())); - Ok((new_fragments, schema, fragments_to_cleanup)) + let preserves_nullability = !merge_introduces_required_field(dataset.schema(), &schema); + + Ok(( + new_fragments, + schema, + fragments_to_cleanup, + preserves_nullability, + )) +} + +/// Whether `merged` introduces a field that data staged against `old` cannot +/// safely omit. The first new node on each path decides: a non-nullable new +/// field beneath an existing ancestor reads as unmasked null for stale rows, +/// which do supply the ancestor, while a nullable new field masks its whole +/// subtree whatever the nullability inside, the same rule the AllNulls +/// transform enforces at the top level. +/// +/// A new node under a non-nullable top-level column claims even when the node +/// itself is nullable: the reader synthesizes missing subcolumns against the +/// column's declared nullability, so a stale fragment cannot be read at all +/// under such a column, nullable child or not. +pub(super) fn merge_introduces_required_field(old: &Schema, merged: &Schema) -> bool { + /// (any node in `merged` is new, any first-new node is non-nullable) + fn subtree_new_nodes(old: &[Field], merged: &[Field]) -> (bool, bool) { + let mut any_new = false; + let mut any_required = false; + for field in merged { + match old.iter().find(|o| o.name == field.name) { + Some(old_field) => { + let (new, required) = subtree_new_nodes(&old_field.children, &field.children); + any_new |= new; + any_required |= required; + } + None => { + any_new = true; + any_required |= !field.nullable; + } + } + } + (any_new, any_required) + } + + merged.fields.iter().any( + |field| match old.fields.iter().find(|o| o.name == field.name) { + Some(old_field) => { + let (any_new, any_required) = + subtree_new_nodes(&old_field.children, &field.children); + any_required || (any_new && !field.nullable) + } + None => !field.nullable, + }, + ) } pub(super) async fn add_columns( @@ -424,16 +475,21 @@ pub(super) async fn add_columns( read_columns: Option>, batch_size: Option, ) -> Result<()> { - let (fragments, schema, _fragments_to_cleanup) = add_columns_to_fragments( - dataset, - transforms, - read_columns, - &dataset.get_fragments(), - batch_size, - ) - .await?; + let (fragments, schema, _fragments_to_cleanup, preserves_nullability) = + add_columns_to_fragments( + dataset, + transforms, + read_columns, + &dataset.get_fragments(), + batch_size, + ) + .await?; - let operation = Operation::Merge { fragments, schema }; + let operation = Operation::Merge { + fragments, + schema, + preserves_nullability, + }; let transaction = Transaction::new(dataset.manifest.version, operation, None); // Once the manifest commit has been attempted, an error does not prove // that the new files are unreferenced: the commit may have landed and only @@ -708,6 +764,7 @@ pub(super) async fn alter_columns( // Mapping of old to new fields that need to be casted. let mut cast_fields: Vec<(Field, Field)> = Vec::new(); + let mut tightens_nullability = false; let mut next_field_id = dataset.manifest.max_field_id() + 1; let version = dataset.manifest.data_storage_format.lance_file_version()?; @@ -725,6 +782,9 @@ pub(super) async fn alter_columns( && !nullable { validate_no_nulls_before_making_non_nullable(dataset, &alteration.path).await?; + // A write since this version can falsify it, so withhold the + // preserves_nullability assertion from the transaction. + tightens_nullability = true; } let field_dest = new_schema.mut_field_by_id(field_src.id).unwrap(); @@ -796,11 +856,21 @@ pub(super) async fn alter_columns( } } + if tightens_nullability && !cast_fields.is_empty() { + return Err(Error::invalid_input( + "cannot make a column non-nullable and cast columns in the same call: \ + apply the cast first, then the nullability change", + )); + } + // If we aren't casting a column, we don't need to touch the fragments. let transaction = if cast_fields.is_empty() { Transaction::new( dataset.manifest.version, - Operation::Project { schema: new_schema }, + Operation::Project { + schema: new_schema, + preserves_nullability: !tightens_nullability, + }, // TODO: Make it possible to alter blob columns /*blob_op= */ None, ) @@ -822,6 +892,12 @@ pub(super) async fn alter_columns( // This schema contains the exact field ids we want to write the new fields with. let new_col_schema = new_schema.project_by_ids(&new_ids, true); + // A cast rewrites the column under a new field id, so data staged + // against the pre-cast schema omits that id and its rows read as null. + // Withhold the assertion when any recast field is non-nullable, at any + // depth: a nested field sits under parent values stale rows do supply. + let cast_touches_required = cast_fields.iter().any(|(_old, new)| !new.nullable); + let mapper = move |batch: &RecordBatch| { let mut fields = Vec::with_capacity(cast_fields.len()); let mut columns = Vec::with_capacity(batch.num_columns()); @@ -875,6 +951,7 @@ pub(super) async fn alter_columns( Operation::Merge { schema: new_schema, fragments, + preserves_nullability: !cast_touches_required, }, /*blob_op= */ None, ) @@ -918,7 +995,10 @@ pub(super) async fn drop_columns(dataset: &mut Dataset, columns: &[&str]) -> Res let transaction = Transaction::new( dataset.manifest.version, - Operation::Project { schema: new_schema }, + Operation::Project { + schema: new_schema, + preserves_nullability: true, + }, /*blob_op= */ None, ); @@ -956,6 +1036,86 @@ pub fn exclude(source: &Schema, other: &Schema, version: &LanceFileVersion) -> R mod test { use std::{collections::HashMap, fs, num::NonZero, path::Path as StdPath, sync::Mutex}; + #[test] + fn test_merge_introduces_required_field() { + let schema = |fields: Vec| Schema::try_from(&ArrowSchema::new(fields)).unwrap(); + let strukt = |name: &str, nullable: bool, children: Vec| { + ArrowField::new( + name, + DataType::Struct(ArrowFields::from(children)), + nullable, + ) + }; + let int = |name: &str, nullable: bool| ArrowField::new(name, DataType::Int32, nullable); + + let old = schema(vec![ + strukt("s", true, vec![int("a", true)]), + strukt("r", false, vec![int("a", true)]), + ]); + // The first new node on each path decides, at any depth; any new node + // under a non-nullable top-level column claims regardless. + for (merged, expected) in [ + // A nullable new child under a non-nullable top-level column: the + // reader cannot synthesize the missing subcolumn, so claim. + ( + schema(vec![ + strukt("s", true, vec![int("a", true)]), + strukt("r", false, vec![int("a", true), int("b", true)]), + ]), + true, + ), + // Required new child under an existing parent: stale rows supply + // the parent, so the child would read as unmasked null. + ( + schema(vec![strukt( + "s", + true, + vec![int("a", true), int("b", false)], + )]), + true, + ), + ( + schema(vec![strukt( + "s", + true, + vec![int("a", true), int("b", true)], + )]), + false, + ), + // A wholly new nullable container masks its required inside. + ( + schema(vec![ + strukt("s", true, vec![int("a", true)]), + strukt("t", true, vec![int("c", false)]), + ]), + false, + ), + // Same, when the new container hangs under an existing parent. + ( + schema(vec![strukt( + "s", + true, + vec![int("a", true), strukt("t", true, vec![int("c", false)])], + )]), + false, + ), + ( + schema(vec![ + strukt("s", true, vec![int("a", true)]), + int("b", false), + ]), + true, + ), + (schema(vec![strukt("s", true, vec![int("a", true)])]), false), + ] { + assert_eq!( + merge_introduces_required_field(&old, &merged), + expected, + "merged={merged:?}" + ); + } + } + use crate::dataset::WriteParams; use arrow_array::{ ArrayRef, Int32Array, ListArray, RecordBatchIterator, StringArray, StructArray, diff --git a/rust/lance/src/dataset/tests/dataset_merge_update.rs b/rust/lance/src/dataset/tests/dataset_merge_update.rs index b9de3723098..af83c2fbc57 100644 --- a/rust/lance/src/dataset/tests/dataset_merge_update.rs +++ b/rust/lance/src/dataset/tests/dataset_merge_update.rs @@ -6,9 +6,12 @@ use std::sync::Arc; use std::time::Duration; use std::vec; +use crate::dataset::CommitBuilder; use crate::dataset::ROW_ID; use crate::dataset::WriteDestination; +use crate::dataset::builder::DatasetBuilder; use crate::dataset::optimize::{CompactionOptions, compact_files}; +use crate::dataset::schema_evolution::ColumnAlteration; use crate::dataset::transaction::{DataReplacementGroup, Operation}; use crate::dataset::{AutoCleanupParams, MergeInsertBuilder, ProjectionRequest, UpdateBuilder}; use crate::index::DatasetIndexExt; @@ -829,6 +832,7 @@ async fn test_datafile_partial_replacement() { Operation::Merge { fragments: vec![fragment], schema: extended_schema.as_ref().try_into().unwrap(), + preserves_nullability: true, }, Some(2), None, @@ -1015,6 +1019,7 @@ async fn test_datafile_replacement_error() { Operation::Merge { fragments: vec![fragment], schema: extended_schema.as_ref().try_into().unwrap(), + preserves_nullability: true, }, Some(2), None, @@ -2476,6 +2481,7 @@ async fn test_merge_rewriting_indexed_column_keeps_index_consistent() { Operation::Merge { fragments: vec![overlay], schema: schema.as_ref().try_into().unwrap(), + preserves_nullability: true, }, Some(read_version), None, @@ -4373,3 +4379,506 @@ async fn test_merge_insert_target_all_bases() { assert!(all_rows.contains(&row), "missing row {:?}", row); } } + +/// A write landing between the tightening scan and its commit falsifies the +/// claim, leaving a table that validates but cannot be scanned. +#[rstest] +#[case::tightening_conflicts(true, true)] +#[case::rename_does_not(false, false)] +#[tokio::test] +async fn test_alter_columns_conflicts_only_when_asserting( + #[case] tighten: bool, + #[case] expect_conflict: bool, +) { + let batch = arrow_array::record_batch!(("value", Int32, [1, 2])).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + + // Leave the first handle a version behind, so the alteration commits stale. + let appended = InsertBuilder::new(WriteDestination::Dataset(Arc::new(dataset.clone()))) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute(vec![ + arrow_array::record_batch!(("value", Int32, [3])).unwrap(), + ]) + .await + .unwrap(); + assert_eq!(appended.version().version, 2); + + let mut stale = dataset; + let alteration = if tighten { + ColumnAlteration::new("value".into()).set_nullable(false) + } else { + ColumnAlteration::new("value".into()).rename("renamed".into()) + }; + let result = stale.alter_columns(&[alteration]).await; + assert_eq!( + result.is_err(), + expect_conflict, + "tighten={tighten}: got {result:?}" + ); +} + +/// read_version is the version the data was validated against. Declaring it +/// honestly puts a later tightening inside the conflict window; declaring a +/// later version skips the checks, which is a caller bug, not a guarantee. +#[rstest] +#[case::honest_read_version_conflicts(1, true)] +#[case::misdeclared_read_version_commits(2, false)] +#[tokio::test] +async fn test_stale_append_protection_follows_read_version( + #[case] declared: u64, + #[case] expect_conflict: bool, +) { + let batch = arrow_array::record_batch!(("value", Int32, [1, 2])).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let dataset = Arc::new(Dataset::write(reader, "memory://", None).await.unwrap()); + + let mut tightened = dataset.schema().clone(); + tightened.fields[0].nullable = false; + let tightened = Dataset::commit( + WriteDestination::Dataset(dataset.clone()), + Operation::Project { + schema: tightened, + preserves_nullability: false, + }, + Some(dataset.version().version), + None, + None, + Arc::new(Default::default()), + false, + ) + .await + .unwrap(); + assert_eq!(tightened.version().version, 2); + + let mut append = InsertBuilder::new(WriteDestination::Dataset(dataset.clone())) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted(vec![batch]) + .await + .unwrap(); + append.read_version = declared; + + let result = CommitBuilder::new(Arc::new(tightened)) + .execute(append) + .await; + assert_eq!( + result.is_err(), + expect_conflict, + "declared={declared}: got {result:?}" + ); +} + +/// A field added and tightened after the write snapshot: the tightening's +/// claim is in the honest conflict window, so the operation-wide barrier +/// rejects the stale append. Added-but-nullable commits, since synthesized +/// nulls are legal there. +#[rstest] +#[case::added_then_tightened(true, true)] +#[case::added_still_nullable(false, false)] +#[tokio::test] +async fn test_stale_append_vs_field_added_since( + #[case] tighten: bool, + #[case] expect_conflict: bool, +) { + let batch = arrow_array::record_batch!(("value", Int32, [1, 2])).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let dataset = Arc::new(Dataset::write(reader, "memory://", None).await.unwrap()); + + let append = InsertBuilder::new(WriteDestination::Dataset(dataset.clone())) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted(vec![batch]) + .await + .unwrap(); + + let mut latest = dataset.as_ref().clone(); + latest + .add_columns( + crate::dataset::NewColumnTransform::SqlExpressions(vec![( + "new_value".to_string(), + "value".to_string(), + )]), + None, + None, + ) + .await + .unwrap(); + if tighten { + latest + .alter_columns(&[ColumnAlteration::new("new_value".into()).set_nullable(false)]) + .await + .unwrap(); + } + + let result = CommitBuilder::new(Arc::new(latest)).execute(append).await; + assert_eq!( + result.is_err(), + expect_conflict, + "tighten={tighten}: got {result:?}" + ); +} + +#[tokio::test] +async fn test_alter_columns_rejects_cast_with_tightening() { + let batch = arrow_array::record_batch!(("value", Int32, [1, 2])).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + + let err = dataset + .alter_columns(&[ColumnAlteration::new("value".into()) + .set_nullable(false) + .cast_to(DataType::Int64)]) + .await + .unwrap_err(); + assert!(err.to_string().contains("same call"), "got: {err}"); + + // Separately, both succeed. + dataset + .alter_columns(&[ColumnAlteration::new("value".into()).cast_to(DataType::Int64)]) + .await + .unwrap(); + dataset + .alter_columns(&[ColumnAlteration::new("value".into()).set_nullable(false)]) + .await + .unwrap(); +} + +/// A merge introducing a non-nullable column claims non-null, so a stale +/// append, whose fragments omit the column and would read as null, conflicts. +/// A nullable column keeps the long-standing behavior: the append commits and +/// its rows legally read as null. +#[rstest] +#[case::required_column_conflicts("1", true)] +#[case::nullable_column_commits("value", false)] +#[tokio::test] +async fn test_stale_append_vs_column_added_by_merge( + #[case] expression: &str, + #[case] expect_conflict: bool, +) { + let batch = arrow_array::record_batch!(("value", Int32, [1, 2])).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let dataset = Arc::new(Dataset::write(reader, "memory://", None).await.unwrap()); + + let append = InsertBuilder::new(WriteDestination::Dataset(dataset.clone())) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted(vec![batch]) + .await + .unwrap(); + + // A literal is non-nullable; a projection of a nullable column is nullable. + let mut latest = dataset.as_ref().clone(); + latest + .add_columns( + crate::dataset::NewColumnTransform::SqlExpressions(vec![( + "new_value".to_string(), + expression.to_string(), + )]), + None, + None, + ) + .await + .unwrap(); + + let result = CommitBuilder::new(Arc::new(latest)).execute(append).await; + assert_eq!( + result.is_err(), + expect_conflict, + "expression={expression}: got {result:?}" + ); + if let Ok(committed) = result { + committed.scan().try_into_batch().await.unwrap(); + } +} + +/// A cast rewrites the column under a new field id, so a stale append omits it +/// and its rows read as null. Casting a non-nullable column therefore claims +/// non-null and conflicts; casting a nullable one does not. +#[rstest] +#[case::cast_of_required_conflicts(true, true)] +#[case::cast_of_nullable_commits(false, false)] +#[tokio::test] +async fn test_stale_append_vs_cast(#[case] tighten_first: bool, #[case] expect_conflict: bool) { + let batch = arrow_array::record_batch!(("value", Int32, [1, 2])).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + + if tighten_first { + dataset + .alter_columns(&[ColumnAlteration::new("value".into()).set_nullable(false)]) + .await + .unwrap(); + } + + // Stage against the pre-cast schema, so the cast lands inside the window. + let staged = Arc::new(dataset.clone()); + let append = InsertBuilder::new(WriteDestination::Dataset(staged.clone())) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted(vec![batch]) + .await + .unwrap(); + + dataset + .alter_columns(&[ColumnAlteration::new("value".into()).cast_to(DataType::Int64)]) + .await + .unwrap(); + + let result = CommitBuilder::new(Arc::new(dataset)).execute(append).await; + assert_eq!( + result.is_err(), + expect_conflict, + "tighten_first={tighten_first}: got {result:?}" + ); + if let Ok(committed) = result { + committed.scan().try_into_batch().await.unwrap(); + } +} + +/// A subcolumn addition (V2.2+) merges a new child into an existing struct. +/// Stale rows supply the parent, so a required new child would read as +/// unmasked null: the merge claims and the stale append conflicts. A nullable +/// new child under a nullable parent masks itself and keeps appends flowing. +/// Under a non-nullable parent even a nullable child claims: the reader +/// synthesizes missing subcolumns against the column's declared nullability, +/// so the stale fragment could not be read at all. +#[rstest] +#[case::required_child_conflicts(true, false, true)] +#[case::required_child_nullable_parent_conflicts(true, true, true)] +#[case::nullable_child_nullable_parent_commits(false, true, false)] +#[case::nullable_child_required_parent_conflicts(false, false, true)] +#[tokio::test] +async fn test_stale_append_vs_sub_column_added_by_merge( + #[case] child_required: bool, + #[case] parent_nullable: bool, + #[case] expect_conflict: bool, +) { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "b", + DataType::Struct(Fields::from(vec![ArrowField::new( + "c", + DataType::Int32, + true, + )])), + parent_nullable, + )])); + let struct_batch = |values: Vec| { + RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StructArray::from(vec![( + Arc::new(ArrowField::new("c", DataType::Int32, true)), + Arc::new(Int32Array::from(values)) as ArrayRef, + )]))], + ) + .unwrap() + }; + let dataset = Arc::new( + Dataset::write( + RecordBatchIterator::new(vec![Ok(struct_batch(vec![1, 2]))], schema.clone()), + "memory://", + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_2), + ..Default::default() + }), + ) + .await + .unwrap(), + ); + + let append = InsertBuilder::new(WriteDestination::Dataset(dataset.clone())) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute_uncommitted(vec![struct_batch(vec![3])]) + .await + .unwrap(); + + let new_child = Arc::new(ArrowField::new("d", DataType::Int32, !child_required)); + let sub_schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "b", + DataType::Struct(Fields::from(vec![new_child.as_ref().clone()])), + parent_nullable, + )])); + let sub_batch = RecordBatch::try_new( + sub_schema.clone(), + vec![Arc::new(StructArray::from(vec![( + new_child, + Arc::new(Int32Array::from(vec![10, 20])) as ArrayRef, + )]))], + ) + .unwrap(); + let mut latest = dataset.as_ref().clone(); + latest + .add_columns( + crate::dataset::NewColumnTransform::Reader(Box::new(RecordBatchIterator::new( + vec![Ok(sub_batch)], + sub_schema, + ))), + None, + None, + ) + .await + .unwrap(); + + let result = CommitBuilder::new(Arc::new(latest)).execute(append).await; + assert_eq!( + result.is_err(), + expect_conflict, + "child_required={child_required} parent_nullable={parent_nullable}: got {result:?}" + ); + if let Ok(committed) = result { + // The stale rows keep their parent values; the new child reads null. + let batch = committed.scan().try_into_batch().await.unwrap(); + assert_eq!(batch.num_rows(), 3); + let parent = batch["b"].as_struct(); + assert_eq!(parent.null_count(), 0); + assert_eq!(parent.column_by_name("d").unwrap().null_count(), 1); + } +} + +/// The barrier is operation-wide, so a tightening of a nested field conflicts +/// with a concurrent write exactly like a top-level one. +#[tokio::test] +async fn test_alter_columns_nested_tightening_conflicts() { + let schema = Arc::new(ArrowSchema::new(vec![ArrowField::new( + "b", + DataType::Struct(Fields::from(vec![ArrowField::new( + "c", + DataType::Int32, + true, + )])), + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(StructArray::from(vec![( + Arc::new(ArrowField::new("c", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, + )]))], + ) + .unwrap(); + let dataset = Dataset::write( + RecordBatchIterator::new(vec![Ok(batch.clone())], schema.clone()), + "memory://", + None, + ) + .await + .unwrap(); + + // Leave the first handle a version behind, so the tightening commits stale. + InsertBuilder::new(WriteDestination::Dataset(Arc::new(dataset.clone()))) + .with_params(&WriteParams { + mode: WriteMode::Append, + ..Default::default() + }) + .execute(vec![batch]) + .await + .unwrap(); + + let mut stale = dataset; + stale + .alter_columns(&[ColumnAlteration::new("b.c".into()).set_nullable(false)]) + .await + .unwrap_err(); +} + +/// The invariant every piece of the tightening barrier serves: no interleaving +/// of honest writers and schema changes may commit a dataset that validates +/// but cannot be scanned. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_concurrent_tightening_stress() { + let dir = TempStrDir::default(); + let uri = dir.as_str().to_string(); + let batch = arrow_array::record_batch!(("value", Int32, [1, 2])).unwrap(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + Dataset::write(reader, uri.as_str(), None).await.unwrap(); + + let mut tasks = tokio::task::JoinSet::new(); + + // Appenders: honest read versions, half the batches carry nulls. A null + // append must either land while the column is nullable or be refused -- + // by the writer against a non-null schema, or by the claim barrier when a + // tightening won the race after the write. + for a in 0..4u8 { + let uri = uri.clone(); + tasks.spawn(async move { + let mut outcomes = [0u32; 2]; + for i in 0..12u32 { + let with_null = (a as u32 + i).is_multiple_of(2); + let batch = if with_null { + arrow_array::record_batch!(("value", Int32, [None, Some(3)])).unwrap() + } else { + arrow_array::record_batch!(("value", Int32, [4, 5])).unwrap() + }; + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + let result = Dataset::write( + reader, + uri.as_str(), + Some(WriteParams { + mode: WriteMode::Append, + ..Default::default() + }), + ) + .await; + outcomes[result.is_ok() as usize] += 1; + } + outcomes + }); + } + + // The tightener alternates NOT NULL and back. Either step may lose to + // concurrent writes; losing is an acceptable outcome, corruption is not. + { + let uri = uri.clone(); + tasks.spawn(async move { + let mut outcomes = [0u32; 2]; + for i in 0..10u32 { + let Ok(mut dataset) = DatasetBuilder::from_uri(uri.as_str()).load().await else { + continue; + }; + let result = dataset + .alter_columns( + &[ColumnAlteration::new("value".into()).set_nullable(i % 2 == 1)], + ) + .await; + outcomes[result.is_ok() as usize] += 1; + } + outcomes + }); + } + + let mut totals = [0u32; 2]; + while let Some(res) = tasks.join_next().await { + let [err, ok] = res.unwrap(); + totals[0] += err; + totals[1] += ok; + } + + // The oracle: whatever interleaving happened, the final dataset must be + // internally consistent -- validation and scanning agree. + let dataset = DatasetBuilder::from_uri(uri.as_str()).load().await.unwrap(); + dataset.validate().await.unwrap(); + let scanned = dataset.scan().try_into_batch().await.unwrap(); + assert!(scanned.num_rows() >= 2); + // And every historical version must scan too: a corrupt intermediate + // commit would have been the bug even if later commits papered over it. + for version in 1..=dataset.version().version { + let at = dataset.checkout_version(version).await.unwrap(); + at.validate().await.unwrap(); + at.scan().try_into_batch().await.unwrap(); + } + assert!(totals[1] > 0, "nothing succeeded: {totals:?}"); +} diff --git a/rust/lance/src/dataset/transaction.rs b/rust/lance/src/dataset/transaction.rs index be757867d0e..c4d713e930a 100644 --- a/rust/lance/src/dataset/transaction.rs +++ b/rust/lance/src/dataset/transaction.rs @@ -406,6 +406,12 @@ pub enum Operation { Merge { fragments: Vec, schema: Schema, + /// Set when this merge makes no nullability-affecting schema change: + /// it introduces no field that data staged against an earlier schema + /// could not safely omit. Without the assertion the merge conflicts + /// with concurrent appends in either commit order, since a stale + /// append omits new columns entirely and its rows read as null. + preserves_nullability: bool, }, /// Restore an old version of the database Restore { version: u64 }, @@ -455,8 +461,16 @@ pub enum Operation { updated_fragment_offsets: Option, }, - /// Project to a new schema. This only changes the schema, not the data. - Project { schema: Schema }, + /// Project to a new schema. + Project { + schema: Schema, + /// Set when this projection makes no nullability-affecting schema + /// change, as a rename or a drop does not. A nullability tightening + /// must not set this: its producer proved the claim by scanning at its + /// read version, so a concurrent write can falsify it and the + /// projection conflicts with value-writes in either commit order. + preserves_nullability: bool, + }, /// Update the dataset configuration. UpdateConfig { @@ -649,12 +663,18 @@ impl PartialEq for Operation { Self::Merge { fragments: a_fragments, schema: a_schema, + preserves_nullability: a_preserves, }, Self::Merge { fragments: b_fragments, schema: b_schema, + preserves_nullability: b_preserves, }, - ) => compare_vec(a_fragments, b_fragments) && a_schema == b_schema, + ) => { + compare_vec(a_fragments, b_fragments) + && a_schema == b_schema + && a_preserves == b_preserves + } (Self::Restore { version: a }, Self::Restore { version: b }) => a == b, ( Self::ReserveFragments { num_fragments: a }, @@ -697,7 +717,16 @@ impl PartialEq for Operation { && a_inserted_rows_filter == b_inserted_rows_filter && a_updated_fragment_offsets == b_updated_fragment_offsets } - (Self::Project { schema: a }, Self::Project { schema: b }) => a == b, + ( + Self::Project { + schema: a, + preserves_nullability: a_preserves, + }, + Self::Project { + schema: b, + preserves_nullability: b_preserves, + }, + ) => a == b && a_preserves == b_preserves, ( Self::UpdateConfig { config_updates: a_config, @@ -3468,12 +3497,17 @@ impl TryFrom for Transaction { fragments, schema, schema_metadata: _schema_metadata, // TODO: handle metadata + preserves_nullability, })) => Operation::Merge { fragments: fragments .into_iter() .map(Fragment::try_from) .collect::>>()?, schema: Schema::try_from(&Fields(schema))?, + // False for a writer that predates the field: no assertion, so + // a legacy required-field merge still conflicts and a legacy + // nullable merge over-conflicts, which only retries. + preserves_nullability, }, Some(pb::transaction::Operation::Restore(pb::transaction::Restore { version })) => { Operation::Restore { version } @@ -3525,11 +3559,16 @@ impl TryFrom for Transaction { } }, }, - Some(pb::transaction::Operation::Project(pb::transaction::Project { schema })) => { - Operation::Project { - schema: Schema::try_from(&Fields(schema))?, - } - } + Some(pb::transaction::Operation::Project(pb::transaction::Project { + schema, + preserves_nullability, + })) => Operation::Project { + schema: Schema::try_from(&Fields(schema))?, + // False for a writer that predates the field: no assertion, so + // a legacy tightening still conflicts and a legacy rename + // over-conflicts, which only retries. + preserves_nullability, + }, Some(pb::transaction::Operation::UpdateConfig(update_config)) => { // Check if new-style fields are present let has_new_fields = update_config.config_updates.is_some() @@ -3815,13 +3854,16 @@ impl From<&Transaction> for pb::Transaction { .map(pb::IndexMetadata::from) .collect(), }), - Operation::Merge { fragments, schema } => { - pb::transaction::Operation::Merge(pb::transaction::Merge { - fragments: fragments.iter().map(pb::DataFragment::from).collect(), - schema: Fields::from(schema).0, - schema_metadata: Default::default(), // TODO: handle metadata - }) - } + Operation::Merge { + fragments, + schema, + preserves_nullability, + } => pb::transaction::Operation::Merge(pb::transaction::Merge { + fragments: fragments.iter().map(pb::DataFragment::from).collect(), + schema: Fields::from(schema).0, + schema_metadata: Default::default(), // TODO: handle metadata + preserves_nullability: *preserves_nullability, + }), Operation::Restore { version } => { pb::transaction::Operation::Restore(pb::transaction::Restore { version: *version }) } @@ -3869,11 +3911,13 @@ impl From<&Transaction> for pb::Transaction { }) .unwrap_or_default(), }), - Operation::Project { schema } => { - pb::transaction::Operation::Project(pb::transaction::Project { - schema: Fields::from(schema).0, - }) - } + Operation::Project { + schema, + preserves_nullability, + } => pb::transaction::Operation::Project(pb::transaction::Project { + schema: Fields::from(schema).0, + preserves_nullability: *preserves_nullability, + }), Operation::UpdateConfig { config_updates, table_metadata_updates, @@ -4022,10 +4066,12 @@ pub fn validate_operation(manifest: Option<&Manifest>, operation: &Operation) -> // Fragments must contain all fields in the schema schema_fragments_valid(Some(manifest), &manifest.schema, fragments) } - Operation::Project { schema } => { + Operation::Project { schema, .. } => { schema_fragments_valid(Some(manifest), schema, manifest.fragments.as_ref()) } - Operation::Merge { fragments, schema } => { + Operation::Merge { + fragments, schema, .. + } => { merge_fragments_valid(manifest, fragments)?; schema_fragments_valid(Some(manifest), schema, fragments) } @@ -5718,6 +5764,7 @@ mod tests { Operation::Merge { fragments: vec![merged_fragment], schema: lance_schema, + preserves_nullability: true, }, None, ); @@ -5803,6 +5850,7 @@ mod tests { Operation::Merge { fragments: vec![merged_fragment], schema: lance_schema, + preserves_nullability: true, }, None, ); @@ -5878,6 +5926,7 @@ mod tests { Operation::Merge { fragments: vec![merged_fragment], schema: lance_schema, + preserves_nullability: true, }, None, ); @@ -5957,6 +6006,7 @@ mod tests { Operation::Merge { fragments: vec![existing_fragment, new_fragment], schema: lance_schema, + preserves_nullability: true, }, None, ); @@ -6983,4 +7033,46 @@ mod tests { }; assert_ne!(overlay(1), rewrite); } + + #[test] + fn test_nullability_assertion_defaults_conservative() { + // A writer that predates the field encodes nothing, which decodes as + // false: no assertion, so a legacy tightening or required-field merge + // still conflicts. Only an explicit true skips the barrier. + for encoded in [false, true] { + let txn = Transaction::try_from(pb::Transaction { + read_version: 1, + uuid: "test".to_string(), + operation: Some(pb::transaction::Operation::Project( + pb::transaction::Project { + schema: vec![], + preserves_nullability: encoded, + }, + )), + ..Default::default() + }) + .unwrap(); + assert!( + matches!(txn.operation, Operation::Project { preserves_nullability, .. } if preserves_nullability == encoded), + "encoded={encoded:?}" + ); + + let txn = Transaction::try_from(pb::Transaction { + read_version: 1, + uuid: "test".to_string(), + operation: Some(pb::transaction::Operation::Merge(pb::transaction::Merge { + fragments: vec![], + schema: vec![], + schema_metadata: Default::default(), + preserves_nullability: encoded, + })), + ..Default::default() + }) + .unwrap(); + assert!( + matches!(txn.operation, Operation::Merge { preserves_nullability, .. } if preserves_nullability == encoded), + "encoded={encoded:?}" + ); + } + } } diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index d28b1cc9882..0c4e3d28889 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -39,6 +39,39 @@ pub struct TransactionRebase<'a> { conflicting_mem_wal_compacted_sstables: Vec, } +/// Whether `operation` may make a nullability-affecting schema change: a +/// projection or merge that does not assert `preserves_nullability`. A +/// tightening projection scanned for nulls at its read version, and a merge +/// may introduce a field that data staged against an earlier schema cannot +/// safely omit; either is falsified by a concurrent value-write. +fn may_alter_nullability(operation: &Operation) -> bool { + matches!( + operation, + Operation::Project { + preserves_nullability: false, + .. + } | Operation::Merge { + preserves_nullability: false, + .. + } + ) +} + +/// Whether `operation` can commit rows that falsify such a change: by writing +/// a null into a scanned field, or by omitting a required column entirely (a +/// stale append's fragments read as null for columns they predate). `Delete` +/// only removes rows and `Rewrite` preserves the values it moves; `Project` +/// already conflicts with projections and merges elsewhere. +fn supplies_values(operation: &Operation) -> bool { + matches!( + operation, + Operation::Append { .. } + | Operation::Update { .. } + | Operation::DataReplacement { .. } + | Operation::DataOverlay { .. } + ) +} + impl<'a> TransactionRebase<'a> { pub async fn try_new( dataset: &Dataset, @@ -221,6 +254,15 @@ impl<'a> TransactionRebase<'a> { /// Will return an error if the transaction is not valid. Otherwise, it will /// return Ok(()). pub fn check_txn(&mut self, other_transaction: &Transaction, other_version: u64) -> Result<()> { + // Either order: the claim was checked without the write's data. + let ours = &self.transaction.operation; + let theirs = &other_transaction.operation; + if (may_alter_nullability(ours) && supplies_values(theirs)) + || (supplies_values(ours) && may_alter_nullability(theirs)) + { + return Err(self.retryable_conflict_err(other_transaction, other_version)); + } + let op = &self.transaction.operation; match op { Operation::Delete { .. } => self.check_delete_txn(other_transaction, other_version), @@ -2679,6 +2721,7 @@ mod tests { Operation::Merge { fragments: vec![fragment0.clone(), fragment2.clone()], schema: lance_core::datatypes::Schema::default(), + preserves_nullability: true, }, Operation::Overwrite { fragments: vec![fragment0.clone(), fragment2.clone()], @@ -2874,6 +2917,7 @@ mod tests { Operation::Merge { fragments: vec![fragment0.clone(), fragment2.clone()], schema: lance_core::datatypes::Schema::default(), + preserves_nullability: true, }, // Merge conflicts with everything except CreateIndex and ReserveFragments. [ @@ -3276,6 +3320,7 @@ mod tests { Operation::Merge { fragments: vec![fragment1.clone()], schema: lance_core::datatypes::Schema::default(), + preserves_nullability: true, }, Retryable, ), @@ -3496,6 +3541,134 @@ mod tests { } } + /// An append is the one value-write that rebases across a committed merge, + /// so a merge introducing a required field must claim: the append's + /// fragments omit the new column and its rows would read as null. A merge + /// without the claim keeps the long-standing behavior of appends passing + /// over nullable column adds. The reverse order conflicts regardless of + /// the claim, because a rebasing merge rewrites the whole fragment list. + #[test] + fn test_merge_claim_blocks_stale_append() { + for claims in [true, false] { + let merge = Operation::Merge { + fragments: vec![Fragment::new(0)], + schema: lance_core::datatypes::Schema::default(), + preserves_nullability: !claims, + }; + let append = Operation::Append { + fragments: vec![Fragment::new(1)], + }; + + let mut append_rebase = TransactionRebase { + transaction: Transaction::new(0, append.clone(), None), + initial_fragments: HashMap::new(), + modified_fragment_ids: HashSet::new(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_compacted_sstables: Vec::new(), + }; + let result = append_rebase.check_txn(&Transaction::new(0, merge.clone(), None), 1); + assert_eq!( + matches!(result, Err(Error::RetryableCommitConflict { .. })), + claims, + "append rebasing over merge/claims={claims}: got {result:?}" + ); + + let mut merge_rebase = TransactionRebase { + transaction: Transaction::new(0, merge, None), + initial_fragments: HashMap::new(), + modified_fragment_ids: HashSet::from_iter([0]), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_compacted_sstables: Vec::new(), + }; + let result = merge_rebase.check_txn(&Transaction::new(0, append, None), 1); + assert!( + matches!(result, Err(Error::RetryableCommitConflict { .. })), + "merge rebasing over append/claims={claims}: got {result:?}" + ); + } + } + + /// A claim conflicts with any write that can supply values, either order. + #[test] + fn test_non_null_claim_barrier() { + use crate::dataset::transaction::{DataOverlayGroup, UpdateMode}; + use lance_table::format::overlay::{DataOverlayFile, OverlayCoverage}; + use roaring::RoaringBitmap; + + let file = || DataFile::new_legacy_from_fields("w.lance", vec![0], None); + let writers = [ + ( + "append", + Operation::Append { + fragments: vec![Fragment::new(1)], + }, + ), + ( + "update", + Operation::Update { + removed_fragment_ids: vec![], + updated_fragments: vec![Fragment::new(0)], + new_fragments: vec![], + fields_modified: vec![0], + compacted_sstables: Vec::new(), + fields_for_preserving_frag_bitmap: vec![], + update_mode: Some(UpdateMode::RewriteColumns), + inserted_rows_filter: None, + updated_fragment_offsets: None, + }, + ), + ( + "replacement", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, file())], + }, + ), + ( + "overlay", + Operation::DataOverlay { + groups: vec![DataOverlayGroup { + fragment_id: 0, + overlays: vec![DataOverlayFile { + data_file: file(), + coverage: OverlayCoverage::dense(RoaringBitmap::from_iter([0u32])), + committed_version: 0, + }], + }], + }, + ), + ]; + + for (writer_name, writer) in &writers { + for claims in [true, false] { + let project = Operation::Project { + schema: lance_core::datatypes::Schema::default(), + preserves_nullability: !claims, + }; + for (order, ours, theirs) in [ + ("project-rebasing", project.clone(), writer.clone()), + ("writer-rebasing", writer.clone(), project.clone()), + ] { + let mut rebase = TransactionRebase { + transaction: Transaction::new(0, ours.clone(), None), + initial_fragments: HashMap::new(), + modified_fragment_ids: modified_fragment_ids(&ours).collect::>(), + affected_rows: None, + conflicting_frag_reuse_indices: Vec::new(), + conflicting_mem_wal_compacted_sstables: Vec::new(), + }; + let result = rebase.check_txn(&Transaction::new(0, theirs, None), 1); + assert_eq!( + matches!(result, Err(Error::RetryableCommitConflict { .. })), + claims, + "{writer_name}/claims={claims}/{order}: got {result:?}" + ); + } + } + } + } + #[tokio::test] #[rstest::rstest] #[case::coverage_overlaps_moved_row(vec![0u32], true)] @@ -4429,6 +4602,7 @@ mod tests { Operation::Merge { fragments: vec![Fragment::new(0)], schema: lance_core::datatypes::Schema::default(), + preserves_nullability: true, }, Retryable, ),