diff --git a/hyperbytedb/src/adapters/chdb/native_adapter.rs b/hyperbytedb/src/adapters/chdb/native_adapter.rs index b8e3e19..fdbdb80 100644 --- a/hyperbytedb/src/adapters/chdb/native_adapter.rs +++ b/hyperbytedb/src/adapters/chdb/native_adapter.rs @@ -1437,6 +1437,7 @@ impl ChdbNativeAdapter { ) -> Result<(), HyperbytedbError> { let pool = self.session.pool()?; let series_table = ensured.series_table.clone(); + let series_batch = &pad_series_batch_to_ensured(series_batch, ensured)?; let batch = series_batch.clone(); let insert_options = self.insert_options.clone(); tokio::task::spawn_blocking(move || { @@ -2063,19 +2064,76 @@ fn pad_record_batch_to_ensured( columns.push(batch.column(idx).clone()); } + // Each field column in the table's type, which decides: a prepared + // batch may hold a narrower one (built before the field widened, or + // merged from batches that disagreed), and a column of another type + // failed the whole batch, so the flush retried it forever. Values that + // cannot be cast become null, as at insert time. + let safe = arrow::compute::CastOptions { + safe: true, + ..Default::default() + }; for (_, phys, disc) in &ensured.field_phys { - fields.push(Field::new(phys, field_arrow_type(*disc), true)); - columns.push(if let Ok(idx) = schema.index_of(phys) { - batch.column(idx).clone() - } else { - new_null_array(&field_arrow_type(*disc), n) - }); + let ty = field_arrow_type(*disc); + let column = match schema.index_of(phys) { + Ok(idx) if batch.column(idx).data_type() == &ty => batch.column(idx).clone(), + Ok(idx) => { + arrow::compute::cast_with_options(batch.column(idx), &ty, &safe).map_err(|e| { + HyperbytedbError::Internal(format!("cast prepared column {phys}: {e}").into()) + })? + } + Err(_) => new_null_array(&ty, n), + }; + fields.push(Field::new(phys, ty, true)); + columns.push(column); } RecordBatch::try_new(Arc::new(Schema::new(fields)), columns) .map_err(|e| HyperbytedbError::Internal(format!("pad prepared RecordBatch: {e}").into())) } +/// A prepared series batch with every tag column of `ensured`, in its order +/// and types, then any other tag columns the batch has: a tag the batch +/// lacks is "", as for a series without it. The +/// batch may have been built before the table gained a tag, or merged from +/// series with other tags; the same is done for fact batches in +/// [`pad_record_batch_to_ensured`]. +fn pad_series_batch_to_ensured( + batch: &RecordBatch, + ensured: &EnsuredTable, +) -> Result { + let n = batch.num_rows(); + let schema = batch.schema(); + let err = |e: arrow::error::ArrowError| { + HyperbytedbError::Internal(format!("pad prepared series batch: {e}").into()) + }; + let sid = schema.index_of("series_id").map_err(err)?; + let mut fields = vec![schema.field(sid).clone()]; + let mut columns: Vec = vec![batch.column(sid).clone()]; + for (_, phys, kind) in &ensured.tag_phys { + let ty = tag_arrow_type(*kind); + let column = match schema.index_of(phys) { + Ok(i) if batch.column(i).data_type() == &ty => batch.column(i).clone(), + Ok(i) => arrow::compute::cast(batch.column(i), &ty).map_err(err)?, + Err(_) => { + let empty: ArrayRef = Arc::new(StringArray::from(vec![""; n])); + arrow::compute::cast(&empty, &ty).map_err(err)? + } + }; + fields.push(Field::new(phys, ty, false)); + columns.push(column); + } + // Tags `ensured` does not list are kept: on the prepared path it is + // planned without points, from metadata that may not have them yet. + for (i, field) in schema.fields().iter().enumerate() { + if i != sid && !fields.iter().any(|f| f.name() == field.name()) { + fields.push(field.as_ref().clone()); + columns.push(batch.column(i).clone()); + } + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).map_err(err) +} + /// Arrow logical type for a tag column, aligned with [`ColumnKind::ch_column_type`]. fn tag_arrow_type(kind: ColumnKind) -> DataType { match kind { @@ -3058,6 +3116,104 @@ mod tests { assert!(sql.contains(&format!("({}, 'b')", id2))); } + #[test] + fn a_prepared_column_is_cast_to_the_tables_type() { + // Prepared as an integer, the field is a float in the table now: the + // padded batch casts it, where a mismatch failed the whole batch. + let ensured = EnsuredTable { + table: "`t`".to_string(), + series_table: "`t_series`".to_string(), + tag_phys: vec![], + field_phys: vec![("v".to_string(), "v".to_string(), 0)], + }; + let schema = Arc::new(Schema::new(vec![ + Field::new( + "time", + DataType::Timestamp(TimeUnit::Nanosecond, Some("UTC".into())), + false, + ), + Field::new("origin_node_id", DataType::UInt64, false), + Field::new("ingest_seq", DataType::UInt64, false), + Field::new("series_id", DataType::UInt64, false), + Field::new("v", DataType::Int64, true), + ])); + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(TimestampNanosecondArray::from(vec![1i64]).with_timezone("UTC")), + Arc::new(UInt64Array::from(vec![0u64])), + Arc::new(UInt64Array::from(vec![0u64])), + Arc::new(UInt64Array::from(vec![7u64])), + Arc::new(arrow::array::Int64Array::from(vec![3i64])), + ], + ) + .unwrap(); + let padded = pad_record_batch_to_ensured(&batch, &ensured).expect("padded"); + let v = padded + .column(4) + .as_any() + .downcast_ref::() + .expect("float"); + assert_eq!(v.value(0), 3.0); + } + + #[test] + fn a_series_batch_is_padded_to_the_tables_tags() { + // Prepared before the table had `host` (or merged from series + // without it): `host` is filled with "", the columns come in the + // table's order, and its metadata entries can be read. + let ensured = EnsuredTable { + table: "`db_rp_m`".to_string(), + series_table: "`db_rp_m_series`".to_string(), + tag_phys: vec![ + ( + "host".to_string(), + "host".to_string(), + ColumnKind::TagLowCardinality, + ), + ( + "region".to_string(), + "region".to_string(), + ColumnKind::TagString, + ), + ], + field_phys: vec![], + }; + let older = EnsuredTable { + table: String::new(), + series_table: String::new(), + tag_phys: vec![( + "region".to_string(), + "region".to_string(), + ColumnKind::TagString, + )], + field_phys: vec![], + }; + let p = make_point(0, &[("region", "us")], &[]); + let batch = build_series_record_batch(&older, &[(1u64, &p)]).expect("batch"); + let padded = pad_series_batch_to_ensured(&batch, &ensured).expect("padded"); + let schema = padded.schema(); + let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert_eq!(names, ["series_id", "host", "region"]); + + // A tag the table plan does not list yet is kept, not dropped. + let unplanned = EnsuredTable { + tag_phys: vec![], + ..older + }; + let kept = pad_series_batch_to_ensured(&batch, &unplanned).expect("kept"); + let schema = kept.schema(); + let names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect(); + assert_eq!(names, ["series_id", "region"]); + assert_eq!( + padded.column(1).data_type(), + &tag_arrow_type(ColumnKind::TagLowCardinality) + ); + let entries = series_entries_from_batch(&padded, &ensured.tag_phys).expect("entries"); + assert_eq!(entries[0].1.get("region"), Some(&"us".to_string())); + assert!(entries[0].1.get("host").is_none_or(|h| h.is_empty())); + } + #[test] fn series_entries_from_batch_decodes_logical_tags() { let ensured = EnsuredTable { diff --git a/hyperbytedb/src/domain/arrow_coalesce.rs b/hyperbytedb/src/domain/arrow_coalesce.rs index 0ad0ef1..920439c 100644 --- a/hyperbytedb/src/domain/arrow_coalesce.rs +++ b/hyperbytedb/src/domain/arrow_coalesce.rs @@ -2,17 +2,19 @@ use std::sync::Arc; -use arrow::array::ArrayRef; -use arrow::compute::{cast, concat_batches}; -use arrow::datatypes::DataType; +use arrow::array::{ArrayRef, StringArray, new_null_array}; +use arrow::compute::{CastOptions, cast, cast_with_options, concat_batches}; +use arrow::datatypes::{DataType, Field, Schema}; use arrow::record_batch::RecordBatch; use crate::domain::prepared_wal::PreparedMeasurementBatch; use crate::error::HyperbytedbError; /// Concatenate multiple prepared batches for the same measurement. Row-level -/// partial-line merge is handled at ingest when possible; here we only stack -/// batches with identical schemas. +/// partial-line merge is handled at ingest when possible; here the batches +/// are stacked by column name. They need not share a schema: two writes to a +/// measurement can carry different fields (`used`, then `free`), and stacking +/// them by position put `free`'s values under `used`. pub fn coalesce_prepared_batches( mut batches: Vec, ) -> Result, HyperbytedbError> { @@ -24,9 +26,8 @@ pub fn coalesce_prepared_batches( } let template = &batches[0]; - let schema = template.batch.schema(); let refs: Vec<&RecordBatch> = batches.iter().map(|b| b.batch.as_ref()).collect(); - let merged = concat_batches(&schema, refs) + let merged = concat_by_name(&refs, Missing::Null) .map_err(|e| HyperbytedbError::Internal(format!("concat prepared batches: {e}").into()))?; let min_time = batches.iter().map(|b| b.min_time).min().unwrap_or(0); @@ -63,9 +64,9 @@ fn merge_series_batches( 0 => Ok(None), 1 => Ok(batches.into_iter().next()), _ => { - let schema = batches[0].schema(); + // New series of one measurement can carry different tag keys. let refs: Vec<&RecordBatch> = batches.iter().map(|b| b.as_ref()).collect(); - let merged = concat_batches(&schema, refs).map_err(|e| { + let merged = concat_by_name(&refs, Missing::EmptyString).map_err(|e| { HyperbytedbError::Internal(format!("concat series batches: {e}").into()) })?; Ok(Some(Arc::new(normalize_dictionary_columns(merged)?))) @@ -73,6 +74,96 @@ fn merge_series_batches( } } +/// What a batch without one of the merged columns gets in its place. +#[derive(Clone, Copy)] +pub(crate) enum Missing { + /// A field the rows did not have. + Null, + /// A tag the series did not have: the series table stores "". + EmptyString, +} + +/// Stack `batches` by column name: the result has every column of any batch, +/// in the order they first appear, and each batch's columns go under their +/// own names. A column one batch lacks is filled per `missing`; a column whose +/// type differs between batches is cast to the wider type (see +/// [`wider_type`]). +pub(crate) fn concat_by_name( + batches: &[&RecordBatch], + missing: Missing, +) -> Result { + let mut fields: Vec = Vec::new(); + for batch in batches { + for field in batch.schema().fields() { + match fields.iter_mut().find(|f| f.name() == field.name()) { + Some(existing) => { + let wider = wider_type(existing.data_type(), field.data_type()); + *existing = existing + .clone() + .with_data_type(wider) + .with_nullable(existing.is_nullable() || field.is_nullable()); + } + None => fields.push(field.as_ref().clone()), + } + } + } + // A column some batch lacks is filled with nulls (or ""), so it must be + // nullable when nulls are what fill it. + for field in &mut fields { + let everywhere = batches + .iter() + .all(|b| b.schema().index_of(field.name()).is_ok()); + if !everywhere && matches!(missing, Missing::Null) { + *field = field.clone().with_nullable(true); + } + } + let schema = Arc::new(Schema::new(fields)); + let safe = CastOptions { + safe: true, + ..Default::default() + }; + let mut aligned = Vec::with_capacity(batches.len()); + for batch in batches { + let n = batch.num_rows(); + let mut columns: Vec = Vec::with_capacity(schema.fields().len()); + for field in schema.fields() { + let column = match batch.schema().index_of(field.name()) { + Ok(i) if batch.column(i).data_type() == field.data_type() => { + Arc::clone(batch.column(i)) + } + Ok(i) => cast_with_options(batch.column(i), field.data_type(), &safe)?, + Err(_) => match missing { + Missing::Null => new_null_array(field.data_type(), n), + Missing::EmptyString => { + let empty: ArrayRef = Arc::new(StringArray::from(vec![""; n])); + cast(&empty, field.data_type())? + } + }, + }; + columns.push(column); + } + aligned.push(RecordBatch::try_new(Arc::clone(&schema), columns)?); + } + concat_batches(&schema, &aligned) +} + +/// The type a column takes when batches disagree, as ingest widens a field +/// (`widen_field_disc`): float over signed integer, unsigned over signed, and +/// a plain string over a dictionary-encoded one. Anything else (float and +/// unsigned included, which ingest does not widen) keeps the first type; the +/// other batch's values are cast to it, and those that cannot be become null. +/// The insert then casts to the table's type, which decides. +fn wider_type(a: &DataType, b: &DataType) -> DataType { + use DataType::*; + match (a, b) { + _ if a == b => a.clone(), + (Float64, Int64) | (Int64, Float64) => Float64, + (Int64, UInt64) | (UInt64, Int64) => UInt64, + (Dictionary(_, _) | Utf8, Dictionary(_, _) | Utf8) => Utf8, + _ => a.clone(), + } +} + /// Re-encode every dictionary-typed column so its dictionary contains only /// unique values. /// @@ -175,6 +266,140 @@ mod tests { } } + use arrow::array::{Float64Array, Int64Array}; + + fn fact(columns: Vec<(&str, ArrayRef)>) -> PreparedMeasurementBatch { + let n = columns[0].1.len(); + let schema = Arc::new(Schema::new( + columns + .iter() + .map(|(name, a)| Field::new(*name, a.data_type().clone(), true)) + .collect::>(), + )); + let batch = + RecordBatch::try_new(schema, columns.into_iter().map(|(_, a)| a).collect()).unwrap(); + PreparedMeasurementBatch { + measurement: "m".into(), + table_name: "`t`".into(), + series_table_name: "`t_series`".into(), + batch: Arc::new(batch), + row_count: n, + min_time: 0, + max_time: 0, + new_series_batch: None, + } + } + + fn floats(v: &[Option]) -> ArrayRef { + Arc::new(Float64Array::from(v.to_vec())) + } + + /// Two writes with different fields, flushed together: each value stays + /// under its own field. Stacked by position, `free`'s value landed under + /// `used` (same column count, same type) and `free` was null. + #[test] + fn batches_with_different_fields_stack_by_name() { + let sid = |v: u64| -> ArrayRef { Arc::new(UInt64Array::from(vec![v])) }; + let a = fact(vec![("series_id", sid(1)), ("used", floats(&[Some(1.0)]))]); + let b = fact(vec![("series_id", sid(2)), ("free", floats(&[Some(2.0)]))]); + let merged = coalesce_prepared_batches(vec![a, b]) + .unwrap() + .unwrap() + .batch; + let col = |name: &str| { + let i = merged.schema().index_of(name).unwrap(); + merged + .column(i) + .as_any() + .downcast_ref::() + .unwrap() + .clone() + }; + assert_eq!(col("used"), Float64Array::from(vec![Some(1.0), None])); + assert_eq!(col("free"), Float64Array::from(vec![None, Some(2.0)])); + } + + /// A field written as an integer in one batch and a float in another is + /// a float, as ingest widens it. + #[test] + fn a_field_whose_type_differs_takes_the_wider_type() { + let a = fact(vec![( + "v", + Arc::new(Int64Array::from(vec![1i64])) as ArrayRef, + )]); + let b = fact(vec![("v", floats(&[Some(2.5)]))]); + let merged = coalesce_prepared_batches(vec![a, b]) + .unwrap() + .unwrap() + .batch; + let v = merged + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(v, &Float64Array::from(vec![1.0, 2.5])); + } + + /// Float and unsigned are not widened into each other, as at ingest: the + /// first type stays and the table's type decides at insert. + #[test] + fn float_and_unsigned_keep_the_first_type() { + assert_eq!( + wider_type(&DataType::Float64, &DataType::UInt64), + DataType::Float64 + ); + assert_eq!( + wider_type(&DataType::UInt64, &DataType::Float64), + DataType::UInt64 + ); + assert_eq!( + wider_type(&DataType::Int64, &DataType::Float64), + DataType::Float64 + ); + assert_eq!( + wider_type(&DataType::Int64, &DataType::UInt64), + DataType::UInt64 + ); + } + + /// New series with different tags, flushed together: each tag value + /// stays under its own tag, and a tag a series lacks is "". + #[test] + fn series_with_different_tags_stack_by_name() { + let series = |sid: u64, tag: &str, value: &str| -> Arc { + let schema = Arc::new(Schema::new(vec![ + Field::new("series_id", DataType::UInt64, false), + Field::new(tag, DataType::Utf8, false), + ])); + Arc::new( + RecordBatch::try_new( + schema, + vec![ + Arc::new(UInt64Array::from(vec![sid])), + Arc::new(StringArray::from(vec![value])), + ], + ) + .unwrap(), + ) + }; + let merged = merge_series_batches(vec![series(1, "host", "a"), series(2, "region", "x")]) + .unwrap() + .unwrap(); + let tag = |name: &str| { + let i = merged.schema().index_of(name).unwrap(); + let a = merged + .column(i) + .as_any() + .downcast_ref::() + .unwrap(); + (0..a.len()) + .map(|r| a.value(r).to_string()) + .collect::>() + }; + assert_eq!(tag("host"), ["a", ""]); + assert_eq!(tag("region"), ["", "x"]); + } + /// Coalescing series batches whose dictionaries overlap must yield a /// dictionary with unique values, even though `concat` alone leaves /// duplicates for low-cardinality columns (the chDB Code 117 trigger). diff --git a/hyperbytedb/tests/compat/prepared_wal_tests.rs b/hyperbytedb/tests/compat/prepared_wal_tests.rs index b2fc2f4..0bb0ebc 100644 --- a/hyperbytedb/tests/compat/prepared_wal_tests.rs +++ b/hyperbytedb/tests/compat/prepared_wal_tests.rs @@ -213,3 +213,94 @@ async fn an_insert_missing_a_column_repairs_the_table_on_retry() { .await .expect("the retry re-plans the table and adds the column back"); } + +/// Writes to a new measurement with different fields and tags, flushed +/// together through prepared Arrow batches: every value stays under its own +/// field and tag. The batches were stacked by position, so `free`'s values +/// were stored under `used` (and `region`'s under `host`). +#[tokio::test] +#[serial(chdb)] +async fn writes_with_different_fields_flushed_together_keep_their_columns() { + use hyperbytedb::adapters::chdb::query_adapter::ChdbQueryAdapter; + use hyperbytedb::adapters::wal::rocksdb_wal::RocksDbWal; + use hyperbytedb::application::flush_service::FlushServiceImpl; + use hyperbytedb::application::ingest_metadata::IngestCardinalityLimits; + use hyperbytedb::application::ingestion_service::IngestionServiceImpl; + use hyperbytedb::application::query_service::QueryServiceImpl; + use hyperbytedb::ports::ingestion::{IngestionPort, WritePayloadFormat}; + use hyperbytedb::ports::points_sink::PointsSinkPort; + use hyperbytedb::ports::query::{QueryPort, QueryService}; + use std::sync::Arc; + + let tmpdir = tempfile::tempdir().unwrap(); + let dir = |d: &str| { + let p = tmpdir.path().join(d); + std::fs::create_dir_all(&p).unwrap(); + p + }; + let wal = Arc::new(RocksDbWal::open(dir("wal")).unwrap()); + let metadata = Arc::new(RocksDbMetadata::open(dir("meta")).unwrap()); + metadata.create_database("db").await.unwrap(); + let shared = SharedSession::new_eager(dir("chdb").to_str().unwrap(), 1).unwrap(); + let sink: Arc = Arc::new(ChdbNativeAdapter::with_metadata( + shared.clone(), + Some(metadata.clone()), + )); + let query_port: Arc = Arc::new(ChdbQueryAdapter::from_shared(shared, 0)); + let ingestion = IngestionServiceImpl::with_sink( + wal.clone(), + Some(sink.clone()), + metadata.clone(), + IngestCardinalityLimits::default(), + 0, + ); + let flush = FlushServiceImpl::new(wal.clone(), 0, sink.clone()); + let queries = QueryServiceImpl::new(query_port, metadata, wal, 30, sink); + + for body in [ + "m,host=a used=1 1000000000", + "m,host=a free=2 2000000000", + "m,region=x free=3 3000000000", + ] { + ingestion + .ingest( + "db", + None, + None, + body.as_bytes(), + WritePayloadFormat::LineProtocol, + ) + .await + .unwrap(); + } + flush.flush().await.unwrap(); + + let resp = queries + .execute_query("db", "SELECT * FROM m", Some("s"), None, None) + .await + .unwrap(); + let series = &resp.results[0].series.as_ref().unwrap()[0]; + let rows: Vec = series + .values + .iter() + .map(|v| { + serde_json::Value::Object( + series + .columns + .iter() + .cloned() + .zip(v.iter().cloned()) + .collect(), + ) + }) + .collect(); + assert_eq!( + serde_json::Value::Array(rows), + // A series without a tag reads as ""; a whole float is written `1`. + serde_json::json!([ + {"time": 1, "free": null, "host": "a", "region": "", "used": 1}, + {"time": 2, "free": 2, "host": "a", "region": "", "used": null}, + {"time": 3, "free": 3, "host": "", "region": "x", "used": null}, + ]) + ); +}