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

exoware_sql/
writer.rs
