This is an automated email from the ASF dual-hosted git repository. CritasWang pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/iotdb-client-rust.git
commit b56c6c3b7d2689edd4576122273208d18669c86d Author: CritasWang <[email protected]> AuthorDate: Fri Jul 10 17:07:03 2026 +0800 Phase 2-3: connection + session layer (open/query/insert RPC), data layer (Tablet serialization, TsBlock decoder) --- src/client/mod.rs | 19 ++ src/client/session.rs | 443 ++++++++++++++++++++++++++++++++++++++++-- src/connection/mod.rs | 183 +++++++++++++++++- src/data/bitmap.rs | 111 +++++++++++ src/data/mod.rs | 114 ++++++++++- src/data/tablet.rs | 505 ++++++++++++++++++++++++++++++++++++++++++++++++ src/data/tsblock.rs | 521 ++++++++++++++++++++++++++++++++++++++++++++++++++ src/data/value.rs | 97 ++++++++++ src/error.rs | 3 + src/lib.rs | 3 +- src/protocol/mod.rs | 18 +- 11 files changed, 1992 insertions(+), 25 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 29ec6ac..34b96db 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -1,4 +1,23 @@ +// 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. + //! Session-layer API: tree model (`Session`) and, later, table model //! (`TableSession`) plus pooled variants — mirroring the Node.js and C# SDKs. pub mod session; + +pub use session::{QueryHandle, Session, SessionConfig}; diff --git a/src/client/session.rs b/src/client/session.rs index c504e25..a0c0a68 100644 --- a/src/client/session.rs +++ b/src/client/session.rs @@ -1,7 +1,45 @@ +// 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. + //! Tree-model session (device/timeseries paths). +use std::collections::BTreeMap; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + use crate::connection::{Connection, Endpoint}; use crate::error::{Error, Result}; +use crate::protocol::client::{ + TIClientRPCServiceSyncClient, TSCloseOperationReq, TSCloseSessionReq, TSExecuteStatementReq, + TSFetchResultsReq, TSInsertTabletReq, TSOpenSessionReq, TSProtocolVersion, +}; +use crate::protocol::common::TSStatus; + +/// TSStatus codes the client special-cases (see protocol spec §2). +pub mod status_code { + pub const SUCCESS_STATUS: i32 = 200; + pub const MULTIPLE_ERROR: i32 = 302; + /// The write succeeded; `redirectNode` merely recommends a better endpoint. + pub const REDIRECTION_RECOMMEND: i32 = 400; +} + +/// Rotating start index shared by all sessions so connections spread across +/// nodes (mirrors the C# SDK's round-robin-with-failover endpoint selection). +static ENDPOINT_START_INDEX: AtomicUsize = AtomicUsize::new(0); /// Configuration for opening a [`Session`]. #[derive(Debug, Clone)] @@ -13,6 +51,12 @@ pub struct SessionConfig { pub sql_dialect: String, pub fetch_size: i32, pub zone_id: String, + /// TCP connect timeout per endpoint attempt. + pub connect_timeout: Duration, + /// Per-query server-side timeout in milliseconds (request body field). + pub query_timeout_ms: i64, + /// Database to select at open time (table dialect; sent as config key `db`). + pub database: Option<String>, } impl Default for SessionConfig { @@ -24,31 +68,122 @@ impl Default for SessionConfig { sql_dialect: "tree".into(), fetch_size: 1024, zone_id: "UTC+8".into(), + connect_timeout: Duration::from_secs(10), + query_timeout_ms: 60_000, + database: None, } } } +impl SessionConfig { + /// Set endpoints from `"host:port"` node-url strings (IPv6 hosts may be + /// bracketed, e.g. `"[::1]:6667"`). + pub fn with_node_urls<S: AsRef<str>>(mut self, node_urls: &[S]) -> Result<Self> { + self.endpoints = node_urls + .iter() + .map(|u| Endpoint::parse(u.as_ref())) + .collect::<Result<Vec<_>>>()?; + Ok(self) + } +} + +/// Raw result of a query statement: response metadata plus undecoded TsBlocks. +/// +/// TsBlock decoding lives in the data layer; this handle only carries the +/// bytes and the bookkeeping needed to fetch more pages and close the query. +#[derive(Debug)] +pub struct QueryHandle { + pub query_id: i64, + pub statement: String, + pub columns: Vec<String>, + pub data_type_list: Vec<String>, + pub ignore_time_stamp: bool, + /// Serialized TsBlocks (decoded by the data layer). + pub query_result: Vec<Vec<u8>>, + pub more_data: bool, + /// Output column ordinal → physical TsBlock column index (`-1` = time + /// column); identity mapping when absent. + pub column_index2_ts_block_column_index_list: Option<Vec<i32>>, +} + /// A tree-model session against an IoTDB cluster. pub struct Session { config: SessionConfig, connection: Option<Connection>, + session_id: i64, + statement_id: i64, + /// Current database, tracked from `USE <db>` responses (table dialect). + database: Option<String>, } impl Session { pub fn new(config: SessionConfig) -> Self { - Self { config, connection: None } + let database = config.database.clone(); + Self { + config, + connection: None, + session_id: -1, + statement_id: -1, + database, + } } + /// Connect (trying endpoints from a rotating start index; first success + /// wins), open the session and request the connection's statement id. pub fn open(&mut self) -> Result<()> { - let endpoint = self - .config - .endpoints - .first() - .ok_or_else(|| Error::Client("no endpoints configured".into()))? - .clone(); - self.connection = Some(Connection::open(endpoint)?); - // TODO(codegen): call openSession RPC (username/password/zoneId/sql_dialect), - // store sessionId + statementId. + if self.connection.is_some() { + return Err(Error::Client("session already open".into())); + } + let n = self.config.endpoints.len(); + if n == 0 { + return Err(Error::Client("no endpoints configured".into())); + } + + let start = ENDPOINT_START_INDEX.fetch_add(1, Ordering::Relaxed) % n; + let mut connection = None; + let mut last_err: Option<Error> = None; + for i in 0..n { + let endpoint = self.config.endpoints[(start + i) % n].clone(); + match Connection::open(endpoint, self.config.connect_timeout) { + Ok(c) => { + connection = Some(c); + break; + } + Err(e) => last_err = Some(e), + } + } + let mut connection = connection.ok_or_else(|| { + last_err.unwrap_or_else(|| Error::Client("no endpoints configured".into())) + })?; + + let mut configuration = BTreeMap::new(); + configuration.insert("sql_dialect".to_string(), self.config.sql_dialect.clone()); + if let Some(db) = &self.database { + // ⚠️ config key is literally "db", not "database". + configuration.insert("db".to_string(), db.clone()); + } + let req = TSOpenSessionReq::new( + TSProtocolVersion::IotdbServiceProtocolV3, + self.config.zone_id.clone(), + self.config.username.clone(), + self.config.password.clone(), + configuration, + ); + let resp = connection.client_mut().open_session(req)?; + check_status(&resp.status)?; + if resp.server_protocol_version != TSProtocolVersion::IotdbServiceProtocolV3 { + log::warn!( + "server protocol version mismatch: expected V3, got {:?}", + resp.server_protocol_version + ); + } + self.session_id = resp + .session_id + .ok_or_else(|| Error::Client("openSession response missing sessionId".into()))?; + self.statement_id = connection + .client_mut() + .request_statement_id(self.session_id)?; + self.connection = Some(connection); Ok(()) } @@ -56,28 +191,308 @@ impl Session { self.connection.is_some() } + /// The database currently selected on this session, if any. + pub fn database(&self) -> Option<&str> { + self.database.as_deref() + } + + fn connection_mut(&mut self) -> Result<&mut Connection> { + self.connection + .as_mut() + .ok_or_else(|| Error::Client("session is not open".into())) + } + + /// Execute a non-query statement (DDL/DML). Tracks `USE <db>` via the + /// response's `database` field. + pub fn execute_non_query(&mut self, sql: &str) -> Result<()> { + let req = self.statement_req(sql); + let resp = self + .connection_mut()? + .client_mut() + .execute_update_statement_v2(req)?; + check_status(&resp.status)?; + if let Some(db) = resp.database { + self.database = Some(db); + } + Ok(()) + } + + /// Execute a query statement, returning raw TsBlock bytes plus metadata. + /// Decoding happens in the data layer. + pub fn execute_query(&mut self, sql: &str) -> Result<QueryHandle> { + let req = self.statement_req(sql); + let resp = self + .connection_mut()? + .client_mut() + .execute_query_statement_v2(req)?; + check_status(&resp.status)?; + let query_id = resp + .query_id + .ok_or_else(|| Error::Client("query response missing queryId".into()))?; + Ok(QueryHandle { + query_id, + statement: sql.to_string(), + columns: resp.columns.unwrap_or_default(), + data_type_list: resp.data_type_list.unwrap_or_default(), + ignore_time_stamp: resp.ignore_time_stamp.unwrap_or(false), + query_result: resp.query_result.unwrap_or_default(), + more_data: resp.more_data.unwrap_or(false), + column_index2_ts_block_column_index_list: resp.column_index2_ts_block_column_index_list, + }) + } + + /// Fetch the next page of TsBlocks for an open query. Returns the raw + /// blocks and whether more data remains; empty when the set is exhausted. + pub fn fetch_results(&mut self, query_id: i64, sql: &str) -> Result<(Vec<Vec<u8>>, bool)> { + let req = TSFetchResultsReq::new( + self.session_id, + sql.to_string(), + self.config.fetch_size, + query_id, + true, // isAlign — always true on the V2/TsBlock path + self.config.query_timeout_ms, + self.statement_id, + ); + let resp = self.connection_mut()?.client_mut().fetch_results_v2(req)?; + check_status(&resp.status)?; + if !resp.has_result_set { + return Ok((Vec::new(), false)); + } + Ok(( + resp.query_result.unwrap_or_default(), + resp.more_data.unwrap_or(false), + )) + } + + /// Close an open query result set. Best-effort: errors are swallowed, + /// matching the Node.js and C# SDKs. + pub fn close_query(&mut self, query_id: i64) { + let (session_id, statement_id) = (self.session_id, self.statement_id); + if let Ok(connection) = self.connection_mut() { + let req = TSCloseOperationReq::new(session_id, query_id, statement_id, None); + if let Err(e) = connection.client_mut().close_operation(req) { + log::debug!("closeOperation for query {query_id} failed (ignored): {e}"); + } + } + } + + /// Insert a tablet from pre-serialized buffers (see protocol spec §3). + /// + /// `values` is the column-major value buffer with trailing null bitmaps; + /// `timestamps` is the `size × 8-byte i64 BE` buffer. Serialization from a + /// `Tablet` lives in the data layer; a typed `insert_tablet(&Tablet)` is + /// wired in the next phase. + #[allow(clippy::too_many_arguments)] + pub fn insert_tablet_raw( + &mut self, + prefix_path: &str, + measurements: Vec<String>, + types: Vec<i32>, + values: Vec<u8>, + timestamps: Vec<u8>, + size: i32, + is_aligned: bool, + write_to_table: Option<bool>, + column_categories: Option<Vec<i8>>, + ) -> Result<()> { + let req = TSInsertTabletReq::new( + self.session_id, + prefix_path.to_string(), + measurements, + values, + timestamps, + types, + size, + is_aligned, + write_to_table, + column_categories, + None, + None, + None, + ); + let status = self.connection_mut()?.client_mut().insert_tablet(req)?; + check_status(&status) + } + + /// Close the session: best-effort `closeSession` RPC, then drop the + /// connection. pub fn close(&mut self) -> Result<()> { - // TODO(codegen): call closeSession RPC before dropping the connection. - self.connection = None; + if let Some(mut connection) = self.connection.take() { + let req = TSCloseSessionReq::new(self.session_id); + if let Err(e) = connection.client_mut().close_session(req) { + log::debug!("closeSession failed (ignored): {e}"); + } + } + self.session_id = -1; + self.statement_id = -1; Ok(()) } + + fn statement_req(&self, sql: &str) -> TSExecuteStatementReq { + TSExecuteStatementReq::new( + self.session_id, + sql.to_string(), + self.statement_id, + self.config.fetch_size, + self.config.query_timeout_ms, + true, // enableRedirectQuery + false, // jdbcQuery=false forces TsBlock queryResult responses + ) + } +} + +impl Drop for Session { + fn drop(&mut self) { + let _ = self.close(); + } +} + +/// Map a `TSStatus` to success or [`Error::Server`] (protocol spec §2): +/// 200 OK; 400 is a **successful** write with a redirect hint; 302 succeeds +/// iff every subStatus is itself OK. +pub fn check_status(status: &TSStatus) -> Result<()> { + match status.code { + status_code::SUCCESS_STATUS | status_code::REDIRECTION_RECOMMEND => Ok(()), + status_code::MULTIPLE_ERROR => { + let failed = status + .sub_status + .as_deref() + .unwrap_or_default() + .iter() + .find(|s| check_status(s).is_err()); + match failed { + None => Ok(()), + Some(s) => Err(Error::Server { + code: s.code, + message: s.message.clone().unwrap_or_default(), + }), + } + } + code => Err(Error::Server { + code, + message: status.message.clone().unwrap_or_default(), + }), + } } #[cfg(test)] mod tests { use super::*; + fn status(code: i32) -> TSStatus { + TSStatus::new(code, None, None, None, None, None) + } + #[test] - fn default_config_targets_localhost() { + fn default_config() { let cfg = SessionConfig::default(); assert_eq!(cfg.endpoints[0], Endpoint::new("localhost", 6667)); + assert_eq!(cfg.username, "root"); + assert_eq!(cfg.password, "root"); assert_eq!(cfg.sql_dialect, "tree"); + assert_eq!(cfg.fetch_size, 1024); + assert_eq!(cfg.query_timeout_ms, 60_000); + assert_eq!(cfg.connect_timeout, Duration::from_secs(10)); + assert!(cfg.database.is_none()); + } + + #[test] + fn config_from_node_urls() { + let cfg = SessionConfig::default() + .with_node_urls(&["10.0.0.1:6667", "[::1]:6668"]) + .unwrap(); + assert_eq!( + cfg.endpoints, + vec![Endpoint::new("10.0.0.1", 6667), Endpoint::new("::1", 6668)] + ); + assert!(SessionConfig::default() + .with_node_urls(&["nohost"]) + .is_err()); + } + + #[test] + fn status_200_is_ok() { + assert!(check_status(&status(200)).is_ok()); + } + + #[test] + fn status_400_redirect_is_success() { + assert!(check_status(&status(400)).is_ok()); + } + + #[test] + fn status_302_all_sub_ok_is_success() { + let mut s = status(302); + s.sub_status = Some(vec![Box::new(status(200)), Box::new(status(400))]); + assert!(check_status(&s).is_ok()); + } + + #[test] + fn status_302_mixed_sub_is_error() { + let mut s = status(302); + s.sub_status = Some(vec![Box::new(status(200)), Box::new(status(500))]); + match check_status(&s) { + Err(Error::Server { code, .. }) => assert_eq!(code, 500), + other => panic!("expected server error, got {other:?}"), + } + } + + #[test] + fn status_500_is_error() { + let mut s = status(500); + s.message = Some("boom".into()); + match check_status(&s) { + Err(Error::Server { code, message }) => { + assert_eq!(code, 500); + assert_eq!(message, "boom"); + } + other => panic!("expected server error, got {other:?}"), + } + } + + #[test] + fn calls_on_closed_session_fail() { + let mut session = Session::new(SessionConfig::default()); + assert!(!session.is_open()); + assert!(session.execute_non_query("SHOW DATABASES").is_err()); + assert!(session.execute_query("SELECT 1").is_err()); + assert!(session.close().is_ok()); // close on never-opened is fine } #[test] fn open_without_endpoints_fails() { - let mut session = Session::new(SessionConfig { endpoints: vec![], ..Default::default() }); + let mut session = Session::new(SessionConfig { + endpoints: vec![], + ..Default::default() + }); assert!(session.open().is_err()); assert!(!session.is_open()); } + + /// End-to-end smoke test against a live server; skipped when no IoTDB + /// instance is reachable on localhost:6667. + #[test] + fn live_server_roundtrip() { + use std::net::TcpStream; + if TcpStream::connect_timeout( + &"127.0.0.1:6667".parse().unwrap(), + Duration::from_millis(300), + ) + .is_err() + { + eprintln!("skipping live_server_roundtrip: no IoTDB server on 127.0.0.1:6667"); + return; + } + + let mut session = Session::new(SessionConfig::default()); + session.open().expect("open session"); + assert!(session.is_open()); + + let handle = session.execute_query("SHOW DATABASES").expect("query"); + assert!(!handle.columns.is_empty()); + session.close_query(handle.query_id); + + session.close().expect("close session"); + assert!(!session.is_open()); + } } diff --git a/src/connection/mod.rs b/src/connection/mod.rs index 430478d..c6cdc89 100644 --- a/src/connection/mod.rs +++ b/src/connection/mod.rs @@ -1,9 +1,44 @@ +// 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. + //! Low-level Thrift connection to an IoTDB node. //! //! Mirrors `src/connection/Connection.ts` (Node.js) and the Thrift client setup //! in the C# SDK: TCP → TFramedTransport → TBinaryProtocol → IClientRPCService client. -use crate::error::Result; +use std::net::{TcpStream, ToSocketAddrs}; +use std::time::Duration; + +use thrift::protocol::{TBinaryInputProtocol, TBinaryOutputProtocol}; +use thrift::transport::{ + ReadHalf, TFramedReadTransport, TFramedWriteTransport, TIoChannel, TTcpChannel, WriteHalf, +}; + +use crate::error::{Error, Result}; +use crate::protocol::client::IClientRPCServiceSyncClient; + +/// Default IoTDB DataNode RPC port. +pub const DEFAULT_PORT: u16 = 6667; + +/// The concrete generated RPC client over framed transport + strict binary protocol. +pub type RpcClient = IClientRPCServiceSyncClient< + TBinaryInputProtocol<TFramedReadTransport<ReadHalf<TTcpChannel>>>, + TBinaryOutputProtocol<TFramedWriteTransport<WriteHalf<TTcpChannel>>>, +>; /// A single endpoint `host:port` of an IoTDB DataNode (default port 6667). #[derive(Debug, Clone, PartialEq, Eq)] @@ -14,24 +49,158 @@ pub struct Endpoint { impl Endpoint { pub fn new(host: impl Into<String>, port: u16) -> Self { - Self { host: host.into(), port } + Self { + host: host.into(), + port, + } + } + + /// Parse a `"host:port"` node-url string. + /// + /// Splits on the **last** `:` so IPv6 literals work; surrounding `[]` + /// brackets on the host part are stripped (e.g. `"[::1]:6667"` → host `::1`). + pub fn parse(s: &str) -> Result<Self> { + let s = s.trim(); + let idx = s + .rfind(':') + .ok_or_else(|| Error::Client(format!("invalid node url '{s}': expected host:port")))?; + let (host_part, port_part) = (&s[..idx], &s[idx + 1..]); + let port: u16 = port_part.parse().map_err(|_| { + Error::Client(format!("invalid node url '{s}': bad port '{port_part}'")) + })?; + let host = host_part + .strip_prefix('[') + .and_then(|h| h.strip_suffix(']')) + .unwrap_or(host_part); + if host.is_empty() { + return Err(Error::Client(format!("invalid node url '{s}': empty host"))); + } + Ok(Self::new(host, port)) + } +} + +impl std::fmt::Display for Endpoint { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.host.contains(':') { + write!(f, "[{}]:{}", self.host, self.port) + } else { + write!(f, "{}:{}", self.host, self.port) + } } } /// Low-level connection wrapper. Owns the Thrift transport/protocol pair -/// and the generated `IClientRPCService` client once codegen lands. +/// and the generated `IClientRPCService` client. pub struct Connection { endpoint: Endpoint, - // TODO(codegen): hold IClientRPCServiceSyncClient over TFramedTransport + TBinaryProtocol + client: RpcClient, } impl Connection { - pub fn open(endpoint: Endpoint) -> Result<Self> { - // TODO(codegen): establish TTcpChannel, wrap in framed transport + binary protocol - Ok(Self { endpoint }) + /// Establish a TCP connection to `endpoint` (bounded by `connect_timeout`) + /// and wrap it in framed transport + strict binary protocol. + pub fn open(endpoint: Endpoint, connect_timeout: Duration) -> Result<Self> { + let stream = connect_stream(&endpoint, connect_timeout)?; + stream.set_nodelay(true).map_err(thrift::Error::from)?; + + let channel = TTcpChannel::with_stream(stream); + let (read_half, write_half) = channel.split()?; + let read_transport = TFramedReadTransport::new(read_half); + let write_transport = TFramedWriteTransport::new(write_half); + let input_protocol = TBinaryInputProtocol::new(read_transport, true); + let output_protocol = TBinaryOutputProtocol::new(write_transport, true); + let client = IClientRPCServiceSyncClient::new(input_protocol, output_protocol); + + Ok(Self { endpoint, client }) + } + + /// Mutable access to the generated RPC client for issuing calls. + pub fn client_mut(&mut self) -> &mut RpcClient { + &mut self.client } pub fn endpoint(&self) -> &Endpoint { &self.endpoint } } + +/// Resolve the endpoint and try each resolved address with the connect timeout. +fn connect_stream(endpoint: &Endpoint, connect_timeout: Duration) -> Result<TcpStream> { + let addrs = (endpoint.host.as_str(), endpoint.port) + .to_socket_addrs() + .map_err(thrift::Error::from)?; + let mut last_err: Option<std::io::Error> = None; + for addr in addrs { + match TcpStream::connect_timeout(&addr, connect_timeout) { + Ok(stream) => return Ok(stream), + Err(e) => last_err = Some(e), + } + } + Err(match last_err { + Some(e) => Error::Thrift(thrift::Error::from(e)), + None => Error::Client(format!("could not resolve endpoint {endpoint}")), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_ipv4() { + let ep = Endpoint::parse("127.0.0.1:6667").unwrap(); + assert_eq!(ep, Endpoint::new("127.0.0.1", 6667)); + } + + #[test] + fn parse_hostname() { + let ep = Endpoint::parse("iotdb.example.com:1234").unwrap(); + assert_eq!(ep, Endpoint::new("iotdb.example.com", 1234)); + } + + #[test] + fn parse_ipv6_bracketed() { + let ep = Endpoint::parse("[::1]:6667").unwrap(); + assert_eq!(ep, Endpoint::new("::1", 6667)); + + let ep = Endpoint::parse("[2001:db8::1]:6668").unwrap(); + assert_eq!(ep, Endpoint::new("2001:db8::1", 6668)); + } + + #[test] + fn parse_trims_whitespace() { + let ep = Endpoint::parse(" localhost:6667 ").unwrap(); + assert_eq!(ep, Endpoint::new("localhost", 6667)); + } + + #[test] + fn parse_no_port_is_error() { + assert!(Endpoint::parse("localhost").is_err()); + } + + #[test] + fn parse_bad_port_is_error() { + assert!(Endpoint::parse("localhost:abc").is_err()); + assert!(Endpoint::parse("localhost:99999").is_err()); + assert!(Endpoint::parse("localhost:").is_err()); + } + + #[test] + fn parse_empty_host_is_error() { + assert!(Endpoint::parse(":6667").is_err()); + assert!(Endpoint::parse("[]:6667").is_err()); + } + + #[test] + fn display_roundtrip() { + assert_eq!( + Endpoint::new("localhost", 6667).to_string(), + "localhost:6667" + ); + assert_eq!(Endpoint::new("::1", 6667).to_string(), "[::1]:6667"); + assert_eq!( + Endpoint::parse(&Endpoint::new("::1", 6667).to_string()).unwrap(), + Endpoint::new("::1", 6667) + ); + } +} diff --git a/src/data/bitmap.rs b/src/data/bitmap.rs new file mode 100644 index 0000000..b03d163 --- /dev/null +++ b/src/data/bitmap.rs @@ -0,0 +1,111 @@ +// 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. + +//! Bitmap helpers for the two — deliberately separate — bit conventions in +//! the IoTDB wire protocol (spec gotcha #4): +//! +//! * **Write side** (tablet null bitmaps, §3.3): LSB-first, bit=1 ⇒ null. +//! * **Read side** (TsBlock null indicators and BOOLEAN value arrays, §5.3): +//! MSB-first, bit=1 ⇒ set (null, or `true` for booleans). +//! +//! Do not mix them. + +/// Packs `bits` LSB-first into `ceil(len/8)` bytes: bit `i` → byte `i >> 3`, +/// mask `1 << (i & 7)`. Used for tablet write-side null bitmaps. +pub fn pack_bits_lsb_first(bits: &[bool]) -> Vec<u8> { + let mut out = vec![0u8; bits.len().div_ceil(8)]; + for (i, &b) in bits.iter().enumerate() { + if b { + out[i >> 3] |= 1 << (i & 7); + } + } + out +} + +/// Unpacks `count` bits MSB-first from `bytes`: bit `i` → byte `i >> 3`, +/// mask `0x80 >> (i & 7)`. Used for TsBlock read-side null indicators and +/// BOOLEAN value arrays. Returns `None` if `bytes` is too short. +pub fn unpack_bits_msb_first(bytes: &[u8], count: usize) -> Option<Vec<bool>> { + if bytes.len() < count.div_ceil(8) { + return None; + } + Some( + (0..count) + .map(|i| bytes[i >> 3] & (0x80 >> (i & 7)) != 0) + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn pack_lsb_first_bit_positions() { + // Position 0 → 0x01, position 7 → 0x80. + assert_eq!(pack_bits_lsb_first(&[true]), vec![0x01]); + let mut bits = [false; 8]; + bits[7] = true; + assert_eq!(pack_bits_lsb_first(&bits), vec![0x80]); + } + + #[test] + fn pack_lsb_first_multi_byte_and_padding() { + // 9 bits → 2 bytes; positions 0 and 8 set. + let mut bits = [false; 9]; + bits[0] = true; + bits[8] = true; + assert_eq!(pack_bits_lsb_first(&bits), vec![0x01, 0x01]); + // Padding bits stay 0. + assert_eq!(pack_bits_lsb_first(&[false; 3]), vec![0x00]); + assert_eq!(pack_bits_lsb_first(&[]), Vec::<u8>::new()); + } + + #[test] + fn unpack_msb_first_bit_positions() { + // Position 0 → 0x80, position 7 → 0x01. + assert_eq!(unpack_bits_msb_first(&[0x80], 1), Some(vec![true])); + let bits = unpack_bits_msb_first(&[0x01], 8).unwrap(); + assert!(bits[7]); + assert!(bits[..7].iter().all(|&b| !b)); + } + + #[test] + fn unpack_msb_first_multi_byte() { + // 0b10100000 0b01000000 over 10 bits → positions 0, 2, 9 set. + let bits = unpack_bits_msb_first(&[0xA0, 0x40], 10).unwrap(); + let set: Vec<usize> = (0..10).filter(|&i| bits[i]).collect(); + assert_eq!(set, vec![0, 2, 9]); + } + + #[test] + fn unpack_msb_first_short_buffer_is_none() { + assert_eq!(unpack_bits_msb_first(&[0xFF], 9), None); + assert_eq!(unpack_bits_msb_first(&[], 0), Some(vec![])); + } + + #[test] + fn the_two_conventions_differ() { + // Same logical bit vector, different byte images — guards against + // accidentally unifying the two implementations. + let mut bits = [false; 8]; + bits[0] = true; + let written = pack_bits_lsb_first(&bits); // 0x01 + let read_back = unpack_bits_msb_first(&written, 8).unwrap(); + assert_ne!(read_back, bits.to_vec()); + } +} diff --git a/src/data/mod.rs b/src/data/mod.rs index 82a8cab..c0b0a76 100644 --- a/src/data/mod.rs +++ b/src/data/mod.rs @@ -1,7 +1,33 @@ -//! Data structures: TSDataType codes, Tablet, SessionDataSet, BitMap. +// 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. + +//! Data structures: TSDataType codes, Value, Tablet, TsBlock, bitmap helpers. //! Data-type codes must match the official TSFile spec (0–11), identical //! across all IoTDB client SDKs. +pub mod bitmap; +pub mod tablet; +pub mod tsblock; +pub mod value; + +pub use tablet::Tablet; +pub use tsblock::TsBlock; +pub use value::Value; + /// Official TSFile data type codes. #[derive(Debug, Clone, Copy, PartialEq, Eq)] #[repr(i32)] @@ -19,3 +45,89 @@ pub enum TSDataType { Blob = 10, String = 11, } + +impl TSDataType { + /// The wire code as sent in insert `types` lists (i32) and TsBlock + /// headers (1 byte). + pub fn code(self) -> i32 { + self as i32 + } + + /// Parses a 1-byte type code as found in TsBlock headers. + pub fn from_code(code: u8) -> Option<TSDataType> { + Some(match code { + 0 => TSDataType::Boolean, + 1 => TSDataType::Int32, + 2 => TSDataType::Int64, + 3 => TSDataType::Float, + 4 => TSDataType::Double, + 5 => TSDataType::Text, + 6 => TSDataType::Vector, + 7 => TSDataType::Unknown, + 8 => TSDataType::Timestamp, + 9 => TSDataType::Date, + 10 => TSDataType::Blob, + 11 => TSDataType::String, + _ => return None, + }) + } +} + +/// Table-model column category, sent as a signed byte in +/// `TSInsertTabletReq.columnCategories`. +/// +/// The server also knows an internal `TIME = 3` category, but it is never +/// sent on the wire — the tablet's timestamps travel in the separate +/// `timestamps` buffer. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(i8)] +pub enum ColumnCategory { + Tag = 0, + Field = 1, + Attribute = 2, +} + +impl ColumnCategory { + /// The wire code as sent in `columnCategories` (signed byte). + pub fn code(self) -> i8 { + self as i8 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn data_type_codes_match_spec() { + assert_eq!(TSDataType::Boolean.code(), 0); + assert_eq!(TSDataType::Int32.code(), 1); + assert_eq!(TSDataType::Int64.code(), 2); + assert_eq!(TSDataType::Float.code(), 3); + assert_eq!(TSDataType::Double.code(), 4); + assert_eq!(TSDataType::Text.code(), 5); + assert_eq!(TSDataType::Vector.code(), 6); + assert_eq!(TSDataType::Unknown.code(), 7); + assert_eq!(TSDataType::Timestamp.code(), 8); + assert_eq!(TSDataType::Date.code(), 9); + assert_eq!(TSDataType::Blob.code(), 10); + assert_eq!(TSDataType::String.code(), 11); + } + + #[test] + fn data_type_from_code_round_trips() { + for code in 0u8..=11 { + let ty = TSDataType::from_code(code).expect("valid code"); + assert_eq!(ty.code(), i32::from(code)); + } + assert_eq!(TSDataType::from_code(12), None); + assert_eq!(TSDataType::from_code(255), None); + } + + #[test] + fn column_category_codes_match_spec() { + assert_eq!(ColumnCategory::Tag.code(), 0); + assert_eq!(ColumnCategory::Field.code(), 1); + assert_eq!(ColumnCategory::Attribute.code(), 2); + } +} diff --git a/src/data/tablet.rs b/src/data/tablet.rs new file mode 100644 index 0000000..66ab5ba --- /dev/null +++ b/src/data/tablet.rs @@ -0,0 +1,505 @@ +// 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. + +//! A batch of rows for one device (tree model) or one table (table model), +//! serialized column-major for `insertTablet` (protocol spec §3). + +use super::bitmap::pack_bits_lsb_first; +use super::value::Value; +use super::{ColumnCategory, TSDataType}; +use crate::error::{Error, Result}; + +/// A column-major batch of rows for `insertTablet`. +/// +/// Tree model: `target` is the device id (e.g. `root.sg.d1`) and +/// `column_categories` is `None`. Table model (via [`Tablet::new_table`]): +/// `target` is the table name and every column carries a +/// [`ColumnCategory`]. +/// +/// Rows are stably sorted by timestamp before serialization (spec §3.5); +/// the server fast path assumes sorted input. +#[derive(Debug, Clone)] +pub struct Tablet { + /// Device id (tree model) or table name (table model). + target: String, + /// Measurement names (tree model) or column names (table model). + measurements: Vec<String>, + types: Vec<TSDataType>, + timestamps: Vec<i64>, + /// Column-major: `values[col][row]`, `None` = null cell. + values: Vec<Vec<Option<Value>>>, + /// Table model only: one category per column. + column_categories: Option<Vec<ColumnCategory>>, +} + +impl Tablet { + /// Creates an empty tree-model tablet for `device_id`. + pub fn new( + device_id: impl Into<String>, + measurements: Vec<String>, + types: Vec<TSDataType>, + ) -> Result<Tablet> { + if measurements.len() != types.len() { + return Err(Error::Client(format!( + "measurement count ({}) != type count ({})", + measurements.len(), + types.len() + ))); + } + let columns = measurements.len(); + Ok(Tablet { + target: device_id.into(), + measurements, + types, + timestamps: Vec::new(), + values: vec![Vec::new(); columns], + column_categories: None, + }) + } + + /// Creates an empty table-model tablet for `table_name` with one + /// [`ColumnCategory`] per column. + pub fn new_table( + table_name: impl Into<String>, + column_names: Vec<String>, + types: Vec<TSDataType>, + column_categories: Vec<ColumnCategory>, + ) -> Result<Tablet> { + if column_categories.len() != types.len() { + return Err(Error::Client(format!( + "category count ({}) != type count ({})", + column_categories.len(), + types.len() + ))); + } + let mut tablet = Tablet::new(table_name, column_names, types)?; + tablet.column_categories = Some(column_categories); + Ok(tablet) + } + + /// Device id (tree model) or table name (table model). + pub fn target(&self) -> &str { + &self.target + } + + /// Table name — alias of [`Tablet::target`] for table-model tablets. + pub fn table_name(&self) -> &str { + &self.target + } + + pub fn measurements(&self) -> &[String] { + &self.measurements + } + + pub fn types(&self) -> &[TSDataType] { + &self.types + } + + pub fn timestamps(&self) -> &[i64] { + &self.timestamps + } + + /// `Some` iff this is a table-model tablet. + pub fn column_categories(&self) -> Option<&[ColumnCategory]> { + self.column_categories.as_deref() + } + + pub fn row_count(&self) -> usize { + self.timestamps.len() + } + + pub fn is_table_model(&self) -> bool { + self.column_categories.is_some() + } + + /// Appends one row. `row[i]` must be `None` (null) or a [`Value`] whose + /// type matches `types[i]`. + pub fn add_row(&mut self, timestamp: i64, row: Vec<Option<Value>>) -> Result<()> { + if row.len() != self.types.len() { + return Err(Error::Client(format!( + "row has {} cells, tablet has {} columns", + row.len(), + self.types.len() + ))); + } + for (i, cell) in row.iter().enumerate() { + match cell { + None | Some(Value::Null) => {} + Some(v) if v.data_type() == Some(self.types[i]) => {} + Some(v) => { + return Err(Error::Client(format!( + "column {i} ({}) expects {:?}, got {v:?}", + self.measurements[i], self.types[i] + ))); + } + } + } + self.timestamps.push(timestamp); + for (col, cell) in self.values.iter_mut().zip(row) { + // Normalize Some(Null) to None so null detection is uniform. + col.push(cell.filter(|v| !v.is_null())); + } + Ok(()) + } + + /// Stably sorts rows by timestamp, reordering all value columns in step. + pub fn sort_by_timestamp(&mut self) { + let n = self.timestamps.len(); + let mut order: Vec<usize> = (0..n).collect(); + order.sort_by_key(|&i| self.timestamps[i]); // stable + if order.iter().enumerate().all(|(pos, &i)| pos == i) { + return; + } + self.timestamps = order.iter().map(|&i| self.timestamps[i]).collect(); + for col in &mut self.values { + *col = order.iter().map(|&i| col[i].clone()).collect(); + } + } + + /// Serializes the value buffer per spec §3.2–3.3: all columns + /// column-major (nulls occupy placeholder slots), then one trailing + /// bitmap entry per column — a flag byte (1 = column has nulls) followed, + /// when flagged, by a `ceil(rows/8)`-byte LSB-first bitmap with bit=1 for + /// null rows. All multi-byte values big-endian. + /// + /// Sorts rows by timestamp first (spec §3.5). + pub fn serialize_values(&mut self) -> Vec<u8> { + self.sort_by_timestamp(); + let rows = self.row_count(); + let mut buf = Vec::new(); + for (col, &ty) in self.values.iter().zip(&self.types) { + for cell in col { + write_cell(&mut buf, ty, cell.as_ref()); + } + } + // Trailing per-column null bitmap section. + for col in &self.values { + let nulls: Vec<bool> = col.iter().map(Option::is_none).collect(); + if nulls.iter().any(|&n| n) { + buf.push(1); + buf.extend_from_slice(&pack_bits_lsb_first(&nulls)); + } else { + buf.push(0); + } + } + debug_assert!(self.values.iter().all(|c| c.len() == rows)); + buf + } + + /// Serializes the timestamp buffer per spec §3.4: `row_count` contiguous + /// i64 big-endian values, no count prefix. + /// + /// Sorts rows by timestamp first (spec §3.5). + pub fn serialize_timestamps(&mut self) -> Vec<u8> { + self.sort_by_timestamp(); + let mut buf = Vec::with_capacity(self.timestamps.len() * 8); + for ts in &self.timestamps { + buf.extend_from_slice(&ts.to_be_bytes()); + } + buf + } +} + +/// Writes one cell, using the C#-style sentinel placeholders for nulls +/// (spec §3.2 — any placeholder works, the server masks by bitmap). +fn write_cell(buf: &mut Vec<u8>, ty: TSDataType, cell: Option<&Value>) { + match (ty, cell) { + (TSDataType::Boolean, Some(Value::Boolean(b))) => buf.push(u8::from(*b)), + (TSDataType::Boolean, None) => buf.push(0), + (TSDataType::Int32, Some(Value::Int32(v))) => buf.extend_from_slice(&v.to_be_bytes()), + (TSDataType::Int32, None) => buf.extend_from_slice(&i32::MIN.to_be_bytes()), + (TSDataType::Int64, Some(Value::Int64(v))) => buf.extend_from_slice(&v.to_be_bytes()), + (TSDataType::Timestamp, Some(Value::Timestamp(v))) => { + buf.extend_from_slice(&v.to_be_bytes()) + } + (TSDataType::Int64 | TSDataType::Timestamp, None) => { + buf.extend_from_slice(&i64::MIN.to_be_bytes()) + } + (TSDataType::Float, Some(Value::Float(v))) => buf.extend_from_slice(&v.to_be_bytes()), + (TSDataType::Float, None) => buf.extend_from_slice(&f32::MIN.to_be_bytes()), + (TSDataType::Double, Some(Value::Double(v))) => buf.extend_from_slice(&v.to_be_bytes()), + (TSDataType::Double, None) => buf.extend_from_slice(&f64::MIN.to_be_bytes()), + (TSDataType::Text, Some(Value::Text(s))) | (TSDataType::String, Some(Value::String(s))) => { + write_binary(buf, s.as_bytes()); + } + (TSDataType::Blob, Some(Value::Blob(b))) => write_binary(buf, b), + (TSDataType::Text | TSDataType::String | TSDataType::Blob, None) => write_binary(buf, &[]), + // Null sentinel 10000101 = 1000-01-01 (yyyyMMdd), per C#/Java. + (TSDataType::Date, Some(Value::Date(v))) => buf.extend_from_slice(&v.to_be_bytes()), + (TSDataType::Date, None) => buf.extend_from_slice(&10000101i32.to_be_bytes()), + // add_row validates cell types; Vector/Unknown are not insertable. + (ty, cell) => unreachable!("cell {cell:?} does not match tablet column type {ty:?}"), + } +} + +/// 4-byte big-endian length prefix + raw bytes (TEXT/STRING/BLOB, §3.2). +fn write_binary(buf: &mut Vec<u8>, bytes: &[u8]) { + buf.extend_from_slice(&(bytes.len() as i32).to_be_bytes()); + buf.extend_from_slice(bytes); +} + +#[cfg(test)] +mod tests { + use super::*; + + fn tree_tablet(types: Vec<TSDataType>) -> Tablet { + let measurements = (0..types.len()).map(|i| format!("s{i}")).collect(); + Tablet::new("root.sg.d1", measurements, types).unwrap() + } + + #[test] + fn int32_column_with_null_known_bytes() { + let mut t = tree_tablet(vec![TSDataType::Int32]); + t.add_row(1, vec![Some(Value::Int32(5))]).unwrap(); + t.add_row(2, vec![None]).unwrap(); + let expected = [ + 0x00, 0x00, 0x00, 0x05, // row 0: 5 + 0x80, 0x00, 0x00, 0x00, // row 1: i32::MIN placeholder + 0x01, // bitmap flag: has nulls + 0x02, // LSB-first: row 1 null → bit 1 → 0x02 + ]; + assert_eq!(t.serialize_values(), expected); + assert_eq!( + t.serialize_timestamps(), + [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 2] + ); + } + + #[test] + fn no_nulls_writes_zero_flag_only() { + let mut t = tree_tablet(vec![TSDataType::Boolean]); + t.add_row(1, vec![Some(Value::Boolean(true))]).unwrap(); + t.add_row(2, vec![Some(Value::Boolean(false))]).unwrap(); + assert_eq!(t.serialize_values(), [0x01, 0x00, 0x00]); + } + + #[test] + fn all_types_known_bytes() { + let types = vec![ + TSDataType::Boolean, + TSDataType::Int32, + TSDataType::Int64, + TSDataType::Float, + TSDataType::Double, + TSDataType::Text, + TSDataType::Timestamp, + TSDataType::Date, + TSDataType::Blob, + TSDataType::String, + ]; + let mut t = tree_tablet(types); + t.add_row( + 100, + vec![ + Some(Value::Boolean(true)), + Some(Value::Int32(7)), + Some(Value::Int64(-2)), + Some(Value::Float(1.5)), + Some(Value::Double(-0.5)), + Some(Value::Text("ab".into())), + Some(Value::Timestamp(3)), + Some(Value::Date(20260710)), + Some(Value::Blob(vec![0xDE, 0xAD])), + Some(Value::String("é".into())), + ], + ) + .unwrap(); + + let mut expected: Vec<u8> = Vec::new(); + expected.push(0x01); // true + expected.extend_from_slice(&7i32.to_be_bytes()); + expected.extend_from_slice(&(-2i64).to_be_bytes()); + expected.extend_from_slice(&[0x3F, 0xC0, 0x00, 0x00]); // 1.5f32 IEEE 754 BE + expected.extend_from_slice(&(-0.5f64).to_be_bytes()); + expected.extend_from_slice(&[0, 0, 0, 2, b'a', b'b']); + expected.extend_from_slice(&3i64.to_be_bytes()); + expected.extend_from_slice(&20260710i32.to_be_bytes()); + expected.extend_from_slice(&[0, 0, 0, 2, 0xDE, 0xAD]); + expected.extend_from_slice(&[0, 0, 0, 2, 0xC3, 0xA9]); // "é" UTF-8 + expected.extend_from_slice(&[0; 10]); // 10 columns, no nulls + assert_eq!(t.serialize_values(), expected); + } + + #[test] + fn all_types_null_placeholders() { + let types = vec![ + TSDataType::Boolean, + TSDataType::Int32, + TSDataType::Int64, + TSDataType::Float, + TSDataType::Double, + TSDataType::Text, + TSDataType::Timestamp, + TSDataType::Date, + TSDataType::Blob, + TSDataType::String, + ]; + let n = types.len(); + let mut t = tree_tablet(types); + t.add_row(1, vec![None; n]).unwrap(); + + let mut expected: Vec<u8> = Vec::new(); + expected.push(0x00); // false + expected.extend_from_slice(&i32::MIN.to_be_bytes()); + expected.extend_from_slice(&i64::MIN.to_be_bytes()); + expected.extend_from_slice(&f32::MIN.to_be_bytes()); + expected.extend_from_slice(&f64::MIN.to_be_bytes()); + expected.extend_from_slice(&[0, 0, 0, 0]); // empty text + expected.extend_from_slice(&i64::MIN.to_be_bytes()); + expected.extend_from_slice(&10000101i32.to_be_bytes()); // 1000-01-01 + expected.extend_from_slice(&[0, 0, 0, 0]); // empty blob + expected.extend_from_slice(&[0, 0, 0, 0]); // empty string + for _ in 0..n { + expected.extend_from_slice(&[0x01, 0x01]); // flag + bitmap (row 0 null) + } + assert_eq!(t.serialize_values(), expected); + } + + #[test] + fn unsorted_rows_are_sorted_by_timestamp() { + let mut t = tree_tablet(vec![TSDataType::Int32]); + t.add_row(3, vec![Some(Value::Int32(30))]).unwrap(); + t.add_row(1, vec![None]).unwrap(); + t.add_row(2, vec![Some(Value::Int32(20))]).unwrap(); + + let ts = t.serialize_timestamps(); + assert_eq!( + ts, + [1i64, 2, 3] + .iter() + .flat_map(|v| v.to_be_bytes()) + .collect::<Vec<u8>>() + ); + // Values reordered with their rows; the null moved to row 0. + let expected = [ + 0x80, 0x00, 0x00, 0x00, // null placeholder (was ts=1) + 0x00, 0x00, 0x00, 0x14, // 20 (ts=2) + 0x00, 0x00, 0x00, 0x1E, // 30 (ts=3) + 0x01, 0x01, // flag + LSB-first bitmap: row 0 null + ]; + assert_eq!(t.serialize_values(), expected); + assert_eq!(t.timestamps(), &[1, 2, 3]); + } + + #[test] + fn sort_is_stable_for_equal_timestamps() { + let mut t = tree_tablet(vec![TSDataType::Text]); + t.add_row(5, vec![Some(Value::Text("first".into()))]) + .unwrap(); + t.add_row(1, vec![Some(Value::Text("zero".into()))]) + .unwrap(); + t.add_row(5, vec![Some(Value::Text("second".into()))]) + .unwrap(); + t.sort_by_timestamp(); + assert_eq!(t.timestamps(), &[1, 5, 5]); + assert_eq!(t.values[0][1], Some(Value::Text("first".into()))); + assert_eq!(t.values[0][2], Some(Value::Text("second".into()))); + } + + #[test] + fn bitmap_is_ceil_rows_over_8() { + // 8 rows with a null: exactly 1 bitmap byte (ceil, not rows/8 + 1). + let mut t = tree_tablet(vec![TSDataType::Int32]); + for i in 0..8 { + let cell = if i == 7 { None } else { Some(Value::Int32(i)) }; + t.add_row(i64::from(i), vec![cell]).unwrap(); + } + let buf = t.serialize_values(); + assert_eq!(buf.len(), 8 * 4 + 1 + 1); + assert_eq!(buf[8 * 4], 0x01); // flag + assert_eq!(buf[8 * 4 + 1], 0x80); // row 7 null, LSB-first → bit 7 + + // 9 rows: 2 bitmap bytes. + let mut t = tree_tablet(vec![TSDataType::Int32]); + for i in 0..9 { + let cell = if i == 8 { None } else { Some(Value::Int32(i)) }; + t.add_row(i64::from(i), vec![cell]).unwrap(); + } + let buf = t.serialize_values(); + assert_eq!(buf.len(), 9 * 4 + 1 + 2); + assert_eq!(&buf[9 * 4..], [0x01, 0x00, 0x01]); // row 8 → byte 1 bit 0 + } + + #[test] + fn multi_column_is_column_major() { + let mut t = tree_tablet(vec![TSDataType::Int32, TSDataType::Boolean]); + t.add_row(1, vec![Some(Value::Int32(1)), Some(Value::Boolean(true))]) + .unwrap(); + t.add_row(2, vec![Some(Value::Int32(2)), Some(Value::Boolean(false))]) + .unwrap(); + let expected = [ + 0, 0, 0, 1, 0, 0, 0, 2, // int32 column, both rows + 1, 0, // boolean column, both rows + 0, 0, // both flags: no nulls + ]; + assert_eq!(t.serialize_values(), expected); + } + + #[test] + fn empty_tablet_serializes_flags_only() { + let mut t = tree_tablet(vec![TSDataType::Int64]); + assert_eq!(t.serialize_values(), [0x00]); + assert_eq!(t.serialize_timestamps(), Vec::<u8>::new()); + } + + #[test] + fn table_model_tablet() { + let t = Tablet::new_table( + "sensors", + vec!["tag1".into(), "f1".into(), "attr1".into()], + vec![TSDataType::String, TSDataType::Double, TSDataType::String], + vec![ + ColumnCategory::Tag, + ColumnCategory::Field, + ColumnCategory::Attribute, + ], + ) + .unwrap(); + assert!(t.is_table_model()); + assert_eq!(t.table_name(), "sensors"); + assert_eq!( + t.column_categories().unwrap(), + &[ + ColumnCategory::Tag, + ColumnCategory::Field, + ColumnCategory::Attribute + ] + ); + } + + #[test] + fn constructor_and_add_row_validation() { + assert!(Tablet::new("d", vec!["s1".into()], vec![]).is_err()); + assert!(Tablet::new_table( + "t", + vec!["c1".into()], + vec![TSDataType::Int32], + vec![ColumnCategory::Tag, ColumnCategory::Field], + ) + .is_err()); + + let mut t = tree_tablet(vec![TSDataType::Int32]); + // Wrong arity. + assert!(t.add_row(1, vec![]).is_err()); + // Wrong type. + assert!(t.add_row(1, vec![Some(Value::Int64(1))]).is_err()); + // Explicit Value::Null behaves like None. + t.add_row(1, vec![Some(Value::Null)]).unwrap(); + assert_eq!(&t.serialize_values()[4..], [0x01, 0x01]); // flag + null bitmap + assert_eq!(t.row_count(), 1); + } +} diff --git a/src/data/tsblock.rs b/src/data/tsblock.rs new file mode 100644 index 0000000..a393cf2 --- /dev/null +++ b/src/data/tsblock.rs @@ -0,0 +1,521 @@ +// 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). +const ENCODING_BYTE_ARRAY: u8 = 0; +const ENCODING_INT32_ARRAY: u8 = 1; +const ENCODING_INT64_ARRAY: u8 = 2; +const ENCODING_BINARY_ARRAY: u8 = 3; +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), + TSDataType::Text => Value::Text(decode_utf8(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], + ])) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Builds a TsBlock header: valueColumnCount, type bytes, positionCount, + /// time encoding (Int64Array), value encodings. + 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. + 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 + } + + #[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(_)))); + } +} diff --git a/src/data/value.rs b/src/data/value.rs new file mode 100644 index 0000000..947b073 --- /dev/null +++ b/src/data/value.rs @@ -0,0 +1,97 @@ +// 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. + +//! A dynamically-typed IoTDB cell value. + +use super::TSDataType; + +/// One cell of an IoTDB row: a typed scalar or `Null`. +/// +/// `Date` carries an `i32` in `yyyyMMdd` form (e.g. 2026-07-10 → +/// `20260710`), matching the C#/Java wire encoding. `Timestamp` is epoch +/// milliseconds. +#[derive(Debug, Clone, PartialEq)] +pub enum Value { + Boolean(bool), + Int32(i32), + Int64(i64), + Float(f32), + Double(f64), + Text(String), + Timestamp(i64), + /// Date as i32 `yyyyMMdd` (e.g. `20260710`). + Date(i32), + Blob(Vec<u8>), + String(String), + Null, +} + +impl Value { + /// The [`TSDataType`] this value carries, or `None` for [`Value::Null`]. + pub fn data_type(&self) -> Option<TSDataType> { + Some(match self { + Value::Boolean(_) => TSDataType::Boolean, + Value::Int32(_) => TSDataType::Int32, + Value::Int64(_) => TSDataType::Int64, + Value::Float(_) => TSDataType::Float, + Value::Double(_) => TSDataType::Double, + Value::Text(_) => TSDataType::Text, + Value::Timestamp(_) => TSDataType::Timestamp, + Value::Date(_) => TSDataType::Date, + Value::Blob(_) => TSDataType::Blob, + Value::String(_) => TSDataType::String, + Value::Null => return None, + }) + } + + /// The wire type code (§8 of the protocol spec), or `None` for `Null`. + pub fn type_code(&self) -> Option<i32> { + self.data_type().map(TSDataType::code) + } + + /// True iff this is [`Value::Null`]. + pub fn is_null(&self) -> bool { + matches!(self, Value::Null) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn type_codes() { + assert_eq!(Value::Boolean(true).type_code(), Some(0)); + assert_eq!(Value::Int32(1).type_code(), Some(1)); + assert_eq!(Value::Int64(1).type_code(), Some(2)); + assert_eq!(Value::Float(1.0).type_code(), Some(3)); + assert_eq!(Value::Double(1.0).type_code(), Some(4)); + assert_eq!(Value::Text("t".into()).type_code(), Some(5)); + assert_eq!(Value::Timestamp(0).type_code(), Some(8)); + assert_eq!(Value::Date(20260710).type_code(), Some(9)); + assert_eq!(Value::Blob(vec![0]).type_code(), Some(10)); + assert_eq!(Value::String("s".into()).type_code(), Some(11)); + assert_eq!(Value::Null.type_code(), None); + } + + #[test] + fn data_type_of_null_is_none() { + assert_eq!(Value::Null.data_type(), None); + assert!(Value::Null.is_null()); + assert!(!Value::Int32(7).is_null()); + } +} diff --git a/src/error.rs b/src/error.rs index 6762061..6da4412 100644 --- a/src/error.rs +++ b/src/error.rs @@ -12,6 +12,8 @@ pub enum Error { Server { code: i32, message: String }, /// Client-side usage or state error (e.g. session not open). Client(String), + /// Malformed binary payload received from the server (e.g. truncated TsBlock). + Decode(String), } impl fmt::Display for Error { @@ -20,6 +22,7 @@ impl fmt::Display for Error { Error::Thrift(e) => write!(f, "thrift error: {e}"), Error::Server { code, message } => write!(f, "server error {code}: {message}"), Error::Client(msg) => write!(f, "client error: {msg}"), + Error::Decode(msg) => write!(f, "decode error: {msg}"), } } } diff --git a/src/lib.rs b/src/lib.rs index bfae2e2..130d7a8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,5 +13,6 @@ pub mod data; pub mod error; pub mod protocol; -pub use client::session::Session; +pub use client::session::{QueryHandle, Session, SessionConfig}; +pub use connection::Endpoint; pub use error::{Error, Result}; diff --git a/src/protocol/mod.rs b/src/protocol/mod.rs index 4b61da4..13d8478 100644 --- a/src/protocol/mod.rs +++ b/src/protocol/mod.rs @@ -8,8 +8,22 @@ //! downloaded by the IoTDB Maven build (`iotdb-protocol/*/target/thrift/bin/thrift`, //! version pinned by the IoTDB pom). Never hand-edit generated files. -#[allow(clippy::all, unused_imports, dead_code, deprecated, unused_variables, unreachable_patterns)] +#[allow( + clippy::all, + unused_imports, + dead_code, + deprecated, + unused_variables, + unreachable_patterns +)] pub mod common; -#[allow(clippy::all, unused_imports, dead_code, deprecated, unused_variables, unreachable_patterns)] +#[allow( + clippy::all, + unused_imports, + dead_code, + deprecated, + unused_variables, + unreachable_patterns +)] pub mod client;
