velesdb_core/column_store/
batch.rs

1//! Batch operations for `ColumnStore`.
2//!
3//! This module provides batch update, TTL expiration, and upsert operations.
4
5use std::collections::HashMap;
6
7use super::types::{
8    BatchUpdate, BatchUpdateResult, BatchUpsertResult, ColumnStoreError, ColumnValue, ExpireResult,
9    TypedColumn, UpsertResult,
10};
11use super::ColumnStore;
12
13impl ColumnStore {
14    /// Performs batch updates with optimized cache locality.
15    pub fn batch_update(&mut self, updates: &[BatchUpdate]) -> BatchUpdateResult {
16        let mut result = BatchUpdateResult::default();
17        let by_column = self.partition_batch_updates(updates, &mut result);
18        let row_to_pk = self.build_row_to_pk_map(updates);
19        self.apply_column_updates(by_column, &row_to_pk, &mut result);
20        result
21    }
22
23    /// Partitions batch updates by column, rejecting invalid ones into `result.failed`.
24    fn partition_batch_updates<'a>(
25        &self,
26        updates: &'a [BatchUpdate],
27        result: &mut BatchUpdateResult,
28    ) -> HashMap<&'a str, Vec<(usize, ColumnValue)>> {
29        let mut by_column: HashMap<&str, Vec<(usize, ColumnValue)>> = HashMap::new();
30
31        for update in updates {
32            if self
33                .primary_key_column
34                .as_ref()
35                .is_some_and(|pk_col| pk_col == &update.column)
36            {
37                result
38                    .failed
39                    .push((update.pk, ColumnStoreError::PrimaryKeyUpdate));
40                continue;
41            }
42
43            match self.primary_index.get(&update.pk) {
44                Some(&row_idx) if !self.is_row_deleted_bitmap(row_idx) => {
45                    by_column
46                        .entry(update.column.as_str())
47                        .or_default()
48                        .push((row_idx, update.value.clone()));
49                }
50                _ => {
51                    result
52                        .failed
53                        .push((update.pk, ColumnStoreError::RowNotFound(update.pk)));
54                }
55            }
56        }
57
58        by_column
59    }
60
61    /// Builds a reverse mapping from row index to primary key.
62    fn build_row_to_pk_map(&self, updates: &[BatchUpdate]) -> HashMap<usize, i64> {
63        let mut row_to_pk: HashMap<usize, i64> = HashMap::new();
64        for update in updates {
65            if let Some(&row_idx) = self.primary_index.get(&update.pk) {
66                row_to_pk.insert(row_idx, update.pk);
67            }
68        }
69        row_to_pk
70    }
71
72    /// Applies partitioned column updates, recording successes and failures.
73    fn apply_column_updates(
74        &mut self,
75        by_column: HashMap<&str, Vec<(usize, ColumnValue)>>,
76        row_to_pk: &HashMap<usize, i64>,
77        result: &mut BatchUpdateResult,
78    ) {
79        for (col_name, col_updates) in by_column {
80            if let Some(col) = self.columns.get_mut(col_name) {
81                for (row_idx, value) in col_updates {
82                    let actual_type = Self::value_type_name(&value);
83                    if Self::set_column_value(col, row_idx, value).is_ok() {
84                        result.successful += 1;
85                    } else {
86                        let pk = row_to_pk.get(&row_idx).copied().unwrap_or(0);
87                        result.failed.push((
88                            pk,
89                            ColumnStoreError::TypeMismatch {
90                                expected: Self::column_type_name(col),
91                                actual: actual_type,
92                            },
93                        ));
94                    }
95                }
96            } else {
97                for (row_idx, _) in col_updates {
98                    let pk = row_to_pk.get(&row_idx).copied().unwrap_or(0);
99                    result
100                        .failed
101                        .push((pk, ColumnStoreError::ColumnNotFound(col_name.to_string())));
102                }
103            }
104        }
105    }
106
107    /// Batch update with same value for multiple primary keys.
108    pub fn batch_update_same_value(
109        &mut self,
110        pks: &[i64],
111        column: &str,
112        value: &ColumnValue,
113    ) -> BatchUpdateResult {
114        let updates: Vec<BatchUpdate> = pks
115            .iter()
116            .map(|&pk| BatchUpdate {
117                pk,
118                column: column.to_string(),
119                value: value.clone(),
120            })
121            .collect();
122        self.batch_update(&updates)
123    }
124
125    /// Sets a TTL (Time To Live) on a row.
126    ///
127    /// # Errors
128    ///
129    /// Returns an error when `pk` is missing or points to a deleted row.
130    pub fn set_ttl(&mut self, pk: i64, ttl_seconds: u64) -> Result<(), ColumnStoreError> {
131        let row_idx = self.resolve_live_row(pk)?;
132        let expiry_ts = Self::now_timestamp() + ttl_seconds;
133        self.row_expiry.insert(row_idx, expiry_ts);
134        Ok(())
135    }
136
137    /// Expires all rows that have passed their TTL.
138    pub fn expire_rows(&mut self) -> ExpireResult {
139        let now = Self::now_timestamp();
140        let mut result = ExpireResult::default();
141
142        let expired_rows: Vec<usize> = self
143            .row_expiry
144            .iter()
145            .filter(|(_, &expiry)| expiry <= now)
146            .map(|(&row_idx, _)| row_idx)
147            .collect();
148
149        for row_idx in expired_rows {
150            if let Some(&pk) = self.row_idx_to_pk.get(&row_idx) {
151                if let Ok(idx) = u32::try_from(row_idx) {
152                    self.deletion_bitmap.insert(idx);
153                }
154                self.row_expiry.remove(&row_idx);
155                result.pks.push(pk);
156                result.expired_count += 1;
157            }
158        }
159
160        result
161    }
162
163    /// Upsert: inserts a new row or updates an existing one.
164    ///
165    /// # Errors
166    ///
167    /// Returns an error when primary key constraints are violated,
168    /// when a referenced column does not exist, or when types mismatch.
169    pub fn upsert(
170        &mut self,
171        values: &[(&str, ColumnValue)],
172    ) -> Result<UpsertResult, ColumnStoreError> {
173        let pk_col = self
174            .primary_key_column
175            .clone()
176            .ok_or(ColumnStoreError::MissingPrimaryKey)?;
177
178        let pk_value = Self::extract_pk_value(values, &pk_col)?;
179        Self::validate_columns_exist(&self.columns, values, &pk_col)?;
180
181        if let Some(&row_idx) = self.primary_index.get(&pk_value) {
182            self.upsert_existing_row(values, row_idx, &pk_col)
183        } else {
184            self.insert_row(values)?;
185            Ok(UpsertResult::Inserted)
186        }
187    }
188
189    /// Validates that all non-pk columns referenced in values actually exist.
190    fn validate_columns_exist(
191        columns: &HashMap<String, TypedColumn>,
192        values: &[(&str, ColumnValue)],
193        pk_col: &str,
194    ) -> Result<(), ColumnStoreError> {
195        for (col_name, _) in values {
196            if *col_name != pk_col && !columns.contains_key(*col_name) {
197                return Err(ColumnStoreError::ColumnNotFound((*col_name).to_string()));
198            }
199        }
200        Ok(())
201    }
202
203    /// Handles upsert for an existing row (either deleted or live).
204    fn upsert_existing_row(
205        &mut self,
206        values: &[(&str, ColumnValue)],
207        row_idx: usize,
208        pk_col: &str,
209    ) -> Result<UpsertResult, ColumnStoreError> {
210        if self.is_row_deleted_bitmap(row_idx) {
211            Self::validate_value_types(&self.columns, values, Some(pk_col))?;
212            if let Ok(idx) = u32::try_from(row_idx) {
213                self.deletion_bitmap.remove(idx);
214            }
215            self.row_expiry.remove(&row_idx);
216            self.set_row_values(values, row_idx, Some(pk_col))?;
217            return Ok(UpsertResult::Inserted);
218        }
219
220        Self::validate_value_types(&self.columns, values, Some(pk_col))?;
221        self.update_non_pk_values(values, row_idx, pk_col)?;
222        Ok(UpsertResult::Updated)
223    }
224
225    /// Updates only the non-pk columns that are provided in values.
226    fn update_non_pk_values(
227        &mut self,
228        values: &[(&str, ColumnValue)],
229        row_idx: usize,
230        pk_col: &str,
231    ) -> Result<(), ColumnStoreError> {
232        for (col_name, value) in values {
233            if *col_name == pk_col {
234                continue;
235            }
236            if let Some(col) = self.columns.get_mut(*col_name) {
237                Self::set_column_value(col, row_idx, value.clone())?;
238            }
239        }
240        Ok(())
241    }
242
243    /// Batch upsert: inserts or updates multiple rows.
244    pub fn batch_upsert(&mut self, rows: &[Vec<(&str, ColumnValue)>]) -> BatchUpsertResult {
245        let mut result = BatchUpsertResult::default();
246
247        for row in rows {
248            match self.upsert(row) {
249                Ok(UpsertResult::Inserted) => result.inserted += 1,
250                Ok(UpsertResult::Updated) => result.updated += 1,
251                Err(e) => {
252                    let pk = row
253                        .iter()
254                        .find(|(name, _)| {
255                            self.primary_key_column
256                                .as_ref()
257                                .is_some_and(|pk| pk.as_str() == *name)
258                        })
259                        .and_then(|(_, v)| {
260                            if let ColumnValue::Int(pk) = v {
261                                Some(*pk)
262                            } else {
263                                None
264                            }
265                        })
266                        .unwrap_or(0);
267                    result.failed.push((pk, e));
268                }
269            }
270        }
271
272        result
273    }
274
275    pub(super) fn validate_type_match(
276        col: &TypedColumn,
277        value: &ColumnValue,
278    ) -> Result<(), ColumnStoreError> {
279        let type_matches = matches!(
280            (col, value),
281            (TypedColumn::Int(_), ColumnValue::Int(_))
282                | (TypedColumn::Float(_), ColumnValue::Float(_))
283                | (TypedColumn::String(_), ColumnValue::String(_))
284                | (TypedColumn::Bool(_), ColumnValue::Bool(_))
285                | (TypedColumn::Array { .. }, ColumnValue::Array(_))
286                | (TypedColumn::GeoPoint(_), ColumnValue::GeoPoint(_, _))
287                | (_, ColumnValue::Null)
288        );
289
290        if type_matches {
291            Ok(())
292        } else {
293            Err(ColumnStoreError::TypeMismatch {
294                expected: Self::column_type_name(col),
295                actual: Self::value_type_name(value),
296            })
297        }
298    }
299
300    pub(super) fn set_column_value(
301        col: &mut TypedColumn,
302        row_idx: usize,
303        value: ColumnValue,
304    ) -> Result<(), ColumnStoreError> {
305        if matches!(value, ColumnValue::Null) {
306            return Self::set_column_null(col, row_idx);
307        }
308
309        match (col, value) {
310            (TypedColumn::Int(vec), ColumnValue::Int(v)) => {
311                Self::checked_set(vec, row_idx, Some(v))
312            }
313            (TypedColumn::Float(vec), ColumnValue::Float(v)) => {
314                Self::checked_set(vec, row_idx, Some(v))
315            }
316            (TypedColumn::String(vec), ColumnValue::String(v)) => {
317                Self::checked_set(vec, row_idx, Some(v))
318            }
319            (TypedColumn::Bool(vec), ColumnValue::Bool(v)) => {
320                Self::checked_set(vec, row_idx, Some(v))
321            }
322            (TypedColumn::Array { data, .. }, ColumnValue::Array(arr)) => {
323                Self::checked_set_array(data, row_idx, arr)
324            }
325            (TypedColumn::GeoPoint(vec), ColumnValue::GeoPoint(lat, lng)) => {
326                Self::checked_set_geopoint(vec, row_idx, lat, lng)
327            }
328            (col, value) => Err(ColumnStoreError::TypeMismatch {
329                expected: Self::column_type_name(col),
330                actual: Self::value_type_name(&value),
331            }),
332        }
333    }
334
335    /// Sets an array column cell at `row_idx` with bounds checking.
336    fn checked_set_array(
337        data: &mut [Option<smallvec::SmallVec<[ColumnValue; 8]>>],
338        row_idx: usize,
339        arr: Vec<ColumnValue>,
340    ) -> Result<(), ColumnStoreError> {
341        if row_idx >= data.len() {
342            return Err(ColumnStoreError::IndexOutOfBounds(row_idx));
343        }
344        data[row_idx] = Some(smallvec::SmallVec::from_vec(arr));
345        Ok(())
346    }
347
348    /// Sets a GeoPoint column cell at `row_idx` with bounds checking.
349    fn checked_set_geopoint(
350        vec: &mut [Option<(f64, f64)>],
351        row_idx: usize,
352        lat: f64,
353        lng: f64,
354    ) -> Result<(), ColumnStoreError> {
355        if row_idx >= vec.len() {
356            return Err(ColumnStoreError::IndexOutOfBounds(row_idx));
357        }
358        vec[row_idx] = Some((lat, lng));
359        Ok(())
360    }
361
362    /// Sets a column cell to null at the given row index.
363    fn set_column_null(col: &mut TypedColumn, row_idx: usize) -> Result<(), ColumnStoreError> {
364        match col {
365            TypedColumn::Int(vec) => Self::checked_set(vec, row_idx, None),
366            TypedColumn::Float(vec) => Self::checked_set(vec, row_idx, None),
367            TypedColumn::String(vec) => Self::checked_set(vec, row_idx, None),
368            TypedColumn::Bool(vec) => Self::checked_set(vec, row_idx, None),
369            TypedColumn::Array { data, .. } => {
370                if row_idx >= data.len() {
371                    return Err(ColumnStoreError::IndexOutOfBounds(row_idx));
372                }
373                data[row_idx] = None;
374                Ok(())
375            }
376            TypedColumn::GeoPoint(vec) => Self::checked_set(vec, row_idx, None),
377        }
378    }
379
380    /// Sets a value at `row_idx` with bounds checking.
381    fn checked_set<T>(
382        vec: &mut [Option<T>],
383        row_idx: usize,
384        value: Option<T>,
385    ) -> Result<(), ColumnStoreError> {
386        if row_idx >= vec.len() {
387            return Err(ColumnStoreError::IndexOutOfBounds(row_idx));
388        }
389        vec[row_idx] = value;
390        Ok(())
391    }
392
393    pub(super) fn column_type_name(col: &TypedColumn) -> String {
394        match col {
395            TypedColumn::Int(_) => "Int".to_string(),
396            TypedColumn::Float(_) => "Float".to_string(),
397            TypedColumn::String(_) => "String".to_string(),
398            TypedColumn::Bool(_) => "Bool".to_string(),
399            TypedColumn::Array { .. } => "Array".to_string(),
400            TypedColumn::GeoPoint(_) => "GeoPoint".to_string(),
401        }
402    }
403
404    pub(super) fn value_type_name(value: &ColumnValue) -> String {
405        match value {
406            ColumnValue::Int(_) => "Int".to_string(),
407            ColumnValue::Float(_) => "Float".to_string(),
408            ColumnValue::String(_) => "String".to_string(),
409            ColumnValue::Bool(_) => "Bool".to_string(),
410            ColumnValue::Null => "Null".to_string(),
411            ColumnValue::Array(_) => "Array".to_string(),
412            ColumnValue::GeoPoint(_, _) => "GeoPoint".to_string(),
413        }
414    }
415
416    pub(super) fn now_timestamp() -> u64 {
417        std::time::SystemTime::now()
418            .duration_since(std::time::UNIX_EPOCH)
419            .map_or(0, |d| d.as_secs())
420    }
421}
422
423// Tests moved to batch_tests.rs per project rules
velesdb_core/column_store/batch.rs

velesdb_core/column_store/
batch.rs
