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.deleted_rows.contains(&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                self.deleted_rows.insert(row_idx);
152                // BUG-2 FIX: Also update RoaringBitmap to keep both in sync
153                if let Ok(idx) = u32::try_from(row_idx) {
154                    self.deletion_bitmap.insert(idx);
155                }
156                self.row_expiry.remove(&row_idx);
157                result.pks.push(pk);
158                result.expired_count += 1;
159            }
160        }
161
162        result
163    }
164
165    /// Upsert: inserts a new row or updates an existing one.
166    ///
167    /// # Errors
168    ///
169    /// Returns an error when primary key constraints are violated,
170    /// when a referenced column does not exist, or when types mismatch.
171    pub fn upsert(
172        &mut self,
173        values: &[(&str, ColumnValue)],
174    ) -> Result<UpsertResult, ColumnStoreError> {
175        let pk_col = self
176            .primary_key_column
177            .clone()
178            .ok_or(ColumnStoreError::MissingPrimaryKey)?;
179
180        let pk_value = Self::extract_pk_value(values, &pk_col)?;
181        Self::validate_columns_exist(&self.columns, values, &pk_col)?;
182
183        if let Some(&row_idx) = self.primary_index.get(&pk_value) {
184            self.upsert_existing_row(values, row_idx, &pk_col)
185        } else {
186            self.insert_row(values)?;
187            Ok(UpsertResult::Inserted)
188        }
189    }
190
191    /// Validates that all non-pk columns referenced in values actually exist.
192    fn validate_columns_exist(
193        columns: &HashMap<String, TypedColumn>,
194        values: &[(&str, ColumnValue)],
195        pk_col: &str,
196    ) -> Result<(), ColumnStoreError> {
197        for (col_name, _) in values {
198            if *col_name != pk_col && !columns.contains_key(*col_name) {
199                return Err(ColumnStoreError::ColumnNotFound((*col_name).to_string()));
200            }
201        }
202        Ok(())
203    }
204
205    /// Handles upsert for an existing row (either deleted or live).
206    fn upsert_existing_row(
207        &mut self,
208        values: &[(&str, ColumnValue)],
209        row_idx: usize,
210        pk_col: &str,
211    ) -> Result<UpsertResult, ColumnStoreError> {
212        if self.deleted_rows.contains(&row_idx) {
213            Self::validate_value_types(&self.columns, values, Some(pk_col))?;
214            self.deleted_rows.remove(&row_idx);
215            if let Ok(idx) = u32::try_from(row_idx) {
216                self.deletion_bitmap.remove(idx);
217            }
218            self.row_expiry.remove(&row_idx);
219            self.set_row_values(values, row_idx, Some(pk_col))?;
220            return Ok(UpsertResult::Inserted);
221        }
222
223        Self::validate_value_types(&self.columns, values, Some(pk_col))?;
224        self.update_non_pk_values(values, row_idx, pk_col)?;
225        Ok(UpsertResult::Updated)
226    }
227
228    /// Updates only the non-pk columns that are provided in values.
229    fn update_non_pk_values(
230        &mut self,
231        values: &[(&str, ColumnValue)],
232        row_idx: usize,
233        pk_col: &str,
234    ) -> Result<(), ColumnStoreError> {
235        for (col_name, value) in values {
236            if *col_name == pk_col {
237                continue;
238            }
239            if let Some(col) = self.columns.get_mut(*col_name) {
240                Self::set_column_value(col, row_idx, value.clone())?;
241            }
242        }
243        Ok(())
244    }
245
246    /// Batch upsert: inserts or updates multiple rows.
247    pub fn batch_upsert(&mut self, rows: &[Vec<(&str, ColumnValue)>]) -> BatchUpsertResult {
248        let mut result = BatchUpsertResult::default();
249
250        for row in rows {
251            match self.upsert(row) {
252                Ok(UpsertResult::Inserted) => result.inserted += 1,
253                Ok(UpsertResult::Updated) => result.updated += 1,
254                Err(e) => {
255                    let pk = row
256                        .iter()
257                        .find(|(name, _)| {
258                            self.primary_key_column
259                                .as_ref()
260                                .is_some_and(|pk| pk.as_str() == *name)
261                        })
262                        .and_then(|(_, v)| {
263                            if let ColumnValue::Int(pk) = v {
264                                Some(*pk)
265                            } else {
266                                None
267                            }
268                        })
269                        .unwrap_or(0);
270                    result.failed.push((pk, e));
271                }
272            }
273        }
274
275        result
276    }
277
278    pub(super) fn validate_type_match(
279        col: &TypedColumn,
280        value: &ColumnValue,
281    ) -> Result<(), ColumnStoreError> {
282        let type_matches = matches!(
283            (col, value),
284            (TypedColumn::Int(_), ColumnValue::Int(_))
285                | (TypedColumn::Float(_), ColumnValue::Float(_))
286                | (TypedColumn::String(_), ColumnValue::String(_))
287                | (TypedColumn::Bool(_), ColumnValue::Bool(_))
288                | (TypedColumn::Array { .. }, ColumnValue::Array(_))
289                | (TypedColumn::GeoPoint(_), ColumnValue::GeoPoint(_, _))
290                | (_, ColumnValue::Null)
291        );
292
293        if type_matches {
294            Ok(())
295        } else {
296            Err(ColumnStoreError::TypeMismatch {
297                expected: Self::column_type_name(col),
298                actual: Self::value_type_name(value),
299            })
300        }
301    }
302
303    pub(super) fn set_column_value(
304        col: &mut TypedColumn,
305        row_idx: usize,
306        value: ColumnValue,
307    ) -> Result<(), ColumnStoreError> {
308        if matches!(value, ColumnValue::Null) {
309            return Self::set_column_null(col, row_idx);
310        }
311
312        match (col, value) {
313            (TypedColumn::Int(vec), ColumnValue::Int(v)) => {
314                Self::checked_set(vec, row_idx, Some(v))
315            }
316            (TypedColumn::Float(vec), ColumnValue::Float(v)) => {
317                Self::checked_set(vec, row_idx, Some(v))
318            }
319            (TypedColumn::String(vec), ColumnValue::String(v)) => {
320                Self::checked_set(vec, row_idx, Some(v))
321            }
322            (TypedColumn::Bool(vec), ColumnValue::Bool(v)) => {
323                Self::checked_set(vec, row_idx, Some(v))
324            }
325            (TypedColumn::Array { data, .. }, ColumnValue::Array(arr)) => {
326                Self::checked_set_array(data, row_idx, arr)
327            }
328            (TypedColumn::GeoPoint(vec), ColumnValue::GeoPoint(lat, lng)) => {
329                Self::checked_set_geopoint(vec, row_idx, lat, lng)
330            }
331            (col, value) => Err(ColumnStoreError::TypeMismatch {
332                expected: Self::column_type_name(col),
333                actual: Self::value_type_name(&value),
334            }),
335        }
336    }
337
338    /// Sets an array column cell at `row_idx` with bounds checking.
339    fn checked_set_array(
340        data: &mut [Option<smallvec::SmallVec<[ColumnValue; 8]>>],
341        row_idx: usize,
342        arr: Vec<ColumnValue>,
343    ) -> Result<(), ColumnStoreError> {
344        if row_idx >= data.len() {
345            return Err(ColumnStoreError::IndexOutOfBounds(row_idx));
346        }
347        data[row_idx] = Some(smallvec::SmallVec::from_vec(arr));
348        Ok(())
349    }
350
351    /// Sets a GeoPoint column cell at `row_idx` with bounds checking.
352    fn checked_set_geopoint(
353        vec: &mut [Option<(f64, f64)>],
354        row_idx: usize,
355        lat: f64,
356        lng: f64,
357    ) -> Result<(), ColumnStoreError> {
358        if row_idx >= vec.len() {
359            return Err(ColumnStoreError::IndexOutOfBounds(row_idx));
360        }
361        vec[row_idx] = Some((lat, lng));
362        Ok(())
363    }
364
365    /// Sets a column cell to null at the given row index.
366    fn set_column_null(col: &mut TypedColumn, row_idx: usize) -> Result<(), ColumnStoreError> {
367        match col {
368            TypedColumn::Int(vec) => Self::checked_set(vec, row_idx, None),
369            TypedColumn::Float(vec) => Self::checked_set(vec, row_idx, None),
370            TypedColumn::String(vec) => Self::checked_set(vec, row_idx, None),
371            TypedColumn::Bool(vec) => Self::checked_set(vec, row_idx, None),
372            TypedColumn::Array { data, .. } => {
373                if row_idx >= data.len() {
374                    return Err(ColumnStoreError::IndexOutOfBounds(row_idx));
375                }
376                data[row_idx] = None;
377                Ok(())
378            }
379            TypedColumn::GeoPoint(vec) => Self::checked_set(vec, row_idx, None),
380        }
381    }
382
383    /// Sets a value at `row_idx` with bounds checking.
384    fn checked_set<T>(
385        vec: &mut [Option<T>],
386        row_idx: usize,
387        value: Option<T>,
388    ) -> Result<(), ColumnStoreError> {
389        if row_idx >= vec.len() {
390            return Err(ColumnStoreError::IndexOutOfBounds(row_idx));
391        }
392        vec[row_idx] = value;
393        Ok(())
394    }
395
396    pub(super) fn column_type_name(col: &TypedColumn) -> String {
397        match col {
398            TypedColumn::Int(_) => "Int".to_string(),
399            TypedColumn::Float(_) => "Float".to_string(),
400            TypedColumn::String(_) => "String".to_string(),
401            TypedColumn::Bool(_) => "Bool".to_string(),
402            TypedColumn::Array { .. } => "Array".to_string(),
403            TypedColumn::GeoPoint(_) => "GeoPoint".to_string(),
404        }
405    }
406
407    pub(super) fn value_type_name(value: &ColumnValue) -> String {
408        match value {
409            ColumnValue::Int(_) => "Int".to_string(),
410            ColumnValue::Float(_) => "Float".to_string(),
411            ColumnValue::String(_) => "String".to_string(),
412            ColumnValue::Bool(_) => "Bool".to_string(),
413            ColumnValue::Null => "Null".to_string(),
414            ColumnValue::Array(_) => "Array".to_string(),
415            ColumnValue::GeoPoint(_, _) => "GeoPoint".to_string(),
416        }
417    }
418
419    pub(super) fn now_timestamp() -> u64 {
420        std::time::SystemTime::now()
421            .duration_since(std::time::UNIX_EPOCH)
422            .map_or(0, |d| d.as_secs())
423    }
424}
425
426// Tests moved to batch_tests.rs per project rules
velesdb_core/column_store/batch.rs

velesdb_core/column_store/
batch.rs
