diff --git a/src/handlers/http/mod.rs b/src/handlers/http/mod.rs index 6bf903095..e1878a298 100644 --- a/src/handlers/http/mod.rs +++ b/src/handlers/http/mod.rs @@ -139,6 +139,13 @@ pub async fn fetch_schema( stream_name: &str, tenant_id: &Option, ) -> anyhow::Result { + if let Some(global) = PARSEABLE + .metastore + .get_global_schema(stream_name, tenant_id) + .await? + { + return Ok(global.schema); + } let res: Vec = PARSEABLE .metastore .get_all_schemas(stream_name, tenant_id) diff --git a/src/lib.rs b/src/lib.rs index 5ab009230..8ffefe1bc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -50,6 +50,7 @@ pub mod prism; pub mod query; pub mod rbac; mod response; +pub mod schema_registry; pub mod sse; mod static_schema; mod stats; diff --git a/src/metastore/metastore_traits.rs b/src/metastore/metastore_traits.rs index 267428d7a..fb7e6e20f 100644 --- a/src/metastore/metastore_traits.rs +++ b/src/metastore/metastore_traits.rs @@ -23,6 +23,7 @@ use bytes::Bytes; use chrono::{DateTime, Utc}; use dashmap::DashMap; use erased_serde::Serialize as ErasedSerialize; +use relative_path::RelativePathBuf; use tonic::async_trait; use ulid::Ulid; @@ -36,6 +37,7 @@ use crate::{ handlers::http::modal::NodeType, metastore::MetastoreError, option::Mode, + parseable::GlobalSchema, users::filters::Filter, }; @@ -291,6 +293,37 @@ pub trait Metastore: std::fmt::Debug + Send + Sync { stream_name: &str, tenant_id: &Option, ) -> Result<(), MetastoreError>; + async fn get_global_schema_locks( + &self, + stream_name: &str, + tenant_id: &Option, + ) -> Result, MetastoreError>; + async fn put_global_schema_lock( + &self, + lock_file_path: RelativePathBuf, + tenant_id: &Option, + ) -> Result<(), MetastoreError>; + async fn delete_global_schema_lock( + &self, + path: RelativePathBuf, + tenant_id: &Option, + ) -> Result<(), MetastoreError>; + /// Global schema operations used to detect and resolve schema conflicts. + /// A missing expected hash creates the record; an existing hash is checked + /// before updating. Callers coordinate updates with global-schema lock files; + /// the hash check and write are not atomic. `false` indicates a conflict. + async fn get_global_schema( + &self, + stream_name: &str, + tenant_id: &Option, + ) -> Result, MetastoreError>; + async fn put_global_schema( + &self, + obj: GlobalSchema, + stream_name: &str, + tenant_id: &Option, + expected: Option, + ) -> Result; /// parseable metadata async fn get_parseable_metadata( diff --git a/src/metastore/metastores/object_store_metastore.rs b/src/metastore/metastores/object_store_metastore.rs index 867544601..1cec5bc76 100644 --- a/src/metastore/metastores/object_store_metastore.rs +++ b/src/metastore/metastores/object_store_metastore.rs @@ -51,16 +51,17 @@ use crate::{ metastore_traits::{Metastore, MetastoreObject}, }, option::Mode, - parseable::{DEFAULT_TENANT, PARSEABLE}, + parseable::{DEFAULT_TENANT, GlobalSchema, PARSEABLE}, storage::{ - ALERTS_ROOT_DIRECTORY, ObjectStorage, ObjectStorageError, PARSEABLE_ROOT_DIRECTORY, - SETTINGS_ROOT_DIRECTORY, STREAM_METADATA_FILE_NAME, STREAM_ROOT_DIRECTORY, - TARGETS_ROOT_DIRECTORY, TOMBSTONE_ROOT_DIRECTORY, + ALERTS_ROOT_DIRECTORY, GLOBAL_SCHEMA_FILE_NAME, ObjectStorage, ObjectStorageError, + PARSEABLE_ROOT_DIRECTORY, SETTINGS_ROOT_DIRECTORY, STREAM_METADATA_FILE_NAME, + STREAM_ROOT_DIRECTORY, TARGETS_ROOT_DIRECTORY, TOMBSTONE_ROOT_DIRECTORY, object_storage::{ - alert_json_path, alert_state_json_path, filter_path, manifest_path, mttr_json_path, - outbound_http_policy_json_path, parseable_json_path, schema_path, stream_json_path, - to_bytes, + alert_json_path, alert_state_json_path, filter_path, global_schema_lock_file_path, + global_schema_path, manifest_path, mttr_json_path, outbound_http_policy_json_path, + parseable_json_path, schema_path, stream_json_path, to_bytes, }, + to_object_store_path, }, users::filters::{Filter, migrate_v1_v2}, }; @@ -79,6 +80,10 @@ pub fn is_missing_optional_dir(err: &ObjectStorageError) -> bool { } } +fn is_legacy_schema_file(file_name: &str) -> bool { + file_name.contains(".schema") && !file_name.ends_with(GLOBAL_SCHEMA_FILE_NAME) +} + #[async_trait] impl Metastore for ObjectStoreMetastore { /// Since Parseable already starts with a connection to an object store, no need to implement this @@ -1160,7 +1165,7 @@ impl Metastore for ObjectStoreMetastore { .storage .get_objects( Some(&path_prefix), - Box::new(|file_name: String| file_name.contains(".schema")), + Box::new(|file_name: String| is_legacy_schema_file(&file_name)), tenant_id, ) .await? @@ -1178,6 +1183,11 @@ impl Metastore for ObjectStoreMetastore { stream_name: &str, tenant_id: &Option, ) -> Result { + if let Some(global_schema) = self.get_global_schema(stream_name, tenant_id).await? { + let schema = to_bytes(&global_schema.schema); + return Ok(schema); + } + Ok(self .storage .get_object(&schema_path(stream_name, tenant_id), tenant_id) @@ -1196,6 +1206,87 @@ impl Metastore for ObjectStoreMetastore { .put_object(&path, to_bytes(&obj), tenant_id) .await?) } + async fn get_global_schema_locks( + &self, + stream_name: &str, + tenant_id: &Option, + ) -> Result, MetastoreError> { + let path = to_object_store_path(&global_schema_lock_file_path(stream_name, tenant_id)); + let mut result_file_list = Vec::new(); + let resp = self.storage.list_with_delimiter(Some(path)).await?; + for obj in resp.objects.iter() { + let name = obj.location.filename(); + if let Some(name) = name + && name.ends_with(".lock") + { + result_file_list.push(name.to_owned()); + } + } + Ok(result_file_list) + } + async fn put_global_schema_lock( + &self, + lock_file_path: RelativePathBuf, + tenant_id: &Option, + ) -> Result<(), MetastoreError> { + let data = to_bytes(""); + self.storage + .put_object(&lock_file_path, data, tenant_id) + .await?; + Ok(()) + } + async fn delete_global_schema_lock( + &self, + path: RelativePathBuf, + tenant_id: &Option, + ) -> Result<(), MetastoreError> { + match self.storage.delete_object(&path, tenant_id).await { + Ok(()) | Err(ObjectStorageError::NoSuchKey(_)) => Ok(()), + Err(ObjectStorageError::IoError(error)) + if error.kind() == std::io::ErrorKind::NotFound => + { + Ok(()) + } + Err(error) => Err(error.into()), + } + } + /// Retrieve the persisted global schema and its stored content hash. + async fn get_global_schema( + &self, + stream_name: &str, + tenant_id: &Option, + ) -> Result, MetastoreError> { + // LocalFS does not participate in the shared registry. Only bypass + // explicitly unsupported backends; propagate errors from cloud stores. + let path = global_schema_path(stream_name, tenant_id); + let global_schema = match self.storage.get_object(&path, tenant_id).await { + Ok(obj) => serde_json::from_slice(&obj)?, + Err(ObjectStorageError::NoSuchKey(_) | ObjectStorageError::PathError(_)) => { + return Ok(None); + } + Err(e) => { + tracing::error!("{e}"); + return Err(MetastoreError::ObjectStorageError(e)); + } + }; + Ok(Some(global_schema)) + } + + /// Create a missing global schema or update one whose stored hash matches. + /// Callers coordinate the non-atomic hash check and update with lock files. + async fn put_global_schema( + &self, + obj: GlobalSchema, + stream_name: &str, + tenant_id: &Option, + expected: Option, + ) -> Result { + let path = global_schema_path(stream_name, tenant_id); + Ok(self + .storage + .put_object_if(&path, to_bytes(&obj), expected, tenant_id) + .await?) + } async fn get_parseable_metadata( &self, @@ -1469,4 +1560,40 @@ mod tests { Err(ObjectStorageError::NoSuchKey(_)) )); } + + #[test] + fn global_schema_paths_normalize_default_tenant_and_isolate_real_tenants() { + let no_tenant = None; + let default_tenant = Some(DEFAULT_TENANT.to_owned()); + let tenant_a = Some("tenant-a".to_owned()); + let tenant_b = Some("tenant-b".to_owned()); + + assert_eq!( + global_schema_path("events", &no_tenant).as_str(), + "events/.stream/.global.schema" + ); + assert_eq!( + global_schema_path("events", &default_tenant), + global_schema_path("events", &no_tenant) + ); + assert_eq!( + global_schema_lock_file_path("events", &default_tenant), + global_schema_lock_file_path("events", &no_tenant) + ); + assert_eq!( + global_schema_path("events", &tenant_a).as_str(), + "tenant-a/events/.stream/.global.schema" + ); + assert_ne!( + global_schema_path("events", &tenant_a), + global_schema_path("events", &tenant_b) + ); + assert_ne!( + global_schema_lock_file_path("events", &tenant_a), + global_schema_lock_file_path("events", &tenant_b) + ); + assert!(is_legacy_schema_file(".schema")); + assert!(is_legacy_schema_file(".ingestor.node-1.schema")); + assert!(!is_legacy_schema_file(GLOBAL_SCHEMA_FILE_NAME)); + } } diff --git a/src/metastore/mod.rs b/src/metastore/mod.rs index acf7b7a97..dfb8e2f74 100644 --- a/src/metastore/mod.rs +++ b/src/metastore/mod.rs @@ -49,6 +49,12 @@ pub enum MetastoreError { #[error("JSON schema validation error: {message}")] JsonSchemaError { message: String }, + #[error("CAS convergence failure: {message}")] + CASConvergenceFailure { message: String }, + + #[error("LockFile failure: {message}")] + LockFileFailure { message: String }, + #[error("Invalid JSON structure: expected {expected}, found {found}")] InvalidJsonStructure { expected: String, found: String }, @@ -100,6 +106,24 @@ impl MetastoreError { metadata: std::collections::HashMap::new(), status_code: 400, }, + MetastoreError::CASConvergenceFailure { message } => MetastoreErrorDetail { + operation: "CASConvergenceFailure".to_string(), + message: message.to_string(), + stream_name: None, + file_path: None, + timestamp: Some(chrono::Utc::now()), + metadata: std::collections::HashMap::new(), + status_code: 500, + }, + MetastoreError::LockFileFailure { message } => MetastoreErrorDetail { + operation: "LockFileFailure".to_string(), + message: message.to_string(), + stream_name: None, + file_path: None, + timestamp: Some(chrono::Utc::now()), + metadata: std::collections::HashMap::new(), + status_code: 500, + }, MetastoreError::ManifestCodecError(e) => MetastoreErrorDetail { operation: "ManifestCodecError".to_string(), message: e.to_string(), @@ -164,9 +188,11 @@ impl MetastoreError { MetastoreError::JsonParseError(..) => StatusCode::INTERNAL_SERVER_ERROR, MetastoreError::ManifestCodecError(..) => StatusCode::INTERNAL_SERVER_ERROR, MetastoreError::JsonSchemaError { .. } => StatusCode::INTERNAL_SERVER_ERROR, - MetastoreError::InvalidJsonStructure { .. } => StatusCode::INTERNAL_SERVER_ERROR, + MetastoreError::InvalidJsonStructure { .. } => StatusCode::BAD_REQUEST, MetastoreError::MissingJsonField { .. } => StatusCode::INTERNAL_SERVER_ERROR, MetastoreError::InvalidJsonValue { .. } => StatusCode::INTERNAL_SERVER_ERROR, + MetastoreError::CASConvergenceFailure { .. } + | MetastoreError::LockFileFailure { .. } => StatusCode::INTERNAL_SERVER_ERROR, MetastoreError::Error { status_code, .. } => *status_code, } } diff --git a/src/migration/mod.rs b/src/migration/mod.rs index 6821e0ab1..9c2453ad5 100644 --- a/src/migration/mod.rs +++ b/src/migration/mod.rs @@ -485,9 +485,32 @@ pub async fn setup_logstream_metadata( tenant_id.as_deref().unwrap_or(DEFAULT_TENANT), ); - let schema = PARSEABLE - .get_or_create_stream(stream, tenant_id) - .updated_schema(arrow_schema.clone()); + let stream_ref = PARSEABLE.get_or_create_stream(stream, tenant_id); + // Once a dataset is coordinated, its persisted schema uses canonical names + // while surviving IPC files still contain raw ingest names. Resolve each + // file independently before merging, rather than merging those raw types + // directly into the canonical schema during recovery. + // Disabling bootstrap must not revert an already canonical dataset to raw + // conversion: existing coordinated streams still require their mappings. + let schema = if PARSEABLE.options.mode == Mode::Ingest + && !static_schema_flag + && PARSEABLE + .metastore + .get_global_schema(stream, tenant_id) + .await? + .is_some() + { + stream_ref + .restore_global_schema( + arrow_schema.clone(), + tenant_id, + time_partition.as_ref(), + custom_partition.as_ref(), + ) + .await? + } else { + stream_ref.updated_schema(arrow_schema.clone()) + }; let schema = HashMap::from_iter( schema .fields diff --git a/src/parseable/mod.rs b/src/parseable/mod.rs index cd127adbf..d400603d3 100644 --- a/src/parseable/mod.rs +++ b/src/parseable/mod.rs @@ -34,7 +34,7 @@ use regex::Regex; use relative_path::RelativePathBuf; pub use staging::StagingError; use streams::StreamRef; -pub use streams::{Stream, StreamNotFound, Streams}; +pub use streams::{GlobalSchema, Stream, StreamNotFound, Streams}; use tokio::try_join; use tracing::error; diff --git a/src/parseable/staging/mod.rs b/src/parseable/staging/mod.rs index 823f95c45..0bc1a9fd1 100644 --- a/src/parseable/staging/mod.rs +++ b/src/parseable/staging/mod.rs @@ -42,4 +42,6 @@ pub enum StagingError { PoisonError(#[from] PoisonError), #[error("JSON Error {0}")] Json(#[from] serde_json::Error), + #[error("Global schema resolution failed: {0}")] + GlobalSchema(String), } diff --git a/src/parseable/staging/reader.rs b/src/parseable/staging/reader.rs index 00cd399da..e1cab0ebb 100644 --- a/src/parseable/staging/reader.rs +++ b/src/parseable/staging/reader.rs @@ -37,7 +37,7 @@ use tracing::{error, info_span}; use crate::{ event::DEFAULT_TIMESTAMP_KEY, - utils::arrow::{adapt_batch, reverse}, + utils::arrow::{batch_adapter::try_adapt_batch, reverse}, }; #[derive(Debug)] @@ -205,14 +205,17 @@ impl MergedForwardRecordReader { self.file_plans.into_iter() } + /// Returns the schemas exposed by each readable Arrow stream. + pub fn source_schemas(&self) -> Vec { + self.file_plans + .iter() + .map(|plan| plan.schema.as_ref().clone()) + .collect() + } + /// Returns the union of schemas exposed by all readable Arrow streams. - pub fn merged_schema(&self) -> Schema { - Schema::try_merge( - self.file_plans - .iter() - .map(|plan| plan.schema.as_ref().clone()), - ) - .unwrap() + pub fn merged_schema(&self) -> Result { + Schema::try_merge(self.source_schemas()) } } @@ -251,11 +254,29 @@ impl MergedReverseRecordReader { } } - /// Merges all readers in reverse timestamp order and adapts each batch to `schema`. + /// Identity-resolution convenience for reader tests. + #[cfg(test)] pub fn merged_iter( self, schema: Arc, time_partition: Option, + ) -> impl Iterator> { + self.merged_iter_with_resolution( + schema.clone(), + schema, + time_partition, + Arc::new(crate::schema_registry::Resolution::identity()), + ) + } + + /// Merges in reverse timestamp order, adapts batches to the local merged schema, + /// resolves field names, and adapts them to the final canonical schema. + pub fn merged_iter_with_resolution( + self, + schema: Arc, + local_schema: Arc, + time_partition: Option, + resolution: Arc, ) -> impl Iterator> { let adapted_readers = self.readers; kmerge_by( @@ -275,17 +296,25 @@ impl MergedReverseRecordReader { }, ) .map(|batch| batch.map(|batch| reverse(&batch))) - .map(move |batch| batch.map(|batch| adapt_batch(schema.clone(), &batch))) + .map(move |batch| { + batch + .and_then(|batch| try_adapt_batch(local_schema.clone(), &batch)) + .and_then(|batch| resolution.rename_batch(&batch)) + .and_then(|batch| try_adapt_batch(schema.clone(), &batch)) + }) + } + + /// Returns the schemas exposed by each readable Arrow stream. + pub fn source_schemas(&self) -> Vec { + self.readers + .iter() + .map(|reader| reader.schema().as_ref().clone()) + .collect() } /// Returns the union of schemas exposed by all readable Arrow streams. - pub fn merged_schema(&self) -> Schema { - Schema::try_merge( - self.readers - .iter() - .map(|reader| reader.schema().as_ref().clone()), - ) - .unwrap() + pub fn merged_schema(&self) -> Result { + Schema::try_merge(self.source_schemas()) } } @@ -635,14 +664,15 @@ fn find_limit_and_type( #[cfg(test)] mod tests { use std::{ - collections::VecDeque, + collections::{HashMap, VecDeque}, io::{self, Cursor, Read}, path::Path, sync::Arc, }; use arrow_array::{ - Array, DictionaryArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray, + Array, DictionaryArray, Float64Array, Int32Array, Int64Array, NullArray, RecordBatch, + StringArray, TimestampMillisecondArray, cast::AsArray, types::{Int32Type, Int64Type}, }; @@ -653,7 +683,7 @@ mod tests { write_message, }, }; - use arrow_schema::{DataType, Field, Schema}; + use arrow_schema::{DataType, Field, Schema, TimeUnit}; use chrono::Utc; use temp_dir::TempDir; @@ -663,6 +693,7 @@ mod tests { reader::{MergedForwardRecordReader, MergedReverseRecordReader, OffsetReader}, writer::DiskWriter, }, + schema_registry::Resolution, utils::time::TimeRange, }; @@ -996,6 +1027,163 @@ mod tests { Ok(()) } + #[test] + fn schema_resolution_rejects_incompatible_local_raw_schemas() { + let dir = TempDir::new().unwrap(); + let int_path = dir.path().join("int.data.arrows"); + let text_path = dir.path().join("text.data.arrows"); + let timestamp_type = DataType::Timestamp(TimeUnit::Millisecond, None); + let int_schema = Arc::new(Schema::new(vec![ + Field::new("ts", timestamp_type.clone(), false), + Field::new("value", DataType::Int32, false), + ])); + let text_schema = Arc::new(Schema::new(vec![ + Field::new("ts", timestamp_type.clone(), false), + Field::new("value", DataType::Utf8, false), + ])); + let int_batch = RecordBatch::try_new( + int_schema, + vec![ + Arc::new(TimestampMillisecondArray::from(vec![200])), + Arc::new(Int32Array::from(vec![7])), + ], + ) + .unwrap(); + let text_batch = RecordBatch::try_new( + text_schema, + vec![ + Arc::new(TimestampMillisecondArray::from(vec![100])), + Arc::new(StringArray::from(vec!["low"])), + ], + ) + .unwrap(); + std::fs::write(&int_path, write_mem(&[int_batch])).unwrap(); + std::fs::write(&text_path, write_mem(&[text_batch])).unwrap(); + let file_paths = vec![int_path, text_path]; + + let mut targets = HashMap::new(); + targets.insert( + ("ts".to_string(), timestamp_type.clone()), + Arc::new(Field::new("ts", timestamp_type, false)), + ); + targets.insert( + ("value".to_string(), DataType::Int32), + Arc::new(Field::new("value_i32", DataType::Int32, true)), + ); + targets.insert( + ("value".to_string(), DataType::Utf8), + Arc::new(Field::new("value_utf8", DataType::Utf8, true)), + ); + let resolution = Resolution::from_targets(targets); + + let forward_reader = MergedForwardRecordReader::try_new(&file_paths); + assert!(forward_reader.merged_schema().is_err()); + assert!( + forward_reader + .merged_schema() + .and_then(|schema| resolution.rename_schema(&schema)) + .is_err() + ); + + let reverse_reader = MergedReverseRecordReader::try_new(&file_paths); + assert!(reverse_reader.merged_schema().is_err()); + assert!( + reverse_reader + .merged_schema() + .and_then(|schema| resolution.rename_schema(&schema)) + .is_err() + ); + } + + #[test] + fn schema_resolution_adapts_compatible_local_batches_before_renaming() { + let dir = TempDir::new().unwrap(); + let null_path = dir.path().join("null.data.arrows"); + let int_path = dir.path().join("int.data.arrows"); + let timestamp_type = DataType::Timestamp(TimeUnit::Millisecond, None); + let null_schema = Arc::new(Schema::new(vec![ + Field::new("ts", timestamp_type.clone(), false), + Field::new("value", DataType::Null, true), + ])); + let int_schema = Arc::new(Schema::new(vec![ + Field::new("ts", timestamp_type.clone(), false), + Field::new("value", DataType::Int32, false), + ])); + let null_batch = RecordBatch::try_new( + null_schema, + vec![ + Arc::new(TimestampMillisecondArray::from(vec![100, 101])), + Arc::new(NullArray::new(2)), + ], + ) + .unwrap(); + let int_batch = RecordBatch::try_new( + int_schema, + vec![ + Arc::new(TimestampMillisecondArray::from(vec![200, 201])), + Arc::new(Int32Array::from(vec![7, 8])), + ], + ) + .unwrap(); + std::fs::write(&null_path, write_mem(&[null_batch])).unwrap(); + std::fs::write(&int_path, write_mem(&[int_batch])).unwrap(); + let file_paths = vec![null_path, int_path]; + + let mut targets = HashMap::new(); + targets.insert( + ("ts".to_string(), timestamp_type.clone()), + Arc::new(Field::new("ts", timestamp_type, false)), + ); + targets.insert( + ("value".to_string(), DataType::Int32), + Arc::new(Field::new("value_i32", DataType::Int32, true)), + ); + let resolution = Arc::new(Resolution::from_targets(targets)); + + let forward_reader = MergedForwardRecordReader::try_new(&file_paths); + let forward_local_schema = forward_reader.merged_schema().unwrap(); + assert_eq!( + forward_local_schema + .field_with_name("value") + .unwrap() + .data_type(), + &DataType::Int32 + ); + assert!( + forward_reader + .merged_schema() + .and_then(|schema| resolution.rename_schema(&schema)) + .unwrap() + .field_with_name("value_i32") + .is_ok() + ); + + let reverse_reader = MergedReverseRecordReader::try_new(&file_paths); + let local_schema = Arc::new(reverse_reader.merged_schema().unwrap()); + let schema = Arc::new(resolution.rename_schema(&local_schema).unwrap()); + let batches = reverse_reader + .merged_iter_with_resolution(schema, local_schema, None, resolution) + .collect::, _>>() + .unwrap(); + + assert_eq!(batches.len(), 2); + let int_values = batches[0] + .column_by_name("value_i32") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(int_values.values(), &[8, 7]); + let null_values = batches[1] + .column_by_name("value_i32") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(null_values.len(), 2); + assert_eq!(null_values.null_count(), 2); + } + #[test] fn forward_reader_lanes_preserve_original_batch_order() -> io::Result<()> { let dir = TempDir::new().unwrap(); diff --git a/src/parseable/streams.rs b/src/parseable/streams.rs index 6832bd7f6..8f8024bcd 100644 --- a/src/parseable/streams.rs +++ b/src/parseable/streams.rs @@ -51,7 +51,8 @@ use rayon::{ slice::ParallelSliceMut, }; use relative_path::RelativePathBuf; -use std::sync::PoisonError; +use rustc_hash::FxHasher; +use serde::{Deserialize, Serialize}; use std::{ collections::{HashMap, HashSet, VecDeque}, fs::{self, File, OpenOptions, remove_file, write}, @@ -61,6 +62,10 @@ use std::{ sync::{Arc, Mutex, RwLock}, time::{Duration, Instant, SystemTime, UNIX_EPOCH}, }; +use std::{ + hash::{Hash, Hasher}, + sync::PoisonError, +}; use tokio::task::JoinSet; use tracing::{error, info, info_span, instrument, trace, warn}; use ulid::Ulid; @@ -78,7 +83,8 @@ use crate::{ metrics, option::Mode, parseable::{DEFAULT_TENANT, PARSEABLE}, - storage::{StreamType, object_storage::to_bytes, retention::Retention}, + schema_registry::Resolution, + storage::{ComparableBytes, StreamType, object_storage::to_bytes, retention::Retention}, sync::FLUSH_AND_CONVERT_RUNTIME, utils::{ arrow::adapt_batch, @@ -372,6 +378,27 @@ fn minute_from_system_time(time: SystemTime) -> u128 { / 60000 } +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct GlobalSchema { + pub schema: Schema, + pub hash: u64, +} + +impl GlobalSchema { + pub fn from_schema(schema: Schema) -> Self { + let mut hasher = FxHasher::with_seed(0); + schema.hash(&mut hasher); + let hash = hasher.finish(); + Self { schema, hash } + } +} + +impl ComparableBytes for GlobalSchema { + fn get_hash(&self) -> u64 { + self.hash + } +} + /// All state associated with a single logstream in Parseable. pub struct Stream { pub stream_name: String, @@ -380,6 +407,8 @@ pub struct Stream { pub options: Arc, pub writer: Mutex, schema_writer: Mutex<()>, + global_schema: RwLock, Arc)>>, + global_schema_resolution: Mutex<()>, pub ingestor_id: Option, } @@ -506,17 +535,36 @@ struct MetricForwardRecordIterator { active_files: VecDeque, readers_in_flight: usize, setup_error_queued: bool, - schema: Arc, + final_schema: Arc, + local_schema: Arc, + resolution: Arc, } impl MetricForwardRecordIterator { + #[cfg(test)] fn new(reader: MergedForwardRecordReader, schema: Arc) -> Self { + Self::with_resolution( + reader, + schema.clone(), + schema, + Arc::new(Resolution::identity()), + ) + } + + fn with_resolution( + reader: MergedForwardRecordReader, + final_schema: Arc, + local_schema: Arc, + resolution: Arc, + ) -> Self { Self { plans: reader.into_file_plans(), active_files: VecDeque::new(), readers_in_flight: 0, setup_error_queued: false, - schema, + final_schema, + local_schema, + resolution, } } @@ -571,7 +619,19 @@ impl Iterator for MetricForwardRecordIterator { match &mut queued.state { MetricForwardFileState::Ready(active) => { if let Some(batch) = active.next_batch() { - return Some(batch.map(|batch| adapt_batch(self.schema.clone(), &batch))); + return Some(batch.and_then(|batch| { + crate::utils::arrow::batch_adapter::try_adapt_batch( + self.local_schema.clone(), + &batch, + ) + .and_then(|batch| self.resolution.rename_batch(&batch)) + .and_then(|batch| { + crate::utils::arrow::batch_adapter::try_adapt_batch( + self.final_schema.clone(), + &batch, + ) + }) + })); } } MetricForwardFileState::Failed(_) => {} @@ -624,21 +684,33 @@ impl ConversionRecordReader { } } - fn merged_schema(&self) -> Schema { + fn merged_schema(&self) -> Result { match self { Self::Forward(reader) => reader.merged_schema(), Self::Reverse(reader) => reader.merged_schema(), } } - fn merged_iter( + fn merged_iter_with_resolution( self, - schema: Arc, + final_schema: Arc, + local_schema: Arc, time_partition: Option, + resolution: Arc, ) -> Box> + Send> { match self { - Self::Forward(reader) => Box::new(MetricForwardRecordIterator::new(reader, schema)), - Self::Reverse(reader) => Box::new(reader.merged_iter(schema, time_partition)), + Self::Forward(reader) => Box::new(MetricForwardRecordIterator::with_resolution( + reader, + final_schema, + local_schema, + resolution, + )), + Self::Reverse(reader) => Box::new(reader.merged_iter_with_resolution( + final_schema, + local_schema, + time_partition, + resolution, + )), } } } @@ -686,6 +758,8 @@ impl Stream { options, writer: Mutex::new(Writer::default()), schema_writer: Mutex::new(()), + global_schema: RwLock::new(None), + global_schema_resolution: Mutex::new(()), ingestor_id, }) } @@ -1112,11 +1186,20 @@ impl Stream { let path = RelativePathBuf::from_iter([file_name]).to_path(&self.data_path); let tmp_path = path.with_extension("schema.tmp"); - let staging_schemas = self.get_schemas_if_present()?; - if !staging_schemas.is_empty() { - let mut staging_schemas = staging_schemas; - staging_schemas.push(schema); - schema = Schema::try_merge(staging_schemas)?; + let canonical = self.global_schema.read().expect(LOCK_EXPECT).clone(); + if self.uses_global_schema() + && let Some((global, _)) = canonical + { + // The registry is authoritative. Raw schemas staged before the + // first coordinated conversion may contain the losing field type. + schema = global.as_ref().to_owned(); + } else { + let staging_schemas = self.get_schemas_if_present()?; + if !staging_schemas.is_empty() { + let mut staging_schemas = staging_schemas; + staging_schemas.push(schema); + schema = Schema::try_merge(staging_schemas)?; + } } // save the merged schema on staging disk @@ -1140,7 +1223,40 @@ impl Stream { })?; if let Some(mem) = writer.mem.as_mut() { - mem.recordbatch_cloned(schema) + let global = self.global_schema.read().expect(LOCK_EXPECT).clone(); + if self.uses_global_schema() + && let Some((_, resolution)) = global + { + let raw_fields = self + .get_schema_raw() + .into_values() + .sorted_by_key(|field| field.name().clone()) + .collect::(); + let batches = mem.recordbatch_cloned(&Arc::new(Schema::new(raw_fields)))?; + batches + .into_iter() + .map(|batch| { + // Newly ingested, unclaimed columns are not yet part of + // the published schema; retain them in the raw buffers. + let indices: Vec<_> = batch + .schema() + .fields() + .iter() + .enumerate() + .filter_map(|(index, field)| { + resolution + .target(field.name(), field.data_type()) + .map(|_| index) + }) + .collect(); + let projected = batch.project(&indices)?; + let renamed = resolution.rename_batch(&projected)?; + Ok(adapt_batch(schema.clone(), &renamed)) + }) + .collect() + } else { + mem.recordbatch_cloned(schema) + } } else { Ok(Vec::new()) } @@ -1229,7 +1345,7 @@ impl Stream { // queried name still gets read. The bloom answers membership exactly, // so a row group that never saw the metric is rejected outright. let column_path = ColumnPath::new(vec!["metric_name".to_string()]); - let bloom_filter_position = if PARSEABLE.options.bloom_filter_default_position { + let bloom_filter_position = if self.options.bloom_filter_default_position { BloomFilterPosition::AfterRowGroup } else { BloomFilterPosition::End @@ -1402,6 +1518,8 @@ impl Stream { time_partition: Option, time_partition_field: String, target_rows: usize, + local_schema: Arc, + resolution: Arc, ) -> std::sync::mpsc::Receiver { // Metric batches are decoded by independent per-file reader lanes and // reassembled in source order. Concat and sort are also dispatched @@ -1415,7 +1533,12 @@ impl Stream { let mut buffered_record_batches = 0; let mut pending_preparations = VecDeque::with_capacity(*METRIC_ROW_GROUP_PREP_IN_FLIGHT); - let mut merged_iter = record_reader.merged_iter(schema.clone(), time_partition); + let mut merged_iter = record_reader.merged_iter_with_resolution( + schema.clone(), + local_schema, + time_partition, + resolution, + ); loop { let read_started = Instant::now(); @@ -1765,6 +1888,148 @@ impl Stream { ) } + fn can_use_global_schema(&self) -> bool { + self.options.mode == Mode::Ingest && !self.get_static_schema_flag() + } + + fn uses_global_schema(&self) -> bool { + self.can_use_global_schema() + && (crate::schema_registry::enabled() + || self.global_schema.read().expect(LOCK_EXPECT).is_some()) + } + + /// Merges the raw schemas of readable local Arrow files. Canonical stream + /// metadata is intentionally excluded: it may contain an already renamed + /// field that conflicts with an unconverted local Arrow field. + fn merge_local_arrow_schema( + &self, + staging_files: &[(PathBuf, Vec)], + ) -> Result, StagingError> { + let mut schemas = Vec::new(); + for (_, files) in staging_files { + let reader = ConversionRecordReader::for_stream(files, self.is_otel_metrics()); + if reader.reader_count() != 0 { + schemas.push(reader.merged_schema()?); + } + } + if schemas.is_empty() { + Ok(None) + } else { + Ok(Some(Schema::try_merge(schemas)?)) + } + } + + /// Resolves the one locally merged raw schema against a shared, append-only + /// schema before any Parquet writer is created. This runs on the conversion + /// blocking pool. + fn resolve_global_schema( + &self, + local_schema: &Schema, + tenant_id: &Option, + ) -> Result, StagingError> { + let tenant_id = if let Some(tenant) = tenant_id + && tenant.eq(DEFAULT_TENANT) + { + &None + } else { + tenant_id + }; + + let handle = tokio::runtime::Handle::try_current() + .map_err(|err| StagingError::GlobalSchema(err.to_string()))?; + // The opt-out prevents bootstrap, not safe handling of a registry that + // another node already created. Reverting to raw writes would violate + // that dataset's published schema. Probe even for newly loaded streams. + if !crate::schema_registry::enabled() + && self.global_schema.read().expect(LOCK_EXPECT).is_none() + && handle + .block_on( + PARSEABLE + .metastore + .get_global_schema(&self.stream_name, tenant_id), + ) + .map_err(|err| StagingError::GlobalSchema(err.to_string()))? + .is_none() + { + return Ok(Arc::new(Resolution::identity())); + } + let _guard = self.global_schema_resolution.lock().map_err(|err| { + StagingError::GlobalSchema(format!("global-schema lock poisoned: {err}")) + })?; + let mut pinned = HashSet::from([DEFAULT_TIMESTAMP_KEY.to_owned()]); + pinned.extend(self.get_time_partition()); + if let Some(custom) = self.get_custom_partition() { + pinned.extend(custom.split(',').map(|name| name.trim().to_owned())); + } + let (mut resolution, canonical) = handle.block_on(async { + let resolution = crate::schema_registry::resolve( + PARSEABLE.metastore.as_ref(), + &self.stream_name, + tenant_id, + std::slice::from_ref(local_schema), + &pinned, + ) + .await + .map_err(|err| StagingError::GlobalSchema(err.to_string()))?; + let global = PARSEABLE + .metastore + .get_global_schema(&self.stream_name, tenant_id) + .await + .map_err(|err| StagingError::GlobalSchema(err.to_string()))? + .ok_or_else(|| { + StagingError::GlobalSchema( + "global schema disappeared after resolution".to_owned(), + ) + })?; + let canonical = global.schema; + Ok::<_, StagingError>((resolution, Arc::new(canonical))) + })?; + resolution.include_identity_schema(&canonical); + let resolution = Arc::new(resolution); + *self.global_schema.write().expect(LOCK_EXPECT) = Some((canonical, resolution.clone())); + Ok(resolution) + } + + fn requeue_unresolved_arrow_files(&self, staging_files: &[(PathBuf, Vec)]) { + for (_, files) in staging_files { + for source in files { + let Some(name) = source.file_name().and_then(|name| name.to_str()) else { + continue; + }; + let Some((_, suffix)) = name.split_once('.') else { + continue; + }; + let destination = self.data_path.join(format!("{}.{}", Ulid::new(), suffix)); + // Creating a link is exclusive: never overwrite a live writer's + // finalized file. Source and destination share the staging FS. + if let Err(err) = fs::hard_link(source, &destination) { + error!( + "Could not requeue unresolved Arrow file {}: {err}", + source.display() + ); + continue; + } + if let Err(err) = fs::remove_file(source) { + // Do not leave a duplicate for the next conversion cycle. + if let Err(cleanup) = fs::remove_file(&destination) { + error!( + "Could not remove duplicate requeued file {}: {cleanup}", + destination.display() + ); + } + error!( + "Could not remove original requeued file {}: {err}", + source.display() + ); + continue; + } + if let Some(directory) = source.parent() { + let _ = fs::remove_dir(directory); + } + } + } + } + fn convert_arrow_file_groups_to_parquet( &self, staging_files: Vec<(PathBuf, Vec)>, @@ -1779,6 +2044,32 @@ impl Stream { } self.update_staging_metrics(&staging_files, tenant_id); + // Coordinate once for the entire cycle, before any group writes Parquet. + // A failure leaves every source Arrow file available for a later retry. + let (local_schema, resolution) = if self.can_use_global_schema() { + match self.merge_local_arrow_schema(&staging_files) { + Ok(Some(schema)) => { + let local_schema = Arc::new(schema); + let resolution = match self.resolve_global_schema(&local_schema, tenant_id) { + Ok(resolution) => resolution, + Err(err) => { + // Periodic sync claims only root files. Return this snapshot to + // the root so a metastore outage is retried without a restart. + self.requeue_unresolved_arrow_files(&staging_files); + return Err(err); + } + }; + (Some(local_schema), resolution) + } + Ok(None) => (None, Arc::new(Resolution::identity())), + Err(err) => { + self.requeue_unresolved_arrow_files(&staging_files); + return Err(err); + } + } + } else { + (None, Arc::new(Resolution::identity())) + }; let mut schemas = Vec::new(); let mut first_error = None; @@ -1789,6 +2080,8 @@ impl Stream { time_partition, custom_partition, tenant_id, + local_schema.clone(), + &resolution, ); (parquet_path, result) }; @@ -1861,6 +2154,7 @@ impl Stream { /// Converts one Parquet output group and removes its Arrow sources only /// after the final Parquet rename succeeds. + #[allow(clippy::too_many_arguments)] fn convert_arrow_group( &self, parquet_path: PathBuf, @@ -1868,6 +2162,8 @@ impl Stream { time_partition: Option<&String>, custom_partition: Option<&String>, tenant_id: &Option, + local_schema: Option>, + resolution: &Arc, ) -> Result, StagingError> { // Metrics sort every output row group explicitly, so reading their IPC // streams forward avoids reverse seeks and a full-column row reversal. @@ -1879,7 +2175,12 @@ impl Stream { return Ok(None); } - let merged_schema = record_reader.merged_schema(); + let local_schema = match local_schema { + Some(schema) => schema, + None => Arc::new(record_reader.merged_schema()?), + }; + let merged_schema = resolution.rename_schema(&local_schema)?; + let props = self.parquet_writer_props(&merged_schema, time_partition, custom_partition); let schema = Arc::new(merged_schema.clone()); let mut part_path = parquet_path.clone(); @@ -1891,6 +2192,8 @@ impl Stream { &schema, &props, time_partition, + local_schema, + resolution.clone(), ); match write_result { Ok(true) => {} @@ -1931,6 +2234,7 @@ impl Stream { #[cfg_attr(feature = "hotpath", hotpath::measure)] /// Writes one conversion group to a temporary Parquet file. + #[allow(clippy::too_many_arguments)] fn write_parquet_part_file( &self, part_path: &Path, @@ -1938,6 +2242,8 @@ impl Stream { schema: &Arc, props: &WriterProperties, time_partition: Option<&String>, + local_schema: Arc, + resolution: Arc, ) -> Result { let _span = info_span!( "write_parquet_part_file", @@ -1958,11 +2264,18 @@ impl Stream { schema, props, time_partition, + local_schema, + resolution, )?; } else { let mut writer = ArrowWriter::try_new(&mut part_file, schema.clone(), Some(props.clone()))?; - for record in record_reader.merged_iter(schema.clone(), time_partition.cloned()) { + for record in record_reader.merged_iter_with_resolution( + schema.clone(), + local_schema, + time_partition.cloned(), + resolution, + ) { writer.write(&record?)?; } writer.close()?; @@ -1984,6 +2297,7 @@ impl Stream { /// Writes a metric conversion group with pipelined row-group preparation /// and encoding while preserving row-group output order. + #[allow(clippy::too_many_arguments)] fn write_metric_parquet_part_file( &self, part_path: &Path, @@ -1992,6 +2306,8 @@ impl Stream { schema: &Arc, props: &WriterProperties, time_partition: Option<&String>, + local_schema: Arc, + resolution: Arc, ) -> Result<(), StagingError> { let conversion_started = Instant::now(); let source_arrow_files = record_reader.readable_files().len(); @@ -2033,6 +2349,8 @@ impl Stream { time_partition.cloned(), time_partition_field, target, + local_schema, + resolution, ); let mut preparation_complete = false; let mut next_row_group_index = writer.flushed_row_groups().len(); @@ -2287,6 +2605,51 @@ impl Stream { self.cleanup_arrow_files_and_dir(arrow_files, tenant_id); } + /// Restore an already coordinated dataset without merging raw, pre-flush + /// Arrow field names directly into its canonical schema. + pub async fn restore_global_schema( + &self, + current_schema: Schema, + tenant_id: &Option, + time_partition: Option<&String>, + custom_partition: Option<&String>, + ) -> Result { + let reader = MergedReverseRecordReader::try_new(&self.arrow_files()); + let mut schemas = vec![current_schema]; + if !reader.readers.is_empty() { + // Raw Arrow files must agree locally before they can be resolved + // alongside persisted canonical metadata. + schemas.insert(0, reader.merged_schema()?); + } + let mut pinned = HashSet::from([DEFAULT_TIMESTAMP_KEY.to_owned()]); + pinned.extend(time_partition.cloned()); + if let Some(custom) = custom_partition { + pinned.extend(custom.split(',').map(|name| name.trim().to_owned())); + } + let mut resolution = crate::schema_registry::resolve( + PARSEABLE.metastore.as_ref(), + &self.stream_name, + tenant_id, + &schemas, + &pinned, + ) + .await + .map_err(|err| StagingError::GlobalSchema(err.to_string()))?; + let global = PARSEABLE + .metastore + .get_global_schema(&self.stream_name, tenant_id) + .await + .map_err(|err| StagingError::GlobalSchema(err.to_string()))? + .ok_or_else(|| { + StagingError::GlobalSchema("global schema disappeared during startup".to_owned()) + })?; + let canonical = global.schema; + resolution.include_identity_schema(&canonical); + *self.global_schema.write().expect(LOCK_EXPECT) = + Some((Arc::new(canonical.clone()), Arc::new(resolution))); + Ok(canonical) + } + pub fn updated_schema(&self, current_schema: Schema) -> Schema { let staging_files = self.arrow_files(); let record_reader = MergedReverseRecordReader::try_new(&staging_files); @@ -2294,7 +2657,7 @@ impl Stream { return current_schema; } - let schema = record_reader.merged_schema(); + let schema = record_reader.merged_schema().unwrap(); Schema::try_merge(vec![schema, current_schema]).unwrap() } @@ -2352,6 +2715,11 @@ impl Stream { } pub fn get_schema(&self) -> Arc { + if self.uses_global_schema() + && let Some((schema, _)) = self.global_schema.read().expect(LOCK_EXPECT).as_ref() + { + return schema.clone(); + } let metadata = self.metadata.read().expect(LOCK_EXPECT); // sort fields on read from hashmap as order of fields can differ. @@ -2967,6 +3335,190 @@ mod tests { use super::*; + #[test] + fn resolved_conversion_preserves_conflicting_columns_for_logs_and_metrics() { + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + + for forward in [false, true] { + let temp = TempDir::new().unwrap(); + let options = Arc::new(Options { + local_staging_path: temp.path().to_path_buf(), + row_group_size: 2, + ..Default::default() + }); + let mut metadata = LogStreamMetadata::default(); + if forward { + metadata.log_source = + vec![LogSourceEntry::new(LogSource::OtelMetrics, HashSet::new())]; + } + let stream = Stream::new(options, "resolved", metadata, None, &None); + fs::create_dir_all(&stream.data_path).unwrap(); + let timestamp_type = DataType::Timestamp(TimeUnit::Millisecond, None); + let schema = Arc::new(Schema::new(vec![ + Field::new(DEFAULT_TIMESTAMP_KEY, timestamp_type.clone(), true), + Field::new("metric_name", DataType::Utf8, true), + Field::new("value", DataType::Int32, true), + ])); + let mut files = Vec::new(); + for (index, values) in [[11, 22], [33, 44]].into_iter().enumerate() { + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(TimestampMillisecondArray::from(vec![ + (index * 2 + 1) as i64, + (index * 2 + 2) as i64, + ])), + Arc::new(StringArray::from(vec!["cpu", "cpu"])), + Arc::new(Int32Array::from(values.to_vec())), + ], + ) + .unwrap(); + let path = stream.data_path.join(format!("node{index}.arrows")); + let mut writer = + ArrowStreamWriter::try_new(File::create(&path).unwrap(), &schema).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + files.push(path); + } + let local_schema = Arc::new(Schema::try_merge([schema.as_ref().clone()]).unwrap()); + let resolution = Arc::new(Resolution::from_targets(HashMap::from([ + ( + (DEFAULT_TIMESTAMP_KEY.to_owned(), timestamp_type.clone()), + Arc::new(Field::new(DEFAULT_TIMESTAMP_KEY, timestamp_type, true)), + ), + ( + ("metric_name".to_owned(), DataType::Utf8), + Arc::new(Field::new("metric_name", DataType::Utf8, true)), + ), + ( + ("value".to_owned(), DataType::Int32), + Arc::new(Field::new("value_int32", DataType::Int32, true)), + ), + ]))); + let parquet_path = stream.data_path.join("resolved.parquet"); + let written_schema = stream + .convert_arrow_group( + parquet_path.clone(), + files.clone(), + None, + None, + &None, + Some(local_schema), + &resolution, + ) + .unwrap() + .unwrap(); + assert!(files.iter().all(|path| !path.exists())); + let reader = + ParquetRecordBatchReaderBuilder::try_new(File::open(&parquet_path).unwrap()) + .unwrap(); + Schema::try_merge([written_schema, reader.schema().as_ref().clone()]).unwrap(); + let batches = reader + .build() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 4); + let mut integers = Vec::new(); + for batch in batches { + assert!(batch.column_by_name("value").is_none()); + integers.extend( + batch + .column_by_name("value_int32") + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .flatten(), + ); + } + integers.sort(); + assert_eq!(integers, vec![11, 22, 33, 44]); + } + } + + #[test] + fn local_raw_schema_conflicts_fail_before_resolution() { + let temp = TempDir::new().unwrap(); + let stream = Stream::new( + Arc::new(Options { + local_staging_path: temp.path().to_path_buf(), + ..Default::default() + }), + "local-conflict", + LogStreamMetadata::default(), + None, + &None, + ); + fs::create_dir_all(&stream.data_path).unwrap(); + let mut files = Vec::new(); + for (index, field) in [ + Field::new("value", DataType::Int32, true), + Field::new("value", DataType::Utf8, true), + ] + .into_iter() + .enumerate() + { + let schema = Arc::new(Schema::new(vec![field])); + let path = stream.data_path.join(format!("conflict-{index}.arrows")); + let mut writer = + ArrowStreamWriter::try_new(File::create(&path).unwrap(), &schema).unwrap(); + let values: ArrayRef = if index == 0 { + Arc::new(Int32Array::from(vec![1])) + } else { + Arc::new(StringArray::from(vec!["one"])) + }; + writer + .write(&RecordBatch::try_new(schema, vec![values]).unwrap()) + .unwrap(); + writer.finish().unwrap(); + files.push(path); + } + + assert!( + stream + .merge_local_arrow_schema(&[(stream.data_path.join("conflict.parquet"), files)]) + .is_err() + ); + } + + #[test] + fn unresolved_sources_are_requeued_without_overwriting_new_files() { + let temp = TempDir::new().unwrap(); + let stream = Stream::new( + Arc::new(Options { + local_staging_path: temp.path().to_path_buf(), + ..Default::default() + }), + "retry", + LogStreamMetadata::default(), + None, + &None, + ); + let processing = stream.data_path.join("processing_retry"); + fs::create_dir_all(&processing).unwrap(); + let name = "key.date=2026-10-06.hour=04.minute=01.host.data.arrows"; + let original = processing.join(name); + let live = stream.data_path.join(name); + fs::write(&original, b"unresolved").unwrap(); + fs::write(&live, b"live").unwrap(); + stream.requeue_unresolved_arrow_files(&[( + stream.data_path.join("retry.parquet"), + vec![original.clone()], + )]); + assert!(!original.exists()); + assert_eq!(fs::read(&live).unwrap(), b"live"); + let queued = stream.arrow_files(); + assert_eq!(queued.len(), 2); + assert!( + queued + .iter() + .any(|path| fs::read(path).unwrap() == b"unresolved") + ); + assert!(!processing.exists()); + } + #[test] fn parallel_metric_writer_creates_valid_ordered_row_groups() { let temp_dir = TempDir::new().unwrap(); @@ -3107,7 +3659,7 @@ mod tests { let record_reader = ConversionRecordReader::Forward(MergedForwardRecordReader::try_new(&arrow_files)); - let merged_schema = Arc::new(record_reader.merged_schema()); + let merged_schema = Arc::new(record_reader.merged_schema().unwrap()); let props = WriterProperties::builder() .set_max_row_group_row_count(Some(5)) .build(); @@ -3120,6 +3672,8 @@ mod tests { &merged_schema, &props, None, + merged_schema.clone(), + Arc::new(Resolution::identity()), ) .unwrap() ); diff --git a/src/schema_registry.rs b/src/schema_registry.rs new file mode 100644 index 000000000..6243d55e6 --- /dev/null +++ b/src/schema_registry.rs @@ -0,0 +1,1492 @@ +/* + * Parseable Server (C) 2022 - 2025 Parseable, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as + * published by the Free Software Foundation, either version 3 of the + * License, or (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +//! Append-only, per-stream canonical schema resolution for distributed ingest. +//! +//! Canonical fields are keyed by their published name. Once a name has been +//! published, its datatype is immutable; a conflicting local field is given a +//! deterministic datatype suffix instead. Resolutions only replace schema +//! metadata and retain the original Arrow arrays. + +use std::{ + collections::{BTreeMap, HashMap, HashSet}, + sync::Arc, + time::Duration, +}; + +use arrow_array::{RecordBatch, RecordBatchOptions}; +use arrow_schema::{ArrowError, DataType, Field, FieldRef, Schema}; +use chrono::{DateTime, Utc}; +use tokio::time::sleep; +use tonic::async_trait; + +use crate::{ + event::format::get_datatype_suffix, + handlers::http::modal::ingest_server::INGESTOR_META, + metastore::{MetastoreError, metastore_traits::Metastore}, + option::Mode, + parseable::{GlobalSchema, PARSEABLE}, + storage::object_storage::global_schema_lock_file_path, +}; + +const MAX_CAS_ATTEMPTS: usize = 16; +const MAX_NAME_ATTEMPTS: usize = 64; +static P_CANONICAL_SCHEMA_ENV_KEY: &str = "P_CANONICAL_SCHEMA"; + +#[async_trait] +trait RegistryBackend: Send + Sync { + async fn read_global(&self) -> Result, MetastoreError>; + async fn read_legacy_schemas(&self) -> Result, MetastoreError>; + async fn compare_and_swap( + &self, + global_schema: GlobalSchema, + expected: Option, + ) -> Result; +} + +struct MetastoreRegistryBackend<'a> { + metastore: &'a dyn Metastore, + stream_name: &'a str, + tenant_id: &'a Option, +} + +#[async_trait] +trait GlobalSchemaLockBackend: Send + Sync { + async fn list_locks(&self) -> Result, MetastoreError>; + async fn delete_lock(&self, lock: &str) -> Result<(), MetastoreError>; + async fn wait_for_lock_release(&self); +} + +#[async_trait] +impl GlobalSchemaLockBackend for MetastoreRegistryBackend<'_> { + async fn list_locks(&self) -> Result, MetastoreError> { + self.metastore + .get_global_schema_locks(self.stream_name, self.tenant_id) + .await + } + + async fn delete_lock(&self, lock: &str) -> Result<(), MetastoreError> { + self.metastore + .delete_global_schema_lock( + global_schema_lock_file_path(self.stream_name, self.tenant_id).join(lock), + self.tenant_id, + ) + .await + } + + async fn wait_for_lock_release(&self) { + sleep(Duration::from_secs(5)).await; + } +} + +#[async_trait] +impl RegistryBackend for MetastoreRegistryBackend<'_> { + async fn read_global(&self) -> Result, MetastoreError> { + self.metastore + .get_global_schema(self.stream_name, self.tenant_id) + .await + } + + async fn read_legacy_schemas(&self) -> Result, MetastoreError> { + self.metastore + .get_all_schemas(self.stream_name, self.tenant_id) + .await + } + + async fn compare_and_swap( + &self, + global_schema: GlobalSchema, + expected: Option, + ) -> Result { + self.metastore + .put_global_schema(global_schema, self.stream_name, self.tenant_id, expected) + .await + } +} + +/// Maps each local `(field name, datatype)` to its canonical Arrow field. +#[derive(Debug, Clone)] +pub struct Resolution { + targets: HashMap<(String, DataType), FieldRef>, + identity: bool, +} + +impl Default for Resolution { + fn default() -> Self { + Self { + targets: HashMap::new(), + identity: true, + } + } +} + +impl Resolution { + /// A passthrough resolution. The original schema and batch are retained. + pub fn identity() -> Self { + Self::default() + } + + /// Construct a resolution from the mapping maintained by the stream layer. + pub fn from_targets(targets: HashMap<(String, DataType), FieldRef>) -> Self { + let targets = targets + .into_iter() + .map(|(key, field)| { + ( + key, + Arc::new(Field::new(field.name(), field.data_type().clone(), true)), + ) + }) + .collect(); + Self { + targets, + identity: false, + } + } + + pub fn is_identity(&self) -> bool { + self.identity + } + + /// Find the canonical target for a local field. + pub fn target(&self, name: &str, data_type: &DataType) -> Option<&FieldRef> { + self.targets.get(&(name.to_owned(), data_type.clone())) + } + + /// Retain already published columns in the cached in-memory view. Canonical + /// metadata is not a raw local schema and must not enter the local merge. + pub(crate) fn include_identity_schema(&mut self, schema: &Schema) { + for field in schema + .fields() + .iter() + .filter(|field| !field.data_type().is_null()) + { + self.targets + .entry((field.name().clone(), field.data_type().clone())) + .or_insert_with(|| { + Arc::new(Field::new(field.name(), field.data_type().clone(), true)) + }); + } + } + + /// Resolve a schema without touching its arrays. + pub fn rename_schema(&self, schema: &Schema) -> Result { + if self.identity { + return Ok(schema.clone()); + } + + let mut fields = Vec::with_capacity(schema.fields().len()); + let mut target_names = HashSet::with_capacity(schema.fields().len()); + for field in schema.fields() { + let Some(target) = self.target_field(field)? else { + // Top-level Null fields carry no values and may be omitted. + continue; + }; + if !target_names.insert(target.name().clone()) { + return Err(ArrowError::SchemaError(format!( + "multiple source columns resolve to canonical field '{}'", + target.name() + ))); + } + fields.push(target); + } + Ok(Schema::new(fields)) + } + + /// Rename field metadata while retaining every non-Null array by Arc clone. + pub fn rename_batch(&self, batch: &RecordBatch) -> Result { + if self.identity { + return Ok(batch.clone()); + } + + let mut fields = Vec::with_capacity(batch.num_columns()); + let mut columns = Vec::with_capacity(batch.num_columns()); + let mut target_names = HashSet::with_capacity(batch.num_columns()); + for (field, column) in batch.schema().fields().iter().zip(batch.columns()) { + let Some(target) = self.target_field(field)? else { + continue; + }; + if !target_names.insert(target.name().clone()) { + return Err(ArrowError::SchemaError(format!( + "multiple source columns resolve to canonical field '{}'", + target.name() + ))); + } + fields.push(target); + columns.push(Arc::clone(column)); + } + + RecordBatch::try_new_with_options( + Arc::new(Schema::new(fields)), + columns, + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + ) + } + + fn target_field(&self, field: &Field) -> Result, ArrowError> { + if field.data_type().is_null() { + return Ok(None); + } + self.target(field.name(), field.data_type()) + .filter(|target| target.data_type() == field.data_type()) + .cloned() + .map(Some) + .ok_or_else(|| { + ArrowError::SchemaError(format!( + "no canonical resolution for field '{}' of type {}", + field.name(), + field.data_type() + )) + }) + } +} + +/// Whether staging may bootstrap a new global canonical schema registry. +/// +/// Only the explicit value `false` disables bootstrap; malformed values do not +/// silently disable it. Existing coordinated streams must keep resolving against +/// their registry even when bootstrap is disabled, to preserve published names. +/// The caller remains responsible for applying this only to the staging path. +pub fn enabled() -> bool { + PARSEABLE.options.mode == Mode::Ingest + && !std::env::var(P_CANONICAL_SCHEMA_ENV_KEY) + .is_ok_and(|value| value.trim().eq_ignore_ascii_case("false")) +} + +/// Resolve all schemas in one operation against an append-only global schema. +/// +/// On every compare-and-swap contention, the global schema is re-read and the +/// complete local set is resolved again. No mapping computed against a losing +/// snapshot is returned. +pub async fn resolve( + metastore: &dyn Metastore, + stream_name: &str, + tenant_id: &Option, + schemas: &[Schema], + pinned: &HashSet, +) -> Result { + let backend = MetastoreRegistryBackend { + metastore, + stream_name, + tenant_id, + }; + // Before beginning, clear only stale locks. Fresh locks are never deleted by + // another resolver; after one bounded wait they must have disappeared. + ensure_clean_lock_state(&backend).await?; + let lock_file_path = global_schema_lock_file_path(stream_name, tenant_id).join(format!( + "{}-{}.lock", + INGESTOR_META + .get() + .expect("Ingestor node metadata should exist") + .node_id, + Utc::now().timestamp_micros() + )); + backend + .metastore + .put_global_schema_lock(lock_file_path.clone(), tenant_id) + .await?; + + // Always attempt release after resolution. Preserve the resolution error if + // both resolution and cleanup fail; otherwise surface a cleanup failure. + let resolution = resolve_backend(&backend, stream_name, schemas, pinned).await; + let cleanup = backend + .metastore + .delete_global_schema_lock(lock_file_path, tenant_id) + .await; + finish_resolution(resolution, cleanup) +} + +async fn ensure_clean_lock_state( + backend: &dyn GlobalSchemaLockBackend, +) -> Result<(), MetastoreError> { + let locks = backend.list_locks().await?; + let now = Utc::now(); + let parsed_locks = locks + .iter() + .map(|lock| parse_lock_timestamp(lock).map(|timestamp| (lock, timestamp))) + .collect::, _>>()?; + + let mut fresh_lock_exists = false; + for (lock, timestamp) in parsed_locks { + if now.signed_duration_since(timestamp).num_seconds() > 180 { + tracing::warn!("removing stale global schema lock file {lock}"); + backend.delete_lock(lock).await?; + } else { + fresh_lock_exists = true; + } + } + + if fresh_lock_exists { + backend.wait_for_lock_release().await; + if !backend.list_locks().await?.is_empty() { + return Err(lock_file_failure( + "Unable to obtain a clean lock file state", + )); + } + } + + Ok(()) +} + +fn parse_lock_timestamp(lock: &str) -> Result, MetastoreError> { + let timestamp = lock + .strip_suffix(".lock") + .and_then(|name| name.rsplit_once('-')) + .filter(|(node, timestamp)| !node.is_empty() && !timestamp.is_empty()) + .and_then(|(_, timestamp)| timestamp.parse::().ok()) + .and_then(DateTime::from_timestamp_micros) + .ok_or_else(|| lock_file_failure(format!("Malformed global schema lock file '{lock}'")))?; + Ok(timestamp) +} + +fn lock_file_failure(message: impl Into) -> MetastoreError { + MetastoreError::LockFileFailure { + message: message.into(), + } +} + +fn finish_resolution( + resolution: Result, + cleanup: Result<(), MetastoreError>, +) -> Result { + match (resolution, cleanup) { + (Err(resolution), Err(cleanup)) => { + tracing::error!( + "failed to release global schema lock after resolution error: {cleanup}" + ); + Err(resolution) + } + (Err(resolution), Ok(())) => Err(resolution), + (Ok(_), Err(cleanup)) => Err(cleanup), + (Ok(resolution), Ok(())) => Ok(resolution), + } +} + +async fn resolve_backend( + backend: &dyn RegistryBackend, + stream_name: &str, + schemas: &[Schema], + pinned: &HashSet, +) -> Result { + for attempt in 0..MAX_CAS_ATTEMPTS { + let global = backend.read_global().await?; + let (canonical, original_schema) = match global.as_ref() { + Some(global) => { + let original = global.schema.to_owned(); + let normalized = normalize_registry(&original).map_err(schema_error)?; + (normalized, Some(original)) + } + None => { + // Legacy schemas are consulted only while bootstrapping an + // absent global registry. Conflicts are handled leniently. + let legacy = backend.read_legacy_schemas().await?; + (legacy_seed(&legacy), None) + } + }; + + let (resolution, proposed) = + build_resolution(&canonical, schemas, pinned).map_err(schema_error)?; + let needs_publish = match (&original_schema, global.is_some()) { + (Some(original), true) => original != &proposed, + (None, false) => !proposed.fields().is_empty(), + _ => true, + }; + + if !needs_publish { + return Ok(resolution); + } + + // `None` is explicitly create-only; a snapshot version is an + // update-only CAS. A failed CAS invalidates this whole resolution. + let expected = global.as_ref().map(|snapshot| snapshot.hash); + if backend + .compare_and_swap(GlobalSchema::from_schema(proposed), expected) + .await? + { + return Ok(resolution); + } + + if attempt + 1 < MAX_CAS_ATTEMPTS { + let backoff_ms = 2_u64.saturating_mul(1_u64 << attempt.min(6)); + tokio::time::sleep(Duration::from_millis(backoff_ms)).await; + } + } + + Err(schema_error(format!( + "global schema compare-and-swap did not converge after {MAX_CAS_ATTEMPTS} attempts for stream '{stream_name}'" + ))) +} + +fn schema_error(message: impl Into) -> MetastoreError { + MetastoreError::CASConvergenceFailure { + message: message.into(), + } +} + +/// Make a field match the registry's canonical representation. +fn canonical_field(field: &Field) -> FieldRef { + Arc::new(Field::new(field.name(), field.data_type().clone(), true)) +} + +/// Normalize a persisted registry and reject entries that violate its identity +/// rules. Existing `(name, datatype)` claims are never rewritten to a new type. +fn normalize_registry(schema: &Schema) -> Result { + let mut fields = BTreeMap::::new(); + for field in schema.fields() { + if field.data_type().is_null() { + return Err(format!( + "global schema contains unsupported Null field '{}'", + field.name() + )); + } + if fields + .insert(field.name().clone(), canonical_field(field)) + .is_some() + { + return Err(format!( + "global schema contains duplicate field name '{}'", + field.name() + )); + } + } + Ok(Schema::new(fields.into_values().collect::>())) +} + +/// Build a stable, tolerant union of the old per-node schemas for bootstrapping. +fn legacy_seed(schemas: &[Schema]) -> Schema { + let mut chosen = BTreeMap::::new(); + for field in schemas.iter().flat_map(|schema| schema.fields()) { + if field.data_type().is_null() { + continue; + } + let replace = chosen + .get(field.name()) + .is_none_or(|current| legacy_type_precedes(field.data_type(), current.data_type())); + if replace { + chosen.insert(field.name().clone(), canonical_field(field)); + } + } + Schema::new(chosen.into_values().collect::>()) +} + +fn legacy_type_precedes(candidate: &DataType, current: &DataType) -> bool { + match (is_utf8_family(candidate), is_utf8_family(current)) { + (true, false) => true, + (false, true) => false, + _ => format!("{candidate:?}") < format!("{current:?}"), + } +} + +fn is_utf8_family(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View + ) +} + +/// Resolve a set of schemas against the canonical snapshot and return the +/// resulting canonical schema. This is intentionally storage-independent. +fn build_resolution( + canonical: &Schema, + schemas: &[Schema], + pinned: &HashSet, +) -> Result<(Resolution, Schema), String> { + let normalized = normalize_registry(canonical)?; + let mut fields = normalized + .fields() + .iter() + .map(|field| (field.name().clone(), field.clone())) + .collect(); + let keys = source_keys(schemas); + let cooccurring_names = cooccurring_source_names(schemas); + + reserve_original_names(&mut fields, &keys, pinned)?; + let targets = resolve_targets(&mut fields, &keys, &cooccurring_names, pinned)?; + validate_unique_targets(schemas, &targets)?; + + let canonical = Schema::new(fields.into_values().collect::>()); + Ok((Resolution::from_targets(targets), canonical)) +} + +fn source_keys(schemas: &[Schema]) -> Vec<(String, DataType)> { + let mut keys: Vec<_> = schemas + .iter() + .flat_map(|schema| schema.fields()) + .filter(|field| !field.data_type().is_null()) + .map(|field| (field.name().clone(), field.data_type().clone())) + .collect::>() + .into_iter() + .collect(); + keys.sort_by(|left, right| { + left.0 + .cmp(&right.0) + .then_with(|| format!("{:?}", left.1).cmp(&format!("{:?}", right.1))) + }); + keys +} + +// Record only names of distinct columns that co-occur with each exact source +// key. Schemas processed separately must not reserve names from one another, +// since compatible canonical suffix fields are reusable. +fn cooccurring_source_names(schemas: &[Schema]) -> HashMap<(String, DataType), HashSet> { + let mut cooccurring_names = HashMap::<(String, DataType), HashSet>::new(); + for schema in schemas { + let fields: Vec<_> = schema + .fields() + .iter() + .filter(|field| !field.data_type().is_null()) + .collect(); + for (index, field) in fields.iter().enumerate() { + let names = cooccurring_names + .entry((field.name().clone(), field.data_type().clone())) + .or_default(); + for (other_index, other) in fields.iter().enumerate() { + if index != other_index { + names.insert(other.name().clone()); + } + } + } + } + cooccurring_names +} + +// Reserve free source names first so conflict suffixes do not capture fields +// that require their original canonical names. +fn reserve_original_names( + fields: &mut BTreeMap, + keys: &[(String, DataType)], + pinned: &HashSet, +) -> Result<(), String> { + let mut types_by_name = BTreeMap::>::new(); + for (name, data_type) in keys { + types_by_name + .entry(name.clone()) + .or_default() + .push(data_type.clone()); + } + + for (name, types) in types_by_name { + if fields.contains_key(&name) { + continue; + } + if pinned.contains(&name) && types.len() > 1 { + return Err(format!( + "pinned field '{name}' has incompatible local datatypes" + )); + } + if let Some(data_type) = types.first() { + fields.insert( + name.clone(), + Arc::new(Field::new(&name, data_type.clone(), true)), + ); + } + } + Ok(()) +} + +fn resolve_targets( + fields: &mut BTreeMap, + keys: &[(String, DataType)], + cooccurring_names: &HashMap<(String, DataType), HashSet>, + pinned: &HashSet, +) -> Result, String> { + let mut targets = HashMap::with_capacity(keys.len()); + for (name, data_type) in keys { + let target = choose_canonical_target(fields, name, data_type, cooccurring_names, pinned)?; + targets.insert((name.clone(), data_type.clone()), target); + } + Ok(targets) +} + +fn choose_canonical_target( + fields: &mut BTreeMap, + name: &str, + data_type: &DataType, + cooccurring_names: &HashMap<(String, DataType), HashSet>, + pinned: &HashSet, +) -> Result { + if let Some(field) = fields + .get(name) + .filter(|field| field.data_type() == data_type) + { + return Ok(field.clone()); + } + + if pinned.contains(name) { + let established_type = fields + .get(name) + .map_or_else(|| data_type.clone(), |field| field.data_type().clone()); + return Err(format!( + "pinned field '{name}' has local datatype {data_type} but canonical datatype {established_type}" + )); + } + + choose_suffix_target(fields, name, data_type, cooccurring_names, pinned) +} + +fn choose_suffix_target( + fields: &mut BTreeMap, + name: &str, + data_type: &DataType, + cooccurring_names: &HashMap<(String, DataType), HashSet>, + pinned: &HashSet, +) -> Result { + let suffix = get_datatype_suffix(data_type); + let mut candidate = format!("{name}_{suffix}"); + let source_key = (name.to_owned(), data_type.clone()); + for _ in 0..MAX_NAME_ATTEMPTS { + let conflicts_with_local_column = cooccurring_names + .get(&source_key) + .is_some_and(|names| names.contains(&candidate)); + if !pinned.contains(&candidate) && !conflicts_with_local_column { + match fields.get(&candidate) { + Some(field) if field.data_type() == data_type => return Ok(field.clone()), + None => { + let field = Arc::new(Field::new(&candidate, data_type.clone(), true)); + fields.insert(candidate, field.clone()); + return Ok(field); + } + Some(_) => {} + } + } + candidate.push_str(&format!("_{suffix}")); + } + Err(format!( + "could not find an available canonical name for '{name}' after {MAX_NAME_ATTEMPTS} suffix attempts" + )) +} + +// Ensure each input schema remains lossless: every retained source column must +// have a unique target name in that schema. +fn validate_unique_targets( + schemas: &[Schema], + targets: &HashMap<(String, DataType), FieldRef>, +) -> Result<(), String> { + for schema in schemas { + let mut names = HashSet::with_capacity(schema.fields().len()); + for field in schema + .fields() + .iter() + .filter(|field| !field.data_type().is_null()) + { + let key = (field.name().clone(), field.data_type().clone()); + let target = targets.get(&key).ok_or_else(|| { + format!( + "no canonical resolution for field '{}' of type {}", + field.name(), + field.data_type() + ) + })?; + if !names.insert(target.name().clone()) { + return Err(format!( + "multiple source columns in one schema resolve to canonical field '{}'", + target.name() + )); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::{ + collections::{HashMap, HashSet}, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + }; + + use arrow_array::{ArrayRef, Int64Array, RecordBatch, StringArray}; + use arrow_schema::{DataType, Field, FieldRef, Schema, TimeUnit}; + use chrono::Utc; + use tokio::sync::Barrier; + use tonic::async_trait; + + use crate::{metastore::MetastoreError, parseable::GlobalSchema}; + + use super::{ + GlobalSchemaLockBackend, MAX_CAS_ATTEMPTS, RegistryBackend, Resolution, build_resolution, + canonical_field, ensure_clean_lock_state, finish_resolution, legacy_seed, + normalize_registry, resolve_backend, + }; + + fn field(name: &str, data_type: DataType) -> FieldRef { + Arc::new(Field::new(name, data_type, true)) + } + + fn resolve_local( + canonical: Schema, + schemas: &[Schema], + pinned: &[&str], + ) -> Result<(Resolution, Schema), String> { + let pinned = pinned.iter().map(|name| (*name).to_owned()).collect(); + build_resolution(&canonical, schemas, &pinned) + } + + struct MemoryState { + schema: Option, + legacy: Vec, + } + + struct InMemoryBackend { + state: Mutex, + initial_read_barrier: Option>, + barrier_readers: usize, + global_reads: AtomicUsize, + legacy_reads: AtomicUsize, + cas_attempts: AtomicUsize, + force_contention: bool, + fail_read: bool, + fail_write: bool, + } + + impl InMemoryBackend { + fn new(schema: Option, legacy: Vec, barrier_readers: usize) -> Self { + Self { + state: Mutex::new(MemoryState { schema, legacy }), + initial_read_barrier: (barrier_readers > 0) + .then(|| Arc::new(Barrier::new(barrier_readers))), + barrier_readers, + global_reads: AtomicUsize::new(0), + legacy_reads: AtomicUsize::new(0), + cas_attempts: AtomicUsize::new(0), + force_contention: false, + fail_read: false, + fail_write: false, + } + } + + fn with_contention(mut self) -> Self { + self.force_contention = true; + self + } + + fn with_read_error(mut self) -> Self { + self.fail_read = true; + self + } + + fn with_write_error(mut self) -> Self { + self.fail_write = true; + self + } + + fn global_read_count(&self) -> usize { + self.global_reads.load(Ordering::SeqCst) + } + + fn legacy_read_count(&self) -> usize { + self.legacy_reads.load(Ordering::SeqCst) + } + + fn cas_attempt_count(&self) -> usize { + self.cas_attempts.load(Ordering::SeqCst) + } + + fn canonical_schema(&self) -> Option { + self.state.lock().unwrap().schema.clone() + } + } + + fn injected_storage_error() -> MetastoreError { + MetastoreError::Error { + status_code: actix_web::http::StatusCode::INTERNAL_SERVER_ERROR, + message: "injected registry storage failure".to_owned(), + flow: "schema_registry_test".to_owned(), + } + } + + fn is_injected_storage_error(error: &MetastoreError) -> bool { + matches!( + error, + MetastoreError::Error { message, .. } + if message == "injected registry storage failure" + ) + } + + #[async_trait] + impl RegistryBackend for InMemoryBackend { + async fn read_global(&self) -> Result, MetastoreError> { + let read_number = self.global_reads.fetch_add(1, Ordering::SeqCst); + if self.fail_read { + return Err(injected_storage_error()); + } + + let snapshot = { + let state = self.state.lock().unwrap(); + state.schema.clone().map(GlobalSchema::from_schema) + }; + if read_number < self.barrier_readers { + self.initial_read_barrier + .as_ref() + .expect("barrier configured for the first reads") + .wait() + .await; + } + Ok(snapshot) + } + + async fn read_legacy_schemas(&self) -> Result, MetastoreError> { + self.legacy_reads.fetch_add(1, Ordering::SeqCst); + Ok(self.state.lock().unwrap().legacy.clone()) + } + + async fn compare_and_swap( + &self, + global_schema: GlobalSchema, + expected: Option, + ) -> Result { + self.cas_attempts.fetch_add(1, Ordering::SeqCst); + if self.fail_write { + return Err(injected_storage_error()); + } + if self.force_contention { + return Ok(false); + } + + let mut state = self.state.lock().unwrap(); + let matches_expected = match (&state.schema, expected) { + (None, None) => true, + (Some(schema), Some(expected)) => { + expected == GlobalSchema::from_schema(schema.clone()).hash + } + _ => false, + }; + if !matches_expected { + return Ok(false); + } + + state.schema = Some(global_schema.schema); + Ok(true) + } + } + + struct InMemoryLockBackend { + locks: Mutex>, + deleted: Mutex>, + delete_fails: bool, + clear_on_wait: bool, + waits: AtomicUsize, + } + + impl InMemoryLockBackend { + fn new(locks: impl IntoIterator) -> Self { + Self { + locks: Mutex::new(locks.into_iter().collect()), + deleted: Mutex::new(Vec::new()), + delete_fails: false, + clear_on_wait: false, + waits: AtomicUsize::new(0), + } + } + + fn with_delete_failure(mut self) -> Self { + self.delete_fails = true; + self + } + + fn clear_on_wait(mut self) -> Self { + self.clear_on_wait = true; + self + } + } + + #[async_trait] + impl GlobalSchemaLockBackend for InMemoryLockBackend { + async fn list_locks(&self) -> Result, MetastoreError> { + Ok(self.locks.lock().unwrap().clone()) + } + + async fn delete_lock(&self, lock: &str) -> Result<(), MetastoreError> { + if self.delete_fails { + return Err(injected_storage_error()); + } + self.deleted.lock().unwrap().push(lock.to_owned()); + self.locks.lock().unwrap().retain(|entry| entry != lock); + Ok(()) + } + + async fn wait_for_lock_release(&self) { + self.waits.fetch_add(1, Ordering::SeqCst); + if self.clear_on_wait { + self.locks.lock().unwrap().clear(); + } + } + } + + fn lock_name(node: &str, timestamp_micros: i64) -> String { + format!("{node}-{timestamp_micros}.lock") + } + + fn assert_schemas_merge_into(canonical: &Schema, resolved: impl IntoIterator) { + let mut schemas = vec![canonical.clone()]; + schemas.extend(resolved); + assert_eq!(Schema::try_merge(schemas).unwrap(), canonical.clone()); + } + + #[tokio::test] + async fn lock_cleanup_accepts_zero_locks_and_removes_all_stale_locks() { + let empty = InMemoryLockBackend::new([]); + ensure_clean_lock_state(&empty).await.unwrap(); + assert_eq!(empty.waits.load(Ordering::SeqCst), 0); + + let first_stale = lock_name("first-node", Utc::now().timestamp_micros() - 181_000_000); + let second_stale = lock_name("second-node", Utc::now().timestamp_micros() - 182_000_000); + let backend = InMemoryLockBackend::new([first_stale.clone(), second_stale.clone()]); + ensure_clean_lock_state(&backend).await.unwrap(); + assert_eq!( + *backend.deleted.lock().unwrap(), + vec![first_stale, second_stale] + ); + assert!(backend.locks.lock().unwrap().is_empty()); + } + + #[tokio::test] + async fn lock_cleanup_deletes_only_stale_locks_and_rejects_fresh_contention() { + let stale = lock_name("stale-node", Utc::now().timestamp_micros() - 181_000_000); + let fresh = lock_name("fresh-node", Utc::now().timestamp_micros()); + let backend = InMemoryLockBackend::new([stale.clone(), fresh.clone()]); + + let error = ensure_clean_lock_state(&backend).await.unwrap_err(); + assert!(matches!(error, MetastoreError::LockFileFailure { .. })); + assert_eq!(*backend.deleted.lock().unwrap(), vec![stale]); + assert_eq!(*backend.locks.lock().unwrap(), vec![fresh]); + assert_eq!(backend.waits.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn lock_cleanup_allows_a_fresh_lock_that_is_released_during_wait() { + let fresh = lock_name("node", Utc::now().timestamp_micros()); + let backend = InMemoryLockBackend::new([fresh]).clear_on_wait(); + + ensure_clean_lock_state(&backend).await.unwrap(); + assert!(backend.deleted.lock().unwrap().is_empty()); + assert_eq!(backend.waits.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn lock_cleanup_rejects_malformed_locks_without_deleting_anything() { + for malformed in [ + "not-a-lock", + "node-invalid.lock", + "node-999999999999999999999.lock", + "node-9223372036854775807.lock", + "-123.lock", + ] { + let stale = lock_name("stale-node", Utc::now().timestamp_micros() - 181_000_000); + let backend = InMemoryLockBackend::new([stale, malformed.to_owned()]); + + let error = ensure_clean_lock_state(&backend).await.unwrap_err(); + assert!(matches!(error, MetastoreError::LockFileFailure { .. })); + assert!(backend.deleted.lock().unwrap().is_empty()); + assert_eq!(backend.waits.load(Ordering::SeqCst), 0); + } + } + + #[tokio::test] + async fn lock_cleanup_stops_when_stale_lock_deletion_fails() { + let stale = lock_name("node", Utc::now().timestamp_micros() - 181_000_000); + let backend = InMemoryLockBackend::new([stale]).with_delete_failure(); + + let error = ensure_clean_lock_state(&backend).await.unwrap_err(); + assert!(is_injected_storage_error(&error)); + assert_eq!(backend.locks.lock().unwrap().len(), 1); + assert_eq!(backend.waits.load(Ordering::SeqCst), 0); + } + + #[test] + fn resolution_error_is_preserved_when_lock_release_also_fails() { + let result = finish_resolution::<()>( + Err(injected_storage_error()), + Err(MetastoreError::LockFileFailure { + message: "injected release failure".to_owned(), + }), + ); + assert!(is_injected_storage_error(&result.unwrap_err())); + + let release_error = finish_resolution::<()>( + Ok(()), + Err(MetastoreError::LockFileFailure { + message: "injected release failure".to_owned(), + }), + ) + .unwrap_err(); + assert!(matches!( + release_error, + MetastoreError::LockFileFailure { .. } + )); + } + + #[test] + fn default_resolution_is_identity() { + let resolution = Resolution::default(); + let schema = Schema::new(vec![field("a", DataType::Int64)]); + assert!(resolution.is_identity()); + assert_eq!(resolution.rename_schema(&schema).unwrap(), schema); + } + + #[test] + fn canonical_identity_targets_preserve_local_renames() { + let canonical = Schema::new(vec![ + field("a", DataType::Utf8), + field("a_int64", DataType::Int64), + ]); + let local = Schema::new(vec![field("a", DataType::Int64)]); + let (mut resolution, _) = resolve_local(canonical.clone(), &[local.clone()], &[]).unwrap(); + resolution.include_identity_schema(&canonical); + + assert_eq!( + resolution.target("a", &DataType::Int64).unwrap().name(), + "a_int64" + ); + assert_eq!(resolution.rename_schema(&canonical).unwrap(), canonical); + assert_eq!( + resolution.rename_schema(&local).unwrap().field(0).name(), + "a_int64" + ); + assert!(resolution.target("unclaimed", &DataType::Utf8).is_none()); + } + + #[test] + fn existing_type_is_retained_and_local_conflict_is_renamed() { + let canonical = Schema::new(vec![field("a", DataType::Utf8)]); + let local = [Schema::new(vec![field("a", DataType::Int64)])]; + let (resolution, updated) = resolve_local(canonical, &local, &[]).unwrap(); + assert_eq!( + resolution.target("a", &DataType::Int64).unwrap().name(), + "a_int64" + ); + assert_eq!( + updated.field_with_name("a").unwrap().data_type(), + &DataType::Utf8 + ); + } + + #[test] + fn suffix_resolution_repeats_past_occupied_names() { + let canonical = Schema::new(vec![ + field("a", DataType::Utf8), + field("a_int64", DataType::Utf8), + ]); + let local = [Schema::new(vec![field("a", DataType::Int64)])]; + let (resolution, _) = resolve_local(canonical, &local, &[]).unwrap(); + assert_eq!( + resolution.target("a", &DataType::Int64).unwrap().name(), + "a_int64_int64" + ); + } + + #[test] + fn nested_datatype_equality_is_conservative() { + let item_nullable = DataType::List(field("item", DataType::Int64)); + let item_non_nullable = + DataType::List(Arc::new(Field::new("item", DataType::Int64, false))); + let canonical = Schema::new(vec![field("nested", item_nullable)]); + let local = [Schema::new(vec![field( + "nested", + item_non_nullable.clone(), + )])]; + let (resolution, _) = resolve_local(canonical, &local, &[]).unwrap(); + assert_eq!( + resolution + .target("nested", &item_non_nullable) + .unwrap() + .name(), + "nested_list" + ); + } + + #[test] + fn pinned_fields_fail_on_type_conflict() { + let canonical = Schema::new(vec![field("p_timestamp", DataType::Utf8)]); + let local = [Schema::new(vec![field("p_timestamp", DataType::Int64)])]; + assert!(resolve_local(canonical, &local, &["p_timestamp"]).is_err()); + } + + #[test] + fn source_name_reservation_prevents_suffix_aliasing() { + let canonical = Schema::new(vec![field("a", DataType::Utf8)]); + let local = [Schema::new(vec![ + field("a", DataType::Int64), + field("a_int64", DataType::Int64), + ])]; + let (resolution, _) = resolve_local(canonical, &local, &[]).unwrap(); + assert_eq!( + resolution.target("a", &DataType::Int64).unwrap().name(), + "a_int64_int64" + ); + assert_eq!( + resolution + .target("a_int64", &DataType::Int64) + .unwrap() + .name(), + "a_int64" + ); + } + + #[test] + fn separate_schemas_can_reuse_compatible_canonical_suffix_name() { + let canonical = Schema::new(vec![ + field("x", DataType::Utf8), + field("x_int64", DataType::Int64), + ]); + let old_raw_schema = Schema::new(vec![field("x", DataType::Int64)]); + let canonical_shaped_schema = Schema::new(vec![ + field("x", DataType::Utf8), + field("x_int64", DataType::Int64), + ]); + let local = [old_raw_schema, canonical_shaped_schema]; + + let (resolution, updated) = resolve_local(canonical.clone(), &local, &[]).unwrap(); + assert_eq!( + resolution.target("x", &DataType::Int64).unwrap().name(), + "x_int64" + ); + assert_eq!(updated, canonical); + } + + #[test] + fn pinned_name_is_never_claimed_by_a_generated_suffix() { + let canonical = Schema::new(vec![ + field("foo", DataType::Int64), + field("foo_utf8", DataType::Utf8), + ]); + let local = [Schema::new(vec![field("foo", DataType::Utf8)])]; + let (resolution, updated) = resolve_local(canonical, &local, &["foo_utf8"]).unwrap(); + + assert_eq!( + resolution.target("foo", &DataType::Utf8).unwrap().name(), + "foo_utf8_utf8" + ); + assert_eq!( + updated.field_with_name("foo_utf8").unwrap().data_type(), + &DataType::Utf8 + ); + } + + #[test] + fn duplicate_target_columns_are_rejected_instead_of_dropped() { + let target = field("a_x", DataType::Utf8); + let resolution = Resolution::from_targets( + [ + (("a".to_owned(), DataType::Utf8), target.clone()), + (("a_x".to_owned(), DataType::Utf8), target), + ] + .into_iter() + .collect(), + ); + let schema = Schema::new(vec![ + field("a", DataType::Utf8), + field("a_x", DataType::Utf8), + ]); + assert!(resolution.rename_schema(&schema).is_err()); + + let batch = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(StringArray::from(vec![Some("first")])), + Arc::new(StringArray::from(vec![Some("second")])), + ], + ) + .unwrap(); + assert!(resolution.rename_batch(&batch).is_err()); + } + + #[test] + fn top_level_null_columns_drop_but_empty_batch_keeps_rows() { + let null_column = Arc::new(arrow_array::NullArray::new(4)) as ArrayRef; + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![field("nulls", DataType::Null)])), + vec![null_column], + ) + .unwrap(); + let resolution = Resolution::from_targets(HashMap::new()); + let out = resolution.rename_batch(&batch).unwrap(); + assert_eq!(out.num_rows(), 4); + assert_eq!(out.num_columns(), 0); + assert!( + resolution + .rename_schema(batch.schema().as_ref()) + .unwrap() + .fields() + .is_empty() + ); + } + + #[test] + fn renamed_batch_keeps_column_order_row_count_and_array_arcs() { + let resolution = Resolution::from_targets( + [ + (("b".to_owned(), DataType::Utf8), field("b", DataType::Utf8)), + ( + ("a".to_owned(), DataType::Int64), + field("a_int64", DataType::Int64), + ), + ] + .into_iter() + .collect(), + ); + let first: ArrayRef = Arc::new(StringArray::from(vec![Some("x"), Some("y")])); + let second: ArrayRef = Arc::new(Int64Array::from(vec![1, 2])); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + field("b", DataType::Utf8), + field("a", DataType::Int64), + ])), + vec![first.clone(), second.clone()], + ) + .unwrap(); + let out = resolution.rename_batch(&batch).unwrap(); + assert_eq!(out.num_rows(), 2); + assert_eq!(out.schema().field(0).name(), "b"); + assert_eq!(out.schema().field(1).name(), "a_int64"); + assert!(Arc::ptr_eq(&first, out.column(0))); + assert!(Arc::ptr_eq(&second, out.column(1))); + } + + #[test] + fn identity_resolution_retains_schema_metadata_and_original_fields() { + let original = Schema::new_with_metadata( + vec![Arc::new( + Field::new("a", DataType::Int64, false) + .with_metadata([("origin".to_owned(), "source".to_owned())].into()), + )], + [("schema".to_owned(), "metadata".to_owned())].into(), + ); + let renamed = Resolution::identity().rename_schema(&original).unwrap(); + assert_eq!(renamed, original); + assert_eq!(renamed.field(0).metadata(), original.field(0).metadata()); + assert_eq!(renamed.metadata(), original.metadata()); + } + + #[test] + fn canonical_fields_and_registry_are_normalized_and_sorted() { + let noncanonical_target = Arc::new( + Field::new("target", DataType::Utf8, false) + .with_metadata([("source".to_owned(), "legacy".to_owned())].into()), + ); + let resolution = Resolution::from_targets( + [(("source".to_owned(), DataType::Utf8), noncanonical_target)] + .into_iter() + .collect(), + ); + let target = resolution.target("source", &DataType::Utf8).unwrap(); + assert!(target.is_nullable()); + assert!(target.metadata().is_empty()); + + let schema = Schema::new(vec![ + Arc::new( + Field::new("z", DataType::Int64, false) + .with_metadata([("k".to_owned(), "v".to_owned())].into()), + ), + Arc::new(Field::new("a", DataType::Utf8, false)), + ]); + let normalized = normalize_registry(&schema).unwrap(); + assert_eq!(normalized.field(0).name(), "a"); + assert_eq!(normalized.field(1).name(), "z"); + assert!(normalized.fields().iter().all(|field| field.is_nullable())); + assert!( + normalized + .fields() + .iter() + .all(|field| field.metadata().is_empty()) + ); + assert!(normalized.metadata().is_empty()); + assert_eq!(canonical_field(schema.field(1)), field("a", DataType::Utf8)); + } + + #[test] + fn canonical_registry_rejects_duplicate_names_and_null_fields() { + let duplicates = Schema::new(vec![field("a", DataType::Utf8), field("a", DataType::Utf8)]); + assert!(normalize_registry(&duplicates).is_err()); + let nulls = Schema::new(vec![field("n", DataType::Null)]); + assert!(normalize_registry(&nulls).is_err()); + } + + #[test] + fn legacy_union_is_stable_and_prefers_utf8_for_conflicts() { + let first = Schema::new(vec![ + field("x", DataType::Int64), + field("b", DataType::Boolean), + ]); + let second = Schema::new(vec![field("x", DataType::Utf8)]); + let forward = legacy_seed(&[first.clone(), second.clone()]); + let reverse = legacy_seed(&[second, first]); + assert_eq!(forward, reverse); + assert_eq!( + forward.field_with_name("x").unwrap().data_type(), + &DataType::Utf8 + ); + assert_eq!(forward.field(0).name(), "b"); + } + + #[test] + fn resolution_is_deterministic_across_schema_order() { + let a = Schema::new(vec![field("value", DataType::Int64)]); + let b = Schema::new(vec![field("value", DataType::Utf8)]); + let (first, first_schema) = + resolve_local(Schema::empty(), &[a.clone(), b.clone()], &[]).unwrap(); + let (second, second_schema) = resolve_local(Schema::empty(), &[b, a], &[]).unwrap(); + assert_eq!(first_schema, second_schema); + assert_eq!( + first.target("value", &DataType::Int64).unwrap().name(), + second.target("value", &DataType::Int64).unwrap().name() + ); + assert_eq!( + first.target("value", &DataType::Utf8).unwrap().name(), + second.target("value", &DataType::Utf8).unwrap().name() + ); + } + + #[test] + fn timestamp_types_are_not_merged_by_suffix_equivalence() { + let s = DataType::Timestamp(TimeUnit::Second, None); + let ms = DataType::Timestamp(TimeUnit::Millisecond, None); + let canonical = Schema::new(vec![field("time", s)]); + let local = [Schema::new(vec![field("time", ms.clone())])]; + let (resolution, _) = resolve_local(canonical, &local, &[]).unwrap(); + assert_eq!( + resolution.target("time", &ms).unwrap().name(), + "time_timestamp_ms" + ); + } + + #[tokio::test] + async fn concurrent_absent_conflicts_retry_against_the_winner() { + let backend = InMemoryBackend::new(None, Vec::new(), 2); + let pinned = HashSet::new(); + let int_schema = Schema::new(vec![field("value", DataType::Int64)]); + let utf8_schema = Schema::new(vec![field("value", DataType::Utf8)]); + let int_schemas = [int_schema.clone()]; + let utf8_schemas = [utf8_schema.clone()]; + + let (int_resolution, utf8_resolution) = tokio::join!( + resolve_backend(&backend, "concurrent", &int_schemas, &pinned), + resolve_backend(&backend, "concurrent", &utf8_schemas, &pinned), + ); + let int_resolution = int_resolution.unwrap(); + let utf8_resolution = utf8_resolution.unwrap(); + let int_target = int_resolution.target("value", &DataType::Int64).unwrap(); + let utf8_target = utf8_resolution.target("value", &DataType::Utf8).unwrap(); + assert_ne!(int_target.name(), utf8_target.name()); + assert!(int_target.name() == "value" || utf8_target.name() == "value"); + + let renamed_int = int_resolution.rename_schema(&int_schema).unwrap(); + let renamed_utf8 = utf8_resolution.rename_schema(&utf8_schema).unwrap(); + let canonical = backend.canonical_schema().unwrap(); + assert_schemas_merge_into(&canonical, [renamed_int, renamed_utf8]); + assert_eq!(canonical.fields().len(), 2); + assert!(backend.cas_attempt_count() >= 3); + assert!(backend.global_read_count() >= 3); + } + + #[tokio::test] + async fn concurrent_existing_version_appends_independent_fields_without_loss() { + let initial = Schema::new(vec![field("base", DataType::Boolean)]); + let backend = InMemoryBackend::new(Some(initial), Vec::new(), 2); + let pinned = HashSet::new(); + let left_schema = Schema::new(vec![field("left", DataType::Int64)]); + let right_schema = Schema::new(vec![field("right", DataType::Utf8)]); + let left_schemas = [left_schema.clone()]; + let right_schemas = [right_schema.clone()]; + + let (left_resolution, right_resolution) = tokio::join!( + resolve_backend(&backend, "append", &left_schemas, &pinned), + resolve_backend(&backend, "append", &right_schemas, &pinned), + ); + let left = left_resolution + .unwrap() + .rename_schema(&left_schema) + .unwrap(); + let right = right_resolution + .unwrap() + .rename_schema(&right_schema) + .unwrap(); + let canonical = backend.canonical_schema().unwrap(); + + assert!(canonical.field_with_name("base").is_ok()); + assert!(canonical.field_with_name("left").is_ok()); + assert!(canonical.field_with_name("right").is_ok()); + assert_schemas_merge_into(&canonical, [left, right]); + assert_eq!(backend.legacy_read_count(), 0); + assert!(backend.cas_attempt_count() >= 3); + } + + #[tokio::test] + async fn legacy_schemas_are_loaded_only_when_global_schema_is_absent() { + let legacy = Schema::new(vec![field("legacy", DataType::Utf8)]); + let backend = InMemoryBackend::new(None, vec![legacy], 0); + let local = [Schema::new(vec![field("new", DataType::Int64)])]; + let pinned = HashSet::new(); + + resolve_backend(&backend, "bootstrap", &local, &pinned) + .await + .unwrap(); + assert_eq!(backend.legacy_read_count(), 1); + let second = [Schema::new(vec![field("legacy", DataType::Utf8)])]; + resolve_backend(&backend, "bootstrap", &second, &pinned) + .await + .unwrap(); + assert_eq!(backend.legacy_read_count(), 1); + assert!( + backend + .canonical_schema() + .unwrap() + .field_with_name("legacy") + .is_ok() + ); + } + + #[tokio::test] + async fn repeated_cas_contention_returns_bounded_exhaustion_error() { + let initial = Schema::new(vec![field("base", DataType::Boolean)]); + let backend = InMemoryBackend::new(Some(initial), Vec::new(), 0).with_contention(); + let local = [Schema::new(vec![field("new", DataType::Int64)])]; + let error = resolve_backend(&backend, "contended", &local, &HashSet::new()) + .await + .unwrap_err(); + + assert!(matches!( + error, + MetastoreError::CASConvergenceFailure { message } + if message.contains("did not converge") + )); + assert_eq!(backend.cas_attempt_count(), MAX_CAS_ATTEMPTS); + assert_eq!(backend.global_read_count(), MAX_CAS_ATTEMPTS); + } + + #[tokio::test] + async fn backend_storage_errors_propagate_without_becoming_contention() { + let local = [Schema::new(vec![field("new", DataType::Int64)])]; + let read_failure = InMemoryBackend::new(None, Vec::new(), 0).with_read_error(); + let read_error = resolve_backend(&read_failure, "read-error", &local, &HashSet::new()) + .await + .unwrap_err(); + assert!(is_injected_storage_error(&read_error)); + assert_eq!(read_failure.cas_attempt_count(), 0); + assert_eq!(read_failure.global_read_count(), 1); + + let write_failure = InMemoryBackend::new(None, Vec::new(), 0).with_write_error(); + let write_error = resolve_backend(&write_failure, "write-error", &local, &HashSet::new()) + .await + .unwrap_err(); + assert!(is_injected_storage_error(&write_error)); + assert_eq!(write_failure.cas_attempt_count(), 1); + assert_eq!(write_failure.global_read_count(), 1); + } +} diff --git a/src/storage/azure_blob.rs b/src/storage/azure_blob.rs index 6c0bc776b..b9ef17423 100644 --- a/src/storage/azure_blob.rs +++ b/src/storage/azure_blob.rs @@ -686,6 +686,29 @@ impl ObjectStorage for BlobStore { Ok(self._get_object(path, tenant_id).await?) } + async fn get_object_versioned( + &self, + path: &RelativePath, + tenant_id: &Option, + ) -> Result, ObjectStorageError> { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let date = Utc::now().date_naive().to_string(); + let result = + crate::storage::get_object_versioned(&*self.client, &to_object_store_path(path)).await; + increment_object_store_calls_by_date("GET", &date, tenant); + let result = result?; + if let Some((body, _)) = &result { + increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant); + increment_bytes_scanned_in_object_store_calls_by_date( + "GET", + body.len() as u64, + &date, + tenant, + ); + } + Ok(result) + } + async fn get_object_ranged( &self, path: &RelativePath, @@ -819,6 +842,30 @@ impl ObjectStorage for BlobStore { Ok(()) } + async fn put_object_if( + &self, + path: &RelativePath, + data: Bytes, + expected: Option, + tenant_id: &Option, + ) -> Result { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let date = Utc::now().date_naive().to_string(); + let result = crate::storage::put_schema_if( + &*self.client, + &to_object_store_path(path), + data, + expected, + ) + .await; + increment_object_store_calls_by_date("PUT", &date, tenant); + let written = result?; + if written { + increment_files_scanned_in_object_store_calls_by_date("PUT", 1, &date, tenant); + } + Ok(written) + } + async fn delete_prefix( &self, path: &RelativePath, diff --git a/src/storage/gcs.rs b/src/storage/gcs.rs index 8ac87f763..17c78dc3d 100644 --- a/src/storage/gcs.rs +++ b/src/storage/gcs.rs @@ -668,6 +668,29 @@ impl ObjectStorage for Gcs { Ok(self._get_object(path, tenant_id).await?) } + async fn get_object_versioned( + &self, + path: &RelativePath, + tenant_id: &Option, + ) -> Result, ObjectStorageError> { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let date = Utc::now().date_naive().to_string(); + let result = + crate::storage::get_object_versioned(&*self.client, &to_object_store_path(path)).await; + increment_object_store_calls_by_date("GET", &date, tenant); + let result = result?; + if let Some((body, _)) = &result { + increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant); + increment_bytes_scanned_in_object_store_calls_by_date( + "GET", + body.len() as u64, + &date, + tenant, + ); + } + Ok(result) + } + async fn get_object_ranged( &self, path: &RelativePath, @@ -801,6 +824,30 @@ impl ObjectStorage for Gcs { Ok(()) } + async fn put_object_if( + &self, + path: &RelativePath, + data: Bytes, + expected: Option, + tenant_id: &Option, + ) -> Result { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let date = Utc::now().date_naive().to_string(); + let result = crate::storage::put_schema_if( + &*self.client, + &to_object_store_path(path), + data, + expected, + ) + .await; + increment_object_store_calls_by_date("PUT", &date, tenant); + let written = result?; + if written { + increment_files_scanned_in_object_store_calls_by_date("PUT", 1, &date, tenant); + } + Ok(written) + } + async fn delete_prefix( &self, path: &RelativePath, diff --git a/src/storage/localfs.rs b/src/storage/localfs.rs index d66714023..7d1749a19 100644 --- a/src/storage/localfs.rs +++ b/src/storage/localfs.rs @@ -916,3 +916,119 @@ impl From for ObjectStorageError { ObjectStorageError::UnhandledError(Box::new(e)) } } + +#[cfg(test)] +mod tests { + use arrow_schema::{DataType, Field, Schema}; + + use crate::{ + metastore::{ + metastore_traits::Metastore, metastores::object_store_metastore::ObjectStoreMetastore, + }, + parseable::GlobalSchema, + storage::{ + ObjectStorage, SCHEMA_FILE_NAME, STREAM_ROOT_DIRECTORY, + object_storage::{global_schema_path, to_bytes}, + }, + }; + use relative_path::RelativePathBuf; + + use super::LocalFS; + use std::sync::Arc; + + #[tokio::test] + async fn global_schema_lock_deletion_is_idempotent_and_propagates_failures() { + let temp_dir = tempfile::tempdir().unwrap(); + let storage = Arc::new(LocalFS::new(temp_dir.path().to_path_buf())); + let metastore = ObjectStoreMetastore { storage }; + let lock = RelativePathBuf::from("events/.stream/node-1.lock"); + + metastore + .put_global_schema_lock(lock.clone(), &None) + .await + .unwrap(); + metastore + .delete_global_schema_lock(lock.clone(), &None) + .await + .unwrap(); + metastore + .delete_global_schema_lock(lock, &None) + .await + .expect("deleting an absent lock is idempotent"); + + let blocked = RelativePathBuf::from("events/.stream/blocked.lock"); + let directory = temp_dir.path().join(blocked.as_str()); + std::fs::create_dir_all(&directory).unwrap(); + metastore + .delete_global_schema_lock(blocked, &None) + .await + .expect_err("a failed lock deletion must not be swallowed"); + assert!(directory.is_dir()); + } + + #[tokio::test] + async fn localfs_skips_global_schema_and_discovers_legacy_schema() { + let temp_dir = tempfile::tempdir().unwrap(); + let storage = Arc::new(LocalFS::new(temp_dir.path().to_path_buf())); + let metastore = ObjectStoreMetastore { + storage: storage.clone(), + }; + let tenant_id = None; + let stream_name = "events"; + + assert!( + metastore + .get_global_schema(stream_name, &tenant_id) + .await + .unwrap() + .is_none() + ); + + let schema = Schema::new(vec![Field::new("value", DataType::Utf8, true)]); + let legacy_schema_path = + RelativePathBuf::from_iter([stream_name, STREAM_ROOT_DIRECTORY, SCHEMA_FILE_NAME]); + let serialized_schema = to_bytes(&schema); + storage + .put_object(&legacy_schema_path, serialized_schema.clone(), &tenant_id) + .await + .unwrap(); + + let stored_schema = storage + .get_object(&legacy_schema_path, &tenant_id) + .await + .unwrap(); + assert_eq!(stored_schema, serialized_schema); + assert_eq!( + serde_json::from_slice::(&stored_schema).unwrap(), + schema + ); + assert_eq!( + metastore + .get_all_schemas(stream_name, &tenant_id) + .await + .unwrap(), + vec![schema.clone()] + ); + + let error = metastore + .put_global_schema( + GlobalSchema::from_schema(schema), + stream_name, + &tenant_id, + None, + ) + .await + .expect_err("LocalFS must reject unsupported conditional global-schema writes"); + assert!( + error + .to_string() + .contains("conditional object writes are not supported") + ); + assert!( + storage + .get_object(&global_schema_path(stream_name, &tenant_id), &tenant_id) + .await + .is_err() + ); + } +} diff --git a/src/storage/mod.rs b/src/storage/mod.rs index c376a487f..a42cb0311 100644 --- a/src/storage/mod.rs +++ b/src/storage/mod.rs @@ -16,7 +16,8 @@ * */ -use object_store::path::Path; +use bytes::Bytes; +use object_store::{ObjectStoreExt, path::Path}; use relative_path::RelativePath; use serde::{Deserialize, Serialize}; use tokio::task::JoinError; @@ -29,7 +30,7 @@ use crate::{ metadata::SchemaVersion, metastore::{MetastoreErrorDetail, metastore_traits::MetastoreObject}, option::StandaloneWithDistributed, - parseable::StreamNotFound, + parseable::{GlobalSchema, StreamNotFound}, stats::FullStats, utils::json::{deserialize_string_as_true, serialize_bool_as_true}, }; @@ -52,6 +53,10 @@ pub mod retention; mod s3; pub mod store_metadata; +pub trait ComparableBytes: Sized + serde::Serialize { + fn get_hash(&self) -> u64; +} + /// Size of the first request issued by [`get_object_ranged`], which doubles as /// the size probe. Anything at or below this is fetched as one plain GET, with /// no generation pinning and no retry path. Sized so that a compressed manifest @@ -73,6 +78,95 @@ pub(crate) const RANGED_GET_MAX_ATTEMPTS: u32 = 3; /// round trip latency, so overlapping them is close to a pure win. pub(crate) const GET_OBJECTS_CONCURRENCY: usize = 32; +/// Fetch an object's body and conditional-write token from the same GET request. +pub(crate) async fn get_object_versioned( + client: &S, + path: &Path, +) -> object_store::Result> { + let result = match client + .get_opts(path, object_store::GetOptions::default()) + .await + { + Ok(result) => result, + Err(object_store::Error::NotFound { .. }) => return Ok(None), + Err(error) => return Err(error), + }; + + // Use the version returned by this same GET; a separate HEAD could describe + // a different generation from the bytes below. + let version = object_store::UpdateVersion { + e_tag: result.meta.e_tag.clone(), + version: result.meta.version.clone(), + }; + if version.e_tag.is_none() && version.version.is_none() { + return Err(object_store::Error::NotSupported { + source: Box::new(std::io::Error::new( + std::io::ErrorKind::Unsupported, + "object store returned no validator for conditional updates", + )), + }); + } + + Ok(Some((result.bytes().await?, version))) +} + +/// Atomically create or update `path`; return `false` only for a conditional +/// precondition conflict. Every other storage failure is preserved as an error. +/// Specific function for comparing global schema hashes +/// (the read path doesn't know what to Deserialize the Bytes to for comparison) +pub(crate) async fn put_schema_if( + client: &S, + path: &Path, + data: Bytes, + expected: Option, +) -> object_store::Result { + use object_store::{PutMode, PutOptions}; + + let res = match expected { + None => { + // create + client + .put_opts( + path, + data.into(), + PutOptions { + mode: PutMode::Create, + ..Default::default() + }, + ) + .await + } + Some(version) => { + // Update if version matches + // get the object first + match client.get(path).await { + Ok(obj) => { + let obj = serde_json::from_slice::(&obj.bytes().await?).map_err( + |e| object_store::Error::Generic { + store: "object-store hash", + source: Box::new(e), + }, + )?; + if obj.hash.ne(&version) { + return Ok(false); + } + } + Err(object_store::Error::NotFound { .. }) => return Ok(false), + Err(error) => return Err(error), + } + client.put(path, data.into()).await + } + }; + + match res { + Ok(_) => Ok(true), + Err( + object_store::Error::AlreadyExists { .. } | object_store::Error::Precondition { .. }, + ) => Ok(false), + Err(error) => Err(error), + } +} + /// Fetch an object as a series of bounded byte ranges rather than one stream. /// /// A single multi-hundred-MB response body is regularly killed mid transfer by @@ -307,6 +401,7 @@ pub const PARSEABLE_METADATA_FILE_NAME: &str = ".parseable.json"; pub const STREAM_ROOT_DIRECTORY: &str = ".stream"; pub const PARSEABLE_ROOT_DIRECTORY: &str = ".parseable"; pub const SCHEMA_FILE_NAME: &str = ".schema"; +pub const GLOBAL_SCHEMA_FILE_NAME: &str = ".global.schema"; pub const ALERTS_ROOT_DIRECTORY: &str = ".alerts"; pub const SETTINGS_ROOT_DIRECTORY: &str = ".settings"; pub const TARGETS_ROOT_DIRECTORY: &str = ".targets"; @@ -676,3 +771,134 @@ mod tests { ); } } + +#[cfg(test)] +mod global_schema_storage_tests { + use arrow_schema::{DataType, Field, Schema}; + use bytes::Bytes; + use object_store::{ObjectStoreExt, memory::InMemory, path::Path}; + + use crate::parseable::GlobalSchema; + + use super::{get_object_versioned, put_schema_if}; + + #[tokio::test] + async fn versioned_get_reports_absence_and_returns_body_with_validator() { + let store = InMemory::new(); + let path = Path::from("tenant/stream/.stream/.global.schema"); + + assert!( + get_object_versioned(&store, &path) + .await + .expect("missing objects are not errors") + .is_none() + ); + + assert!( + put_schema_if(&store, &path, Bytes::from_static(b"schema"), None) + .await + .expect("create-only write should succeed") + ); + let (body, version) = get_object_versioned(&store, &path) + .await + .expect("read should succeed") + .expect("object was just created"); + assert_eq!(body, Bytes::from_static(b"schema")); + assert!(version.e_tag.is_some() || version.version.is_some()); + } + + #[tokio::test] + async fn concurrent_create_only_writes_have_exactly_one_winner() { + let store = InMemory::new(); + let path = Path::from("tenant/stream/.stream/.global.schema"); + let (first, second) = tokio::join!( + put_schema_if(&store, &path, Bytes::from_static(b"first"), None), + put_schema_if(&store, &path, Bytes::from_static(b"second"), None), + ); + + let outcomes = [ + first.expect("first write must resolve"), + second.expect("second write must resolve"), + ]; + assert_eq!(outcomes.iter().filter(|written| **written).count(), 1); + assert_eq!(outcomes.iter().filter(|written| !**written).count(), 1); + } + + #[tokio::test] + async fn stale_update_is_rejected_without_overwriting_newer_schema() { + let store = InMemory::new(); + let path = Path::from("tenant/stream/.stream/.global.schema"); + let initial = + GlobalSchema::from_schema(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + let newer = GlobalSchema::from_schema(Schema::new(vec![ + Field::new("value", DataType::Utf8, true), + Field::new( + "timestamp", + DataType::Timestamp(arrow_schema::TimeUnit::Millisecond, None), + true, + ), + ])); + + assert!( + put_schema_if( + &store, + &path, + Bytes::from(serde_json::to_vec(&initial).unwrap()), + None, + ) + .await + .unwrap() + ); + let stored_initial: GlobalSchema = + serde_json::from_slice(&store.get(&path).await.unwrap().bytes().await.unwrap()) + .unwrap(); + assert_eq!(stored_initial.hash, initial.hash); + + assert!( + put_schema_if( + &store, + &path, + Bytes::from(serde_json::to_vec(&newer).unwrap()), + Some(stored_initial.hash), + ) + .await + .unwrap() + ); + assert!( + !put_schema_if( + &store, + &path, + Bytes::from(serde_json::to_vec(&initial).unwrap()), + Some(initial.hash), + ) + .await + .unwrap() + ); + + let actual: GlobalSchema = + serde_json::from_slice(&store.get(&path).await.unwrap().bytes().await.unwrap()) + .unwrap(); + assert_eq!(actual.hash, newer.hash); + assert_eq!(actual.schema, newer.schema); + } + + #[tokio::test] + async fn update_with_expected_hash_for_missing_schema_is_a_conflict() { + let store = InMemory::new(); + let path = Path::from("tenant/stream/.stream/.global.schema"); + let schema = + GlobalSchema::from_schema(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + + assert!( + !put_schema_if( + &store, + &path, + Bytes::from(serde_json::to_vec(&schema).unwrap()), + Some(schema.hash), + ) + .await + .unwrap() + ); + assert!(store.get(&path).await.is_err()); + } +} diff --git a/src/storage/object_storage.rs b/src/storage/object_storage.rs index 368326a67..acdf5283a 100644 --- a/src/storage/object_storage.rs +++ b/src/storage/object_storage.rs @@ -31,11 +31,11 @@ use crate::option::Mode; use crate::parseable::DEFAULT_TENANT; use crate::parseable::{LogStream, PARSEABLE, Stream}; use crate::stats::FullStats; -use crate::storage::SETTINGS_ROOT_DIRECTORY; use crate::storage::TARGETS_ROOT_DIRECTORY; use crate::storage::field_stats::DATASET_STATS_STREAM_NAME; use crate::storage::field_stats::calculate_field_stats; use crate::storage::field_stats::extract_datetime_from_parquet_path_regex; +use crate::storage::{GLOBAL_SCHEMA_FILE_NAME, SETTINGS_ROOT_DIRECTORY}; use crate::sync::{ACTIVE_OBJECT_STORE_SYNC_FILES, FLUSH_AND_CONVERT_RUNTIME}; use arrow_schema::Schema; use async_trait::async_trait; @@ -353,6 +353,34 @@ pub trait ObjectStorage: Debug + Send + Sync + 'static { path: &RelativePath, tenant_id: &Option, ) -> Result; + /// Fetch bytes and the exact version token returned by the same GET. + /// + /// Implementations must not emulate this with a separate HEAD request: the + /// metadata and body need to describe one object generation. + async fn get_object_versioned( + &self, + _path: &RelativePath, + _tenant_id: &Option, + ) -> Result, ObjectStorageError> { + Err(ObjectStorageError::Custom( + "versioned object reads are not supported by this storage backend".into(), + )) + } + /// Atomically create an object when `expected` is `None`, or replace it only + /// if it still has the supplied version. `false` means a precondition + /// conflict; unsupported backends must return an error rather than falling + /// back to an unconditional write. + async fn put_object_if( + &self, + _path: &RelativePath, + _data: Bytes, + _expected: Option, + _tenant_id: &Option, + ) -> Result { + Err(ObjectStorageError::Custom( + "conditional object writes are not supported by this storage backend".into(), + )) + } /// Fetch an object as concurrent bounded byte ranges instead of one stream. /// /// Intended for the few objects that are large enough (hundreds of MB) that @@ -1430,6 +1458,24 @@ pub async fn commit_schema_to_storage( schema: Schema, tenant_id: &Option, ) -> Result<(), ObjectStorageError> { + if let Some(global_schema) = PARSEABLE + .metastore + .get_global_schema(stream_name, tenant_id) + .await + .map_err(|e| ObjectStorageError::MetastoreError(Box::new(e.to_detail())))? + { + // The global schema already contains canonical names and resolved types. + // Do not merge it with this ingestor's raw schema: equal field names may + // intentionally have different raw types before registry resolution. + let canonical_schema = global_schema.schema; + return PARSEABLE + .metastore + .put_schema(canonical_schema, stream_name, tenant_id) + .await + .map_err(|e| ObjectStorageError::MetastoreError(Box::new(e.to_detail()))); + } + + // Preserve the legacy strict-merge behavior while no global registry exists. let stream_schema = PARSEABLE .metastore .get_schema(stream_name, tenant_id) @@ -1471,6 +1517,31 @@ pub fn schema_path(stream_name: &str, tenant_id: &Option) -> RelativePat } } +/// Shared directory for global-schema data and lock files. +/// +/// `None` and `Some(DEFAULT_TENANT)` both represent the legacy global tenant, +/// whose objects intentionally have no tenant prefix. +fn global_schema_directory_path(stream_name: &str, tenant_id: &Option) -> RelativePathBuf { + match tenant_id + .as_deref() + .filter(|tenant| *tenant != DEFAULT_TENANT) + { + Some(tenant) => RelativePathBuf::from_iter([tenant, stream_name, STREAM_ROOT_DIRECTORY]), + None => RelativePathBuf::from_iter([stream_name, STREAM_ROOT_DIRECTORY]), + } +} + +pub fn global_schema_path(stream_name: &str, tenant_id: &Option) -> RelativePathBuf { + global_schema_directory_path(stream_name, tenant_id).join(GLOBAL_SCHEMA_FILE_NAME) +} + +pub fn global_schema_lock_file_path( + stream_name: &str, + tenant_id: &Option, +) -> RelativePathBuf { + global_schema_directory_path(stream_name, tenant_id) +} + #[inline(always)] pub fn stream_json_path(stream_name: &str, tenant_id: &Option) -> RelativePathBuf { let tenant = tenant_id.as_deref().unwrap_or(""); diff --git a/src/storage/s3.rs b/src/storage/s3.rs index a92923172..00accfab7 100644 --- a/src/storage/s3.rs +++ b/src/storage/s3.rs @@ -41,7 +41,7 @@ use futures::{StreamExt, TryStreamExt, stream::FuturesUnordered}; use object_store::{ BackoffConfig, ClientOptions, ListResult, ObjectMeta, ObjectStore, ObjectStoreExt, PutPayload, RetryConfig, - aws::{AmazonS3, AmazonS3Builder, AmazonS3ConfigKey, Checksum}, + aws::{AmazonS3, AmazonS3Builder, AmazonS3ConfigKey, Checksum, S3ConditionalPut}, buffered::BufReader, limit::LimitStore, path::Path as StorePath, @@ -277,6 +277,9 @@ impl S3Config { .with_region(&self.region) .with_endpoint(&self.endpoint_url) .with_bucket_name(&self.bucket_name) + // Global schema writes rely on conditional PUT, so enable S3's + // If-Match / If-None-Match support explicitly. + .with_conditional_put(S3ConditionalPut::ETagMatch) .with_virtual_hosted_style_request(!self.use_path_style) .with_allow_http(true) .with_retry(retry_config); @@ -907,6 +910,29 @@ impl ObjectStorage for S3 { self._get_object(path, tenant_id).await } + async fn get_object_versioned( + &self, + path: &RelativePath, + tenant_id: &Option, + ) -> Result, ObjectStorageError> { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let date = Utc::now().date_naive().to_string(); + let result = + crate::storage::get_object_versioned(&*self.client, &to_object_store_path(path)).await; + increment_object_store_calls_by_date("GET", &date, tenant); + let result = result?; + if let Some((body, _)) = &result { + increment_files_scanned_in_object_store_calls_by_date("GET", 1, &date, tenant); + increment_bytes_scanned_in_object_store_calls_by_date( + "GET", + body.len() as u64, + &date, + tenant, + ); + } + Ok(result) + } + async fn get_object_ranged( &self, path: &RelativePath, @@ -1050,6 +1076,30 @@ impl ObjectStorage for S3 { Ok(()) } + async fn put_object_if( + &self, + path: &RelativePath, + data: Bytes, + expected: Option, + tenant_id: &Option, + ) -> Result { + let tenant = tenant_id.as_deref().unwrap_or(DEFAULT_TENANT); + let date = Utc::now().date_naive().to_string(); + let result = crate::storage::put_schema_if( + &*self.client, + &to_object_store_path(path), + data, + expected, + ) + .await; + increment_object_store_calls_by_date("PUT", &date, tenant); + let written = result?; + if written { + increment_files_scanned_in_object_store_calls_by_date("PUT", 1, &date, tenant); + } + Ok(written) + } + async fn delete_prefix( &self, path: &RelativePath, diff --git a/src/utils/arrow/batch_adapter.rs b/src/utils/arrow/batch_adapter.rs index 6780ed2b1..18998ece3 100644 --- a/src/utils/arrow/batch_adapter.rs +++ b/src/utils/arrow/batch_adapter.rs @@ -16,13 +16,15 @@ * */ -use datafusion::arrow::array::ArrayRef; -use datafusion::arrow::array::new_null_array; -use datafusion::arrow::datatypes::Schema; -use datafusion::arrow::record_batch::RecordBatch; - use std::sync::Arc; +use arrow::{ + array::{ArrayRef, RecordBatch, RecordBatchOptions, new_null_array}, + compute::cast, + datatypes::Schema, +}; +use arrow_schema::ArrowError; + // This function takes a new event's record batch and the // current schema of the log stream. It returns a new record // with nulls added to the fields that don't exist @@ -31,17 +33,137 @@ use std::sync::Arc; // This is necessary because all the record batches in a log // stream need to have all the fields. pub fn adapt_batch(table_schema: Arc, batch: &RecordBatch) -> RecordBatch { - let batch_schema = &*batch.schema(); - let batch_cols = batch.columns().to_vec(); + let batch_schema = batch.schema(); + let mut cols = Vec::with_capacity(table_schema.fields().len()); + for table_field in table_schema.fields() { + if let Some((batch_idx, _)) = batch_schema.column_with_name(table_field.name()) { + cols.push(Arc::clone(batch.column(batch_idx))); + } else { + cols.push(new_null_array(table_field.data_type(), batch.num_rows())); + } + } + RecordBatch::try_new(table_schema, cols).unwrap() +} + +/// Adapts a batch to `table_schema`, preserving arrays where their types already match. +/// +/// Missing fields are represented by null arrays. Fields with compatible but differing +/// types are cast to the merged schema type before being returned. +pub fn try_adapt_batch( + table_schema: Arc, + batch: &RecordBatch, +) -> Result { + let batch_schema = batch.schema(); + let batch_cols = batch.columns(); let mut cols: Vec = Vec::with_capacity(table_schema.fields().len()); for table_field in table_schema.fields() { - if let Some((batch_idx, _)) = batch_schema.column_with_name(table_field.name().as_str()) { - cols.push(Arc::clone(&batch_cols[batch_idx])); + if let Some((batch_idx, batch_field)) = batch_schema.column_with_name(table_field.name()) { + let column = &batch_cols[batch_idx]; + if batch_field.data_type() == table_field.data_type() { + cols.push(Arc::clone(column)); + } else { + // Only normalize types compatible with the local schema merge; + // never cast away a concrete field-type conflict. + let mut merged_field = table_field.as_ref().clone(); + merged_field.try_merge(batch_field)?; + if merged_field.data_type() != table_field.data_type() { + return Err(ArrowError::SchemaError(format!( + "field '{}' is not covered by the merged local schema", + table_field.name() + ))); + } + cols.push(cast(column.as_ref(), table_field.data_type())?); + } } else { - cols.push(new_null_array(table_field.data_type(), batch.num_rows())) + cols.push(new_null_array(table_field.data_type(), batch.num_rows())); } } - RecordBatch::try_new(table_schema, cols).unwrap() + RecordBatch::try_new_with_options( + table_schema, + cols, + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + ) +} + +#[cfg(test)] +mod tests { + use super::try_adapt_batch; + use arrow::{ + array::{Array, Int32Array, NullArray, StructArray}, + datatypes::{DataType, Field, Schema}, + }; + use std::sync::Arc; + + #[test] + fn preserves_row_count_when_adapting_to_an_empty_schema() { + let batch = arrow::array::RecordBatch::try_from_iter([( + "value", + Arc::new(Int32Array::from(vec![1, 2])) as arrow::array::ArrayRef, + )]) + .unwrap(); + + let adapted = try_adapt_batch(Arc::new(Schema::empty()), &batch).unwrap(); + + assert_eq!(adapted.num_columns(), 0); + assert_eq!(adapted.num_rows(), 2); + } + + #[test] + fn rejects_concrete_conflicts_instead_of_casting_them() { + let batch = arrow::array::RecordBatch::try_from_iter([( + "value", + Arc::new(Int32Array::from(vec![1, 2])) as arrow::array::ArrayRef, + )]) + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("value", DataType::Utf8, true)])); + assert!(try_adapt_batch(schema, &batch).is_err()); + } + + #[test] + fn exact_datatypes_retain_the_original_array() { + let values: arrow::array::ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let batch = arrow::array::RecordBatch::try_from_iter([("value", values.clone())]).unwrap(); + let adapted = try_adapt_batch(batch.schema(), &batch).unwrap(); + assert!(Arc::ptr_eq(adapted.column(0), &values)); + } + + #[test] + fn casts_nested_null_struct_fields_to_the_merged_type() { + let null_field = Arc::new(Field::new("value", DataType::Null, true)); + let source_type = DataType::Struct(vec![null_field.clone()].into()); + let source = StructArray::try_new( + vec![null_field].into(), + vec![Arc::new(NullArray::new(2))], + None, + ) + .unwrap(); + let batch = arrow::array::RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("payload", source_type, true)])), + vec![Arc::new(source)], + ) + .unwrap(); + + let value_field = Arc::new(Field::new("value", DataType::Int32, true)); + let target_type = DataType::Struct(vec![value_field].into()); + let adapted = try_adapt_batch( + Arc::new(Schema::new(vec![Field::new("payload", target_type, true)])), + &batch, + ) + .unwrap(); + + let payload = adapted + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let values = payload + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(values.len(), 2); + assert_eq!(values.null_count(), 2); + } }