pub(crate) mod append;
pub(crate) mod archival;
mod create;
mod keygen;
pub mod lock;
pub(crate) mod markers;
pub(crate) mod metadata;
mod metadata_compaction;
mod rewrite;
pub(crate) mod rollback;
pub(crate) mod sizing;
pub(crate) use append::append_batches;
pub(crate) use append::set_commit_timezone;
pub use create::TableCreateBuilder;
pub(crate) use keygen::{
complex_keygen_validation_enabled, ensure_complex_keygen_validated,
ensure_supported_key_generator,
};
pub use lock::{InProcessLockProvider, LockLease, LockProvider};
pub(crate) use rewrite::ensure_writable_table_version;
pub use rewrite::{CompactionResult, UpsertOptions};
pub(crate) use rewrite::{
await_async_compaction, compact_table, delete_filter, delete_keys,
dynamic_partition_overwrite_batches, maybe_compact_inline, overwrite_batches, update_filter,
upsert_batches,
};
use std::collections::HashMap;
use std::str::FromStr;
use apache_avro::to_avro_datum;
use apache_avro::types::Value as AvroValue;
use crate::Result;
use crate::error::CoreError;
use crate::file_group::log_file::writer::LogFileWriter;
use crate::file_group::log_file::{BlockMetadataKey, BlockType};
use crate::metadata::meta_field::MetaField;
use crate::schema::delete::{
avro_schema_for_delete_record_list, delete_record_list_avro_schema_json,
};
use crate::storage::Storage;
use crate::table::partition::PARTITION_METAFIELD_PREFIX;
use crate::timeline::instant::{Action, Instant};
pub(crate) async fn ensure_partition_metadata(
storage: &Storage,
partition_path: &str,
instant: &str,
) -> Result<()> {
if partition_path.is_empty() {
return Ok(());
}
let meta_rel = format!("{partition_path}/{PARTITION_METAFIELD_PREFIX}");
if storage.exists(&meta_rel).await? {
return Ok(());
}
let depth = partition_path.matches('/').count() + 1;
let body = format!(
"#partition metadata\n\
commitTime={instant}\n\
partitionDepth={depth}\n"
);
storage.put_file(&meta_rel, body.into_bytes()).await?;
Ok(())
}
pub(crate) fn align_batch_to_schema(
batch: &arrow_array::RecordBatch,
target: &std::sync::Arc<arrow_schema::Schema>,
) -> Option<arrow_array::RecordBatch> {
if &batch.schema() == target {
return Some(batch.clone());
}
if batch.num_columns() != target.fields().len() {
return None;
}
for (batch_field, target_field) in batch.schema().fields().iter().zip(target.fields()) {
if batch_field.name() != target_field.name()
|| batch_field.data_type() != target_field.data_type()
{
return None;
}
}
arrow_array::RecordBatch::try_new(target.clone(), batch.columns().to_vec()).ok()
}
pub(crate) fn widen_to_table_nullability(
batches: &[arrow_array::RecordBatch],
table_schema: &arrow_schema::Schema,
) -> Result<Vec<arrow_array::RecordBatch>> {
batches
.iter()
.map(|batch| {
let source = batch.schema();
let fields: Vec<arrow_schema::Field> = source
.fields()
.iter()
.map(|field| {
let table_nullable = table_schema
.field_with_name(field.name())
.is_ok_and(|table_field| table_field.is_nullable());
field
.as_ref()
.clone()
.with_nullable(field.is_nullable() || table_nullable)
})
.collect();
let schema = arrow_schema::Schema::new_with_metadata(fields, source.metadata().clone());
arrow_array::RecordBatch::try_new(std::sync::Arc::new(schema), batch.columns().to_vec())
.map_err(CoreError::ArrowError)
})
.collect()
}
pub(crate) fn strip_meta_columns(
batch: &arrow_array::RecordBatch,
) -> Result<arrow_array::RecordBatch> {
let keep: Vec<usize> = (0..batch.num_columns())
.filter(|&i| !MetaField::is_meta_field(batch.schema().field(i).name()))
.collect();
batch.project(&keep).map_err(CoreError::ArrowError)
}
pub(crate) fn normalize_for_write(
batches: &[arrow_array::RecordBatch],
) -> Result<Vec<arrow_array::RecordBatch>> {
batches.iter().map(normalize_write_batch).collect()
}
pub(crate) fn normalize_write_batch(
batch: &arrow_array::RecordBatch,
) -> Result<arrow_array::RecordBatch> {
use arrow_schema::Field;
let batch = strip_meta_columns(batch)?;
let source = batch.schema();
let read_back = crate::schema::resolver::avro_json_to_arrow_schema(
&append::arrow_schema_to_avro_json(&source)?,
)?;
let mut changed = false;
let mut fields = Vec::with_capacity(batch.num_columns());
let mut columns = Vec::with_capacity(batch.num_columns());
for ((field, column), read_back) in source
.fields()
.iter()
.zip(batch.columns())
.zip(read_back.fields())
{
let target = read_back.data_type();
if target == field.data_type() {
fields.push(field.clone());
columns.push(column.clone());
continue;
}
changed = true;
let options = arrow::compute::CastOptions {
safe: false,
..Default::default()
};
let casted =
arrow::compute::cast_with_options(column.as_ref(), target, &options).map_err(|e| {
CoreError::Schema(format!(
"cannot convert column '{}' from {} to {target} for writing: {e}",
field.name(),
field.data_type()
))
})?;
fields.push(std::sync::Arc::new(
Field::new(field.name(), target.clone(), field.is_nullable())
.with_metadata(field.metadata().clone()),
));
columns.push(casted);
}
if !changed {
return Ok(batch);
}
let schema = std::sync::Arc::new(arrow_schema::Schema::new_with_metadata(
fields,
source.metadata().clone(),
));
arrow_array::RecordBatch::try_new(schema, columns)
.map_err(|e| CoreError::Schema(format!("cannot rebuild batch for writing: {e}")))
}
pub(crate) fn new_file_id() -> String {
format!("{}-0", uuid::Uuid::new_v4())
}
pub(crate) async fn fence_timeline_instant(
storage: &Storage,
timeline_dir: &str,
timestamp: &str,
action: Action,
requested_bytes: Vec<u8>,
inflight_bytes: Vec<u8>,
) -> Result<()> {
let requested = Instant::from_str(&format!("{}.{}.requested", timestamp, action.as_ref()))
.map_err(|e| CoreError::Write(format!("invalid fencing instant: {e}")))?;
let inflight = if action == Action::Commit {
Instant::from_str(&format!("{timestamp}.inflight"))
} else {
Instant::from_str(&format!("{}.{}.inflight", timestamp, action.as_ref()))
}
.map_err(|e| CoreError::Write(format!("invalid fencing instant: {e}")))?;
storage
.put_file_if_absent(
&requested.relative_path_with_base(timeline_dir)?,
requested_bytes,
)
.await?;
storage
.put_file_if_absent(
&inflight.relative_path_with_base(timeline_dir)?,
inflight_bytes,
)
.await?;
Ok(())
}
pub(crate) fn write_task_parallelism(table: &crate::table::Table) -> usize {
table
.hudi_configs
.as_options()
.get("hoodie.write.task.parallelism")
.and_then(|v| v.parse().ok())
.filter(|n: &usize| *n >= 1)
.unwrap_or_else(|| {
2 * std::thread::available_parallelism()
.map(std::num::NonZeroUsize::get)
.unwrap_or(4)
})
}
pub(crate) async fn run_write_tasks<T: Send + 'static>(
tasks: Vec<futures::future::BoxFuture<'static, Result<T>>>,
parallelism: usize,
) -> Vec<Result<T>> {
use futures::StreamExt;
let count = tasks.len();
let mut results: Vec<Option<Result<T>>> = std::iter::repeat_with(|| None).take(count).collect();
let indexed: Vec<_> = tasks
.into_iter()
.enumerate()
.map(|(index, task)| async move {
let result = AbortOnDrop(tokio::spawn(task))
.await
.unwrap_or_else(|e| Err(CoreError::Write(format!("write task failed: {e}"))));
(index, result)
})
.collect();
let mut stream = futures::stream::iter(indexed).buffer_unordered(parallelism.max(1));
while let Some((index, result)) = stream.next().await {
results[index] = Some(result);
}
results
.into_iter()
.map(|r| r.unwrap_or_else(|| Err(CoreError::Write("write task vanished".to_string()))))
.collect()
}
struct AbortOnDrop<T>(tokio::task::JoinHandle<T>);
impl<T> Drop for AbortOnDrop<T> {
fn drop(&mut self) {
self.0.abort();
}
}
impl<T> std::future::Future for AbortOnDrop<T> {
type Output = std::result::Result<T, tokio::task::JoinError>;
fn poll(
mut self: std::pin::Pin<&mut Self>,
cx: &mut std::task::Context<'_>,
) -> std::task::Poll<Self::Output> {
std::pin::Pin::new(&mut self.0).poll(cx)
}
}
pub(crate) struct FileTaskOutput {
pub size: i64,
pub ranges: Option<Vec<crate::metadata::table::column_stats::ColumnRangeStats>>,
}
pub(crate) fn parquet_file_task(
storage: std::sync::Arc<Storage>,
props: parquet::file::properties::WriterProperties,
bloom: Option<crate::bloom::BloomFilterConfig>,
prepared: Vec<arrow_array::RecordBatch>,
relative_path: String,
collect_ranges: bool,
) -> futures::future::BoxFuture<'static, Result<FileTaskOutput>> {
Box::pin(async move {
let (bytes, ranges) = tokio::task::spawn_blocking(move || -> Result<_> {
let bytes =
crate::write::append::write_parquet_bytes_with_props(props, bloom, &prepared)?;
let ranges = if collect_ranges {
Some(
crate::metadata::table::column_stats::column_ranges_from_parquet_bytes(&bytes)?,
)
} else {
None
};
Ok((bytes, ranges))
})
.await
.map_err(|e| CoreError::Write(format!("write task panicked: {e}")))??;
let size = bytes.len() as i64;
storage.put_file(&relative_path, bytes).await?;
Ok(FileTaskOutput { size, ranges })
})
}
pub(crate) fn log_file_task(
storage: std::sync::Arc<Storage>,
props: parquet::file::properties::WriterProperties,
prepared: Vec<arrow_array::RecordBatch>,
relative_path: String,
instant: String,
schema_json: String,
collect_ranges: bool,
) -> futures::future::BoxFuture<'static, Result<FileTaskOutput>> {
Box::pin(async move {
let (block, ranges) = tokio::task::spawn_blocking(move || -> Result<_> {
let parquet =
crate::write::append::write_parquet_bytes_with_props(props, None, &prepared)?;
let ranges = if collect_ranges {
Some(
crate::metadata::table::column_stats::column_ranges_from_parquet_bytes(
&parquet,
)?,
)
} else {
None
};
let block = LogFileWriter::write_log_block(
BlockType::ParquetData,
HashMap::from([
(BlockMetadataKey::InstantTime, instant),
(BlockMetadataKey::Schema, schema_json),
]),
&parquet,
);
Ok((block, ranges))
})
.await
.map_err(|e| CoreError::Write(format!("write task panicked: {e}")))??;
let size = block.len() as i64;
storage.put_file(&relative_path, block).await?;
Ok(FileTaskOutput { size, ranges })
})
}
pub(crate) async fn post_complete_bookkeeping(
table: &crate::table::Table,
storage: &Storage,
instant: &str,
) -> Result<()> {
markers::delete_marker_dir(storage, instant).await;
let (min_keep, max_keep) = archival::validated_keep_range(&table.hudi_configs)?;
archival::archive_timeline_if_needed(
storage,
&crate::write::append::timeline_dir(table),
&crate::timeline::loader::lsm_history_dir(&table.hudi_configs),
min_keep,
max_keep,
)
.await?;
Ok(())
}
pub(crate) fn inflight_commit_metadata_bytes(operation: &str, layout_two: bool) -> Result<Vec<u8>> {
let metadata = crate::metadata::commit::HoodieCommitMetadata {
version: Some(1),
operation_type: Some(operation.to_string()),
partition_to_write_stats: Some(std::collections::HashMap::new()),
compacted: Some(false),
extra_metadata: Some(std::collections::HashMap::new()),
partition_to_replace_file_ids: None,
};
if layout_two {
metadata.to_avro_bytes()
} else {
metadata.to_json_bytes()
}
}
pub(crate) fn build_delete_log_block(
instant: &str,
keys: &[(String, String)],
ordering_val: i64,
) -> Result<Vec<u8>> {
let records = keys
.iter()
.map(|(record_key, partition_path)| {
AvroValue::Record(vec![
(
"recordKey".to_string(),
AvroValue::Union(1, Box::new(AvroValue::String(record_key.clone()))),
),
(
"partitionPath".to_string(),
AvroValue::Union(1, Box::new(AvroValue::String(partition_path.clone()))),
),
(
"orderingVal".to_string(),
if ordering_val == 0 {
AvroValue::Union(
2,
Box::new(AvroValue::Record(vec![(
"value".to_string(),
AvroValue::Int(0),
)])),
)
} else {
AvroValue::Union(
3,
Box::new(AvroValue::Record(vec![(
"value".to_string(),
AvroValue::Long(ordering_val),
)])),
)
},
),
])
})
.collect();
let payload = to_avro_datum(
avro_schema_for_delete_record_list()?,
AvroValue::Record(vec![(
"deleteRecordList".to_string(),
AvroValue::Array(records),
)]),
)?;
let mut content = Vec::with_capacity(8 + payload.len());
content.extend_from_slice(&3u32.to_be_bytes());
content.extend_from_slice(&(payload.len() as u32).to_be_bytes());
content.extend_from_slice(&payload);
Ok(LogFileWriter::write_log_block(
BlockType::Delete,
HashMap::from([
(BlockMetadataKey::InstantTime, instant.to_string()),
(
BlockMetadataKey::Schema,
delete_record_list_avro_schema_json().to_string(),
),
]),
&content,
))
}
#[derive(Default)]
pub(crate) struct MetadataTableUpdates {
pub files: Vec<(String, String, i64, bool)>,
pub record_index: Vec<crate::metadata::table::encode::RecordIndexEntry>,
pub column_stats: Vec<metadata::StatsFileUpdate>,
pub replaced_paths: Vec<String>,
}
pub(crate) fn stats_by_partition(
stats: Vec<crate::metadata::commit::HoodieWriteStat>,
) -> HashMap<String, Vec<crate::metadata::commit::HoodieWriteStat>> {
let mut by_partition: HashMap<String, Vec<_>> = HashMap::new();
for stat in stats {
let partition = stat.partition_path.clone().unwrap_or_default();
by_partition.entry(partition).or_default().push(stat);
}
by_partition
}
pub(crate) fn commit_metadata_bytes(
table: &crate::table::Table,
operation: &str,
partition_to_write_stats: HashMap<String, Vec<crate::metadata::commit::HoodieWriteStat>>,
partition_to_replace_file_ids: Option<HashMap<String, Vec<String>>>,
schema: &arrow_schema::Schema,
) -> Result<Vec<u8>> {
data_instant_metadata_bytes(
table,
operation,
partition_to_write_stats,
partition_to_replace_file_ids,
schema,
false,
)
}
pub(crate) fn compaction_commit_metadata_bytes(
table: &crate::table::Table,
partition_to_write_stats: HashMap<String, Vec<crate::metadata::commit::HoodieWriteStat>>,
schema: &arrow_schema::Schema,
) -> Result<Vec<u8>> {
data_instant_metadata_bytes(
table,
"COMPACT",
partition_to_write_stats,
None,
schema,
true,
)
}
fn data_instant_metadata_bytes(
table: &crate::table::Table,
operation: &str,
partition_to_write_stats: HashMap<String, Vec<crate::metadata::commit::HoodieWriteStat>>,
partition_to_replace_file_ids: Option<HashMap<String, Vec<String>>>,
schema: &arrow_schema::Schema,
compacted: bool,
) -> Result<Vec<u8>> {
let schema_json =
append::arrow_schema_to_avro_json(&append::strip_meta_fields_from_schema(schema))?;
let extra_metadata = Some(HashMap::from([("schema".to_string(), schema_json)]));
let layout_two = append::is_layout_two(table);
match partition_to_replace_file_ids {
Some(partition_to_replace_file_ids) => {
let metadata = crate::metadata::replace_commit::HoodieReplaceCommitMetadata {
version: Some(1),
operation_type: Some(operation.to_string()),
partition_to_write_stats: Some(partition_to_write_stats),
compacted: Some(compacted),
extra_metadata,
partition_to_replace_file_ids: Some(partition_to_replace_file_ids),
};
if layout_two {
metadata.to_avro_bytes()
} else {
metadata.to_json_bytes()
}
}
None => {
let metadata = crate::metadata::commit::HoodieCommitMetadata {
version: Some(1),
operation_type: Some(operation.to_string()),
partition_to_write_stats: Some(partition_to_write_stats),
compacted: Some(compacted),
extra_metadata,
partition_to_replace_file_ids: None,
};
if layout_two {
metadata.to_avro_bytes()
} else {
metadata.to_json_bytes()
}
}
}
}
pub(crate) async fn complete_write(
table: &mut crate::table::Table,
instant: &str,
action: Action,
commit_bytes: Vec<u8>,
schema: &arrow_schema::Schema,
mut updates: MetadataTableUpdates,
written_paths: &[String],
) -> Result<CompletedWrite> {
let storage = table.file_system_view.storage.clone();
let mut mdt_stats = HashMap::new();
if table.is_metadata_table_enabled() {
let mdt_result = async {
if !updates.files.is_empty() {
mdt_stats.insert(
"files".to_string(),
metadata::update_files_partition_entries(
storage.as_ref(),
instant,
&updates.files,
)
.await?,
);
}
if crate::index::is_record_index_enabled(table) && !updates.record_index.is_empty() {
let stats = metadata::update_record_index(
storage.as_ref(),
instant,
std::mem::take(&mut updates.record_index),
write_task_parallelism(table),
)
.await?;
if !stats.is_empty() {
mdt_stats.insert("record_index".to_string(), stats);
}
}
if metadata::is_column_stats_enabled(table) && !updates.column_stats.is_empty() {
let stats = metadata::update_column_stats_partitions(
table,
instant,
&updates.column_stats,
schema,
&updates.replaced_paths,
)
.await?;
mdt_stats.extend(stats);
}
Ok::<(), CoreError>(())
}
.await;
if let Err(error) = mdt_result {
for path in written_paths {
let _ = storage.delete_file(path).await;
}
return Err(error);
}
}
let lock = lock::lock_provider_for(table);
let critical_section = lock.lock().await?;
if table.is_metadata_table_enabled()
&& let Err(error) =
metadata::write_metadata_commit(storage.as_ref(), instant, mdt_stats).await
{
for path in written_paths {
let _ = storage.delete_file(path).await;
}
return Err(error);
}
let completion = if append::is_layout_two(table) {
Some(append::generate_instant_time().await)
} else {
None
};
let completed = Instant::new_completed(instant.to_string(), action, completion.clone())?;
let commit_relative_path = completed.relative_path_with_base(&append::timeline_dir(table))?;
storage
.put_file_if_absent(&commit_relative_path, commit_bytes)
.await?;
if let Err(error) = post_complete_bookkeeping(table, storage.as_ref(), instant).await {
log::warn!("post-commit bookkeeping for {instant} failed: {error}");
}
if let Err(error) = metadata_compaction::compact_metadata_table_if_needed(table).await {
log::warn!("metadata table compaction after {instant} failed: {error}");
}
drop(critical_section);
let _ = table.timeline.reload_completed_commits().await;
table.clear_read_caches();
Ok(CompletedWrite {
completion_time: completion,
})
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct WriteResult {
pub requested_time: Option<String>,
pub completion_time: Option<String>,
pub num_writes: usize,
pub num_updates: usize,
pub num_inserts: usize,
pub num_deletes: usize,
}
pub(crate) struct CompletedWrite {
pub completion_time: Option<String>,
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_fence_timeline_instant_conflicts_on_same_timestamp() {
let dir = tempfile::tempdir().unwrap();
let base_url = url::Url::from_directory_path(dir.path()).unwrap();
let storage = Storage::new_with_base_url(base_url).unwrap();
fence_timeline_instant(
&storage,
".hoodie/timeline",
"20260101000000000",
Action::Commit,
Vec::new(),
Vec::new(),
)
.await
.unwrap();
let err = fence_timeline_instant(
&storage,
".hoodie/timeline",
"20260101000000000",
Action::Commit,
Vec::new(),
Vec::new(),
)
.await
.unwrap_err();
assert!(
err.to_string().contains("already exists"),
"expected instant conflict, got: {err}"
);
}
#[test]
fn test_build_delete_log_block_both_ordering_branches() {
for ordering in [0i64, 42i64] {
let block = build_delete_log_block(
"20260101000000000",
&[("k1".to_string(), "p".to_string())],
ordering,
)
.unwrap();
assert!(!block.is_empty());
assert_eq!(&block[..6], b"#HUDI#");
}
}
#[tokio::test]
async fn test_ensure_partition_metadata_idempotent() {
let dir = tempfile::tempdir().unwrap();
let base_url = url::Url::from_directory_path(dir.path()).unwrap();
let storage = Storage::new_with_base_url(base_url).unwrap();
ensure_partition_metadata(&storage, "", "20260101000000000")
.await
.unwrap();
ensure_partition_metadata(&storage, "city=sf", "20260101000000000")
.await
.unwrap();
let meta = dir.path().join("city=sf/.hoodie_partition_metadata");
let first = std::fs::read_to_string(&meta).unwrap();
assert!(first.contains("commitTime=20260101000000000"));
assert!(first.contains("partitionDepth=1"));
ensure_partition_metadata(&storage, "city=sf", "20270101000000000")
.await
.unwrap();
assert_eq!(std::fs::read_to_string(&meta).unwrap(), first);
}
#[test]
fn test_inflight_commit_metadata_bytes_both_layouts() {
let avro = inflight_commit_metadata_bytes("UPSERT", true).unwrap();
assert_eq!(&avro[..4], b"Obj\x01");
let json = inflight_commit_metadata_bytes("UPSERT", false).unwrap();
assert!(std::str::from_utf8(&json).unwrap().contains("UPSERT"));
}
fn list_of(name: &str, item: arrow_schema::DataType) -> arrow_schema::DataType {
arrow_schema::DataType::List(std::sync::Arc::new(arrow_schema::Field::new(
name, item, true,
)))
}
#[test]
fn test_normalize_write_batch_renames_items_and_widens_lists_at_any_depth() {
use arrow_schema::{DataType, Field, Fields, Schema};
let item = std::sync::Arc::new(Field::new("item", DataType::Float32, true));
let element = list_of("element", DataType::Float32);
let in_struct =
|vec: DataType| DataType::Struct(Fields::from(vec![Field::new("vec", vec, true)]));
let cases = [
(DataType::FixedSizeList(item.clone(), 2), element.clone()),
(list_of("item", DataType::Float32), element.clone()),
(DataType::LargeList(item.clone()), element.clone()),
(
list_of("item", DataType::FixedSizeList(item.clone(), 2)),
list_of("element", element.clone()),
),
(
in_struct(DataType::FixedSizeList(item, 2)),
in_struct(element.clone()),
),
(DataType::Int64, DataType::Int64),
(element.clone(), element),
];
let schema = Schema::new(
cases
.iter()
.enumerate()
.map(|(index, (input, _))| Field::new(format!("c{index}"), input.clone(), true))
.collect::<Vec<_>>(),
);
let batch = arrow_array::RecordBatch::new_empty(std::sync::Arc::new(schema));
let normalized = normalize_write_batch(&batch).unwrap().schema();
for ((_, expected), field) in cases.iter().zip(normalized.fields()) {
assert_eq!(field.data_type(), expected, "{}", field.name());
}
}
#[test]
fn test_normalize_write_batch_casts_list_columns_and_keeps_the_rest() {
use arrow_array::{Array, ArrayRef, FixedSizeListArray, Float32Array, Int64Array};
use arrow_schema::{DataType, Field, Schema};
let item = std::sync::Arc::new(Field::new("item", DataType::Float32, true));
let embedding = FixedSizeListArray::new(
item.clone(),
2,
std::sync::Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0, 4.0])),
Some(vec![true, false].into()),
);
let field_metadata = HashMap::from([("k".to_string(), "v".to_string())]);
let schema = std::sync::Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("embedding", DataType::FixedSizeList(item, 2), true)
.with_metadata(field_metadata.clone()),
]));
let batch = arrow_array::RecordBatch::try_new(
schema,
vec![
std::sync::Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef,
std::sync::Arc::new(embedding),
],
)
.unwrap();
let normalized = normalize_write_batch(&batch).unwrap();
let field = normalized.schema().field(1).clone();
assert_eq!(field.data_type(), &list_of("element", DataType::Float32));
assert_eq!(field.metadata(), &field_metadata);
assert_eq!(normalized.schema().field(0), batch.schema().field(0));
let list = normalized
.column(1)
.as_any()
.downcast_ref::<arrow_array::ListArray>()
.unwrap();
assert!(list.is_null(1));
let first = list.value(0);
let first = first.as_any().downcast_ref::<Float32Array>().unwrap();
assert_eq!(first.values(), &[1.0, 2.0]);
assert_eq!(normalize_write_batch(&normalized).unwrap(), normalized);
let ids = batch.project(&[0]).unwrap();
assert_eq!(normalize_write_batch(&ids).unwrap(), ids);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn test_run_write_tasks_runs_on_workers_in_input_order() {
let caller = std::thread::current().id();
let tasks: Vec<futures::future::BoxFuture<'static, Result<std::thread::ThreadId>>> = (0
..8u64)
.map(|i| {
let task: futures::future::BoxFuture<'static, _> = Box::pin(async move {
let spin = std::time::Duration::from_millis(5 * (8 - i));
let start = std::time::Instant::now();
while start.elapsed() < spin {
std::hint::spin_loop();
}
if i == 3 {
Err(CoreError::Write("task 3".to_string()))
} else {
Ok(std::thread::current().id())
}
});
task
})
.collect();
let results = run_write_tasks(tasks, 4).await;
assert_eq!(results.len(), 8);
let mut threads = std::collections::HashSet::new();
for (i, result) in results.into_iter().enumerate() {
match result {
Err(CoreError::Write(msg)) if i == 3 => assert_eq!(msg, "task 3"),
Ok(thread) if i != 3 => {
assert_ne!(thread, caller, "task {i} ran on the awaiting thread");
threads.insert(thread);
}
other => panic!("unexpected result for task {i}: {other:?}"),
}
}
assert!(
threads.len() > 1,
"tasks ran on {} thread(s)",
threads.len()
);
}
#[tokio::test]
async fn test_run_write_tasks_maps_panic_to_write_error() {
let tasks: Vec<futures::future::BoxFuture<'static, Result<usize>>> = vec![
Box::pin(async { Ok(0) }),
Box::pin(async { panic!("boom") }),
Box::pin(async { Ok(2) }),
];
let results = run_write_tasks(tasks, 2).await;
assert!(matches!(results[0], Ok(0)));
assert!(
matches!(&results[1], Err(CoreError::Write(msg)) if msg.contains("panicked")),
"{:?}",
results[1]
);
assert!(matches!(results[2], Ok(2)));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn test_run_write_tasks_aborts_spawned_tasks_on_drop() {
use futures::FutureExt;
let (tx, mut rx) = tokio::sync::oneshot::channel::<()>();
let task: futures::future::BoxFuture<'static, Result<()>> = Box::pin(async move {
let _tx = tx;
std::future::pending::<()>().await;
Ok(())
});
assert!(run_write_tasks(vec![task], 1).now_or_never().is_none());
for _ in 0..500 {
if matches!(
rx.try_recv(),
Err(tokio::sync::oneshot::error::TryRecvError::Closed)
) {
return;
}
std::thread::sleep(std::time::Duration::from_millis(10));
}
panic!("spawned write task kept running after the write was dropped");
}
}