exoware_sql/
writer.rs

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

exoware_sql/
writer.rs
