| // 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. |
| |
| //! TsBlock decoder — the binary query-result format returned in |
| //! `queryResult` by the V2 RPCs (protocol spec §5). All integers big-endian. |
| |
| use super::bitmap::unpack_bits_msb_first; |
| use super::value::Value; |
| use super::TSDataType; |
| use crate::error::{Error, Result}; |
| |
| /// TsBlock column encodings (spec §5.2). `pub(crate)` so the dataset tests |
| /// can craft synthetic blocks with the same constants. |
| pub(crate) const ENCODING_BYTE_ARRAY: u8 = 0; |
| pub(crate) const ENCODING_INT32_ARRAY: u8 = 1; |
| pub(crate) const ENCODING_INT64_ARRAY: u8 = 2; |
| pub(crate) const ENCODING_BINARY_ARRAY: u8 = 3; |
| pub(crate) const ENCODING_RLE: u8 = 4; |
| |
| /// One decoded TsBlock: `position_count` rows of a time column plus |
| /// `columns.len()` value columns. Null cells are [`Value::Null`]. |
| #[derive(Debug, Clone, PartialEq)] |
| pub struct TsBlock { |
| pub position_count: usize, |
| pub column_types: Vec<TSDataType>, |
| /// Time column, always present on the wire — the caller drops it when |
| /// `ignoreTimeStamp` is set (spec gotcha #10). |
| pub timestamps: Vec<i64>, |
| /// `columns[col][row]`; nulls decoded as [`Value::Null`]. |
| pub columns: Vec<Vec<Value>>, |
| } |
| |
| impl TsBlock { |
| /// Decodes one serialized TsBlock (one element of `queryResult`). |
| /// |
| /// Layout (spec §5.1): `valueColumnCount` i32, per-column type bytes, |
| /// `positionCount` i32, time-column encoding byte, per-column encoding |
| /// bytes, then the time column followed by each value column. |
| pub fn decode(bytes: &[u8]) -> Result<TsBlock> { |
| let mut r = Reader::new(bytes); |
| |
| let value_column_count = r.read_i32()?; |
| if value_column_count < 0 { |
| return Err(Error::Decode(format!( |
| "negative valueColumnCount: {value_column_count}" |
| ))); |
| } |
| let value_column_count = value_column_count as usize; |
| |
| let mut column_types = Vec::with_capacity(value_column_count); |
| for _ in 0..value_column_count { |
| let code = r.read_u8()?; |
| column_types.push( |
| TSDataType::from_code(code) |
| .ok_or_else(|| Error::Decode(format!("unknown TSDataType code {code}")))?, |
| ); |
| } |
| |
| let position_count = r.read_i32()?; |
| if position_count < 0 { |
| return Err(Error::Decode(format!( |
| "negative positionCount: {position_count}" |
| ))); |
| } |
| let position_count = position_count as usize; |
| |
| let time_encoding = r.read_u8()?; |
| let mut value_encodings = Vec::with_capacity(value_column_count); |
| for _ in 0..value_column_count { |
| value_encodings.push(r.read_u8()?); |
| } |
| |
| // Time column is always present (gotcha #10); decoded as INT64. |
| let time_column = decode_column(&mut r, time_encoding, TSDataType::Int64, position_count)?; |
| let mut timestamps = Vec::with_capacity(position_count); |
| for v in time_column { |
| match v { |
| Value::Int64(t) => timestamps.push(t), |
| other => { |
| return Err(Error::Decode(format!( |
| "time column contains non-INT64 value: {other:?}" |
| ))) |
| } |
| } |
| } |
| |
| let mut columns = Vec::with_capacity(value_column_count); |
| for (i, &ty) in column_types.iter().enumerate() { |
| columns.push(decode_column( |
| &mut r, |
| value_encodings[i], |
| ty, |
| position_count, |
| )?); |
| } |
| |
| Ok(TsBlock { |
| position_count, |
| column_types, |
| timestamps, |
| columns, |
| }) |
| } |
| } |
| |
| /// Decodes one column body: null-indicator section, then encoding-specific |
| /// values (spec §5.3). |
| fn decode_column( |
| r: &mut Reader<'_>, |
| encoding: u8, |
| ty: TSDataType, |
| position_count: usize, |
| ) -> Result<Vec<Value>> { |
| match encoding { |
| ENCODING_RLE => { |
| // Inner-encoding byte, decode ONE value, replicate. |
| let inner = r.read_u8()?; |
| if inner == ENCODING_RLE { |
| return Err(Error::Decode("nested RLE column encoding".into())); |
| } |
| let single = decode_column(r, inner, ty, 1)?; |
| let v = single |
| .into_iter() |
| .next() |
| .ok_or_else(|| Error::Decode("empty RLE inner column".into()))?; |
| Ok(vec![v; position_count]) |
| } |
| ENCODING_BYTE_ARRAY => { |
| let nulls = read_null_indicators(r, position_count)?; |
| // BOOLEAN values: full positionCount bit array (MSB-first), |
| // NOT null-skipped — unlike the numeric encodings. |
| let bits = read_packed_bits(r, position_count)?; |
| Ok((0..position_count) |
| .map(|i| { |
| if nulls[i] { |
| Value::Null |
| } else { |
| Value::Boolean(bits[i]) |
| } |
| }) |
| .collect()) |
| } |
| ENCODING_INT32_ARRAY => { |
| let nulls = read_null_indicators(r, position_count)?; |
| let mut out = Vec::with_capacity(position_count); |
| for &is_null in &nulls { |
| if is_null { |
| out.push(Value::Null); |
| continue; // dense: null positions consume no bytes |
| } |
| let raw = r.read_i32()?; |
| out.push(match ty { |
| TSDataType::Int32 => Value::Int32(raw), |
| TSDataType::Date => Value::Date(raw), |
| TSDataType::Float => Value::Float(f32::from_bits(raw as u32)), |
| _ => { |
| return Err(Error::Decode(format!( |
| "Int32Array encoding with incompatible type {ty:?}" |
| ))) |
| } |
| }); |
| } |
| Ok(out) |
| } |
| ENCODING_INT64_ARRAY => { |
| let nulls = read_null_indicators(r, position_count)?; |
| let mut out = Vec::with_capacity(position_count); |
| for &is_null in &nulls { |
| if is_null { |
| out.push(Value::Null); |
| continue; |
| } |
| let raw = r.read_i64()?; |
| out.push(match ty { |
| TSDataType::Int64 => Value::Int64(raw), |
| TSDataType::Timestamp => Value::Timestamp(raw), |
| TSDataType::Double => Value::Double(f64::from_bits(raw as u64)), |
| _ => { |
| return Err(Error::Decode(format!( |
| "Int64Array encoding with incompatible type {ty:?}" |
| ))) |
| } |
| }); |
| } |
| Ok(out) |
| } |
| ENCODING_BINARY_ARRAY => { |
| let nulls = read_null_indicators(r, position_count)?; |
| let mut out = Vec::with_capacity(position_count); |
| for &is_null in &nulls { |
| if is_null { |
| out.push(Value::Null); |
| continue; |
| } |
| let len = r.read_i32()?; |
| if len < 0 { |
| return Err(Error::Decode(format!("negative binary length: {len}"))); |
| } |
| let bytes = r.read_bytes(len as usize)?.to_vec(); |
| out.push(match ty { |
| TSDataType::Blob => Value::Blob(bytes), |
| // TsBlock headers carry the *physical* type: BLOB columns |
| // arrive typed TEXT here. Keep non-UTF-8 payloads as Blob |
| // so row assembly can re-tag them via the response's |
| // logical dataTypeList (same pattern as DATE/INT32). |
| TSDataType::Text => match String::from_utf8(bytes) { |
| Ok(s) => Value::Text(s), |
| Err(e) => Value::Blob(e.into_bytes()), |
| }, |
| TSDataType::String => Value::String(decode_utf8(bytes)?), |
| _ => { |
| return Err(Error::Decode(format!( |
| "BinaryArray encoding with incompatible type {ty:?}" |
| ))) |
| } |
| }); |
| } |
| Ok(out) |
| } |
| other => Err(Error::Decode(format!("unknown column encoding {other}"))), |
| } |
| } |
| |
| /// Null-indicator section: 1 byte `mayHaveNull`; if non-zero, a |
| /// `ceil(positionCount/8)`-byte MSB-first bitmap with bit=1 ⇒ null. |
| fn read_null_indicators(r: &mut Reader<'_>, position_count: usize) -> Result<Vec<bool>> { |
| let may_have_null = r.read_u8()?; |
| if may_have_null == 0 { |
| return Ok(vec![false; position_count]); |
| } |
| read_packed_bits(r, position_count) |
| } |
| |
| /// Reads `count` MSB-first packed bits (`ceil(count/8)` bytes). |
| fn read_packed_bits(r: &mut Reader<'_>, count: usize) -> Result<Vec<bool>> { |
| let bytes = r.read_bytes(count.div_ceil(8))?; |
| unpack_bits_msb_first(bytes, count) |
| .ok_or_else(|| Error::Decode("truncated packed bit array".into())) |
| } |
| |
| fn decode_utf8(bytes: Vec<u8>) -> Result<String> { |
| String::from_utf8(bytes) |
| .map_err(|e| Error::Decode(format!("invalid UTF-8 in text column: {e}"))) |
| } |
| |
| /// Bounds-checked big-endian cursor over a byte slice. |
| struct Reader<'a> { |
| buf: &'a [u8], |
| pos: usize, |
| } |
| |
| impl<'a> Reader<'a> { |
| fn new(buf: &'a [u8]) -> Reader<'a> { |
| Reader { buf, pos: 0 } |
| } |
| |
| fn read_bytes(&mut self, n: usize) -> Result<&'a [u8]> { |
| let end = self |
| .pos |
| .checked_add(n) |
| .filter(|&e| e <= self.buf.len()) |
| .ok_or_else(|| { |
| Error::Decode(format!( |
| "truncated TsBlock: need {n} bytes at offset {}, have {}", |
| self.pos, |
| self.buf.len() - self.pos |
| )) |
| })?; |
| let slice = &self.buf[self.pos..end]; |
| self.pos = end; |
| Ok(slice) |
| } |
| |
| fn read_u8(&mut self) -> Result<u8> { |
| Ok(self.read_bytes(1)?[0]) |
| } |
| |
| fn read_i32(&mut self) -> Result<i32> { |
| let b = self.read_bytes(4)?; |
| Ok(i32::from_be_bytes([b[0], b[1], b[2], b[3]])) |
| } |
| |
| fn read_i64(&mut self) -> Result<i64> { |
| let b = self.read_bytes(8)?; |
| Ok(i64::from_be_bytes([ |
| b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7], |
| ])) |
| } |
| } |
| |
| /// Encoding helpers for crafting synthetic TsBlocks in unit tests (shared |
| /// with the dataset state-machine tests). |
| #[cfg(test)] |
| pub(crate) mod test_util { |
| use super::*; |
| |
| /// Builds a TsBlock header: valueColumnCount, type bytes, positionCount, |
| /// time encoding (Int64Array), value encodings. |
| pub(crate) fn header(types: &[TSDataType], position_count: i32, encodings: &[u8]) -> Vec<u8> { |
| let mut b = Vec::new(); |
| b.extend_from_slice(&(types.len() as i32).to_be_bytes()); |
| for &t in types { |
| b.push(t.code() as u8); |
| } |
| b.extend_from_slice(&position_count.to_be_bytes()); |
| b.push(ENCODING_INT64_ARRAY); // time column |
| b.extend_from_slice(encodings); |
| b |
| } |
| |
| /// Dense Int64Array time column with no nulls. |
| pub(crate) fn time_column(ts: &[i64]) -> Vec<u8> { |
| let mut b = vec![0u8]; // mayHaveNull = 0 |
| for t in ts { |
| b.extend_from_slice(&t.to_be_bytes()); |
| } |
| b |
| } |
| |
| /// One-Int32-column block: timestamps `ts`, dense non-null values `vals`. |
| pub(crate) fn int32_block(ts: &[i64], vals: &[i32]) -> Vec<u8> { |
| assert_eq!(ts.len(), vals.len()); |
| let mut b = header( |
| &[TSDataType::Int32], |
| ts.len() as i32, |
| &[ENCODING_INT32_ARRAY], |
| ); |
| b.extend_from_slice(&time_column(ts)); |
| b.push(0); // mayHaveNull |
| for v in vals { |
| b.extend_from_slice(&v.to_be_bytes()); |
| } |
| b |
| } |
| } |
| |
| #[cfg(test)] |
| mod tests { |
| use super::test_util::{header, time_column}; |
| use super::*; |
| |
| #[test] |
| fn int32_column_no_nulls() { |
| let mut b = header(&[TSDataType::Int32], 2, &[ENCODING_INT32_ARRAY]); |
| b.extend_from_slice(&time_column(&[10, 20])); |
| b.push(0); // mayHaveNull |
| b.extend_from_slice(&5i32.to_be_bytes()); |
| b.extend_from_slice(&(-1i32).to_be_bytes()); |
| |
| let block = TsBlock::decode(&b).unwrap(); |
| assert_eq!(block.position_count, 2); |
| assert_eq!(block.timestamps, vec![10, 20]); |
| assert_eq!(block.column_types, vec![TSDataType::Int32]); |
| assert_eq!(block.columns, vec![vec![Value::Int32(5), Value::Int32(-1)]]); |
| } |
| |
| #[test] |
| fn nulls_are_dense_skipped_in_numeric_columns() { |
| // 3 rows, row 1 null: only 2 i32 values on the wire. |
| let mut b = header(&[TSDataType::Int32], 3, &[ENCODING_INT32_ARRAY]); |
| b.extend_from_slice(&time_column(&[1, 2, 3])); |
| b.push(1); // mayHaveNull |
| b.push(0b0100_0000); // MSB-first: position 1 null |
| b.extend_from_slice(&7i32.to_be_bytes()); |
| b.extend_from_slice(&9i32.to_be_bytes()); |
| |
| let block = TsBlock::decode(&b).unwrap(); |
| assert_eq!( |
| block.columns[0], |
| vec![Value::Int32(7), Value::Null, Value::Int32(9)] |
| ); |
| } |
| |
| #[test] |
| fn int64_double_and_float_bit_patterns() { |
| let types = [TSDataType::Double, TSDataType::Float, TSDataType::Timestamp]; |
| let mut b = header( |
| &types, |
| 1, |
| &[ |
| ENCODING_INT64_ARRAY, |
| ENCODING_INT32_ARRAY, |
| ENCODING_INT64_ARRAY, |
| ], |
| ); |
| b.extend_from_slice(&time_column(&[42])); |
| b.push(0); |
| b.extend_from_slice(&(-2.5f64).to_be_bytes()); // f64 via i64 bits |
| b.push(0); |
| b.extend_from_slice(&1.5f32.to_be_bytes()); |
| b.push(0); |
| b.extend_from_slice(&123456789012345i64.to_be_bytes()); |
| |
| let block = TsBlock::decode(&b).unwrap(); |
| assert_eq!(block.columns[0], vec![Value::Double(-2.5)]); |
| assert_eq!(block.columns[1], vec![Value::Float(1.5)]); |
| assert_eq!(block.columns[2], vec![Value::Timestamp(123456789012345)]); |
| } |
| |
| #[test] |
| fn boolean_bytearray_full_bit_array_with_nulls() { |
| // 10 rows: nulls at 0 and 9; values MSB-first over ALL positions. |
| let mut b = header(&[TSDataType::Boolean], 10, &[ENCODING_BYTE_ARRAY]); |
| b.extend_from_slice(&time_column(&(0..10).collect::<Vec<i64>>())); |
| b.push(1); // mayHaveNull |
| b.extend_from_slice(&[0b1000_0000, 0b0100_0000]); // nulls at 0, 9 |
| // Value bits (positions 1..=8 meaningful): true at 1, 3, 8. |
| b.extend_from_slice(&[0b0101_0000, 0b1000_0000]); |
| |
| let block = TsBlock::decode(&b).unwrap(); |
| let expect: Vec<Value> = vec![ |
| Value::Null, |
| Value::Boolean(true), |
| Value::Boolean(false), |
| Value::Boolean(true), |
| Value::Boolean(false), |
| Value::Boolean(false), |
| Value::Boolean(false), |
| Value::Boolean(false), |
| Value::Boolean(true), |
| Value::Null, |
| ]; |
| assert_eq!(block.columns[0], expect); |
| } |
| |
| #[test] |
| fn binary_array_text_blob_and_nulls() { |
| let types = [TSDataType::Text, TSDataType::Blob]; |
| let mut b = header(&types, 2, &[ENCODING_BINARY_ARRAY, ENCODING_BINARY_ARRAY]); |
| b.extend_from_slice(&time_column(&[1, 2])); |
| // Text column: ["héllo", null] — null consumes no bytes. |
| b.push(1); |
| b.push(0b0100_0000); |
| let s = "héllo".as_bytes(); |
| b.extend_from_slice(&(s.len() as i32).to_be_bytes()); |
| b.extend_from_slice(s); |
| // Blob column, no nulls. |
| b.push(0); |
| b.extend_from_slice(&2i32.to_be_bytes()); |
| b.extend_from_slice(&[0xCA, 0xFE]); |
| b.extend_from_slice(&0i32.to_be_bytes()); // empty blob |
| |
| let block = TsBlock::decode(&b).unwrap(); |
| assert_eq!( |
| block.columns[0], |
| vec![Value::Text("héllo".into()), Value::Null] |
| ); |
| assert_eq!( |
| block.columns[1], |
| vec![Value::Blob(vec![0xCA, 0xFE]), Value::Blob(vec![])] |
| ); |
| } |
| |
| #[test] |
| fn rle_replicates_single_value() { |
| let mut b = header(&[TSDataType::Int64], 4, &[ENCODING_RLE]); |
| b.extend_from_slice(&time_column(&[1, 2, 3, 4])); |
| b.push(ENCODING_INT64_ARRAY); // inner encoding |
| b.push(0); // inner mayHaveNull |
| b.extend_from_slice(&99i64.to_be_bytes()); |
| |
| let block = TsBlock::decode(&b).unwrap(); |
| assert_eq!(block.columns[0], vec![Value::Int64(99); 4]); |
| } |
| |
| #[test] |
| fn rle_of_null_replicates_null() { |
| let mut b = header(&[TSDataType::Int32], 3, &[ENCODING_RLE]); |
| b.extend_from_slice(&time_column(&[1, 2, 3])); |
| b.push(ENCODING_INT32_ARRAY); |
| b.push(1); // inner mayHaveNull |
| b.push(0b1000_0000); // single position, null → zero dense values |
| let block = TsBlock::decode(&b).unwrap(); |
| assert_eq!(block.columns[0], vec![Value::Null; 3]); |
| } |
| |
| #[test] |
| fn rle_time_column() { |
| // Constant time column via RLE (encoding byte in header slot). |
| let mut b = Vec::new(); |
| b.extend_from_slice(&0i32.to_be_bytes()); // no value columns |
| b.extend_from_slice(&3i32.to_be_bytes()); // positionCount |
| b.push(ENCODING_RLE); // time encoding |
| b.push(ENCODING_INT64_ARRAY); // inner |
| b.push(0); |
| b.extend_from_slice(&7i64.to_be_bytes()); |
| |
| let block = TsBlock::decode(&b).unwrap(); |
| assert_eq!(block.timestamps, vec![7, 7, 7]); |
| assert!(block.columns.is_empty()); |
| } |
| |
| #[test] |
| fn date_decodes_as_yyyymmdd_i32() { |
| let mut b = header(&[TSDataType::Date], 1, &[ENCODING_INT32_ARRAY]); |
| b.extend_from_slice(&time_column(&[0])); |
| b.push(0); |
| b.extend_from_slice(&20260710i32.to_be_bytes()); |
| let block = TsBlock::decode(&b).unwrap(); |
| assert_eq!(block.columns[0], vec![Value::Date(20260710)]); |
| } |
| |
| #[test] |
| fn empty_block() { |
| let mut b = header(&[TSDataType::Int32], 0, &[ENCODING_INT32_ARRAY]); |
| b.extend_from_slice(&time_column(&[])); |
| b.push(0); // value column: mayHaveNull, zero values |
| let block = TsBlock::decode(&b).unwrap(); |
| assert_eq!(block.position_count, 0); |
| assert!(block.timestamps.is_empty()); |
| assert_eq!(block.columns, vec![Vec::<Value>::new()]); |
| } |
| |
| #[test] |
| fn truncated_and_malformed_inputs_error() { |
| // Truncated header. |
| assert!(matches!(TsBlock::decode(&[0, 0]), Err(Error::Decode(_)))); |
| // Unknown type code. |
| let mut b = Vec::new(); |
| b.extend_from_slice(&1i32.to_be_bytes()); |
| b.push(200); |
| assert!(matches!(TsBlock::decode(&b), Err(Error::Decode(_)))); |
| // Unknown encoding. |
| let mut b = header(&[TSDataType::Int32], 1, &[9]); |
| b.extend_from_slice(&time_column(&[1])); |
| assert!(matches!(TsBlock::decode(&b), Err(Error::Decode(_)))); |
| // Truncated value payload. |
| let mut b = header(&[TSDataType::Int64], 1, &[ENCODING_INT64_ARRAY]); |
| b.extend_from_slice(&time_column(&[1])); |
| b.push(0); |
| b.extend_from_slice(&[0, 0]); // only 2 of 8 bytes |
| assert!(matches!(TsBlock::decode(&b), Err(Error::Decode(_)))); |
| // Nested RLE is rejected. |
| let mut b = header(&[TSDataType::Int32], 1, &[ENCODING_RLE]); |
| b.extend_from_slice(&time_column(&[1])); |
| b.push(ENCODING_RLE); |
| assert!(matches!(TsBlock::decode(&b), Err(Error::Decode(_)))); |
| } |
| } |
