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, ),