exoware_sql/
writer.rs

1use std::any::Any;
2use std::collections::HashMap;
3use std::fmt;
4use std::sync::Arc;
5use std::sync::Mutex;
6
7use async_trait::async_trait;
8use bytes::Bytes;
9use commonware_codec::Encode;
10use datafusion::arrow::array::{
11    ArrayRef, BooleanArray, Date32Array, Date64Array, Decimal128Array, Decimal256Array,
12    FixedSizeBinaryArray, Float64Array, Int64Array, LargeStringArray, ListArray, StringArray,
13    StringViewArray, TimestampMicrosecondArray, UInt64Array,
14};
15use datafusion::arrow::datatypes::{i256, SchemaRef};
16use datafusion::arrow::record_batch::RecordBatch;
17use datafusion::common::{DataFusionError, Result as DataFusionResult};
18use datafusion::datasource::sink::DataSink;
19use datafusion::execution::context::TaskContext;
20use datafusion::physical_plan::{DisplayAs, DisplayFormatType, SendableRecordBatchStream};
21use exoware_sdk::keys::Key;
22#[cfg(test)]
23use exoware_sdk::kv_codec::decode_stored_row;
24use exoware_sdk::kv_codec::{StoredRow, StoredValue};
25use exoware_sdk::{PrefixedStoreClient, StoreBatchUpload, StoreWriteBatch};
26use futures::{future::BoxFuture, TryStreamExt};
27
28use crate::builder::archived_non_pk_value_is_valid;
29use crate::codec::*;
30use crate::types::*;
31
32#[derive(Debug)]
33pub struct TableWriter {
34    model: Arc<TableModel>,
35    index_specs: Arc<Vec<ResolvedIndexSpec>>,
36}
37
38impl TableWriter {
39    pub fn encode_row(&self, values: Vec<CellValue>) -> Result<Vec<(Key, Vec<u8>)>, String> {
40        let row = KvRow { values };
41        if row.values.len() != self.model.columns.len() {
42            return Err(format!(
43                "expected {} values, got {}",
44                self.model.columns.len(),
45                row.values.len()
46            ));
47        }
48        let base_key = encode_primary_key_from_row(self.model.table_prefix, &row, &self.model)?;
49        let base_value = encode_base_row_value(&row, &self.model).map_err(|e| format!("{e}"))?;
50        let mut out = vec![(base_key, base_value)];
51        for spec in self.index_specs.iter() {
52            let idx_key =
53                encode_secondary_index_key(self.model.table_prefix, spec, &self.model, &row)?;
54            let idx_value = encode_secondary_index_value(&row, &self.model, spec)
55                .map_err(|e| format!("{e}"))?;
56            out.push((idx_key, idx_value));
57        }
58        Ok(out)
59    }
60}
61
62#[derive(Debug)]
63pub struct BatchWriter {
64    client: PrefixedStoreClient,
65    tables: HashMap<String, TableWriter>,
66    next_request_id: u64,
67    failed_prepared: Mutex<Vec<PreparedBatch>>,
68    pub(crate) pending_keys: Vec<Key>,
69    pub(crate) pending_values: Vec<Bytes>,
70}
71
72#[derive(Debug)]
73#[must_use]
74pub struct PreparedBatch {
75    request_id: u64,
76    entry_count: usize,
77    keys: Vec<Key>,
78    values: Vec<Bytes>,
79}
80
81impl PreparedBatch {
82    pub fn request_id(&self) -> u64 {
83        self.request_id
84    }
85
86    pub fn entry_count(&self) -> usize {
87        self.entry_count
88    }
89
90    pub fn is_empty(&self) -> bool {
91        self.keys.is_empty()
92    }
93}
94
95#[derive(Clone, Copy, Debug, PartialEq, Eq)]
96pub struct BatchReceipt {
97    pub writer_request_id: u64,
98    pub entry_count: usize,
99    pub store_sequence_number: u64,
100}
101
102impl BatchWriter {
103    pub(crate) fn new(
104        client: PrefixedStoreClient,
105        table_configs: &[(String, KvTableConfig)],
106    ) -> Self {
107        let mut tables = HashMap::new();
108        for (name, config) in table_configs {
109            let model = Arc::new(
110                TableModel::from_config(config).expect("config already validated by KvSchema"),
111            );
112            let index_specs = Arc::new(
113                model
114                    .resolve_index_specs(&config.index_specs)
115                    .expect("specs already validated by KvSchema"),
116            );
117            tables.insert(name.clone(), TableWriter { model, index_specs });
118        }
119        Self {
120            client,
121            tables,
122            next_request_id: 0,
123            failed_prepared: Mutex::new(Vec::new()),
124            pending_keys: Vec::new(),
125            pending_values: Vec::new(),
126        }
127    }
128
129    pub fn insert(
130        &mut self,
131        table_name: &str,
132        values: Vec<CellValue>,
133    ) -> Result<&mut Self, String> {
134        let writer = self
135            .tables
136            .get(table_name)
137            .ok_or_else(|| format!("unknown table '{table_name}'"))?;
138        let entries = writer.encode_row(values)?;
139        for (key, value) in entries {
140            self.pending_keys.push(key);
141            self.pending_values.push(value.into());
142        }
143        Ok(self)
144    }
145
146    pub fn pending_count(&self) -> usize {
147        self.pending_keys.len()
148            + self
149                .failed_prepared
150                .lock()
151                .expect("failed prepared mutex poisoned")
152                .iter()
153                .map(PreparedBatch::entry_count)
154                .sum::<usize>()
155    }
156
157    /// Flush pending rows to ingest and return the post-ingest consistency token.
158    pub async fn flush(&mut self) -> DataFusionResult<u64> {
159        Ok(self
160            .flush_with_receipt()
161            .await?
162            .map(|receipt| receipt.store_sequence_number)
163            .unwrap_or(0))
164    }
165
166    /// Flush pending rows to ingest and return metadata for the persisted batch.
167    pub async fn flush_with_receipt(&mut self) -> DataFusionResult<Option<BatchReceipt>> {
168        let Some(prepared) = self.prepare_flush()? else {
169            return Ok(None);
170        };
171        Ok(Some(self.commit_upload(prepared).await?))
172    }
173
174    pub fn prepare_flush(&mut self) -> DataFusionResult<Option<PreparedBatch>> {
175        if let Some(prepared) = self.take_failed_prepared() {
176            return Ok(Some(prepared));
177        }
178        if self.pending_keys.is_empty() {
179            return Ok(None);
180        }
181        let request_id = self.next_request_id;
182        self.next_request_id += 1;
183        Ok(Some(PreparedBatch {
184            request_id,
185            entry_count: self.pending_keys.len(),
186            keys: std::mem::take(&mut self.pending_keys),
187            values: std::mem::take(&mut self.pending_values),
188        }))
189    }
190
191    pub fn stage_flush(
192        &self,
193        prepared: &PreparedBatch,
194        batch: &mut StoreWriteBatch,
195    ) -> DataFusionResult<()> {
196        for (key, value) in prepared.keys.iter().zip(prepared.values.iter()) {
197            batch
198                .push(&self.client, key, value)
199                .map_err(|e| DataFusionError::External(Box::new(e)))?;
200        }
201        Ok(())
202    }
203
204    pub fn mark_flush_persisted(
205        &self,
206        prepared: PreparedBatch,
207        sequence_number: u64,
208    ) -> BatchReceipt {
209        BatchReceipt {
210            writer_request_id: prepared.request_id,
211            entry_count: prepared.entry_count(),
212            store_sequence_number: sequence_number,
213        }
214    }
215
216    pub fn mark_flush_failed(&self, prepared: PreparedBatch) {
217        self.failed_prepared
218            .lock()
219            .expect("failed prepared mutex poisoned")
220            .push(prepared);
221    }
222
223    fn take_failed_prepared(&self) -> Option<PreparedBatch> {
224        let mut failed = self
225            .failed_prepared
226            .lock()
227            .expect("failed prepared mutex poisoned");
228        let (idx, _) = failed
229            .iter()
230            .enumerate()
231            .min_by_key(|(_, prepared)| prepared.request_id)?;
232        Some(failed.remove(idx))
233    }
234}
235
236impl StoreBatchUpload for BatchWriter {
237    type Prepared = PreparedBatch;
238    type Receipt = BatchReceipt;
239    type Error = DataFusionError;
240
241    fn store_client(&self) -> &PrefixedStoreClient {
242        &self.client
243    }
244
245    fn stage_upload(
246        &self,
247        prepared: &mut Self::Prepared,
248        batch: &mut StoreWriteBatch,
249    ) -> Result<(), Self::Error> {
250        self.stage_flush(prepared, batch)
251    }
252
253    fn commit_error(&self, error: exoware_sdk::ClientError) -> Self::Error {
254        DataFusionError::External(Box::new(error))
255    }
256
257    fn mark_upload_persisted<'a>(
258        &'a self,
259        prepared: Self::Prepared,
260        sequence_number: u64,
261    ) -> BoxFuture<'a, Self::Receipt>
262    where
263        Self: Sync + 'a,
264        Self::Prepared: 'a,
265    {
266        Box::pin(async move { self.mark_flush_persisted(prepared, sequence_number) })
267    }
268
269    fn mark_upload_failed<'a>(
270        &'a self,
271        prepared: Self::Prepared,
272        _error: String,
273    ) -> BoxFuture<'a, ()>
274    where
275        Self: Sync + 'a,
276        Self::Prepared: 'a,
277    {
278        Box::pin(async move {
279            self.mark_flush_failed(prepared);
280        })
281    }
282}
283
284#[derive(Debug)]
285pub(crate) struct KvIngestSink {
286    pub(crate) client: PrefixedStoreClient,
287    pub(crate) schema: SchemaRef,
288    pub(crate) model: Arc<TableModel>,
289    pub(crate) index_specs: Arc<Vec<ResolvedIndexSpec>>,
290}
291
292impl KvIngestSink {
293    pub(crate) fn new(
294        client: PrefixedStoreClient,
295        schema: SchemaRef,
296        model: Arc<TableModel>,
297        index_specs: Arc<Vec<ResolvedIndexSpec>>,
298    ) -> Self {
299        Self {
300            client,
301            schema,
302            model,
303            index_specs,
304        }
305    }
306}
307
308impl DisplayAs for KvIngestSink {
309    fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter<'_>) -> fmt::Result {
310        write!(f, "KvIngestSink")
311    }
312}
313
314#[async_trait]
315impl DataSink for KvIngestSink {
316    fn as_any(&self) -> &dyn Any {
317        self
318    }
319
320    fn schema(&self) -> &SchemaRef {
321        &self.schema
322    }
323
324    async fn write_all(
325        &self,
326        data: SendableRecordBatchStream,
327        _context: &Arc<TaskContext>,
328    ) -> DataFusionResult<u64> {
329        let mut data = data;
330        let mut pending_keys: Vec<Key> = Vec::new();
331        let mut pending_values: Vec<Bytes> = Vec::new();
332        let mut logical_rows_written = 0u64;
333
334        while let Some(batch) = data.try_next().await? {
335            let encoded_entries = encode_insert_entries(&batch, &self.model, &self.index_specs)?;
336            logical_rows_written += batch.num_rows() as u64;
337            for (key, value) in encoded_entries {
338                pending_keys.push(key);
339                pending_values.push(value.into());
340            }
341        }
342
343        if !pending_keys.is_empty() {
344            flush_ingest_batch(&self.client, &mut pending_keys, &mut pending_values).await?;
345        }
346        Ok(logical_rows_written)
347    }
348}
349
350pub(crate) fn encode_insert_entries(
351    batch: &RecordBatch,
352    model: &TableModel,
353    index_specs: &[ResolvedIndexSpec],
354) -> DataFusionResult<Vec<(Key, Vec<u8>)>> {
355    let mut out = Vec::with_capacity(batch.num_rows() * (1 + index_specs.len()));
356    for row_idx in 0..batch.num_rows() {
357        let row = extract_row_from_batch(batch, row_idx, model)?;
358        let base_key = encode_primary_key_from_row(model.table_prefix, &row, model)
359            .map_err(DataFusionError::Execution)?;
360        let base_value = encode_base_row_value(&row, model)?;
361        out.push((base_key, base_value));
362
363        for spec in index_specs {
364            let secondary_key = encode_secondary_index_key(model.table_prefix, spec, model, &row)
365                .map_err(DataFusionError::Execution)?;
366            let secondary_value = encode_secondary_index_value(&row, model, spec)?;
367            out.push((secondary_key, secondary_value));
368        }
369    }
370    Ok(out)
371}
372
373pub(crate) fn extract_row_from_batch(
374    batch: &RecordBatch,
375    row_idx: usize,
376    model: &TableModel,
377) -> DataFusionResult<KvRow> {
378    let mut values = Vec::with_capacity(model.columns.len());
379    for col in &model.columns {
380        let array = required_column(batch, &col.name)?;
381        if col.nullable && array.is_null(row_idx) {
382            values.push(CellValue::Null);
383            continue;
384        }
385        let value = match col.kind {
386            ColumnKind::Int64 => CellValue::Int64(i64_value_at(array, row_idx, &col.name)?),
387            ColumnKind::UInt64 => CellValue::UInt64(uint64_value_at(array, row_idx, &col.name)?),
388            ColumnKind::Float64 => CellValue::Float64(f64_value_at(array, row_idx, &col.name)?),
389            ColumnKind::Boolean => CellValue::Boolean(bool_value_at(array, row_idx, &col.name)?),
390            ColumnKind::Date32 => CellValue::Date32(date32_value_at(array, row_idx, &col.name)?),
391            ColumnKind::Date64 => CellValue::Date64(date64_value_at(array, row_idx, &col.name)?),
392            ColumnKind::Timestamp => {
393                CellValue::Timestamp(timestamp_micros_value_at(array, row_idx, &col.name)?)
394            }
395            ColumnKind::Decimal128 => {
396                CellValue::Decimal128(decimal128_value_at(array, row_idx, &col.name)?)
397            }
398            ColumnKind::Decimal256 => {
399                CellValue::Decimal256(decimal256_value_at(array, row_idx, &col.name)?)
400            }
401            ColumnKind::Utf8 => CellValue::Utf8(string_value_at(array, row_idx, &col.name)?),
402            ColumnKind::FixedSizeBinary(_) => {
403                CellValue::FixedBinary(fixed_binary_value_at(array, row_idx, &col.name)?)
404            }
405            ColumnKind::List(elem) => list_value_at(array, row_idx, &col.name, elem)?,
406        };
407        values.push(value);
408    }
409    Ok(KvRow { values })
410}
411
412pub(crate) fn encode_base_row_value(row: &KvRow, model: &TableModel) -> DataFusionResult<Vec<u8>> {
413    let mut values = Vec::with_capacity(model.columns.len());
414    for (idx, col) in model.columns.iter().enumerate() {
415        if model.is_pk_column(idx) {
416            values.push(None);
417            continue;
418        }
419        values.push(encode_non_pk_cell_value(row.value_at(idx), col)?);
420    }
421    let stored_row = StoredRow { values };
422    Ok(stored_row.encode().to_vec())
423}
424
425pub(crate) fn encode_secondary_index_value(
426    row: &KvRow,
427    model: &TableModel,
428    spec: &ResolvedIndexSpec,
429) -> DataFusionResult<Vec<u8>> {
430    let mut values = Vec::with_capacity(model.columns.len());
431    for (idx, col) in model.columns.iter().enumerate() {
432        if model.is_pk_column(idx) || !spec.value_column_mask[idx] {
433            values.push(None);
434            continue;
435        }
436        values.push(encode_non_pk_cell_value(row.value_at(idx), col)?);
437    }
438    let stored_row = StoredRow { values };
439    Ok(stored_row.encode().to_vec())
440}
441
442pub(crate) fn encode_secondary_index_value_from_archived(
443    archived: &StoredRow,
444    model: &TableModel,
445    spec: &ResolvedIndexSpec,
446) -> DataFusionResult<Vec<u8>> {
447    if archived.values.len() != model.columns.len() {
448        return Err(DataFusionError::Execution(
449            "archived row column count mismatch".to_string(),
450        ));
451    }
452    let mut values = Vec::with_capacity(model.columns.len());
453    for (idx, col) in model.columns.iter().enumerate() {
454        if model.is_pk_column(idx) || !spec.value_column_mask[idx] {
455            values.push(None);
456            continue;
457        }
458        let stored_opt = archived.values.get(idx).and_then(|value| value.as_ref());
459        if !archived_non_pk_value_is_valid(col, stored_opt) {
460            return Err(DataFusionError::Execution(format!(
461                "invalid archived value for secondary index column '{}'",
462                col.name
463            )));
464        }
465        values.push(owned_stored_value_from_archived(stored_opt)?);
466    }
467    let stored_row = StoredRow { values };
468    Ok(stored_row.encode().to_vec())
469}
470
471pub(crate) fn encode_non_pk_cell_value(
472    value: &CellValue,
473    col: &ResolvedColumn,
474) -> DataFusionResult<Option<StoredValue>> {
475    match (col.kind, value) {
476        (_, CellValue::Null) => {
477            if !col.nullable {
478                return Err(DataFusionError::Execution(format!(
479                    "column '{}' is not nullable but received NULL",
480                    col.name
481                )));
482            }
483            Ok(None)
484        }
485        (ColumnKind::Int64, CellValue::Int64(v)) => Ok(Some(StoredValue::Int64(*v))),
486        (ColumnKind::UInt64, CellValue::UInt64(v)) => Ok(Some(StoredValue::UInt64(*v))),
487        (ColumnKind::Float64, CellValue::Float64(v)) => Ok(Some(StoredValue::Float64(*v))),
488        (ColumnKind::Boolean, CellValue::Boolean(v)) => Ok(Some(StoredValue::Boolean(*v))),
489        (ColumnKind::Date32, CellValue::Date32(v)) => Ok(Some(StoredValue::Int64(*v as i64))),
490        (ColumnKind::Date64, CellValue::Date64(v)) => Ok(Some(StoredValue::Int64(*v))),
491        (ColumnKind::Timestamp, CellValue::Timestamp(v)) => Ok(Some(StoredValue::Int64(*v))),
492        (ColumnKind::Decimal128, CellValue::Decimal128(v)) => {
493            Ok(Some(StoredValue::Bytes(v.to_le_bytes().to_vec())))
494        }
495        (ColumnKind::Decimal256, CellValue::Decimal256(v)) => {
496            Ok(Some(StoredValue::Bytes(v.to_le_bytes().to_vec())))
497        }
498        (ColumnKind::Utf8, CellValue::Utf8(v)) => Ok(Some(StoredValue::Utf8(v.clone()))),
499        (ColumnKind::FixedSizeBinary(n), CellValue::FixedBinary(v)) => {
500            if v.len() != n {
501                return Err(DataFusionError::Execution(format!(
502                    "column '{}' expects FixedSizeBinary({n}) value with exactly {n} bytes, got {}",
503                    col.name,
504                    v.len()
505                )));
506            }
507            Ok(Some(StoredValue::Bytes(v.clone())))
508        }
509        (ColumnKind::List(elem), CellValue::List(items)) => {
510            let mut stored_items = Vec::with_capacity(items.len());
511            for item in items {
512                let stored_item = match (elem, item) {
513                    (ListElementKind::Int64, CellValue::Int64(v)) => StoredValue::Int64(*v),
514                    (ListElementKind::Float64, CellValue::Float64(v)) => StoredValue::Float64(*v),
515                    (ListElementKind::Boolean, CellValue::Boolean(v)) => StoredValue::Boolean(*v),
516                    (ListElementKind::Utf8, CellValue::Utf8(v)) => StoredValue::Utf8(v.clone()),
517                    _ => {
518                        return Err(DataFusionError::Execution(format!(
519                            "column '{}' list element type mismatch (expected {:?}, got {:?})",
520                            col.name, elem, item
521                        )))
522                    }
523                };
524                stored_items.push(stored_item);
525            }
526            Ok(Some(StoredValue::List(stored_items)))
527        }
528        _ => Err(DataFusionError::Execution(format!(
529            "column '{}' type mismatch (expected {:?}, got {:?})",
530            col.name, col.kind, value
531        ))),
532    }
533}
534
535pub(crate) fn owned_stored_value_from_archived(
536    stored_opt: Option<&StoredValue>,
537) -> DataFusionResult<Option<StoredValue>> {
538    let Some(stored) = stored_opt else {
539        return Ok(None);
540    };
541    Ok(Some(match stored {
542        StoredValue::Int64(v) => StoredValue::Int64(*v),
543        StoredValue::UInt64(v) => StoredValue::UInt64(*v),
544        StoredValue::Float64(v) => StoredValue::Float64(*v),
545        StoredValue::Boolean(v) => StoredValue::Boolean(*v),
546        StoredValue::Utf8(v) => StoredValue::Utf8(v.as_str().to_string()),
547        StoredValue::Bytes(v) => StoredValue::Bytes(v.as_slice().to_vec()),
548        StoredValue::List(items) => {
549            let mut out = Vec::with_capacity(items.len());
550            for item in items.iter() {
551                let owned = owned_stored_value_from_archived(Some(item))?.ok_or_else(|| {
552                    DataFusionError::Execution(
553                        "archived list item unexpectedly decoded as NULL".to_string(),
554                    )
555                })?;
556                out.push(owned);
557            }
558            StoredValue::List(out)
559        }
560    }))
561}
562
563#[cfg(test)]
564pub(crate) fn decode_base_row(
565    pk_values: Vec<CellValue>,
566    value: &[u8],
567    model: &TableModel,
568) -> Option<KvRow> {
569    if pk_values.len() != model.primary_key_indices.len() {
570        return None;
571    }
572    let archived = decode_stored_row(value).ok()?;
573    if archived.values.len() != model.columns.len() {
574        return None;
575    }
576    let mut values = vec![CellValue::Null; model.columns.len()];
577    for (pk_pos, pk_value) in pk_values.into_iter().enumerate() {
578        let col_idx = *model.primary_key_indices.get(pk_pos)?;
579        values[col_idx] = pk_value;
580    }
581
582    for (idx, col) in model.columns.iter().enumerate() {
583        if model.is_pk_column(idx) {
584            continue;
585        }
586        let Some(stored) = archived.values[idx].as_ref() else {
587            if col.nullable {
588                continue;
589            }
590            return None;
591        };
592        values[idx] = match (col.kind, stored) {
593            (ColumnKind::Int64, StoredValue::Int64(v)) => CellValue::Int64(*v),
594            (ColumnKind::UInt64, StoredValue::UInt64(v)) => CellValue::UInt64(*v),
595            (ColumnKind::Float64, StoredValue::Float64(v)) => CellValue::Float64(*v),
596            (ColumnKind::Float64, StoredValue::Int64(v)) => CellValue::Float64(*v as f64),
597            (ColumnKind::Boolean, StoredValue::Boolean(v)) => CellValue::Boolean(*v),
598            (ColumnKind::Date32, StoredValue::Int64(v)) => CellValue::Date32(*v as i32),
599            (ColumnKind::Date64, StoredValue::Int64(v)) => CellValue::Date64(*v),
600            (ColumnKind::Timestamp, StoredValue::Int64(v)) => CellValue::Timestamp(*v),
601            (ColumnKind::Decimal128, StoredValue::Bytes(bytes)) => {
602                let arr: [u8; 16] = bytes.as_slice().try_into().ok()?;
603                CellValue::Decimal128(i128::from_le_bytes(arr))
604            }
605            (ColumnKind::Decimal256, StoredValue::Bytes(bytes)) => {
606                let arr: [u8; 32] = bytes.as_slice().try_into().ok()?;
607                CellValue::Decimal256(i256::from_le_bytes(arr))
608            }
609            (ColumnKind::Utf8, StoredValue::Utf8(v)) => CellValue::Utf8(v.as_str().to_string()),
610            (ColumnKind::FixedSizeBinary(_), StoredValue::Bytes(v)) => {
611                CellValue::FixedBinary(v.as_slice().to_vec())
612            }
613            (ColumnKind::List(elem), StoredValue::List(items)) => {
614                let mut cells = Vec::with_capacity(items.len());
615                for item in items.iter() {
616                    cells.push(decode_list_element_archived(elem, item)?);
617                }
618                CellValue::List(cells)
619            }
620            _ => return None,
621        };
622    }
623    Some(KvRow { values })
624}
625
626pub(crate) fn decode_list_element_archived(
627    elem: ListElementKind,
628    stored: &StoredValue,
629) -> Option<CellValue> {
630    Some(match (elem, stored) {
631        (ListElementKind::Int64, StoredValue::Int64(v)) => CellValue::Int64(*v),
632        (ListElementKind::Float64, StoredValue::Float64(v)) => CellValue::Float64(*v),
633        (ListElementKind::Float64, StoredValue::Int64(v)) => CellValue::Float64(*v as f64),
634        (ListElementKind::Boolean, StoredValue::Boolean(v)) => CellValue::Boolean(*v),
635        (ListElementKind::Utf8, StoredValue::Utf8(v)) => CellValue::Utf8(v.as_str().to_string()),
636        _ => return None,
637    })
638}
639
640pub(crate) fn required_column<'a>(
641    batch: &'a RecordBatch,
642    name: &str,
643) -> DataFusionResult<&'a ArrayRef> {
644    batch.column_by_name(name).ok_or_else(|| {
645        DataFusionError::Execution(format!("insert batch is missing required column '{name}'"))
646    })
647}
648
649pub(crate) fn i64_value_at(
650    array: &ArrayRef,
651    row_idx: usize,
652    column_name: &str,
653) -> DataFusionResult<i64> {
654    if array.is_null(row_idx) {
655        return Err(DataFusionError::Execution(format!(
656            "column '{column_name}' cannot be NULL for kv table insert"
657        )));
658    }
659    let values = array.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
660        DataFusionError::Execution(format!(
661            "column '{column_name}' expected Int64, got {:?}",
662            array.data_type()
663        ))
664    })?;
665    Ok(values.value(row_idx))
666}
667
668pub(crate) fn string_value_at(
669    array: &ArrayRef,
670    row_idx: usize,
671    column_name: &str,
672) -> DataFusionResult<String> {
673    if array.is_null(row_idx) {
674        return Err(DataFusionError::Execution(format!(
675            "column '{column_name}' cannot be NULL for kv table insert"
676        )));
677    }
678    if let Some(values) = array.as_any().downcast_ref::<StringArray>() {
679        return Ok(values.value(row_idx).to_string());
680    }
681    if let Some(values) = array.as_any().downcast_ref::<LargeStringArray>() {
682        return Ok(values.value(row_idx).to_string());
683    }
684    if let Some(values) = array.as_any().downcast_ref::<StringViewArray>() {
685        return Ok(values.value(row_idx).to_string());
686    }
687    Err(DataFusionError::Execution(format!(
688        "column '{column_name}' expected string, got {:?}",
689        array.data_type()
690    )))
691}
692
693pub(crate) fn f64_value_at(
694    array: &ArrayRef,
695    row_idx: usize,
696    column_name: &str,
697) -> DataFusionResult<f64> {
698    if array.is_null(row_idx) {
699        return Err(DataFusionError::Execution(format!(
700            "column '{column_name}' cannot be NULL for kv table insert"
701        )));
702    }
703    let values = array
704        .as_any()
705        .downcast_ref::<Float64Array>()
706        .ok_or_else(|| {
707            DataFusionError::Execution(format!(
708                "column '{column_name}' expected Float64, got {:?}",
709                array.data_type()
710            ))
711        })?;
712    Ok(values.value(row_idx))
713}
714
715pub(crate) fn bool_value_at(
716    array: &ArrayRef,
717    row_idx: usize,
718    column_name: &str,
719) -> DataFusionResult<bool> {
720    if array.is_null(row_idx) {
721        return Err(DataFusionError::Execution(format!(
722            "column '{column_name}' cannot be NULL for kv table insert"
723        )));
724    }
725    let values = array
726        .as_any()
727        .downcast_ref::<BooleanArray>()
728        .ok_or_else(|| {
729            DataFusionError::Execution(format!(
730                "column '{column_name}' expected Boolean, got {:?}",
731                array.data_type()
732            ))
733        })?;
734    Ok(values.value(row_idx))
735}
736
737pub(crate) fn date32_value_at(
738    array: &ArrayRef,
739    row_idx: usize,
740    column_name: &str,
741) -> DataFusionResult<i32> {
742    if array.is_null(row_idx) {
743        return Err(DataFusionError::Execution(format!(
744            "column '{column_name}' cannot be NULL for kv table insert"
745        )));
746    }
747    let values = array
748        .as_any()
749        .downcast_ref::<Date32Array>()
750        .ok_or_else(|| {
751            DataFusionError::Execution(format!(
752                "column '{column_name}' expected Date32, got {:?}",
753                array.data_type()
754            ))
755        })?;
756    Ok(values.value(row_idx))
757}
758
759pub(crate) fn date64_value_at(
760    array: &ArrayRef,
761    row_idx: usize,
762    column_name: &str,
763) -> DataFusionResult<i64> {
764    if array.is_null(row_idx) {
765        return Err(DataFusionError::Execution(format!(
766            "column '{column_name}' cannot be NULL for kv table insert"
767        )));
768    }
769    let values = array
770        .as_any()
771        .downcast_ref::<Date64Array>()
772        .ok_or_else(|| {
773            DataFusionError::Execution(format!(
774                "column '{column_name}' expected Date64, got {:?}",
775                array.data_type()
776            ))
777        })?;
778    Ok(values.value(row_idx))
779}
780
781pub(crate) fn timestamp_micros_value_at(
782    array: &ArrayRef,
783    row_idx: usize,
784    column_name: &str,
785) -> DataFusionResult<i64> {
786    if array.is_null(row_idx) {
787        return Err(DataFusionError::Execution(format!(
788            "column '{column_name}' cannot be NULL for kv table insert"
789        )));
790    }
791    let values = array
792        .as_any()
793        .downcast_ref::<TimestampMicrosecondArray>()
794        .ok_or_else(|| {
795            DataFusionError::Execution(format!(
796                "column '{column_name}' expected TimestampMicrosecond, got {:?}",
797                array.data_type()
798            ))
799        })?;
800    Ok(values.value(row_idx))
801}
802
803pub(crate) fn decimal128_value_at(
804    array: &ArrayRef,
805    row_idx: usize,
806    column_name: &str,
807) -> DataFusionResult<i128> {
808    if array.is_null(row_idx) {
809        return Err(DataFusionError::Execution(format!(
810            "column '{column_name}' cannot be NULL for kv table insert"
811        )));
812    }
813    let values = array
814        .as_any()
815        .downcast_ref::<Decimal128Array>()
816        .ok_or_else(|| {
817            DataFusionError::Execution(format!(
818                "column '{column_name}' expected Decimal128, got {:?}",
819                array.data_type()
820            ))
821        })?;
822    Ok(values.value(row_idx))
823}
824
825pub(crate) fn uint64_value_at(
826    array: &ArrayRef,
827    row_idx: usize,
828    column_name: &str,
829) -> DataFusionResult<u64> {
830    if array.is_null(row_idx) {
831        return Err(DataFusionError::Execution(format!(
832            "column '{column_name}' cannot be NULL for kv table insert"
833        )));
834    }
835    let values = array
836        .as_any()
837        .downcast_ref::<UInt64Array>()
838        .ok_or_else(|| {
839            DataFusionError::Execution(format!(
840                "column '{column_name}' expected UInt64, got {:?}",
841                array.data_type()
842            ))
843        })?;
844    Ok(values.value(row_idx))
845}
846
847pub(crate) fn decimal256_value_at(
848    array: &ArrayRef,
849    row_idx: usize,
850    column_name: &str,
851) -> DataFusionResult<i256> {
852    if array.is_null(row_idx) {
853        return Err(DataFusionError::Execution(format!(
854            "column '{column_name}' cannot be NULL for kv table insert"
855        )));
856    }
857    let values = array
858        .as_any()
859        .downcast_ref::<Decimal256Array>()
860        .ok_or_else(|| {
861            DataFusionError::Execution(format!(
862                "column '{column_name}' expected Decimal256, got {:?}",
863                array.data_type()
864            ))
865        })?;
866    Ok(values.value(row_idx))
867}
868
869pub(crate) fn fixed_binary_value_at(
870    array: &ArrayRef,
871    row_idx: usize,
872    column_name: &str,
873) -> DataFusionResult<Vec<u8>> {
874    if array.is_null(row_idx) {
875        return Err(DataFusionError::Execution(format!(
876            "column '{column_name}' cannot be NULL for kv table insert"
877        )));
878    }
879    let values = array
880        .as_any()
881        .downcast_ref::<FixedSizeBinaryArray>()
882        .ok_or_else(|| {
883            DataFusionError::Execution(format!(
884                "column '{column_name}' expected FixedSizeBinary, got {:?}",
885                array.data_type()
886            ))
887        })?;
888    Ok(values.value(row_idx).to_vec())
889}
890
891pub(crate) fn list_value_at(
892    array: &ArrayRef,
893    row_idx: usize,
894    column_name: &str,
895    elem: ListElementKind,
896) -> DataFusionResult<CellValue> {
897    if array.is_null(row_idx) {
898        return Err(DataFusionError::Execution(format!(
899            "column '{column_name}' cannot be NULL for kv table insert"
900        )));
901    }
902    let list_array = array.as_any().downcast_ref::<ListArray>().ok_or_else(|| {
903        DataFusionError::Execution(format!(
904            "column '{column_name}' expected List, got {:?}",
905            array.data_type()
906        ))
907    })?;
908    let child = list_array.value(row_idx);
909    let mut items = Vec::with_capacity(child.len());
910    for i in 0..child.len() {
911        let item = match elem {
912            ListElementKind::Int64 => {
913                let arr = child.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
914                    DataFusionError::Execution(format!(
915                        "column '{column_name}' list element expected Int64"
916                    ))
917                })?;
918                CellValue::Int64(arr.value(i))
919            }
920            ListElementKind::Float64 => {
921                let arr = child
922                    .as_any()
923                    .downcast_ref::<Float64Array>()
924                    .ok_or_else(|| {
925                        DataFusionError::Execution(format!(
926                            "column '{column_name}' list element expected Float64"
927                        ))
928                    })?;
929                CellValue::Float64(arr.value(i))
930            }
931            ListElementKind::Boolean => {
932                let arr = child
933                    .as_any()
934                    .downcast_ref::<BooleanArray>()
935                    .ok_or_else(|| {
936                        DataFusionError::Execution(format!(
937                            "column '{column_name}' list element expected Boolean"
938                        ))
939                    })?;
940                CellValue::Boolean(arr.value(i))
941            }
942            ListElementKind::Utf8 => {
943                let arr = child
944                    .as_any()
945                    .downcast_ref::<StringArray>()
946                    .ok_or_else(|| {
947                        DataFusionError::Execution(format!(
948                            "column '{column_name}' list element expected Utf8"
949                        ))
950                    })?;
951                CellValue::Utf8(arr.value(i).to_string())
952            }
953        };
954        items.push(item);
955    }
956    Ok(CellValue::List(items))
957}
958
959pub(crate) async fn flush_ingest_batch(
960    client: &PrefixedStoreClient,
961    keys: &mut Vec<Key>,
962    values: &mut Vec<Bytes>,
963) -> DataFusionResult<u64> {
964    if keys.is_empty() {
965        return Ok(0);
966    }
967    let mut batch = StoreWriteBatch::new();
968    for (key, value) in keys.iter().zip(values.iter()) {
969        batch
970            .push(client, key, value)
971            .map_err(|e| DataFusionError::External(Box::new(e)))?;
972    }
973    let token = batch
974        .commit(client.client())
975        .await
976        .map_err(|e| DataFusionError::External(Box::new(e)))?;
977    keys.clear();
978    values.clear();
979    Ok(token)
980}
981
982#[cfg(test)]
983mod tests {
984    use bytes::Bytes;
985    use exoware_sdk::StoreClient;
986
987    use super::*;
988
989    #[test]
990    fn store_batch_upload_stage_preserves_rows_for_failed_retry() {
991        let writer = BatchWriter::new(
992            PrefixedStoreClient::empty(StoreClient::new("http://127.0.0.1:1")),
993            &[],
994        );
995        let mut prepared = PreparedBatch {
996            request_id: 7,
997            entry_count: 2,
998            keys: vec![Bytes::from_static(b"a"), Bytes::from_static(b"b")],
999            values: vec![Bytes::from_static(&[1]), Bytes::from_static(&[2, 3])],
1000        };
1001        let mut batch = StoreWriteBatch::new();
1002
1003        StoreBatchUpload::stage_upload(&writer, &mut prepared, &mut batch).expect("stage flush");
1004
1005        assert_eq!(batch.len(), 2);
1006        assert_eq!(prepared.entry_count(), 2);
1007        assert_eq!(prepared.keys.len(), 2);
1008        assert_eq!(prepared.values.len(), 2);
1009
1010        writer.mark_flush_failed(prepared);
1011        let mut retry = writer.take_failed_prepared().expect("failed batch queued");
1012        assert_eq!(retry.entry_count(), 2);
1013
1014        let mut retry_batch = StoreWriteBatch::new();
1015        StoreBatchUpload::stage_upload(&writer, &mut retry, &mut retry_batch).expect("stage retry");
1016        assert_eq!(retry_batch.len(), 2);
1017    }
1018}
exoware_sql/writer.rs

exoware_sql/
writer.rs
