1pub(crate) mod append;
31pub(crate) mod archival;
32mod create;
33mod keygen;
34pub mod lock;
35pub(crate) mod markers;
36pub(crate) mod metadata;
37mod metadata_compaction;
38mod rewrite;
39pub(crate) mod rollback;
40pub(crate) mod sizing;
41
42pub(crate) use append::append_batches;
43pub(crate) use append::set_commit_timezone;
44pub use create::TableCreateBuilder;
45pub(crate) use keygen::{
46 complex_keygen_validation_enabled, ensure_complex_keygen_validated,
47 ensure_supported_key_generator,
48};
49pub use lock::{InProcessLockProvider, LockLease, LockProvider};
50pub(crate) use rewrite::ensure_writable_table_version;
51pub use rewrite::{CompactionResult, UpsertOptions};
52pub(crate) use rewrite::{
53 await_async_compaction, compact_table, delete_filter, delete_keys,
54 dynamic_partition_overwrite_batches, maybe_compact_inline, overwrite_batches, update_filter,
55 upsert_batches,
56};
57
58use std::collections::HashMap;
59use std::str::FromStr;
60
61use apache_avro::to_avro_datum;
62use apache_avro::types::Value as AvroValue;
63
64use crate::Result;
65use crate::error::CoreError;
66use crate::file_group::log_file::writer::LogFileWriter;
67use crate::file_group::log_file::{BlockMetadataKey, BlockType};
68use crate::metadata::meta_field::MetaField;
69use crate::schema::delete::{
70 avro_schema_for_delete_record_list, delete_record_list_avro_schema_json,
71};
72use crate::storage::Storage;
73use crate::table::partition::PARTITION_METAFIELD_PREFIX;
74use crate::timeline::instant::{Action, Instant};
75
76pub(crate) async fn ensure_partition_metadata(
78 storage: &Storage,
79 partition_path: &str,
80 instant: &str,
81) -> Result<()> {
82 if partition_path.is_empty() {
83 return Ok(());
84 }
85 let meta_rel = format!("{partition_path}/{PARTITION_METAFIELD_PREFIX}");
86 if storage.exists(&meta_rel).await? {
87 return Ok(());
88 }
89 let depth = partition_path.matches('/').count() + 1;
90 let body = format!(
92 "#partition metadata\n\
93 commitTime={instant}\n\
94 partitionDepth={depth}\n"
95 );
96 storage.put_file(&meta_rel, body.into_bytes()).await?;
97 Ok(())
98}
99
100pub(crate) fn align_batch_to_schema(
106 batch: &arrow_array::RecordBatch,
107 target: &std::sync::Arc<arrow_schema::Schema>,
108) -> Option<arrow_array::RecordBatch> {
109 if &batch.schema() == target {
110 return Some(batch.clone());
111 }
112 if batch.num_columns() != target.fields().len() {
113 return None;
114 }
115 for (batch_field, target_field) in batch.schema().fields().iter().zip(target.fields()) {
116 if batch_field.name() != target_field.name()
117 || batch_field.data_type() != target_field.data_type()
118 {
119 return None;
120 }
121 }
122 arrow_array::RecordBatch::try_new(target.clone(), batch.columns().to_vec()).ok()
123}
124
125pub(crate) fn widen_to_table_nullability(
130 batches: &[arrow_array::RecordBatch],
131 table_schema: &arrow_schema::Schema,
132) -> Result<Vec<arrow_array::RecordBatch>> {
133 batches
134 .iter()
135 .map(|batch| {
136 let source = batch.schema();
137 let fields: Vec<arrow_schema::Field> = source
138 .fields()
139 .iter()
140 .map(|field| {
141 let table_nullable = table_schema
142 .field_with_name(field.name())
143 .is_ok_and(|table_field| table_field.is_nullable());
144 field
145 .as_ref()
146 .clone()
147 .with_nullable(field.is_nullable() || table_nullable)
148 })
149 .collect();
150 let schema = arrow_schema::Schema::new_with_metadata(fields, source.metadata().clone());
151 arrow_array::RecordBatch::try_new(std::sync::Arc::new(schema), batch.columns().to_vec())
152 .map_err(CoreError::ArrowError)
153 })
154 .collect()
155}
156
157pub(crate) fn strip_meta_columns(
160 batch: &arrow_array::RecordBatch,
161) -> Result<arrow_array::RecordBatch> {
162 let keep: Vec<usize> = (0..batch.num_columns())
163 .filter(|&i| !MetaField::is_meta_field(batch.schema().field(i).name()))
164 .collect();
165 batch.project(&keep).map_err(CoreError::ArrowError)
166}
167
168pub(crate) fn normalize_for_write(
178 batches: &[arrow_array::RecordBatch],
179) -> Result<Vec<arrow_array::RecordBatch>> {
180 batches.iter().map(normalize_write_batch).collect()
181}
182
183pub(crate) fn normalize_write_batch(
184 batch: &arrow_array::RecordBatch,
185) -> Result<arrow_array::RecordBatch> {
186 use arrow_schema::Field;
187 let batch = strip_meta_columns(batch)?;
188 let source = batch.schema();
189 let read_back = crate::schema::resolver::avro_json_to_arrow_schema(
190 &append::arrow_schema_to_avro_json(&source)?,
191 )?;
192 let mut changed = false;
193 let mut fields = Vec::with_capacity(batch.num_columns());
194 let mut columns = Vec::with_capacity(batch.num_columns());
195 for ((field, column), read_back) in source
196 .fields()
197 .iter()
198 .zip(batch.columns())
199 .zip(read_back.fields())
200 {
201 let target = read_back.data_type();
202 if target == field.data_type() {
203 fields.push(field.clone());
204 columns.push(column.clone());
205 continue;
206 }
207 changed = true;
208 let options = arrow::compute::CastOptions {
210 safe: false,
211 ..Default::default()
212 };
213 let casted =
214 arrow::compute::cast_with_options(column.as_ref(), target, &options).map_err(|e| {
215 CoreError::Schema(format!(
216 "cannot convert column '{}' from {} to {target} for writing: {e}",
217 field.name(),
218 field.data_type()
219 ))
220 })?;
221 fields.push(std::sync::Arc::new(
222 Field::new(field.name(), target.clone(), field.is_nullable())
223 .with_metadata(field.metadata().clone()),
224 ));
225 columns.push(casted);
226 }
227 if !changed {
228 return Ok(batch);
229 }
230 let schema = std::sync::Arc::new(arrow_schema::Schema::new_with_metadata(
231 fields,
232 source.metadata().clone(),
233 ));
234 arrow_array::RecordBatch::try_new(schema, columns)
235 .map_err(|e| CoreError::Schema(format!("cannot rebuild batch for writing: {e}")))
236}
237
238pub(crate) fn new_file_id() -> String {
239 format!("{}-0", uuid::Uuid::new_v4())
240}
241
242pub(crate) async fn fence_timeline_instant(
249 storage: &Storage,
250 timeline_dir: &str,
251 timestamp: &str,
252 action: Action,
253 requested_bytes: Vec<u8>,
254 inflight_bytes: Vec<u8>,
255) -> Result<()> {
256 let requested = Instant::from_str(&format!("{}.{}.requested", timestamp, action.as_ref()))
257 .map_err(|e| CoreError::Write(format!("invalid fencing instant: {e}")))?;
258 let inflight = if action == Action::Commit {
259 Instant::from_str(&format!("{timestamp}.inflight"))
260 } else {
261 Instant::from_str(&format!("{}.{}.inflight", timestamp, action.as_ref()))
262 }
263 .map_err(|e| CoreError::Write(format!("invalid fencing instant: {e}")))?;
264
265 storage
268 .put_file_if_absent(
269 &requested.relative_path_with_base(timeline_dir)?,
270 requested_bytes,
271 )
272 .await?;
273 storage
274 .put_file_if_absent(
275 &inflight.relative_path_with_base(timeline_dir)?,
276 inflight_bytes,
277 )
278 .await?;
279 Ok(())
280}
281
282pub(crate) fn write_task_parallelism(table: &crate::table::Table) -> usize {
286 table
287 .hudi_configs
288 .as_options()
289 .get("hoodie.write.task.parallelism")
290 .and_then(|v| v.parse().ok())
291 .filter(|n: &usize| *n >= 1)
292 .unwrap_or_else(|| {
293 2 * std::thread::available_parallelism()
294 .map(std::num::NonZeroUsize::get)
295 .unwrap_or(4)
296 })
297}
298
299pub(crate) async fn run_write_tasks<T: Send + 'static>(
302 tasks: Vec<futures::future::BoxFuture<'static, Result<T>>>,
303 parallelism: usize,
304) -> Vec<Result<T>> {
305 use futures::StreamExt;
306 let count = tasks.len();
307 let mut results: Vec<Option<Result<T>>> = std::iter::repeat_with(|| None).take(count).collect();
308 let indexed: Vec<_> = tasks
311 .into_iter()
312 .enumerate()
313 .map(|(index, task)| async move {
314 let result = AbortOnDrop(tokio::spawn(task))
315 .await
316 .unwrap_or_else(|e| Err(CoreError::Write(format!("write task failed: {e}"))));
317 (index, result)
318 })
319 .collect();
320 let mut stream = futures::stream::iter(indexed).buffer_unordered(parallelism.max(1));
321 while let Some((index, result)) = stream.next().await {
322 results[index] = Some(result);
323 }
324 results
325 .into_iter()
326 .map(|r| r.unwrap_or_else(|| Err(CoreError::Write("write task vanished".to_string()))))
327 .collect()
328}
329
330struct AbortOnDrop<T>(tokio::task::JoinHandle<T>);
333
334impl<T> Drop for AbortOnDrop<T> {
335 fn drop(&mut self) {
336 self.0.abort();
337 }
338}
339
340impl<T> std::future::Future for AbortOnDrop<T> {
341 type Output = std::result::Result<T, tokio::task::JoinError>;
342
343 fn poll(
344 mut self: std::pin::Pin<&mut Self>,
345 cx: &mut std::task::Context<'_>,
346 ) -> std::task::Poll<Self::Output> {
347 std::pin::Pin::new(&mut self.0).poll(cx)
348 }
349}
350
351pub(crate) struct FileTaskOutput {
353 pub size: i64,
354 pub ranges: Option<Vec<crate::metadata::table::column_stats::ColumnRangeStats>>,
356}
357
358pub(crate) fn parquet_file_task(
360 storage: std::sync::Arc<Storage>,
361 props: parquet::file::properties::WriterProperties,
362 bloom: Option<crate::bloom::BloomFilterConfig>,
363 prepared: Vec<arrow_array::RecordBatch>,
364 relative_path: String,
365 collect_ranges: bool,
366) -> futures::future::BoxFuture<'static, Result<FileTaskOutput>> {
367 Box::pin(async move {
368 let (bytes, ranges) = tokio::task::spawn_blocking(move || -> Result<_> {
369 let bytes =
370 crate::write::append::write_parquet_bytes_with_props(props, bloom, &prepared)?;
371 let ranges = if collect_ranges {
372 Some(
373 crate::metadata::table::column_stats::column_ranges_from_parquet_bytes(&bytes)?,
374 )
375 } else {
376 None
377 };
378 Ok((bytes, ranges))
379 })
380 .await
381 .map_err(|e| CoreError::Write(format!("write task panicked: {e}")))??;
382 let size = bytes.len() as i64;
383 storage.put_file(&relative_path, bytes).await?;
384 Ok(FileTaskOutput { size, ranges })
385 })
386}
387
388pub(crate) fn log_file_task(
391 storage: std::sync::Arc<Storage>,
392 props: parquet::file::properties::WriterProperties,
393 prepared: Vec<arrow_array::RecordBatch>,
394 relative_path: String,
395 instant: String,
396 schema_json: String,
397 collect_ranges: bool,
398) -> futures::future::BoxFuture<'static, Result<FileTaskOutput>> {
399 Box::pin(async move {
400 let (block, ranges) = tokio::task::spawn_blocking(move || -> Result<_> {
401 let parquet =
402 crate::write::append::write_parquet_bytes_with_props(props, None, &prepared)?;
403 let ranges = if collect_ranges {
404 Some(
405 crate::metadata::table::column_stats::column_ranges_from_parquet_bytes(
406 &parquet,
407 )?,
408 )
409 } else {
410 None
411 };
412 let block = LogFileWriter::write_log_block(
413 BlockType::ParquetData,
414 HashMap::from([
415 (BlockMetadataKey::InstantTime, instant),
416 (BlockMetadataKey::Schema, schema_json),
417 ]),
418 &parquet,
419 );
420 Ok((block, ranges))
421 })
422 .await
423 .map_err(|e| CoreError::Write(format!("write task panicked: {e}")))??;
424 let size = block.len() as i64;
425 storage.put_file(&relative_path, block).await?;
426 Ok(FileTaskOutput { size, ranges })
427 })
428}
429
430pub(crate) async fn post_complete_bookkeeping(
434 table: &crate::table::Table,
435 storage: &Storage,
436 instant: &str,
437) -> Result<()> {
438 markers::delete_marker_dir(storage, instant).await;
439 let (min_keep, max_keep) = archival::validated_keep_range(&table.hudi_configs)?;
440 archival::archive_timeline_if_needed(
444 storage,
445 &crate::write::append::timeline_dir(table),
446 &crate::timeline::loader::lsm_history_dir(&table.hudi_configs),
447 min_keep,
448 max_keep,
449 )
450 .await?;
451 Ok(())
455}
456
457pub(crate) fn inflight_commit_metadata_bytes(operation: &str, layout_two: bool) -> Result<Vec<u8>> {
461 let metadata = crate::metadata::commit::HoodieCommitMetadata {
462 version: Some(1),
463 operation_type: Some(operation.to_string()),
464 partition_to_write_stats: Some(std::collections::HashMap::new()),
465 compacted: Some(false),
466 extra_metadata: Some(std::collections::HashMap::new()),
467 partition_to_replace_file_ids: None,
468 };
469 if layout_two {
470 metadata.to_avro_bytes()
471 } else {
472 metadata.to_json_bytes()
473 }
474}
475
476pub(crate) fn build_delete_log_block(
481 instant: &str,
482 keys: &[(String, String)],
483 ordering_val: i64,
484) -> Result<Vec<u8>> {
485 let records = keys
486 .iter()
487 .map(|(record_key, partition_path)| {
488 AvroValue::Record(vec![
489 (
490 "recordKey".to_string(),
491 AvroValue::Union(1, Box::new(AvroValue::String(record_key.clone()))),
492 ),
493 (
494 "partitionPath".to_string(),
495 AvroValue::Union(1, Box::new(AvroValue::String(partition_path.clone()))),
496 ),
497 (
498 "orderingVal".to_string(),
499 if ordering_val == 0 {
506 AvroValue::Union(
507 2,
508 Box::new(AvroValue::Record(vec![(
509 "value".to_string(),
510 AvroValue::Int(0),
511 )])),
512 )
513 } else {
514 AvroValue::Union(
515 3,
516 Box::new(AvroValue::Record(vec![(
517 "value".to_string(),
518 AvroValue::Long(ordering_val),
519 )])),
520 )
521 },
522 ),
523 ])
524 })
525 .collect();
526 let payload = to_avro_datum(
527 avro_schema_for_delete_record_list()?,
528 AvroValue::Record(vec![(
529 "deleteRecordList".to_string(),
530 AvroValue::Array(records),
531 )]),
532 )?;
533 let mut content = Vec::with_capacity(8 + payload.len());
534 content.extend_from_slice(&3u32.to_be_bytes());
535 content.extend_from_slice(&(payload.len() as u32).to_be_bytes());
536 content.extend_from_slice(&payload);
537 Ok(LogFileWriter::write_log_block(
538 BlockType::Delete,
539 HashMap::from([
540 (BlockMetadataKey::InstantTime, instant.to_string()),
541 (
542 BlockMetadataKey::Schema,
543 delete_record_list_avro_schema_json().to_string(),
544 ),
545 ]),
546 &content,
547 ))
548}
549
550#[derive(Default)]
552pub(crate) struct MetadataTableUpdates {
553 pub files: Vec<(String, String, i64, bool)>,
555 pub record_index: Vec<crate::metadata::table::encode::RecordIndexEntry>,
556 pub column_stats: Vec<metadata::StatsFileUpdate>,
557 pub replaced_paths: Vec<String>,
563}
564
565pub(crate) fn stats_by_partition(
567 stats: Vec<crate::metadata::commit::HoodieWriteStat>,
568) -> HashMap<String, Vec<crate::metadata::commit::HoodieWriteStat>> {
569 let mut by_partition: HashMap<String, Vec<_>> = HashMap::new();
570 for stat in stats {
571 let partition = stat.partition_path.clone().unwrap_or_default();
572 by_partition.entry(partition).or_default().push(stat);
573 }
574 by_partition
575}
576
577pub(crate) fn commit_metadata_bytes(
583 table: &crate::table::Table,
584 operation: &str,
585 partition_to_write_stats: HashMap<String, Vec<crate::metadata::commit::HoodieWriteStat>>,
586 partition_to_replace_file_ids: Option<HashMap<String, Vec<String>>>,
587 schema: &arrow_schema::Schema,
588) -> Result<Vec<u8>> {
589 data_instant_metadata_bytes(
590 table,
591 operation,
592 partition_to_write_stats,
593 partition_to_replace_file_ids,
594 schema,
595 false,
596 )
597}
598
599pub(crate) fn compaction_commit_metadata_bytes(
603 table: &crate::table::Table,
604 partition_to_write_stats: HashMap<String, Vec<crate::metadata::commit::HoodieWriteStat>>,
605 schema: &arrow_schema::Schema,
606) -> Result<Vec<u8>> {
607 data_instant_metadata_bytes(
608 table,
609 "COMPACT",
610 partition_to_write_stats,
611 None,
612 schema,
613 true,
614 )
615}
616
617fn data_instant_metadata_bytes(
618 table: &crate::table::Table,
619 operation: &str,
620 partition_to_write_stats: HashMap<String, Vec<crate::metadata::commit::HoodieWriteStat>>,
621 partition_to_replace_file_ids: Option<HashMap<String, Vec<String>>>,
622 schema: &arrow_schema::Schema,
623 compacted: bool,
624) -> Result<Vec<u8>> {
625 let schema_json =
626 append::arrow_schema_to_avro_json(&append::strip_meta_fields_from_schema(schema))?;
627 let extra_metadata = Some(HashMap::from([("schema".to_string(), schema_json)]));
628 let layout_two = append::is_layout_two(table);
629 match partition_to_replace_file_ids {
630 Some(partition_to_replace_file_ids) => {
631 let metadata = crate::metadata::replace_commit::HoodieReplaceCommitMetadata {
632 version: Some(1),
633 operation_type: Some(operation.to_string()),
634 partition_to_write_stats: Some(partition_to_write_stats),
635 compacted: Some(compacted),
636 extra_metadata,
637 partition_to_replace_file_ids: Some(partition_to_replace_file_ids),
638 };
639 if layout_two {
640 metadata.to_avro_bytes()
641 } else {
642 metadata.to_json_bytes()
643 }
644 }
645 None => {
646 let metadata = crate::metadata::commit::HoodieCommitMetadata {
647 version: Some(1),
648 operation_type: Some(operation.to_string()),
649 partition_to_write_stats: Some(partition_to_write_stats),
650 compacted: Some(compacted),
651 extra_metadata,
652 partition_to_replace_file_ids: None,
653 };
654 if layout_two {
655 metadata.to_avro_bytes()
656 } else {
657 metadata.to_json_bytes()
658 }
659 }
660 }
661}
662
663pub(crate) async fn complete_write(
681 table: &mut crate::table::Table,
682 instant: &str,
683 action: Action,
684 commit_bytes: Vec<u8>,
685 schema: &arrow_schema::Schema,
686 mut updates: MetadataTableUpdates,
687 written_paths: &[String],
688) -> Result<CompletedWrite> {
689 let storage = table.file_system_view.storage.clone();
690 let mut mdt_stats = HashMap::new();
691 if table.is_metadata_table_enabled() {
692 let mdt_result = async {
693 if !updates.files.is_empty() {
694 mdt_stats.insert(
695 "files".to_string(),
696 metadata::update_files_partition_entries(
697 storage.as_ref(),
698 instant,
699 &updates.files,
700 )
701 .await?,
702 );
703 }
704 if crate::index::is_record_index_enabled(table) && !updates.record_index.is_empty() {
705 let stats = metadata::update_record_index(
706 storage.as_ref(),
707 instant,
708 std::mem::take(&mut updates.record_index),
709 write_task_parallelism(table),
710 )
711 .await?;
712 if !stats.is_empty() {
713 mdt_stats.insert("record_index".to_string(), stats);
714 }
715 }
716 if metadata::is_column_stats_enabled(table) && !updates.column_stats.is_empty() {
717 let stats = metadata::update_column_stats_partitions(
718 table,
719 instant,
720 &updates.column_stats,
721 schema,
722 &updates.replaced_paths,
723 )
724 .await?;
725 mdt_stats.extend(stats);
726 }
727 Ok::<(), CoreError>(())
728 }
729 .await;
730 if let Err(error) = mdt_result {
731 for path in written_paths {
732 let _ = storage.delete_file(path).await;
733 }
734 return Err(error);
735 }
736 }
737
738 let lock = lock::lock_provider_for(table);
739 let critical_section = lock.lock().await?;
740 if table.is_metadata_table_enabled()
741 && let Err(error) =
742 metadata::write_metadata_commit(storage.as_ref(), instant, mdt_stats).await
743 {
744 for path in written_paths {
745 let _ = storage.delete_file(path).await;
746 }
747 return Err(error);
748 }
749 let completion = if append::is_layout_two(table) {
750 Some(append::generate_instant_time().await)
751 } else {
752 None
753 };
754 let completed = Instant::new_completed(instant.to_string(), action, completion.clone())?;
755 let commit_relative_path = completed.relative_path_with_base(&append::timeline_dir(table))?;
756 storage
757 .put_file_if_absent(&commit_relative_path, commit_bytes)
758 .await?;
759 if let Err(error) = post_complete_bookkeeping(table, storage.as_ref(), instant).await {
761 log::warn!("post-commit bookkeeping for {instant} failed: {error}");
762 }
763 if let Err(error) = metadata_compaction::compact_metadata_table_if_needed(table).await {
766 log::warn!("metadata table compaction after {instant} failed: {error}");
767 }
768 drop(critical_section);
769
770 let _ = table.timeline.reload_completed_commits().await;
771 table.clear_read_caches();
772 Ok(CompletedWrite {
773 completion_time: completion,
774 })
775}
776
777#[derive(Clone, Debug, Default, PartialEq, Eq)]
779pub struct WriteResult {
780 pub requested_time: Option<String>,
783 pub completion_time: Option<String>,
787 pub num_writes: usize,
791 pub num_updates: usize,
793 pub num_inserts: usize,
796 pub num_deletes: usize,
798}
799
800pub(crate) struct CompletedWrite {
802 pub completion_time: Option<String>,
805}
806
807#[cfg(test)]
808mod tests {
809 use super::*;
810
811 #[tokio::test]
814 async fn test_fence_timeline_instant_conflicts_on_same_timestamp() {
815 let dir = tempfile::tempdir().unwrap();
816 let base_url = url::Url::from_directory_path(dir.path()).unwrap();
817 let storage = Storage::new_with_base_url(base_url).unwrap();
818
819 fence_timeline_instant(
820 &storage,
821 ".hoodie/timeline",
822 "20260101000000000",
823 Action::Commit,
824 Vec::new(),
825 Vec::new(),
826 )
827 .await
828 .unwrap();
829
830 let err = fence_timeline_instant(
831 &storage,
832 ".hoodie/timeline",
833 "20260101000000000",
834 Action::Commit,
835 Vec::new(),
836 Vec::new(),
837 )
838 .await
839 .unwrap_err();
840 assert!(
841 err.to_string().contains("already exists"),
842 "expected instant conflict, got: {err}"
843 );
844 }
845
846 #[test]
850 fn test_build_delete_log_block_both_ordering_branches() {
851 for ordering in [0i64, 42i64] {
852 let block = build_delete_log_block(
853 "20260101000000000",
854 &[("k1".to_string(), "p".to_string())],
855 ordering,
856 )
857 .unwrap();
858 assert!(!block.is_empty());
859 assert_eq!(&block[..6], b"#HUDI#");
860 }
861 }
862
863 #[tokio::test]
866 async fn test_ensure_partition_metadata_idempotent() {
867 let dir = tempfile::tempdir().unwrap();
868 let base_url = url::Url::from_directory_path(dir.path()).unwrap();
869 let storage = Storage::new_with_base_url(base_url).unwrap();
870 ensure_partition_metadata(&storage, "", "20260101000000000")
871 .await
872 .unwrap();
873 ensure_partition_metadata(&storage, "city=sf", "20260101000000000")
874 .await
875 .unwrap();
876 let meta = dir.path().join("city=sf/.hoodie_partition_metadata");
877 let first = std::fs::read_to_string(&meta).unwrap();
878 assert!(first.contains("commitTime=20260101000000000"));
879 assert!(first.contains("partitionDepth=1"));
880 ensure_partition_metadata(&storage, "city=sf", "20270101000000000")
881 .await
882 .unwrap();
883 assert_eq!(std::fs::read_to_string(&meta).unwrap(), first);
884 }
885
886 #[test]
888 fn test_inflight_commit_metadata_bytes_both_layouts() {
889 let avro = inflight_commit_metadata_bytes("UPSERT", true).unwrap();
890 assert_eq!(&avro[..4], b"Obj\x01");
891 let json = inflight_commit_metadata_bytes("UPSERT", false).unwrap();
892 assert!(std::str::from_utf8(&json).unwrap().contains("UPSERT"));
893 }
894
895 fn list_of(name: &str, item: arrow_schema::DataType) -> arrow_schema::DataType {
896 arrow_schema::DataType::List(std::sync::Arc::new(arrow_schema::Field::new(
897 name, item, true,
898 )))
899 }
900
901 #[test]
902 fn test_normalize_write_batch_renames_items_and_widens_lists_at_any_depth() {
903 use arrow_schema::{DataType, Field, Fields, Schema};
904 let item = std::sync::Arc::new(Field::new("item", DataType::Float32, true));
905 let element = list_of("element", DataType::Float32);
906 let in_struct =
907 |vec: DataType| DataType::Struct(Fields::from(vec![Field::new("vec", vec, true)]));
908 let cases = [
909 (DataType::FixedSizeList(item.clone(), 2), element.clone()),
910 (list_of("item", DataType::Float32), element.clone()),
911 (DataType::LargeList(item.clone()), element.clone()),
912 (
913 list_of("item", DataType::FixedSizeList(item.clone(), 2)),
914 list_of("element", element.clone()),
915 ),
916 (
917 in_struct(DataType::FixedSizeList(item, 2)),
918 in_struct(element.clone()),
919 ),
920 (DataType::Int64, DataType::Int64),
921 (element.clone(), element),
922 ];
923 let schema = Schema::new(
924 cases
925 .iter()
926 .enumerate()
927 .map(|(index, (input, _))| Field::new(format!("c{index}"), input.clone(), true))
928 .collect::<Vec<_>>(),
929 );
930 let batch = arrow_array::RecordBatch::new_empty(std::sync::Arc::new(schema));
931 let normalized = normalize_write_batch(&batch).unwrap().schema();
932 for ((_, expected), field) in cases.iter().zip(normalized.fields()) {
933 assert_eq!(field.data_type(), expected, "{}", field.name());
934 }
935 }
936
937 #[test]
938 fn test_normalize_write_batch_casts_list_columns_and_keeps_the_rest() {
939 use arrow_array::{Array, ArrayRef, FixedSizeListArray, Float32Array, Int64Array};
940 use arrow_schema::{DataType, Field, Schema};
941 let item = std::sync::Arc::new(Field::new("item", DataType::Float32, true));
942 let embedding = FixedSizeListArray::new(
943 item.clone(),
944 2,
945 std::sync::Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, 4.0])),
946 Some(vec![true, false].into()),
947 );
948 let field_metadata = HashMap::from([("k".to_string(), "v".to_string())]);
949 let schema = std::sync::Arc::new(Schema::new(vec![
950 Field::new("id", DataType::Int64, false),
951 Field::new("embedding", DataType::FixedSizeList(item, 2), true)
952 .with_metadata(field_metadata.clone()),
953 ]));
954 let batch = arrow_array::RecordBatch::try_new(
955 schema,
956 vec![
957 std::sync::Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
958 std::sync::Arc::new(embedding),
959 ],
960 )
961 .unwrap();
962
963 let normalized = normalize_write_batch(&batch).unwrap();
964 let field = normalized.schema().field(1).clone();
965 assert_eq!(field.data_type(), &list_of("element", DataType::Float32));
966 assert_eq!(field.metadata(), &field_metadata);
967 assert_eq!(normalized.schema().field(0), batch.schema().field(0));
968 let list = normalized
969 .column(1)
970 .as_any()
971 .downcast_ref::<arrow_array::ListArray>()
972 .unwrap();
973 assert!(list.is_null(1));
974 let first = list.value(0);
975 let first = first.as_any().downcast_ref::<Float32Array>().unwrap();
976 assert_eq!(first.values(), &[1.0, 2.0]);
977
978 assert_eq!(normalize_write_batch(&normalized).unwrap(), normalized);
980 let ids = batch.project(&[0]).unwrap();
981 assert_eq!(normalize_write_batch(&ids).unwrap(), ids);
982 }
983
984 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
987 async fn test_run_write_tasks_runs_on_workers_in_input_order() {
988 let caller = std::thread::current().id();
989 let tasks: Vec<futures::future::BoxFuture<'static, Result<std::thread::ThreadId>>> = (0
990 ..8u64)
991 .map(|i| {
992 let task: futures::future::BoxFuture<'static, _> = Box::pin(async move {
993 let spin = std::time::Duration::from_millis(5 * (8 - i));
995 let start = std::time::Instant::now();
996 while start.elapsed() < spin {
997 std::hint::spin_loop();
998 }
999 if i == 3 {
1000 Err(CoreError::Write("task 3".to_string()))
1001 } else {
1002 Ok(std::thread::current().id())
1003 }
1004 });
1005 task
1006 })
1007 .collect();
1008
1009 let results = run_write_tasks(tasks, 4).await;
1010
1011 assert_eq!(results.len(), 8);
1012 let mut threads = std::collections::HashSet::new();
1013 for (i, result) in results.into_iter().enumerate() {
1014 match result {
1015 Err(CoreError::Write(msg)) if i == 3 => assert_eq!(msg, "task 3"),
1016 Ok(thread) if i != 3 => {
1017 assert_ne!(thread, caller, "task {i} ran on the awaiting thread");
1018 threads.insert(thread);
1019 }
1020 other => panic!("unexpected result for task {i}: {other:?}"),
1021 }
1022 }
1023 assert!(
1024 threads.len() > 1,
1025 "tasks ran on {} thread(s)",
1026 threads.len()
1027 );
1028 }
1029
1030 #[tokio::test]
1032 async fn test_run_write_tasks_maps_panic_to_write_error() {
1033 let tasks: Vec<futures::future::BoxFuture<'static, Result<usize>>> = vec![
1034 Box::pin(async { Ok(0) }),
1035 Box::pin(async { panic!("boom") }),
1036 Box::pin(async { Ok(2) }),
1037 ];
1038
1039 let results = run_write_tasks(tasks, 2).await;
1040
1041 assert!(matches!(results[0], Ok(0)));
1042 assert!(
1043 matches!(&results[1], Err(CoreError::Write(msg)) if msg.contains("panicked")),
1044 "{:?}",
1045 results[1]
1046 );
1047 assert!(matches!(results[2], Ok(2)));
1048 }
1049
1050 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1052 async fn test_run_write_tasks_aborts_spawned_tasks_on_drop() {
1053 use futures::FutureExt;
1054 let (tx, mut rx) = tokio::sync::oneshot::channel::<()>();
1055 let task: futures::future::BoxFuture<'static, Result<()>> = Box::pin(async move {
1056 let _tx = tx;
1057 std::future::pending::<()>().await;
1058 Ok(())
1059 });
1060
1061 assert!(run_write_tasks(vec![task], 1).now_or_never().is_none());
1062
1063 for _ in 0..500 {
1064 if matches!(
1065 rx.try_recv(),
1066 Err(tokio::sync::oneshot::error::TryRecvError::Closed)
1067 ) {
1068 return;
1069 }
1070 std::thread::sleep(std::time::Duration::from_millis(10));
1071 }
1072 panic!("spawned write task kept running after the write was dropped");
1073 }
1074}