diff --git a/Cargo.toml b/Cargo.toml index 60184a0b..d897cf5c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -196,6 +196,8 @@ arrow = "58.2.0" # Required by `record_batch!()` arrow-schema = "58.2.0" +insta = "1.48.0" + # Only used in testing and examples if the `opentelemetry` feature is enabled, # but dev-dependencies cannot be optional. tracing-subscriber = { version = "0.3.23", features = ["env-filter", "json"] } diff --git a/src/error.rs b/src/error.rs index 916a0db4..d3ee71ed 100644 --- a/src/error.rs +++ b/src/error.rs @@ -7,7 +7,7 @@ use std::{error::Error as StdError, fmt, io, result, str::Utf8Error}; /// A result with a specified [`Error`] type. pub type Result = result::Result; -type BoxedError = Box; +pub(crate) type BoxedError = Box; /// Represents all possible errors. #[derive(Debug, thiserror::Error)] diff --git a/src/insert_native.rs b/src/insert_native.rs new file mode 100644 index 00000000..566d0e31 --- /dev/null +++ b/src/insert_native.rs @@ -0,0 +1,47 @@ +use crate::native::Block; +use crate::native::writer::BlockWriter; +use crate::{Client, Compression, insert_formatted, sql}; + +pub struct InsertNative { + writer: BlockWriter, +} + +impl InsertNative { + pub(crate) fn new(client: &Client, table_name: &str, escape: bool) -> Self { + let mut sql = "INSERT INTO ".to_string(); + + if escape { + sql::escape::identifier(table_name, &mut sql).expect("error escaping table name"); + } else { + sql.push_str(table_name); + } + + sql.push_str(" FORMAT Native"); + + Self { + writer: BlockWriter::new(insert_formatted::InsertFormatted::new( + // FIXME: use HTTP body compression instead of block-level compression + &client.clone().with_compression(Compression::None), + sql, + Some(table_name), + )), + } + } + + /// Send a block of data. + /// + /// # NOT Cancel Safe + /// If this `async` method is canceled (i.e. by dropping the resulting `Future`), + /// the insert is automatically aborted. + /// + /// This is because the block data is not sent in a single write, since that would require + /// copying it into a separate buffer. There is no way to resynchronize the stream + /// once a block has been partially sent. Resuming a write would corrupt the stream. + pub async fn write(&mut self, block: &Block) -> crate::Result<()> { + self.writer.write(block).await + } + + pub async fn end(self) -> crate::Result<()> { + self.writer.end().await + } +} diff --git a/src/lib.rs b/src/lib.rs index b86a66c6..7e335150 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -21,6 +21,7 @@ use tokio::sync::RwLock; pub mod error; pub mod insert; pub mod insert_formatted; +pub mod insert_native; #[cfg(feature = "inserter")] pub mod inserter; pub mod native; @@ -543,6 +544,14 @@ impl Client { insert_formatted::InsertFormatted::new(self, sql.into(), None) } + pub fn insert_native(&self, table_name: &str) -> insert_native::InsertNative { + insert_native::InsertNative::new(self, table_name, true) + } + + pub fn insert_native_unescaped(&self, raw_table_name: &str) -> insert_native::InsertNative { + insert_native::InsertNative::new(self, raw_table_name, false) + } + /// Starts a new SELECT/DDL query. pub fn query(&self, query: &str) -> query::Query { query::Query::new(self, query) diff --git a/src/native/builder.rs b/src/native/builder.rs new file mode 100644 index 00000000..7a4ee406 --- /dev/null +++ b/src/native/builder.rs @@ -0,0 +1,768 @@ +use crate::error::BoxedError; +use crate::native::encode::{Encode, ValueWriter}; +use crate::native::string::MaybeUtf8; +use crate::native::utils::{DebugFixedData, DebugNullMap, DebugVariableData, type_fixed_width}; +use crate::native::{Block, Column, Layout, LayoutKind}; +use bytes::{BufMut, BytesMut}; +use clickhouse_types::DataTypeNode; +use hashbrown::{HashMap, hash_map}; +use std::collections::VecDeque; +use std::fmt::{Debug, Formatter}; +use std::marker::PhantomData; +use std::mem; + +#[derive(Default)] +pub struct BlockBuilder { + column_names: HashMap, + columns: Vec, +} + +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum BlockBuilderError { + #[error( + "attempting to overwrite existing column `{name} {existing_type}` with a different type: {new_type}" + )] + ColumnExists { + name: String, + existing_type: DataTypeNode, + new_type: DataTypeNode, + }, + #[error("unsupported type or subtype of column `{name}`: `{data_type}`")] + UnsupportedType { + name: String, + data_type: DataTypeNode, + }, + #[error( + "block contains columns of mismatched lengths; \ + longest column: `{longest_column}` (len: {longest_len}), \ + shortest column: `{shortest_column}` (len: {shortest_len})" + )] + MismatchedLengths { + longest_column: String, + longest_len: usize, + shortest_column: String, + shortest_len: usize, + }, + #[error("column `{column_name} {column_type}` contains invalid data: {message}")] + ColumnDataInvalid { + column_name: String, + column_type: DataTypeNode, + message: String, + }, +} + +impl BlockBuilder { + pub fn new() -> Self { + Self::default() + } + + /// Add an empty column to the block, or get a reference to an existing one. + /// + /// The given data type will have any `LowCardinality(_)` or `SimpleAggregateFunction(...)` + /// wrappers erased for ease of implementation. + /// + /// # Errors + /// * If a column with the same name already exists, but with a different type. + /// * If the given type is not currently supported by the implementation. + pub fn upsert_column( + &mut self, + name: impl Into, + ) -> Result, Box> { + self.upsert_column_with(name, T::produces()) + .map(|inner| ColumnBuilder { + inner, + _marker: PhantomData, + }) + } + + fn upsert_column_with( + &mut self, + name: impl Into, + data_type: DataTypeNode, + ) -> Result<&mut ColumnBuilderRaw, Box> { + let data_type = erase_wrappers(data_type); + + match self.column_names.entry(MaybeUtf8::from_string(name)) { + hash_map::Entry::Occupied(existing) => { + let col = &mut self.columns[*existing.get()]; + + if col.data_type != data_type { + return Err(BlockBuilderError::ColumnExists { + name: col.name.to_string(), + existing_type: col.data_type.clone(), + new_type: data_type, + } + .into()); + } + + Ok(col) + } + hash_map::Entry::Vacant(vacant) => { + let col = ColumnBuilderRaw { + layout: LayoutBuilder::new(vacant.key(), &data_type)?, + data_type, + name: vacant.key().clone(), + }; + + vacant.insert(self.columns.len()); + // FIXME: replace with `Vec::push_mut()` after Rust 1.95 + self.columns.push(col); + Ok(self.columns.last_mut().unwrap()) + } + } + } + + pub fn build(&mut self) -> Result> { + let mut num_rows = 0; + + // Check that all the columns have the same length + if let Some((mut longest_col, columns)) = self.columns.split_first() { + let mut len_mismatch = false; + + let mut shortest_col = longest_col; + + num_rows = longest_col.num_values(); + + for col in columns { + if col.num_values() > longest_col.num_values() { + longest_col = col; + len_mismatch = true; + } + + if col.num_values() < shortest_col.num_values() { + shortest_col = col; + len_mismatch = true; + } + } + + if len_mismatch { + return Err(BlockBuilderError::MismatchedLengths { + longest_column: longest_col.name.to_string(), + longest_len: longest_col.num_values(), + shortest_column: shortest_col.name.to_string(), + shortest_len: shortest_col.num_values(), + } + .into()); + } + } + + // Note: try to perform as much validation as possible before consuming `self` + for col in &self.columns { + col.layout.validate(&col.data_type).map_err(|message| { + BlockBuilderError::ColumnDataInvalid { + column_name: col.name.to_string(), + column_type: col.data_type.clone(), + message, + } + })?; + } + + let columns = self + .columns + .drain(..) + .map(|col| Column { + name: col.name, + data_type: col.data_type, + layout: col.layout.into_layout(), + }) + .collect(); + + Ok(Block { + columns, + column_names: mem::take(&mut self.column_names), + num_rows, + }) + } +} + +impl Debug for BlockBuilder { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BlockBuilder") + // Ignore `column_names`, implementation detail + .field("columns", &self.columns) + .finish() + } +} + +pub struct ColumnBuilder<'a, T> { + inner: &'a mut ColumnBuilderRaw, + _marker: PhantomData, +} + +impl ColumnBuilder<'_, T> +where + T: Encode, +{ + pub fn num_values(&self) -> usize { + self.inner.layout.num_values() + } + + pub fn add(&mut self, value: T) -> Result<&mut Self, BoxedError> { + // Compatibility checked when this was created + self.inner.add_unchecked(value)?; + Ok(self) + } + + pub fn add_all(&mut self, values: I) -> Result<&mut Self, BoxedError> + where + I: IntoIterator, + { + self.inner.add_all_unchecked(values)?; + Ok(self) + } +} + +impl Debug for ColumnBuilder<'_, T> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + self.inner.fmt(f) + } +} + +#[derive(Debug)] // Derived impl works for us here +struct ColumnBuilderRaw { + name: MaybeUtf8, + data_type: DataTypeNode, + layout: LayoutBuilder, +} + +impl ColumnBuilderRaw { + fn num_values(&self) -> usize { + self.layout.num_values() + } + + fn add_unchecked(&mut self, value: T) -> Result<&mut Self, BoxedError> + where + T: Encode, + { + value.encode(&mut ValueWriter { + data_type: &self.data_type, + layout: &mut self.layout, + })?; + + Ok(self) + } + + fn add_all_unchecked(&mut self, values: I) -> Result<&mut Self, BoxedError> + where + I: IntoIterator, + I::Item: Encode, + { + let mut values = values.into_iter(); + + while let Some(value) = values.next() { + // Catches an infinite-length iterator that returns `usize::MAX` for its size hint + // This is comparable to the default behavior of `impl Extend for Vec` + let (lower_bound, _) = values.size_hint(); + self.layout.reserve(lower_bound.saturating_add(1)); + + self.add_unchecked(value)?; + } + + Ok(self) + } +} + +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum ColumnBuilderError { + #[error("incompatible value type for column `{name} {data_type}")] + IncompatibleType { + name: String, + data_type: DataTypeNode, + }, + #[error("error encoding value")] + Encode(#[source] BoxedError), +} + +// These types may be identical to `Layout` but they need to use growable containers. +// `LowCardinality` is also deliberately omitted in the initial implementation +// since the server can do the transformation automatically. +pub(super) struct LayoutBuilder { + pub(super) kind: LayoutBuilderKind, + pub(super) nulls: Option, +} + +pub(super) enum LayoutBuilderKind { + /// Fixed layout. Width of each cell depends only on [`DataTypeNode`]. + Fixed { + type_width: usize, + data: BytesMut, + }, + /// Variable-length data (namely strings) + Variable { + /// Ending offset of each string in `data`. + /// + /// The offset of the first string is always `0` unless this is empty. + end_offsets: Vec, + // Each `Bytes` instance is 4 `usizes` (32 bytes on 64-bit), + // so we save 24 bytes per string by linearizing the string data and storing offsets, + // assuming many small strings instead of fewer big ones, which also amortizes allocations + data: BytesMut, + }, + /// Array data. Element data governed by `elem_layout`. + Array { + /// Ending index of each array in `elem_layout`. + end_indices: Vec, + elem_layout: Box, + }, + Tuple { + layouts: Box<[LayoutBuilder]>, + }, + Map { + key_val_layouts: Box<[LayoutBuilder; 2]>, + end_indices: Vec, + }, +} + +impl LayoutBuilder { + fn new( + column_name: &MaybeUtf8, + data_type: &DataTypeNode, + ) -> Result> { + let (non_nullable, nulls) = if let DataTypeNode::Nullable(inner) = data_type { + (&**inner, Some(BytesMut::new())) + } else { + (data_type, None) + }; + + if let Some(type_width) = type_fixed_width(non_nullable) { + return Ok(Self { + nulls, + kind: LayoutBuilderKind::Fixed { + type_width, + data: Default::default(), + }, + }); + }; + + match non_nullable { + DataTypeNode::String => Ok(Self { + nulls, + kind: LayoutBuilderKind::Variable { + end_offsets: vec![], + data: Default::default(), + }, + }), + DataTypeNode::Tuple(types) => Ok(Self { + nulls, + kind: LayoutBuilderKind::Tuple { + layouts: types + .iter() + .map(|ty| LayoutBuilder::new(column_name, ty)) + .collect::>()?, + }, + }), + DataTypeNode::Array(elem_type) => Ok(Self { + nulls, + kind: LayoutBuilderKind::Array { + end_indices: vec![], + elem_layout: Box::new(LayoutBuilder::new(column_name, elem_type)?), + }, + }), + DataTypeNode::Map(key_val_types) => Ok(Self { + nulls, + kind: LayoutBuilderKind::Map { + key_val_layouts: Box::new([ + LayoutBuilder::new(column_name, &key_val_types[0])?, + LayoutBuilder::new(column_name, &key_val_types[1])?, + ]), + end_indices: vec![], + }, + }), + _ => Err(Box::new(BlockBuilderError::UnsupportedType { + name: column_name.to_string(), + data_type: data_type.clone(), + })), + } + } + + pub(super) fn num_values(&self) -> usize { + match &self.kind { + LayoutBuilderKind::Fixed { type_width, data } => data.len() / type_width, + LayoutBuilderKind::Variable { end_offsets, .. } => end_offsets.len(), + LayoutBuilderKind::Array { end_indices, .. } => end_indices.len(), + LayoutBuilderKind::Tuple { layouts, .. } => { + layouts.first().map_or(0, |layout| layout.num_values()) + } + LayoutBuilderKind::Map { end_indices, .. } => end_indices.len(), + } + } + + pub(super) fn reserve(&mut self, additional: usize) { + match &mut self.kind { + LayoutBuilderKind::Fixed { type_width, data } => { + data.reserve(type_width.saturating_mul(additional)); + } + LayoutBuilderKind::Variable { end_offsets, .. } => { + end_offsets.reserve(additional); + // Don't reserve in `data` because we don't know the total additional size + } + LayoutBuilderKind::Array { .. } => {} + LayoutBuilderKind::Tuple { .. } => {} + LayoutBuilderKind::Map { .. } => {} + } + } + + /// Push a valid placeholder value + pub(super) fn push_placeholder(&mut self) { + match &mut self.kind { + LayoutBuilderKind::Fixed { type_width, data } => { + data.put_bytes(0, *type_width); + } + LayoutBuilderKind::Variable { end_offsets, data } => { + end_offsets.push(data.len()); + } + LayoutBuilderKind::Array { end_indices, .. } => { + let end_index = end_indices.last().copied().unwrap_or(0); + end_indices.push(end_index); + } + // This is only needed for `Nullable(Tuple(...))` which is currently experimental + LayoutBuilderKind::Tuple { layouts } => { + for layout in layouts { + layout.push_placeholder(); + } + } + LayoutBuilderKind::Map { end_indices, .. } => { + let end_index = end_indices.last().copied().unwrap_or(0); + end_indices.push(end_index); + } + } + } + + /// Truncate to the given number of values. + /// + /// For arrays and maps, this truncates to the length of the array at `num_values`. + pub(super) fn truncate(&mut self, num_values: usize) { + match &mut self.kind { + LayoutBuilderKind::Fixed { type_width, data } => { + data.truncate(type_width.saturating_mul(num_values)); + } + LayoutBuilderKind::Variable { end_offsets, data } => { + end_offsets.truncate(num_values); + + let last_offset = end_offsets.last().copied().unwrap_or(0); + data.truncate(last_offset); + } + LayoutBuilderKind::Array { + end_indices, + elem_layout, + } => { + end_indices.truncate(num_values); + + let last_index = end_indices.last().copied().unwrap_or(0); + elem_layout.truncate(last_index); + } + LayoutBuilderKind::Tuple { layouts } => { + for layout in layouts { + layout.truncate(num_values); + } + } + LayoutBuilderKind::Map { + key_val_layouts, + end_indices, + } => { + end_indices.truncate(num_values); + + let last_index = end_indices.last().copied().unwrap_or(0); + + key_val_layouts[0].truncate(last_index); + key_val_layouts[1].truncate(last_index); + } + } + } + + fn validate(&self, data_type: &DataTypeNode) -> Result<(), String> { + self.validate_nulls(data_type)?; + + let non_nullable = if let DataTypeNode::Nullable(inner) = data_type { + inner + } else { + data_type + }; + + match &self.kind { + LayoutBuilderKind::Fixed { type_width, data } => { + let expected_width = type_fixed_width(non_nullable) + .ok_or_else(|| format!("data type {non_nullable} is not fixed-width but we encoded {} bytes of {type_width}-byte values", data.len()))?; + + if expected_width != *type_width { + return Err(format!( + "data type {non_nullable} has a fixed width of {expected_width} but we encoded {} bytes of {type_width}-byte values", + data.len() + )); + } + + if !data.len().is_multiple_of(*type_width) { + return Err(format!( + "data length ({}) is not a multiple of type_width ({type_width})", + data.len() + )); + } + + Ok(()) + } + LayoutBuilderKind::Variable { end_offsets, data } => { + for (i, &end_offset) in end_offsets.iter().enumerate() { + if end_offset > data.len() { + return Err(format!( + "string {i} end offset {end_offset} is out of bounds: {}", + data.len() + )); + } + } + + Ok(()) + } + LayoutBuilderKind::Array { + end_indices, + elem_layout, + } => { + let DataTypeNode::Array(elem_type) = non_nullable else { + return Err(format!("expected type Array(_), got {non_nullable}")); + }; + + let num_elements = elem_layout.num_values(); + + for (i, &end_index) in end_indices.iter().enumerate() { + if end_index > num_elements { + return Err(format!( + "array {i} end index ({end_index}) out of bounds: {num_elements}" + )); + } + } + + let last_index = end_indices.last().copied().unwrap_or(0); + + if last_index != num_elements { + // Most likely cause of this error is a leaked `ArrayWriter` + return Err(format!( + "last array index ({last_index}) out of sync with total elements: {num_elements}" + )); + } + + elem_layout.validate(elem_type) + } + LayoutBuilderKind::Tuple { layouts } => { + let DataTypeNode::Tuple(types) = non_nullable else { + return Err(format!("expected type Tuple(...), got {non_nullable}")); + }; + + let expected_len = layouts.first().map_or(0, LayoutBuilder::num_values); + + for (i, (ty, layout)) in types.iter().zip(layouts).enumerate() { + layout.validate(ty)?; + + let actual_len = layout.num_values(); + + if layout.num_values() != expected_len { + // Most likely cause of this error is a leaked `TupleWriter` + return Err(format!( + "tuple index {i} (type {ty}) total elements out of sync: {actual_len} vs {expected_len}" + )); + } + } + + Ok(()) + } + LayoutBuilderKind::Map { + key_val_layouts, + end_indices, + } => { + let DataTypeNode::Map([key_ty, val_ty]) = non_nullable else { + return Err(format!("expected type Map(...), got {non_nullable}")); + }; + + let keys_len = key_val_layouts[0].num_values(); + let values_len = key_val_layouts[1].num_values(); + + if keys_len != values_len { + return Err(format!( + "number of keys and values is out of sync: {keys_len} vs {values_len}" + )); + } + + for (i, &end_index) in end_indices.iter().enumerate() { + if end_index > keys_len { + return Err(format!( + "map {i} end index ({end_index}) out of bounds: {keys_len}" + )); + } + } + + let last_index = end_indices.last().copied().unwrap_or(0); + + if last_index != keys_len { + // Most likely cause of this error is a leaked `MapWriter` + return Err(format!( + "last map index ({last_index}) out of sync with total elements: {keys_len}" + )); + } + + key_val_layouts[0].validate(key_ty)?; + key_val_layouts[1].validate(val_ty)?; + + Ok(()) + } + } + } + + fn validate_nulls(&self, data_type: &DataTypeNode) -> Result<(), String> { + match (&self.nulls, data_type) { + (Some(nulls), DataTypeNode::Nullable(_)) => { + if nulls.len() != self.num_values() { + return Err(format!( + "null bitmap length invalid: {}; expected: {}", + nulls.len(), + self.num_values() + )); + } + } + (Some(nulls), _) => { + return Err(format!( + "null bitmap of length {} created for non-nullable type {data_type}", + nulls.len() + )); + } + (None, DataTypeNode::Nullable(_)) => { + return Err(format!("nullable type {data_type} missing null bitmap")); + } + _ => (), + } + + Ok(()) + } + + fn into_layout(self) -> Layout { + Layout { + num_values: self.num_values(), + nulls: self.nulls.map(BytesMut::freeze), + kind: match self.kind { + LayoutBuilderKind::Fixed { type_width, data } => LayoutKind::Fixed { + type_width, + data: data.freeze(), + }, + LayoutBuilderKind::Variable { end_offsets, data } => LayoutKind::Variable { + end_offsets: end_offsets.into(), + data: data.freeze(), + }, + LayoutBuilderKind::Array { + end_indices, + elem_layout, + } => LayoutKind::Array { + end_indices: end_indices.into(), + elem_layout: Box::new(elem_layout.into_layout()), + }, + LayoutBuilderKind::Tuple { layouts } => LayoutKind::Tuple { + layouts: layouts + .into_iter() + .map(LayoutBuilder::into_layout) + .collect(), + }, + LayoutBuilderKind::Map { + key_val_layouts, + end_indices, + } => LayoutKind::Map { + key_val_layouts: Box::new(key_val_layouts.map(LayoutBuilder::into_layout)), + end_indices: end_indices.into(), + }, + }, + } + } +} +impl Debug for LayoutBuilder { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LayoutBuilder") + .field("kind", &self.kind) + .field("nulls", &self.nulls.as_deref().map(DebugNullMap)) + .finish() + } +} + +impl Debug for LayoutBuilderKind { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + use LayoutBuilderKind::*; + + match self { + Fixed { type_width, data } => f + .debug_struct("Fixed") + .field( + "data", + &DebugFixedData { + type_width: *type_width, + data, + }, + ) + .finish(), + Variable { end_offsets, data } => f + .debug_struct("Variable") + .field("data", &DebugVariableData { end_offsets, data }) + .finish(), + Array { + elem_layout, + end_indices, + } => f + .debug_struct("Array") + .field("elem_layout", elem_layout) + .field("end_indices", end_indices) + .finish(), + Tuple { layouts } => { + let mut tuple = f.debug_tuple("Tuple"); + + for layout in layouts { + tuple.field(layout); + } + + tuple.finish() + } + Map { + key_val_layouts, + end_indices, + } => f + .debug_struct("Map") + .field("keys", &key_val_layouts[0]) + .field("values", &key_val_layouts[1]) + .field("end_indices", end_indices) + .finish(), + } + } +} + +/// Erase `LowCardinality` and `SimpleAggregateFunction` from the column type as inserts can be done +/// without them. +/// +/// TODO: encode `LowCardinality` +fn erase_wrappers(data_type: DataTypeNode) -> DataTypeNode { + match data_type { + DataTypeNode::LowCardinality(inner) | DataTypeNode::SimpleAggregateFunction(_, inner) => { + erase_wrappers(*inner) + } + DataTypeNode::Nullable(mut inner) => { + *inner = erase_wrappers(*inner); + DataTypeNode::Nullable(inner) + } + DataTypeNode::Array(mut inner) => { + *inner = erase_wrappers(*inner); + DataTypeNode::Array(inner) + } + DataTypeNode::Tuple(types) => { + // Converting to a `VeqDeque` is an `O(1)` operation that then lets us + // iterate through `types` by-value and push them back into the same allocation. + let mut types = VecDeque::from(types); + + for _ in 0..types.len() { + let ty = types.pop_front().unwrap(); + types.push_back(erase_wrappers(ty)); + } + + // The vector should be linear again, so this conversion should also be trivial. + DataTypeNode::Tuple(types.into()) + } + DataTypeNode::Map([mut key_ty, mut val_ty]) => { + *key_ty = erase_wrappers(*key_ty); + *val_ty = erase_wrappers(*val_ty); + DataTypeNode::Map([key_ty, val_ty]) + } + other => other, + } +} diff --git a/src/native/decode.rs b/src/native/decode.rs index 431b3e55..228f0977 100644 --- a/src/native/decode.rs +++ b/src/native/decode.rs @@ -1,7 +1,7 @@ +use crate::error::BoxedError; use crate::native::array::{ArrayData, TupleIter}; use clickhouse_types::DataTypeNode; use std::collections::{BTreeMap, HashMap}; -use std::error::Error; use std::hash::{BuildHasher, Hash}; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; @@ -31,6 +31,7 @@ impl<'a> ValueReader<'a> { } #[derive(Debug, thiserror::Error)] +#[non_exhaustive] pub enum ValueReadError { #[error("expected {expected} bytes, got {actual}")] InvalidLength { expected: usize, actual: usize }, @@ -39,27 +40,21 @@ pub enum ValueReadError { pub trait Decode<'a>: 'a + Sized { fn compatible(data_type: &DataTypeNode) -> bool; - fn decode(reader: &mut ValueReader<'a>) - -> Result>; + fn decode(reader: &mut ValueReader<'a>) -> Result; - fn decode_null( - data_type: &DataTypeNode, - ) -> Result> { + fn decode_null(data_type: &DataTypeNode) -> Result { Err(format!("data type {data_type:?} cannot be NULL").into()) } - fn decode_array(data: ArrayData<'a>) -> Result> { + fn decode_array(data: ArrayData<'a>) -> Result { Err(format!("unexpected data type Array({})", data.elem_type).into()) } - fn decode_tuple(data: TupleIter<'a>) -> Result> { + fn decode_tuple(data: TupleIter<'a>) -> Result { Err(format!("unexpected data type Tuple({:?})", data.types.as_slice()).into()) } - fn decode_map( - key_data: ArrayData<'a>, - value_data: ArrayData<'a>, - ) -> Result> { + fn decode_map(key_data: ArrayData<'a>, value_data: ArrayData<'a>) -> Result { Err(format!( "unexpected data type Map({}, {})", key_data.elem_type, value_data.elem_type @@ -93,7 +88,7 @@ macro_rules! impl_from_le_bytes { fn decode( reader: &mut ValueReader<'a>, - ) -> Result> { + ) -> Result { Ok($ty::from_le_bytes(*reader.read_bytes_fixed()?)) } } @@ -124,9 +119,7 @@ impl Decode<'_> for bool { type_matches!(data_type, DataTypeNode::Bool) } - fn decode( - reader: &mut ValueReader<'_>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'_>) -> Result { // https://clickhouse.com/docs/interfaces/specs/NativeFormat#bool let [b] = reader.read_bytes_fixed()?; Ok(*b != 0) @@ -138,9 +131,7 @@ impl<'a> Decode<'a> for &'a str { <&[u8] as Decode>::compatible(data_type) } - fn decode( - reader: &mut ValueReader<'a>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'a>) -> Result { Ok(str::from_utf8(<&[u8] as Decode>::decode(reader)?)?) } } @@ -150,9 +141,7 @@ impl<'a> Decode<'a> for String { <&str as Decode>::compatible(data_type) } - fn decode( - reader: &mut ValueReader<'a>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'a>) -> Result { Ok(<&str as Decode>::decode(reader)?.into()) } } @@ -165,9 +154,7 @@ impl<'a> Decode<'a> for &'a [u8] { ) } - fn decode( - reader: &mut ValueReader<'a>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'a>) -> Result { Ok(reader.native_bytes) } } @@ -181,9 +168,7 @@ impl<'a, T: Decode<'a>> Decode<'a> for Option { } } - fn decode( - reader: &mut ValueReader<'a>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'a>) -> Result { let DataTypeNode::Nullable(inner_type) = reader.data_type else { return Err(format!("expected `Nullable(_)`, got {:?}", reader.data_type).into()); }; @@ -194,17 +179,15 @@ impl<'a, T: Decode<'a>> Decode<'a> for Option { })?)) } - fn decode_null( - _data_type: &DataTypeNode, - ) -> Result> { + fn decode_null(_data_type: &DataTypeNode) -> Result { Ok(None) } - fn decode_array(data: ArrayData<'a>) -> Result> { + fn decode_array(data: ArrayData<'a>) -> Result { T::decode_array(data).map(Some) } - fn decode_tuple(data: TupleIter<'a>) -> Result> { + fn decode_tuple(data: TupleIter<'a>) -> Result { T::decode_tuple(data).map(Some) } } @@ -219,13 +202,11 @@ impl<'a, T: Decode<'a> + 'a> Decode<'a> for Vec { } } - fn decode( - reader: &mut ValueReader<'a>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'a>) -> Result { Err(format!("expected array type, got {}", reader.data_type).into()) } - fn decode_array(data: ArrayData<'a>) -> Result> { + fn decode_array(data: ArrayData<'a>) -> Result { data.into_reader::()? .collect::, crate::Error>>() .map_err(Into::into) @@ -256,11 +237,11 @@ macro_rules! tuple_impl { fn decode( reader: &mut ValueReader<'a>, - ) -> Result> { + ) -> Result { Err(format!("expected array type, got {}", reader.data_type).into()) } - fn decode_tuple(mut data: TupleIter<'a>) -> Result> { + fn decode_tuple(mut data: TupleIter<'a>) -> Result { Ok(( data.decode_next::<$ty1>()?, $(data.decode_next::<$ty>()?),* @@ -294,21 +275,16 @@ where && V::compatible(val_ty.remove_low_cardinality()) } - fn decode( - reader: &mut ValueReader<'a>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'a>) -> Result { Err(format!("expected map, got {}", reader.data_type).into()) } - fn decode_map( - key_data: ArrayData<'a>, - value_data: ArrayData<'a>, - ) -> Result> { + fn decode_map(key_data: ArrayData<'a>, value_data: ArrayData<'a>) -> Result { key_data .into_reader::()? .zip(value_data.into_reader::()?) .map(|(k, v)| Ok((k?, v?))) - .collect::>>() + .collect::>() } } @@ -326,21 +302,16 @@ where && V::compatible(val_ty.remove_low_cardinality()) } - fn decode( - reader: &mut ValueReader<'a>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'a>) -> Result { Err(format!("expected map, got {}", reader.data_type).into()) } - fn decode_map( - key_data: ArrayData<'a>, - value_data: ArrayData<'a>, - ) -> Result> { + fn decode_map(key_data: ArrayData<'a>, value_data: ArrayData<'a>) -> Result { key_data .into_reader::()? .zip(value_data.into_reader::()?) .map(|(k, v)| Ok((k?, v?))) - .collect::>>() + .collect::>() } } @@ -349,9 +320,7 @@ impl Decode<'_> for Ipv4Addr { type_matches!(data_type, DataTypeNode::IPv4) } - fn decode( - reader: &mut ValueReader<'_>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'_>) -> Result { // https://clickhouse.com/docs/interfaces/specs/NativeFormat#ipv4-and-ipv6 // IPv4 is byte-reversed, so little-endian let bytes_le = u32::from_le_bytes(*reader.read_bytes_fixed()?); @@ -365,9 +334,7 @@ impl Decode<'_> for Ipv6Addr { type_matches!(data_type, DataTypeNode::IPv6) } - fn decode( - reader: &mut ValueReader<'_>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'_>) -> Result { // https://clickhouse.com/docs/interfaces/specs/NativeFormat#ipv4-and-ipv6 // IPv6 uses canonical (big-endian) encoding Ok(Ipv6Addr::from(*reader.read_bytes_fixed::<16>()?)) @@ -379,9 +346,7 @@ impl Decode<'_> for IpAddr { type_matches!(data_type, DataTypeNode::IPv4 | DataTypeNode::IPv6) } - fn decode( - reader: &mut ValueReader<'_>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'_>) -> Result { match reader.data_type { DataTypeNode::IPv4 => Ipv4Addr::decode(reader).map(Into::into), DataTypeNode::IPv6 => Ipv6Addr::decode(reader).map(Into::into), @@ -393,8 +358,8 @@ impl Decode<'_> for IpAddr { #[cfg(feature = "uuid")] mod uuid { use super::{Decode, ValueReader}; + use crate::error::BoxedError; use clickhouse_types::DataTypeNode; - use std::error::Error; use uuid::Uuid; impl Decode<'_> for Uuid { @@ -402,9 +367,7 @@ mod uuid { type_matches!(data_type, DataTypeNode::UUID) } - fn decode( - reader: &mut ValueReader<'_>, - ) -> Result> { + fn decode(reader: &mut ValueReader<'_>) -> Result { // https://clickhouse.com/docs/interfaces/specs/NativeFormat#uuid // Wire bytes 0..7 = canonical bytes 0..7 reversed. // Wire bytes 8..15 = canonical bytes 8..15 reversed. diff --git a/src/native/encode.rs b/src/native/encode.rs new file mode 100644 index 00000000..0d4e1d07 --- /dev/null +++ b/src/native/encode.rs @@ -0,0 +1,831 @@ +use crate::error::BoxedError; +use crate::native::builder::{LayoutBuilder, LayoutBuilderKind}; +use bytes::{BufMut, BytesMut}; +use clickhouse_types::DataTypeNode; +use std::cmp; +use std::collections::{BTreeMap, HashMap}; +use std::marker::PhantomData; +use std::net::{Ipv4Addr, Ipv6Addr}; + +pub trait Encode { + fn produces() -> DataTypeNode; + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError>; + + fn compatible(column_type: &DataTypeNode) -> bool { + let produced_type = Self::produces(); + + default_compatible(&produced_type, column_type) + } +} + +fn default_compatible(produced_type: &DataTypeNode, column_type: &DataTypeNode) -> bool { + recursive_compatible(column_type, |column_type| { + if produced_type == column_type { + return true; + } + + match (produced_type, column_type) { + // SimpleAggregateFunction has the same wire image as the underlying type + // and the server should implicitly expand LowCardinality + ( + DataTypeNode::LowCardinality(left) | DataTypeNode::SimpleAggregateFunction(_, left), + right, + ) => default_compatible(left, right), + // Not-null value can be written to nullable column but not vice versa + (DataTypeNode::Nullable(left), DataTypeNode::Nullable(right)) => { + default_compatible(left, right) + } + ( + left, + DataTypeNode::LowCardinality(right) + | DataTypeNode::SimpleAggregateFunction(_, right) + | DataTypeNode::Nullable(right), + ) => default_compatible(left, right), + _ => false, + } + }) +} + +fn recursive_compatible bool>( + column_type: &DataTypeNode, + compatible: F, +) -> bool { + if compatible(column_type) { + return true; + } + + match column_type { + // SimpleAggregateFunction has the same wire image as the underlying type + // and the server should implicitly expand LowCardinality + // Not-null value can be written to nullable column but not vice versa + DataTypeNode::LowCardinality(inner) + | DataTypeNode::SimpleAggregateFunction(_, inner) + | DataTypeNode::Nullable(inner) => recursive_compatible(inner, compatible), + _ => false, + } +} + +pub struct ValueWriter<'a> { + pub(super) data_type: &'a DataTypeNode, + pub(super) layout: &'a mut LayoutBuilder, +} + +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum ValueWriteError { + #[error("attempting to use incorrect writer method for this type")] + IncorrectMethod, + + #[error("column does not allow nullable values here")] + UnexpectedNull, + + #[error("expected {expected} bytes, got {actual}")] + InvalidLength { expected: usize, actual: usize }, +} + +impl<'a> ValueWriter<'a> { + pub fn column_type(&self) -> &'a DataTypeNode { + self.data_type + } + + pub fn write_fixed(&mut self, bytes: &[u8]) -> Result<(), ValueWriteError> { + let LayoutBuilderKind::Fixed { + type_width, + ref mut data, + } = self.layout.kind + else { + return Err(ValueWriteError::IncorrectMethod); + }; + + if bytes.len() != type_width { + return Err(ValueWriteError::InvalidLength { + expected: type_width, + actual: bytes.len(), + }); + } + + data.extend_from_slice(bytes); + self.write_not_null(); + + Ok(()) + } + + pub fn write_string(&mut self, string_bytes: &[u8]) -> Result<(), ValueWriteError> { + let LayoutBuilderKind::Variable { end_offsets, data } = &mut self.layout.kind else { + return Err(ValueWriteError::IncorrectMethod); + }; + + data.extend_from_slice(string_bytes); + end_offsets.push(data.len()); + + self.write_not_null(); + + Ok(()) + } + + pub fn write_null(&mut self) -> Result<(), ValueWriteError> { + let nulls = self + .layout + .nulls + .as_mut() + .ok_or(ValueWriteError::UnexpectedNull)?; + + nulls.put_u8(1); + + self.layout.push_placeholder(); + + Ok(()) + } + + pub fn write_array(&mut self) -> Result, ArrayWriteError> + where + T: Encode, + { + let DataTypeNode::Array(elem_type) = &self.data_type else { + return Err(ArrayWriteError::NotAnArray { + data_type: self.data_type.clone(), + }); + }; + + if !T::compatible(elem_type) { + return Err(ArrayWriteError::IncompatibleType { + expected_type: (**elem_type).clone(), + }); + } + + let LayoutBuilderKind::Array { + end_indices, + elem_layout, + } = &mut self.layout.kind + else { + // Technically a bug if we reach this point + unreachable!("BUG: expected LayoutBuilderKind::Array") + }; + + Ok(ArrayWriter { + elem_type, + elem_layout, + end_indices, + outer_nulls: self.layout.nulls.as_mut(), + finished: false, + _marker: PhantomData, + }) + } + + pub fn write_tuple(&mut self) -> Result, ValueWriteError> { + let LayoutBuilderKind::Tuple { layouts } = &mut self.layout.kind else { + return Err(ValueWriteError::IncorrectMethod); + }; + + let DataTypeNode::Tuple(types) = &self.data_type else { + return Err(ValueWriteError::IncorrectMethod); + }; + + Ok(TupleWriter { + index: 0, + elem_layouts: layouts, + elem_types: types, + finished: false, + }) + } + + pub fn write_map(&mut self) -> Result, ValueWriteError> { + let LayoutBuilderKind::Map { + key_val_layouts, + end_indices, + } = &mut self.layout.kind + else { + return Err(ValueWriteError::IncorrectMethod); + }; + + let DataTypeNode::Map([key_ty, val_ty]) = &self.data_type else { + return Err(ValueWriteError::IncorrectMethod); + }; + + let [key_layout, val_layout] = &mut **key_val_layouts; + + Ok(MapWriter { + key_ty, + val_ty, + key_layout, + val_layout, + end_indices, + finished: false, + _marker: PhantomData, + }) + } + + fn write_not_null(&mut self) { + if let Some(nulls) = &mut self.layout.nulls { + nulls.put_u8(0); + } + } +} + +#[must_use = "rolls back the written tuple elements on-drop if `.finish()` is not called"] +pub struct ArrayWriter<'a, T> { + elem_type: &'a DataTypeNode, + elem_layout: &'a mut LayoutBuilder, + end_indices: &'a mut Vec, + outer_nulls: Option<&'a mut BytesMut>, + finished: bool, + _marker: PhantomData, +} + +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum ArrayWriteError { + #[error("attempted to write an array to a non-array column: {data_type}")] + NotAnArray { data_type: DataTypeNode }, + + #[error("value type is not compatible with expected type {expected_type}")] + IncompatibleType { expected_type: DataTypeNode }, + + #[error("error writing value at array index {index}")] + ValueWriteError { + index: usize, + #[source] + error: BoxedError, + }, +} + +impl ArrayWriter<'_, T> { + pub fn write(&mut self, value: T) -> Result<&mut Self, ArrayWriteError> + where + T: Encode, + { + value + .encode(&mut ValueWriter { + data_type: self.elem_type, + layout: self.elem_layout, + }) + .map_err(|error| ArrayWriteError::ValueWriteError { + index: self.written_len(), + error, + })?; + + Ok(self) + } + + pub fn finish(mut self) { + self.finish_mut() + } + + fn written_len(&self) -> usize { + let last_array_end = self.end_indices.last().copied().unwrap_or(0); + self.elem_layout.num_values().saturating_sub(last_array_end) + } + + fn finish_mut(&mut self) { + if self.finished { + return; + } + + self.end_indices.push(self.elem_layout.num_values()); + + if let Some(nulls) = &mut self.outer_nulls { + nulls.put_u8(0); + } + + self.finished = true; + } +} + +impl Drop for ArrayWriter<'_, T> { + fn drop(&mut self) { + if self.finished { + return; + } + + let last_array_end = self.end_indices.last().copied().unwrap_or(0); + self.elem_layout.truncate(last_array_end); + } +} + +#[must_use = "rolls back the written tuple elements on-drop if `.finish()` is not called"] +pub struct TupleWriter<'a> { + index: usize, + elem_layouts: &'a mut [LayoutBuilder], + elem_types: &'a [DataTypeNode], + finished: bool, +} + +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum TupleWriteError { + #[error("attempting to write to a full tuple")] + TupleFull, + + #[error( + "value type is not compatible with expected type {expected_type} at tuple index {index}" + )] + IncompatibleType { + index: usize, + expected_type: DataTypeNode, + }, + + #[error("error writing value at tuple index {index}")] + ValueWriteError { + index: usize, + #[source] + error: BoxedError, + }, +} + +#[derive(Debug, thiserror::Error)] +#[error("tuple not fully written; expected {expected_len} values, got {written_len}")] +pub struct IncompleteTupleError { + expected_len: usize, + written_len: usize, +} + +impl TupleWriter<'_> { + pub fn write(&mut self, value: T) -> Result<&mut Self, TupleWriteError> + where + T: Encode, + { + let data_type = self + .elem_types + .get(self.index) + .ok_or(TupleWriteError::TupleFull)?; + + if !T::compatible(data_type) { + return Err(TupleWriteError::IncompatibleType { + index: self.index, + expected_type: data_type.clone(), + }); + } + + value + .encode(&mut ValueWriter { + layout: &mut self.elem_layouts[self.index], + data_type, + }) + .map_err(|error| TupleWriteError::ValueWriteError { + index: self.index, + error, + })?; + + // Overflow here is likely to be a bug since + // `self.elem_types` would have to be `usize::MAX` long + self.index = self.index.checked_add(1).expect("tuple index overflowed"); + + Ok(self) + } + + pub fn finish(mut self) -> Result<(), IncompleteTupleError> { + if self.index < self.elem_types.len() { + return Err(IncompleteTupleError { + expected_len: self.elem_types.len(), + written_len: self.index, + }); + } + + self.finished = true; + + Ok(()) + } + + fn abort_mut(&mut self) { + if self.finished { + return; + } + + let written_len = cmp::min(self.index, self.elem_layouts.len()); + + for layout in &mut self.elem_layouts[..written_len] { + let len = layout.num_values(); + layout.truncate(len.saturating_sub(1)); + } + } +} + +impl Drop for TupleWriter<'_> { + fn drop(&mut self) { + self.abort_mut(); + } +} + +#[must_use = "rolls back the written map elements on-drop if `.finish()` is not called"] +pub struct MapWriter<'a, K, V> { + key_ty: &'a DataTypeNode, + val_ty: &'a DataTypeNode, + key_layout: &'a mut LayoutBuilder, + val_layout: &'a mut LayoutBuilder, + end_indices: &'a mut Vec, + finished: bool, + _marker: PhantomData, +} + +impl MapWriter<'_, K, V> +where + K: Encode, + V: Encode, +{ + pub fn write(&mut self, key: K, value: V) -> Result<&mut Self, MapWriteError> { + key.encode(&mut ValueWriter { + layout: self.key_layout, + data_type: self.key_ty, + }) + .map_err(|error| MapWriteError::ValueWriteError { + error, + index: self.key_layout.num_values(), + })?; + + value + .encode(&mut ValueWriter { + layout: self.val_layout, + data_type: self.val_ty, + }) + .map_err(|error| MapWriteError::ValueWriteError { + error, + index: self.val_layout.num_values(), + })?; + + Ok(self) + } + + pub fn finish(mut self) { + if self.finished { + return; + } + + let end_index = self.key_layout.num_values(); + self.end_indices.push(end_index); + + self.finished = true; + } +} + +impl Drop for MapWriter<'_, K, V> { + fn drop(&mut self) { + if self.finished { + return; + } + + let truncate_len = self.end_indices.last().copied().unwrap_or(0); + + self.key_layout.truncate(truncate_len); + self.val_layout.truncate(truncate_len); + } +} + +#[derive(Debug, thiserror::Error)] +#[non_exhaustive] +pub enum MapWriteError { + #[error("type is not compatible with expected type {expected_type} at entry index {index}")] + IncompatibleType { + index: usize, + expected_type: DataTypeNode, + }, + + #[error("error writing value at entry index {index}")] + ValueWriteError { + index: usize, + #[source] + error: BoxedError, + }, +} + +impl Encode for Option +where + T: Encode, +{ + fn produces() -> DataTypeNode { + DataTypeNode::Nullable(Box::new(T::produces())) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + match self { + Some(inner) => inner.encode(writer), + None => Ok(writer.write_null()?), + } + } + + fn compatible(column_type: &DataTypeNode) -> bool { + // Make sure we forward to `T::compatible()` + recursive_compatible(column_type, T::compatible) + } +} + +impl Encode for &'_ T +where + T: Encode + ?Sized, +{ + fn produces() -> DataTypeNode { + T::produces() + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + (**self).encode(writer) + } + + fn compatible(column_type: &DataTypeNode) -> bool { + T::compatible(column_type) + } +} + +macro_rules! impl_to_le_bytes { + ($($dataty:ident: $ty:ident),* $(,)?) => { + $( + impl Encode for $ty { + fn produces() -> DataTypeNode { + DataTypeNode::$dataty + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + Ok(writer.write_fixed(&self.to_le_bytes())?) + } + } + )* + }; +} + +// All scalar primitives are in little-endian +impl_to_le_bytes!( + // 8-bit ints don't have a concept of "endianness" but they still implement `to_le_bytes()` + // for the express purpose of being included in macros like this + Int8: i8, + Int16: i16, + Int32: i32, + Int64: i64, + Int128: i128, + UInt8: u8, + UInt16: u16, + UInt32: u32, + UInt64: u64, + UInt128: u128, + Float32: f32, + Float64: f64, +); + +impl Encode for bool { + fn produces() -> DataTypeNode { + DataTypeNode::Bool + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + writer.write_fixed(&[*self as u8])?; + Ok(()) + } +} + +impl Encode for str { + fn produces() -> DataTypeNode { + DataTypeNode::String + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + if let DataTypeNode::FixedString(fixed_len) = writer.data_type { + if self.len() != *fixed_len { + // Give a more informative error + return Err(format!( + "attempting to write a string of length {} to FixedString({fixed_len})", + self.len() + ) + .into()); + } + + writer.write_fixed(self.as_bytes())?; + } else { + writer.write_string(self.as_bytes())?; + } + + Ok(()) + } + + fn compatible(column_type: &DataTypeNode) -> bool { + recursive_compatible(column_type, |column_type| { + matches!( + column_type, + DataTypeNode::String | DataTypeNode::FixedString(_) + ) + }) + } +} + +impl Encode for String { + fn produces() -> DataTypeNode { + str::produces() + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + self.as_str().encode(writer) + } + + fn compatible(column_type: &DataTypeNode) -> bool { + str::compatible(column_type) + } +} + +impl Encode for [T] +where + T: Encode, +{ + fn produces() -> DataTypeNode { + DataTypeNode::Array(Box::new(T::produces())) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + let mut writer = writer.write_array()?; + + for val in self { + writer.write(val)?; + } + + writer.finish(); + + Ok(()) + } + + fn compatible(column_type: &DataTypeNode) -> bool { + recursive_compatible(column_type, |column_type| match column_type { + DataTypeNode::Array(elem_type) => T::compatible(elem_type), + _ => false, + }) + } +} + +impl Encode for Vec +where + T: Encode, +{ + fn produces() -> DataTypeNode { + <[T]>::produces() + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + self.as_slice().encode(writer) + } + + fn compatible(column_type: &DataTypeNode) -> bool { + <[T]>::compatible(column_type) + } +} + +macro_rules! tuple_impl { + ($var1:ident: $ty1:ident $(, $var:ident: $ty:ident)*) => { + impl<'a, $ty1 $(, $ty)* > Encode for ($ty1, $($ty),*) + where + $ty1: Encode, + $($ty: Encode,)* + { + fn produces() -> DataTypeNode { + DataTypeNode::Tuple(vec![$ty1::produces() $(, $ty::produces())*]) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + let mut writer = writer.write_tuple()?; + + let ($var1, $($var),*) = self; + + writer.write($var1)?; + $( + writer.write($var)?; + )* + + writer.finish()?; + Ok(()) + } + + fn compatible(column_type: &DataTypeNode) -> bool { + recursive_compatible(column_type, |column_type| { + let DataTypeNode::Tuple(types) = column_type else { + return false; + }; + + let [$var1, $($var),*] = &types[..] else { + return false; + }; + + $ty1::compatible($var1) + $(&& $ty::compatible($var))* + }) + } + } + + tuple_impl!($($var: $ty),*); + }; + () => {} +} + +tuple_impl!( + t1: T1, t2: T2, t3: T3, t4: T4, t5: T5, t6: T6, t7: T7, t8: T8, t9: T9, + t10: T10, t11: T11, t12: T12, t13: T13, t14: T14, t15: T15, t16: T16 +); + +impl Encode for Ipv4Addr { + fn produces() -> DataTypeNode { + DataTypeNode::IPv4 + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + writer.write_fixed(&self.to_bits().to_le_bytes())?; + Ok(()) + } +} + +impl Encode for Ipv6Addr { + fn produces() -> DataTypeNode { + DataTypeNode::IPv6 + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + writer.write_fixed(&self.octets())?; + Ok(()) + } +} + +impl Encode for HashMap +where + K: Encode, + V: Encode, +{ + fn produces() -> DataTypeNode { + DataTypeNode::Map([Box::new(K::produces()), Box::new(V::produces())]) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + let mut writer = writer.write_map()?; + + for (k, v) in self { + writer.write(k, v)?; + } + + writer.finish(); + + Ok(()) + } + + fn compatible(column_type: &DataTypeNode) -> bool { + recursive_compatible(column_type, |column_type| { + let DataTypeNode::Map([key_ty, val_ty]) = column_type else { + return false; + }; + + K::compatible(key_ty) && V::compatible(val_ty) + }) + } +} + +impl Encode for BTreeMap +where + K: Encode, + V: Encode, +{ + fn produces() -> DataTypeNode { + DataTypeNode::Map([Box::new(K::produces()), Box::new(V::produces())]) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + let mut writer = writer.write_map()?; + + for (k, v) in self { + writer.write(k, v)?; + } + + writer.finish(); + + Ok(()) + } + + fn compatible(column_type: &DataTypeNode) -> bool { + recursive_compatible(column_type, |column_type| { + let DataTypeNode::Map([key_ty, val_ty]) = column_type else { + return false; + }; + + K::compatible(key_ty) && V::compatible(val_ty) + }) + } +} + +#[cfg(feature = "uuid")] +mod uuid { + use super::{Encode, ValueWriter}; + use crate::error::BoxedError; + use clickhouse_types::DataTypeNode; + use uuid::Uuid; + + impl Encode for Uuid { + fn produces() -> DataTypeNode { + DataTypeNode::UUID + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + // https://clickhouse.com/docs/interfaces/specs/NativeFormat#uuid + // Wire bytes 0..7 = canonical bytes 0..7 reversed. + // Wire bytes 8..15 = canonical bytes 8..15 reversed. + let (lo_bytes, hi_bytes) = self.as_u64_pair(); + + let mut bytes = [0u8; 16]; + bytes[..8].copy_from_slice(&lo_bytes.to_le_bytes()); + bytes[8..].copy_from_slice(&hi_bytes.to_le_bytes()); + + writer.write_fixed(&bytes)?; + + Ok(()) + } + } +} diff --git a/src/native/mod.rs b/src/native/mod.rs index 5cf92e9e..690ea256 100644 --- a/src/native/mod.rs +++ b/src/native/mod.rs @@ -1,25 +1,32 @@ use crate::error::Error; use crate::native::string::MaybeUtf8; use bytes::Bytes; -use clickhouse_types::data_types::{DecimalType, EnumType}; use std::ops::Index; use hashbrown::HashMap; pub use array::{ArrayData, ArrayReader}; -pub use decode::Decode; pub use reader::BlockReadError; +use crate::native::decode::Decode; pub use clickhouse_types::DataTypeNode; pub(crate) mod array; -pub(crate) mod decode; +pub mod builder; +pub mod decode; +pub mod encode; pub(crate) mod reader; pub(crate) mod string; +mod utils; +mod varuint; +pub(crate) mod writer; + +#[cfg(test)] +mod tests; pub struct Block { column_names: HashMap, - columns: Vec, + columns: Box<[Column]>, num_rows: usize, } @@ -31,7 +38,7 @@ impl Block { .enumerate() .map(|(i, column)| (column.name.clone(), i)) .collect(), - columns, + columns: columns.into(), num_rows, } } @@ -185,67 +192,3 @@ impl<'a, T> ColumnIter<'a, T> { self.column } } - -fn type_fixed_width(data_type: &DataTypeNode) -> Option { - match data_type { - DataTypeNode::Bool => Some(1), - DataTypeNode::UInt8 => Some(1), - DataTypeNode::UInt16 => Some(2), - DataTypeNode::UInt32 => Some(4), - DataTypeNode::UInt64 => Some(8), - DataTypeNode::UInt128 => Some(16), - DataTypeNode::UInt256 => Some(32), - DataTypeNode::Int8 => Some(1), - DataTypeNode::Int16 => Some(2), - DataTypeNode::Int32 => Some(4), - DataTypeNode::Int64 => Some(8), - DataTypeNode::Int128 => Some(16), - DataTypeNode::Int256 => Some(32), - DataTypeNode::Float32 => Some(4), - DataTypeNode::Float64 => Some(8), - DataTypeNode::BFloat16 => Some(2), - DataTypeNode::Decimal(_, _, type_) => match type_ { - DecimalType::Decimal32 => Some(4), - DecimalType::Decimal64 => Some(8), - DecimalType::Decimal128 => Some(16), - DecimalType::Decimal256 => Some(32), - }, - DataTypeNode::String => None, - DataTypeNode::FixedString(len) => Some(*len), - DataTypeNode::UUID => Some(16), - DataTypeNode::Date => Some(2), - DataTypeNode::Date32 => Some(4), - DataTypeNode::DateTime(_) => Some(4), - DataTypeNode::DateTime64(_, _) => Some(8), - DataTypeNode::Time => Some(4), - DataTypeNode::Time64(_) => Some(8), - DataTypeNode::Interval(_) => Some(8), - DataTypeNode::IPv4 => Some(4), - DataTypeNode::IPv6 => Some(16), - // Nullable needs to be handled specially - DataTypeNode::Nullable(_) => None, - // Type width determined by metadata that comes before column data. - DataTypeNode::LowCardinality(_) => None, - DataTypeNode::Array(_) => None, - // Tuples are serialized column-by-column and need a structural layout. - DataTypeNode::Tuple(_) => None, - DataTypeNode::Enum(type_, _) => match type_ { - EnumType::Enum8 => Some(1), - EnumType::Enum16 => Some(2), - }, - DataTypeNode::Map(_) => None, - DataTypeNode::AggregateFunction(_, _) => None, - DataTypeNode::SimpleAggregateFunction(_, inner) => type_fixed_width(inner), - DataTypeNode::Variant(_) => None, - DataTypeNode::Dynamic => None, - DataTypeNode::JSON => None, - DataTypeNode::JsonWithHint(_) => None, - DataTypeNode::Point => Some(16), // Tuple(Float64, Float64) - DataTypeNode::Ring => None, - DataTypeNode::LineString => None, - DataTypeNode::MultiLineString => None, - DataTypeNode::Polygon => None, - DataTypeNode::MultiPolygon => None, - _ => None, - } -} diff --git a/src/native/reader.rs b/src/native/reader.rs index 008f335e..514a6010 100644 --- a/src/native/reader.rs +++ b/src/native/reader.rs @@ -1,6 +1,8 @@ use crate::error::Error; use crate::native::string::MaybeUtf8; -use crate::native::{Block, Column, Layout, LayoutKind, LayoutLowCardinality, type_fixed_width}; +use crate::native::utils::type_fixed_width; +use crate::native::varuint::ParseVarUInt; +use crate::native::{Block, Column, Layout, LayoutKind, LayoutLowCardinality}; use crate::response::Chunks; use bytes::{Buf, Bytes, BytesMut}; use clickhouse_types::DataTypeNode; @@ -465,7 +467,10 @@ impl ReaderInner { let mut parser = ParseVarUInt::default(); loop { - if let ControlFlow::Break(val) = parser.feed(&mut self.last_chunk)? { + if let ControlFlow::Break(val) = parser + .feed(&mut self.last_chunk) + .map_err(|e| read_error!("error parsing VarUInt").with_source(e))? + { return Ok(val); } @@ -624,43 +629,6 @@ impl ReaderInner { } } -#[derive(Default)] -struct ParseVarUInt { - accumulator: u64, - shift: u32, -} - -impl ParseVarUInt { - fn feed(&mut self, mut buf: impl Buf) -> Result, Error> { - const MAX_LEN: usize = 10; - - for _ in 0..MAX_LEN { - let Ok(b) = buf.try_get_u8() else { - return Ok(ControlFlow::Continue(())); - }; - - self.accumulator |= (b as u64 & 0x7F).checked_shl(self.shift).ok_or_else(|| { - read_error!( - "VarUInt repr overflowed: {:016x} byte: {b:02x}", - self.accumulator - ) - })?; - - if b <= 0x7F { - return Ok(ControlFlow::Break(self.accumulator)); - } - - self.shift += 7; - } - - Err(read_error!( - "terminating byte missing in VarUInt encoding: {:016x}", - self.accumulator, - ) - .into()) - } -} - impl LcKeyType { const MASK: u64 = 0xFF; // Low 8 bits @@ -780,79 +748,3 @@ impl ResultExt for Result { self.map_err(|e| e.with_column(name, data_type)) } } - -#[cfg(test)] -mod tests { - use super::ParseVarUInt; - use bytes::Buf; - use std::ops::ControlFlow; - - #[test] - fn parse_varuint() { - let encoded_and_decoded: &[(&[u8], u64)] = &[ - (&[0u8][..], 0u64), - (&[1], 1), - (&[127], 127), - (&[0x80, 0x01], 1 << 7), - (&[0x80, 0x80, 0x01], 1 << 14), - (&[0x80, 0x80, 0x80, 0x01], 1 << 21), - (&[0x80, 0x80, 0x80, 0x80, 0x01], 1 << 28), - (&[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x1F], 0xFF_FF_FF_FF_FF), - ]; - - let pad_len = 16usize; - - for (encoded, decoded) in encoded_and_decoded { - // Pad with junk data that must be ignored - let padded: Vec<_> = encoded - .iter() - .copied() - .chain((0..).cycle()) - .take(pad_len) - .collect(); - - // Test feeding slices in different size chunks - for chunk_size in 1..=padded.len() { - let mut parser = ParseVarUInt::default(); - - let mut slice = &padded[..]; - - let mut last_remaining = slice.len(); - - loop { - match parser.feed((&mut slice).take(chunk_size)) { - Ok(ControlFlow::Break(res)) => { - assert_eq!( - res, *decoded, - "invalid decoding; chunk_size: {chunk_size}, padded: {padded:?}, remaining: {slice:?}" - ); - assert_eq!( - slice.len(), - padded.len() - encoded.len(), - "extra data consumed: {slice:?}" - ); - break; - } - Ok(ControlFlow::Continue(())) => { - assert!( - !slice.is_empty(), - "full slice consumed without giving a result" - ); - assert_ne!( - slice.len(), - last_remaining, - "parser failed to make progress" - ); - last_remaining = slice.len(); - } - Err(e) => { - panic!( - "error: {e:?}, chunk_size: {chunk_size}, padded: {padded:?}, remaining: {slice:?}" - ); - } - } - } - } - } - } -} diff --git a/src/native/string.rs b/src/native/string.rs index a3e66254..02c9b96d 100644 --- a/src/native/string.rs +++ b/src/native/string.rs @@ -20,6 +20,10 @@ impl MaybeUtf8 { pub fn from_string(string: impl Into) -> Self { Self::from(string.into()) } + + pub fn len(&self) -> usize { + self.bytes.len() + } } impl From<&'static [u8]> for MaybeUtf8 { diff --git a/src/native/tests/builder.rs b/src/native/tests/builder.rs new file mode 100644 index 00000000..39f7cc38 --- /dev/null +++ b/src/native/tests/builder.rs @@ -0,0 +1,132 @@ +use crate::native::builder::{BlockBuilder, BlockBuilderError}; +use clickhouse_types::DataTypeNode; +use std::collections::BTreeMap; + +#[test] +fn forbids_incompatible_upsert() { + let mut builder = BlockBuilder::new(); + + builder.upsert_column::("foo").unwrap(); + + let Err(e) = builder.upsert_column::("foo") else { + panic!("expected error") + }; + + let BlockBuilderError::ColumnExists { + name, + existing_type, + new_type, + } = *e + else { + panic!("unexpected error variant: {e:?}") + }; + + assert_eq!(name, "foo"); + assert_eq!(existing_type, DataTypeNode::Int32); + assert_eq!(new_type, DataTypeNode::UInt32); +} + +#[test] +fn forbids_mismatched_lengths() { + let mut builder = BlockBuilder::new(); + + builder + .upsert_column::("foo") + .unwrap() + .add_all([0, 1, 2, 3, 4, 5]) + .unwrap(); + + builder + .upsert_column::("bar") + .unwrap() + .add_all([0, 1, 2, 3]) + .unwrap(); + + let Err(err) = builder.build() else { + panic!("expected error"); + }; + + let BlockBuilderError::MismatchedLengths { + longest_column, + longest_len, + shortest_column, + shortest_len, + } = *err + else { + panic!("unexpected error variant: {err:?}") + }; + + assert_eq!(longest_column, "foo"); + assert_eq!(longest_len, 6); + + assert_eq!(shortest_column, "bar"); + assert_eq!(shortest_len, 4); +} + +#[test] +fn debug() { + let mut builder = BlockBuilder::new(); + + assert_eq!(format!("{builder:?}"), "BlockBuilder { columns: [] }"); + + builder + .upsert_column::("foo") + .unwrap() + .add_all([0, 1, 2, 3, 4]) + .unwrap(); + + // Saves us having to write these out by hand + insta::assert_debug_snapshot!(builder); + + builder + .upsert_column::<&str>("bar") + .unwrap() + .add_all(["lorem", "ipsum", "dolor", "sit", "amet"]) + .unwrap(); + + insta::assert_debug_snapshot!(builder); + + builder + .upsert_column::<&[u64]>("baz") + .unwrap() + .add_all([ + &[][..], + &[0], + &[0, 1], + &[0, 1, 2], + &[0, 1, 2, 3], + &[0, 1, 2, 3, 4], + ]) + .unwrap(); + + insta::assert_debug_snapshot!(builder); + + builder + .upsert_column("quux") + .unwrap() + .add_all((0u32..5).map(|i| (i, i.to_string()))) + .unwrap(); + + insta::assert_debug_snapshot!(builder); + + builder + .upsert_column("foobar") + .unwrap() + // `HashMap` order is not deterministic + .add_all((0i64..5).map(|i| { + (0..i) + .map(|j| (j, j.to_string())) + .collect::>() + })) + .unwrap(); + + insta::assert_debug_snapshot!(builder); + + builder + .upsert_column("foo_with_nulls") + .unwrap() + .add_all((0..5).map(|i| (i % 2 != 0).then_some(i))) + .unwrap(); + + insta::assert_debug_snapshot!(builder); +} diff --git a/src/native/tests/encode.rs b/src/native/tests/encode.rs new file mode 100644 index 00000000..84fa4f92 --- /dev/null +++ b/src/native/tests/encode.rs @@ -0,0 +1,324 @@ +use crate::error::BoxedError; +use crate::native::builder::{BlockBuilder, BlockBuilderError}; +use crate::native::encode::{Encode, ValueWriter}; +use clickhouse_types::DataTypeNode; + +use std::collections::HashMap; +use std::mem; + +#[test] +fn array_writer_rolls_back() { + struct BadArray<'a>(&'a [u32]); + + impl Encode for BadArray<'_> { + fn produces() -> DataTypeNode { + DataTypeNode::Array(Box::new(DataTypeNode::UInt32)) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + let mut writer = writer.write_array()?; + + for val in self.0 { + writer.write(val)?; + } + + // Deliberately don't call `writer.finish()` + drop(writer); + + Ok(()) + } + } + + let mut builder = BlockBuilder::new(); + + builder + .upsert_column("foo") + .unwrap() + .add(&[0u32, 1, 2, 3, 4, 5][..]) + .unwrap(); + + builder + .upsert_column("foo") + .unwrap() + .add(BadArray(&[6, 7, 8, 9, 10])) + .unwrap(); + + let block = builder.build().unwrap(); + + assert_eq!(block.num_rows(), 1); +} + +#[test] +fn tuple_writer_rolls_back() { + struct BadTuple(i32, String, #[expect(dead_code)] Vec); + + impl Encode for BadTuple { + fn produces() -> DataTypeNode { + DataTypeNode::Tuple(vec![ + DataTypeNode::Int32, + DataTypeNode::String, + DataTypeNode::Array(Box::new(DataTypeNode::Int64)), + ]) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + let mut writer = writer.write_tuple()?; + + writer.write(self.0)?; + writer.write(&self.1)?; + + // Deliberately don't finish, this would put the block out of sync + drop(writer); + + Ok(()) + } + } + + let mut builder = BlockBuilder::new(); + + builder + .upsert_column("foo") + .unwrap() + .add((0i32, "0".to_string(), vec![0i64; 16])) + .unwrap(); + + builder + .upsert_column("foo") + .unwrap() + .add(BadTuple(1, "1".to_string(), vec![1i64; 16])) + .unwrap(); + + let block = builder.build().unwrap(); + + assert_eq!(block.num_rows(), 1); +} + +#[test] +fn map_writer_rolls_back() { + struct BadMap(Vec<(u32, String)>); + + impl Encode for BadMap { + fn produces() -> DataTypeNode { + DataTypeNode::Map([ + Box::new(DataTypeNode::UInt32), + Box::new(DataTypeNode::String), + ]) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + let mut writer = writer.write_map()?; + + for (key, val) in &self.0 { + writer.write(key, val)?; + } + + drop(writer); + + Ok(()) + } + } + + let mut builder = BlockBuilder::new(); + + builder + .upsert_column("foo") + .unwrap() + .add( + (0u32..5) + .map(|i| (i, i.to_string())) + .collect::>(), + ) + .unwrap(); + + builder + .upsert_column("foo") + .unwrap() + .add(BadMap((5..10).map(|i| (i, i.to_string())).collect())) + .unwrap(); + + let block = builder.build().unwrap(); + + assert_eq!(block.num_rows(), 1); +} + +// A leak writer could put a block out of sync and result in data corruption; +// it's better if we catch it during validation. +#[test] +fn leaked_array_writer_fails_validation() { + struct LeakWriter<'a>(&'a [u32]); + + impl Encode for LeakWriter<'_> { + fn produces() -> DataTypeNode { + DataTypeNode::Array(Box::new(DataTypeNode::UInt32)) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + let mut writer = writer.write_array()?; + + for val in self.0 { + writer.write(val)?; + } + + // Deliberately leak the writer to put the block out of sync + mem::forget(writer); + + Ok(()) + } + } + + let mut builder = BlockBuilder::new(); + + builder + .upsert_column("foo") + .unwrap() + .add(&[0u32, 1, 2, 3, 4, 5][..]) + .unwrap(); + + builder + .upsert_column("foo") + .unwrap() + .add(LeakWriter(&[6, 7, 8, 9, 10])) + .unwrap(); + + let err = builder.build().err().expect("expected block builder error"); + + let BlockBuilderError::ColumnDataInvalid { + column_name, + column_type, + message, + } = *err + else { + panic!("unexpected error kind: {err}"); + }; + + assert_eq!(column_name, "foo"); + assert_eq!(column_type, <[u32] as Encode>::produces()); + + assert_eq!( + message, + "last array index (6) out of sync with total elements: 11" + ); +} + +#[test] +fn leaked_tuple_writer_fails_validation() { + struct LeakWriter(i32, String, #[expect(dead_code)] Vec); + + impl Encode for LeakWriter { + fn produces() -> DataTypeNode { + DataTypeNode::Tuple(vec![ + DataTypeNode::Int32, + DataTypeNode::String, + DataTypeNode::Array(Box::new(DataTypeNode::Int64)), + ]) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + let mut writer = writer.write_tuple()?; + + writer.write(self.0)?; + writer.write(&self.1)?; + + // Deliberately leak writer, this would put the block out of sync + mem::forget(writer); + + Ok(()) + } + } + + let mut builder = BlockBuilder::new(); + + builder + .upsert_column("foo") + .unwrap() + .add((0i32, "0".to_string(), vec![0i64; 16])) + .unwrap(); + + builder + .upsert_column("foo") + .unwrap() + .add(LeakWriter(1, "1".to_string(), vec![1i64; 16])) + .unwrap(); + + let err = builder.build().err().expect("expected block builder error"); + + let BlockBuilderError::ColumnDataInvalid { + column_name, + column_type, + message, + } = *err + else { + panic!("unexpected error kind: {err}"); + }; + + assert_eq!(column_name, "foo"); + assert_eq!(column_type, <(i32, String, Vec) as Encode>::produces()); + + assert_eq!( + message, + "tuple index 2 (type Array(Int64)) total elements out of sync: 1 vs 2" + ); +} + +#[test] +fn leaked_map_writer_fails_validation() { + struct LeakWriter(Vec<(u32, String)>); + + impl Encode for LeakWriter { + fn produces() -> DataTypeNode { + DataTypeNode::Map([ + Box::new(DataTypeNode::UInt32), + Box::new(DataTypeNode::String), + ]) + } + + fn encode(&self, writer: &mut ValueWriter<'_>) -> Result<(), BoxedError> { + let mut writer = writer.write_map()?; + + for (key, val) in &self.0 { + writer.write(key, val)?; + } + + mem::forget(writer); + + Ok(()) + } + } + + let mut builder = BlockBuilder::new(); + + builder + .upsert_column("foo") + .unwrap() + .add( + (0u32..5) + .map(|i| (i, i.to_string())) + .collect::>(), + ) + .unwrap(); + + builder + .upsert_column("foo") + .unwrap() + .add(LeakWriter((5..10).map(|i| (i, i.to_string())).collect())) + .unwrap(); + + let err = builder.build().err().expect("expected block builder error"); + + let BlockBuilderError::ColumnDataInvalid { + column_name, + column_type, + message, + } = *err + else { + panic!("unexpected error kind: {err}"); + }; + + assert_eq!(column_name, "foo"); + assert_eq!(column_type, as Encode>::produces()); + + assert_eq!( + message, + "last map index (5) out of sync with total elements: 10" + ); +} diff --git a/src/native/tests/mod.rs b/src/native/tests/mod.rs new file mode 100644 index 00000000..60868466 --- /dev/null +++ b/src/native/tests/mod.rs @@ -0,0 +1,2 @@ +mod builder; +mod encode; diff --git a/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-2.snap b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-2.snap new file mode 100644 index 00000000..3558ddfb --- /dev/null +++ b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-2.snap @@ -0,0 +1,41 @@ +--- +source: src/native/tests/builder.rs +assertion_line: 76 +expression: builder +--- +BlockBuilder { + columns: [ + ColumnBuilderRaw { + name: "foo", + data_type: Int32, + layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x00000000, + 0x01000000, + 0x02000000, + 0x03000000, + 0x04000000, + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "bar", + data_type: String, + layout: LayoutBuilder { + kind: Variable { + data: [ + "lorem", + "ipsum", + "dolor", + "sit", + "amet", + ], + }, + nulls: None, + }, + }, + ], +} diff --git a/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-3.snap b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-3.snap new file mode 100644 index 00000000..743030da --- /dev/null +++ b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-3.snap @@ -0,0 +1,82 @@ +--- +source: src/native/tests/builder.rs +assertion_line: 84 +expression: builder +--- +BlockBuilder { + columns: [ + ColumnBuilderRaw { + name: "foo", + data_type: Int32, + layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x00000000, + 0x01000000, + 0x02000000, + 0x03000000, + 0x04000000, + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "bar", + data_type: String, + layout: LayoutBuilder { + kind: Variable { + data: [ + "lorem", + "ipsum", + "dolor", + "sit", + "amet", + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "baz", + data_type: Array( + UInt64, + ), + layout: LayoutBuilder { + kind: Array { + elem_layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x0000000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0300000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0300000000000000, + 0x0400000000000000, + ], + }, + nulls: None, + }, + end_indices: [ + 0, + 1, + 3, + 6, + 10, + 15, + ], + }, + nulls: None, + }, + }, + ], +} diff --git a/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-4.snap b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-4.snap new file mode 100644 index 00000000..ee4aa806 --- /dev/null +++ b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-4.snap @@ -0,0 +1,120 @@ +--- +source: src/native/tests/builder.rs +assertion_line: 92 +expression: builder +--- +BlockBuilder { + columns: [ + ColumnBuilderRaw { + name: "foo", + data_type: Int32, + layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x00000000, + 0x01000000, + 0x02000000, + 0x03000000, + 0x04000000, + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "bar", + data_type: String, + layout: LayoutBuilder { + kind: Variable { + data: [ + "lorem", + "ipsum", + "dolor", + "sit", + "amet", + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "baz", + data_type: Array( + UInt64, + ), + layout: LayoutBuilder { + kind: Array { + elem_layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x0000000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0300000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0300000000000000, + 0x0400000000000000, + ], + }, + nulls: None, + }, + end_indices: [ + 0, + 1, + 3, + 6, + 10, + 15, + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "quux", + data_type: Tuple( + [ + UInt32, + String, + ], + ), + layout: LayoutBuilder { + kind: Tuple( + LayoutBuilder { + kind: Fixed { + data: [ + 0x00000000, + 0x01000000, + 0x02000000, + 0x03000000, + 0x04000000, + ], + }, + nulls: None, + }, + LayoutBuilder { + kind: Variable { + data: [ + "0", + "1", + "2", + "3", + "4", + ], + }, + nulls: None, + }, + ), + nulls: None, + }, + }, + ], +} diff --git a/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-5.snap b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-5.snap new file mode 100644 index 00000000..f0edbfa4 --- /dev/null +++ b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-5.snap @@ -0,0 +1,175 @@ +--- +source: src/native/tests/builder.rs +assertion_line: 101 +expression: builder +--- +BlockBuilder { + columns: [ + ColumnBuilderRaw { + name: "foo", + data_type: Int32, + layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x00000000, + 0x01000000, + 0x02000000, + 0x03000000, + 0x04000000, + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "bar", + data_type: String, + layout: LayoutBuilder { + kind: Variable { + data: [ + "lorem", + "ipsum", + "dolor", + "sit", + "amet", + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "baz", + data_type: Array( + UInt64, + ), + layout: LayoutBuilder { + kind: Array { + elem_layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x0000000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0300000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0300000000000000, + 0x0400000000000000, + ], + }, + nulls: None, + }, + end_indices: [ + 0, + 1, + 3, + 6, + 10, + 15, + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "quux", + data_type: Tuple( + [ + UInt32, + String, + ], + ), + layout: LayoutBuilder { + kind: Tuple( + LayoutBuilder { + kind: Fixed { + data: [ + 0x00000000, + 0x01000000, + 0x02000000, + 0x03000000, + 0x04000000, + ], + }, + nulls: None, + }, + LayoutBuilder { + kind: Variable { + data: [ + "0", + "1", + "2", + "3", + "4", + ], + }, + nulls: None, + }, + ), + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "foobar", + data_type: Map( + [ + Int64, + String, + ], + ), + layout: LayoutBuilder { + kind: Map { + keys: LayoutBuilder { + kind: Fixed { + data: [ + 0x0000000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0300000000000000, + ], + }, + nulls: None, + }, + values: LayoutBuilder { + kind: Variable { + data: [ + "0", + "0", + "1", + "0", + "1", + "2", + "0", + "1", + "2", + "3", + ], + }, + nulls: None, + }, + end_indices: [ + 0, + 1, + 3, + 6, + 10, + ], + }, + nulls: None, + }, + }, + ], +} diff --git a/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-6.snap b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-6.snap new file mode 100644 index 00000000..f319b72c --- /dev/null +++ b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug-6.snap @@ -0,0 +1,195 @@ +--- +source: src/native/tests/builder.rs +assertion_line: 109 +expression: builder +--- +BlockBuilder { + columns: [ + ColumnBuilderRaw { + name: "foo", + data_type: Int32, + layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x00000000, + 0x01000000, + 0x02000000, + 0x03000000, + 0x04000000, + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "bar", + data_type: String, + layout: LayoutBuilder { + kind: Variable { + data: [ + "lorem", + "ipsum", + "dolor", + "sit", + "amet", + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "baz", + data_type: Array( + UInt64, + ), + layout: LayoutBuilder { + kind: Array { + elem_layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x0000000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0300000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0300000000000000, + 0x0400000000000000, + ], + }, + nulls: None, + }, + end_indices: [ + 0, + 1, + 3, + 6, + 10, + 15, + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "quux", + data_type: Tuple( + [ + UInt32, + String, + ], + ), + layout: LayoutBuilder { + kind: Tuple( + LayoutBuilder { + kind: Fixed { + data: [ + 0x00000000, + 0x01000000, + 0x02000000, + 0x03000000, + 0x04000000, + ], + }, + nulls: None, + }, + LayoutBuilder { + kind: Variable { + data: [ + "0", + "1", + "2", + "3", + "4", + ], + }, + nulls: None, + }, + ), + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "foobar", + data_type: Map( + [ + Int64, + String, + ], + ), + layout: LayoutBuilder { + kind: Map { + keys: LayoutBuilder { + kind: Fixed { + data: [ + 0x0000000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0000000000000000, + 0x0100000000000000, + 0x0200000000000000, + 0x0300000000000000, + ], + }, + nulls: None, + }, + values: LayoutBuilder { + kind: Variable { + data: [ + "0", + "0", + "1", + "0", + "1", + "2", + "0", + "1", + "2", + "3", + ], + }, + nulls: None, + }, + end_indices: [ + 0, + 1, + 3, + 6, + 10, + ], + }, + nulls: None, + }, + }, + ColumnBuilderRaw { + name: "foo_with_nulls", + data_type: Nullable( + Int32, + ), + layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x00000000, + 0x01000000, + 0x00000000, + 0x03000000, + 0x00000000, + ], + }, + nulls: Some( + 10101, + ), + }, + }, + ], +} diff --git a/src/native/tests/snapshots/clickhouse__native__tests__builder__debug.snap b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug.snap new file mode 100644 index 00000000..020c044f --- /dev/null +++ b/src/native/tests/snapshots/clickhouse__native__tests__builder__debug.snap @@ -0,0 +1,25 @@ +--- +source: src/native/tests/builder.rs +assertion_line: 68 +expression: builder +--- +BlockBuilder { + columns: [ + ColumnBuilderRaw { + name: "foo", + data_type: Int32, + layout: LayoutBuilder { + kind: Fixed { + data: [ + 0x00000000, + 0x01000000, + 0x02000000, + 0x03000000, + 0x04000000, + ], + }, + nulls: None, + }, + }, + ], +} diff --git a/src/native/utils.rs b/src/native/utils.rs new file mode 100644 index 00000000..a23936a3 --- /dev/null +++ b/src/native/utils.rs @@ -0,0 +1,137 @@ +use clickhouse_types::DataTypeNode; +use clickhouse_types::data_types::{DecimalType, EnumType}; +use std::fmt::{Debug, Formatter, Write}; + +pub(super) struct DebugNullMap<'a>(pub &'a [u8]); + +impl Debug for DebugNullMap<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + for b in self.0 { + match b { + 0 => f.write_char('0')?, + 1 => f.write_char('1')?, + // Flag invalid bytes with brackets + _ => write!(f, "{{{b:x}}}")?, + } + } + + Ok(()) + } +} + +pub(super) struct DebugFixedData<'a> { + pub(super) type_width: usize, + pub(super) data: &'a [u8], +} + +pub(super) struct DebugHex<'a>(&'a [u8]); + +impl Debug for DebugFixedData<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_list() + .entries(self.data.chunks(self.type_width).map(DebugHex)) + .finish() + } +} + +impl Debug for DebugHex<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.write_str("0x")?; + + // Write the data as it appears so there's no confusion; + // for little-endian values this may be in reverse + for b in self.0.iter() { + write!(f, "{b:02x}")?; + } + + Ok(()) + } +} + +pub(super) struct DebugVariableData<'a> { + pub end_offsets: &'a [usize], + pub data: &'a [u8], +} + +impl Debug for DebugVariableData<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + let mut list = f.debug_list(); + + let mut start_offset = 0; + + for &end_offset in self.end_offsets { + // `impl Debug for fmt::Arguments` forwards to `Display` + list.entry(&format_args!( + "\"{}\"", + &self.data[start_offset..end_offset].escape_ascii() + )); + start_offset = end_offset; + } + + list.finish() + } +} + +pub(super) fn type_fixed_width(data_type: &DataTypeNode) -> Option { + match data_type { + DataTypeNode::Bool => Some(1), + DataTypeNode::UInt8 => Some(1), + DataTypeNode::UInt16 => Some(2), + DataTypeNode::UInt32 => Some(4), + DataTypeNode::UInt64 => Some(8), + DataTypeNode::UInt128 => Some(16), + DataTypeNode::UInt256 => Some(32), + DataTypeNode::Int8 => Some(1), + DataTypeNode::Int16 => Some(2), + DataTypeNode::Int32 => Some(4), + DataTypeNode::Int64 => Some(8), + DataTypeNode::Int128 => Some(16), + DataTypeNode::Int256 => Some(32), + DataTypeNode::Float32 => Some(4), + DataTypeNode::Float64 => Some(8), + DataTypeNode::BFloat16 => Some(2), + DataTypeNode::Decimal(_, _, type_) => match type_ { + DecimalType::Decimal32 => Some(4), + DecimalType::Decimal64 => Some(8), + DecimalType::Decimal128 => Some(16), + DecimalType::Decimal256 => Some(32), + }, + DataTypeNode::String => None, + DataTypeNode::FixedString(len) => Some(*len), + DataTypeNode::UUID => Some(16), + DataTypeNode::Date => Some(2), + DataTypeNode::Date32 => Some(4), + DataTypeNode::DateTime(_) => Some(4), + DataTypeNode::DateTime64(_, _) => Some(8), + DataTypeNode::Time => Some(4), + DataTypeNode::Time64(_) => Some(8), + DataTypeNode::Interval(_) => Some(8), + DataTypeNode::IPv4 => Some(4), + DataTypeNode::IPv6 => Some(16), + // Nullable needs to be handled specially + DataTypeNode::Nullable(_) => None, + // Type width determined by metadata that comes before column data. + DataTypeNode::LowCardinality(_) => None, + DataTypeNode::Array(_) => None, + // Tuples are serialized column-by-column and need a structural layout. + DataTypeNode::Tuple(_) => None, + DataTypeNode::Enum(type_, _) => match type_ { + EnumType::Enum8 => Some(1), + EnumType::Enum16 => Some(2), + }, + DataTypeNode::Map(_) => None, + DataTypeNode::AggregateFunction(_, _) => None, + DataTypeNode::SimpleAggregateFunction(_, inner) => type_fixed_width(inner), + DataTypeNode::Variant(_) => None, + DataTypeNode::Dynamic => None, + DataTypeNode::JSON => None, + DataTypeNode::JsonWithHint(_) => None, + DataTypeNode::Point => Some(16), // Tuple(Float64, Float64) + DataTypeNode::Ring => None, + DataTypeNode::LineString => None, + DataTypeNode::MultiLineString => None, + DataTypeNode::Polygon => None, + DataTypeNode::MultiPolygon => None, + _ => None, + } +} diff --git a/src/native/varuint.rs b/src/native/varuint.rs new file mode 100644 index 00000000..b7d15bb6 --- /dev/null +++ b/src/native/varuint.rs @@ -0,0 +1,149 @@ +use bytes::{Buf, BufMut}; +use std::ops::ControlFlow; + +/// Write a `usize` as a native `VarUInt` +#[allow(clippy::cast_possible_truncation)] // truncation is intentional here +pub(super) fn write(mut buf: impl BufMut, mut uint: usize) { + while uint > 0x7F { + let b = (uint as u8) | 0x80; + buf.put_u8(b); + uint >>= 7; + } + + buf.put_u8(uint as u8); +} + +#[derive(Default)] +pub(super) struct ParseVarUInt { + accumulator: u64, + shift: u32, +} + +#[derive(Debug, thiserror::Error)] +pub(super) enum ParseVarUIntError { + #[error("VarUInt repr overflowed: {accumulator:#x} byte: {byte:#02x}")] + Overflow { accumulator: u64, byte: u8 }, + #[error("terminating byte missing in VarUInt encoding: {accumulator:#x}")] + MissingTerminator { accumulator: u64 }, +} + +impl ParseVarUInt { + pub(super) fn feed( + &mut self, + mut buf: impl Buf, + ) -> Result, ParseVarUIntError> { + const MAX_LEN: usize = 10; + + for _ in 0..MAX_LEN { + let Ok(b) = buf.try_get_u8() else { + return Ok(ControlFlow::Continue(())); + }; + + self.accumulator |= + (b as u64 & 0x7F) + .checked_shl(self.shift) + .ok_or(ParseVarUIntError::Overflow { + accumulator: self.accumulator, + byte: b, + })?; + + if b <= 0x7F { + return Ok(ControlFlow::Break(self.accumulator)); + } + + self.shift += 7; + } + + Err(ParseVarUIntError::MissingTerminator { + accumulator: self.accumulator, + }) + } +} + +#[cfg(test)] +mod tests { + use crate::native::varuint::ParseVarUInt; + use bytes::Buf; + use std::ops::ControlFlow; + const ENCODED_AND_DECODED: &[(&[u8], u64)] = &[ + (&[0], 0u64), + (&[1], 1), + (&[127], 127), + (&[0x80, 0x01], 1 << 7), + (&[0x80, 0x80, 0x01], 1 << 14), + (&[0x80, 0x80, 0x80, 0x01], 1 << 21), + (&[0x80, 0x80, 0x80, 0x80, 0x01], 1 << 28), + (&[0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0x1F], 0xFF_FF_FF_FF_FF), + ]; + + #[test] + fn parse_varuint() { + let pad_len = 16usize; + + for (encoded, decoded) in ENCODED_AND_DECODED { + // Pad with junk data that must be ignored + let padded: Vec<_> = encoded + .iter() + .copied() + .chain((0..).cycle()) + .take(pad_len) + .collect(); + + // Test feeding slices in different size chunks + for chunk_size in 1..=padded.len() { + let mut parser = ParseVarUInt::default(); + + let mut slice = &padded[..]; + + let mut last_remaining = slice.len(); + + loop { + match parser.feed((&mut slice).take(chunk_size)) { + Ok(ControlFlow::Break(res)) => { + assert_eq!( + res, *decoded, + "invalid decoding; chunk_size: {chunk_size}, padded: {padded:?}, remaining: {slice:?}" + ); + assert_eq!( + slice.len(), + padded.len() - encoded.len(), + "extra data consumed: {slice:?}" + ); + break; + } + Ok(ControlFlow::Continue(())) => { + assert!( + !slice.is_empty(), + "full slice consumed without giving a result" + ); + assert_ne!( + slice.len(), + last_remaining, + "parser failed to make progress" + ); + last_remaining = slice.len(); + } + Err(e) => { + panic!( + "error: {e:?}, chunk_size: {chunk_size}, padded: {padded:?}, remaining: {slice:?}" + ); + } + } + } + } + } + } + + #[test] + fn write_varuint() { + let mut buf = Vec::with_capacity(10); + + for (encoded, decoded) in ENCODED_AND_DECODED { + buf.clear(); + + super::write(&mut buf, (*decoded).try_into().unwrap()); + + assert_eq!(*buf, **encoded, "decoded: {decoded}"); + } + } +} diff --git a/src/native/writer.rs b/src/native/writer.rs new file mode 100644 index 00000000..f4c68fa0 --- /dev/null +++ b/src/native/writer.rs @@ -0,0 +1,178 @@ +use crate::error::Error; +use crate::insert_formatted::InsertFormatted; +use crate::native::{Block, Column, Layout, LayoutKind, varuint}; +use bytes::{BufMut, Bytes, BytesMut}; + +pub(crate) struct BlockWriter { + insert: InsertFormatted, + buf: BytesMut, +} + +impl BlockWriter { + pub(crate) fn new(insert: InsertFormatted) -> Self { + Self { + insert, + buf: BytesMut::with_capacity(8192), + } + } + + pub(crate) async fn write(&mut self, block: &Block) -> Result<(), Error> { + // If canceled while writing a block, we have no way to recover. + // We have to abort the request instead. + let mut guard = WriteGuard { + insert: &mut self.insert, + buf: &mut self.buf, + finished: false, + }; + + varuint::write(&mut guard.buf, block.columns.len()); + varuint::write(&mut guard.buf, block.num_rows); + + for column in &block.columns { + guard.write_column(column).await?; + } + + // We deliberately don't send the block header right away + // since the first column header can share the same buffer. + // + // If `write_column()` wrote any data, this should be a no-op. + guard.flush().await?; + + guard.finished = true; + + Ok(()) + } + + pub(crate) async fn end(mut self) -> Result<(), Error> { + let mut guard = WriteGuard { + insert: &mut self.insert, + buf: &mut self.buf, + finished: false, + }; + + guard.flush().await?; + + guard.finished = true; + + drop(guard); + + self.insert.end().await + } +} + +struct WriteGuard<'a> { + insert: &'a mut InsertFormatted, + buf: &'a mut BytesMut, + finished: bool, +} + +impl WriteGuard<'_> { + async fn write_column(&mut self, column: &Column) -> Result<(), Error> { + varuint::write(&mut self.buf, column.name.len()); + self.buf.extend_from_slice(column.name.as_bytes()); + + // We have to format the data type to know exactly how long the string is, + // though in most cases we can use a static string. + let data_type_str = column.data_type.to_str(); + + varuint::write(&mut self.buf, data_type_str.len()); + self.buf.extend_from_slice(data_type_str.as_bytes()); + + self.write_layout(&column.layout).await + } + + async fn write_layout(&mut self, layout: &Layout) -> Result<(), Error> { + if let Some(nulls) = &layout.nulls { + self.send(nulls).await?; + } + + match &layout.kind { + LayoutKind::Fixed { data, .. } => { + self.send(data).await?; + } + LayoutKind::Variable { end_offsets, data } => { + let mut start_offset = 0; + + // Convert from offsets *back* to lengths and subslices + for &end_offset in end_offsets { + let len = end_offset + .checked_sub(start_offset) + .ok_or_else(|| Error::Other(format!("BUG: string length underflow in encoding block: {end_offset} - {start_offset}").into()))?; + + varuint::write(&mut self.buf, len); + + self.send(&data.slice(start_offset..end_offset)).await?; + + start_offset = end_offset; + } + } + LayoutKind::LowCardinality(_) => { + return Err(Error::Other( + "inserting LowCardinality data not yet implemented".into(), + )); + } + LayoutKind::Array { + end_indices, + elem_layout, + } => { + for index in end_indices { + self.buf.put_u64_le(u64::try_from(*index).map_err(|_| { + Error::Other(format!("array end index out of range: {index}").into()) + })?); + } + + Box::pin(self.write_layout(elem_layout)).await?; + } + LayoutKind::Tuple { layouts } => { + for layout in layouts { + Box::pin(self.write_layout(layout)).await?; + } + } + LayoutKind::Map { + key_val_layouts, + end_indices, + } => { + for index in end_indices { + self.buf.put_u64_le(u64::try_from(*index).map_err(|_| { + Error::Other(format!("map end index out of range: {index}").into()) + })?); + } + + Box::pin(self.write_layout(&key_val_layouts[0])).await?; + Box::pin(self.write_layout(&key_val_layouts[1])).await?; + } + } + + Ok(()) + } + + async fn send(&mut self, data: &Bytes) -> Result<(), Error> { + /// If a data buffer is smaller than this threshold, copy it to `self.buf` instead. + const COPY_THRESHOLD: usize = 128; + + if data.len() < COPY_THRESHOLD { + self.buf.extend_from_slice(data); + Ok(()) + } else { + self.flush().await?; + + self.insert.send(data.clone()).await + } + } + + async fn flush(&mut self) -> Result<(), Error> { + if !self.buf.is_empty() { + self.insert.send(self.buf.split().freeze()).await?; + } + + Ok(()) + } +} + +impl Drop for WriteGuard<'_> { + fn drop(&mut self) { + if !self.finished { + self.insert.abort(); + } + } +} diff --git a/tests/it/fetch_native.rs b/tests/it/fetch_native.rs index 140854b7..6c8691f5 100644 --- a/tests/it/fetch_native.rs +++ b/tests/it/fetch_native.rs @@ -1,5 +1,5 @@ use crate::get_client; -use clickhouse::native::{Column, Decode}; +use clickhouse::native::{Column, decode::Decode}; use std::collections::{BTreeMap, HashMap}; use std::fmt::Debug; @@ -23,41 +23,6 @@ async fn mixed_types_1000() { mixed_types(1000).await } -#[tokio::test] -async fn fixed_width_tuples() { - let client = get_client(); - - let mut cursor = client - .query( - "SELECT - tuple(toUInt64(number), toUInt32(number + 1)) AS fixed_tuple, - arrayMap( - x -> tuple(toUInt64(x), toUInt32(x + 1)), - range(number) - ) AS fixed_tuple_array - FROM system.numbers - LIMIT 3", - ) - .fetch_native() - .unwrap(); - - let block = cursor.next().await.unwrap().expect("expected one block"); - - let tuples = block["fixed_tuple"] - .iter::<(u64, u32)>() - .unwrap() - .collect::, _>>() - .unwrap(); - assert_eq!(tuples, [(0, 1), (1, 2), (2, 3)]); - - let tuple_arrays = block["fixed_tuple_array"] - .iter::>() - .unwrap() - .collect::, _>>() - .unwrap(); - assert_eq!(tuple_arrays, [vec![], vec![(0, 1)], vec![(0, 1), (1, 2)]]); -} - // NOTE: requires >50GiB RAM on the server, likely because all the strings have to live in memory #[ignore] #[tokio::test] @@ -492,3 +457,38 @@ async fn nested_arrays_100() { async fn nested_arrays_1000() { nested_arrays(1000).await; } + +#[tokio::test] +async fn fixed_width_tuples() { + let client = get_client(); + + let mut cursor = client + .query( + "SELECT + tuple(toUInt64(number), toUInt32(number + 1)) AS fixed_tuple, + arrayMap( + x -> tuple(toUInt64(x), toUInt32(x + 1)), + range(number) + ) AS fixed_tuple_array + FROM system.numbers + LIMIT 3", + ) + .fetch_native() + .unwrap(); + + let block = cursor.next().await.unwrap().expect("expected one block"); + + let tuples = block["fixed_tuple"] + .iter::<(u64, u32)>() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(tuples, [(0, 1), (1, 2), (2, 3)]); + + let tuple_arrays = block["fixed_tuple_array"] + .iter::>() + .unwrap() + .collect::, _>>() + .unwrap(); + assert_eq!(tuple_arrays, [vec![], vec![(0, 1)], vec![(0, 1), (1, 2)]]); +} diff --git a/tests/it/insert.rs b/tests/it/insert.rs index 631b048b..5cd4423f 100644 --- a/tests/it/insert.rs +++ b/tests/it/insert.rs @@ -1,7 +1,9 @@ -use crate::{SimpleRow, create_simple_table, fetch_rows, flush_query_log, get_client}; +use crate::{ + SimpleRow, create_simple_table, fetch_rows, flush_query_log, get_client, + get_client_with_session, +}; use clickhouse::insert::Insert; use clickhouse::{Row, sql::Identifier}; -use rand::distr::{Alphanumeric, SampleString}; use serde::{Deserialize, Serialize}; use serde_json::Value; use std::panic::AssertUnwindSafe; @@ -617,10 +619,7 @@ async fn insert_into_temp_table() { baz: Option, } - let client = get_client().with_setting( - "session_id", - Alphanumeric.sample_string(&mut rand::rng(), 16), - ); + let client = get_client_with_session(); client .query("CREATE TEMPORARY TABLE foo(bar Int32, baz Nullable(String))") diff --git a/tests/it/insert_native.rs b/tests/it/insert_native.rs new file mode 100644 index 00000000..8b9e8581 --- /dev/null +++ b/tests/it/insert_native.rs @@ -0,0 +1,358 @@ +use crate::get_client_with_session; +use clickhouse::native::builder::BlockBuilder; +use std::collections::HashMap; +#[tokio::test] +async fn mixed_types_empty() { + mixed_types(0).await +} + +#[tokio::test] +async fn mixed_types_1() { + mixed_types(1).await +} + +#[tokio::test] +async fn mixed_types_10() { + mixed_types(10).await +} + +#[tokio::test] +async fn mixed_types_100() { + mixed_types(100).await +} + +#[tokio::test] +async fn mixed_types_1000() { + mixed_types(1000).await +} + +async fn mixed_types(num_rows: u64) { + let client = get_client_with_session(); + + client + .query( + "CREATE TEMPORARY TABLE foo( + number Int32, + text String, + nullable_number Nullable(UInt64), + nullable_text Nullable(String), + number_tuple Tuple(UInt32, Int64), + number_text_tuple Tuple(Int64, String), + nullable_tuple Tuple(Nullable(Int64), Nullable(String)), + low_cardinality_text LowCardinality(String), + number_array Array(Int32), + text_array Array(String), + nullable_text_array Array(Nullable(String)), + number_text_map Map(Int32, String), + )", + ) + .execute() + .await + .unwrap(); + + let numbers = (0..(num_rows as i32)).collect::>(); + let texts = (0..(num_rows as usize)) + .map(|i| { + if i == 0 { + return "".to_string(); + } + + // Left-pad with dots to `i` width + format!("{i:.>(); + + let number_arrays = (0..(num_rows as usize)) + .map(|len| (0..len as i32).collect::>()) + .collect::>(); + + let text_arrays = (0..(num_rows as usize)) + .map(|len| { + (0..len) + .map(|i| { + if i == 0 { + return "".to_string(); + } + + // Right-pad with asterisks to `i` width + format!("{i:*>i$}") + }) + .collect::>() + }) + .collect::>(); + + let maps = (0..(num_rows as usize)) + .map(|len| { + numbers[..len] + .iter() + .zip(&texts) + .map(|(&n, t)| (n, t.clone())) + .collect::>() + }) + .collect::>(); + + let mut builder = BlockBuilder::new(); + + builder + .upsert_column("number") + .unwrap() + .add_all(&numbers) + .unwrap(); + + builder + .upsert_column("text") + .unwrap() + .add_all(&texts) + .unwrap(); + + builder + .upsert_column("nullable_number") + .unwrap() + .add_all(numbers.iter().map(|&i| (i % 2 == 0).then_some(i))) + .unwrap(); + + builder + .upsert_column("nullable_text") + .unwrap() + .add_all( + numbers + .iter() + .zip(&texts) + .map(|(&number, text)| (number % 2 != 0).then_some(text)), + ) + .unwrap(); + + builder + .upsert_column::<(i32, i64)>("number_tuple") + .unwrap() + .add_all(numbers.iter().map(|&i| (i, i as i64))) + .unwrap(); + + builder + .upsert_column("number_text_tuple") + .unwrap() + .add_all( + numbers + .iter() + .zip(&texts) + .map(|(&number, text)| (number as i64, text)), + ) + .unwrap(); + + builder + .upsert_column("nullable_tuple") + .unwrap() + .add_all(numbers.iter().zip(&texts).map(|(&number, text)| { + ( + (number % 2 == 0).then_some(number as i64), + (number % 2 != 0).then_some(text), + ) + })) + .unwrap(); + + // Verifying the assumption that `LowCardinality` can accept a regular data stream + builder + .upsert_column("low_cardinality_text") + .unwrap() + .add_all(&texts) + .unwrap(); + + builder + .upsert_column("number_array") + .unwrap() + .add_all(&number_arrays) + .unwrap(); + + builder + .upsert_column("text_array") + .unwrap() + .add_all(&text_arrays) + .unwrap(); + + builder + .upsert_column("nullable_text_array") + .unwrap() + .add_all(text_arrays.iter().map(|array| { + array + .iter() + .enumerate() + .map(|(i, text)| i.is_multiple_of(2).then_some(text)) + .collect::>() + })) + .unwrap(); + + builder + .upsert_column("number_text_map") + .unwrap() + .add_all(&maps) + .unwrap(); + + let block_in = builder.build().unwrap(); + + let mut insert = client.insert_native("foo"); + + insert.write(&block_in).await.unwrap(); + + insert.end().await.unwrap(); + + // This is going to be similar to the `fetch_native` test, + // but we need to make sure the data actually got inserted correctly. + let mut cursor = client.query("SELECT * FROM foo").fetch_native().unwrap(); + + let Some(block_out) = cursor.next().await.unwrap() else { + assert_eq!(num_rows, 0, "expected block, got none"); + return; + }; + + let mut number_iter = block_out["number"].iter::().unwrap(); + + for (res, &expected) in number_iter.by_ref().zip(&numbers) { + let actual = res.unwrap(); + assert_eq!(actual, expected); + } + + if let Some(res) = number_iter.next() { + panic!("unexpected value {res:?}"); + } + + let mut text_iter = block_out["text"].iter::().unwrap(); + + for (res, expected) in text_iter.by_ref().zip(&texts) { + let actual = res.unwrap(); + assert_eq!(actual, *expected); + } + + if let Some(res) = text_iter.next() { + panic!("unexpected value {res:?}"); + } + + let mut nullable_text_iter = block_out["nullable_text"].iter::>().unwrap(); + + for (i, (res, expected)) in nullable_text_iter.by_ref().zip(&texts).enumerate() { + let actual = res.unwrap(); + + if !i.is_multiple_of(2) { + assert_eq!(actual.as_ref(), Some(expected)); + } else { + assert_eq!(actual, None); + } + } + + if let Some(res) = nullable_text_iter.next() { + panic!("unexpected value {res:?}"); + } + + let mut number_tuple_iter = block_out["number_tuple"].iter::<(u32, i64)>().unwrap(); + + for (res, &expected) in number_tuple_iter.by_ref().zip(&numbers) { + let (actual_uint32, actual_int64) = res.unwrap(); + assert_eq!(actual_uint32, expected as u32); + assert_eq!(actual_int64, expected as i64); + } + + if let Some(res) = number_tuple_iter.next() { + panic!("unexpected value {res:?}"); + } + + let mut number_text_tuple_iter = block_out["number_text_tuple"] + .iter::<(i64, String)>() + .unwrap(); + + for ((res, &expected_number), expected_text) in + number_text_tuple_iter.by_ref().zip(&numbers).zip(&texts) + { + let (actual_number, actual_text) = res.unwrap(); + assert_eq!(actual_number, expected_number as i64); + assert_eq!(actual_text, *expected_text); + } + + if let Some(res) = number_text_tuple_iter.next() { + panic!("unexpected value {res:?}"); + } + + let mut nullable_tuple_iter = block_out["nullable_tuple"] + .iter::<(Option, Option)>() + .unwrap(); + + for ((res, &expected_number), expected_text) in + nullable_tuple_iter.by_ref().zip(&numbers).zip(&texts) + { + let (actual_number, actual_text) = res.unwrap(); + + assert_eq!( + actual_number, + (expected_number % 2 == 0).then_some(expected_number as i64) + ); + assert_eq!( + actual_text.as_ref(), + (expected_number % 2 != 0).then_some(expected_text) + ); + } + + if let Some(res) = nullable_tuple_iter.next() { + panic!("unexpected value {res:?}"); + } + + let mut lc_text_iter = block_out["low_cardinality_text"].iter::().unwrap(); + + for (res, expected) in lc_text_iter.by_ref().zip(&texts) { + let actual = res.unwrap(); + assert_eq!(actual, *expected); + } + + if let Some(res) = lc_text_iter.next() { + panic!("unexpected value {res:?}"); + } + + let mut number_array_iter = block_out["number_array"].iter::>().unwrap(); + + for (res, expected) in number_array_iter.by_ref().zip(&number_arrays) { + let actual = res.unwrap(); + assert_eq!(actual, *expected); + } + + if let Some(res) = number_array_iter.next() { + panic!("unexpected value {res:?}"); + } + + let mut text_array_iter = block_out["text_array"].iter::>().unwrap(); + + for (res, expected) in text_array_iter.by_ref().zip(&text_arrays) { + let actual = res.unwrap(); + assert_eq!(actual, *expected); + } + + if let Some(res) = text_array_iter.next() { + panic!("unexpected value {res:?}"); + } + + let mut nullable_text_array_iter = block_out["nullable_text_array"] + .iter::>>() + .unwrap(); + + for (res, expected) in nullable_text_array_iter.by_ref().zip(&text_arrays) { + let actual = res.unwrap(); + + for (i, (actual, expected)) in actual.iter().zip(expected).enumerate() { + assert_eq!(actual.as_ref(), i.is_multiple_of(2).then_some(expected)); + } + } + + if let Some(res) = nullable_text_array_iter.next() { + panic!("unexpected value {res:?}"); + } + + let mut map_iter = block_out["number_text_map"] + .iter::>() + .unwrap(); + + for (res, expected) in map_iter.by_ref().zip(&maps) { + let actual = res.unwrap(); + + assert_eq!(actual, *expected); + } + + if let Some(res) = map_iter.next() { + panic!("unexpected value {res:?}"); + } +} diff --git a/tests/it/main.rs b/tests/it/main.rs index cf227581..f76808bd 100644 --- a/tests/it/main.rs +++ b/tests/it/main.rs @@ -26,6 +26,7 @@ //! clean up outdated databases based on its creation time. use clickhouse::{Client, Row, RowOwned, RowRead, RowWrite, sql::Identifier}; +use rand::distr::{Alphanumeric, SampleString}; use serde::{Deserialize, Serialize}; use std::sync::LazyLock; @@ -138,6 +139,13 @@ pub(crate) fn get_client() -> Client { } } +pub(crate) fn get_client_with_session() -> Client { + get_client().with_setting( + "session_id", + Alphanumeric.sample_string(&mut rand::rng(), 16), + ) +} + pub(crate) fn require_env_var(name: &str) -> String { std::env::var(name).unwrap_or_else(|_| panic!("{name} environment variable is not set")) } @@ -258,6 +266,7 @@ mod fetch_native; mod https_errors; mod insert; mod insert_formatted; +mod insert_native; #[cfg(feature = "inserter")] mod inserter; mod int128; diff --git a/tests/it/native_types.rs b/tests/it/native_types.rs index cd471b4d..cee08141 100644 --- a/tests/it/native_types.rs +++ b/tests/it/native_types.rs @@ -1,4 +1,4 @@ -use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; +use std::net::{Ipv4Addr, Ipv6Addr}; macro_rules! test_type { ( @@ -10,39 +10,118 @@ macro_rules! test_type { $(#[$attr])* #[tokio::test] async fn $fn_name() { - let client = $crate::get_client(); + use clickhouse::native::builder::BlockBuilder; - let mut cursor = client - .query(concat!( - "SELECT c1 FROM Values (" - $(,"'c1 ", $sqlty, "', ")? - $(, $sql, )","* - ")" + let db_name = test_database_name!(); + + let client = $crate::_priv::prepare_database(&db_name).await; + + client + .query(&format!( + // Temporary tables apparently can't be the source of `CREATE TABLE .. AS ..`, + // you get a confusing "unknown table" error + "CREATE TABLE {db_name}.sql_values ENGINE = Memory AS {}", + // In case any values contain `{}` + concat!("SELECT * FROM Values(" $(, "'c1 ", $sqlty, "', ")? $(, $sql, )","* ")") + )) + .execute() + .await + .expect("error creating value source table"); + + client + .query(&format!( + "CREATE TABLE {db_name}.insert_values ENGINE = Memory AS \ + {db_name}.sql_values", )) + .execute() + .await + .expect("error creating insert table"); + + let mut insert_block = BlockBuilder::new(); + + let mut column = insert_block.upsert_column::<$ty>("c1").expect("error from upsert_column"); + + $( + column.add($rust).expect(concat!("error writing expression `", stringify!($rust), "`")); + )* + + let insert_block = insert_block.build().expect("error from insert_block.build()"); + + let mut insert = client.insert_native("insert_values"); + + insert.write(&insert_block).await.expect("error from insert.write()"); + + insert.end().await.expect("error from insert.end()"); + + let mut cursor = client + .query( + "SELECT \ + sql_values.c1 AS sql_value, insert_values.c1 AS insert_value, \ + toBool(sql_value == insert_value) AS values_equal \ + FROM sql_values PASTE JOIN insert_values", + ) .fetch_native() - .expect("error from `.fetch_native()`"); + .expect("error from fetch_native"); - let block = cursor.next().await + let block = cursor + .next() + .await .expect("error from `cursor.next()`") .expect("expected block, got none"); - let mut iter = block["c1"] + let mut sql_iter = block["sql_value"] .iter::<$ty>() - .expect("error from `.iter()`"); + .expect("error from `block[\"sql_value\"].iter()`"); $( - let expected = $rust; + let expected: $ty = $rust; - let val = iter + let val = sql_iter .next() - .expect("expected another value, got none") + .expect("expected another value from `sql_iter`") .unwrap_or_else(|e| panic!("error decoding SQL `{}` as Rust value `{expected:?}`: {e:?}", $sql)); - assert_eq!(val, expected); + assert_eq!(val, expected, "SQL value does not equal Rust value"); )* - if let Some(next) = iter.next() { - panic!("unexpected value: {next:?}"); + if let Some(next) = sql_iter.next() { + panic!("`unexpected value from `sql_iter.next(): {next:?}`"); + } + + let mut insert_iter = block["insert_value"] + .iter::<$ty>() + .expect("error from `block[\"insert_value\"].iter()`"); + + $( + let expected: $ty = $rust; + + let val = insert_iter + .next() + .expect("expected another value from `insert_iter`") + .unwrap_or_else(|e| panic!("error round-tripping Rust value `{expected:?}`: {e:?}")); + + assert_eq!(val, expected, "Rust value did not round-trip correctly"); + )* + + if let Some(next) = sql_iter.next() { + panic!("unexpected value from `insert_iter.next()`: {next:?}"); + } + + let mut equals_iter = block["values_equal"] + .iter::() + .expect("error from `block[\"values_equal\"].iter()`"); + + $( + let equals = equals_iter + .next() + .expect("expected another value from `equals_iter`") + .expect("error decoding value from `equals_iter`"); + + assert!(equals, "values not equal in SQL: {:?} vs `{}`", $sql, stringify!($rust)); + )* + + if let Some(next) = equals_iter.next() { + panic!("`unexpected value from `equals_iter.next(): {next:?}`"); } } }; @@ -117,17 +196,6 @@ test_type!( } ); -test_type!( - test_ipaddr_from_v4(IpAddr, "IPv4") { - "'0.0.0.0'" == Ipv4Addr::UNSPECIFIED, - "'1.1.1.1'" == Ipv4Addr::new(1, 1, 1, 1), - "'127.0.0.1'" == Ipv4Addr::LOCALHOST, - "'192.168.2.1'" == Ipv4Addr::new(192, 168, 2, 1), - "'255.255.255.0'" == Ipv4Addr::new(255, 255, 255, 0), - "'255.255.255.255'" == Ipv4Addr::BROADCAST, - } -); - test_type!( test_ipv6(Ipv6Addr, "IPv6") { "'::1'" == Ipv6Addr::LOCALHOST, @@ -137,15 +205,6 @@ test_type!( } ); -test_type!( - test_ipaddr_from_v6(IpAddr, "IPv6") { - "'::1'" == Ipv6Addr::LOCALHOST, - // IPv6 addresses for ClickHouse.com - "'2606:4700:3108::ac42:2b07'" == "2606:4700:3108::ac42:2b07".parse::().unwrap(), - "'2606:4700:3108::ac42:28f9'" == "2606:4700:3108::ac42:28f9".parse::().unwrap(), - } -); - #[cfg(feature = "uuid")] mod uuid { use uuid::Uuid; diff --git a/types/src/data_types.rs b/types/src/data_types.rs index 681be0ec..4a0bf5cb 100644 --- a/types/src/data_types.rs +++ b/types/src/data_types.rs @@ -1,4 +1,5 @@ use crate::error::TypesError; +use std::borrow::Cow; use std::collections::HashMap; use std::fmt::{Display, Formatter}; @@ -196,6 +197,55 @@ impl DataTypeNode { _ => self, } } + + /// If `self` has a static string representation (e.g. `"UInt8"`), return it. + /// + /// Returns `None` for polymorphic types (e.g. `Array(T)` or `Decimal(P, S)`). + pub fn as_str(&self) -> Option<&'static str> { + use DataTypeNode::*; + + Some(match self { + UInt8 => "UInt8", + UInt16 => "UInt16", + UInt32 => "UInt32", + UInt64 => "UInt64", + UInt128 => "UInt128", + UInt256 => "UInt256", + Int8 => "Int8", + Int16 => "Int16", + Int32 => "Int32", + Int64 => "Int64", + Int128 => "Int128", + Int256 => "Int256", + Float32 => "Float32", + Float64 => "Float64", + BFloat16 => "BFloat16", + String => "String", + UUID => "UUID", + Date => "Date", + Date32 => "Date32", + DateTime(None) => "DateTime", + Time => "Time", + IPv4 => "IPv4", + IPv6 => "IPv6", + Bool => "Bool", + JSON => "JSON", + Dynamic => "Dynamic", + Point => "Point", + Ring => "Ring", + LineString => "LineString", + MultiLineString => "MultiLineString", + Polygon => "Polygon", + MultiPolygon => "MultiPolygon", + _ => return None, + }) + } + + /// Return the string representation for this type, without allocating if possible. + pub fn to_str(&self) -> Cow<'static, str> { + self.as_str() + .map_or_else(|| self.to_string().into(), Cow::Borrowed) + } } impl From for String {