diff --git a/fluss-rust/crates/fluss/src/metadata/table.rs b/fluss-rust/crates/fluss/src/metadata/table.rs index 79755d4cf9d..b7362d730b8 100644 --- a/fluss-rust/crates/fluss/src/metadata/table.rs +++ b/fluss-rust/crates/fluss/src/metadata/table.rs @@ -22,6 +22,7 @@ use crate::metadata::DataLakeFormat; use crate::metadata::datatype::{ DataField, DataType, RowType, UNASSIGNED_FIELD_ID, reassign_field_ids, }; +use crate::record::is_supported_statistics_type; use crate::{BucketId, PartitionId, SnapshotId, TableId}; use core::fmt; use serde::{Deserialize, Serialize}; @@ -33,6 +34,10 @@ use strum_macros::EnumString; /// Sentinel for a column whose stable id has not yet been assigned. pub const UNKNOWN_COLUMN_ID: i32 = -1; +/// Table property selecting the columns that written batches collect statistics +/// for, either `*` or a comma-separated list. +pub const TABLE_STATISTICS_COLUMNS: &str = "table.statistics.columns"; + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct Column { name: String, @@ -1199,6 +1204,41 @@ impl TableConfig { pub fn get_auto_partition_strategy(&self) -> AutoPartitionStrategy { AutoPartitionStrategy::from(&self.properties) } + + /// Reads `table.statistics.columns`, which decides whether written batches + /// carry the statistics the server prunes by. + pub fn get_statistics_columns(&self) -> StatisticsColumns { + match self.properties.get(TABLE_STATISTICS_COLUMNS) { + None => StatisticsColumns::Disabled, + Some(value) if value == "*" => StatisticsColumns::All, + Some(value) => StatisticsColumns::Specified( + value + .split(',') + .map(str::trim) + .filter(|name| !name.is_empty()) + .map(str::to_string) + .collect(), + ), + } + } +} + +/// Which columns a table collects statistics for, mirroring Java's +/// `StatisticsColumnsConfig`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum StatisticsColumns { + /// The property is unset, so batches stay in the V0 format. + Disabled, + /// `*`, meaning every column whose type supports statistics. + All, + /// An explicit column list, taken as given. + Specified(Vec), +} + +impl StatisticsColumns { + pub fn is_enabled(&self) -> bool { + !matches!(self, StatisticsColumns::Disabled) + } } impl TableInfo { @@ -1355,6 +1395,45 @@ impl TableInfo { &self.properties } + /// Column indices, in order, that written batches collect statistics for. + /// + /// Empty when the table has not enabled statistics. `*` keeps only the + /// columns whose type supports statistics, while an explicit list is taken + /// as given: the server already rejects an unsupported type when the table + /// is created or altered, so the client trusts it as Java's does. + /// + /// # Errors + /// Returns an error if a named column is absent from the table schema. + pub fn get_stats_index_mapping(&self) -> Result> { + let names = match self.table_config.get_statistics_columns() { + StatisticsColumns::Disabled => return Ok(Vec::new()), + StatisticsColumns::All => { + return Ok(self + .row_type + .fields() + .iter() + .enumerate() + .filter(|(_, field)| is_supported_statistics_type(field.data_type())) + .map(|(index, _)| index) + .collect()); + } + StatisticsColumns::Specified(names) => names, + }; + + names + .iter() + .map(|name| { + self.row_type + .fields() + .iter() + .position(|field| field.name() == name) + .ok_or_else(|| Error::IllegalArgument { + message: format!("Statistics column '{name}' not found in table schema"), + }) + }) + .collect() + } + pub fn get_table_config(&self) -> &TableConfig { &self.table_config } @@ -1672,4 +1751,74 @@ mod tests { ); assert!(table_info.is_auto_partitioned()); } + + fn stats_table(property: Option<&str>) -> TableInfo { + let schema = Schema::builder() + .column("id", DataTypes::int()) + .column("name", DataTypes::string()) + .column("payload", DataTypes::bytes()) + .build() + .expect("schema"); + let mut descriptor = TableDescriptor::builder() + .schema(schema) + .distributed_by(Some(1), vec![]); + if let Some(value) = property { + descriptor = descriptor.property(TABLE_STATISTICS_COLUMNS, value); + } + TableInfo::of( + TablePath::new("db", "tbl"), + 1, + 1, + descriptor.build().expect("descriptor"), + 0, + 0, + ) + } + + #[test] + fn statistics_are_disabled_without_the_property() { + let table = stats_table(None); + assert_eq!( + table.get_table_config().get_statistics_columns(), + StatisticsColumns::Disabled + ); + assert!(table.get_stats_index_mapping().expect("mapping").is_empty()); + } + + #[test] + fn star_keeps_only_columns_whose_type_supports_statistics() { + let table = stats_table(Some("*")); + assert_eq!( + table.get_table_config().get_statistics_columns(), + StatisticsColumns::All + ); + // BYTES has no statistics support, so the payload column drops out. + assert_eq!( + table.get_stats_index_mapping().expect("mapping"), + vec![0, 1] + ); + } + + #[test] + fn a_named_list_is_taken_as_given_and_trimmed() { + let table = stats_table(Some(" name , id ")); + assert_eq!( + table.get_table_config().get_statistics_columns(), + StatisticsColumns::Specified(vec!["name".to_string(), "id".to_string()]) + ); + // Order follows the property, not the schema. + assert_eq!( + table.get_stats_index_mapping().expect("mapping"), + vec![1, 0] + ); + } + + #[test] + fn an_unknown_statistics_column_is_rejected() { + let table = stats_table(Some("nope")); + assert!(matches!( + table.get_stats_index_mapping(), + Err(Error::IllegalArgument { .. }) + )); + } } diff --git a/fluss-rust/crates/fluss/src/record/mod.rs b/fluss-rust/crates/fluss/src/record/mod.rs index 7e548b2cace..c30d49a20ed 100644 --- a/fluss-rust/crates/fluss/src/record/mod.rs +++ b/fluss-rust/crates/fluss/src/record/mod.rs @@ -24,6 +24,11 @@ use std::collections::HashMap; mod arrow; mod error; pub mod kv; +// Reachable once the Arrow builder emits V1 batches. +#[allow(dead_code, reason = "consumed by the V1 batch builder")] +mod statistics; + +pub(crate) use statistics::is_supported_statistics_type; pub use arrow::*; diff --git a/fluss-rust/crates/fluss/src/record/statistics.rs b/fluss-rust/crates/fluss/src/record/statistics.rs new file mode 100644 index 00000000000..aa7797f03a2 --- /dev/null +++ b/fluss-rust/crates/fluss/src/record/statistics.rs @@ -0,0 +1,908 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Per-column statistics carried by a V1 log record batch, which the server +//! uses to skip whole batches that a pushed-down filter cannot match. +//! +//! Statistics are derived from the finished Arrow batch rather than row by row +//! as Java does, so both the row-append and batch-append writer paths are +//! covered by one implementation. +//! +//! Collecting two columns, with `mapping = [1, 0]`, lays out as: +//! +//! ```text +//! offset 0 1 3 7 15 +//! |-----|---------|------------|-------------| +//! | ver | count=2 | indexes | null counts | +//! | 0x01| i16 | i16 x 2 | i32 x 2 | +//! |-----|---------|------------|-------------| +//! [1, 0] [n1, n0] +//! +//! 15 19 ... +//! |-------------|---------------------| +//! | min row len | min row (AlignedRow)| +//! | i32 | 2 fields | +//! |-------------|---------------------| +//! | max row len | max row (AlignedRow)| +//! | i32 | 2 fields | +//! |-------------|---------------------| +//! ``` +//! +//! Everything is little-endian, and the two rows carry one field per collected +//! column in the same order as the index array. + +use crate::error::{Error, Result}; +use crate::metadata::{DataType, RowType}; +use crate::row::aligned::AlignedRowWriter; +use crate::row::binary::BinaryWriter; +use crate::row::{Decimal, TimestampLtz, TimestampNtz}; +use arrow::array::{Array, RecordBatch}; +use arrow::compute::kernels::aggregate; +use arrow::datatypes::{ + Date32Type, Decimal128Type, Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, + Int64Type, Time32MillisecondType, TimestampMicrosecondType, TimestampMillisecondType, + TimestampNanosecondType, TimestampSecondType, +}; + +/// Version byte leading the statistics block, matching Java's +/// `LogRecordBatchFormat.STATISTICS_VERSION`. +const STATISTICS_VERSION: u8 = 1; + +/// Whether statistics can be collected for `data_type`, mirroring Java's +/// `DataTypeChecks.isSupportedStatisticsType`. +pub(crate) fn is_supported_statistics_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::Boolean(_) + | DataType::TinyInt(_) + | DataType::SmallInt(_) + | DataType::Int(_) + | DataType::BigInt(_) + | DataType::Float(_) + | DataType::Double(_) + | DataType::String(_) + | DataType::Char(_) + | DataType::Decimal(_) + | DataType::Date(_) + | DataType::Time(_) + | DataType::Timestamp(_) + | DataType::TimestampLTz(_) + ) +} + +/// The minimum and maximum of one column, or `None` when every value is null. +enum ColumnBounds { + Bool(bool, bool), + Int8(i8, i8), + Int16(i16, i16), + Int32(i32, i32), + Int64(i64, i64), + Float32(f32, f32), + Float64(f64, f64), + Str(String, String), + Decimal(Decimal, Decimal), + TimestampNtz(TimestampNtz, TimestampNtz), + TimestampLtz(TimestampLtz, TimestampLtz), +} + +/// Serialises the statistics of `batch` for the columns named by `mapping`. +/// +/// `mapping[i]` is the table column whose min, max and null count go into the +/// block's `i`th position, so `[1, 0]` collects column 1 before column 0. +/// +/// Returns `None` when there is nothing worth sending, which is how an empty +/// batch or an empty mapping is signalled to the caller. +pub(crate) fn serialize_statistics( + batch: &RecordBatch, + row_type: &RowType, + mapping: &[usize], +) -> Result>> { + if mapping.is_empty() || batch.num_rows() == 0 { + return Ok(None); + } + + // Indexing the batch and the row type below would panic rather than error, + // so make both bounds explicit. Column types are checked lazily by + // `column_bounds` as each one is read. + let field_count = row_type.fields().len(); + if batch.num_columns() != field_count { + return Err(Error::IllegalArgument { + message: format!( + "Statistics need a batch matching the table schema, got {} columns for {field_count} fields", + batch.num_columns() + ), + }); + } + if let Some(&index) = mapping.iter().find(|&&index| index >= field_count) { + return Err(Error::IllegalArgument { + message: format!( + "Statistics column index {index} is out of range for {field_count} fields" + ), + }); + } + + let mut null_counts = Vec::with_capacity(mapping.len()); + let mut bounds = Vec::with_capacity(mapping.len()); + for &column_index in mapping { + let column = batch.column(column_index); + null_counts.push(column.null_count() as i32); + bounds.push(column_bounds( + column, + row_type.fields()[column_index].data_type(), + )?); + } + + // Everything below is little-endian, as Java's memory segments are. + let mut out = Vec::new(); + out.push(STATISTICS_VERSION); + // Append the column count, which sizes the two arrays that follow. + out.extend_from_slice(&(mapping.len() as i16).to_le_bytes()); + // Append the table column index each position describes. + for &column_index in mapping { + out.extend_from_slice(&(column_index as i16).to_le_bytes()); + } + // Append the null count of each of those columns. + for count in &null_counts { + out.extend_from_slice(&count.to_le_bytes()); + } + + let types: Vec<&DataType> = mapping + .iter() + .map(|&i| row_type.fields()[i].data_type()) + .collect(); + append_row(&mut out, &bounds, &types, Bound::Min); + append_row(&mut out, &bounds, &types, Bound::Max); + Ok(Some(out)) +} + +#[derive(Clone, Copy)] +enum Bound { + Min, + Max, +} + +/// Writes one aligned row of bounds, length-prefixed as Java's +/// `LogRecordBatchStatisticsWriter.writeRowData`. +fn append_row(out: &mut Vec, bounds: &[Option], types: &[&DataType], b: Bound) { + let mut writer = AlignedRowWriter::new(bounds.len()); + for (index, bound) in bounds.iter().enumerate() { + match bound { + // An all-null column has no bound, so the slot itself is null. + None => writer.set_null_at(index), + Some(bound) => write_bound(&mut writer, bound, types[index], b), + } + } + writer.complete(); + let row = writer.to_bytes(); + out.extend_from_slice(&(row.len() as i32).to_le_bytes()); + out.extend_from_slice(&row); +} + +/// Picks the requested end of a bound pair. +fn pick(b: Bound, min: T, max: T) -> T { + match b { + Bound::Min => min, + Bound::Max => max, + } +} + +fn write_bound(writer: &mut AlignedRowWriter, bound: &ColumnBounds, ty: &DataType, b: Bound) { + match bound { + ColumnBounds::Bool(min, max) => writer.write_boolean(pick(b, *min, *max)), + ColumnBounds::Int8(min, max) => writer.write_byte(pick(b, *min, *max) as u8), + ColumnBounds::Int16(min, max) => writer.write_short(pick(b, *min, *max)), + ColumnBounds::Int32(min, max) => writer.write_int(pick(b, *min, *max)), + ColumnBounds::Int64(min, max) => writer.write_long(pick(b, *min, *max)), + ColumnBounds::Float32(min, max) => writer.write_float(pick(b, *min, *max)), + ColumnBounds::Float64(min, max) => writer.write_double(pick(b, *min, *max)), + ColumnBounds::Str(min, max) => writer.write_string(pick(b, min, max)), + ColumnBounds::Decimal(min, max) => { + let value = pick(b, min, max); + let precision = match ty { + DataType::Decimal(decimal_type) => decimal_type.precision(), + _ => value.precision(), + }; + writer.write_decimal(value, precision); + } + ColumnBounds::TimestampNtz(min, max) => { + writer.write_timestamp_ntz(pick(b, min, max), precision_of(ty)); + } + ColumnBounds::TimestampLtz(min, max) => { + writer.write_timestamp_ltz(pick(b, min, max), precision_of(ty)); + } + } +} + +fn precision_of(ty: &DataType) -> u32 { + match ty { + DataType::Timestamp(t) => t.precision(), + DataType::TimestampLTz(t) => t.precision(), + _ => 6, + } +} + +/// Reduces one Arrow column to its bounds, returning `None` when it is entirely +/// null and therefore has none. +fn column_bounds(column: &dyn Array, data_type: &DataType) -> Result> { + use arrow::array::*; + use arrow::datatypes::DataType as ArrowType; + + macro_rules! primitive { + ($arrow_ty:ty, $variant:ident) => {{ + let array = column + .as_any() + .downcast_ref::>() + .ok_or_else(|| unexpected_array(column, data_type))?; + match (aggregate::min(array), aggregate::max(array)) { + (Some(min), Some(max)) => Ok(Some(ColumnBounds::$variant(min, max))), + _ => Ok(None), + } + }}; + } + + match data_type { + DataType::Boolean(_) => { + let array = column + .as_any() + .downcast_ref::() + .ok_or_else(|| unexpected_array(column, data_type))?; + match (aggregate::min_boolean(array), aggregate::max_boolean(array)) { + (Some(min), Some(max)) => Ok(Some(ColumnBounds::Bool(min, max))), + _ => Ok(None), + } + } + DataType::TinyInt(_) => primitive!(Int8Type, Int8), + DataType::SmallInt(_) => primitive!(Int16Type, Int16), + DataType::Int(_) => primitive!(Int32Type, Int32), + DataType::Date(_) => primitive!(Date32Type, Int32), + DataType::BigInt(_) => primitive!(Int64Type, Int64), + DataType::Float(_) => primitive!(Float32Type, Float32), + DataType::Double(_) => primitive!(Float64Type, Float64), + // Fluss stores TIME as millis of day, so every unit but millisecond + // has to be converted back from what the Arrow array holds. + DataType::Time(_) => match column.data_type() { + ArrowType::Time32(arrow::datatypes::TimeUnit::Second) => { + let array = column + .as_any() + .downcast_ref::() + .ok_or_else(|| unexpected_array(column, data_type))?; + match (aggregate::min(array), aggregate::max(array)) { + (Some(min), Some(max)) => { + Ok(Some(ColumnBounds::Int32(min * 1_000, max * 1_000))) + } + _ => Ok(None), + } + } + ArrowType::Time32(_) => primitive!(Time32MillisecondType, Int32), + ArrowType::Time64(arrow::datatypes::TimeUnit::Microsecond) => time64_as_millis( + column.as_any().downcast_ref::(), + 1_000, + ), + ArrowType::Time64(_) => time64_as_millis( + column.as_any().downcast_ref::(), + 1_000_000, + ), + _ => Err(unexpected_array(column, data_type)), + }, + DataType::String(_) | DataType::Char(_) => { + let array = column + .as_any() + .downcast_ref::() + .ok_or_else(|| unexpected_array(column, data_type))?; + match (aggregate::min_string(array), aggregate::max_string(array)) { + (Some(min), Some(max)) => { + Ok(Some(ColumnBounds::Str(min.to_string(), max.to_string()))) + } + _ => Ok(None), + } + } + DataType::Decimal(decimal_type) => { + let array = column + .as_any() + .downcast_ref::>() + .ok_or_else(|| unexpected_array(column, data_type))?; + let (precision, scale) = (decimal_type.precision(), decimal_type.scale()); + match (aggregate::min(array), aggregate::max(array)) { + (Some(min), Some(max)) => Ok(Some(ColumnBounds::Decimal( + decimal_from_i128(min, precision, scale)?, + decimal_from_i128(max, precision, scale)?, + ))), + _ => Ok(None), + } + } + DataType::Timestamp(_) => { + let (min, max) = match timestamp_bounds(column, data_type)? { + Some(bounds) => bounds, + None => return Ok(None), + }; + Ok(Some(ColumnBounds::TimestampNtz( + TimestampNtz::from_millis_nanos(min.0, min.1)?, + TimestampNtz::from_millis_nanos(max.0, max.1)?, + ))) + } + DataType::TimestampLTz(_) => { + let (min, max) = match timestamp_bounds(column, data_type)? { + Some(bounds) => bounds, + None => return Ok(None), + }; + Ok(Some(ColumnBounds::TimestampLtz( + TimestampLtz::from_millis_nanos(min.0, min.1)?, + TimestampLtz::from_millis_nanos(max.0, max.1)?, + ))) + } + other => Err(Error::IllegalArgument { + message: format!("Statistics are not supported for column type {other:?}"), + }), + } +} + +/// Fluss stores TIME as milliseconds of day, so a finer Arrow unit is scaled +/// down by `divisor` before it becomes a bound. +fn time64_as_millis( + array: Option<&arrow::array::PrimitiveArray>, + divisor: i64, +) -> Result> +where + T: arrow::datatypes::ArrowPrimitiveType, +{ + let array = array.ok_or_else(|| Error::IllegalArgument { + message: "TIME column is not backed by a Time64 array".to_string(), + })?; + match (aggregate::min(array), aggregate::max(array)) { + (Some(min), Some(max)) => Ok(Some(ColumnBounds::Int32( + (min / divisor) as i32, + (max / divisor) as i32, + ))), + _ => Ok(None), + } +} + +/// Returns the (millis, nano-of-milli) bounds of a timestamp column. +#[allow(clippy::type_complexity)] +fn timestamp_bounds( + column: &dyn Array, + data_type: &DataType, +) -> Result> { + use arrow::array::PrimitiveArray; + use arrow::datatypes::DataType as ArrowType; + use arrow::datatypes::TimeUnit; + + macro_rules! bounds { + ($arrow_ty:ty, $to_parts:expr) => {{ + let array = column + .as_any() + .downcast_ref::>() + .ok_or_else(|| unexpected_array(column, data_type))?; + match (aggregate::min(array), aggregate::max(array)) { + (Some(min), Some(max)) => Ok(Some(($to_parts(min), $to_parts(max)))), + _ => Ok(None), + } + }}; + } + + match column.data_type() { + ArrowType::Timestamp(TimeUnit::Second, _) => { + bounds!(TimestampSecondType, |v: i64| (v * 1_000, 0)) + } + ArrowType::Timestamp(TimeUnit::Millisecond, _) => { + bounds!(TimestampMillisecondType, |v: i64| (v, 0)) + } + ArrowType::Timestamp(TimeUnit::Microsecond, _) => { + bounds!(TimestampMicrosecondType, |v: i64| ( + v.div_euclid(1_000), + (v.rem_euclid(1_000) * 1_000) as i32 + )) + } + ArrowType::Timestamp(TimeUnit::Nanosecond, _) => { + bounds!(TimestampNanosecondType, |v: i64| ( + v.div_euclid(1_000_000), + v.rem_euclid(1_000_000) as i32 + )) + } + _ => Err(unexpected_array(column, data_type)), + } +} + +fn decimal_from_i128(value: i128, precision: u32, scale: u32) -> Result { + if Decimal::is_compact_precision(precision) { + Decimal::from_unscaled_long(value as i64, precision, scale) + } else { + Decimal::from_unscaled_bytes(&value.to_be_bytes(), precision, scale) + } +} + +fn unexpected_array(column: &dyn Array, data_type: &DataType) -> Error { + Error::IllegalArgument { + message: format!( + "Column of Fluss type {data_type:?} is backed by unexpected Arrow type {:?}", + column.data_type() + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::metadata::{DataField, DataTypes}; + use arrow::array::{ + BooleanArray, Date32Array, Decimal128Array, Float32Array, Float64Array, Int8Array, + Int16Array, Int32Array, Int64Array, StringArray, Time32MillisecondArray, Time32SecondArray, + Time64MicrosecondArray, Time64NanosecondArray, TimestampMicrosecondArray, + TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, + }; + use arrow::datatypes::{DataType as ArrowType, Field, Schema}; + use std::sync::Arc; + + fn row_type() -> RowType { + RowType::new(vec![ + DataField::new("id", DataTypes::int(), None), + DataField::new("name", DataTypes::string(), None), + DataField::new("score", DataTypes::bigint(), None), + ]) + } + + fn batch( + ids: Vec>, + names: Vec>, + scores: Vec>, + ) -> RecordBatch { + let schema = Schema::new(vec![ + Field::new("id", ArrowType::Int32, true), + Field::new("name", ArrowType::Utf8, true), + Field::new("score", ArrowType::Int64, true), + ]); + RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Int32Array::from(ids)), + Arc::new(StringArray::from(names)), + Arc::new(Int64Array::from(scores)), + ], + ) + .expect("batch") + } + + /// Reads back the fixed prefix so a layout change cannot pass silently. + fn parse_prefix(bytes: &[u8], columns: usize) -> (u8, i16, Vec, Vec) { + let version = bytes[0]; + let count = i16::from_le_bytes(bytes[1..3].try_into().unwrap()); + let mut at = 3; + let mut indexes = Vec::new(); + for _ in 0..columns { + indexes.push(i16::from_le_bytes(bytes[at..at + 2].try_into().unwrap())); + at += 2; + } + let mut nulls = Vec::new(); + for _ in 0..columns { + nulls.push(i32::from_le_bytes(bytes[at..at + 4].try_into().unwrap())); + at += 4; + } + (version, count, indexes, nulls) + } + + #[test] + fn serializes_the_documented_prefix() { + let rt = row_type(); + let batch = batch( + vec![Some(3), Some(1), Some(2)], + vec![Some("c"), Some("a"), None], + vec![Some(30), Some(10), Some(20)], + ); + let bytes = serialize_statistics(&batch, &rt, &[0, 1, 2]) + .expect("serialize") + .expect("statistics"); + + let (version, count, indexes, nulls) = parse_prefix(&bytes, 3); + assert_eq!(version, 1); + assert_eq!(count, 3); + assert_eq!(indexes, vec![0, 1, 2]); + assert_eq!(nulls, vec![0, 1, 0]); + } + + #[test] + fn collects_bounds_over_the_whole_column() { + let rt = row_type(); + let batch = batch( + vec![Some(3), Some(1), Some(2)], + vec![Some("c"), Some("a"), Some("b")], + vec![Some(30), Some(10), Some(20)], + ); + let bytes = serialize_statistics(&batch, &rt, &[0]) + .expect("serialize") + .expect("statistics"); + + // version(1) + count(2) + index(2) + nullCount(4) = 9 bytes of prefix. + let min_size = i32::from_le_bytes(bytes[9..13].try_into().unwrap()) as usize; + let min_row = &bytes[13..13 + min_size]; + // One field: 8 null-bit bytes then the value slot. + assert_eq!(i32::from_le_bytes(min_row[8..12].try_into().unwrap()), 1); + + let max_at = 13 + min_size; + let max_size = i32::from_le_bytes(bytes[max_at..max_at + 4].try_into().unwrap()) as usize; + let max_row = &bytes[max_at + 4..max_at + 4 + max_size]; + assert_eq!(i32::from_le_bytes(max_row[8..12].try_into().unwrap()), 3); + } + + #[test] + fn marks_an_all_null_column_as_having_no_bounds() { + let rt = row_type(); + let batch = batch(vec![None, None], vec![None, None], vec![None, None]); + let bytes = serialize_statistics(&batch, &rt, &[0]) + .expect("serialize") + .expect("statistics"); + + let (_, _, _, nulls) = parse_prefix(&bytes, 1); + assert_eq!(nulls, vec![2]); + let min_size = i32::from_le_bytes(bytes[9..13].try_into().unwrap()) as usize; + let min_row = &bytes[13..13 + min_size]; + // Field 0's null bit is bit 8, the first after the reserved header bits. + assert_eq!(min_row[1] & 0x01, 0x01); + } + + #[test] + fn skips_an_empty_batch_and_an_empty_mapping() { + let rt = row_type(); + let empty = batch(vec![], vec![], vec![]); + assert!( + serialize_statistics(&empty, &rt, &[0]) + .expect("serialize") + .is_none() + ); + + let populated = batch(vec![Some(1)], vec![Some("a")], vec![Some(1)]); + assert!( + serialize_statistics(&populated, &rt, &[]) + .expect("serialize") + .is_none() + ); + } + + #[test] + fn rejects_a_batch_that_does_not_match_the_schema() { + // Two Arrow columns against a three field row type: indexing the batch + // by a mapping built from the schema would otherwise panic. + let schema = Schema::new(vec![ + Field::new("id", ArrowType::Int32, true), + Field::new("name", ArrowType::Utf8, true), + ]); + let narrow = RecordBatch::try_new( + Arc::new(schema), + vec![ + Arc::new(Int32Array::from(vec![Some(1)])) as arrow::array::ArrayRef, + Arc::new(StringArray::from(vec![Some("a")])), + ], + ) + .expect("batch"); + + assert!(matches!( + serialize_statistics(&narrow, &row_type(), &[0, 1, 2]), + Err(Error::IllegalArgument { .. }) + )); + } + + #[test] + fn rejects_a_mapping_beyond_the_row_type() { + let batch = batch(vec![Some(1)], vec![Some("a")], vec![Some(1)]); + assert!(matches!( + serialize_statistics(&batch, &row_type(), &[3]), + Err(Error::IllegalArgument { .. }) + )); + } + + #[test] + fn rejects_a_column_type_without_statistics_support() { + assert!(!is_supported_statistics_type(&DataTypes::bytes())); + assert!(is_supported_statistics_type(&DataTypes::string())); + assert!(is_supported_statistics_type(&DataTypes::timestamp())); + } + + /// Serialises a one-column batch and returns its min and max aligned rows. + fn single_column_rows( + data_type: DataType, + arrow_type: ArrowType, + array: arrow::array::ArrayRef, + ) -> (Vec, Vec) { + let rt = RowType::new(vec![DataField::new("v", data_type, None)]); + let schema = Schema::new(vec![Field::new("v", arrow_type, true)]); + let batch = RecordBatch::try_new(Arc::new(schema), vec![array]).expect("batch"); + let bytes = serialize_statistics(&batch, &rt, &[0]) + .expect("serialize") + .expect("statistics"); + + // One column means a 9 byte prefix before the length-prefixed rows. + let min_len = i32::from_le_bytes(bytes[9..13].try_into().unwrap()) as usize; + let min = bytes[13..13 + min_len].to_vec(); + let max_at = 13 + min_len; + let max_len = i32::from_le_bytes(bytes[max_at..max_at + 4].try_into().unwrap()) as usize; + let max = bytes[max_at + 4..max_at + 4 + max_len].to_vec(); + (min, max) + } + + /// The 8-byte slot of the single column, which starts after the null bits. + fn only_slot(row: &[u8]) -> &[u8] { + &row[8..16] + } + + #[test] + fn collects_bounds_for_boolean() { + let (min, max) = single_column_rows( + DataTypes::boolean(), + ArrowType::Boolean, + Arc::new(BooleanArray::from(vec![ + Some(true), + Some(false), + Some(true), + ])), + ); + assert_eq!(only_slot(&min)[0], 0); + assert_eq!(only_slot(&max)[0], 1); + } + + #[test] + fn collects_bounds_for_the_narrow_integers() { + let (min, max) = single_column_rows( + DataTypes::tinyint(), + ArrowType::Int8, + Arc::new(Int8Array::from(vec![Some(7), Some(-3)])), + ); + assert_eq!(only_slot(&min)[0] as i8, -3); + assert_eq!(only_slot(&max)[0] as i8, 7); + + let (min, max) = single_column_rows( + DataTypes::smallint(), + ArrowType::Int16, + Arc::new(Int16Array::from(vec![Some(300), Some(-300)])), + ); + assert_eq!(i16::from_le_bytes(min[8..10].try_into().unwrap()), -300); + assert_eq!(i16::from_le_bytes(max[8..10].try_into().unwrap()), 300); + } + + #[test] + fn collects_bounds_for_the_floating_types() { + let (min, max) = single_column_rows( + DataTypes::float(), + ArrowType::Float32, + Arc::new(Float32Array::from(vec![Some(2.5), Some(-1.5)])), + ); + assert_eq!(f32::from_le_bytes(min[8..12].try_into().unwrap()), -1.5); + assert_eq!(f32::from_le_bytes(max[8..12].try_into().unwrap()), 2.5); + + let (min, max) = single_column_rows( + DataTypes::double(), + ArrowType::Float64, + Arc::new(Float64Array::from(vec![Some(2.5), Some(-1.5)])), + ); + assert_eq!(f64::from_le_bytes(min[8..16].try_into().unwrap()), -1.5); + assert_eq!(f64::from_le_bytes(max[8..16].try_into().unwrap()), 2.5); + } + + #[test] + fn collects_bounds_for_char_like_a_string() { + let (min, max) = single_column_rows( + DataTypes::char(2), + ArrowType::Utf8, + Arc::new(StringArray::from(vec![Some("bb"), Some("aa")])), + ); + // Two bytes inline, with the length marker in the slot's top byte. + assert_eq!(&only_slot(&min)[..2], b"aa"); + assert_eq!(only_slot(&min)[7], 0x82); + assert_eq!(&only_slot(&max)[..2], b"bb"); + } + + #[test] + fn collects_bounds_for_date_as_epoch_days() { + let (min, max) = single_column_rows( + DataTypes::date(), + ArrowType::Date32, + Arc::new(Date32Array::from(vec![Some(19_000), Some(18_000)])), + ); + assert_eq!(i32::from_le_bytes(min[8..12].try_into().unwrap()), 18_000); + assert_eq!(i32::from_le_bytes(max[8..12].try_into().unwrap()), 19_000); + } + + /// The Arrow array holds seconds at precision 0, but the statistics format + /// is always millis of day. + #[test] + fn scales_a_second_precision_time_to_millis() { + let (min, max) = single_column_rows( + DataTypes::time_with_precision(0), + ArrowType::Time32(arrow::datatypes::TimeUnit::Second), + Arc::new(Time32SecondArray::from(vec![Some(7_200), Some(3_600)])), + ); + assert_eq!( + i32::from_le_bytes(min[8..12].try_into().unwrap()), + 3_600_000 + ); + assert_eq!( + i32::from_le_bytes(max[8..12].try_into().unwrap()), + 7_200_000 + ); + } + + #[test] + fn collects_bounds_for_time_as_millis_of_day() { + let (min, max) = single_column_rows( + DataTypes::time(), + ArrowType::Time32(arrow::datatypes::TimeUnit::Millisecond), + Arc::new(Time32MillisecondArray::from(vec![ + Some(7_200_000), + Some(3_600_000), + ])), + ); + assert_eq!( + i32::from_le_bytes(min[8..12].try_into().unwrap()), + 3_600_000 + ); + assert_eq!( + i32::from_le_bytes(max[8..12].try_into().unwrap()), + 7_200_000 + ); + } + + #[test] + fn scales_a_microsecond_time_to_millis() { + let (min, max) = single_column_rows( + DataTypes::time_with_precision(6), + ArrowType::Time64(arrow::datatypes::TimeUnit::Microsecond), + Arc::new(Time64MicrosecondArray::from(vec![ + Some(7_200_000_000), + Some(3_600_000_000), + ])), + ); + assert_eq!( + i32::from_le_bytes(min[8..12].try_into().unwrap()), + 3_600_000 + ); + assert_eq!( + i32::from_le_bytes(max[8..12].try_into().unwrap()), + 7_200_000 + ); + } + + #[test] + fn scales_a_nanosecond_time_to_millis() { + let (min, max) = single_column_rows( + DataTypes::time_with_precision(9), + ArrowType::Time64(arrow::datatypes::TimeUnit::Nanosecond), + Arc::new(Time64NanosecondArray::from(vec![ + Some(7_200_000_000_000), + Some(3_600_000_000_000), + ])), + ); + assert_eq!( + i32::from_le_bytes(min[8..12].try_into().unwrap()), + 3_600_000 + ); + assert_eq!( + i32::from_le_bytes(max[8..12].try_into().unwrap()), + 7_200_000 + ); + } + + #[test] + fn scales_a_second_precision_timestamp_to_millis() { + let (min, max) = single_column_rows( + DataTypes::timestamp_with_precision(0), + ArrowType::Timestamp(arrow::datatypes::TimeUnit::Second, None), + Arc::new(TimestampSecondArray::from(vec![Some(2), Some(1)])), + ); + // Precision 0 is compact, so the millis sit in the slot. + assert_eq!(i64::from_le_bytes(min[8..16].try_into().unwrap()), 1_000); + assert_eq!(i64::from_le_bytes(max[8..16].try_into().unwrap()), 2_000); + } + + /// Decodes a non-compact timestamp field into its (millis, nanos) pair. + fn split_timestamp(row: &[u8]) -> (i64, i32) { + let packed = i64::from_le_bytes(row[8..16].try_into().unwrap()); + let offset = (packed >> 32) as usize; + let nanos = (packed & 0xFFFF_FFFF) as i32; + let millis = i64::from_le_bytes(row[offset..offset + 8].try_into().unwrap()); + (millis, nanos) + } + + #[test] + fn splits_a_nanosecond_timestamp_into_millis_and_nanos() { + let (min, max) = single_column_rows( + DataTypes::timestamp_with_precision(9), + ArrowType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), + Arc::new(TimestampNanosecondArray::from(vec![ + Some(2_000_456_789), + Some(1_000_654_321), + ])), + ); + assert_eq!(split_timestamp(&min), (1_000, 654_321)); + assert_eq!(split_timestamp(&max), (2_000, 456_789)); + } + + #[test] + fn spills_a_non_compact_decimal_bound_to_the_tail() { + let array = Decimal128Array::from(vec![Some(555_000_i128), Some(100_000_i128)]) + .with_precision_and_scale(25, 5) + .expect("decimal array"); + let (min, _) = single_column_rows( + DataTypes::decimal(25, 5), + ArrowType::Decimal128(25, 5), + Arc::new(array), + ); + // Precision 25 is not compact, so the slot points into the tail. + let packed = i64::from_le_bytes(min[8..16].try_into().unwrap()); + let (offset, size) = ((packed >> 32) as usize, (packed & 0xFFFF_FFFF) as usize); + assert_eq!(offset, 16); + let unscaled = Decimal::from_unscaled_bytes(&min[offset..offset + size], 25, 5) + .expect("decimal") + .to_big_decimal(); + assert_eq!(unscaled.to_string(), "1.00000"); + } + + #[test] + fn collects_bounds_for_a_compact_decimal() { + let array = Decimal128Array::from(vec![Some(12_345_i128), Some(500_i128)]) + .with_precision_and_scale(10, 2) + .expect("decimal array"); + let (min, max) = single_column_rows( + DataTypes::decimal(10, 2), + ArrowType::Decimal128(10, 2), + Arc::new(array), + ); + // Precision 10 is compact, so the unscaled value sits in the slot. + assert_eq!(i64::from_le_bytes(min[8..16].try_into().unwrap()), 500); + assert_eq!(i64::from_le_bytes(max[8..16].try_into().unwrap()), 12_345); + } + + #[test] + fn keeps_a_millisecond_timestamp_in_its_slot() { + let (min, max) = single_column_rows( + DataTypes::timestamp_with_precision(3), + ArrowType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, None), + Arc::new(TimestampMillisecondArray::from(vec![ + Some(2_000), + Some(1_000), + ])), + ); + assert_eq!(i64::from_le_bytes(min[8..16].try_into().unwrap()), 1_000); + assert_eq!(i64::from_le_bytes(max[8..16].try_into().unwrap()), 2_000); + } + + #[test] + fn splits_a_microsecond_timestamp_into_millis_and_nanos() { + let (min, max) = single_column_rows( + DataTypes::timestamp_with_precision(6), + ArrowType::Timestamp(arrow::datatypes::TimeUnit::Microsecond, None), + Arc::new(TimestampMicrosecondArray::from(vec![ + Some(2_000_500), + Some(1_000_456), + ])), + ); + // Precision 6 is not compact, so millis move to the tail and the slot + // carries the offset with the nano-of-millisecond. + assert_eq!(split_timestamp(&min), (1_000, 456_000)); + assert_eq!(split_timestamp(&max), (2_000, 500_000)); + } + + #[test] + fn collects_bounds_for_a_local_zoned_timestamp() { + let array = + TimestampMillisecondArray::from(vec![Some(2_000), Some(1_000)]).with_timezone("UTC"); + let (min, max) = single_column_rows( + DataTypes::timestamp_ltz_with_precision(3), + ArrowType::Timestamp(arrow::datatypes::TimeUnit::Millisecond, Some("UTC".into())), + Arc::new(array), + ); + assert_eq!(i64::from_le_bytes(min[8..16].try_into().unwrap()), 1_000); + assert_eq!(i64::from_le_bytes(max[8..16].try_into().unwrap()), 2_000); + } +} diff --git a/fluss-rust/crates/fluss/src/row/aligned/aligned_row_writer.rs b/fluss-rust/crates/fluss/src/row/aligned/aligned_row_writer.rs new file mode 100644 index 00000000000..4228701bedb --- /dev/null +++ b/fluss-rust/crates/fluss/src/row/aligned/aligned_row_writer.rs @@ -0,0 +1,623 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Writer for Fluss's aligned binary row, the format carrying the min and max +//! values of a V1 log record batch's statistics. + +use crate::row::binary::BinaryWriter; +use crate::row::datum::{TimestampLtz, TimestampNtz}; +use crate::row::{Decimal, FlussArray, FlussMap}; +use bytes::Bytes; + +/// Bits reserved ahead of the null bitset, matching Java's +/// `AlignedRow.HEADER_SIZE_IN_BITS`. +const HEADER_SIZE_IN_BITS: usize = 8; + +/// Longest payload that still fits inside an 8-byte field slot alongside its +/// length marker. +const MAX_FIX_PART_DATA_SIZE: usize = 7; + +/// Builds Java's `AlignedRow` byte for byte: a fixed part of null bits plus one +/// 8-byte slot per field, then an 8-byte-aligned variable-length part. +/// +/// Laying out `(id INT = 7, name STRING = "hello world", score BIGINT = 42)`: +/// +/// ```text +/// byte 0 8 16 24 32 43 48 +/// |------------|---------|---------|---------|-----------|----| +/// | null bits | slot 0 | slot 1 | slot 2 | "hello world" | +/// | (8 bytes) | int 7 | ptr | long 42 | + zero padding | +/// |------------|---------|---------|---------|----------------| +/// <--------- fixed part (32) ----------> <-- variable tail --> +/// | ^ +/// +-- (offset=32, len=11) packed into 8 bytes +/// ``` +/// +/// The null bits reserve 8 header bits, so field `i` owns bit `i + 8`, and the +/// region is padded to whole 8-byte words. A value of 8 bytes or less lives in +/// its slot, while anything longer stores `(offset << 32) | len` there and puts +/// its payload in the tail. +/// +/// Fields must be written in order; each `write_*` consumes the next position. +pub struct AlignedRowWriter { + buffer: Vec, + null_bits_size_in_bytes: usize, + /// Size of the null bits plus one 8-byte slot per field, which is also + /// where the variable-length part starts. + fixed_size: usize, + /// Byte offset where the next value too large for its slot gets appended. + cursor: usize, + /// Index of the field the next `write_*` fills, which picks its 8-byte slot. + current_pos: usize, +} + +impl AlignedRowWriter { + pub fn new(arity: usize) -> Self { + let null_bits_size_in_bytes = calculate_bit_set_width_in_bytes(arity); + let fixed_size = null_bits_size_in_bytes + 8 * arity; + Self { + buffer: vec![0u8; fixed_size], + null_bits_size_in_bytes, + fixed_size, + cursor: fixed_size, + current_pos: 0, + } + } + + pub fn to_bytes(&self) -> Bytes { + Bytes::copy_from_slice(&self.buffer[..self.cursor]) + } + + pub fn size_in_bytes(&self) -> usize { + self.cursor + } + + fn field_offset(&self, pos: usize) -> usize { + self.null_bits_size_in_bytes + 8 * pos + } + + fn set_null_bit(&mut self, pos: usize) { + let bit = pos + HEADER_SIZE_IN_BITS; + self.buffer[bit / 8] |= 1u8 << (bit % 8); + } + + fn put_long_le(&mut self, offset: usize, value: i64) { + self.buffer[offset..offset + 8].copy_from_slice(&value.to_le_bytes()); + } + + fn put_int_le(&mut self, offset: usize, value: i32) { + self.buffer[offset..offset + 4].copy_from_slice(&value.to_le_bytes()); + } + + fn put_short_le(&mut self, offset: usize, value: i16) { + self.buffer[offset..offset + 2].copy_from_slice(&value.to_le_bytes()); + } + + /// Packs `(offset << 32) | size` into the field slot, how the reader finds a + /// value that lives in the variable-length part. + /// + /// Timestamps pass the nano-of-millisecond as `size`, since their payload is + /// always 8 bytes and the low half would otherwise go to waste. + fn set_offset_and_size(&mut self, pos: usize, offset: usize, size: u64) { + let packed = ((offset as i64) << 32) | (size as i64); + let field_offset = self.field_offset(pos); + self.put_long_le(field_offset, packed); + } + + /// Inlines a payload of at most 7 bytes into the field slot, with + /// `len | 0x80` in the slot's high byte as Java's `writeBytesToFixLenPart`. + fn write_bytes_to_fix_len_part(&mut self, pos: usize, bytes: &[u8]) { + let len = bytes.len(); + debug_assert!(len <= MAX_FIX_PART_DATA_SIZE); + let field_offset = self.field_offset(pos); + self.put_long_le(field_offset, 0); + self.buffer[field_offset..field_offset + len].copy_from_slice(bytes); + self.buffer[field_offset + 7] = (len as u8) | 0x80; + } + + fn ensure_capacity(&mut self, needed_size: usize) { + let length = self.cursor + needed_size; + if self.buffer.len() < length { + let old_capacity = self.buffer.len(); + let new_capacity = (old_capacity + (old_capacity >> 1)).max(length); + self.buffer.resize(new_capacity, 0); + } + } + + /// Zeroes the tail of the word a value only partly fills, so the padding is + /// deterministic rather than whatever the buffer last held. + fn zero_out_padding_bytes(&mut self, num_bytes: usize) { + if (num_bytes & 0x07) > 0 { + let off = self.cursor + ((num_bytes >> 3) << 3); + for b in &mut self.buffer[off..off + 8] { + *b = 0; + } + } + } + + fn write_bytes_to_var_len_part(&mut self, pos: usize, bytes: &[u8]) { + let len = bytes.len(); + let rounded_size = round_number_of_bytes_to_nearest_word(len); + + self.ensure_capacity(rounded_size); + self.zero_out_padding_bytes(len); + self.buffer[self.cursor..self.cursor + len].copy_from_slice(bytes); + self.set_offset_and_size(pos, self.cursor, len as u64); + self.cursor += rounded_size; + } + + fn write_bytes_internal(&mut self, pos: usize, bytes: &[u8]) { + if bytes.len() <= MAX_FIX_PART_DATA_SIZE { + self.write_bytes_to_fix_len_part(pos, bytes); + } else { + self.write_bytes_to_var_len_part(pos, bytes); + } + } +} + +/// Null bits are padded to whole 8-byte words, after the reserved header bits. +fn calculate_bit_set_width_in_bytes(arity: usize) -> usize { + ((arity + 63 + HEADER_SIZE_IN_BITS) / 64) * 8 +} + +fn round_number_of_bytes_to_nearest_word(num_bytes: usize) -> usize { + let remainder = num_bytes & 0x07; + if remainder == 0 { + num_bytes + } else { + num_bytes + (8 - remainder) + } +} + +impl BinaryWriter for AlignedRowWriter { + fn reset(&mut self) { + self.cursor = self.fixed_size; + self.current_pos = 0; + for b in &mut self.buffer[..self.fixed_size] { + *b = 0; + } + } + + fn set_null_at(&mut self, pos: usize) { + self.set_null_bit(pos); + let field_offset = self.field_offset(pos); + self.put_long_le(field_offset, 0); + self.current_pos = pos + 1; + } + + fn write_boolean(&mut self, value: bool) { + let off = self.field_offset(self.current_pos); + self.put_long_le(off, 0); + self.buffer[off] = u8::from(value); + self.current_pos += 1; + } + + fn write_byte(&mut self, value: u8) { + let off = self.field_offset(self.current_pos); + self.put_long_le(off, 0); + self.buffer[off] = value; + self.current_pos += 1; + } + + fn write_bytes(&mut self, value: &[u8]) { + let pos = self.current_pos; + self.write_bytes_internal(pos, value); + self.current_pos = pos + 1; + } + + fn write_char(&mut self, value: &str, _length: usize) { + self.write_string(value); + } + + fn write_string(&mut self, value: &str) { + let pos = self.current_pos; + self.write_bytes_internal(pos, value.as_bytes()); + self.current_pos = pos + 1; + } + + fn write_short(&mut self, value: i16) { + let off = self.field_offset(self.current_pos); + self.put_long_le(off, 0); + self.put_short_le(off, value); + self.current_pos += 1; + } + + fn write_int(&mut self, value: i32) { + let off = self.field_offset(self.current_pos); + self.put_long_le(off, 0); + self.put_int_le(off, value); + self.current_pos += 1; + } + + fn write_long(&mut self, value: i64) { + let off = self.field_offset(self.current_pos); + self.put_long_le(off, value); + self.current_pos += 1; + } + + fn write_float(&mut self, value: f32) { + let off = self.field_offset(self.current_pos); + self.put_long_le(off, 0); + self.buffer[off..off + 4].copy_from_slice(&value.to_le_bytes()); + self.current_pos += 1; + } + + fn write_double(&mut self, value: f64) { + let off = self.field_offset(self.current_pos); + self.buffer[off..off + 8].copy_from_slice(&value.to_le_bytes()); + self.current_pos += 1; + } + + fn write_binary(&mut self, bytes: &[u8], length: usize) { + let pos = self.current_pos; + let slice = &bytes[..length.min(bytes.len())]; + self.write_bytes_internal(pos, slice); + self.current_pos = pos + 1; + } + + fn write_decimal(&mut self, value: &Decimal, precision: u32) { + assert_eq!( + value.precision(), + precision, + "decimal was built at a different precision than the column's" + ); + let pos = self.current_pos; + if Decimal::is_compact_precision(precision) { + let unscaled = value + .to_unscaled_long() + .expect("a compact precision guarantees the unscaled value fits in i64"); + let off = self.field_offset(pos); + self.put_long_le(off, unscaled); + } else { + // Java always reserves 16 bytes here, whatever the unscaled length. + self.ensure_capacity(16); + for b in &mut self.buffer[self.cursor..self.cursor + 16] { + *b = 0; + } + let bytes = value.to_unscaled_bytes(); + debug_assert!(bytes.len() <= 16, "decimal unscaled bytes exceed 16"); + self.buffer[self.cursor..self.cursor + bytes.len()].copy_from_slice(&bytes); + self.set_offset_and_size(pos, self.cursor, bytes.len() as u64); + self.cursor += 16; + } + self.current_pos = pos + 1; + } + + fn write_time(&mut self, value: i32, _precision: u32) { + self.write_int(value); + } + + fn write_timestamp_ntz(&mut self, value: &TimestampNtz, precision: u32) { + let pos = self.current_pos; + if TimestampNtz::is_compact(precision) { + let off = self.field_offset(pos); + self.put_long_le(off, value.get_millisecond()); + } else { + self.ensure_capacity(8); + self.put_long_le(self.cursor, value.get_millisecond()); + self.set_offset_and_size(pos, self.cursor, value.get_nano_of_millisecond() as u64); + self.cursor += 8; + } + self.current_pos = pos + 1; + } + + fn write_timestamp_ltz(&mut self, value: &TimestampLtz, precision: u32) { + let pos = self.current_pos; + if TimestampLtz::is_compact(precision) { + let off = self.field_offset(pos); + self.put_long_le(off, value.get_epoch_millisecond()); + } else { + self.ensure_capacity(8); + self.put_long_le(self.cursor, value.get_epoch_millisecond()); + self.set_offset_and_size(pos, self.cursor, value.get_nano_of_millisecond() as u64); + self.cursor += 8; + } + self.current_pos = pos + 1; + } + + fn write_array(&mut self, _value: &FlussArray) { + panic!("statistics are never collected for ARRAY columns"); + } + + fn write_map(&mut self, _value: &FlussMap) { + panic!("statistics are never collected for MAP columns"); + } + + fn complete(&mut self) { + // `to_bytes` already trims to the cursor. + } +} + +#[cfg(test)] +mod tests { + use super::*; + use bigdecimal::BigDecimal; + use std::str::FromStr; + + #[test] + fn fixed_part_matches_java_sizing() { + // arity 4 -> ceil((4 + 63 + 8) / 64) * 8 = 8 null bytes, + 8 per field. + let writer = AlignedRowWriter::new(4); + assert_eq!(writer.null_bits_size_in_bytes, 8); + assert_eq!(writer.fixed_size, 40); + assert_eq!(writer.cursor, 40); + + // 57 fields cross into a second word of null bits. + assert_eq!(calculate_bit_set_width_in_bytes(56), 8); + assert_eq!(calculate_bit_set_width_in_bytes(57), 16); + } + + #[test] + fn writes_int_little_endian_into_its_slot() { + let mut writer = AlignedRowWriter::new(1); + writer.write_int(0x01020304); + let bytes = writer.to_bytes(); + assert_eq!(bytes.len(), 16); + assert_eq!(&bytes[8..12], &0x01020304_i32.to_le_bytes()); + assert_eq!(&bytes[12..16], &[0u8; 4]); + } + + #[test] + fn inlines_a_short_string_with_its_length_marker() { + let mut writer = AlignedRowWriter::new(1); + writer.write_string("abc"); + let bytes = writer.to_bytes(); + // No variable part is needed, so the row is just the fixed part. + assert_eq!(bytes.len(), 16); + assert_eq!(&bytes[8..11], b"abc"); + assert_eq!(bytes[15], 0x83); + } + + #[test] + fn spills_a_long_string_to_the_variable_part() { + let mut writer = AlignedRowWriter::new(1); + writer.write_string("abcdefghij"); + let bytes = writer.to_bytes(); + // 10 bytes of payload round up to two 8-byte words. + assert_eq!(bytes.len(), 16 + 16); + let packed = i64::from_le_bytes(bytes[8..16].try_into().unwrap()); + assert_eq!((packed >> 32) as usize, 16); + assert_eq!((packed & 0xFFFF_FFFF) as usize, 10); + assert_eq!(&bytes[16..26], b"abcdefghij"); + // The padding to the word boundary is zeroed. + assert_eq!(&bytes[26..32], &[0u8; 6]); + } + + #[test] + fn marks_null_without_disturbing_other_fields() { + let mut writer = AlignedRowWriter::new(2); + writer.set_null_at(0); + writer.write_long(7); + let bytes = writer.to_bytes(); + // Field 0's null bit sits at bit 8, the first bit after the header. + assert_eq!(bytes[1], 0x01); + assert_eq!(&bytes[8..16], &[0u8; 8]); + assert_eq!(i64::from_le_bytes(bytes[16..24].try_into().unwrap()), 7); + } + + #[test] + fn reset_clears_the_fixed_part_and_rewinds() { + let mut writer = AlignedRowWriter::new(1); + writer.write_string("abcdefghij"); + writer.reset(); + writer.write_int(5); + let bytes = writer.to_bytes(); + assert_eq!(bytes.len(), 16); + assert_eq!(&bytes[8..12], &5_i32.to_le_bytes()); + } + + /// Reads the 8-byte slot of field `pos` as the packed `(offset, size)` pair + /// used for values that live in the variable-length tail. + fn packed_slot(bytes: &[u8], null_bits: usize, pos: usize) -> (usize, usize) { + let at = null_bits + 8 * pos; + let packed = i64::from_le_bytes(bytes[at..at + 8].try_into().unwrap()); + ((packed >> 32) as usize, (packed & 0xFFFF_FFFF) as usize) + } + + /// Slices out the 8-byte fixed-part slot belonging to field `pos`. + fn slot(bytes: &[u8], null_bits: usize, pos: usize) -> &[u8] { + let at = null_bits + 8 * pos; + &bytes[at..at + 8] + } + + #[test] + fn writes_each_numeric_width_into_the_low_bytes_of_its_slot() { + let mut writer = AlignedRowWriter::new(7); + writer.write_boolean(true); + writer.write_byte(0xAB); + writer.write_short(-2); + writer.write_int(-3); + writer.write_long(-4); + writer.write_float(1.5); + writer.write_double(2.5); + let bytes = writer.to_bytes(); + + assert_eq!(slot(&bytes, 8, 0), &[1, 0, 0, 0, 0, 0, 0, 0]); + assert_eq!(slot(&bytes, 8, 1), &[0xAB, 0, 0, 0, 0, 0, 0, 0]); + assert_eq!(&slot(&bytes, 8, 2)[..2], &(-2i16).to_le_bytes()); + assert_eq!(&slot(&bytes, 8, 3)[..4], &(-3i32).to_le_bytes()); + assert_eq!(slot(&bytes, 8, 4), &(-4i64).to_le_bytes()); + assert_eq!(&slot(&bytes, 8, 5)[..4], &1.5f32.to_le_bytes()); + assert_eq!(slot(&bytes, 8, 6), &2.5f64.to_le_bytes()); + // The unused high bytes of the narrower slots are zeroed. + assert_eq!(&slot(&bytes, 8, 2)[2..], &[0u8; 6]); + assert_eq!(&slot(&bytes, 8, 5)[4..], &[0u8; 4]); + } + + #[test] + fn keeps_a_compact_decimal_in_its_slot() { + let precision = 4; + let decimal = Decimal::from_unscaled_long(5, precision, 2).expect("decimal"); + let mut writer = AlignedRowWriter::new(2); + writer.write_decimal(&decimal, precision); + writer.set_null_at(1); + let bytes = writer.to_bytes(); + + // Precision 4 is compact, so the unscaled value sits inline. + assert_eq!(bytes.len(), 24); + assert_eq!(slot(&bytes, 8, 0), &5i64.to_le_bytes()); + // Field 1's null bit is bit 9, so byte 1 bit 1. + assert_eq!(bytes[1] & 0x02, 0x02); + } + + #[test] + fn spills_a_non_compact_decimal_into_sixteen_tail_bytes() { + let precision = 25; + let decimal = + Decimal::from_big_decimal(BigDecimal::from_str("5.55").unwrap(), precision, 5) + .expect("decimal"); + let unscaled = decimal.to_unscaled_bytes(); + + let mut writer = AlignedRowWriter::new(1); + writer.write_decimal(&decimal, precision); + let bytes = writer.to_bytes(); + + // Java always reserves 16 tail bytes here, whatever the unscaled length. + assert_eq!(bytes.len(), 16 + 16); + let (offset, size) = packed_slot(&bytes, 8, 0); + assert_eq!(offset, 16); + assert_eq!(size, unscaled.len()); + assert_eq!(&bytes[16..16 + unscaled.len()], &unscaled[..]); + // The unscaled value rarely fills all 16, so the rest must be zeroed. + assert_eq!( + &bytes[16 + unscaled.len()..32], + &vec![0u8; 16 - unscaled.len()][..] + ); + } + + #[test] + #[should_panic(expected = "assertion")] + fn rejects_a_decimal_built_at_another_precision() { + let decimal = Decimal::from_big_decimal(BigDecimal::from_str("5.55").unwrap(), 25, 5) + .expect("decimal"); + let mut writer = AlignedRowWriter::new(1); + writer.write_decimal(&decimal, 10); + } + + #[test] + fn keeps_a_compact_timestamp_in_its_slot() { + let value = TimestampNtz::from_millis_nanos(123, 0).expect("timestamp"); + let mut writer = AlignedRowWriter::new(1); + writer.write_timestamp_ntz(&value, 3); + let bytes = writer.to_bytes(); + + assert_eq!(bytes.len(), 16); + assert_eq!(slot(&bytes, 8, 0), &123i64.to_le_bytes()); + } + + #[test] + fn splits_a_non_compact_timestamp_between_slot_and_tail() { + let value = TimestampNtz::from_millis_nanos(123, 456_000).expect("timestamp"); + let mut writer = AlignedRowWriter::new(1); + writer.write_timestamp_ntz(&value, 6); + let bytes = writer.to_bytes(); + + // Millis go to the tail; the slot carries the offset and the nanos. + // Timestamps are the one case where the packed low half is not a + // length, since the tail is always exactly 8 bytes. + assert_eq!(bytes.len(), 16 + 8); + let (offset, nanos) = packed_slot(&bytes, 8, 0); + assert_eq!(offset, 16); + assert_eq!(nanos, 456_000); + assert_eq!(i64::from_le_bytes(bytes[16..24].try_into().unwrap()), 123); + } + + #[test] + fn encodes_a_local_zoned_timestamp_the_same_way() { + let compact = TimestampLtz::from_millis_nanos(99, 0).expect("timestamp"); + let mut writer = AlignedRowWriter::new(1); + writer.write_timestamp_ltz(&compact, 3); + assert_eq!(slot(&writer.to_bytes(), 8, 0), &99i64.to_le_bytes()); + + let wide = TimestampLtz::from_millis_nanos(99, 1_000).expect("timestamp"); + let mut writer = AlignedRowWriter::new(1); + writer.write_timestamp_ltz(&wide, 9); + let bytes = writer.to_bytes(); + let (offset, nanos) = packed_slot(&bytes, 8, 0); + assert_eq!((offset, nanos), (16, 1_000)); + assert_eq!(i64::from_le_bytes(bytes[16..24].try_into().unwrap()), 99); + } + + #[test] + fn inlines_short_binary_and_spills_longer_binary() { + let mut writer = AlignedRowWriter::new(2); + writer.write_bytes(&[1, 0xFF, 5]); + writer.write_bytes(&[1, 0xFF, 5, 5, 1, 5, 1, 5]); + let bytes = writer.to_bytes(); + + // Three bytes fit inline with the length marker in the slot's top byte. + assert_eq!(&slot(&bytes, 8, 0)[..3], &[1, 0xFF, 5]); + assert_eq!(slot(&bytes, 8, 0)[7], 0x83); + // Eight bytes exceed the seven that fit, so they move to the tail. + let (offset, size) = packed_slot(&bytes, 8, 1); + assert_eq!((offset, size), (24, 8)); + assert_eq!(&bytes[24..32], &[1, 0xFF, 5, 5, 1, 5, 1, 5]); + } + + #[test] + fn writes_char_like_a_string() { + let mut writer = AlignedRowWriter::new(1); + writer.write_char("ab", 2); + let bytes = writer.to_bytes(); + assert_eq!(&slot(&bytes, 8, 0)[..2], b"ab"); + assert_eq!(slot(&bytes, 8, 0)[7], 0x82); + } + + #[test] + fn tracks_null_bits_across_a_second_word() { + // 60 fields need two 8-byte words of null bits. + let mut writer = AlignedRowWriter::new(60); + assert_eq!(writer.null_bits_size_in_bytes, 16); + for pos in 0..60 { + if pos == 56 { + writer.set_null_at(pos); + } else { + writer.write_int(pos as i32); + } + } + let bytes = writer.to_bytes(); + + // Field 56 owns bit 56 + 8 = 64, which is byte 8 bit 0, so the null + // landed in the first byte of the second word rather than overflowing + // the first. + assert_eq!(bytes[8], 0x01); + // The first word stays clear: no other field is null and the 8 reserved + // header bits are never set. + assert_eq!(&bytes[..8], &[0u8; 8]); + // Slots start after both words, so the last field sits at 16 + 8 * 59 + // rather than the 8 + 8 * 59 a single-word row would use. + assert_eq!(&slot(&bytes, 16, 59)[..4], &59i32.to_le_bytes()); + } + + #[test] + fn grows_the_buffer_across_many_spilled_values() { + let mut writer = AlignedRowWriter::new(8); + let long = "0123456789abcdef"; + for _ in 0..8 { + writer.write_string(long); + } + let bytes = writer.to_bytes(); + + // Eight 16-byte payloads follow the 8 + 64 byte fixed part. + assert_eq!(bytes.len(), 72 + 8 * 16); + for pos in 0..8 { + let (offset, size) = packed_slot(&bytes, 8, pos); + assert_eq!(size, 16); + assert_eq!(&bytes[offset..offset + 16], long.as_bytes()); + } + } +} diff --git a/fluss-rust/crates/fluss/src/row/aligned/mod.rs b/fluss-rust/crates/fluss/src/row/aligned/mod.rs new file mode 100644 index 00000000000..4ccb5b760d8 --- /dev/null +++ b/fluss-rust/crates/fluss/src/row/aligned/mod.rs @@ -0,0 +1,20 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod aligned_row_writer; + +pub use aligned_row_writer::AlignedRowWriter; diff --git a/fluss-rust/crates/fluss/src/row/mod.rs b/fluss-rust/crates/fluss/src/row/mod.rs index 7b483b4ebe4..5d606deae13 100644 --- a/fluss-rust/crates/fluss/src/row/mod.rs +++ b/fluss-rust/crates/fluss/src/row/mod.rs @@ -25,6 +25,7 @@ pub mod view; pub(crate) mod datum; mod decimal; +pub mod aligned; pub mod binary; pub(crate) mod column_writer; pub mod compacted;