diff --git a/docs/src/guide/read_and_write.md b/docs/src/guide/read_and_write.md index 40062b6ce33..ffeab642308 100644 --- a/docs/src/guide/read_and_write.md +++ b/docs/src/guide/read_and_write.md @@ -497,8 +497,9 @@ ordering of the data will be preserved. !!! note - Compaction creates a new version of the table. It does not delete the old - version of the table and the files referenced by it. + Compaction creates a new version of the table, or two when it both + rewrites fragments and repacks columns (see below). It does not delete the + old versions of the table and the files referenced by them. ```python import lance @@ -517,6 +518,41 @@ of this, it's recommended to rewrite files before re-building indices. +#### Repack columns + +Each `add_columns` backfill gives every fragment one more data file. A large +fragment with few deletions is never rewritten, so those files pile up and slow +down scans and random reads. Compaction can repack a fragment's columns into +fewer data files instead. A repack moves no rows, so fragment ids, row +addresses, deletions, data overlays and index coverage stay as they are. + +Repacking is off by default. Set either option, as a `compact_files` argument or +as table config: + +- `max_data_files_per_fragment` (`lance.compaction.max_data_files_per_fragment`): + repack a fragment that holds its columns in more files than this. +- `column_groups` (`lance.compaction.column_groups`, groups separated by `;` and + columns by `,`): keep these top-level columns together in a file of their + own. Fragments that compaction rewrites are written this way too. + +`scope` limits a run to one kind of work: `"rewrite_fragments"`, +`"repack_columns"`, or `"all"` (the default). + +```python +dataset.optimize.compact_files(max_data_files_per_fragment=2) + +# Only repack, and keep the embedding in a file of its own. +dataset.optimize.compact_files(column_groups=[["embedding"]], scope="repack_columns") + +# Per fragment: live data files, file sizes, fields per file, dead field slots. +dataset.stats.column_layout_stats() +``` + +A repack only merges files it can empty. Files holding a blob column, a column +only partly present in the fragment, or spilled row lineage stay as they are, +so a fragment can stay above the limit. Fragments with legacy (V1) data files +are not repacked. + ### Cleanup old versions Lance is an immutable format — every write creates a new version. The new version diff --git a/java/lance-jni/src/blocking_dataset.rs b/java/lance-jni/src/blocking_dataset.rs index 0659f1cb88c..a4571bc1724 100644 --- a/java/lance-jni/src/blocking_dataset.rs +++ b/java/lance-jni/src/blocking_dataset.rs @@ -10,9 +10,9 @@ use crate::namespace::{ use crate::session::{handle_from_session, session_from_handle}; use crate::traits::{FromJObjectWithEnv, FromJString, export_vec, import_vec, import_vec_to_rust}; use crate::utils::{ - build_compaction_options, extract_base_store_params, extract_storage_options, - extract_write_params, get_scalar_index_params, get_vector_index_params, to_java_map, - to_rust_map, + apply_repack_options, build_compaction_options, extract_base_store_params, + extract_storage_options, extract_write_params, get_scalar_index_params, + get_vector_index_params, to_java_map, to_rust_map, }; use crate::{block_on, traits::IntoJava}; use arrow::array::RecordBatchReader; @@ -1964,6 +1964,86 @@ fn inner_get_fragment_statistics<'local>( )?) } +#[unsafe(no_mangle)] +pub extern "system" fn Java_org_lance_Dataset_nativeGetColumnLayoutStatistics<'a>( + mut env: JNIEnv<'a>, + jdataset: JObject, +) -> JObject<'a> { + ok_or_throw!(env, inner_get_column_layout_statistics(&mut env, jdataset)) +} + +/// `Dataset::column_layout_stats` as parallel Java arrays. +fn inner_get_column_layout_statistics<'local>( + env: &mut JNIEnv<'local>, + jdataset: JObject, +) -> Result> { + let stats = { + let dataset = + unsafe { env.get_rust_field::<_, _, BlockingDataset>(jdataset, NATIVE_DATASET) }?; + dataset.inner.column_layout_stats() + }; + let to_int = |value: usize| { + i32::try_from(value) + .map_err(|_| Error::runtime_error(format!("{value} does not fit in a Java int"))) + }; + let len = to_int(stats.len())?; + let fragment_ids: Vec = stats.iter().map(|s| s.fragment_id as i64).collect(); + let live_file_counts = stats + .iter() + .map(|s| to_int(s.live_file_count)) + .collect::>>()?; + let ratios: Vec = stats.iter().map(|s| s.tombstoned_field_ratio).collect(); + let overlay_counts = stats + .iter() + .map(|s| to_int(s.overlay_count)) + .collect::>>()?; + + let jfragment_ids = env.new_long_array(len)?; + let jlive_file_counts = env.new_int_array(len)?; + let jratios = env.new_double_array(len)?; + let joverlay_counts = env.new_int_array(len)?; + env.set_long_array_region(&jfragment_ids, 0, &fragment_ids)?; + env.set_int_array_region(&jlive_file_counts, 0, &live_file_counts)?; + env.set_double_array_region(&jratios, 0, &ratios)?; + env.set_int_array_region(&joverlay_counts, 0, &overlay_counts)?; + + let jfile_sizes = env.new_object_array(len, "[J", JObject::null())?; + let jfields_per_file = env.new_object_array(len, "[I", JObject::null())?; + for (index, s) in stats.iter().enumerate() { + let sizes: Vec = s + .file_sizes + .iter() + .map(|size| size.map_or(-1, |size| size as i64)) + .collect(); + let fields = s + .fields_per_file + .iter() + .map(|count| to_int(*count)) + .collect::>>()?; + let jsizes = env.new_long_array(to_int(sizes.len())?)?; + env.set_long_array_region(&jsizes, 0, &sizes)?; + env.set_object_array_element(&jfile_sizes, index as i32, &jsizes)?; + env.delete_local_ref(jsizes)?; + let jfields = env.new_int_array(to_int(fields.len())?)?; + env.set_int_array_region(&jfields, 0, &fields)?; + env.set_object_array_element(&jfields_per_file, index as i32, &jfields)?; + env.delete_local_ref(jfields)?; + } + + Ok(env.new_object( + "org/lance/ColumnLayoutStatistics", + "([J[I[[J[[I[D[I)V", + &[ + JValue::Object(&jfragment_ids), + JValue::Object(&jlive_file_counts), + JValue::Object(&jfile_sizes), + JValue::Object(&jfields_per_file), + JValue::Object(&jratios), + JValue::Object(&joverlay_counts), + ], + )?) +} + #[unsafe(no_mangle)] pub extern "system" fn Java_org_lance_Dataset_getFragmentNative<'a>( mut env: JNIEnv<'a>, @@ -3628,7 +3708,22 @@ fn convert_java_compaction_options_to_rust( )? .l()?; - build_compaction_options( + let max_data_files_per_fragment = env + .call_method( + &java_options, + "getMaxDataFilesPerFragment", + "()Ljava/util/Optional;", + &[], + )? + .l()?; + let column_groups = env + .call_method(&java_options, "getColumnGroups", "()Ljava/util/List;", &[])? + .l()?; + let scope = env + .call_method(&java_options, "getScope", "()Ljava/util/Optional;", &[])? + .l()?; + + let mut options = build_compaction_options( env, &target_rows_per_fragment, &max_rows_per_group, @@ -3646,7 +3741,15 @@ fn convert_java_compaction_options_to_rust( &excluded_fragment_ids, &data_storage_version, config, - ) + )?; + apply_repack_options( + env, + &mut options, + &max_data_files_per_fragment, + &column_groups, + &scope, + )?; + Ok(options) } #[unsafe(no_mangle)] diff --git a/java/lance-jni/src/optimize.rs b/java/lance-jni/src/optimize.rs index ce0151d202a..edc630cc3a3 100644 --- a/java/lance-jni/src/optimize.rs +++ b/java/lance-jni/src/optimize.rs @@ -11,8 +11,9 @@ use jni::{ use lance::dataset::{ index::DatasetIndexRemapperOptions, optimize::{ - CompactionMetrics, CompactionMode, CompactionOptions, CompactionPlan, CompactionTask, - IndexRemapperOptions, RewriteResult, TaskData, commit_compaction, plan_compaction, + CompactionMetrics, CompactionMode, CompactionOptions, CompactionPlan, CompactionScope, + CompactionTask, CompactionTaskKind, IndexRemapperOptions, RepackedFiles, RewriteResult, + TaskData, commit_compaction, plan_compaction, }, }; @@ -23,12 +24,13 @@ use crate::{ FromJObjectWithEnv, IntoJava, export_vec, import_vec_from_method, import_vec_to_rust, }, utils::{ - build_compaction_options, to_java_boolean_obj, to_java_float_obj, to_java_list, - to_java_long_obj, to_java_optional, + apply_repack_options, build_compaction_options, to_java_boolean_obj, to_java_float_obj, + to_java_list, to_java_long_obj, to_java_optional, }, }; use crate::error::Result; +use crate::ffi::JNIEnvExt; #[unsafe(no_mangle)] pub extern "system" fn Java_org_lance_compaction_Compaction_nativePlanCompaction<'local>( @@ -50,6 +52,9 @@ pub extern "system" fn Java_org_lance_compaction_Compaction_nativePlanCompaction max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List data_storage_version: JObject, // Optional + max_data_files_per_fragment: JObject, // Optional + column_groups: JObject, // List> + scope: JObject, // Optional ) -> JObject<'local> { ok_or_throw_with_return!( env, @@ -70,7 +75,10 @@ pub extern "system" fn Java_org_lance_compaction_Compaction_nativePlanCompaction max_source_rows, max_source_bytes, excluded_fragment_ids, - data_storage_version + data_storage_version, + max_data_files_per_fragment, + column_groups, + scope ), JObject::null() ) @@ -95,13 +103,16 @@ fn inner_plan_compaction<'local>( max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List data_storage_version: JObject, // Optional + max_data_files_per_fragment: JObject, // Optional + column_groups: JObject, // List> + scope: JObject, // Optional ) -> Result> { let config = { let dataset = unsafe { env.get_rust_field::<_, _, BlockingDataset>(&java_dataset, NATIVE_DATASET) }?; dataset.inner.manifest.config.clone() }; - let compaction_options = build_compaction_options( + let mut compaction_options = build_compaction_options( env, &target_rows_per_fragment, &max_rows_per_group, @@ -120,6 +131,13 @@ fn inner_plan_compaction<'local>( &data_storage_version, &config, )?; + apply_repack_options( + env, + &mut compaction_options, + &max_data_files_per_fragment, + &column_groups, + &scope, + )?; let plan = { let dataset = @@ -258,6 +276,9 @@ pub extern "system" fn Java_org_lance_compaction_CompactionTask_nativeExecute<'l max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List data_storage_version: JObject, // Optional + max_data_files_per_fragment: JObject, // Optional + column_groups: JObject, // List> + scope: JObject, // Optional ) -> JObject<'local> { ok_or_throw_with_return!( env, @@ -280,7 +301,10 @@ pub extern "system" fn Java_org_lance_compaction_CompactionTask_nativeExecute<'l max_source_rows, max_source_bytes, excluded_fragment_ids, - data_storage_version + data_storage_version, + max_data_files_per_fragment, + column_groups, + scope ), JObject::null() ) @@ -307,6 +331,9 @@ fn inner_execute_task<'local>( max_source_bytes: JObject, // Optional excluded_fragment_ids: JObject, // List data_storage_version: JObject, // Optional + max_data_files_per_fragment: JObject, // Optional + column_groups: JObject, // List> + scope: JObject, // Optional ) -> Result> { let task_data: TaskData = task_data.extract_object(env)?; let config = { @@ -314,7 +341,7 @@ fn inner_execute_task<'local>( unsafe { env.get_rust_field::<_, _, BlockingDataset>(&java_dataset, NATIVE_DATASET) }?; dataset.inner.manifest.config.clone() }; - let compaction_options = build_compaction_options( + let mut compaction_options = build_compaction_options( env, &target_rows_per_fragment, &max_rows_per_group, @@ -333,6 +360,13 @@ fn inner_execute_task<'local>( &data_storage_version, &config, )?; + apply_repack_options( + env, + &mut compaction_options, + &max_data_files_per_fragment, + &column_groups, + &scope, + )?; let compaction_task = CompactionTask { task: task_data, read_version: read_version as u64, @@ -347,26 +381,48 @@ fn inner_execute_task<'local>( } const TASK_DATA_CLASS: &str = "org/lance/compaction/TaskData"; -const TASK_DATA_CONSTRUCTOR_SIG: &str = "(Ljava/util/List;)V"; +const TASK_DATA_CONSTRUCTOR_SIG: &str = "(Ljava/util/List;Ljava/util/List;)V"; const COMPACTION_METRICS_CLASS: &str = "org/lance/compaction/CompactionMetrics"; const COMPACTION_METRICS_CONSTRUCTOR_SIG: &str = "(JJJJ)V"; const COMPACTION_PLAN_CLASS: &str = "org/lance/compaction/CompactionPlan"; const COMPACTION_PLAN_CONSTRUCTOR_SIG: &str = "(Ljava/util/List;JLorg/lance/compaction/CompactionOptions;)V"; const REWRITE_RESULT_CLASS: &str = "org/lance/compaction/RewriteResult"; -const REWRITE_RESULT_CONSTRUCTOR_SIG: &str = - "(Lorg/lance/compaction/CompactionMetrics;Ljava/util/List;Ljava/util/List;J[B)V"; +const REWRITE_RESULT_CONSTRUCTOR_SIG: &str = "(Lorg/lance/compaction/CompactionMetrics;Ljava/util/List;Ljava/util/List;J[BLjava/lang/Long;Ljava/util/List;)V"; const COMPACTION_OPTIONS_CLASS: &str = "org/lance/compaction/CompactionOptions"; const COMPACTION_MODE_CLASS: &str = "org/lance/compaction/CompactionMode"; -const COMPACTION_OPTIONS_CONSTRUCTOR_SIG: &str = "(Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/List;Ljava/util/Optional;)V"; +const COMPACTION_OPTIONS_CONSTRUCTOR_SIG: &str = "(Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/List;Ljava/util/Optional;Ljava/util/Optional;Ljava/util/List;Ljava/util/Optional;)V"; impl IntoJava for &TaskData { fn into_java<'a>(self, env: &mut JNIEnv<'a>) -> Result> { let fragments = export_vec(env, &self.fragments)?; + let repack_files = match &self.kind { + CompactionTaskKind::RewriteFragments => JObject::null(), + CompactionTaskKind::RepackColumns { files } => { + let mut lists = Vec::with_capacity(files.len()); + for file in files { + let ids = file + .iter() + .map(|id| { + Ok(env.new_object( + "java/lang/Integer", + "(I)V", + &[JValueGen::Int(*id)], + )?) + }) + .collect::>>()?; + lists.push(to_java_list(env, &ids)?); + } + to_java_list(env, &lists)? + } + }; Ok(env.new_object( TASK_DATA_CLASS, TASK_DATA_CONSTRUCTOR_SIG, - &[JValueGen::Object(&fragments)], + &[ + JValueGen::Object(&fragments), + JValueGen::Object(&repack_files), + ], )?) } } @@ -445,6 +501,25 @@ impl IntoJava for &CompactionOptions { None => JObject::null(), }; let data_storage_version_opt = to_java_optional(env, data_storage_version)?; + let max_data_files_per_fragment = + to_java_long_obj(env, self.max_data_files_per_fragment.map(|v| v as i64))?; + let max_data_files_per_fragment_opt = to_java_optional(env, max_data_files_per_fragment)?; + let mut column_groups = Vec::with_capacity(self.column_groups.len()); + for group in &self.column_groups { + let names = group + .iter() + .map(|name| Ok(JObject::from(env.new_string(name)?))) + .collect::>>()?; + column_groups.push(to_java_list(env, &names)?); + } + let column_groups = to_java_list(env, &column_groups)?; + let scope = match self.scope { + CompactionScope::All => "all", + CompactionScope::RewriteFragments => "rewrite_fragments", + CompactionScope::RepackColumns => "repack_columns", + }; + let scope: JObject = env.new_string(scope)?.into(); + let scope_opt = to_java_optional(env, scope)?; Ok(env.new_object( COMPACTION_OPTIONS_CLASS, @@ -465,6 +540,9 @@ impl IntoJava for &CompactionOptions { JValueGen::Object(&max_source_bytes_opt), JValueGen::Object(&excluded_fragment_ids), JValueGen::Object(&data_storage_version_opt), + JValueGen::Object(&max_data_files_per_fragment_opt), + JValueGen::Object(&column_groups), + JValueGen::Object(&scope_opt), ], )?) } @@ -496,6 +574,13 @@ impl IntoJava for &RewriteResult { } else { JObject::null() }; + let (repacked_fragment_id, repacked_files) = match &self.repacked_files { + Some(repacked) => ( + to_java_long_obj(env, Some(repacked.fragment_id as i64))?, + export_vec(env, &repacked.files)?, + ), + None => (JObject::null(), JObject::null()), + }; Ok(env.new_object( REWRITE_RESULT_CLASS, REWRITE_RESULT_CONSTRUCTOR_SIG, @@ -505,6 +590,8 @@ impl IntoJava for &RewriteResult { JValueGen::Object(&original_fragments), JValueGen::Long(self.read_version as i64), JValueGen::Object(&row_addrs), + JValueGen::Object(&repacked_fragment_id), + JValueGen::Object(&repacked_files), ], )?) } @@ -534,8 +621,19 @@ impl FromJObjectWithEnv for JObject<'_> { let task_data = import_vec_from_method(env, self, "getFragments", |env, fragment| { fragment.extract_object(env) })?; + let repack_files = env + .call_method(self, "getRepackFiles", "()Ljava/util/List;", &[])? + .l()?; + let kind = if repack_files.is_null() { + CompactionTaskKind::RewriteFragments + } else { + let files = + import_vec_to_rust(env, &repack_files, |env, file| env.get_integers(&file))?; + CompactionTaskKind::RepackColumns { files } + }; Ok(TaskData { fragments: task_data, + kind, }) } } @@ -569,12 +667,27 @@ impl FromJObjectWithEnv for JObject<'_> { } else { Some(env.convert_byte_array(row_addrs_obj)?) }; + let repacked_fragment_id = env + .call_method(self, "getRepackedFragmentId", "()Ljava/lang/Long;", &[])? + .l()?; + let repacked_files = if repacked_fragment_id.is_null() { + None + } else { + let fragment_id = env + .call_method(&repacked_fragment_id, "longValue", "()J", &[])? + .j()? as u64; + let files = import_vec_from_method(env, self, "getRepackedFiles", |env, file| { + file.extract_object(env) + })?; + Some(RepackedFiles { fragment_id, files }) + }; Ok(RewriteResult { metrics, new_fragments, read_version, original_fragments, row_addrs, + repacked_files, }) } } diff --git a/java/lance-jni/src/transaction.rs b/java/lance-jni/src/transaction.rs index eb06d599f2d..40f0df3522a 100644 --- a/java/lance-jni/src/transaction.rs +++ b/java/lance-jni/src/transaction.rs @@ -1125,13 +1125,19 @@ fn convert_to_java_operation_inner<'local>( )?; Ok(java_operation) } - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { + replacements, + data_change, + } => { let java_replacements = export_vec(env, &replacements)?; Ok(env.new_object( "org/lance/operation/DataReplacement", - "(Ljava/util/List;)V", - &[JValue::Object(&java_replacements)], + "(Ljava/util/List;Z)V", + &[ + JValue::Object(&java_replacements), + JValue::Bool(data_change as u8), + ], )?) } Operation::DataOverlay { groups } => { @@ -1998,7 +2004,13 @@ fn convert_to_rust_operation( import_vec_from_method(env, java_operation, "replacements", |env, replacement| { replacement.extract_object(env) })?; - Operation::DataReplacement { replacements } + let data_change = env + .call_method(java_operation, "dataChange", "()Z", &[])? + .z()?; + Operation::DataReplacement { + replacements, + data_change, + } } "DataOverlay" => { let groups = import_vec_from_method(env, java_operation, "getGroups", |env, group| { diff --git a/java/lance-jni/src/utils.rs b/java/lance-jni/src/utils.rs index 5a3f795880c..e9d950c7aa1 100644 --- a/java/lance-jni/src/utils.rs +++ b/java/lance-jni/src/utils.rs @@ -8,7 +8,7 @@ use arrow_schema::{DataType, Field}; use jni::JNIEnv; use jni::objects::{JFloatArray, JMap, JObject, JString, JValue, JValueGen}; use jni::sys::{jboolean, jfloat, jlong}; -use lance::dataset::optimize::{CompactionMode, CompactionOptions}; +use lance::dataset::optimize::{CompactionMode, CompactionOptions, CompactionScope}; use lance::dataset::{WriteMode, WriteParams}; use lance::index::vector::{IndexFileVersion, StageParams, VectorIndexParams}; use lance::io::ObjectStoreParams; @@ -286,6 +286,34 @@ pub fn build_compaction_options( Ok(compaction_options) } +/// Apply the column repack options a Java `CompactionOptions` carries. An empty +/// `column_groups` keeps the table config's groups. +pub fn apply_repack_options( + env: &mut JNIEnv, + options: &mut CompactionOptions, + max_data_files_per_fragment: &JObject, // Optional + column_groups: &JObject, // List> + scope: &JObject, // Optional +) -> Result<()> { + if let Some(max) = env.get_long_opt(max_data_files_per_fragment)? { + options.max_data_files_per_fragment = Some(usize::try_from(max).map_err(|_| { + Error::input_error(format!( + "max_data_files_per_fragment must be positive, got {max}" + )) + })?); + } + let groups = crate::traits::import_vec_to_rust(env, column_groups, |env, group| { + env.get_strings(&group) + })?; + if !groups.is_empty() { + options.column_groups = groups; + } + if let Some(scope) = env.get_string_opt(scope)? { + options.scope = CompactionScope::try_from(scope.as_str())?; + } + Ok(()) +} + // Convert from Java Optional to Rust Option pub fn get_query(env: &mut JNIEnv, query_obj: JObject) -> Result> { let query = env.get_optional(&query_obj, |env, java_obj| { diff --git a/java/src/main/java/org/lance/ColumnLayoutStatistics.java b/java/src/main/java/org/lance/ColumnLayoutStatistics.java new file mode 100644 index 00000000000..2c7031d7552 --- /dev/null +++ b/java/src/main/java/org/lance/ColumnLayoutStatistics.java @@ -0,0 +1,90 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance; + +/** + * Per-fragment column-layout statistics of a dataset version as parallel arrays: index {@code i} of + * every getter describes the same fragment. Returned by {@link + * Dataset#getColumnLayoutStatistics()}. Compaction repacks a fragment's columns into fewer files + * when its live file count is above {@code maxDataFilesPerFragment}. + */ +public final class ColumnLayoutStatistics { + private final long[] fragmentIds; + private final int[] liveFileCounts; + private final long[][] fileSizes; + private final int[][] fieldsPerFile; + private final double[] tombstonedFieldRatios; + private final int[] overlayCounts; + + ColumnLayoutStatistics( + long[] fragmentIds, + int[] liveFileCounts, + long[][] fileSizes, + int[][] fieldsPerFile, + double[] tombstonedFieldRatios, + int[] overlayCounts) { + this.fragmentIds = fragmentIds; + this.liveFileCounts = liveFileCounts; + this.fileSizes = fileSizes; + this.fieldsPerFile = fieldsPerFile; + this.tombstonedFieldRatios = tombstonedFieldRatios; + this.overlayCounts = overlayCounts; + } + + /** Fragment IDs in manifest order. */ + public long[] getFragmentIds() { + return fragmentIds; + } + + /** + * Number of data files holding at least one column of the schema, per fragment, aligned with + * {@link #getFragmentIds()}. A file left holding only tombstones or dropped columns, or kept only + * for the fragment's spilled row lineage, is not counted. + */ + public int[] getLiveFileCounts() { + return liveFileCounts; + } + + /** + * Recorded size in bytes of each data file, per fragment, in the fragment's file order. {@code + * -1} when the manifest has no size for the file. + */ + public long[][] getFileSizes() { + return fileSizes; + } + + /** Number of schema fields each data file holds, per fragment, in the fragment's file order. */ + public int[][] getFieldsPerFile() { + return fieldsPerFile; + } + + /** + * Share of each fragment's field slots holding no live data: slots tombstoned by a column update + * or repack, or left by a dropped column. Spilled row lineage is not counted. A fragment rewrite + * reclaims these slots. + */ + public double[] getTombstonedFieldRatios() { + return tombstonedFieldRatios; + } + + /** Number of data overlay files per fragment, aligned with {@link #getFragmentIds()}. */ + public int[] getOverlayCounts() { + return overlayCounts; + } + + /** Number of fragments described. */ + public int size() { + return fragmentIds.length; + } +} diff --git a/java/src/main/java/org/lance/Dataset.java b/java/src/main/java/org/lance/Dataset.java index 99b4e3cff9b..13503f2c54a 100644 --- a/java/src/main/java/org/lance/Dataset.java +++ b/java/src/main/java/org/lance/Dataset.java @@ -1577,6 +1577,22 @@ public FragmentStatistics getFragmentStatistics() { private native FragmentStatistics nativeGetFragmentStatistics(); + /** + * Get per-fragment column-layout statistics for this dataset version: how many data files hold a + * column of the schema, and how many overlay files each fragment carries. Only manifest metadata + * is read. + * + * @return column-layout statistics as parallel arrays, in manifest order + */ + public ColumnLayoutStatistics getColumnLayoutStatistics() { + try (LockManager.ReadLock readLock = lockManager.acquireReadLock()) { + Preconditions.checkArgument(nativeDatasetHandle != 0, "Dataset is closed"); + return nativeGetColumnLayoutStatistics(); + } + } + + private native ColumnLayoutStatistics nativeGetColumnLayoutStatistics(); + /** * Gets the arrow schema of the dataset. * diff --git a/java/src/main/java/org/lance/compaction/Compaction.java b/java/src/main/java/org/lance/compaction/Compaction.java index e739d09dac8..7bd561f2446 100644 --- a/java/src/main/java/org/lance/compaction/Compaction.java +++ b/java/src/main/java/org/lance/compaction/Compaction.java @@ -50,7 +50,10 @@ public static CompactionPlan planCompaction( compactionOptions.getMaxSourceRows(), compactionOptions.getMaxSourceBytes(), compactionOptions.getExcludedFragmentIds(), - compactionOptions.getDataStorageVersion()); + compactionOptions.getDataStorageVersion(), + compactionOptions.getMaxDataFilesPerFragment(), + compactionOptions.getColumnGroups(), + compactionOptions.getScope()); } } @@ -155,5 +158,8 @@ private static native CompactionPlan nativePlanCompaction( Optional maxSourceRows, Optional maxSourceBytes, List excludedFragmentIds, - Optional dataStorageVersion); + Optional dataStorageVersion, + Optional maxDataFilesPerFragment, + List> columnGroups, + Optional scope); } diff --git a/java/src/main/java/org/lance/compaction/CompactionOptions.java b/java/src/main/java/org/lance/compaction/CompactionOptions.java index d83e618da5b..3dd7ed2750b 100644 --- a/java/src/main/java/org/lance/compaction/CompactionOptions.java +++ b/java/src/main/java/org/lance/compaction/CompactionOptions.java @@ -22,6 +22,7 @@ import java.io.ObjectOutputStream; import java.io.OptionalDataException; import java.io.Serializable; +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Objects; @@ -54,7 +55,11 @@ public class CompactionOptions implements Serializable { private Optional maxSourceBytes; private List excludedFragmentIds; private Optional dataStorageVersion; + private Optional maxDataFilesPerFragment; + private List> columnGroups; + private Optional scope; + // Also called from the native layer. private CompactionOptions( Optional targetRowsPerFragment, Optional maxRowsPerGroup, @@ -70,7 +75,10 @@ private CompactionOptions( Optional maxSourceRows, Optional maxSourceBytes, List excludedFragmentIds, - Optional dataStorageVersion) { + Optional dataStorageVersion, + Optional maxDataFilesPerFragment, + List> columnGroups, + Optional scope) { this.targetRowsPerFragment = targetRowsPerFragment; this.maxRowsPerGroup = maxRowsPerGroup; this.maxBytesPerFile = maxBytesPerFile; @@ -86,6 +94,30 @@ private CompactionOptions( this.maxSourceBytes = maxSourceBytes; this.excludedFragmentIds = List.copyOf(excludedFragmentIds); this.dataStorageVersion = dataStorageVersion.map(DataStorageVersion::fromRustString); + this.maxDataFilesPerFragment = maxDataFilesPerFragment; + this.columnGroups = copyGroups(columnGroups); + this.scope = scope.map(CompactionOptions::scopeFromValue); + } + + private static List> copyGroups(List> groups) { + List> copy = new ArrayList<>(groups.size()); + for (List group : groups) { + copy.add(List.copyOf(group)); + } + return Collections.unmodifiableList(copy); + } + + public Optional getMaxDataFilesPerFragment() { + return maxDataFilesPerFragment; + } + + public List> getColumnGroups() { + return columnGroups; + } + + /** Returns the scope as its string value for the native layer. */ + public Optional getScope() { + return scope.map(CompactionScope::getValue); } public Optional getDeferIndexRemap() { @@ -172,6 +204,9 @@ public String toString() { .add("maxSourceBytes", maxSourceBytes.orElse(null)) .add("excludedFragmentIds", excludedFragmentIds) .add("dataStorageVersion", dataStorageVersion.orElse(null)) + .add("maxDataFilesPerFragment", maxDataFilesPerFragment.orElse(null)) + .add("columnGroups", columnGroups) + .add("scope", scope.orElse(null)) .toString(); } @@ -191,6 +226,9 @@ private void writeObject(ObjectOutputStream output) throws IOException { output.writeObject(maxSourceBytes.orElse(null)); output.writeObject(excludedFragmentIds); output.writeObject(getDataStorageVersion().orElse(null)); + output.writeObject(maxDataFilesPerFragment.orElse(null)); + output.writeObject(new ArrayList<>(columnGroups)); + output.writeObject(getScope().orElse(null)); } private void readObject(ObjectInputStream input) throws IOException, ClassNotFoundException { @@ -218,6 +256,32 @@ private void readObject(ObjectInputStream input) throws IOException, ClassNotFou this.maxSourceBytes = readTrailingLong(input); this.excludedFragmentIds = readTrailingLongList(input); this.dataStorageVersion = readTrailingString(input).map(DataStorageVersion::fromRustString); + this.maxDataFilesPerFragment = readTrailingLong(input); + this.columnGroups = readTrailingGroups(input); + this.scope = readTrailingString(input).map(CompactionOptions::scopeFromValue); + } + + private static CompactionScope scopeFromValue(String value) { + for (CompactionScope scope : CompactionScope.values()) { + if (scope.getValue().equals(value)) { + return scope; + } + } + throw new IllegalArgumentException("Unknown compaction scope: " + value); + } + + @SuppressWarnings("unchecked") + private static List> readTrailingGroups(ObjectInputStream input) + throws IOException, ClassNotFoundException { + try { + List> groups = (List>) input.readObject(); + return groups == null ? Collections.emptyList() : copyGroups(groups); + } catch (OptionalDataException e) { + if (!e.eof) { + throw e; + } + return Collections.emptyList(); + } } /** @@ -280,6 +344,9 @@ public static class Builder { private Optional maxSourceBytes = Optional.empty(); private List excludedFragmentIds = Collections.emptyList(); private Optional dataStorageVersion = Optional.empty(); + private Optional maxDataFilesPerFragment = Optional.empty(); + private List> columnGroups = Collections.emptyList(); + private Optional scope = Optional.empty(); private Builder() {} @@ -402,6 +469,38 @@ public Builder withExcludedFragmentIds(List excludedFragmentIds) { return this; } + /** + * Maximum number of data files a fragment may hold columns in before its columns are repacked + * into fewer files. Each {@code addColumns} backfill adds one file per fragment. A repack + * rewrites only the columns that move and keeps rows, fragment ids and indices as they are. + * Unset means no file-count trigger. + * + * @throws IllegalArgumentException if {@code maxDataFilesPerFragment} is not positive + */ + public Builder withMaxDataFilesPerFragment(long maxDataFilesPerFragment) { + this.maxDataFilesPerFragment = + Optional.of(positiveBudget("maxDataFilesPerFragment", maxDataFilesPerFragment)); + return this; + } + + /** + * Top-level columns to keep in their own data files. Each inner list becomes one data file per + * fragment; the columns no group names share one file. Fragments the compaction rewrites are + * written this way, and the others have their columns repacked to match. Names that are not + * top-level columns are ignored. An empty list keeps the {@code lance.compaction.column_groups} + * table config, if any. + */ + public Builder withColumnGroups(List> columnGroups) { + this.columnGroups = copyGroups(Objects.requireNonNull(columnGroups, "columnGroups")); + return this; + } + + /** Which kinds of task to plan. Defaults to {@link CompactionScope#ALL}. */ + public Builder withScope(CompactionScope scope) { + this.scope = Optional.of(Objects.requireNonNull(scope, "scope")); + return this; + } + /** * A max source budget of zero admits no work and a negative value would wrap around to an * effectively unlimited budget on the Rust side, so both are rejected here. Leave the option @@ -431,7 +530,10 @@ public CompactionOptions build() { maxSourceRows, maxSourceBytes, excludedFragmentIds, - dataStorageVersion.map(DataStorageVersion::toRustString)); + dataStorageVersion.map(DataStorageVersion::toRustString), + maxDataFilesPerFragment, + columnGroups, + scope.map(CompactionScope::getValue)); } } } diff --git a/java/src/main/java/org/lance/compaction/CompactionScope.java b/java/src/main/java/org/lance/compaction/CompactionScope.java new file mode 100644 index 00000000000..c5d509fbbb6 --- /dev/null +++ b/java/src/main/java/org/lance/compaction/CompactionScope.java @@ -0,0 +1,34 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.lance.compaction; + +/** Which kinds of task a compaction plans. */ +public enum CompactionScope { + /** Rewrite the fragments that need it and repack the columns of the others (default). */ + ALL("all"), + /** Only rewrite fragments. */ + REWRITE_FRAGMENTS("rewrite_fragments"), + /** Only repack columns. */ + REPACK_COLUMNS("repack_columns"); + + private final String value; + + CompactionScope(String value) { + this.value = value; + } + + public String getValue() { + return value; + } +} diff --git a/java/src/main/java/org/lance/compaction/CompactionTask.java b/java/src/main/java/org/lance/compaction/CompactionTask.java index ab4f9c7c7a2..e5cedbb6587 100644 --- a/java/src/main/java/org/lance/compaction/CompactionTask.java +++ b/java/src/main/java/org/lance/compaction/CompactionTask.java @@ -63,7 +63,10 @@ public RewriteResult execute(Dataset dataset) { compactionOptions.getMaxSourceRows(), compactionOptions.getMaxSourceBytes(), compactionOptions.getExcludedFragmentIds(), - compactionOptions.getDataStorageVersion()); + compactionOptions.getDataStorageVersion(), + compactionOptions.getMaxDataFilesPerFragment(), + compactionOptions.getColumnGroups(), + compactionOptions.getScope()); } } @@ -85,7 +88,10 @@ private native RewriteResult nativeExecute( Optional maxSourceRows, Optional maxSourceBytes, List excludedFragmentIds, - Optional dataStorageVersion); + Optional dataStorageVersion, + Optional maxDataFilesPerFragment, + List> columnGroups, + Optional scope); public CompactionOptions getCompactionOptions() { return compactionOptions; diff --git a/java/src/main/java/org/lance/compaction/RewriteResult.java b/java/src/main/java/org/lance/compaction/RewriteResult.java index c4c5d816c2c..f33c52a4f64 100644 --- a/java/src/main/java/org/lance/compaction/RewriteResult.java +++ b/java/src/main/java/org/lance/compaction/RewriteResult.java @@ -14,6 +14,7 @@ package org.lance.compaction; import org.lance.FragmentMetadata; +import org.lance.fragment.DataFile; import javax.annotation.Nullable; @@ -25,6 +26,10 @@ * committed later. */ public class RewriteResult implements Serializable { + // Pinned to the UID generated before the repack fields were added, so that results produced by + // older workers still deserialize during a rolling upgrade. + private static final long serialVersionUID = 4501818269828675274L; + private final CompactionMetrics metrics; private final List newFragments; private final List originalFragments; @@ -34,17 +39,54 @@ public class RewriteResult implements Serializable { // null for stable row IDs. @Nullable private final byte[] rowAddrs; + // Set only for a column repack: the fragment it repacked and the data files it wrote. + @Nullable private final Long repackedFragmentId; + @Nullable private final List repackedFiles; + public RewriteResult( CompactionMetrics metrics, List newFragments, List originalFragments, long readVersion, byte[] rowAddrs) { + this(metrics, newFragments, originalFragments, readVersion, rowAddrs, null, null); + } + + public RewriteResult( + CompactionMetrics metrics, + List newFragments, + List originalFragments, + long readVersion, + byte[] rowAddrs, + @Nullable Long repackedFragmentId, + @Nullable List repackedFiles) { this.metrics = metrics; this.newFragments = newFragments; this.originalFragments = originalFragments; this.readVersion = readVersion; this.rowAddrs = rowAddrs; + this.repackedFragmentId = repackedFragmentId; + this.repackedFiles = repackedFiles; + } + + /** + * The fragment a column repack wrote new data files for. + * + * @return null for the result of a fragment rewrite + */ + @Nullable + public Long getRepackedFragmentId() { + return repackedFragmentId; + } + + /** + * The data files a column repack wrote, empty when the fragment had no rows to write. + * + * @return null for the result of a fragment rewrite + */ + @Nullable + public List getRepackedFiles() { + return repackedFiles; } public long getReadVersion() { diff --git a/java/src/main/java/org/lance/compaction/TaskData.java b/java/src/main/java/org/lance/compaction/TaskData.java index 4ec5958215c..2d9cfa95f23 100644 --- a/java/src/main/java/org/lance/compaction/TaskData.java +++ b/java/src/main/java/org/lance/compaction/TaskData.java @@ -15,18 +15,48 @@ import org.lance.FragmentMetadata; +import javax.annotation.Nullable; + import java.io.Serializable; import java.util.List; -/** Data of compaction task. */ +/** + * Data of compaction task. + * + *

A task either rewrites its fragments into new fragments, or, when {@link #getRepackFiles()} is + * non-null, rewrites some columns of its one fragment into new data files and leaves the fragment + * id, rows, deletions, overlays and index coverage as they are. + */ public class TaskData implements Serializable { + // Pinned to the UID generated before repackFiles was added, so that tasks queued by older workers + // still deserialize (as fragment rewrites) during a rolling upgrade. + private static final long serialVersionUID = -4884632518342713596L; + private final List fragments; + // One entry per new data file: the field ids of the top-level columns it holds. + @Nullable private final List> repackFiles; + public TaskData(List fragments) { + this(fragments, null); + } + + public TaskData(List fragments, @Nullable List> repackFiles) { this.fragments = fragments; + this.repackFiles = repackFiles; } public List getFragments() { return fragments; } + + /** + * The new data files of a column repack, each the field ids of the top-level columns it holds. + * + * @return null for a task that rewrites its fragments + */ + @Nullable + public List> getRepackFiles() { + return repackFiles; + } } diff --git a/java/src/main/java/org/lance/operation/DataReplacement.java b/java/src/main/java/org/lance/operation/DataReplacement.java index 42f0c12615f..b46e123f28a 100644 --- a/java/src/main/java/org/lance/operation/DataReplacement.java +++ b/java/src/main/java/org/lance/operation/DataReplacement.java @@ -21,21 +21,22 @@ import java.util.Objects; /** - * Replace data in a column in the dataset with new data. This is used for null column population - * where we replace an entirely null column with a new column that has data. + * Replace the data files backing some fields of existing fragments with new files, without moving + * rows. Each group names a fragment and one new data file for it. At commit, a file holding exactly + * the new file's fields is swapped for it; otherwise the new file's fields are tombstoned where + * they live and the new file is appended. A fragment can take several groups, one per new file. + * Used for null column population, and by compaction to repack columns into fewer files. * - *

This operation will only allow replacing files that contain the same schema e.g. if the - * original files contain columns A, B, C and the new files contain only columns A, B then the - * operation is not allowed. - * - *

Corollary to the above: the operation will also not allow replacing files unless the affected - * columns all have the same datafile layout across the fragments being replaced. + *

{@code dataChange == false} declares that the new files hold the same values as the files they + * replace: indices keep their coverage, overlays keep shadowing, and no row is reported as updated. */ public class DataReplacement implements Operation { private final List replacements; + private final boolean dataChange; - private DataReplacement(List replacements) { + private DataReplacement(List replacements, boolean dataChange) { this.replacements = replacements; + this.dataChange = dataChange; } /** @@ -47,6 +48,15 @@ public List replacements() { return replacements; } + /** + * Whether the new files change any value. + * + * @return false if the values were only moved to new files + */ + public boolean dataChange() { + return dataChange; + } + @Override public String name() { return "DataReplacement"; @@ -54,7 +64,10 @@ public String name() { @Override public String toString() { - return MoreObjects.toStringHelper(this).add("replacements", replacements).toString(); + return MoreObjects.toStringHelper(this) + .add("replacements", replacements) + .add("dataChange", dataChange) + .toString(); } @Override @@ -62,7 +75,12 @@ public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; DataReplacement that = (DataReplacement) o; - return Objects.equals(replacements, that.replacements); + return dataChange == that.dataChange && Objects.equals(replacements, that.replacements); + } + + @Override + public int hashCode() { + return Objects.hash(replacements, dataChange); } /** @@ -77,6 +95,7 @@ public static Builder builder() { /** Builder for DataReplacement. */ public static class Builder { private List replacements; + private boolean dataChange = true; public Builder() {} @@ -91,13 +110,25 @@ public Builder replacements(List replacements) { return this; } + /** + * Set whether the new files change any value. Defaults to true; pass false only when the new + * files hold the same values as the files they replace. + * + * @param dataChange whether the new files change any value + * @return this builder + */ + public Builder dataChange(boolean dataChange) { + this.dataChange = dataChange; + return this; + } + /** * Build a new DataReplacement. * * @return a new DataReplacement */ public DataReplacement build() { - return new DataReplacement(replacements); + return new DataReplacement(replacements, dataChange); } } diff --git a/java/src/test/java/org/lance/CompactionTest.java b/java/src/test/java/org/lance/CompactionTest.java index 4c874c4c878..e0019af1890 100644 --- a/java/src/test/java/org/lance/CompactionTest.java +++ b/java/src/test/java/org/lance/CompactionTest.java @@ -18,8 +18,10 @@ import org.lance.compaction.CompactionMode; import org.lance.compaction.CompactionOptions; import org.lance.compaction.CompactionPlan; +import org.lance.compaction.CompactionScope; import org.lance.compaction.CompactionTask; import org.lance.compaction.RewriteResult; +import org.lance.schema.SqlExpressions; import org.apache.arrow.memory.RootAllocator; import org.junit.jupiter.api.Test; @@ -33,11 +35,14 @@ import java.io.ObjectInputStream; import java.io.ObjectOutputStream; import java.nio.file.Path; +import java.util.ArrayList; import java.util.Arrays; import java.util.Base64; import java.util.Collections; +import java.util.List; import java.util.Optional; +import static org.junit.jupiter.api.Assertions.assertArrayEquals; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; @@ -268,6 +273,65 @@ public void testCompactionModeRoundTrip(CompactionMode mode, @TempDir Path tempD } } + @Test + public void testRepackColumns(@TempDir Path tempDir) throws Exception { + String datasetPath = tempDir.resolve("test_repack_columns").toString(); + try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + TestUtils.SimpleTestDataset testDataset = + new TestUtils.SimpleTestDataset(allocator, datasetPath); + testDataset.createEmptyDataset().close(); + testDataset.write(1, 10).close(); + try (Dataset dataset = testDataset.write(2, 10)) { + dataset.addColumns( + new SqlExpressions.Builder().withExpression("double_id", "id * 2").build(), + Optional.empty()); + assertEquals(2, dataset.getFragments().get(0).metadata().getFiles().size()); + ColumnLayoutStatistics layout = dataset.getColumnLayoutStatistics(); + assertEquals(2, layout.size()); + assertArrayEquals(new long[] {0, 1}, layout.getFragmentIds()); + assertArrayEquals(new int[] {2, 2}, layout.getLiveFileCounts()); + assertArrayEquals(new int[] {2, 1}, layout.getFieldsPerFile()[0]); + assertEquals(2, layout.getFileSizes()[0].length); + assertTrue(layout.getFileSizes()[0][0] > 0); + assertArrayEquals(new double[] {0.0, 0.0}, layout.getTombstonedFieldRatios()); + assertArrayEquals(new int[] {0, 0}, layout.getOverlayCounts()); + + CompactionOptions options = + CompactionOptions.builder() + .withMaxDataFilesPerFragment(1) + .withScope(CompactionScope.REPACK_COLUMNS) + .build(); + CompactionPlan plan = Compaction.planCompaction(dataset, options); + assertEquals(Optional.of(1L), plan.getCompactionOptions().getMaxDataFilesPerFragment()); + assertEquals( + Optional.of(CompactionScope.REPACK_COLUMNS.getValue()), + plan.getCompactionOptions().getScope()); + assertEquals(2, plan.getCompactionTasks().size()); + + List results = new ArrayList<>(); + for (CompactionTask task : plan.getCompactionTasks()) { + task = serializeAndDeserialize(task); + assertEquals( + Collections.singletonList(Arrays.asList(0, 1, 2)), + task.getTaskData().getRepackFiles()); + RewriteResult result = serializeAndDeserialize(task.execute(dataset)); + assertEquals(1, result.getRepackedFiles().size()); + results.add(result); + } + Compaction.commitCompaction(dataset, results, plan.getCompactionOptions()); + + dataset.checkoutLatest(); + assertEquals(2, dataset.getFragments().size()); + for (Fragment fragment : dataset.getFragments()) { + assertEquals(1, fragment.metadata().getFiles().size()); + } + ColumnLayoutStatistics repacked = dataset.getColumnLayoutStatistics(); + assertArrayEquals(new int[] {1, 1}, repacked.getLiveFileCounts()); + assertArrayEquals(new int[] {3}, repacked.getFieldsPerFile()[0]); + } + } + } + /** * A serialized CompactionOptions produced by the class as it existed before maxSourceRows and * maxSourceBytes were added (no declared serialVersionUID, stream ends after maxSourceFragments), diff --git a/protos/transaction.proto b/protos/transaction.proto index 6e01831b4af..75c164f527a 100644 --- a/protos/transaction.proto +++ b/protos/transaction.proto @@ -364,7 +364,16 @@ message Transaction { // An operation that replaces the data in a region of the table with new data. message DataReplacement { + /* A fragment can be the target of several groups, one per new file. The + * groups apply in order, each to the fragment as the previous ones left it. + */ repeated DataReplacementGroup replacements = 1; + /* Whether the new files change any value. Absent means true, which is what + * a writer that predates this field meant. False means the values were only + * moved to new files: index coverage, overlays, and row update versions are + * left as they are. + */ + optional bool data_change = 2; } /* Overlay files to append to a single fragment, in order (the last entry is diff --git a/python/python/lance/dataset.py b/python/python/lance/dataset.py index 13af2bbe19f..aa45b6fefb2 100644 --- a/python/python/lance/dataset.py +++ b/python/python/lance/dataset.py @@ -6603,9 +6603,15 @@ class DataReplacementGroup: class DataReplacement(BaseOperation): """ Operation that replaces existing datafiles in the dataset. + + ``data_change=False`` declares that the new files hold the same values + as the files they replace (compaction moved them). Indices then keep + their coverage, overlays keep shadowing, and no row is reported as + updated. """ replacements: List[LanceOperation.DataReplacementGroup] + data_change: bool = True @dataclass class DataOverlayFile: @@ -7607,6 +7613,9 @@ def compact_files( max_source_bytes: Optional[int] = None, excluded_fragment_ids: Optional[list[int]] = None, data_storage_version: Optional[str] = None, + max_data_files_per_fragment: Optional[int] = None, + column_groups: Optional[list[list[str]]] = None, + scope: Optional[Literal["all", "rewrite_fragments", "repack_columns"]] = None, ) -> CompactionMetrics: """Compacts small files in the dataset, reducing total number of files. @@ -7614,6 +7623,8 @@ def compact_files( * Removes deleted rows from fragments * Removes dropped columns from fragments * Merges small fragments into larger ones + * Repacks a fragment's columns into fewer data files, when + ``max_data_files_per_fragment`` or ``column_groups`` asks for it This method preserves the insertion order of the dataset. This may mean it leaves small fragments in the dataset if they are not adjacent to @@ -7641,6 +7652,8 @@ def compact_files( ``lance.compaction.max_source_fragments``, ``lance.compaction.max_source_rows``, ``lance.compaction.max_source_bytes``, + ``lance.compaction.max_data_files_per_fragment``, + ``lance.compaction.column_groups``, ``lance.compaction.data_storage_version``. Parameters @@ -7722,6 +7735,24 @@ def compact_files( Uses the compaction config target when set, otherwise the dataset's default write version. Does not change that default or the versions of unselected files. V1/V2 cross-family targets are rejected. + max_data_files_per_fragment: int, optional + Maximum number of data files a fragment may hold columns in before + its columns are repacked into fewer files. Each ``add_columns`` + backfill adds one file per fragment. A repack rewrites only the + columns that move and keeps rows, fragment ids and indices as they + are. If not specified, uses the manifest config value, or no limit. + column_groups: list[list[str]], optional + Top-level columns to keep in their own data files. Each inner list + becomes one data file per fragment; the columns no group names share + one file. Fragments the compaction rewrites are written this way, + and the others have their columns repacked to match. Names that are + not top-level columns are ignored. If not specified, uses the + manifest config value (``"b, c; d"`` is ``[["b", "c"], ["d"]]``). + scope: str, optional + Which tasks to plan: ``"rewrite_fragments"`` only rewrites + fragments, ``"repack_columns"`` only repacks columns, and + ``"all"`` (the default) does both. A run that does both commits + two versions. Returns ------- @@ -7750,6 +7781,9 @@ def compact_files( max_source_bytes=max_source_bytes, excluded_fragment_ids=excluded_fragment_ids, data_storage_version=data_storage_version, + max_data_files_per_fragment=max_data_files_per_fragment, + column_groups=column_groups, + scope=scope, ).items() if v is not None } @@ -7998,6 +8032,24 @@ class DatasetStats(TypedDict): num_small_files: int +class FragmentColumnLayoutStats(TypedDict): + fragment_id: int + #: Data files holding at least one column of the schema. This is the count + #: ``max_data_files_per_fragment`` in compaction is compared with. + live_file_count: int + #: Recorded size in bytes of each data file, in the fragment's file order; + #: ``None`` when the manifest has no size for the file. + file_sizes: List[Optional[int]] + #: Number of schema fields each data file holds, in the fragment's file order. + fields_per_file: List[int] + #: Share of the fragment's field slots holding no live data (tombstoned by a + #: column update or repack, or left by a dropped column). Spilled row + #: lineage is not counted. A fragment rewrite reclaims these slots. + tombstoned_field_ratio: float + #: Data overlay files attached to the fragment. + overlay_count: int + + class LanceStats: """ Statistics about a LanceDataset. @@ -8034,6 +8086,17 @@ def data_stats(self) -> DataStatistics: """ return self._ds.data_stats() + def column_layout_stats(self) -> List[FragmentColumnLayoutStats]: + """ + How each fragment's columns are laid out across data files, in + manifest fragment order. + + Compaction repacks a fragment's columns into fewer files when + ``live_file_count`` is above ``max_data_files_per_fragment``. Only + manifest metadata is read. + """ + return self._ds.column_layout_stats() + def write_dataset( data_obj: ReaderLike, diff --git a/python/python/lance/lance/__init__.pyi b/python/python/lance/lance/__init__.pyi index f338f16c9c7..25ba1e4fc24 100644 --- a/python/python/lance/lance/__init__.pyi +++ b/python/python/lance/lance/__init__.pyi @@ -503,6 +503,7 @@ class _Dataset: data_storage_version: Optional[str] = None, ) -> UpdateResult: ... def count_deleted_rows(self) -> int: ... + def column_layout_stats(self) -> List[Dict[str, Any]]: ... def versions(self) -> List[Version]: ... def version_refs(self) -> List[VersionRef]: ... def version(self) -> int: ... diff --git a/python/python/lance/lance/optimize.pyi b/python/python/lance/lance/optimize.pyi index c4b6b6546e6..f4949a5fe0a 100644 --- a/python/python/lance/lance/optimize.pyi +++ b/python/python/lance/lance/optimize.pyi @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional +from typing import List, Literal, Optional from lance import LanceDataset from lance.fragment import FragmentMetadata @@ -32,6 +32,7 @@ class RewriteResult: class CompactionTask: read_version: int + kind: Literal["rewrite_fragments", "repack_columns"] fragments: List["FragmentMetadata"] def execute(self, dataset: "LanceDataset") -> RewriteResult: ... diff --git a/python/python/lance/optimize.py b/python/python/lance/optimize.py index ab18255c853..cff80ed61a7 100644 --- a/python/python/lance/optimize.py +++ b/python/python/lance/optimize.py @@ -121,6 +121,28 @@ class CompactionOptions(TypedDict, total=False): are not combined into the same task. Duplicate and unknown IDs are ignored. (default: None) """ + max_data_files_per_fragment: Optional[int] + """ + Maximum number of data files a fragment may hold columns in before its + columns are repacked into fewer files. Each ``add_columns`` backfill adds + one file per fragment. A repack rewrites only the columns that move and + keeps rows, fragment ids and indices as they are. (default: None, no + file-count trigger) + """ + column_groups: Optional[list[list[str]]] + """ + Top-level columns to keep in their own data files. Each inner list becomes + one data file per fragment; the columns no group names share one file. + Fragments the compaction rewrites are written this way, and the others + have their columns repacked to match. Names that are not top-level + columns are ignored. (default: None, one file per fragment) + """ + scope: Optional[Literal["all", "rewrite_fragments", "repack_columns"]] + """ + Which tasks to plan: ``"rewrite_fragments"`` only rewrites fragments, + ``"repack_columns"`` only repacks columns, ``"all"`` does both. + (default: "all") + """ data_storage_version: Optional[str] """ Output data file version, such as "2.2", "stable", or "next". If omitted, diff --git a/python/python/tests/test_optimize.py b/python/python/tests/test_optimize.py index 34199956c1e..f8ccd73e601 100644 --- a/python/python/tests/test_optimize.py +++ b/python/python/tests/test_optimize.py @@ -157,6 +157,68 @@ def test_compact_files_max_source_fragments(tmp_path: Path): assert len(dataset.get_fragments()) == 7 +def _backfilled(tmp_path: Path): + dataset = lance.write_dataset( + pa.table({"a": range(8), "b": range(8)}), + tmp_path / "dataset", + max_rows_per_file=4, + ) + dataset.add_columns({"c": "a + 1"}) + dataset.add_columns({"d": "a + 2"}) + return dataset + + +def _files_per_fragment(dataset): + return [len(fragment.metadata.files) for fragment in dataset.get_fragments()] + + +def test_compact_files_repacks_columns(tmp_path: Path): + dataset = _backfilled(tmp_path) + expected = dataset.to_table() + assert _files_per_fragment(dataset) == [3, 3] + stats = dataset.stats.column_layout_stats() + assert [s["fragment_id"] for s in stats] == [0, 1] + for s in stats: + assert s["live_file_count"] == 3 + assert s["fields_per_file"] == [2, 1, 1] + assert all(size is not None for size in s["file_sizes"]) + assert s["tombstoned_field_ratio"] == 0.0 + assert s["overlay_count"] == 0 + + metrics = dataset.optimize.compact_files( + max_data_files_per_fragment=1, scope="repack_columns" + ) + + assert metrics.files_added == 2 + assert metrics.fragments_added == 0 + assert _files_per_fragment(dataset) == [1, 1] + stats = dataset.stats.column_layout_stats() + assert [s["live_file_count"] for s in stats] == [1, 1] + assert [s["fields_per_file"] for s in stats] == [[4], [4]] + assert dataset.to_table() == expected + assert [f.fragment_id for f in dataset.get_fragments()] == [0, 1] + + +def test_distributed_repack(tmp_path: Path): + dataset = _backfilled(tmp_path) + expected = dataset.to_table() + + # c and d sit in separate files, so the group moves them into one. + plan = Compaction.plan( + dataset, options=dict(column_groups=[["c", "d"]], scope="repack_columns") + ) + assert [task.kind for task in plan.tasks] == ["repack_columns"] * 2 + results = [ + pickle.loads(pickle.dumps(pickle.loads(pickle.dumps(task)).execute(dataset))) + for task in plan.tasks + ] + Compaction.commit(dataset, results) + + dataset = lance.dataset(dataset.uri) + assert _files_per_fragment(dataset) == [2, 2] + assert dataset.to_table() == expected + + def test_blob_compaction(tmp_path: Path): base_dir = tmp_path / "blob_dataset" blob_field = pa.field( diff --git a/python/src/dataset.rs b/python/src/dataset.rs index eca9c2e756d..47663948771 100644 --- a/python/src/dataset.rs +++ b/python/src/dataset.rs @@ -2103,6 +2103,24 @@ impl Dataset { .map_err(|err| PyIOError::new_err(err.to_string())) } + /// Per-fragment column-layout stats, in manifest fragment order. + fn column_layout_stats(&self, py: Python<'_>) -> PyResult>> { + self.ds + .column_layout_stats() + .into_iter() + .map(|stats| { + let dict = PyDict::new(py); + dict.set_item("fragment_id", stats.fragment_id)?; + dict.set_item("live_file_count", stats.live_file_count)?; + dict.set_item("file_sizes", stats.file_sizes)?; + dict.set_item("fields_per_file", stats.fields_per_file)?; + dict.set_item("tombstoned_field_ratio", stats.tombstoned_field_ratio)?; + dict.set_item("overlay_count", stats.overlay_count)?; + Ok(dict.unbind()) + }) + .collect() + } + #[pyo3(signature=(new_bases, transaction_properties=None))] fn add_bases( &mut self, diff --git a/python/src/dataset/optimize.rs b/python/src/dataset/optimize.rs index e17023b7450..2fb2e2578c4 100644 --- a/python/src/dataset/optimize.rs +++ b/python/src/dataset/optimize.rs @@ -15,8 +15,9 @@ use lance::dataset::{ index::DatasetIndexRemapperOptions, optimize::{ - CompactionMetrics, CompactionMode, CompactionOptions, CompactionPlan, CompactionTask, - RewriteResult, commit_compaction, compact_files, plan_compaction, + CompactionMetrics, CompactionMode, CompactionOptions, CompactionPlan, CompactionScope, + CompactionTask, CompactionTaskKind, RewriteResult, commit_compaction, compact_files, + plan_compaction, }, }; use pyo3::{exceptions::PyNotImplementedError, pyclass::CompareOp, types::PyTuple}; @@ -92,6 +93,21 @@ fn parse_compaction_options( opts.data_storage_version = Some(version.parse().infer_error()?); } } + "max_data_files_per_fragment" => { + opts.max_data_files_per_fragment = value.extract()?; + } + "column_groups" => { + opts.column_groups = value + .extract::>>>()? + .unwrap_or_default(); + } + "scope" => { + let scope: Option = value.extract()?; + if let Some(scope) = scope { + opts.scope = CompactionScope::try_from(scope.as_str()) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + } + } _ => { return Err(PyValueError::new_err(format!( "Invalid compaction option: {}", @@ -129,8 +145,9 @@ pub struct PyCompactionMetrics { /// int : The number of files that have been removed, including deletion files. #[pyo3(get)] pub files_removed: usize, - /// int : The number of files that have been added, which is always equal to the - /// number of fragments. + /// int : The number of data files that have been added. A rewrite adds one + /// per new fragment (one per group with ``column_groups``), a column repack + /// one per new file. #[pyo3(get)] pub files_added: usize, } @@ -263,8 +280,10 @@ impl PyCompactionTask { .collect::>>()? .join(", "); Ok(format!( - "CompactionTask(read_version={}, fragments=[{}])", - self.0.read_version, fragment_reprs + "CompactionTask(read_version={}, kind={}, fragments=[{}])", + self.0.read_version, + self.kind(), + fragment_reprs )) } @@ -274,6 +293,16 @@ impl PyCompactionTask { self.0.read_version } + /// str : ``"rewrite_fragments"`` for a task that rewrites its fragments, + /// ``"repack_columns"`` for one that repacks its fragment's columns. + #[getter] + pub fn kind(&self) -> &'static str { + match self.0.task.kind { + CompactionTaskKind::RewriteFragments => "rewrite_fragments", + CompactionTaskKind::RepackColumns { .. } => "repack_columns", + } + } + /// List[lance.fragment.FragmentMetadata] : The fragments that will be compacted. #[getter] pub fn fragments<'py>(&self, py: Python<'py>) -> PyResult>> { diff --git a/python/src/transaction.rs b/python/src/transaction.rs index b3f2311b3f4..267353d0f0d 100644 --- a/python/src/transaction.rs +++ b/python/src/transaction.rs @@ -583,8 +583,12 @@ impl FromPyObject<'_, '_> for PyLance { } "DataReplacement" => { let replacements = extract_vec(&ob.getattr("replacements")?)?; + let data_change = ob.getattr("data_change")?.extract::()?; - let op = Operation::DataReplacement { replacements }; + let op = Operation::DataReplacement { + replacements, + data_change, + }; Ok(Self(op)) } @@ -751,12 +755,15 @@ impl<'py> IntoPyObject<'py> for PyLance<&Operation> { updated_fragment_offsets, )) } - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { + replacements, + data_change, + } => { let replacements = export_vec(py, replacements.as_slice())?; let cls = namespace .getattr("DataReplacement") .expect("Failed to get DataReplacement class"); - cls.call1((replacements,)) + cls.call1((replacements, *data_change)) } Operation::DataOverlay { groups } => { let groups = export_vec(py, groups.as_slice())?; diff --git a/rust/lance-table/src/system_index/frag_reuse/gate.rs b/rust/lance-table/src/system_index/frag_reuse/gate.rs index 42d953532c7..04ddbbc3b81 100644 --- a/rust/lance-table/src/system_index/frag_reuse/gate.rs +++ b/rust/lance-table/src/system_index/frag_reuse/gate.rs @@ -373,6 +373,7 @@ mod tests { None, ), )], + data_change: true, }, "data_overlay" => Operation::DataOverlay { groups: vec![DataOverlayGroup { diff --git a/rust/lance-table/src/transaction/conflicts.rs b/rust/lance-table/src/transaction/conflicts.rs index cba1b5fb547..47a8c868621 100644 --- a/rust/lance-table/src/transaction/conflicts.rs +++ b/rust/lance-table/src/transaction/conflicts.rs @@ -198,9 +198,15 @@ impl PartialEq for Operation { && a_field == b_field } ( - Self::DataReplacement { replacements: a }, - Self::DataReplacement { replacements: b }, - ) => a.len() == b.len() && a.iter().all(|r| b.contains(r)), + Self::DataReplacement { + replacements: a, + data_change: a_change, + }, + Self::DataReplacement { + replacements: b, + data_change: b_change, + }, + ) => a_change == b_change && a.len() == b.len() && a.iter().all(|r| b.contains(r)), // Handle all remaining combinations. // We spell out all combinations explicitly to prevent // us accidentally handling a new case in the wrong way. diff --git a/rust/lance-table/src/transaction/index_maintenance.rs b/rust/lance-table/src/transaction/index_maintenance.rs index 6c43eb58e51..63471ec29db 100644 --- a/rust/lance-table/src/transaction/index_maintenance.rs +++ b/rust/lance-table/src/transaction/index_maintenance.rs @@ -98,9 +98,9 @@ impl Transaction { /// id, unexpanded: an update with `fields_modified` rewrites those /// fields in every updated fragment; a merge rewrites, in each fragment /// present in `previous_fragments`, the fields whose backing data file - /// changed (`merge_rewritten_fields`); a data replacement rewrites the - /// fields its new files carry, read through `schema`. Any other - /// operation rewrites nothing. `previous_fragments` is the caller's + /// changed (`merge_rewritten_fields`); a data replacement that changes + /// data rewrites the fields its new files carry, read through `schema`. + /// Any other operation rewrites nothing. `previous_fragments` is the caller's /// "before" list: the current manifest's for a commit, the read /// version's for a rebase. pub fn rewritten_physical_columns( @@ -120,7 +120,11 @@ impl Transaction { Operation::Merge { fragments, .. } => { Self::merge_rewritten_fields(previous_fragments, fragments) } - Operation::DataReplacement { replacements } => replacements + // Values that only moved to new files invalidate no index. + Operation::DataReplacement { + replacements, + data_change: true, + } => replacements .iter() .map(|DataReplacementGroup(fragment_id, new_file)| { let mut fields: Vec = new_file diff --git a/rust/lance-table/src/transaction/manifest_build.rs b/rust/lance-table/src/transaction/manifest_build.rs index ee16618fb1f..adb7351e622 100644 --- a/rust/lance-table/src/transaction/manifest_build.rs +++ b/rust/lance-table/src/transaction/manifest_build.rs @@ -1247,7 +1247,10 @@ impl Transaction { Operation::Restore { .. } => { unreachable!() } - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { + replacements, + data_change, + } => { log::warn!( "Building manifest with DataReplacement operation. This operation is not stable yet, please use with caution." ); @@ -1257,56 +1260,24 @@ impl Transaction { .map(|DataReplacementGroup(fragment_id, new_file)| (fragment_id, new_file)) .unzip(); - // 1. make sure the new files all have the same fields / or empty - // NOTE: arguably this requirement could be relaxed in the future - // for the sake of simplicity, we require the new files to have the same fields - if new_datafiles - .iter() - .map(|f| f.fields.clone()) - .collect::>() - .len() - > 1 - { - let field_info = new_datafiles - .iter() - .enumerate() - .map(|(id, f)| (id, f.fields.clone())) - .fold("".to_string(), |acc, (id, fields)| { - format!("{}File {}: {:?}\n", acc, id, fields) - }); - - return Err(Error::invalid_input(format!( - "All new data files must have the same fields, but found different fields:\n{field_info}" - ))); - } - let existing_fragments = maybe_existing_fragments?; - // Collect replaced field IDs before consuming new_datafiles - let replaced_fields: Vec = new_datafiles - .first() - .map(|f| { - f.schema(&schema) - .field_ids() - .iter() - .chain(f.fields.iter()) - .filter(|&&id| id >= 0) - .map(|&id| id as u32) - .collect() - }) - .unwrap_or_default(); - - // 2. check that the fragments being modified have isomorphic layouts along the columns being replaced - // 3. add modified fragments to final_fragments + // Apply the groups in order. Several groups may name the same + // fragment (one new file each), so each applies to the + // fragment as the previous ones left it. + let mut replaced: HashMap = HashMap::new(); for (frag_id, new_file) in old_fragment_ids.iter().zip(new_datafiles) { - let frag = existing_fragments - .iter() - .find(|f| f.id == **frag_id) - .ok_or_else(|| { - Error::invalid_input( - "Fragment being replaced not found in existing fragments", - ) - })?; + let frag = match replaced.get(*frag_id) { + Some(frag) => frag, + None => existing_fragments + .iter() + .find(|f| f.id == **frag_id) + .ok_or_else(|| { + Error::invalid_input( + "Fragment being replaced not found in existing fragments", + ) + })?, + }; let mut new_frag = frag.clone(); // Physical mappings differ across V2 encodings (a nested @@ -1425,19 +1396,31 @@ impl Transaction { // value though -- the conflict resolver rebases these two // precisely because the overlay wins -- so it stays, and // being newer it stays last, preserving the ordering. - let (mut superseded, newer): (Vec<_>, Vec<_>) = new_frag - .overlays - .drain(..) - .partition(|overlay| overlay.committed_version <= self.read_version); - crate::format::overlay::tombstone_overlay_fields( - &mut superseded, - &replaced_fields, - ); - superseded.extend(newer); - new_frag.overlays = superseded; + // Values that only moved supersede nothing. + if *data_change { + let replaced_fields: Vec = new_file + .schema(&schema) + .field_ids() + .iter() + .chain(new_file.fields.iter()) + .filter(|&&id| id >= 0) + .map(|&id| id as u32) + .collect(); + let (mut superseded, newer): (Vec<_>, Vec<_>) = new_frag + .overlays + .drain(..) + .partition(|overlay| overlay.committed_version <= self.read_version); + crate::format::overlay::tombstone_overlay_fields( + &mut superseded, + &replaced_fields, + ); + superseded.extend(newer); + new_frag.overlays = superseded; + } - final_fragments.push(new_frag); + replaced.insert(**frag_id, new_frag); } + final_fragments.extend(replaced.into_values()); let fragments_changed = old_fragment_ids .iter() @@ -1457,7 +1440,7 @@ impl Transaction { // A replacement changes what its rows read as, so stamp them // updated. Without this, get_updated_rows never reports them and // an incremental consumer skips them for good. - if next_row_id.is_some() { + if next_row_id.is_some() && *data_change { let new_version = current_manifest.map_or(1, |m| m.version + 1); for fragment in final_fragments .iter_mut() @@ -1470,8 +1453,8 @@ impl Transaction { } } - // The replaced fields' coverage of the modified fragments was - // withdrawn by `prepare_indices`. + // When the data changed, the replaced fields' coverage of the + // modified fragments was withdrawn by `prepare_indices`. } Operation::DataOverlay { groups } => { // Stamp each overlay with the version this commit is producing. @@ -2160,6 +2143,7 @@ mod tests { None, ), )], + data_change: true, }, _ => unreachable!(), } @@ -2867,6 +2851,33 @@ mod tests { ); } + /// A replacement whose values only moved to new files keeps every index's + /// coverage of the replaced fragment. + #[test] + fn prepare_indices_keeps_coverage_for_moved_values() { + let mut manifest = manifest_with_file("a.lance"); + manifest.reader_feature_flags &= !FLAG_FRAGMENT_REUSE_INDEX; + manifest.writer_feature_flags &= !FLAG_FRAGMENT_REUSE_INDEX; + let mut segment = sample_index_metadata("id_idx"); + segment.fragment_bitmap = Some([0u32, 3].into_iter().collect()); + let mut operation = in_place_rewrite("data_replacement", &manifest, vec![0]); + let Operation::DataReplacement { data_change, .. } = &mut operation else { + unreachable!() + }; + *data_change = false; + let transaction = Transaction::new(manifest.version, operation, None); + let prepared = transaction + .prepare_indices( + Some(&manifest), + vec![segment.clone()], + &default_build_config(), + None, + FragReuseUpdate::None, + ) + .unwrap(); + assert_eq!(prepared.prepared(), &[segment]); + } + /// Operations that rewrite no column in place, `Restore` included, /// prepare the list unchanged. #[test] @@ -4347,6 +4358,7 @@ mod tests { 0, DataFile::new_legacy_from_fields("f5-new.lance", vec![5], None), )], + data_change: true, }, None, ); @@ -4364,6 +4376,109 @@ mod tests { assert_eq!(frag.overlays[0].data_file.fields.as_ref(), &[3, -2]); } + #[test] + fn test_data_replacement_without_data_change_keeps_overlays() { + // Values that only moved to a new file supersede no overlay: both + // overlays keep shadowing the replaced field. + let mut fragment = Fragment::new(0); + fragment.files = vec![ + DataFile::new_legacy_from_fields("f3.lance", vec![3], None), + DataFile::new_legacy_from_fields("f5.lance", vec![5], None), + ]; + let overlay = |path: &str, fields: Vec| DataOverlayFile { + data_file: DataFile::new_legacy_from_fields(path, fields, None), + coverage: OverlayCoverage::dense(roaring::RoaringBitmap::from_iter([0u32])), + committed_version: 1, + }; + fragment.overlays = vec![overlay("o3.lance", vec![3]), overlay("o5.lance", vec![5])]; + let overlays = fragment.overlays.clone(); + + let schema = ArrowSchema::new(vec![ArrowField::new("id", DataType::Int32, false)]); + let manifest = Manifest::new( + LanceSchema::try_from(&schema).unwrap(), + Arc::new(vec![fragment]), + crate::format::DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + let txn = Transaction::new( + manifest.version, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("f5-new.lance", vec![5], None), + )], + data_change: false, + }, + None, + ); + let (result, _) = txn + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap(); + + let frag = &result.fragments[0]; + assert!(frag.files.iter().any(|f| f.path == "f5-new.lance")); + assert_eq!(frag.overlays, overlays); + } + + #[test] + fn test_data_replacement_applies_several_groups_to_one_fragment() { + // Two groups repack four single-field files of one fragment into two + // files. Each group applies to the fragment as the previous one left + // it, and the files left holding no schema field are dropped. + let schema = ArrowSchema::new(vec![ + ArrowField::new("x", DataType::Int32, true), + ArrowField::new("a", DataType::Int32, true), + ArrowField::new("v", DataType::Int32, true), + ArrowField::new("y", DataType::Int32, true), + ]); + let lance_schema = LanceSchema::try_from(&schema).unwrap(); + let ids: Vec = lance_schema.fields.iter().map(|field| field.id).collect(); + let file = |path: &str, fields: Vec| { + let indices = (0..fields.len() as i32).collect(); + DataFile::new(path, fields, indices, ConcreteFileVersion::V2_0, None, None) + }; + let mut fragment = Fragment::new(0); + fragment.files = ids + .iter() + .map(|id| file(&format!("f{id}.lance"), vec![*id])) + .collect(); + let manifest = Manifest::new( + lance_schema, + Arc::new(vec![fragment, Fragment::new(1)]), + crate::format::DataStorageFormat::new(ConcreteFileVersion::V2_0), + HashMap::new(), + ); + + let txn = Transaction::new( + manifest.version, + Operation::DataReplacement { + replacements: vec![ + DataReplacementGroup(0, file("g0.lance", vec![ids[0], ids[1]])), + DataReplacementGroup(0, file("g1.lance", vec![ids[2], ids[3]])), + ], + data_change: false, + }, + None, + ); + let (result, _) = txn + .build_manifest(Some(&manifest), vec![], "txn", &default_build_config()) + .unwrap(); + + assert_eq!(result.fragments.len(), 2, "the untouched fragment stays"); + let layout: Vec<(&str, Vec)> = result.fragments[0] + .files + .iter() + .map(|file| (file.path.as_str(), file.fields.to_vec())) + .collect(); + assert_eq!( + layout, + vec![ + ("g0.lance", vec![ids[0], ids[1]]), + ("g1.lance", vec![ids[2], ids[3]]), + ] + ); + } + /// Replace `fields` in `fragment` at `read_version`, against a manifest /// at `manifest_version` whose schema declares field ids 3 ("x"), 4 ("a"), /// 5 ("v") and 6 ("y"). @@ -4407,6 +4522,7 @@ mod tests { None, ), )], + data_change: true, }, None, ); diff --git a/rust/lance-table/src/transaction/operation.rs b/rust/lance-table/src/transaction/operation.rs index 09f73cf6155..f7121cc3aea 100644 --- a/rust/lance-table/src/transaction/operation.rs +++ b/rust/lance-table/src/transaction/operation.rs @@ -110,24 +110,25 @@ pub enum Operation { /// `groups`, and the manifest is the durable record of the history. frag_reuse_index: Option, }, - /// Replace data in a column in the dataset with new data. This is used for - /// null column population where we replace an entirely null column with a - /// new column that has data. + /// Replace the data files backing some fields of existing fragments with + /// new files, without moving rows. Each group names a fragment and one new + /// data file for it. At commit, against the fragment as it stands then, a + /// file holding exactly the new file's fields (in the same file version) is + /// swapped for it; otherwise the new file's fields are tombstoned where + /// they live and the new file is appended, and a file left holding no + /// field of the schema is dropped. A fragment can take several groups, one + /// per new file, applied in order. Used for null column population, and by + /// compaction to repack columns into fewer files. /// - /// This operation will only allow replacing files that contain the same schema - /// e.g. if the original files contain columns A, B, C and the new files contain - /// only columns A, B then the operation is not allowed. As we would need to split - /// the original files into two files, one with column A, B and the other with column C. - /// - /// Corollary to the above: the operation will also not allow replacing files unless the - /// affected columns all have the same datafile layout across the fragments being replaced. - /// - /// e.g. if fragments being replaced contain files with different schema layouts on - /// the column being replaced, the operation is not allowed. - /// say `frag_1: [A] [B, C]` and `frag_2: [A, B] [C]` and we are trying to replace column A - /// with a new column A, the operation is not allowed. + /// The fields of a legacy (V1) data file cannot be tombstoned one by one, + /// so a field held by a V1 file can only be replaced by an exact match. DataReplacement { replacements: Vec, + /// Whether the new files change any value. `false` means the values + /// were only moved to new files (a compaction repack): indices keep + /// their coverage of the replaced fields, overlays keep shadowing, and + /// no row is stamped as updated. + data_change: bool, }, /// Attach overlay files to fragments, supplying new values for a subset of /// `(physical offset, field)` cells without rewriting the fragments' base diff --git a/rust/lance-table/src/transaction/proto.rs b/rust/lance-table/src/transaction/proto.rs index 0e9e64853fe..6b3138f93aa 100644 --- a/rust/lance-table/src/transaction/proto.rs +++ b/rust/lance-table/src/transaction/proto.rs @@ -373,12 +373,16 @@ impl TryFrom for Transaction { } } Some(pb::transaction::Operation::DataReplacement( - pb::transaction::DataReplacement { replacements }, + pb::transaction::DataReplacement { + replacements, + data_change, + }, )) => Operation::DataReplacement { replacements: replacements .into_iter() .map(DataReplacementGroup::try_from) .collect::>>()?, + data_change: data_change.unwrap_or(true), }, Some(pb::transaction::Operation::UpdateMemWalState( pb::transaction::UpdateMemWalState { compacted_sstables }, @@ -673,14 +677,18 @@ impl From<&Transaction> for pb::Transaction { schema_metadata: Default::default(), field_metadata: Default::default(), }), - Operation::DataReplacement { replacements } => { - pb::transaction::Operation::DataReplacement(pb::transaction::DataReplacement { - replacements: replacements - .iter() - .map(pb::transaction::DataReplacementGroup::from) - .collect(), - }) - } + Operation::DataReplacement { + replacements, + data_change, + } => pb::transaction::Operation::DataReplacement(pb::transaction::DataReplacement { + replacements: replacements + .iter() + .map(pb::transaction::DataReplacementGroup::from) + .collect(), + // Written only when false, so a transaction file from a data + // changing replacement stays byte-identical to an older writer's. + data_change: (!data_change).then_some(false), + }), Operation::DataOverlay { groups } => { pb::transaction::Operation::DataOverlay(pb::transaction::DataOverlay { groups: groups @@ -882,4 +890,44 @@ mod tests { other => panic!("expected DataOverlay, got {other:?}"), } } + + /// `data_change` survives the round trip, and a transaction written + /// before the field existed reads back as a data change. + #[rstest::rstest] + #[case::moved(false)] + #[case::changed(true)] + fn test_data_replacement_data_change_roundtrips(#[case] data_change: bool) { + let transaction = Transaction::new( + 1, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("new.lance", vec![3], None), + )], + data_change, + }, + None, + ); + let mut message = pb::Transaction::from(&transaction); + let decoded = Transaction::try_from(message.clone()).unwrap(); + assert!(matches!( + decoded.operation, + Operation::DataReplacement { data_change: d, .. } if d == data_change + )); + + if let Some(pb::transaction::Operation::DataReplacement(replacement)) = + message.operation.as_mut() + { + assert_eq!(replacement.data_change, (!data_change).then_some(false)); + replacement.data_change = None; + } + let legacy = Transaction::try_from(message).unwrap(); + assert!(matches!( + legacy.operation, + Operation::DataReplacement { + data_change: true, + .. + } + )); + } } diff --git a/rust/lance/src/dataset.rs b/rust/lance/src/dataset.rs index e2b7cf94b45..0e1b5850f0f 100644 --- a/rust/lance/src/dataset.rs +++ b/rust/lance/src/dataset.rs @@ -69,6 +69,7 @@ pub(crate) mod blob; pub(crate) mod branch_location; pub mod builder; pub mod cleanup; +pub mod compaction_stats; mod data_file; mod data_file_part; pub mod delta; diff --git a/rust/lance/src/dataset/blob.rs b/rust/lance/src/dataset/blob.rs index 823adaa3fd8..da34547687d 100644 --- a/rust/lance/src/dataset/blob.rs +++ b/rust/lance/src/dataset/blob.rs @@ -6819,6 +6819,7 @@ mod tests { uuid: Uuid::new_v4().hyphenated().to_string(), operation: Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file)], + data_change: true, }, tag: None, transaction_properties: None, @@ -7002,6 +7003,7 @@ mod tests { uuid: Uuid::new_v4().hyphenated().to_string(), operation: Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file)], + data_change: true, }, tag: None, transaction_properties: None, diff --git a/rust/lance/src/dataset/compaction_stats.rs b/rust/lance/src/dataset/compaction_stats.rs new file mode 100644 index 00000000000..a690225aa41 --- /dev/null +++ b/rust/lance/src/dataset/compaction_stats.rs @@ -0,0 +1,243 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Per-fragment column-layout statistics. +//! +//! The default compaction planner repacks a fragment's columns when it holds +//! them in more than `max_data_files_per_fragment` files, counted here. +//! Repeated `add_columns` backfills are the usual cause. Exposing the counts +//! via [`Dataset::column_layout_stats`] lets a caller see why a fragment was, +//! or was not, picked, and lets a custom planner make its own call. + +use std::collections::HashSet; + +use lance_table::format::overlay::TOMBSTONE_FIELD_ID; + +use super::Dataset; + +/// Column-layout statistics for a single fragment. +#[derive(Debug, Clone, PartialEq)] +pub struct FragmentColumnLayoutStats { + /// The fragment these stats describe. + pub fragment_id: u64, + /// Number of data files holding at least one column of the dataset schema, + /// the count `max_data_files_per_fragment` is compared with. A file left + /// holding only tombstones or dropped columns, or kept only for the + /// fragment's spilled row lineage, holds no column a read touches and is + /// not counted. + pub live_file_count: usize, + /// Recorded size in bytes of each data file, in the fragment's file order. + /// `None` when the manifest has no size for the file. + pub file_sizes: Vec>, + /// Number of fields of the dataset schema each data file holds, in the + /// fragment's file order. A file counted in `live_file_count` holds at + /// least one. + pub fields_per_file: Vec, + /// The share of the fragment's field slots that hold no live data: slots + /// tombstoned by a column update or repack, or left by a dropped column. + /// The reserved ids of spilled row lineage are not counted on either side. + /// A fragment rewrite reclaims these slots. + pub tombstoned_field_ratio: f64, + /// Number of overlay files attached to the fragment. + pub overlay_count: usize, +} + +impl Dataset { + /// Per-fragment column-layout stats, in manifest fragment order. + /// + /// The default planner's input for column repacks. It reads only fragment + /// metadata (no data files), so it is cheap to call. + pub fn column_layout_stats(&self) -> Vec { + let schema_ids: HashSet = self + .schema() + .fields_pre_order() + .map(|field| field.id) + .collect(); + self.manifest + .fragments + .iter() + .map(|fragment| { + let fields_per_file: Vec = fragment + .files + .iter() + .map(|file| { + file.fields + .iter() + .filter(|id| schema_ids.contains(id)) + .count() + }) + .collect(); + // Field slots of user columns, live or dead; the other negative + // ids are spilled row lineage and are left out. + let user_slots = fragment + .files + .iter() + .flat_map(|file| file.fields.iter()) + .filter(|id| **id >= 0 || **id == TOMBSTONE_FIELD_ID) + .count(); + let live_slots: usize = fields_per_file.iter().sum(); + FragmentColumnLayoutStats { + fragment_id: fragment.id, + live_file_count: fields_per_file.iter().filter(|count| **count > 0).count(), + file_sizes: fragment + .files + .iter() + .map(|file| file.file_size_bytes.get().map(|size| size.get())) + .collect(), + fields_per_file, + tombstoned_field_ratio: if user_slots == 0 { + 0.0 + } else { + (user_slots - live_slots) as f64 / user_slots as f64 + }, + overlay_count: fragment.overlays.len(), + } + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dataset::{NewColumnTransform, WriteParams}; + use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator}; + use arrow_schema::{DataType, Field, Schema as ArrowSchema}; + use std::sync::Arc; + + #[tokio::test] + async fn column_layout_stats_counts_files_per_fragment() { + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "a", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..8))], + ) + .unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], schema); + // Two fragments (max_rows_per_file = 4), one data file each to start. + let mut dataset = Dataset::write( + reader, + "memory://", + Some(WriteParams { + max_rows_per_file: 4, + ..Default::default() + }), + ) + .await + .unwrap(); + + let stats = dataset.column_layout_stats(); + assert_eq!(stats.len(), 2); + assert!( + stats + .iter() + .all(|s| s.live_file_count == 1 && s.overlay_count == 0) + ); + + // add_columns appends a second data file per fragment. + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![("b".into(), "a + 1".into())]), + None, + None, + ) + .await + .unwrap(); + + let stats = dataset.column_layout_stats(); + assert_eq!(stats.len(), 2); + assert!( + stats.iter().all(|s| s.live_file_count == 2), + "each fragment should now have 2 data files: {stats:?}" + ); + for s in &stats { + assert_eq!(s.fields_per_file, vec![1, 1]); + assert!(s.file_sizes.iter().all(Option::is_some), "{s:?}"); + assert_eq!(s.tombstoned_field_ratio, 0.0); + } + } + + /// A dropped column leaves its field id in the file that held it, which + /// counts as a dead slot. + #[tokio::test] + async fn column_layout_stats_counts_dropped_columns_as_dead_slots() { + let schema = Arc::new(ArrowSchema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..4)), + Arc::new(Int32Array::from_iter_values(4..8)), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], schema); + let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + dataset.drop_columns(&["b"]).await.unwrap(); + + let stats = dataset.column_layout_stats(); + assert_eq!(stats[0].live_file_count, 1); + assert_eq!(stats[0].fields_per_file, vec![1]); + assert_eq!(stats[0].tombstoned_field_ratio, 0.5); + } + + /// Only files holding a schema column count. The extra files are added to + /// the manifest by hand: the commit paths that leave such files behind (an + /// in-place column update after a drop, a repack on a table whose lineage + /// was spilled into a data file) are not what this test is about. + #[tokio::test] + async fn column_layout_stats_skips_files_without_schema_columns() { + use lance_file::version::ConcreteFileVersion; + use lance_table::format::{DataFile, ROW_ID_FIELD_ID, RowIdMeta}; + + let schema = Arc::new(ArrowSchema::new(vec![Field::new( + "a", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(Int32Array::from_iter_values(0..4))], + ) + .unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], schema); + let mut dataset = Dataset::write(reader, "memory://", None).await.unwrap(); + + let dropped_id = dataset.schema().max_field_id().unwrap() + 1; + let mut manifest = dataset.manifest.as_ref().clone(); + let mut fragments = manifest.fragments.as_ref().clone(); + let file = |fields: Vec| { + let indices = (0..fields.len() as i32).collect(); + DataFile::new( + "extra.lance", + fields, + indices, + ConcreteFileVersion::V2_0, + None, + None, + ) + }; + fragments[0].files.push(file(vec![TOMBSTONE_FIELD_ID])); + fragments[0].files.push(file(vec![dropped_id])); + fragments[0] + .files + .push(file(vec![TOMBSTONE_FIELD_ID, ROW_ID_FIELD_ID])); + fragments[0].row_id_meta = Some(RowIdMeta::Column); + manifest.fragments = Arc::new(fragments); + dataset.manifest = Arc::new(manifest); + + let stats = dataset.column_layout_stats(); + assert_eq!(stats[0].live_file_count, 1, "{stats:?}"); + assert_eq!(stats[0].fields_per_file, vec![1, 0, 0, 0]); + assert!(stats[0].file_sizes[0].is_some()); + assert_eq!(stats[0].file_sizes[1..], [None, None, None]); + // `a`, two tombstones and the dropped id; the row id slot is lineage. + assert_eq!(stats[0].tombstoned_field_ratio, 0.75); + } +} diff --git a/rust/lance/src/dataset/optimize.rs b/rust/lance/src/dataset/optimize.rs index be0766f67d6..b1cda3dbc13 100644 --- a/rust/lance/src/dataset/optimize.rs +++ b/rust/lance/src/dataset/optimize.rs @@ -146,6 +146,7 @@ use tracing::{info, warn}; pub(super) mod binary_copy; pub mod remapping; +mod repack; use crate::index::frag_reuse::build_new_frag_reuse_index; use crate::io::deletion::read_dataset_deletion_file; @@ -316,6 +317,31 @@ pub struct CompactionOptions { /// carries any overlay, or `None` to disable the overlay-count trigger /// entirely. pub max_overlays_per_fragment: Option, + /// Maximum number of data files a fragment may hold columns in before its + /// columns are repacked into fewer files. Each `add_columns` backfill adds + /// one file per fragment; a large, deletion-free fragment never qualifies + /// for a rewrite, so this is the trigger that collapses those files. The + /// repack rewrites only the columns that move and keeps rows, fragment ids + /// and indices as they are (see [`CompactionTaskKind::RepackColumns`]). + /// The count is [`FragmentColumnLayoutStats::live_file_count`]. Only + /// files the repack empties are merged, at least two at a time. A file + /// stays when it holds a blob column, a column only partly in the + /// fragment's files, or spilled row lineage, or shares a column with a + /// file holding spilled row lineage; without `column_groups`, when every + /// file holding a column could go, at least one stays (the largest when + /// some merge leaves it) unless the limit is 1. So a fragment can stay above the limit; so can one with + /// more `column_groups` than the limit allows. + /// The repack writes each new file whole, with no `max_bytes_per_file` + /// split. Not planned under `ForceBinaryCopy`, since a repack reencodes. + /// + /// Defaults to `None` (no file-count trigger). Must be at least 1. + /// + /// [`FragmentColumnLayoutStats::live_file_count`]: crate::dataset::compaction_stats::FragmentColumnLayoutStats::live_file_count + #[serde(default)] + pub max_data_files_per_fragment: Option, + /// Which kinds of task this run plans. Defaults to [`CompactionScope::All`]. + #[serde(default)] + pub scope: CompactionScope, /// Exact data file version for compacted output. /// /// If omitted, use the dataset's default write version without changing it. @@ -331,6 +357,26 @@ pub struct CompactionOptions { /// }; /// ``` pub data_storage_version: Option, + /// Top-level columns to keep in their own data files. + /// + /// Each inner list becomes one data file per fragment holding exactly + /// those columns. Empty (the default) means no groups. + /// + /// A fragment a compaction rewrites gets one file per group plus one + /// shared file for the columns no group names. A fragment it leaves alone has a + /// group repacked into a file of its own when no single file holds + /// exactly that group's columns (see [`CompactionTaskKind::RepackColumns`]), + /// unless `scope` is [`CompactionScope::RewriteFragments`]; the columns no + /// group names are merged only when the fragment is over + /// `max_data_files_per_fragment`. A group holding a blob column is left + /// as it is. A name that is not a top-level column of the dataset (say, + /// one dropped or renamed since the groups were configured) is ignored + /// with a warning. Binary copy is disabled when groups are set + /// (`ForceBinaryCopy` is rejected), and `max_bytes_per_file` is ignored so + /// every group's files split at the same rows. Set from dataset config + /// with `lance.compaction.column_groups` (see [`Self::from_dataset_config`]). + #[serde(default)] + pub column_groups: Vec>, /// Transaction properties to store with this commit. /// /// These key-value pairs are stored in the transaction file @@ -364,7 +410,10 @@ impl Default for CompactionOptions { max_source_bytes: None, excluded_fragment_ids: Vec::new(), max_overlays_per_fragment: Some(10), + max_data_files_per_fragment: None, + scope: CompactionScope::All, data_storage_version: None, + column_groups: Vec::new(), transaction_properties: None, } } @@ -394,7 +443,12 @@ impl CompactionOptions { /// - `lance.compaction.max_source_rows` /// - `lance.compaction.max_source_bytes` /// - `lance.compaction.max_overlays_per_fragment` + /// - `lance.compaction.max_data_files_per_fragment` /// - `lance.compaction.data_storage_version` + /// - `lance.compaction.column_groups`: groups separated by `;`, columns in + /// a group by `,`, surrounding whitespace trimmed (`"b, c; d"` is + /// `[["b", "c"], ["d"]]`). A column whose name contains `,` or `;`, or + /// starts or ends with whitespace, cannot be named this way. pub fn from_dataset_config(config: &HashMap) -> Result { let mut opts = Self::default(); opts.apply_dataset_config(config)?; @@ -542,6 +596,28 @@ impl CompactionOptions { })?), }; } + "max_data_files_per_fragment" => { + self.max_data_files_per_fragment = Some(value.parse().map_err(|_| { + Error::invalid_input(format!( + "Invalid value for {}: '{}' (expected a positive integer)", + key, value + )) + })?); + } + "column_groups" => { + self.column_groups = value + .split(';') + .map(|group| { + group + .split(',') + .map(str::trim) + .filter(|column| !column.is_empty()) + .map(str::to_owned) + .collect::>() + }) + .filter(|group| !group.is_empty()) + .collect(); + } _ => { warn!("Ignoring unknown compaction config key: {}", key); } @@ -556,6 +632,32 @@ impl CompactionOptions { self.materialize_deletions = false; } + self.column_groups.retain(|group| !group.is_empty()); + let mut seen = HashSet::new(); + for column in self.column_groups.iter().flatten() { + if !seen.insert(column.as_str()) { + return Err(Error::invalid_input(format!( + "CompactionOptions::column_groups lists column \"{column}\" more than once" + ))); + } + } + if !self.column_groups.is_empty() + && matches!(self.compaction_mode(), CompactionMode::ForceBinaryCopy) + { + return Err(Error::invalid_input( + "CompactionOptions::column_groups cannot be combined with \ + compaction_mode=ForceBinaryCopy: binary copy keeps each fragment's \ + file layout, so it cannot split columns into groups", + )); + } + + if self.max_data_files_per_fragment == Some(0) { + return Err(Error::invalid_input( + "CompactionOptions::max_data_files_per_fragment must be at least 1 \ + (use None for no limit)", + )); + } + for (name, value) in [ ( "max_source_fragments", @@ -619,6 +721,10 @@ async fn can_use_binary_copy( options: &CompactionOptions, fragments: &[Fragment], ) -> bool { + if !options.column_groups.is_empty() { + log::debug!("Binary copy disabled: column_groups splits each fragment across files"); + return false; + } let version = options.write_version(dataset); versions::can_use_binary_copy(version, dataset, options, fragments) .await @@ -756,8 +862,9 @@ pub struct CompactionMetrics { pub fragments_added: usize, /// The number of files that have been removed, including deletion files. pub files_removed: usize, - /// The number of files that have been added, which is always equal to the - /// number of fragments. + /// The number of data files that have been added: one per new fragment for + /// a fragment rewrite (one per group with `column_groups`), and one per new + /// file for a column repack. pub files_added: usize, } @@ -787,6 +894,49 @@ pub trait CompactionPlanner: Send + Sync { async fn plan(&self, dataset: &Dataset) -> Result; } +/// Runs one task of a [`CompactionPlan`]. Pass an implementation to +/// [`compact_files_with_executor`] to run the tasks somewhere else (say, on a +/// cluster) or to wrap the default one; [`CompactionTask::execute`] runs the +/// default one for a task executed by hand. +/// +/// The result goes to [`commit_compaction`], which commits it by its kind. +#[async_trait::async_trait] +pub trait CompactionExecutor: Send + Sync { + /// Execute `task` against `dataset`, which is at the plan's read version. + async fn execute( + &self, + dataset: &Dataset, + task: TaskData, + options: &CompactionOptions, + ) -> Result; +} + +/// Runs each kind of task the way this crate does: a +/// [`CompactionTaskKind::RewriteFragments`] task rewrites its fragments into +/// new ones, a [`CompactionTaskKind::RepackColumns`] task writes new data +/// files for its fragment. +#[derive(Debug, Default, Clone, Copy)] +pub struct DefaultCompactionExecutor; + +#[async_trait::async_trait] +impl CompactionExecutor for DefaultCompactionExecutor { + async fn execute( + &self, + dataset: &Dataset, + task: TaskData, + options: &CompactionOptions, + ) -> Result { + match task.kind { + CompactionTaskKind::RewriteFragments => { + rewrite_files(Cow::Borrowed(dataset), task, options).await + } + CompactionTaskKind::RepackColumns { .. } => { + repack::execute_repack(dataset, task, options).await + } + } + } +} + /// Formulate a plan to compact the files in a dataset /// /// The compaction plan will contain a list of tasks to execute. Each task @@ -795,6 +945,9 @@ pub trait CompactionPlanner: Send + Sync { /// tasks may contain a single fragment when that fragment has deletions that /// are being materialized and doesn't have any neighbors that need to be /// compacted. +/// +/// Fragments no rewrite task takes are then checked for a column repack +/// against [`Dataset::column_layout_stats`], one task per fragment. #[derive(Debug, Clone, Default)] pub struct DefaultCompactionPlanner { options: CompactionOptions, @@ -820,6 +973,24 @@ impl CompactionPlanner for DefaultCompactionPlanner { dataset.manifest.data_storage_format.lance_file_format(), write_version, )?; + // Warned about once here; the tasks skip the stale names silently. + for column in self.options.column_groups.iter().flatten() { + if !dataset + .schema() + .fields + .iter() + .any(|field| &field.name == column) + { + warn!( + "Ignoring column_groups entry \"{column}\": it is not a top-level column \ + of the dataset" + ); + } + } + if !self.options.column_groups.is_empty() { + // Surface a bad column-group config here rather than mid-write. + column_group_schemas(dataset.schema(), &self.options.column_groups)?; + } if self.options.defer_index_remap && dataset.manifest.uses_stable_row_ids() { return Err(Error::invalid_input( "defer_index_remap=true is not supported on datasets with stable row IDs: \ @@ -829,6 +1000,33 @@ impl CompactionPlanner for DefaultCompactionPlanner { )); } + let mut all_tasks = if self.options.scope == CompactionScope::RepackColumns { + Vec::new() + } else { + self.plan_rewrites(dataset).await? + }; + if self.options.scope != CompactionScope::RewriteFragments { + let rewritten: HashSet = all_tasks + .iter() + .flat_map(|(task, _)| task.fragments.iter().map(|fragment| fragment.id)) + .collect(); + all_tasks.extend(self.plan_repacks(dataset, &rewritten)?); + } + + let tasks = limit_tasks_to_source_budget(&self.options, dataset.schema(), all_tasks)?; + + let mut options = self.options.clone(); + options.data_storage_version = Some(write_version.to_selector()); + let mut compaction_plan = CompactionPlan::new(dataset.manifest.version, options); + compaction_plan.extend_tasks(tasks); + + Ok(compaction_plan) + } +} + +impl DefaultCompactionPlanner { + /// The tasks that rewrite fragments, each with its live row count. + async fn plan_rewrites(&self, dataset: &Dataset) -> Result> { // get_fragments should be returning fragments in sorted order (by id) // and fragment ids should be unique let fragments = dataset.get_fragments(); @@ -1001,23 +1199,76 @@ impl CompactionPlanner for DefaultCompactionPlanner { .flat_map(|bin| bin.split_for_size(self.options.target_rows_per_fragment)) .map(|bin| { let live_rows = bin.row_counts.iter().sum(); - ( + (TaskData::rewrite_fragments(bin.fragments), live_rows) + }) + .collect(); + Ok(all_tasks) + } + + /// The tasks that repack the columns of a fragment not in `rewritten`, + /// each with its live row count. A fragment is repacked when its files + /// do not match `column_groups`, or when it holds columns in more than + /// `max_data_files_per_fragment` files, as + /// [`Dataset::column_layout_stats`] counts them. + fn plan_repacks( + &self, + dataset: &Dataset, + rewritten: &HashSet, + ) -> Result> { + let groups = if self.options.column_groups.is_empty() { + None + } else { + Some(claimed_column_ids( + dataset.schema(), + &self.options.column_groups, + )) + }; + let max_files = self.options.max_data_files_per_fragment; + if groups.is_none() && max_files.is_none() { + return Ok(Vec::new()); + } + // A repack reencodes the columns it moves. `validate()` already + // rejects `column_groups` with ForceBinaryCopy. + if matches!( + self.options.compaction_mode(), + CompactionMode::ForceBinaryCopy + ) { + log::info!( + "not planning column repacks: compaction_mode=ForceBinaryCopy does not \ + reencode, and a repack does" + ); + return Ok(Vec::new()); + } + let stats = dataset.column_layout_stats(); + let tasks = dataset + .manifest + .fragments + .iter() + .zip(stats) + .filter(|(fragment, _)| { + !rewritten.contains(&fragment.id) + && !u32::try_from(fragment.id) + .is_ok_and(|id| self.excluded_fragment_ids.contains(id)) + }) + .filter_map(|(fragment, stats)| { + let files = repack::plan_fragment_repack( + dataset.schema(), + fragment, + stats.live_file_count, + groups.as_deref(), + max_files, + )?; + let live_rows = fragment.num_rows().unwrap_or_default(); + Some(( TaskData { - fragments: bin.fragments, + fragments: vec![fragment.clone()], + kind: CompactionTaskKind::RepackColumns { files }, }, live_rows, - ) + )) }) .collect(); - - let tasks = limit_tasks_to_source_budget(&self.options, dataset.schema(), all_tasks)?; - - let mut options = self.options.clone(); - options.data_storage_version = Some(write_version.to_selector()); - let mut compaction_plan = CompactionPlan::new(dataset.manifest.version, options); - compaction_plan.extend_tasks(tasks); - - Ok(compaction_plan) + Ok(tasks) } } @@ -1028,6 +1279,14 @@ impl CompactionPlanner for DefaultCompactionPlanner { /// * Removes dropped columns from fragments. /// * Merges fragments that are too small. /// +/// With [`CompactionOptions::max_data_files_per_fragment`] or +/// [`CompactionOptions::column_groups`] set, it also repacks the columns of a +/// fragment spread over too many data files, or laid out other than the +/// groups ask, into new files without moving its rows. A run that both +/// rewrites fragments and repacks columns commits two versions (see +/// [`commit_compaction`]). [`CompactionOptions::scope`] limits a run to one +/// kind of task. +/// /// This method tries to preserve the insertion order of rows in the dataset. /// /// If no compaction is needed, this method will not make a new version of the table. @@ -1038,13 +1297,30 @@ pub async fn compact_files( ) -> Result { info!(target: TRACE_DATASET_EVENTS, event=DATASET_COMPACTING_EVENT, uri = &dataset.uri); let planner = DefaultCompactionPlanner::new(options)?; - compact_files_with_planner(dataset, remap_options, &planner).await + Box::pin(compact_files_with_planner(dataset, remap_options, &planner)).await } pub async fn compact_files_with_planner( dataset: &mut Dataset, remap_options: Option>, // These will be deprecated later planner: &dyn CompactionPlanner, +) -> Result { + Box::pin(compact_files_with_executor( + dataset, + remap_options, + planner, + &DefaultCompactionExecutor, + )) + .await +} + +/// Plan with `planner`, run every task with `executor`, and commit the +/// results with [`commit_compaction`]. +pub async fn compact_files_with_executor( + dataset: &mut Dataset, + remap_options: Option>, // These will be deprecated later + planner: &dyn CompactionPlanner, + executor: &dyn CompactionExecutor, ) -> Result { let compaction_plan: CompactionPlan = planner.plan(dataset).await?; @@ -1052,8 +1328,14 @@ pub async fn compact_files_with_planner( // tagged entry, which only the deferred-remap commit path does; eager // remapping would rewrite provenance the tagged reader depends on. // Checked before any file is rewritten (commit_compaction re-checks for - // callers that commit externally planned results). - if !compaction_plan.options.defer_index_remap + // callers that commit externally planned results). A repack moves no row, + // so only a plan that rewrites fragments is refused. + let rewrites_fragments = compaction_plan + .tasks + .iter() + .any(|task| task.kind == CompactionTaskKind::RewriteFragments); + if rewrites_fragments + && !compaction_plan.options.defer_index_remap && dataset.manifest.writer_feature_flags & lance_table::feature_flags::FLAG_FRAGMENT_REUSE_INDEX != 0 @@ -1074,28 +1356,21 @@ pub async fn compact_files_with_planner( return Ok(CompactionMetrics::default()); } - let dataset_ref = &dataset.clone(); - - let result_stream = futures::stream::iter(compaction_plan.tasks) - .map(|task| rewrite_files(Cow::Borrowed(dataset_ref), task, &compaction_plan.options)) - .buffer_unordered( - compaction_plan - .options - .num_threads - .unwrap_or_else(get_num_compute_intensive_cpus), - ); - - let completed_tasks: Vec = result_stream.try_collect().await?; + let concurrency = compaction_plan + .options + .num_threads + .unwrap_or_else(get_num_compute_intensive_cpus); let remap_options = remap_options.unwrap_or(Arc::new(DatasetIndexRemapperOptions::default())); - let metrics = commit_compaction( - dataset, - completed_tasks, - remap_options, - &compaction_plan.options, - ) - .await?; - - Ok(metrics) + let options = &compaction_plan.options; + // Execute against an immutable snapshot; the commit takes `&mut dataset` + // once every task has finished. + let snapshot = &dataset.clone(); + let results: Vec = futures::stream::iter(compaction_plan.tasks.clone()) + .map(|task| executor.execute(snapshot, task, options)) + .buffer_unordered(concurrency.max(1)) + .try_collect() + .await?; + commit_compaction(dataset, results, remap_options, options).await } /// Information about a fragment used to decide its fate in compaction @@ -1211,7 +1486,8 @@ fn limit_tasks_to_source_budget( total_fragments += task.fragments.len(); total_rows = total_rows.saturating_add(live_rows); if options.max_source_bytes.is_some() { - total_bytes = total_bytes.saturating_add(task_source_bytes(&task, &schema_field_ids)?); + total_bytes = + total_bytes.saturating_add(task_source_bytes(&task, schema, &schema_field_ids)?); } let over_budget = options @@ -1247,19 +1523,44 @@ fn limit_tasks_to_source_budget( /// /// Files whose fields are all absent from `schema_field_ids` only back /// dropped columns; compaction does not read them, so they are neither -/// counted nor required to have a recorded size. +/// counted nor required to have a recorded size. A column repack reads only +/// the files holding the columns it moves, and no overlay, so only those +/// files count. /// Only sizes recorded in the manifest are used: a missing size is an error /// rather than a metadata request against object storage, which would turn /// planning into one round trip per file. Deletion files are not counted. -fn task_source_bytes(task: &TaskData, schema_field_ids: &HashSet) -> Result { +fn task_source_bytes( + task: &TaskData, + schema: &lance_core::datatypes::Schema, + schema_field_ids: &HashSet, +) -> Result { + // The field ids a task reads: every schema field, or a repack's columns. + let read_ids: Cow<'_, HashSet> = match &task.kind { + CompactionTaskKind::RewriteFragments => Cow::Borrowed(schema_field_ids), + CompactionTaskKind::RepackColumns { files } => { + let columns: Vec = files.iter().flatten().copied().collect(); + Cow::Owned( + schema + .project_by_ids(&columns, true) + .field_ids() + .into_iter() + .collect(), + ) + } + }; + let reads_overlays = task.kind == CompactionTaskKind::RewriteFragments; let mut total_bytes = 0_u64; for fragment in &task.fragments { - let overlay_files = fragment.overlays.iter().map(|overlay| &overlay.data_file); + let overlay_files = fragment + .overlays + .iter() + .filter(|_| reads_overlays) + .map(|overlay| &overlay.data_file); for data_file in fragment.files.iter().chain(overlay_files) { if !data_file .fields .iter() - .any(|field_id| schema_field_ids.contains(field_id)) + .any(|field_id| read_ids.contains(field_id)) { continue; } @@ -2120,6 +2421,65 @@ async fn prepare_reader( pub struct TaskData { /// The fragments to compact. pub fragments: Vec, + /// What the task does with them. Absent in a task serialized before + /// repacking existed, which rewrote its fragments. + #[serde(default)] + pub kind: CompactionTaskKind, +} + +impl TaskData { + /// A task that rewrites `fragments` into new fragments. + pub fn rewrite_fragments(fragments: Vec) -> Self { + Self { + fragments, + kind: CompactionTaskKind::RewriteFragments, + } + } +} + +/// What a compaction task does. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum CompactionTaskKind { + /// Rewrite the task's fragments into new fragments, merging small ones, + /// dropping deleted rows and materializing overlays. Moves rows, so + /// indices covering the fragments are remapped. + #[default] + RewriteFragments, + /// Rewrite some columns of the task's one fragment into new data files, + /// leaving the fragment id, rows, deletions, overlays and index coverage + /// as they are. Each entry of `files` is one new file, given as the field + /// ids of the top-level columns it holds. The files the columns leave + /// keep their other columns; a file left with none is dropped. + RepackColumns { files: Vec> }, +} + +/// Which kinds of task a compaction plans. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +pub enum CompactionScope { + /// Rewrite the fragments that need it and repack the columns of the + /// others. + #[default] + All, + /// Only rewrite fragments. + RewriteFragments, + /// Only repack columns. + RepackColumns, +} + +impl TryFrom<&str> for CompactionScope { + type Error = Error; + + fn try_from(value: &str) -> std::result::Result { + match value.to_lowercase().as_str() { + "all" => Ok(Self::All), + "rewrite_fragments" => Ok(Self::RewriteFragments), + "repack_columns" => Ok(Self::RepackColumns), + _ => Err(Error::invalid_input(format!( + "Invalid compaction scope \"{}\". Valid values: \"all\", \"rewrite_fragments\", \"repack_columns\"", + value + ))), + } + } } /// A standalone task that can be serialized and sent to another machine for @@ -2146,7 +2506,9 @@ impl CompactionTask { } else { Cow::Owned(dataset.checkout_version(self.read_version).await?) }; - rewrite_files(dataset, self.task.clone(), &self.options).await + DefaultCompactionExecutor + .execute(dataset.as_ref(), self.task.clone(), &self.options) + .await } } @@ -2399,6 +2761,19 @@ pub struct RewriteResult { /// deferred index remap post-processing, or (2) used with reserved /// fragment IDs to build old-to-new mappings. pub row_addrs: Option>, + /// The new data files of a [`CompactionTaskKind::RepackColumns`] task. + /// `new_fragments` and `original_fragments` are then empty. + #[serde(default)] + pub repacked_files: Option, +} + +/// The data files a column repack wrote for one fragment. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct RepackedFiles { + /// The fragment the files belong to. + pub fragment_id: u64, + /// The new files, empty when the fragment had no rows to write. + pub files: Vec, } async fn reserve_fragment_ids( @@ -2437,6 +2812,241 @@ async fn reserve_fragment_ids( Ok(()) } +/// The schema of each data file a compacted fragment gets when column groups +/// are configured: the columns no group claims first (when any), then one +/// schema per group. Columns keep the dataset's order within each file. +/// +/// A name that is not a top-level column is skipped: the groups usually come +/// from table config, which `drop_columns` and a rename leave as they are, and +/// a stale name must not stop every later compaction. The planner warns about +/// it. +fn column_group_schemas( + schema: &lance_core::datatypes::Schema, + groups: &[Vec], +) -> Result> { + let names = || schema.fields.iter().map(|field| field.name.as_str()); + let mut seen = HashSet::new(); + for column in groups.iter().flatten() { + // Checked here as well as in `validate()`: a custom planner or a + // hand-built `CompactionTask` reaches `rewrite_files` without it, and a + // column in two groups would be written to two files. + if !seen.insert(column.as_str()) { + return Err(Error::invalid_input(format!( + "column_groups lists column \"{column}\" more than once" + ))); + } + } + let claimed = |name: &str| groups.iter().flatten().any(|column| column == name); + let rest: Vec<&str> = names().filter(|name| !claimed(name)).collect(); + std::iter::once(rest) + .chain(groups.iter().map(|group| { + names() + .filter(|name| group.iter().any(|column| column == name)) + .collect() + })) + .filter(|columns: &Vec<&str>| !columns.is_empty()) + .map(|columns| schema.project(&columns)) + .collect() +} + +/// The field ids of each group's top-level columns, without the columns no +/// group names. A name that is not a top-level column is skipped, and so is a +/// group left empty. +fn claimed_column_ids( + schema: &lance_core::datatypes::Schema, + groups: &[Vec], +) -> Vec> { + groups + .iter() + .map(|group| { + schema + .fields + .iter() + .filter(|field| group.contains(&field.name)) + .map(|field| field.id) + .collect::>() + }) + .filter(|group| !group.is_empty()) + .collect() +} + +/// Compaction's grouped write: fan one read of the merged rows out to one +/// writer per column group, so each output fragment gets one data file per +/// group instead of one file holding every column. All writers break at the +/// same `file_row_counts`, so their fragments line up one-to-one and merge by +/// appending each later group's files onto the first group's fragments. Index +/// seeds go to whichever group holds the indexed column. +async fn write_column_group_fragments( + write_version: ConcreteFileVersion, + dataset: &Dataset, + group_schemas: Vec, + mut reader: SendableRecordBatchStream, + mut params: WriteParams, + file_row_counts: Vec, +) -> Result> { + use futures::SinkExt; + + // `file_row_counts` is authoritative for the grouped write: every group must + // break at the same rows so the fragments merge one-to-one. A byte cap would + // roll an extra file in a wide group only (write.rs re-plans the remainder on + // a byte-driven close), leaving groups misaligned, so disable it here. + params.max_bytes_per_file = usize::MAX; + + let arrow_schema = reader.schema(); + let projections = group_schemas + .iter() + .map(|schema| { + schema + .fields + .iter() + .map(|field| arrow_schema.index_of(&field.name)) + .collect::, _>>() + }) + .collect::, _>>()?; + let group_arrow_schemas = projections + .iter() + .map(|projection| arrow_schema.project(projection).map(Arc::new)) + .collect::, _>>()?; + + // One bounded channel per group; the forwarder projects each read batch to + // every group. Bound 1 keeps a slow writer from letting the others buffer + // the whole scan ahead of it in memory. + let (mut senders, group_streams): (Vec<_>, Vec<_>) = group_arrow_schemas + .into_iter() + .map(|group_arrow| { + let (tx, rx) = + futures::channel::mpsc::channel::>(1); + let stream: SendableRecordBatchStream = + Box::pin(RecordBatchStreamAdapter::new(group_arrow, rx)); + (tx, stream) + }) + .unzip(); + + // Blob-check each group, then seed the scalar indexes once and route each + // seed to the group holding its column. Every dataset column lands in + // exactly one group (the unclaimed columns form the first group), so no seed + // is created only to be dropped, and each index is seeded exactly once. + for group_schema in &group_schemas { + versions::validate_write_schema(write_version, group_schema)?; + } + let mut per_group_seeds: Vec>> = + group_schemas.iter().map(|_| Vec::new()).collect(); + for seed in versions::create_seed_writers(write_version, Some(dataset), ¶ms).await? { + match group_schemas + .iter() + .position(|schema| schema.field(seed.column_name()).is_some()) + { + Some(index) => per_group_seeds[index].push(seed), + // `column_group_schemas` partitions every top-level column, so an + // index's keyed column always lands in exactly one group. Guard the + // invariant rather than silently dropping the seed, which would + // rebuild the index from incomplete data. + None => { + return Err(Error::internal(format!( + "column groups do not cover indexed column \"{}\"", + seed.column_name() + ))); + } + } + } + + // Returns which writer stopped reading first, if one did: a writer that + // fails closes its receiver, and its error is the one to report. + let forward = async move { + while let Some(batch) = reader.next().await { + let batch = batch?; + for (index, (sender, projection)) in senders.iter_mut().zip(&projections).enumerate() { + if sender.send(Ok(batch.project(projection)?)).await.is_err() { + return Ok(Some(index)); + } + } + } + Ok::, Error>(None) + }; + let writers = group_streams + .into_iter() + .zip(group_schemas) + .zip(per_group_seeds) + .map(|((group_stream, group_schema), seed_writers)| { + let params = params.clone(); + let file_row_counts = file_row_counts.clone(); + async move { + versions::write_fragments_direct( + write_version, + Some(dataset), + dataset.object_store.clone(), + &dataset.base, + &group_schema, + group_stream, + params, + None, + seed_writers, + Some(file_row_counts), + None, + ) + .await + } + }) + .collect::>(); + + // Run every writer to completion rather than dropping the others on the + // first error: a writer cleans up its own files only when it fails, so a + // group that finished must be cleaned up here if the write as a whole fails. + let (forwarded, per_group) = futures::join!(forward, futures::future::join_all(writers)); + let mut written = Vec::with_capacity(per_group.len()); + let mut errors = Vec::new(); + for (index, result) in per_group.into_iter().enumerate() { + match result { + Ok(fragments) => written.push(fragments), + Err(err) => errors.push((index, err)), + } + } + let failure = match forwarded { + // The read failed; the writers' errors follow from it. + Err(err) => Some(err), + Ok(Some(stopped)) => Some( + match errors.into_iter().find(|(index, _)| *index == stopped) { + Some((_, err)) => err, + // Not expected: a writer only finishes once its input ends, + // which needs the forwarder to have dropped every sender. + None => Error::internal(format!( + "column group writer {stopped} stopped reading before the compacted rows ended" + )), + }, + ), + Ok(None) => errors.into_iter().next().map(|(_, err)| err), + }; + let misaligned = written.iter().skip(1).any(|group| { + group.len() != written[0].len() + || written[0] + .iter() + .zip(group) + .any(|(a, b)| a.physical_rows != b.physical_rows) + }); + let failure = failure.or_else(|| { + misaligned.then(|| { + Error::internal( + "column group writers did not split the compacted rows at the same fragments", + ) + }) + }); + if let Some(err) = failure { + for fragments in &written { + cleanup_data_fragments(&dataset.object_store, &dataset.base, None, fragments).await; + } + return Err(err); + } + + let mut written = written.into_iter(); + let mut fragments = written.next().unwrap_or_default(); + for group in written { + for (fragment, other) in fragments.iter_mut().zip(group) { + fragment.files.extend(other.files); + } + } + Ok(fragments) +} + /// Rewrite the files in a single task. /// /// This assumes that the dataset is the correct read version to be compacted. @@ -2459,6 +3069,7 @@ async fn rewrite_files( read_version: dataset.manifest.version, original_fragments: task.fragments, row_addrs: None, + repacked_files: None, }); } @@ -2668,60 +3279,77 @@ async fn rewrite_files( None } else { let mut stream = reader.expect("reader must be prepared for non-binary-copy path"); - let mut write_schema = dataset.schema().clone(); - // On a table that can spill, the output fragments' lineage is computed - // before anything is written: the row ids and versions carry over from - // the inputs, so a sequence that leaves the manifest can ride along as - // a hidden column of the file being written rather than need a file of - // its own. Every other table carries its lineage over after the write, - // cut at the row counts actually written, and holds none of it while - // writing. Binary copy cannot add columns to the files it copies, so - // it always takes that path. - let spill_budget = if dataset.manifest.uses_stable_row_ids() { - inline_row_lineage_max_bytes(dataset.as_ref())? + if options.column_groups.is_empty() { + let mut write_schema = dataset.schema().clone(); + // On a table that can spill, the output fragments' lineage is computed + // before anything is written: the row ids and versions carry over from + // the inputs, so a sequence that leaves the manifest can ride along as + // a hidden column of the file being written rather than need a file of + // its own. Every other table carries its lineage over after the write, + // cut at the row counts actually written, and holds none of it while + // writing. Binary copy cannot add columns to the files it copies, so + // it always takes that path. + let spill_budget = if dataset.manifest.uses_stable_row_ids() { + inline_row_lineage_max_bytes(dataset.as_ref())? + } else { + None + }; + let lineage_plan = match spill_budget { + Some(limit) => { + let planned_rows = file_row_counts + .iter() + .map(|rows| *rows as u64) + .collect::>(); + let mut lineages = + compute_row_lineage(dataset.as_ref(), &fragments, &planned_rows).await?; + let plan = plan_row_lineage_spill(limit, &lineages); + if let RowLineagePlan::InFile(spill) = plan + && spill.any() + { + let fields = spill.schema_fields()?; + let columns = spill.take_columns(&mut lineages); + stream = append_row_lineage_columns(stream, &fields, columns); + // The writer stores only the columns its schema lists. + write_schema.fields.extend(fields); + } + Some(PlannedRowLineage { + lineages, + planned_rows, + plan, + }) + } + None => None, + }; + let (frags, _) = write_fragments_internal_with_file_row_counts( + write_version, + Some(dataset.as_ref()), + dataset.object_store.clone(), + &dataset.base, + write_schema, + stream, + params, + None, + Some(file_row_counts), + ) + .await?; + new_fragments = frags; + lineage_plan } else { + // column_groups repacks columns into one file per group and does not + // yet carry the in-file row-lineage spill; stable row ids are + // rechunked after the write instead (lineage_plan stays None). + let group_schemas = column_group_schemas(dataset.schema(), &options.column_groups)?; + new_fragments = write_column_group_fragments( + write_version, + dataset.as_ref(), + group_schemas, + stream, + params, + file_row_counts, + ) + .await?; None - }; - let lineage_plan = match spill_budget { - Some(limit) => { - let planned_rows = file_row_counts - .iter() - .map(|rows| *rows as u64) - .collect::>(); - let mut lineages = - compute_row_lineage(dataset.as_ref(), &fragments, &planned_rows).await?; - let plan = plan_row_lineage_spill(limit, &lineages); - if let RowLineagePlan::InFile(spill) = plan - && spill.any() - { - let fields = spill.schema_fields()?; - let columns = spill.take_columns(&mut lineages); - stream = append_row_lineage_columns(stream, &fields, columns); - // The writer stores only the columns its schema lists. - write_schema.fields.extend(fields); - } - Some(PlannedRowLineage { - lineages, - planned_rows, - plan, - }) - } - None => None, - }; - let (frags, _) = write_fragments_internal_with_file_row_counts( - write_version, - Some(dataset.as_ref()), - dataset.object_store.clone(), - &dataset.base, - write_schema, - stream, - params, - None, - Some(file_row_counts), - ) - .await?; - new_fragments = frags; - lineage_plan + } }; log::info!("Compaction task {}: file written", task_id); @@ -2789,6 +3417,7 @@ async fn rewrite_files( read_version: dataset.manifest.version, original_fragments: fragments, row_addrs, + repacked_files: None, }) } @@ -3153,11 +3782,48 @@ fn append_row_lineage_columns( /// they can be omitted and the successful tasks can be committed. However, once /// some of the tasks have been committed, the remainder of the tasks will not /// be able to be committed and should be considered cancelled. +/// +/// Fragment rewrites commit as one `Operation::Rewrite`; column repacks then +/// commit as one `Operation::DataReplacement` that moves values without +/// changing them, so a run with both kinds of task makes two versions. A +/// fragment in both kinds of task is rejected before either commit. If the second +/// commit fails, the first stays committed and the error is returned; the +/// repacks can be planned again. pub async fn commit_compaction( dataset: &mut Dataset, completed_tasks: Vec, remap_options: Arc, options: &CompactionOptions, +) -> Result { + let (repacks, rewrites): (Vec<_>, Vec<_>) = completed_tasks + .into_iter() + .partition(|task| task.repacked_files.is_some()); + // Otherwise the rewrite commits and the repack then conflicts with it. + let rewritten: HashSet = rewrites + .iter() + .flat_map(|task| task.original_fragments.iter().map(|fragment| fragment.id)) + .collect(); + if let Some(fragment_id) = repacks + .iter() + .filter_map(|task| task.repacked_files.as_ref()) + .map(|repacked| repacked.fragment_id) + .find(|id| rewritten.contains(id)) + { + return Err(Error::invalid_input(format!( + "fragment {fragment_id} is both rewritten and repacked by this compaction; \ + a plan may put a fragment in one kind of task only" + ))); + } + let mut metrics = Box::pin(commit_rewrites(dataset, rewrites, remap_options, options)).await?; + metrics += repack::commit_repacked_files(dataset, repacks, options).await?; + Ok(metrics) +} + +async fn commit_rewrites( + dataset: &mut Dataset, + completed_tasks: Vec, + remap_options: Arc, + options: &CompactionOptions, ) -> Result { if completed_tasks.is_empty() { return Ok(CompactionMetrics::default()); @@ -3661,6 +4327,7 @@ async fn cleanup_compaction_files_after_reservation_failure( mod tests { mod binary_copy; + mod repack; use self::remapping::RemappedIndex; use super::*; use crate::dataset::WriteDestination; @@ -8710,18 +9377,87 @@ mod tests { assert!(checked > 0, "expected to check at least one stored vector"); } - /// Build an `id` + `vec` dataset, create the given IVF vector index, - /// optionally delete rows, then run deferred compaction (which materializes - /// the deletions into the fragment-reuse index) and assert that KNN over - /// surviving vectors during the FRI window (a) never returns a deleted row - /// and (b) stays consistent with the pre-compaction answer. - /// - /// The deletion path is the interesting one: materialized deletions drop - /// rows from the quantization storage at load time, which shifts storage - /// positions. Flat storage (FLAT/PQ/SQ/RQ) is scanned linearly so this is - /// fine, but the HNSW graph addresses storage positionally and is not - /// frag-reuse aware, so a desync would surface here as recall collapse or a - /// resurrected/again-deleted row. + /// A column repack moves no row, so a vector index over the repacked + /// column keeps its uuid and coverage and answers the same. + #[tokio::test] + async fn repack_columns_preserves_vector_index() { + use arrow_array::cast::AsArray; + use arrow_array::types::{Float32Type, Int32Type}; + use lance_datagen::Dimension; + + const DIM: u32 = 32; + let mut dataset = lance_datagen::gen_batch() + .col("id", lance_datagen::array::step::()) + .col( + "vec", + lance_datagen::array::rand_vec::(Dimension::from(DIM)), + ) + .into_ram_dataset(FragmentCount::from(4), FragmentRowCount::from(256)) + .await + .unwrap(); + let params = VectorIndexParams::with_ivf_pq_params( + DistanceType::L2, + small_ivf(), + PQBuildParams { + max_iters: 2, + num_sub_vectors: 2, + ..Default::default() + }, + ); + dataset + .create_index( + &["vec"], + IndexType::Vector, + Some("vec_idx".into()), + ¶ms, + false, + ) + .await + .unwrap(); + let before = dataset + .load_index_by_name("vec_idx") + .await + .unwrap() + .unwrap(); + + // A query taken from a real row, and its answer before the rewrite. + let query = { + let mut scanner = dataset.scan(); + scanner.project(&["vec"]).unwrap(); + scanner.limit(Some(1), None).unwrap(); + let batch = scanner.try_into_batch().await.unwrap(); + batch["vec"] + .as_fixed_size_list() + .value(0) + .as_primitive::() + .values() + .to_vec() + }; + let knn_before = vector_knn_ids(&dataset, &query, 5).await; + + repack_columns(&mut dataset, vec![vec!["vec".into()]]).await; + for fragment in dataset.get_fragments() { + assert_eq!(fragment.metadata().files.len(), 2, "vec is split out"); + } + + let after = dataset + .load_index_by_name("vec_idx") + .await + .unwrap() + .unwrap(); + assert_eq!(after.uuid, before.uuid, "index must not be rebuilt"); + assert_eq!( + after.fragment_bitmap, before.fragment_bitmap, + "index coverage must be unchanged" + ); + assert_eq!( + vector_knn_ids(&dataset, &query, 5).await, + knn_before, + "search must return the same rows" + ); + dataset.validate().await.unwrap(); + } + /// Top-k `id`s for a KNN query against the `vec` column. async fn vector_knn_ids(dataset: &Dataset, query: &[f32], k: usize) -> Vec { use arrow_array::cast::AsArray; @@ -8744,6 +9480,18 @@ mod tests { ids } + /// Build an `id` + `vec` dataset, create the given IVF vector index, + /// optionally delete rows, then run deferred compaction (which materializes + /// the deletions into the fragment-reuse index) and assert that KNN over + /// surviving vectors during the FRI window (a) never returns a deleted row + /// and (b) stays consistent with the pre-compaction answer. + /// + /// The deletion path is the interesting one: materialized deletions drop + /// rows from the quantization storage at load time, which shifts storage + /// positions. Flat storage (FLAT/PQ/SQ/RQ) is scanned linearly so this is + /// fine, but the HNSW graph addresses storage positionally and is not + /// frag-reuse aware, so a desync would surface here as recall collapse or a + /// resurrected/again-deleted row. async fn check_vector_defer_compaction( params: VectorIndexParams, delete_predicate: Option<&str>, @@ -11936,4 +12684,287 @@ mod tests { Some(4) ); } + + /// Three-column dataset in two 10-row fragments, all columns in one file. + async fn column_group_dataset() -> (Dataset, RecordBatch) { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..20)), + Arc::new(Int32Array::from_iter_values((0..20).map(|v| v * 10))), + Arc::new(Int32Array::from_iter_values((0..20).map(|v| v * 100))), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], schema); + let dataset = Dataset::write( + reader, + "memory://", + Some(WriteParams { + max_rows_per_file: 10, + data_storage_version: Some(LanceFileVersion::V2_0), + ..Default::default() + }), + ) + .await + .unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + (dataset, before) + } + + #[tokio::test] + async fn compaction_column_groups_write_one_file_per_group() { + let (mut dataset, before) = column_group_dataset().await; + + // Merge the two small fragments; column_groups splits the output columns + // into their own files: {a} unclaimed, then {b} and {c}. + let options = CompactionOptions { + column_groups: vec![vec!["b".into()], vec!["c".into()]], + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1, "the two fragments merged into one"); + let mut groups: Vec> = fragments[0] + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + groups.sort(); + assert_eq!( + groups, + vec![vec![0], vec![1], vec![2]], + "one data file per column group" + ); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } + + #[tokio::test] + async fn compaction_column_groups_ignore_unknown_column() { + // A configured group can outlive the column it names (drop_columns and + // renames leave the config as it is), so the name is skipped and the + // rest of the groups still apply. + let (mut dataset, before) = column_group_dataset().await; + let options = CompactionOptions { + column_groups: vec![vec!["nonexistent".into(), "c".into()]], + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1); + let mut groups: Vec> = fragments[0] + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + groups.sort(); + assert_eq!(groups, vec![vec![0, 1], vec![2]]); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + } + + #[tokio::test] + async fn compaction_column_groups_ignore_byte_limit() { + let (mut dataset, before) = column_group_dataset().await; + // A tiny byte cap plus groups of different widths ({b,c} vs {a}) would, + // if honored, roll the wider group into more files and leave the groups + // misaligned. column_groups splits by planned row counts only, so the + // compaction succeeds and still yields one file per group. + let options = CompactionOptions { + column_groups: vec![vec!["b".into(), "c".into()]], + max_bytes_per_file: Some(64), + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1, "the two fragments merged into one"); + let mut groups: Vec> = fragments[0] + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + groups.sort(); + assert_eq!( + groups, + vec![vec![0], vec![1, 2]], + "one file for {{a}}, one for the {{b, c}} group" + ); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } + + #[tokio::test] + async fn compaction_column_groups_preserve_scalar_index() { + let (mut dataset, before) = column_group_dataset().await; + // ZoneMap is the scalar index that writes an index seed, so it actually + // drives the seed-routing loop in `write_column_group_fragments` (BTree + // writes no seed and would leave that path untested). + // ZoneMap only writes an index seed when `use_seeds` is enabled (for a + // fixed-width column it defaults off), and the seed is what drives the + // routing loop in `write_column_group_fragments`. + let index_params = ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap) + .with_params(&serde_json::json!({"use_seeds": true})); + dataset + .create_index( + &["a"], + IndexType::ZoneMap, + Some("a_idx".into()), + &index_params, + false, + ) + .await + .unwrap(); + // Split the indexed column into its own file group: its seed must be + // routed to the group holding "a" (a miss would trip the + // `Error::internal` guard and fail the compaction). + let options = CompactionOptions { + column_groups: vec![vec!["a".into()]], + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + // The grouped write ran the seed-routing loop (ZoneMap writes a seed, + // unlike BTree) and produced the requested split without corrupting the + // index metadata or the values. + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1, "the two fragments merged into one"); + let mut groups: Vec> = fragments[0] + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + groups.sort(); + assert_eq!( + groups, + vec![vec![0], vec![1, 2]], + "\"a\" split into its own file, {{b, c}} in another" + ); + assert!( + dataset.load_index_by_name("a_idx").await.unwrap().is_some(), + "the ZoneMap index survives compaction" + ); + let filtered = dataset + .scan() + .filter("a >= 15") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(filtered.num_rows(), 5, "a >= 15 matches five rows (15..20)"); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } + + /// Run only the column repacks `column_groups` calls for. + async fn repack_columns(dataset: &mut Dataset, column_groups: Vec>) { + let options = CompactionOptions { + column_groups, + scope: CompactionScope::RepackColumns, + ..Default::default() + }; + compact_files(dataset, options, None).await.unwrap(); + } + + #[tokio::test] + async fn column_groups_preserve_repacked_split_across_compaction() { + let (mut dataset, before) = column_group_dataset().await; + // Split column c into its own file per fragment. + repack_columns(&mut dataset, vec![vec!["c".into()]]).await; + for fragment in dataset.get_fragments() { + assert_eq!( + fragment.metadata().files.len(), + 2, + "c is split into its own file" + ); + } + // Vertical compaction with a matching group keeps c split instead of + // folding it back into one file — the feature's core promise. + let options = CompactionOptions { + column_groups: vec![vec!["c".into()]], + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + let fragments = dataset.get_fragments(); + assert_eq!(fragments.len(), 1, "the two fragments merged"); + let mut groups: Vec> = fragments[0] + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect(); + groups.sort(); + assert_eq!( + groups, + vec![vec![0, 1], vec![2]], + "layout stays split: {{a, b}} and {{c}}" + ); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + } + + #[test] + fn compaction_column_groups_parse_from_config() { + let config = std::collections::HashMap::from([( + "lance.compaction.column_groups".to_string(), + "b, c ; d".to_string(), + )]); + let options = CompactionOptions::from_dataset_config(&config).unwrap(); + assert_eq!( + options.column_groups, + vec![ + vec!["b".to_string(), "c".to_string()], + vec!["d".to_string()] + ], + "';' splits groups, ',' splits columns, whitespace trimmed" + ); + } + + #[test] + fn compaction_column_groups_validate_rejects_duplicate() { + let mut options = CompactionOptions { + column_groups: vec![vec!["b".into()], vec!["b".into()]], + ..Default::default() + }; + let err = options.validate().unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + } + + /// `rewrite_files` builds its group schemas without `validate()` (a custom + /// planner or a hand-built `CompactionTask`), so the duplicate check has to + /// hold there too. + #[test] + fn column_group_schemas_reject_duplicate() { + let schema = lance_core::datatypes::Schema::try_from(&Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + ])) + .unwrap(); + let err = column_group_schemas(&schema, &[vec!["b".into()], vec!["b".into()]]).unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + assert!(err.to_string().contains("more than once"), "{err}"); + } + + #[test] + fn compaction_column_groups_validate_rejects_force_binary_copy() { + let mut options = CompactionOptions { + column_groups: vec![vec!["b".into()]], + compaction_mode: Some(CompactionMode::ForceBinaryCopy), + ..Default::default() + }; + let err = options.validate().unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + assert!(err.to_string().contains("ForceBinaryCopy"), "{err}"); + } } diff --git a/rust/lance/src/dataset/optimize/repack.rs b/rust/lance/src/dataset/optimize/repack.rs new file mode 100644 index 00000000000..2cee6e0cff0 --- /dev/null +++ b/rust/lance/src/dataset/optimize/repack.rs @@ -0,0 +1,505 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Column repacking, the compaction task that rewrites some of one fragment's +//! columns into new data files without moving rows: fewer files, or the files +//! `column_groups` asks for. +//! +//! Every `add_columns` backfill gives each fragment one more data file, and a +//! large fragment with few deletions never qualifies for a rewrite, so its +//! files pile up. A [`CompactionTaskKind::RepackColumns`] task reads some of +//! the fragment's columns, writes them to new files, and commits the files as +//! an `Operation::DataReplacement` with `data_change: false`: the fragment id, +//! row addresses, overlays and index coverage are left as they are. + +use std::collections::{BTreeSet, HashMap, HashSet}; +use std::sync::Arc; + +use lance_core::datatypes::{Field as LanceField, Schema}; +use lance_file::version::ConcreteFileVersion; +use lance_table::format::{DataFile, Fragment}; + +use super::{ + CompactionMetrics, CompactionOptions, CompactionTaskKind, RepackedFiles, RewriteResult, + TaskData, +}; +use crate::dataset::fragment::FileFragment; +use crate::dataset::transaction::{DataReplacementGroup, Operation, TransactionBuilder}; +use crate::dataset::{Dataset, cleanup_data_fragments, versions}; +use crate::{Error, Result}; + +/// The ids of `field` and every field beneath it. +fn subtree_ids(field: &LanceField, ids: &mut HashSet) { + ids.insert(field.id); + for child in &field.children { + subtree_ids(child, ids); + } +} + +/// The ids of the leaf fields at and beneath `field`. The planner places a +/// column by where its leaves are: a V2.0 file can keep a struct's header +/// after the struct's last child in it was dropped. +fn leaf_ids(field: &LanceField, ids: &mut HashSet) { + if field.children.is_empty() { + ids.insert(field.id); + } + for child in &field.children { + leaf_ids(child, ids); + } +} + +fn holds_blob(field: &LanceField) -> bool { + field.is_blob() || field.children.iter().any(holds_blob) +} + +/// Every field id a data file answers for, the way a `DataReplacement` +/// commit counts it. +fn file_coverage(file: &DataFile, schema: &Schema) -> HashSet { + file.schema(schema) + .field_ids() + .into_iter() + .chain(file.fields.iter().copied()) + .filter(|id| *id >= 0) + .collect() +} + +/// The new files a repack of `fragment` should write, each the top-level +/// field ids it holds, or `None` when no repack is due. +/// +/// `groups` are the column groups, each the top-level field ids to keep in a +/// file of their own. A group is moved unless one live file holds exactly its +/// columns; moving it out leaves the other columns where they were. Over +/// `max_files`, the files holding only columns no group names are merged as +/// well, under the rule below. +/// +/// A merge only takes files it empties: files whose columns can all move, +/// which hold no spilled lineage and share no column with a file holding it, +/// and whose other fields (a struct header) belong to columns that move too. +/// The files left keep the columns that do not move. Without groups, when +/// every file holding a column can be emptied, at least one of them stays +/// unless `max_files` is 1. Trying the files largest first by recorded size, +/// the one kept is the first whose merge (of the files sharing no column with +/// it, else of every other file, as long as it keeps a column) leaves the +/// largest file in place, or failing that the first with any merge. Nothing is merged +/// unless at least two files would be emptied, so a merge lowers the file +/// count. +/// +/// A column the fragment has no data for (added as all nulls) is left out of +/// its group, and a group holding a blob column, or a column whose fields are +/// only partly in the fragment's files, is left as it is. The new files carry +/// no row lineage, so a repack that would move every column out of the file +/// holding the fragment's spilled lineage is not planned: that file would +/// stay for the lineage alone until a rewrite of the fragment reclaims it. +/// +/// `live_files` is the fragment's +/// [`FragmentColumnLayoutStats::live_file_count`]. A V2.0 file left holding +/// only a struct header counts there but holds no column here: it goes when +/// a repack moves the struct, and otherwise stays until a rewrite of the +/// fragment. +/// +/// [`FragmentColumnLayoutStats::live_file_count`]: crate::dataset::compaction_stats::FragmentColumnLayoutStats::live_file_count +pub(super) fn plan_fragment_repack( + schema: &Schema, + fragment: &Fragment, + live_files: usize, + groups: Option<&[Vec]>, + max_files: Option, +) -> Option>> { + let over_file_limit = max_files.is_some_and(|max| live_files > max); + if groups.is_none() && !over_file_limit { + return None; + } + // A V1 file cannot tombstone single fields. + if fragment.files.iter().any(|file| { + !file + .file_version() + .is_ok_and(|version| version != ConcreteFileVersion::V1) + }) { + return None; + } + let schema_ids: HashSet = schema.fields_pre_order().map(|field| field.id).collect(); + let spilled = fragment.spilled_row_lineage_field_ids(); + let live: Vec<(&DataFile, HashSet)> = fragment + .files + .iter() + .filter(|file| file.fields.iter().any(|id| schema_ids.contains(id))) + .map(|file| (file, file_coverage(file, schema))) + .collect(); + let covered: HashSet = live.iter().flat_map(|(_, ids)| ids).copied().collect(); + let column_leaves: Vec<(i32, HashSet)> = schema + .fields + .iter() + .map(|field| { + let mut ids = HashSet::new(); + leaf_ids(field, &mut ids); + (field.id, ids) + }) + .collect(); + // Every schema field id under each top-level column. + let column_fields: HashMap> = schema + .fields + .iter() + .map(|field| { + let mut ids = HashSet::new(); + subtree_ids(field, &mut ids); + (field.id, ids) + }) + .collect(); + let present: BTreeSet = column_leaves + .iter() + .filter(|(_, leaves)| !leaves.is_disjoint(&covered)) + .map(|(column, _)| *column) + .collect(); + let movable: HashSet = schema + .fields + .iter() + .zip(&column_leaves) + .filter(|(field, (column, leaves))| { + present.contains(column) && !holds_blob(field) && leaves.is_subset(&covered) + }) + .map(|(field, _)| field.id) + .collect(); + // The top-level columns each live file holds data for. + let file_columns: Vec> = live + .iter() + .map(|(_, coverage)| { + column_leaves + .iter() + .filter(|(_, leaves)| !leaves.is_disjoint(coverage)) + .map(|(column, _)| *column) + .collect() + }) + .collect(); + let holds_lineage = |file: &DataFile| file.fields.iter().any(|id| spilled.contains(id)); + let lineage_columns: BTreeSet = live + .iter() + .zip(&file_columns) + .filter(|((file, _), _)| holds_lineage(file)) + .flat_map(|(_, columns)| columns.iter().copied()) + .collect(); + // The live files a merge may take: every column can move, and moving them + // takes nothing out of a file holding spilled lineage. + let candidates: Vec = (0..live.len()) + .filter(|index| { + let columns = &file_columns[*index]; + !columns.is_empty() + && !holds_lineage(live[*index].0) + && columns.is_disjoint(&lineage_columns) + && columns.iter().all(|column| movable.contains(column)) + }) + .collect(); + let columns_of = |files: &[usize]| -> BTreeSet { + files + .iter() + .flat_map(|index| file_columns[*index].iter().copied()) + .collect() + }; + // Of `files`, the ones a move of all their columns together leaves with no + // schema field, so the commit drops them. A file can also hold a field of + // a column it holds no data for (a struct header), which only moving that + // column from another file takes out. + let emptied = |mut files: Vec| loop { + let moved: HashSet = files + .iter() + .flat_map(|index| &file_columns[*index]) + .flat_map(|column| &column_fields[column]) + .copied() + .collect(); + let next: Vec = files + .iter() + .copied() + .filter(|index| { + live[*index] + .0 + .fields + .iter() + .all(|id| !schema_ids.contains(id) || moved.contains(id)) + }) + .collect(); + if next.len() == files.len() { + break files; + } + files = next; + }; + + let wanted: Vec> = match groups { + Some(groups) => { + let mut wanted: Vec> = groups + .iter() + .map(|group| group.iter().copied().collect()) + .collect(); + let claimed: BTreeSet = wanted.iter().flatten().copied().collect(); + if over_file_limit { + let merged = emptied( + candidates + .iter() + .copied() + .filter(|index| file_columns[*index].is_disjoint(&claimed)) + .collect(), + ); + if merged.len() > 1 { + wanted.push(columns_of(&merged)); + } + } + wanted + } + None => { + let all = emptied(candidates); + let holding = file_columns.iter().filter(|c| !c.is_empty()).count(); + // A file a move cannot empty stays anyway. Otherwise one file + // stays unless the limit is 1. + let merged = if all.len() < 2 || all.len() < holding || max_files == Some(1) { + all + } else { + let mut by_size = all.clone(); + by_size.sort_by_key(|index| { + std::cmp::Reverse(( + live[*index].0.file_size_bytes.get().map(|size| size.get()), + file_columns[*index].len(), + )) + }); + // For each file kept, largest first: merge the files sharing + // no column with it, which leaves its data in place, else every + // other file, as long as the kept file keeps a column. A plan + // that leaves the largest file in place wins. + let plans = by_size.iter().flat_map(|kept| { + let untouched = emptied( + all.iter() + .copied() + .filter(|index| file_columns[*index].is_disjoint(&file_columns[*kept])) + .collect(), + ); + let others = emptied(all.iter().copied().filter(|i| i != kept).collect()); + let keeps_a_column = !file_columns[*kept].is_subset(&columns_of(&others)); + [Some(untouched), keeps_a_column.then_some(others)] + .into_iter() + .flatten() + .filter(|merged| merged.len() > 1) + }); + let largest = by_size[0]; + plans + .clone() + .find(|merged| !merged.contains(&largest)) + .or_else(|| plans.clone().next()) + .unwrap_or_default() + }; + if merged.len() > 1 { + vec![columns_of(&merged)] + } else { + Vec::new() + } + } + }; + + let out_of_place: Vec> = wanted + .into_iter() + .filter_map(|group| { + let group: BTreeSet = group.intersection(&present).copied().collect(); + if group.is_empty() || !group.iter().all(|column| movable.contains(column)) { + return None; + } + let holders: Vec<&BTreeSet> = file_columns + .iter() + .filter(|columns| !columns.is_disjoint(&group)) + .collect(); + let in_place = holders.len() == 1 && *holders[0] == group; + (!in_place).then(|| group.into_iter().collect()) + }) + .collect(); + let moved: HashSet = out_of_place.iter().flatten().copied().collect(); + let strands_lineage = live.iter().zip(&file_columns).any(|((file, _), columns)| { + holds_lineage(file) + && !columns.is_empty() + && columns.iter().all(|column| moved.contains(column)) + }); + (!out_of_place.is_empty() && !strands_lineage).then_some(out_of_place) +} + +/// How many of `fragment`'s data files a commit of `written` would drop: the +/// files left holding no field of the schema and none of the fragment's +/// spilled row lineage, the rule the `DataReplacement` commit applies. +fn files_dropped(fragment: &Fragment, written: &[DataFile], schema: &Schema) -> usize { + let moved: HashSet = written + .iter() + .flat_map(|file| file_coverage(file, schema)) + .collect(); + let schema_ids: HashSet = schema.fields_pre_order().map(|field| field.id).collect(); + let spilled = fragment.spilled_row_lineage_field_ids(); + fragment + .files + .iter() + .filter(|file| { + !file + .fields + .iter() + .any(|id| (schema_ids.contains(id) && !moved.contains(id)) || spilled.contains(id)) + }) + .count() +} + +/// Write the new files of a [`CompactionTaskKind::RepackColumns`] task +/// against `dataset`, which must be at the plan's read version. +pub(super) async fn execute_repack( + dataset: &Dataset, + task: TaskData, + options: &CompactionOptions, +) -> Result { + let CompactionTaskKind::RepackColumns { files } = &task.kind else { + return Err(Error::internal("execute_repack called with a rewrite task")); + }; + let [fragment] = task.fragments.as_slice() else { + return Err(Error::invalid_input(format!( + "a RepackColumns task names exactly one fragment, got {}", + task.fragments.len() + ))); + }; + let write_version = options.write_version(dataset); + versions::validate_write_version( + dataset.manifest.data_storage_format.lance_file_format(), + write_version, + )?; + if write_version == ConcreteFileVersion::V1 { + return Err(Error::not_supported( + "repacking columns needs the V2 file format: a V1 file cannot tombstone \ + single fields", + )); + } + + // Read the base values only: the overlays stay on the fragment and keep + // shadowing them, so every live row reads as it did. + let mut base = fragment.clone(); + base.overlays.clear(); + let source = FileFragment::new(Arc::new(dataset.clone()), base); + let schema = dataset.schema(); + let batch_size = options.batch_size.map(|size| size as u32); + + let mut written: Vec = Vec::with_capacity(files.len()); + for column_ids in files { + let names: Vec<&str> = schema + .fields + .iter() + .filter(|field| column_ids.contains(&field.id)) + .map(|field| field.name.as_str()) + .collect(); + match write_one_file(&source, schema, &names, batch_size, write_version).await { + Ok(Some(file)) => written.push(file), + Ok(None) => {} + Err(err) => { + let partial = Fragment { + files: written, + ..Fragment::new(fragment.id) + }; + cleanup_data_fragments(&dataset.object_store, &dataset.base, None, &[partial]) + .await; + return Err(err); + } + } + } + + let metrics = CompactionMetrics { + files_added: written.len(), + files_removed: files_dropped(fragment, &written, schema), + ..Default::default() + }; + Ok(RewriteResult { + metrics, + new_fragments: Vec::new(), + read_version: dataset.manifest.version, + original_fragments: Vec::new(), + row_addrs: None, + repacked_files: Some(RepackedFiles { + fragment_id: fragment.id, + files: written, + }), + }) +} + +/// Write `columns` of `source` to one new data file. `None` when the +/// fragment has no rows to write. +async fn write_one_file( + source: &FileFragment, + schema: &Schema, + columns: &[&str], + batch_size: Option, + write_version: ConcreteFileVersion, +) -> Result> { + let write_schema = schema.project(columns)?; + let mut updater = source + .updater_with_version( + Some(columns), + Some((write_schema, schema.clone())), + batch_size, + None, + write_version, + ) + .await?; + let finished: Result = async { + while let Some(batch) = updater.next().await?.cloned() { + updater.update(batch).await?; + } + updater.finish().await + } + .await; + let fragment = match finished { + Ok(fragment) => fragment, + Err(err) => { + updater.cleanup_unfinished_writer().await; + return Err(err); + } + }; + // The updater appends the file it wrote to the fragment's files. + Ok((fragment.files.len() > source.metadata().files.len()) + .then(|| fragment.files.last().cloned()) + .flatten()) +} + +/// Commit the files of repack results as one `DataReplacement` that moves +/// values without changing them. Each result applies to its fragment as it +/// stands at commit, so a concurrent change to other columns of the fragment +/// is kept; one that touches the moved columns fails the commit with a +/// retryable conflict. +pub(super) async fn commit_repacked_files( + dataset: &mut Dataset, + results: Vec, + options: &CompactionOptions, +) -> Result { + let mut metrics = CompactionMetrics::default(); + let mut read_version = u64::MAX; + let mut replacements = Vec::new(); + for result in results { + let Some(repacked) = result.repacked_files else { + continue; + }; + if repacked.files.is_empty() { + continue; + } + metrics += result.metrics; + read_version = read_version.min(result.read_version); + replacements.extend( + repacked + .files + .into_iter() + .map(|file| DataReplacementGroup(repacked.fragment_id, file)), + ); + } + if replacements.is_empty() { + return Ok(CompactionMetrics::default()); + } + let transaction = TransactionBuilder::new( + // The earliest version a repack read, so the conflict check covers + // every write since (the same reason vertical compaction uses it). + read_version, + Operation::DataReplacement { + replacements, + data_change: false, + }, + ) + .transaction_properties(options.transaction_properties.clone()) + .build(); + // Like a rewrite's, a repack result can be retried after an ambiguous + // success, so its files are left for cleanup rather than deleted here. + dataset + .apply_commit(transaction, &Default::default(), &Default::default()) + .await?; + Ok(metrics) +} diff --git a/rust/lance/src/dataset/optimize/tests/repack.rs b/rust/lance/src/dataset/optimize/tests/repack.rs new file mode 100644 index 00000000000..6c75e82a249 --- /dev/null +++ b/rust/lance/src/dataset/optimize/tests/repack.rs @@ -0,0 +1,1273 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The Lance Authors + +//! Column repacking through `compact_files`. + +use super::*; +use crate::dataset::NewColumnTransform; +use arrow_array::StringArray; +use lance_index::IndexType; +use lance_index::scalar::ScalarIndexParams; + +fn abc_batch(start: i32, rows: i32) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Utf8, true), + ])); + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(start..start + rows)), + Arc::new(Int32Array::from_iter_values( + (start..start + rows).map(|v| v * 10), + )), + Arc::new(StringArray::from_iter_values( + (start..start + rows).map(|v| format!("row-{v}")), + )), + ], + ) + .unwrap() +} + +/// Two fragments of four rows, each `a, b, c` in one file. +async fn write(uri: &str, version: LanceFileVersion) -> Dataset { + let data = abc_batch(0, 8); + let reader = RecordBatchIterator::new([Ok(data.clone())], data.schema()); + Dataset::write( + reader, + uri, + Some(WriteParams { + max_rows_per_file: 4, + data_storage_version: Some(version), + ..Default::default() + }), + ) + .await + .unwrap() +} + +/// `write`, then two backfills: every fragment holds its columns in three +/// files, `[a, b, c]`, `[d]` and `[e]`. +async fn write_backfilled() -> Dataset { + let mut dataset = write("memory://", LanceFileVersion::V2_0).await; + for (name, expr) in [("d", "a + 1"), ("e", "a + 2")] { + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![(name.into(), expr.into())]), + None, + None, + ) + .await + .unwrap(); + } + dataset +} + +fn repack_options(max_files: Option, groups: Vec>) -> CompactionOptions { + CompactionOptions { + max_data_files_per_fragment: max_files, + column_groups: groups + .into_iter() + .map(|group| group.into_iter().map(String::from).collect()) + .collect(), + scope: CompactionScope::RepackColumns, + ..Default::default() + } +} + +fn layout(dataset: &Dataset) -> Vec>> { + dataset + .get_fragments() + .iter() + .map(|fragment| { + fragment + .metadata() + .files + .iter() + .map(|file| file.fields.to_vec()) + .collect() + }) + .collect() +} + +fn live_file_counts(dataset: &Dataset) -> Vec { + dataset + .column_layout_stats() + .iter() + .map(|stats| stats.live_file_count) + .collect() +} + +#[tokio::test] +async fn repack_collapses_backfilled_files() { + let mut dataset = write_backfilled().await; + let before = dataset.scan().try_into_batch().await.unwrap(); + let version = dataset.manifest.version; + assert_eq!(live_file_counts(&dataset), vec![3, 3]); + + let metrics = compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + + assert_eq!(dataset.manifest.version, version + 1, "one commit"); + assert_eq!(live_file_counts(&dataset), vec![1, 1]); + assert_eq!(layout(&dataset), vec![vec![vec![0, 1, 2, 3, 4]]; 2]); + assert_eq!(metrics.files_added, 2); + assert_eq!(metrics.files_removed, 6); + assert_eq!(metrics.fragments_added, 0); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + + // Nothing is left to do, so nothing is committed. + compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + assert_eq!(dataset.manifest.version, version + 1); +} + +/// Above the limit, the columns of the biggest file stay where they are and +/// only the others are rewritten. +#[tokio::test] +async fn repack_keeps_the_biggest_file() { + let mut dataset = write_backfilled().await; + let before = dataset.scan().try_into_batch().await.unwrap(); + + compact_files(&mut dataset, repack_options(Some(2), vec![]), None) + .await + .unwrap(); + + for files in layout(&dataset) { + assert_eq!(files.len(), 2, "{files:?}"); + assert_eq!(files[0], vec![0, 1, 2], "the base file is untouched"); + assert_eq!(files[1], vec![3, 4]); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); +} + +#[tokio::test] +async fn repack_follows_column_groups() { + let mut dataset = write_backfilled().await; + let before = dataset.scan().try_into_batch().await.unwrap(); + + compact_files( + &mut dataset, + repack_options(None, vec![vec!["c", "e"]]), + None, + ) + .await + .unwrap(); + + for files in layout(&dataset) { + let mut files = files + .into_iter() + .map(|fields| { + fields + .into_iter() + .filter(|id| *id != TOMBSTONE_FIELD_ID) + .collect::>() + }) + .collect::>(); + files.sort(); + // Only the group moves; the other columns stay where they were. + assert_eq!(files, vec![vec![0, 1], vec![2, 4], vec![3]], "{files:?}"); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + + // The layout matches the groups now. + let version = dataset.manifest.version; + compact_files( + &mut dataset, + repack_options(None, vec![vec!["c", "e"]]), + None, + ) + .await + .unwrap(); + assert_eq!(dataset.manifest.version, version); +} + +/// Deleted rows stay deleted and the new files line up with the old ones. +#[tokio::test] +async fn repack_keeps_deletions() { + let mut dataset = write_backfilled().await; + dataset.delete("a = 1 OR a = 6").await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + + let options = CompactionOptions { + // Neither small nor deleted enough to be rewritten. + target_rows_per_fragment: 4, + materialize_deletions_threshold: 0.5, + max_data_files_per_fragment: Some(1), + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + assert_eq!(live_file_counts(&dataset), vec![1, 1]); + assert!( + dataset + .get_fragments() + .iter() + .all(|fragment| fragment.metadata().deletion_file.is_some()) + ); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); +} + +#[tokio::test] +async fn repack_keeps_scalar_index() { + let mut dataset = write_backfilled().await; + dataset + .create_index( + &["d"], + IndexType::BTree, + Some("d_idx".into()), + &ScalarIndexParams::default(), + false, + ) + .await + .unwrap(); + let before = dataset.load_indices().await.unwrap()[0].clone(); + + compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + + let after = dataset.load_indices().await.unwrap()[0].clone(); + assert_eq!(after.uuid, before.uuid, "the index is not rebuilt"); + assert_eq!(after.fragment_bitmap, before.fragment_bitmap); + let filtered = dataset + .scan() + .filter("d = 6") + .unwrap() + .try_into_batch() + .await + .unwrap(); + assert_eq!(filtered.num_rows(), 1); + dataset.validate().await.unwrap(); +} + +/// `drop_columns` leaves a dropped column's id in its file. A repack that +/// moves the other columns out drops that file, as `drop_columns` would. +#[tokio::test] +async fn repack_drops_file_left_with_only_dropped_columns() { + let mut dataset = write_backfilled().await; + dataset.drop_columns(&["c"]).await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + + compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + + assert_eq!(layout(&dataset), vec![vec![vec![0, 1, 3, 4]]; 2]); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); +} + +/// A run that rewrites some fragments and repacks others commits the +/// rewrite first and the repacks second, and never touches one fragment both +/// ways. +#[tokio::test] +async fn compaction_rewrites_and_repacks_in_one_run() { + let mut dataset = write_backfilled().await; + // Fragment 0 qualifies for a rewrite, fragment 1 only for a repack. + dataset.delete("a = 0 OR a = 1").await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + let version = dataset.manifest.version; + + let options = CompactionOptions { + target_rows_per_fragment: 3, + max_data_files_per_fragment: Some(1), + ..Default::default() + }; + let plan = plan_compaction(&dataset, &options).await.unwrap(); + assert_eq!(plan.tasks.len(), 2, "{plan:?}"); + assert_eq!(plan.tasks[0].kind, CompactionTaskKind::RewriteFragments); + assert_eq!(plan.tasks[0].fragments[0].id, 0); + assert!(matches!( + plan.tasks[1].kind, + CompactionTaskKind::RepackColumns { .. } + )); + assert_eq!(plan.tasks[1].fragments[0].id, 1); + + let metrics = compact_files(&mut dataset, options, None).await.unwrap(); + let last = dataset.manifest.version; + assert!(last >= version + 2); + for (version, operation) in [(last - 1, "Rewrite"), (last, "DataReplacement")] { + let transaction = dataset + .checkout_version(version) + .await + .unwrap() + .read_transaction() + .await + .unwrap() + .unwrap(); + assert_eq!(transaction.operation.to_string(), operation); + } + assert_eq!(metrics.fragments_removed, 1); + assert_eq!(metrics.fragments_added, 1); + assert_eq!(live_file_counts(&dataset), vec![1, 1]); + // The rewritten fragment gets a new id, so it now scans last; the repack + // changes no value. + let rewritten = dataset.checkout_version(last - 1).await.unwrap(); + let after_rewrite = rewritten.scan().try_into_batch().await.unwrap(); + assert_eq!(after_rewrite.num_rows(), before.num_rows()); + assert_eq!( + dataset.scan().try_into_batch().await.unwrap(), + after_rewrite + ); + dataset.validate().await.unwrap(); +} + +#[tokio::test] +async fn compaction_scope_selects_task_kinds() { + let mut dataset = write_backfilled().await; + dataset.delete("a = 0 OR a = 1").await.unwrap(); + let options = |scope| CompactionOptions { + target_rows_per_fragment: 2, + max_data_files_per_fragment: Some(1), + scope, + ..Default::default() + }; + + let rewrites = plan_compaction(&dataset, &options(CompactionScope::RewriteFragments)) + .await + .unwrap(); + assert_eq!(rewrites.tasks.len(), 1); + assert_eq!(rewrites.tasks[0].kind, CompactionTaskKind::RewriteFragments); + + let repacks = plan_compaction(&dataset, &options(CompactionScope::RepackColumns)) + .await + .unwrap(); + assert_eq!(repacks.tasks.len(), 2, "both fragments are repacked"); + assert!( + repacks + .tasks + .iter() + .all(|task| matches!(task.kind, CompactionTaskKind::RepackColumns { .. })) + ); +} + +/// Without `max_data_files_per_fragment` or `column_groups`, nothing is +/// repacked. +#[tokio::test] +async fn compaction_plans_no_repack_by_default() { + let dataset = write_backfilled().await; + let plan = plan_compaction(&dataset, &repack_options(None, vec![])) + .await + .unwrap(); + assert!(plan.tasks.is_empty(), "{plan:?}"); +} + +/// Repack tasks run on another machine: they are serialized, executed on +/// their own, and their results committed together. +#[tokio::test] +async fn repack_tasks_run_distributed() { + let mut dataset = write_backfilled().await; + let before = dataset.scan().try_into_batch().await.unwrap(); + let plan = plan_compaction(&dataset, &repack_options(Some(1), vec![])) + .await + .unwrap(); + assert_eq!(plan.tasks.len(), 2); + + let mut results = Vec::new(); + for task in plan.compaction_tasks() { + let task: CompactionTask = + serde_json::from_str(&serde_json::to_string(&task).unwrap()).unwrap(); + let result = task.execute(&dataset).await.unwrap(); + let result: RewriteResult = + serde_json::from_str(&serde_json::to_string(&result).unwrap()).unwrap(); + results.push(result); + } + commit_compaction( + &mut dataset, + results, + Arc::new(DatasetIndexRemapperOptions::default()), + plan.options(), + ) + .await + .unwrap(); + + assert_eq!(live_file_counts(&dataset), vec![1, 1]); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); +} + +/// A task serialized before `kind` existed rewrites its fragments. +#[test] +fn task_data_without_kind_rewrites_fragments() { + let task: TaskData = serde_json::from_str(r#"{"fragments": []}"#).unwrap(); + assert_eq!(task.kind, CompactionTaskKind::RewriteFragments); +} + +/// Run a repack of every fragment to the point of commit. +async fn stale_repack_results( + dataset: &Dataset, + options: &CompactionOptions, +) -> Vec { + let plan = plan_compaction(dataset, options).await.unwrap(); + let mut results = Vec::new(); + for task in plan.compaction_tasks() { + results.push(task.execute(dataset).await.unwrap()); + } + results +} + +async fn commit_results( + dataset: &mut Dataset, + results: Vec, +) -> Result { + commit_compaction( + dataset, + results, + Arc::new(DatasetIndexRemapperOptions::default()), + &CompactionOptions::default(), + ) + .await +} + +/// A delete that removes rows of a repacked fragment leaves the new files +/// aligned (a deletion only marks rows), so the repack still commits; one +/// that removes the whole fragment makes it retry. +#[rstest] +#[case::one_row("a = 1", true)] +#[case::whole_fragment("a < 4", false)] +#[tokio::test] +async fn repack_against_concurrent_delete(#[case] predicate: &str, #[case] commits: bool) { + let mut dataset = write_backfilled().await; + let stale = stale_repack_results(&dataset, &repack_options(Some(1), vec![])).await; + + dataset.delete(predicate).await.unwrap(); + let expected = dataset.scan().try_into_batch().await.unwrap(); + let version = dataset.manifest.version; + + let result = commit_results(&mut dataset, stale).await; + if commits { + result.unwrap(); + assert_eq!(live_file_counts(&dataset), vec![1, 1]); + } else { + let err = result.unwrap_err(); + assert!( + matches!(err, Error::RetryableCommitConflict { .. }), + "{err}" + ); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.manifest.version, version); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); + dataset.validate().await.unwrap(); +} + +/// A concurrent `drop_columns` of a column the repack moves makes it retry. +/// One that drops a column the repack leaves alone does not: the new file is +/// applied to the fragment as the drop left it. +#[rstest] +#[case::moved_column("d", false)] +#[case::other_column("b", true)] +#[tokio::test] +async fn repack_against_concurrent_drop(#[case] dropped: &str, #[case] commits: bool) { + let mut dataset = write_backfilled().await; + // Moves d and e into one file and leaves the base file alone. + let stale = stale_repack_results(&dataset, &repack_options(Some(2), vec![])).await; + + dataset.drop_columns(&[dropped]).await.unwrap(); + let expected = dataset.scan().try_into_batch().await.unwrap(); + let version = dataset.manifest.version; + + let result = commit_results(&mut dataset, stale).await; + if commits { + result.unwrap(); + assert_eq!(live_file_counts(&dataset), vec![2, 2]); + } else { + let err = result.unwrap_err(); + assert!( + matches!(err, Error::RetryableCommitConflict { .. }), + "{err}" + ); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.manifest.version, version); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); + dataset.validate().await.unwrap(); +} + +/// A concurrent update that writes new values into a moved column makes the +/// repack retry rather than publish the old values over the new ones. +#[tokio::test] +async fn repack_against_concurrent_update_of_moved_column() { + let mut dataset = write_backfilled().await; + let stale = stale_repack_results(&dataset, &repack_options(Some(1), vec![])).await; + + crate::dataset::UpdateBuilder::new(Arc::new(dataset.clone())) + .update_where("a = 2") + .unwrap() + .set("d", "100") + .unwrap() + .build() + .unwrap() + .execute() + .await + .unwrap(); + dataset.checkout_latest().await.unwrap(); + let expected = dataset.scan().try_into_batch().await.unwrap(); + + let err = commit_results(&mut dataset, stale).await.unwrap_err(); + assert!( + matches!(err, Error::RetryableCommitConflict { .. }), + "{err}" + ); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), expected); + dataset.validate().await.unwrap(); +} + +#[tokio::test] +async fn repack_skips_legacy_files() { + let dir = TempStrDir::default(); + let dataset = write(&dir, LanceFileVersion::Legacy).await; + let plan = plan_compaction(&dataset, &repack_options(None, vec![vec!["c"]])) + .await + .unwrap(); + assert!(plan.tasks.is_empty(), "{plan:?}"); +} + +#[test] +fn max_data_files_per_fragment_must_be_positive() { + let mut options = CompactionOptions { + max_data_files_per_fragment: Some(0), + ..Default::default() + }; + assert!(options.validate().is_err()); +} + +#[test] +fn repack_options_parse_from_config() { + let config = HashMap::from([( + "lance.compaction.max_data_files_per_fragment".to_string(), + "4".to_string(), + )]); + let options = CompactionOptions::from_dataset_config(&config).unwrap(); + assert_eq!(options.max_data_files_per_fragment, Some(4)); +} + +/// The file holding a fragment's spilled row lineage keeps its columns: the +/// new files carry no lineage, so emptying that file would strand it. +#[test] +fn repack_keeps_the_spilled_lineage_file() { + use lance_table::format::{ROW_ID_FIELD_ID, RowIdMeta}; + let schema = lance_core::datatypes::Schema::try_from(&Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ])) + .unwrap(); + let mut fragment = Fragment::new(0); + fragment.add_file( + "lineage.lance", + vec![0, ROW_ID_FIELD_ID], + vec![0, 1], + lance_file::version::ConcreteFileVersion::V2_0, + None, + ); + for (path, field) in [("b.lance", 1), ("c.lance", 2)] { + fragment.add_file( + path, + vec![field], + vec![0], + lance_file::version::ConcreteFileVersion::V2_0, + None, + ); + } + fragment.row_id_meta = Some(RowIdMeta::Column); + + // Above the limit, the lineage file's column stays and the others move. + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 3, + None, + Some(2) + ), + Some(vec![vec![1, 2]]) + ); + // With a limit of 1 as well: moving every column would leave the lineage + // file behind, so its column stays. + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 3, + None, + Some(1) + ), + Some(vec![vec![1, 2]]) + ); +} + +/// Over the file limit with groups set, the columns no group names are merged +/// into one file as well. +#[tokio::test] +async fn repack_merges_unclaimed_columns_over_the_limit() { + let mut dataset = write_backfilled().await; + let before = dataset.scan().try_into_batch().await.unwrap(); + + compact_files(&mut dataset, repack_options(Some(2), vec![vec!["e"]]), None) + .await + .unwrap(); + + for files in layout(&dataset) { + let mut files = files + .into_iter() + .map(|fields| { + fields + .into_iter() + .filter(|id| *id != TOMBSTONE_FIELD_ID) + .collect::>() + }) + .collect::>(); + files.sort(); + assert_eq!(files, vec![vec![0, 1, 2, 3], vec![4]], "{files:?}"); + } + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); +} + +fn lance_schema(arrow: Schema) -> lance_core::datatypes::Schema { + lance_core::datatypes::Schema::try_from(&arrow).unwrap() +} + +fn v2_0_file(path: &str, fields: Vec) -> lance_table::format::DataFile { + let indices = (0..fields.len() as i32).collect(); + lance_table::format::DataFile::new( + path, + fields, + indices, + lance_file::version::ConcreteFileVersion::V2_0, + None, + None, + ) +} + +/// A V2.0 file keeps a struct's header after the struct's last child in it +/// is dropped. That header holds no data, so the file does not hold the +/// struct: the struct sits alone in the file with its live child, and no +/// repack is planned (one would be planned again on every run). +#[test] +fn repack_ignores_struct_header_without_children() { + use arrow_schema::Fields; + let arrow = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new( + "s", + DataType::Struct(Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ])), + true, + ), + ]); + // a=0, s=1, s.x=2 (dropped), s.y=3. + let schema = lance_schema(arrow).project_by_ids(&[0, 1, 3], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("a.lance", vec![0, 1, 2]), + v2_0_file("s.lance", vec![1, 3]), + ]; + + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 2, + Some(&[vec![1]]), + None + ), + None + ); + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 2, + None, + Some(1) + ), + Some(vec![vec![0, 1]]) + ); +} + +/// A blob column never moves, so with a limit of 1 the other columns are +/// merged next to it instead of nothing being planned. +#[test] +fn repack_merges_around_a_blob_column() { + let mut blob = Field::new("img", DataType::LargeBinary, true); + blob.set_metadata(std::collections::HashMap::from([( + lance_arrow::BLOB_META_KEY.to_string(), + "true".to_string(), + )])); + let schema = lance_schema(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + blob, + Field::new("d", DataType::Int32, true), + Field::new("e", DataType::Int32, true), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("base.lance", vec![0, 1]), + v2_0_file("d.lance", vec![2]), + v2_0_file("e.lance", vec![3]), + ]; + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 3, + None, + Some(1) + ), + Some(vec![vec![2, 3]]) + ); +} + +#[tokio::test] +async fn repack_is_not_planned_under_force_binary_copy() { + let dataset = write_backfilled().await; + let options = CompactionOptions { + compaction_mode: Some(CompactionMode::ForceBinaryCopy), + ..repack_options(Some(1), vec![]) + }; + let plan = plan_compaction(&dataset, &options).await.unwrap(); + assert!(plan.tasks.is_empty(), "{plan:?}"); +} + +/// `max_source_bytes` counts only the files a repack reads. +#[tokio::test] +async fn repack_counts_only_the_files_it_reads_against_the_byte_budget() { + let dataset = write_backfilled().await; + let fragment = &dataset.manifest.fragments[0]; + let size = |index: usize| fragment.files[index].file_size_bytes.get().unwrap().get(); + // Each repack leaves the base file alone and merges d and e, so a budget + // for two of those fits both, though not one whole fragment. + let options = CompactionOptions { + max_source_bytes: Some(2 * (size(1) + size(2))), + ..repack_options(Some(2), vec![]) + }; + assert!(size(0) + size(1) + size(2) > 2 * (size(1) + size(2))); + let plan = plan_compaction(&dataset, &options).await.unwrap(); + assert_eq!(plan.tasks.len(), 2, "{plan:?}"); +} + +/// A fragment both rewritten and repacked is refused before either commit. +#[tokio::test] +async fn commit_rejects_a_fragment_in_both_kinds_of_task() { + let mut dataset = write_backfilled().await; + let version = dataset.manifest.version; + let repacks = stale_repack_results(&dataset, &repack_options(Some(1), vec![])).await; + let rewrite = CompactionTask { + task: TaskData::rewrite_fragments(vec![dataset.manifest.fragments[0].clone()]), + read_version: version, + options: CompactionOptions::default(), + } + .execute(&dataset) + .await + .unwrap(); + + let mut results = repacks; + results.push(rewrite); + let err = commit_results(&mut dataset, results).await.unwrap_err(); + assert!(matches!(err, Error::InvalidInput { .. }), "{err}"); + dataset.checkout_latest().await.unwrap(); + assert_eq!(dataset.manifest.version, version); +} + +/// The lineage file and a blob file both stay as they are; the columns that +/// can move are still merged. +#[test] +fn repack_merges_around_a_blob_outside_the_kept_file() { + use lance_table::format::{ROW_ID_FIELD_ID, RowIdMeta}; + let mut blob = Field::new("img", DataType::LargeBinary, true); + blob.set_metadata(std::collections::HashMap::from([( + lance_arrow::BLOB_META_KEY.to_string(), + "true".to_string(), + )])); + let schema = lance_schema(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + blob, + Field::new("d", DataType::Int32, true), + Field::new("e", DataType::Int32, true), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("lineage.lance", vec![0, ROW_ID_FIELD_ID]), + v2_0_file("img.lance", vec![1]), + v2_0_file("d.lance", vec![2]), + v2_0_file("e.lance", vec![3]), + ]; + fragment.row_id_meta = Some(RowIdMeta::Column); + for max_files in [1, 2] { + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 4, + None, + Some(max_files) + ), + Some(vec![vec![2, 3]]), + "max_files={max_files}" + ); + } +} + +/// A file holding a column that cannot move is never partly emptied: moving +/// its other columns out would add a file instead of removing one. +#[test] +fn repack_never_adds_a_file() { + use lance_table::format::{ROW_ID_FIELD_ID, RowIdMeta}; + let blob = |name: &str| { + let mut field = Field::new(name, DataType::LargeBinary, true); + field.set_metadata(std::collections::HashMap::from([( + lance_arrow::BLOB_META_KEY.to_string(), + "true".to_string(), + )])); + field + }; + let plan = |schema: &lance_core::datatypes::Schema, fragment: &Fragment, max| { + crate::dataset::optimize::repack::plan_fragment_repack( + schema, + fragment, + fragment.files.len(), + None, + Some(max), + ) + }; + + // Lineage in one file, a blob next to a movable column in the other. + let schema = lance_schema(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + blob("img"), + Field::new("d", DataType::Int32, true), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("lineage.lance", vec![0, ROW_ID_FIELD_ID]), + v2_0_file("mixed.lance", vec![1, 2]), + ]; + fragment.row_id_meta = Some(RowIdMeta::Column); + assert_eq!(plan(&schema, &fragment, 1), None); + + // Two files, each a blob next to a movable column. + let schema = lance_schema(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + blob("img1"), + blob("img2"), + Field::new("cap", DataType::Utf8, true), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("base.lance", vec![0, 1]), + v2_0_file("extra.lance", vec![2, 3]), + ]; + assert_eq!(plan(&schema, &fragment, 1), None); + + // A blob next to a movable column, and two movable files: only those two + // merge. + let schema = lance_schema(Schema::new(vec![ + blob("img"), + Field::new("d", DataType::Int32, true), + Field::new("e", DataType::Int32, true), + Field::new("f", DataType::Int32, true), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("mixed.lance", vec![0, 1]), + v2_0_file("e.lance", vec![2]), + v2_0_file("f.lance", vec![3]), + ]; + assert_eq!(plan(&schema, &fragment, 2), Some(vec![vec![2, 3]])); +} + +/// A file holding a struct's header next to another column is emptied only +/// when the struct moves too, so a merge is planned around the file whose +/// staying leaves the others emptied. +#[test] +fn repack_moves_a_struct_with_its_header() { + use arrow_schema::Fields; + let plan = |schema: &lance_core::datatypes::Schema, fragment: &Fragment| { + crate::dataset::optimize::repack::plan_fragment_repack( + schema, + fragment, + fragment.files.len(), + None, + Some(2), + ) + }; + let child = |name: &str| Field::new(name, DataType::Int32, true); + + // a=0, s=1 {x=2 (dropped), y=3}, b=4, c=5. + let arrow = Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y")])), + true, + ), + child("b"), + child("c"), + ]); + let schema = lance_schema(arrow).project_by_ids(&[0, 1, 3, 4, 5], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("a.lance", vec![0, 1, 2]), + v2_0_file("s.lance", vec![1, 3, 5]), + v2_0_file("b.lance", vec![4]), + ]; + // Keeping s.lance would leave a.lance alive on the header; keeping + // a.lance lets s.lance and b.lance both go. + assert_eq!(plan(&schema, &fragment), Some(vec![vec![1, 4, 5]])); + + // a=0, s=1 {x=2, y=3, z=4} (x, z dropped), b=5, c=6: both a.lance and + // b.lance hold the header. + let arrow = Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y"), child("z")])), + true, + ), + child("b"), + child("c"), + ]); + let schema = lance_schema(arrow).project_by_ids(&[0, 1, 3, 5, 6], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("a.lance", vec![0, 1, 2]), + v2_0_file("b.lance", vec![5, 1, 4]), + v2_0_file("s.lance", vec![1, 3, 6]), + ]; + assert_eq!(plan(&schema, &fragment), Some(vec![vec![1, 5, 6]])); +} + +/// A fragment whose first file still holds a dropped struct child reaches one +/// file, and the next run plans nothing. +#[tokio::test] +async fn repack_with_a_dropped_struct_child_converges() { + use arrow_array::{ArrayRef, StructArray}; + use arrow_schema::Fields; + let fields = Fields::from(vec![ + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ]); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("s", DataType::Struct(fields.clone()), true), + ])); + let values = + |offset: i32| Arc::new(Int32Array::from_iter_values(offset..offset + 4)) as ArrayRef; + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + values(0), + Arc::new(StructArray::new(fields, vec![values(10), values(20)], None)), + ], + ) + .unwrap(); + let mut dataset = Dataset::write( + RecordBatchIterator::new([Ok(batch)], schema), + "memory://", + Some(WriteParams { + data_storage_version: Some(LanceFileVersion::V2_0), + ..Default::default() + }), + ) + .await + .unwrap(); + dataset + .add_columns( + NewColumnTransform::SqlExpressions(vec![("b".into(), "a + 1".into())]), + None, + None, + ) + .await + .unwrap(); + dataset.drop_columns(&["s.x"]).await.unwrap(); + let before = dataset.scan().try_into_batch().await.unwrap(); + + compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + assert_eq!(live_file_counts(&dataset), vec![1]); + assert_eq!(dataset.scan().try_into_batch().await.unwrap(), before); + dataset.validate().await.unwrap(); + + let version = dataset.manifest.version; + compact_files(&mut dataset, repack_options(Some(1), vec![]), None) + .await + .unwrap(); + assert_eq!(dataset.manifest.version, version); +} + +/// Struct-header layouts for the merges the other planner tests don't reach: +/// the unclaimed columns under `column_groups`, a struct split with the +/// lineage file, and the file kept when every file could go. +#[test] +fn repack_merges_only_files_it_empties() { + use arrow_schema::Fields; + use lance_table::format::{ROW_ID_FIELD_ID, RowIdMeta}; + let child = |name: &str| Field::new(name, DataType::Int32, true); + let plan = |schema: &lance_core::datatypes::Schema, + fragment: &Fragment, + groups: Option<&[Vec]>, + max| { + crate::dataset::optimize::repack::plan_fragment_repack( + schema, + fragment, + fragment.files.len(), + groups, + Some(max), + ) + }; + + // a=0, s=1 {x=2 (dropped), y=3, z=4 (dropped)}, b=5, groups [[s]]. + // a.lance and b.lance both keep s's header, so merging a and b would + // drop neither. + let arrow = Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y"), child("z")])), + true, + ), + child("b"), + ]); + let schema = lance_schema(arrow).project_by_ids(&[0, 1, 3, 5], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("a.lance", vec![0, 1, 2]), + v2_0_file("b.lance", vec![5, 1, 4]), + v2_0_file("s.lance", vec![1, 3]), + ]; + assert_eq!(plan(&schema, &fragment, Some(&[vec![1]]), 2), None); + + // s=0 {y=1, z=2}, b=3, c=4. s is split between the lineage file and + // c1.lance; moving s would empty the lineage file's columns, so only b + // and c merge. + let arrow = Schema::new(vec![ + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("y"), child("z")])), + true, + ), + child("b"), + child("c"), + ]); + let schema = lance_schema(arrow); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("lineage.lance", vec![0, 1, ROW_ID_FIELD_ID]), + v2_0_file("c1.lance", vec![0, 2]), + v2_0_file("b.lance", vec![3]), + v2_0_file("c.lance", vec![4]), + ]; + fragment.row_id_meta = Some(RowIdMeta::Column); + assert_eq!(plan(&schema, &fragment, None, 2), Some(vec![vec![3, 4]])); + + // a=0, s=1 {x=2, y=3}, b=4, c=5. k.lance is the largest and shares s with + // m1.lance; keeping k leaves s and m1 alone and merges b and c. + let arrow = Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y")])), + true, + ), + child("b"), + child("c"), + ]); + let schema = lance_schema(arrow); + let sized = |path: &str, fields: Vec, size: u64| { + let mut file = v2_0_file(path, fields); + file.file_size_bytes = lance_io::utils::CachedFileSize::new(size); + file + }; + let mut fragment = Fragment::new(0); + fragment.files = vec![ + sized("k.lance", vec![1, 2], 100), + sized("m1.lance", vec![0, 1, 3], 10), + sized("m2.lance", vec![4], 10), + sized("m3.lance", vec![5], 10), + ]; + assert_eq!(plan(&schema, &fragment, None, 2), Some(vec![vec![4, 5]])); +} + +/// When every file holds the struct, the largest file is kept with its other +/// column and the struct's other files merge; and when only some files share +/// a column with the largest, the merge falls back to every other file. +#[test] +fn repack_keeps_the_largest_file_that_shares_a_struct() { + use arrow_schema::Fields; + let child = |name: &str| Field::new(name, DataType::Int32, true); + let sized = |path: &str, fields: Vec, size: u64| { + let mut file = v2_0_file(path, fields); + file.file_size_bytes = lance_io::utils::CachedFileSize::new(size); + file + }; + let plan = |schema: &lance_core::datatypes::Schema, fragment: &Fragment| { + crate::dataset::optimize::repack::plan_fragment_repack( + schema, + fragment, + fragment.files.len(), + None, + Some(2), + ) + }; + // a=0, s=1 {x=2, y=3, z=4}, b=5. + let schema = lance_schema(Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y"), child("z")])), + true, + ), + child("b"), + ])); + + let mut fragment = Fragment::new(0); + fragment.files = vec![ + sized("k.lance", vec![0, 1, 2], 100), + sized("m1.lance", vec![1, 3], 10), + sized("m2.lance", vec![1, 4], 10), + ]; + assert_eq!(plan(&schema, &fragment), Some(vec![vec![1]])); + + fragment.files = vec![ + sized("k.lance", vec![0, 1, 2], 100), + sized("m1.lance", vec![1, 3, 4], 10), + sized("m2.lance", vec![5], 10), + ]; + assert_eq!(plan(&schema, &fragment), Some(vec![vec![1, 5]])); +} + +/// A V2.0 file holding only a struct header holds no column, so it does not +/// stop the largest file from being kept; it goes when the struct moves. +#[test] +fn repack_keeps_the_largest_file_next_to_a_header_only_file() { + use arrow_schema::Fields; + let child = |name: &str| Field::new(name, DataType::Int32, true); + let sized = |path: &str, fields: Vec, size: u64| { + let mut file = v2_0_file(path, fields); + file.file_size_bytes = lance_io::utils::CachedFileSize::new(size); + file + }; + // a=0, s=1 {x=2 (dropped), y=3}, b=4. + let schema = lance_schema(Schema::new(vec![ + child("a"), + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y")])), + true, + ), + child("b"), + ])) + .project_by_ids(&[0, 1, 3, 4], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + sized("h.lance", vec![1, 2], 10), + sized("base.lance", vec![0], 100), + sized("sy.lance", vec![1, 3], 10), + sized("bf.lance", vec![4], 10), + ]; + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 4, + None, + Some(2) + ), + Some(vec![vec![1, 4]]) + ); +} + +/// A merge that leaves the largest file in place wins over one that merges +/// it away, even when a smaller file is the one kept. +#[test] +fn repack_prefers_leaving_the_largest_file_in_place() { + use arrow_schema::Fields; + let child = |name: &str| Field::new(name, DataType::Int32, true); + let sized = |path: &str, fields: Vec, size: u64| { + let mut file = v2_0_file(path, fields); + file.file_size_bytes = lance_io::utils::CachedFileSize::new(size); + file + }; + // s=0 {x=1, y=2, z=3}, t=4 {p=5, q=6}, c=7. + let schema = lance_schema(Schema::new(vec![ + Field::new( + "s", + DataType::Struct(Fields::from(vec![child("x"), child("y"), child("z")])), + true, + ), + Field::new( + "t", + DataType::Struct(Fields::from(vec![child("p"), child("q")])), + true, + ), + child("c"), + ])); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + sized("k.lance", vec![0, 1, 4, 5], 100), + sized("m1.lance", vec![0, 2, 7], 10), + sized("m2.lance", vec![0, 3], 10), + sized("m3.lance", vec![4, 6], 10), + ]; + // Keeping m3 merges m1 and m2 (s and c) and leaves k holding t. + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 4, + None, + Some(3) + ), + Some(vec![vec![0, 7]]) + ); +} + +/// A fragment whose files all hold only a struct header has no column to +/// move, and nothing is planned. +#[test] +fn repack_plans_nothing_for_header_only_files() { + use arrow_schema::Fields; + let child = |name: &str| Field::new(name, DataType::Int32, true); + // s=0 {a=1, b=2, c=3, d=4}; only d is left. + let schema = lance_schema(Schema::new(vec![Field::new( + "s", + DataType::Struct(Fields::from(vec![ + child("a"), + child("b"), + child("c"), + child("d"), + ])), + true, + )])) + .project_by_ids(&[0, 4], false); + let mut fragment = Fragment::new(0); + fragment.files = vec![ + v2_0_file("a.lance", vec![0, 1]), + v2_0_file("b.lance", vec![0, 2]), + v2_0_file("c.lance", vec![0, 3]), + ]; + assert_eq!( + crate::dataset::optimize::repack::plan_fragment_repack( + &schema, + &fragment, + 3, + None, + Some(2) + ), + None + ); +} diff --git a/rust/lance/src/dataset/rowids/spill.rs b/rust/lance/src/dataset/rowids/spill.rs index e9e5fff33a2..72f04f28720 100644 --- a/rust/lance/src/dataset/rowids/spill.rs +++ b/rust/lance/src/dataset/rowids/spill.rs @@ -2149,6 +2149,7 @@ mod tests { let read_version = dataset.manifest.version; let operation = Operation::DataReplacement { replacements: vec![replacement], + data_change: true, }; dataset = Dataset::commit( Arc::new(dataset), diff --git a/rust/lance/src/dataset/schema_evolution.rs b/rust/lance/src/dataset/schema_evolution.rs index b57f12c80b4..d177326394f 100644 --- a/rust/lance/src/dataset/schema_evolution.rs +++ b/rust/lance/src/dataset/schema_evolution.rs @@ -3600,6 +3600,7 @@ mod test { dataset.version().version, Operation::DataReplacement { replacements: vec![replacement], + data_change: true, }, None, ); diff --git a/rust/lance/src/dataset/tests/data_file_part.rs b/rust/lance/src/dataset/tests/data_file_part.rs index 35ae388c7c4..fda33abde5d 100644 --- a/rust/lance/src/dataset/tests/data_file_part.rs +++ b/rust/lance/src/dataset/tests/data_file_part.rs @@ -93,6 +93,7 @@ async fn commit(dataset: &Dataset, replacement: DataReplacementGroup) -> Result< WriteDestination::Dataset(Arc::new(dataset.clone())), Operation::DataReplacement { replacements: vec![replacement], + data_change: true, }, Some(dataset.version_id()), None, diff --git a/rust/lance/src/dataset/tests/dataset_index.rs b/rust/lance/src/dataset/tests/dataset_index.rs index cc44fced886..6dbcf5f2dbc 100644 --- a/rust/lance/src/dataset/tests/dataset_index.rs +++ b/rust/lance/src/dataset/tests/dataset_index.rs @@ -3043,6 +3043,7 @@ async fn test_partial_compound_hybrid_prunes_same_path_different_base_rewrite() WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(1, replacement_file)], + data_change: true, }, Some(read_version), None, diff --git a/rust/lance/src/dataset/tests/dataset_merge_update.rs b/rust/lance/src/dataset/tests/dataset_merge_update.rs index 2a5a148fe79..fbf082a4276 100644 --- a/rust/lance/src/dataset/tests/dataset_merge_update.rs +++ b/rust/lance/src/dataset/tests/dataset_merge_update.rs @@ -690,6 +690,7 @@ async fn test_datafile_replacement() { WriteDestination::Dataset(dataset.clone()), Operation::DataReplacement { replacements: vec![], + data_change: true, }, Some(1), None, @@ -719,6 +720,7 @@ async fn test_datafile_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![], + data_change: true, }, Some(3), None, @@ -773,6 +775,7 @@ async fn test_datafile_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(4), None, @@ -890,6 +893,7 @@ async fn test_datafile_partial_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(3), None, @@ -951,6 +955,7 @@ async fn test_datafile_partial_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(4), None, @@ -1056,6 +1061,7 @@ async fn test_datafile_replacement_error() { WriteDestination::Dataset(Arc::new(dataset.clone())), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, // read at the current version (after the Merge above) Some(dataset.manifest.version), @@ -2621,6 +2627,7 @@ async fn test_data_replacement_advances_row_lineage() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(read_version), None, @@ -2730,6 +2737,7 @@ async fn test_data_replacement_invalidates_index_bitmap() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(read_version), None, @@ -3014,6 +3022,7 @@ async fn test_data_replacement_populates_invalidated_bitmap() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, new_data_file)], + data_change: true, }, Some(read_version), None, @@ -3133,6 +3142,7 @@ async fn test_fts_stale_entries_after_data_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(1, new_data_file)], + data_change: true, }, Some(read_version), None, @@ -3273,6 +3283,7 @@ async fn test_cross_column_fast_search_blocks_column_local_stale_postings() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(1, replacement_file)], + data_change: true, }, Some(read_version), None, @@ -3444,6 +3455,7 @@ async fn test_vector_index_after_data_replacement() { WriteDestination::Dataset(Arc::new(dataset)), Operation::DataReplacement { replacements: vec![DataReplacementGroup(frag1_id, new_data_file)], + data_change: true, }, Some(read_version), None, diff --git a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs index f6eb0886336..6c78aa368ea 100644 --- a/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs +++ b/rust/lance/src/dataset/tests/dataset_overlay_index_masking.rs @@ -255,6 +255,45 @@ async fn test_overlay_stale_drop_and_new_match(#[values(false, true)] stable_row assert_eq!(ids_matching(&dataset, "age = 20").await, vec![2]); } +/// A column repack moves base values without changing them, so an overlay on +/// the fragment keeps shadowing the base it was written over. +#[tokio::test] +async fn repack_columns_preserves_data_overlay() { + let dataset = create_base_dataset().await; + // Fragment 0, offset 1 is id=1, age=10; the overlay changes its age to 999. + let mut dataset = commit_overlay( + dataset, + "age_overlay", + 0, + &[1], + OverlayCoverage::dense(RoaringBitmap::from_iter([1])), + vec![i32_array([Some(999)])], + ) + .await; + assert_eq!(ids_matching(&dataset, "age = 999").await, vec![1]); + + // Split `id` into a file of its own; the overlay on `age` must survive. + let options = crate::dataset::optimize::CompactionOptions { + column_groups: vec![vec!["id".into()]], + scope: crate::dataset::optimize::CompactionScope::RepackColumns, + ..Default::default() + }; + crate::dataset::optimize::compact_files(&mut dataset, options, None) + .await + .unwrap(); + assert_eq!(dataset.get_fragment(0).unwrap().metadata().files.len(), 2); + + assert_eq!( + dataset.get_fragment(0).unwrap().metadata().overlays.len(), + 1, + "the data overlay must survive the rewrite" + ); + // The overlaid value still shadows the base: age=999 for id=1, old 10 gone. + assert_eq!(ids_matching(&dataset, "age = 999").await, vec![1]); + assert_eq!(ids_matching(&dataset, "age = 10").await, Vec::::new()); + dataset.validate().await.unwrap(); +} + /// Row-level BTree precision: when one row in a covered fragment is stale, only that row is /// blocked from the index result and re-evaluated on the stale-Take path. Non-stale rows in /// the same fragment (including one that matches the predicate) remain on the indexed path. diff --git a/rust/lance/src/dataset/tests/fragment_write_columns.rs b/rust/lance/src/dataset/tests/fragment_write_columns.rs index ad0e184d721..4d054b53d02 100644 --- a/rust/lance/src/dataset/tests/fragment_write_columns.rs +++ b/rust/lance/src/dataset/tests/fragment_write_columns.rs @@ -97,7 +97,10 @@ async fn commit(dataset: &Dataset, replacements: Vec) -> R let read_version = dataset.manifest.version; Dataset::commit( WriteDestination::Dataset(Arc::new(dataset.clone())), - Operation::DataReplacement { replacements }, + Operation::DataReplacement { + replacements, + data_change: true, + }, Some(read_version), None, None, diff --git a/rust/lance/src/dataset/versions/mod.rs b/rust/lance/src/dataset/versions/mod.rs index ce57af0e13e..56566f989f5 100644 --- a/rust/lance/src/dataset/versions/mod.rs +++ b/rust/lance/src/dataset/versions/mod.rs @@ -105,7 +105,21 @@ pub fn schema_compare_options(version: ConcreteFileVersion) -> SchemaCompareOpti } } -async fn create_seed_writers( +/// Reject a write schema whose blob columns don't match the file version, the +/// same check [`write_fragments`] runs — exposed for callers that go straight to +/// [`write_fragments_direct`] with a schema they built themselves. +pub(super) fn validate_write_schema(version: ConcreteFileVersion, schema: &Schema) -> Result<()> { + match version { + ConcreteFileVersion::V1 | ConcreteFileVersion::V2_0 | ConcreteFileVersion::V2_1 => { + write::validate_legacy_blob_write_schema(schema, &format!("{version:?}")) + } + ConcreteFileVersion::V2_2 | ConcreteFileVersion::V2_3 => { + write::validate_blob_v2_write_schema(schema) + } + } +} + +pub(super) async fn create_seed_writers( version: ConcreteFileVersion, dataset: Option<&Dataset>, params: &WriteParams, diff --git a/rust/lance/src/index.rs b/rust/lance/src/index.rs index 70209f12ccf..9366343ce95 100644 --- a/rust/lance/src/index.rs +++ b/rust/lance/src/index.rs @@ -13749,6 +13749,7 @@ mod tests { .filter(|fragment| covered.contains(fragment.id as u32)) .cloned() .collect(), + kind: Default::default(), }], read_version: dataset.version().version, options: CompactionOptions::default(), @@ -13863,9 +13864,9 @@ mod tests { // fragment 2 when the remap is missing, which would leave this nothing // to ask about and hide whether the guard is sensitive on its own. let plan = CompactionPlan { - tasks: vec![TaskData { - fragments: dataset.fragments().as_ref().clone(), - }], + tasks: vec![TaskData::rewrite_fragments( + dataset.fragments().as_ref().clone(), + )], read_version: dataset.version().version, options: CompactionOptions::default(), }; diff --git a/rust/lance/src/index/append.rs b/rust/lance/src/index/append.rs index 00bae3e0422..c0dc254616d 100644 --- a/rust/lance/src/index/append.rs +++ b/rust/lance/src/index/append.rs @@ -243,6 +243,10 @@ pub async fn build_per_segment_filters( /// Attempt to read seed buffers for `column_name` from `fragments`' data files. /// +/// A seed is written into the data file that holds the column, which after a +/// `column_groups` compaction need not be the fragment's first file, so each +/// fragment's file is found by `field_id`. +/// /// Returns `Some(vec)` only if every fragment has a seed entry; returns `None` /// if any fragment is missing a seed or its data file cannot be opened. /// Index-type-specific validation (e.g. `rows_per_zone` checks) is left to the @@ -251,6 +255,7 @@ async fn try_harvest_seeds( dataset: &Dataset, fragments: &[Fragment], column_name: &str, + field_id: i32, ) -> Result>> { if fragments.is_empty() { return Ok(Some(Vec::new())); @@ -265,7 +270,11 @@ async fn try_harvest_seeds( ); for fragment in fragments { - let Some(data_file) = fragment.files.first() else { + let Some(data_file) = fragment + .files + .iter() + .find(|file| file.fields.contains(&field_id)) + else { return Ok(None); }; @@ -534,7 +543,13 @@ async fn merge_scalar_indices<'a>( // Only open data files looking for seeds when the plugin confirms this // index type and configuration can actually produce them. let maybe_created = if plugin.might_use_seeds(&index_details) { - if let Some(seeds) = try_harvest_seeds(dataset.as_ref(), unindexed, column_name).await? + if let Some(seeds) = try_harvest_seeds( + dataset.as_ref(), + unindexed, + column_name, + reference_idx.fields[0], + ) + .await? { plugin .update_from_seeds(seeds, reference_index.clone(), &index_details, &new_store) @@ -4945,4 +4960,62 @@ mod tests { .num_rows(); assert_eq!(total, 150, "no rows may be lost across compaction + merge"); } + + /// A `column_groups` compaction writes an indexed column's seed into that + /// column's group file, which need not be the fragment's first file; the + /// harvest has to find it there. + #[tokio::test] + async fn test_harvest_seeds_from_column_group_file() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(Int32Array::from_iter_values(0..20)), + Arc::new(Int32Array::from_iter_values(0..20)), + ], + ) + .unwrap(); + let reader = RecordBatchIterator::new([Ok(batch)], schema); + let mut dataset = Dataset::write( + reader, + "memory://", + Some(WriteParams { + max_rows_per_file: 10, + ..Default::default() + }), + ) + .await + .unwrap(); + let params = ScalarIndexParams::for_builtin(BuiltinIndexType::ZoneMap) + .with_params(&serde_json::json!({"use_seeds": true})); + dataset + .create_index(&["a"], IndexType::ZoneMap, None, ¶ms, false) + .await + .unwrap(); + // {b} is unclaimed and becomes the first file; {a} is the second. + let options = CompactionOptions { + column_groups: vec![vec!["a".into()]], + ..Default::default() + }; + compact_files(&mut dataset, options, None).await.unwrap(); + + let fragments: Vec = dataset + .get_fragments() + .iter() + .map(|fragment| fragment.metadata().clone()) + .collect(); + let a_id = dataset.schema().field("a").unwrap().id; + assert!( + !fragments[0].files[0].fields.contains(&a_id), + "the indexed column is not in the first file: {:?}", + fragments[0].files + ); + let seeds = try_harvest_seeds(&dataset, &fragments, "a", a_id) + .await + .unwrap(); + assert!(seeds.is_some(), "the seed in the group file is found"); + } } diff --git a/rust/lance/src/index/frag_reuse.rs b/rust/lance/src/index/frag_reuse.rs index 25f260229f1..191a37d5b15 100644 --- a/rust/lance/src/index/frag_reuse.rs +++ b/rust/lance/src/index/frag_reuse.rs @@ -3410,6 +3410,7 @@ mod tests { read_version: dataset.manifest.version, original_fragments: vec![overlaid], row_addrs: Some(serialized), + repacked_files: None, }; let fragments_before: Vec = dataset.fragments().iter().map(|f| f.id).collect(); let error = crate::dataset::optimize::commit_compaction( diff --git a/rust/lance/src/io/commit.rs b/rust/lance/src/io/commit.rs index cbec4753876..ce0be3efc5d 100644 --- a/rust/lance/src/io/commit.rs +++ b/rust/lance/src/io/commit.rs @@ -1487,7 +1487,8 @@ async fn prepare_attempt( Operation::Update { fields_modified, .. } => !fields_modified.is_empty(), - Operation::Merge { .. } | Operation::DataReplacement { .. } => true, + Operation::Merge { .. } => true, + Operation::DataReplacement { data_change, .. } => *data_change, _ => false, }; let ledger = if may_rewrite_in_place { diff --git a/rust/lance/src/io/commit/conflict_resolver.rs b/rust/lance/src/io/commit/conflict_resolver.rs index 1538d0a2185..3dc2c3b4173 100644 --- a/rust/lance/src/io/commit/conflict_resolver.rs +++ b/rust/lance/src/io/commit/conflict_resolver.rs @@ -107,14 +107,18 @@ fn may_alter_nullability(operation: &Operation) -> bool { /// 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. +/// only removes rows, and `Rewrite` and a `DataReplacement` without +/// `data_change` preserve the values they move; `Project` already conflicts +/// with projections and merges elsewhere. fn supplies_values(operation: &Operation) -> bool { matches!( operation, Operation::Append { .. } | Operation::Update { .. } - | Operation::DataReplacement { .. } + | Operation::DataReplacement { + data_change: true, + .. + } | Operation::DataOverlay { .. } ) } @@ -256,7 +260,7 @@ impl<'a> TransactionRebase<'a> { reuse, }) } - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { replacements, .. } => { let modified_fragment_ids = replacements.iter().map(|r| r.0).collect::>(); let initial_fragments = @@ -348,6 +352,19 @@ impl<'a> TransactionRebase<'a> { ) } + /// Whether this transaction is a replacement whose values only moved + /// (compaction). Losing its target to a concurrent drop or delete costs + /// nothing a fresh plan cannot redo, so it retries rather than fails. + fn only_moves_values(&self) -> bool { + matches!( + self.transaction.operation, + Operation::DataReplacement { + data_change: false, + .. + } + ) + } + #[track_caller] fn data_replacement_target_removed_err( &self, @@ -355,6 +372,9 @@ impl<'a> TransactionRebase<'a> { other_transaction: &Transaction, other_version: u64, ) -> Error { + if self.only_moves_values() { + return self.retryable_conflict_err(other_transaction, other_version); + } Error::incompatible_transaction_source( format!( "DataReplacement target fragment {} was removed by concurrent {} at version {}.", @@ -372,6 +392,9 @@ impl<'a> TransactionRebase<'a> { other_transaction: &Transaction, other_version: u64, ) -> Error { + if self.only_moves_values() { + return self.retryable_conflict_err(other_transaction, other_version); + } Error::incompatible_transaction_source( format!( "DataReplacement target field {} in fragment {} was dropped by concurrent {} at version {}.", @@ -1177,7 +1200,12 @@ impl<'a> TransactionRebase<'a> { } } Operation::UpdateConfig { .. } => Ok(()), - Operation::DataReplacement { replacements } => { + // Values that only moved to new files leave what the index + // was built from unchanged. + Operation::DataReplacement { + data_change: false, .. + } => Ok(()), + Operation::DataReplacement { replacements, .. } => { // A data replacement only conflicts if it is updating a field the // index depends on -- whether keyed on or merely carried, since // `fields` lists both (see `IndexMetadata::covering_fields`). @@ -1388,7 +1416,7 @@ impl<'a> TransactionRebase<'a> { Ok(()) } } - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { replacements, .. } => { // These conflict if the rewrite touches any of the fragments being replaced. for replacement in replacements { for group in groups { @@ -1616,7 +1644,7 @@ impl<'a> TransactionRebase<'a> { other_transaction: &Transaction, other_version: u64, ) -> Result<()> { - if let Operation::DataReplacement { replacements } = &self.transaction.operation { + if let Operation::DataReplacement { replacements, .. } = &self.transaction.operation { match &other_transaction.operation { Operation::Append { .. } | Operation::Clone { .. } @@ -1715,6 +1743,12 @@ impl<'a> TransactionRebase<'a> { // the index's fragment bitmap, which would lead to fewer conflicts. However // this would introduce fragment bitmaps with holes which may not be well tested // yet. For now, we don't allow this case. + // + // Values that only moved still match what the index was + // built from, so those never conflict. + if self.only_moves_values() { + return Ok(()); + } let newly_depended_fields = new_indices .iter() .flat_map(|idx| idx.fields.iter()) @@ -1747,6 +1781,7 @@ impl<'a> TransactionRebase<'a> { } Operation::DataReplacement { replacements: other_replacements, + .. } => { // These conflict if there is overlap in fragment id && fields. for replacement in replacements { @@ -4295,6 +4330,7 @@ mod tests { 1, DataFile::new_legacy_from_fields("r.lance", vec![0], None), )], + data_change: true, }, Compatible, ), @@ -4693,6 +4729,7 @@ mod tests { "replacement", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, file())], + data_change: true, }, ), ( @@ -5817,13 +5854,40 @@ mod tests { .map(|f| f.id) .chain(removed_fragment_ids.iter().copied()), ), - Operation::DataReplacement { replacements } => { + Operation::DataReplacement { replacements, .. } => { Box::new(replacements.iter().map(|r| r.0)) } Operation::DataOverlay { groups } => Box::new(groups.iter().map(|g| g.fragment_id)), } } + /// A schema with fields 0, 2, and 3: field 1 was dropped. + fn project_schema_without_field_1() -> lance_core::datatypes::Schema { + let arrow_schema = Schema::new( + ["a", "b", "c", "d"] + .map(|name| Field::new(name, DataType::Int32, true)) + .to_vec(), + ); + let schema = lance_core::datatypes::Schema::try_from(&arrow_schema).unwrap(); + schema.project_by_ids(&[0, 2, 3], true) + } + + fn index_on_fields(fields: Vec) -> IndexMetadata { + IndexMetadata { + uuid: Uuid::new_v4(), + name: "idx".to_string(), + fields, + covering_fields: vec![], + dataset_version: 1, + fragment_bitmap: Some(RoaringBitmap::from_iter([0u32])), + index_details: None, + index_version: 0, + created_at: None, + base_id: None, + files: None, + } + } + #[tokio::test] async fn test_conflicts_data_replacement() { use io::commit::conflict_resolver::tests::{ConflictResult::*, modified_fragment_ids}; @@ -5843,9 +5907,11 @@ mod tests { "Different fragments", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::DataReplacement { replacements: vec![DataReplacementGroup(1, data_file_frag1_fields01)], + data_change: true, }, Compatible, ), @@ -5853,9 +5919,11 @@ mod tests { "Same fragment, different fields", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields23)], + data_change: true, }, Compatible, ), @@ -5863,9 +5931,11 @@ mod tests { "Same fragment, same fields", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Retryable, ), @@ -5873,12 +5943,14 @@ mod tests { "Same fragment, overlapping fields", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::DataReplacement { replacements: vec![DataReplacementGroup( 0, DataFile::new_legacy_from_fields("path0_12", vec![1, 2], None), )], + data_change: true, }, Retryable, ), @@ -5886,6 +5958,7 @@ mod tests { "DataReplacement vs Rewrite on same fragment", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Rewrite { groups: vec![RewriteGroup { @@ -5901,6 +5974,7 @@ mod tests { "DataReplacement vs Rewrite on different fragment", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Rewrite { groups: vec![RewriteGroup { @@ -5919,6 +5993,7 @@ mod tests { "DataReplacement vs Update (RewriteColumns) on a different field", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![Fragment::new(0)], @@ -5938,6 +6013,7 @@ mod tests { "DataReplacement vs Update (RewriteColumns) with inserts on a different field", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![Fragment::new(0)], @@ -5956,6 +6032,7 @@ mod tests { "DataReplacement vs Update (RewriteColumns) that rewrote one of our fields", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![Fragment::new(0)], @@ -5974,6 +6051,7 @@ mod tests { "DataReplacement vs Update (RewriteRows) that moved our rows", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![Fragment::new(0)], @@ -5992,6 +6070,7 @@ mod tests { "DataReplacement vs Update that removed our fragment", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![], @@ -6010,6 +6089,7 @@ mod tests { "DataReplacement vs Update (RewriteRows) that moved a different fragment's rows", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Update { updated_fragments: vec![Fragment::new(1)], @@ -6028,6 +6108,7 @@ mod tests { "DataReplacement vs Delete (deletion-vector only) on same fragment", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Delete { deleted_fragment_ids: vec![], @@ -6040,6 +6121,7 @@ mod tests { "DataReplacement vs Delete that removes the fragment", Operation::DataReplacement { replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Delete { deleted_fragment_ids: vec![0], @@ -6052,7 +6134,8 @@ mod tests { ( "DataReplacement vs Merge", Operation::DataReplacement { - replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01)], + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, }, Operation::Merge { fragments: vec![Fragment::new(0)], @@ -6093,6 +6176,86 @@ mod tests { 0, DataFile::new_legacy_from_fields("path0_3", vec![3], None), )], + data_change: true, + }, + Retryable, + ), + // A replacement whose values only moved (compaction) retries where + // a data-changing one fails, and never conflicts with an index. + ( + "Moved DataReplacement vs Delete that removes the fragment", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: false, + }, + Operation::Delete { + deleted_fragment_ids: vec![0], + updated_fragments: vec![], + predicate: "a > 0".to_string(), + }, + Retryable, + ), + ( + "Moved DataReplacement vs Project that dropped one of its fields", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: false, + }, + Operation::Project { + schema: project_schema_without_field_1(), + preserves_nullability: true, + }, + Retryable, + ), + ( + "DataReplacement vs Project that dropped one of its fields", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: true, + }, + Operation::Project { + schema: project_schema_without_field_1(), + preserves_nullability: true, + }, + NotCompatible, + ), + ( + "Moved DataReplacement vs CreateIndex on one of its fields", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: false, + }, + Operation::CreateIndex { + new_indices: vec![index_on_fields(vec![1])], + removed_indices: vec![], + }, + Compatible, + ), + ( + // op1 is CreateIndex for the same reason as the case above. + "CreateIndex on a field vs moved DataReplacement of that field", + Operation::CreateIndex { + new_indices: vec![index_on_fields(vec![1])], + removed_indices: vec![], + }, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01.clone())], + data_change: false, + }, + Compatible, + ), + ( + "Moved DataReplacement vs DataReplacement of an overlapping field", + Operation::DataReplacement { + replacements: vec![DataReplacementGroup(0, data_file_frag0_fields01)], + data_change: false, + }, + Operation::DataReplacement { + replacements: vec![DataReplacementGroup( + 0, + DataFile::new_legacy_from_fields("path0_1", vec![1], None), + )], + data_change: true, }, Retryable, ),