@lix-js/sdk 0.6.0-preview.0 → 0.6.0-preview.2

package/dist-engine-src/src/sql2/write_normalization.rs

@@ -0,0 +1,345 @@
+use std::collections::{BTreeMap, BTreeSet};
+use std::sync::Arc;
+
+use datafusion::arrow::array::ArrayRef;
+use datafusion::arrow::datatypes::DataType;
+use datafusion::arrow::record_batch::RecordBatch;
+use datafusion::common::{DataFusionError, Result, ScalarValue};
+use datafusion::logical_expr::Expr;
+use datafusion::physical_expr::expressions::{CastExpr, Literal};
+use datafusion::physical_expr::PhysicalExpr;
+use datafusion::physical_plan::projection::ProjectionExec;
+use datafusion::physical_plan::ExecutionPlan;
+
+use crate::LixError;
+
+#[derive(Debug, Clone)]
+pub(crate) enum SqlCell {
+    Null,
+    Value(ScalarValue),
+}
+
+impl SqlCell {
+    pub(crate) fn from_scalar(value: ScalarValue) -> Self {
+        if value.is_null() {
+            Self::Null
+        } else {
+            Self::Value(value)
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub(crate) enum InsertCell {
+    Omitted,
+    Provided(SqlCell),
+}
+
+#[derive(Debug, Clone)]
+pub(crate) enum UpdateCell {
+    Unassigned,
+    Assigned(SqlCell),
+}
+
+#[derive(Debug, Clone)]
+pub(crate) struct InsertColumnIntents {
+    explicit_columns: Option<BTreeSet<String>>,
+}
+
+impl InsertColumnIntents {
+    pub(crate) fn all_explicit() -> Self {
+        Self {
+            explicit_columns: None,
+        }
+    }
+
+    pub(crate) fn from_input(input: &Arc<dyn ExecutionPlan>) -> Self {
+        let Some(projection) = input.as_any().downcast_ref::<ProjectionExec>() else {
+            return Self {
+                explicit_columns: None,
+            };
+        };
+
+        let explicit_columns = projection
+            .expr()
+            .iter()
+            .filter(|expr| !is_generated_null_default(expr.expr.as_ref()))
+            .map(|expr| expr.alias.clone())
+            .collect();
+
+        Self {
+            explicit_columns: Some(explicit_columns),
+        }
+    }
+
+    pub(crate) fn includes_column(&self, column_name: &str) -> bool {
+        self.explicit_columns
+            .as_ref()
+            .is_none_or(|columns| columns.contains(column_name))
+    }
+
+    pub(crate) fn cell(
+        &self,
+        batch: &RecordBatch,
+        row_index: usize,
+        column_name: &str,
+    ) -> Result<InsertCell> {
+        if !self.includes_column(column_name) {
+            return Ok(InsertCell::Omitted);
+        }
+
+        optional_scalar_value(batch, row_index, column_name).map(|value| match value {
+            None => InsertCell::Omitted,
+            Some(value) => InsertCell::Provided(SqlCell::from_scalar(value)),
+        })
+    }
+}
+
+pub(crate) fn reject_non_binary_casts_for_insert_column(
+    input: &Arc<dyn ExecutionPlan>,
+    column_name: &str,
+    context: &str,
+) -> Result<()> {
+    reject_non_binary_casts_for_insert_column_in_plan(input.as_ref(), column_name, context)
+}
+
+fn reject_non_binary_casts_for_insert_column_in_plan(
+    input: &dyn ExecutionPlan,
+    column_name: &str,
+    context: &str,
+) -> Result<()> {
+    let Some(projection) = input.as_any().downcast_ref::<ProjectionExec>() else {
+        for child in input.children() {
+            reject_non_binary_casts_for_insert_column_in_plan(
+                child.as_ref(),
+                column_name,
+                context,
+            )?;
+        }
+        return Ok(());
+    };
+
+    let Some(expr) = projection
+        .expr()
+        .iter()
+        .find(|expr| expr.alias == column_name)
+    else {
+        return Ok(());
+    };
+
+    if contains_non_binary_cast_to_binary(expr.expr.as_ref()) {
+        return Err(super::error::lix_error_to_datafusion_error(
+            LixError::new(
+                LixError::CODE_TYPE_MISMATCH,
+                format!("{context} expected binary column '{column_name}'"),
+            )
+            .with_hint("Use X'...' or a binary parameter for file contents."),
+        ));
+    }
+
+    Ok(())
+}
+
+fn contains_non_binary_cast_to_binary(expr: &dyn PhysicalExpr) -> bool {
+    let Some(cast) = expr.as_any().downcast_ref::<CastExpr>() else {
+        return false;
+    };
+
+    if is_binary_type(cast.cast_type()) && !physical_expr_is_binary_or_null(cast.expr().as_ref()) {
+        return true;
+    }
+
+    contains_non_binary_cast_to_binary(cast.expr().as_ref())
+}
+
+fn physical_expr_is_binary_or_null(expr: &dyn PhysicalExpr) -> bool {
+    if let Some(literal) = expr.as_any().downcast_ref::<Literal>() {
+        return scalar_is_binary_or_null(literal.value());
+    }
+
+    if let Some(cast) = expr.as_any().downcast_ref::<CastExpr>() {
+        return is_binary_type(cast.cast_type())
+            && physical_expr_is_binary_or_null(cast.expr().as_ref());
+    }
+
+    false
+}
+
+pub(crate) fn scalar_is_binary_or_null(value: &ScalarValue) -> bool {
+    value.is_null()
+        || matches!(
+            value,
+            ScalarValue::Binary(_)
+                | ScalarValue::LargeBinary(_)
+                | ScalarValue::FixedSizeBinary(_, _)
+        )
+}
+
+pub(crate) fn logical_expr_is_binary_or_null(expr: &Expr) -> bool {
+    match expr {
+        Expr::Literal(value, _) => scalar_is_binary_or_null(value),
+        Expr::Cast(cast) => {
+            is_binary_type(&cast.data_type) && logical_expr_is_binary_or_null(&cast.expr)
+        }
+        Expr::Alias(alias) => logical_expr_is_binary_or_null(&alias.expr),
+        _ => false,
+    }
+}
+
+pub(crate) fn is_binary_type(data_type: &DataType) -> bool {
+    matches!(
+        data_type,
+        DataType::Binary | DataType::LargeBinary | DataType::FixedSizeBinary(_)
+    )
+}
+
+pub(crate) fn lix_file_data_type_lix_error() -> LixError {
+    LixError::new(
+        LixError::CODE_TYPE_MISMATCH,
+        "lix_file.data expects binary data",
+    )
+    .with_hint("Use X'...' or a binary parameter for file contents.")
+}
+
+pub(crate) fn lix_file_data_type_error(
+    context: &str,
+    column_name: &str,
+    instruction: &str,
+) -> DataFusionError {
+    super::error::lix_error_to_datafusion_error(
+        LixError::new(
+            LixError::CODE_TYPE_MISMATCH,
+            format!("{context} expected binary column '{column_name}'"),
+        )
+        .with_hint(instruction),
+    )
+}
+
+pub(crate) fn lix_file_data_type_error_with_value(
+    context: &str,
+    column_name: &str,
+    value: &ScalarValue,
+    instruction: &str,
+) -> DataFusionError {
+    super::error::lix_error_to_datafusion_error(
+        LixError::new(
+            LixError::CODE_TYPE_MISMATCH,
+            format!("{context} expected binary column '{column_name}', got {value:?}"),
+        )
+        .with_hint(instruction),
+    )
+}
+
+pub(crate) struct UpdateAssignmentValues {
+    values: BTreeMap<String, ArrayRef>,
+}
+
+impl UpdateAssignmentValues {
+    pub(crate) fn evaluate(
+        batch: &RecordBatch,
+        assignments: &[(String, Arc<dyn PhysicalExpr>)],
+    ) -> Result<Self> {
+        let mut values = BTreeMap::new();
+        for (column_name, assignment) in assignments {
+            values.insert(
+                column_name.clone(),
+                assignment.evaluate(batch)?.into_array(batch.num_rows())?,
+            );
+        }
+        Ok(Self { values })
+    }
+
+    #[cfg(test)]
+    pub(crate) fn from_batch_columns(batch: &RecordBatch, columns: &[&str]) -> Self {
+        let values = columns
+            .iter()
+            .filter_map(|column_name| {
+                let column_index = batch.schema().index_of(column_name).ok()?;
+                Some((
+                    (*column_name).to_string(),
+                    Arc::clone(batch.column(column_index)),
+                ))
+            })
+            .collect();
+        Self { values }
+    }
+
+    /// Returns only the value explicitly assigned by SQL UPDATE.
+    ///
+    /// Use this for document-patch semantics where `Unassigned` must remain
+    /// distinct from `Assigned(NULL)`.
+    pub(crate) fn assigned_cell(&self, row_index: usize, column_name: &str) -> Result<UpdateCell> {
+        let Some(array) = self.values.get(column_name) else {
+            return Ok(UpdateCell::Unassigned);
+        };
+
+        ScalarValue::try_from_array(array.as_ref(), row_index)
+            .map(SqlCell::from_scalar)
+            .map(UpdateCell::Assigned)
+            .map_err(|error| {
+                DataFusionError::Execution(format!(
+                    "failed to decode SQL UPDATE assignment for column '{column_name}' at row {row_index}: {error}"
+                ))
+            })
+    }
+
+    /// Returns the assigned SQL UPDATE value, or falls back to the existing row
+    /// column value when the column was not assigned.
+    ///
+    /// Use this for scalar row-column semantics. Do not use it to reconstruct
+    /// JSON documents from projected property columns, because projection can
+    /// erase the difference between an absent property and an explicit null.
+    pub(crate) fn assigned_or_existing_cell(
+        &self,
+        batch: &RecordBatch,
+        row_index: usize,
+        column_name: &str,
+    ) -> Result<InsertCell> {
+        match self.assigned_cell(row_index, column_name)? {
+            UpdateCell::Assigned(value) => Ok(InsertCell::Provided(value)),
+            UpdateCell::Unassigned => {
+                optional_scalar_value(batch, row_index, column_name).map(|value| match value {
+                    None => InsertCell::Omitted,
+                    Some(value) => InsertCell::Provided(SqlCell::from_scalar(value)),
+                })
+            }
+        }
+    }
+}
+
+pub(crate) fn optional_scalar_value(
+    batch: &RecordBatch,
+    row_index: usize,
+    column_name: &str,
+) -> Result<Option<ScalarValue>> {
+    let schema = batch.schema();
+    let column_index = match schema.index_of(column_name) {
+        Ok(column_index) => column_index,
+        Err(_) => return Ok(None),
+    };
+    if row_index >= batch.num_rows() {
+        return Err(DataFusionError::Execution(format!(
+            "row index {row_index} out of bounds for SQL write batch with {} rows",
+            batch.num_rows()
+        )));
+    }
+    ScalarValue::try_from_array(batch.column(column_index).as_ref(), row_index)
+        .map(Some)
+        .map_err(|error| {
+            DataFusionError::Execution(format!(
+                "failed to decode SQL write column '{column_name}' at row {row_index}: {error}"
+            ))
+        })
+}
+
+fn is_generated_null_default(expr: &dyn PhysicalExpr) -> bool {
+    if let Some(literal) = expr.as_any().downcast_ref::<Literal>() {
+        return literal.value().is_null();
+    }
+
+    if let Some(cast) = expr.as_any().downcast_ref::<CastExpr>() {
+        return is_generated_null_default(cast.expr().as_ref());
+    }
+
+    false
+}

package/dist-engine-src/src/storage/context.rs

@@ -0,0 +1,356 @@
+use std::sync::Arc;
+
+use async_trait::async_trait;
+
+use crate::backend::{Backend, BackendReadTransaction, BackendWriteTransaction};
+use crate::storage::types::{KvWriteBatch, StorageWriter};
+use crate::storage::{
+    KvEntryPage, KvExistsBatch, KvGetRequest, KvKeyPage, KvScanRequest, KvValueBatch, KvValuePage,
+    KvWriteStats, StorageReadTransaction, StorageReader, StorageWriteTransaction,
+};
+use crate::LixError;
+
+#[derive(Clone)]
+pub(crate) struct StorageContext {
+    backend: Arc<dyn Backend + Send + Sync>,
+}
+
+impl StorageContext {
+    pub(crate) fn new(backend: Arc<dyn Backend + Send + Sync>) -> Self {
+        Self { backend }
+    }
+
+    pub(crate) async fn begin_read_transaction(
+        &self,
+    ) -> Result<Box<dyn StorageReadTransaction + Send + Sync + 'static>, LixError> {
+        let transaction = self.backend.begin_read_transaction().await?;
+        Ok(Box::new(StorageContextReadTransaction { transaction }))
+    }
+
+    pub(crate) async fn begin_write_transaction(
+        &self,
+    ) -> Result<Box<dyn StorageWriteTransaction + Send + Sync + 'static>, LixError> {
+        let transaction = self.backend.begin_write_transaction().await?;
+        Ok(Box::new(StorageContextWriteTransaction { transaction }))
+    }
+
+    pub(crate) async fn close(&self) -> Result<(), LixError> {
+        self.backend.close().await
+    }
+
+    pub(crate) async fn destroy(&self) -> Result<(), LixError> {
+        self.backend.destroy().await
+    }
+}
+
+#[cfg(any(test, feature = "storage-benches"))]
+#[async_trait]
+impl StorageReader for StorageContext {
+    async fn get_values(&mut self, request: KvGetRequest) -> Result<KvValueBatch, LixError> {
+        let mut transaction = self.begin_read_transaction().await?;
+        let result = transaction.get_values(request).await;
+        match result {
+            Ok(result) => {
+                transaction.rollback().await?;
+                Ok(result)
+            }
+            Err(error) => {
+                let _ = transaction.rollback().await;
+                Err(error)
+            }
+        }
+    }
+
+    async fn exists_many(&mut self, request: KvGetRequest) -> Result<KvExistsBatch, LixError> {
+        let mut transaction = self.begin_read_transaction().await?;
+        let result = transaction.exists_many(request).await;
+        match result {
+            Ok(result) => {
+                transaction.rollback().await?;
+                Ok(result)
+            }
+            Err(error) => {
+                let _ = transaction.rollback().await;
+                Err(error)
+            }
+        }
+    }
+
+    async fn scan_keys(&mut self, request: KvScanRequest) -> Result<KvKeyPage, LixError> {
+        let mut transaction = self.begin_read_transaction().await?;
+        let result = transaction.scan_keys(request).await;
+        match result {
+            Ok(result) => {
+                transaction.rollback().await?;
+                Ok(result)
+            }
+            Err(error) => {
+                let _ = transaction.rollback().await;
+                Err(error)
+            }
+        }
+    }
+
+    async fn scan_values(&mut self, request: KvScanRequest) -> Result<KvValuePage, LixError> {
+        let mut transaction = self.begin_read_transaction().await?;
+        let result = transaction.scan_values(request).await;
+        match result {
+            Ok(result) => {
+                transaction.rollback().await?;
+                Ok(result)
+            }
+            Err(error) => {
+                let _ = transaction.rollback().await;
+                Err(error)
+            }
+        }
+    }
+
+    async fn scan_entries(&mut self, request: KvScanRequest) -> Result<KvEntryPage, LixError> {
+        let mut transaction = self.begin_read_transaction().await?;
+        let result = transaction.scan_entries(request).await;
+        match result {
+            Ok(result) => {
+                transaction.rollback().await?;
+                Ok(result)
+            }
+            Err(error) => {
+                let _ = transaction.rollback().await;
+                Err(error)
+            }
+        }
+    }
+}
+
+struct StorageContextReadTransaction {
+    transaction: Box<dyn BackendReadTransaction + Send + Sync + 'static>,
+}
+
+struct StorageContextWriteTransaction {
+    transaction: Box<dyn BackendWriteTransaction + Send + Sync + 'static>,
+}
+
+#[async_trait]
+impl StorageReader for StorageContextReadTransaction {
+    async fn get_values(&mut self, request: KvGetRequest) -> Result<KvValueBatch, LixError> {
+        self.transaction
+            .get_values(request.into())
+            .await
+            .map(Into::into)
+    }
+
+    async fn exists_many(&mut self, request: KvGetRequest) -> Result<KvExistsBatch, LixError> {
+        self.transaction
+            .exists_many(request.into())
+            .await
+            .map(Into::into)
+    }
+
+    async fn scan_keys(&mut self, request: KvScanRequest) -> Result<KvKeyPage, LixError> {
+        self.transaction
+            .scan_keys(request.into())
+            .await
+            .map(Into::into)
+    }
+
+    async fn scan_values(&mut self, request: KvScanRequest) -> Result<KvValuePage, LixError> {
+        self.transaction
+            .scan_values(request.into())
+            .await
+            .map(Into::into)
+    }
+
+    async fn scan_entries(&mut self, request: KvScanRequest) -> Result<KvEntryPage, LixError> {
+        self.transaction
+            .scan_entries(request.into())
+            .await
+            .map(Into::into)
+    }
+}
+
+#[async_trait]
+impl StorageReadTransaction for StorageContextReadTransaction {
+    async fn rollback(self: Box<Self>) -> Result<(), LixError> {
+        self.transaction.rollback().await
+    }
+}
+
+#[async_trait]
+impl StorageReader for StorageContextWriteTransaction {
+    async fn get_values(&mut self, request: KvGetRequest) -> Result<KvValueBatch, LixError> {
+        self.transaction
+            .get_values(request.into())
+            .await
+            .map(Into::into)
+    }
+
+    async fn exists_many(&mut self, request: KvGetRequest) -> Result<KvExistsBatch, LixError> {
+        self.transaction
+            .exists_many(request.into())
+            .await
+            .map(Into::into)
+    }
+
+    async fn scan_keys(&mut self, request: KvScanRequest) -> Result<KvKeyPage, LixError> {
+        self.transaction
+            .scan_keys(request.into())
+            .await
+            .map(Into::into)
+    }
+
+    async fn scan_values(&mut self, request: KvScanRequest) -> Result<KvValuePage, LixError> {
+        self.transaction
+            .scan_values(request.into())
+            .await
+            .map(Into::into)
+    }
+
+    async fn scan_entries(&mut self, request: KvScanRequest) -> Result<KvEntryPage, LixError> {
+        self.transaction
+            .scan_entries(request.into())
+            .await
+            .map(Into::into)
+    }
+}
+
+#[async_trait]
+impl StorageWriter for StorageContextWriteTransaction {
+    async fn write_kv_batch(&mut self, batch: KvWriteBatch) -> Result<KvWriteStats, LixError> {
+        self.transaction
+            .write_kv_batch(batch.into())
+            .await
+            .map(Into::into)
+    }
+}
+
+#[async_trait]
+impl StorageReadTransaction for StorageContextWriteTransaction {
+    async fn rollback(self: Box<Self>) -> Result<(), LixError> {
+        self.transaction.rollback().await
+    }
+}
+
+#[async_trait]
+impl StorageWriteTransaction for StorageContextWriteTransaction {
+    async fn commit(self: Box<Self>) -> Result<(), LixError> {
+        self.transaction.commit().await
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use crate::backend::testing::UnitTestBackend;
+    use crate::storage::types::KvWriteBatch;
+    use crate::storage::{KvGetGroup, KvScanRange, StorageWriteSet};
+
+    use super::*;
+
+    #[tokio::test]
+    async fn storage_context_roundtrips_batched_writes_and_reads() {
+        let backend: Arc<dyn Backend + Send + Sync> = Arc::new(UnitTestBackend::new());
+        let storage = StorageContext::new(backend);
+        let mut tx = storage
+            .begin_write_transaction()
+            .await
+            .expect("transaction opens");
+
+        let mut batch = KvWriteBatch::new();
+        batch.put("ns", b"a".to_vec(), b"1".to_vec());
+        batch.put("ns", b"b".to_vec(), b"2".to_vec());
+        let stats = tx.write_kv_batch(batch).await.expect("batch writes");
+        assert_eq!(stats.puts, 2);
+        tx.commit().await.expect("commit succeeds");
+
+        let mut tx = storage
+            .begin_read_transaction()
+            .await
+            .expect("read transaction opens");
+        let result = tx
+            .get_values(KvGetRequest {
+                groups: vec![KvGetGroup {
+                    namespace: "ns".to_string(),
+                    keys: vec![b"a".to_vec(), b"b".to_vec()],
+                }],
+            })
+            .await
+            .expect("batch reads");
+        assert_eq!(result.groups[0].value(0), Some(Some(b"1".as_slice())));
+        assert_eq!(result.groups[0].value(1), Some(Some(b"2".as_slice())));
+
+        let exists = tx
+            .exists_many(KvGetRequest {
+                groups: vec![KvGetGroup {
+                    namespace: "ns".to_string(),
+                    keys: vec![b"a".to_vec(), b"missing".to_vec()],
+                }],
+            })
+            .await
+            .expect("existence reads");
+        assert_eq!(exists.groups[0].exists, vec![true, false]);
+
+        let result = tx
+            .scan_entries(KvScanRequest {
+                namespace: "ns".to_string(),
+                range: KvScanRange::prefix(Vec::new()),
+                after: Some(b"a".to_vec()),
+                limit: 1,
+            })
+            .await
+            .expect("scan reads");
+        assert_eq!(result.key(0).expect("key exists"), b"b");
+        assert_eq!(result.value(0).expect("value exists"), b"2");
+
+        let key_only = tx
+            .scan_keys(KvScanRequest {
+                namespace: "ns".to_string(),
+                range: KvScanRange::prefix(Vec::new()),
+                after: None,
+                limit: 2,
+            })
+            .await
+            .expect("key-only scan reads");
+        assert_eq!(key_only.keys.iter().collect::<Vec<_>>(), vec![b"a", b"b"]);
+        tx.rollback().await.expect("rollback succeeds");
+    }
+
+    #[tokio::test]
+    async fn storage_write_set_applies_as_one_batch() {
+        let backend: Arc<dyn Backend + Send + Sync> = Arc::new(UnitTestBackend::new());
+        let storage = StorageContext::new(backend);
+        let mut tx = storage
+            .begin_write_transaction()
+            .await
+            .expect("transaction opens");
+
+        let mut writes = StorageWriteSet::new();
+        assert!(writes.is_empty());
+        writes.put("ns", b"a".to_vec(), b"1".to_vec());
+        writes.put("ns", b"b".to_vec(), b"2".to_vec());
+        writes.delete("ns", b"missing".to_vec());
+        assert!(!writes.is_empty());
+
+        let stats = writes.apply(tx.as_mut()).await.expect("write set applies");
+        assert_eq!(stats.puts, 2);
+        assert_eq!(stats.deletes, 1);
+        tx.commit().await.expect("commit succeeds");
+
+        let mut tx = storage
+            .begin_read_transaction()
+            .await
+            .expect("read transaction opens");
+        let result = tx
+            .get_values(KvGetRequest {
+                groups: vec![KvGetGroup {
+                    namespace: "ns".to_string(),
+                    keys: vec![b"a".to_vec(), b"b".to_vec()],
+                }],
+            })
+            .await
+            .expect("batch reads");
+        assert_eq!(result.groups[0].value(0), Some(Some(&b"1"[..])));
+        assert_eq!(result.groups[0].value(1), Some(Some(&b"2"[..])));
+        tx.rollback().await.expect("rollback succeeds");
+    }
+}

package/dist-engine-src/src/storage/mod.rs

@@ -0,0 +1,14 @@
+mod context;
+mod read_scope;
+mod types;
+
+pub(crate) use context::StorageContext;
+pub(crate) use read_scope::{ScopedStorageReader, StorageReadScope};
+pub(crate) use types::{
+    KvEntryPage, KvExistsBatch, KvExistsGroup, KvGetGroup, KvGetRequest, KvKeyPage, KvScanRange,
+    KvScanRequest, KvValueBatch, KvValueGroup, KvValuePage, KvWriteStats, StorageReadTransaction,
+    StorageReader, StorageWriteSet, StorageWriteTransaction,
+};
+
+#[cfg(feature = "storage-benches")]
+pub(crate) use types::{KvWriteBatch, KvWriteGroup};