@lix-js/sdk 0.6.0-preview.4 → 0.6.0

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

@@ -1,148 +0,0 @@
-use std::any::Any;
-use std::fmt::Debug;
-use std::sync::Arc;
-
-use async_trait::async_trait;
-use datafusion::arrow::array::{ArrayRef, UInt64Array};
-use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef};
-use datafusion::arrow::record_batch::RecordBatch;
-use datafusion::common::{DataFusionError, Result};
-use datafusion::execution::TaskContext;
-use datafusion::physical_expr::EquivalenceProperties;
-use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType, PlanProperties};
-use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
-use datafusion::physical_plan::{
-    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream,
-};
-use futures_util::stream;
-
-use super::runtime;
-
-#[async_trait]
-pub(crate) trait InsertSink: Debug + DisplayAs + Send + Sync {
-    async fn write_batches(
-        &self,
-        batches: Vec<RecordBatch>,
-        context: &Arc<TaskContext>,
-    ) -> Result<u64>;
-}
-
-pub(crate) struct InsertExec {
-    input: Arc<dyn ExecutionPlan>,
-    sink: Arc<dyn InsertSink>,
-    result_schema: SchemaRef,
-    properties: Arc<PlanProperties>,
-}
-
-impl InsertExec {
-    pub(crate) fn new(input: Arc<dyn ExecutionPlan>, sink: Arc<dyn InsertSink>) -> Self {
-        let result_schema = dml_count_schema();
-        let properties = PlanProperties::new(
-            EquivalenceProperties::new(Arc::clone(&result_schema)),
-            Partitioning::UnknownPartitioning(1),
-            EmissionType::Final,
-            Boundedness::Bounded,
-        );
-        Self {
-            input,
-            sink,
-            result_schema,
-            properties: Arc::new(properties),
-        }
-    }
-}
-
-impl Debug for InsertExec {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("InsertExec").finish()
-    }
-}
-
-impl DisplayAs for InsertExec {
-    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match t {
-            DisplayFormatType::Default | DisplayFormatType::Verbose => {
-                write!(f, "InsertExec: sink=")?;
-                self.sink.fmt_as(t, f)
-            }
-            DisplayFormatType::TreeRender => write!(f, "InsertExec"),
-        }
-    }
-}
-
-impl ExecutionPlan for InsertExec {
-    fn name(&self) -> &str {
-        "InsertExec"
-    }
-
-    fn as_any(&self) -> &dyn Any {
-        self
-    }
-
-    fn properties(&self) -> &Arc<PlanProperties> {
-        &self.properties
-    }
-
-    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
-        vec![&self.input]
-    }
-
-    fn with_new_children(
-        self: Arc<Self>,
-        mut children: Vec<Arc<dyn ExecutionPlan>>,
-    ) -> Result<Arc<dyn ExecutionPlan>> {
-        if children.len() != 1 {
-            return Err(DataFusionError::Execution(format!(
-                "InsertExec expects one input child, got {}",
-                children.len()
-            )));
-        }
-        Ok(Arc::new(Self::new(
-            children.swap_remove(0),
-            Arc::clone(&self.sink),
-        )))
-    }
-
-    fn execute(
-        &self,
-        partition: usize,
-        context: Arc<TaskContext>,
-    ) -> Result<SendableRecordBatchStream> {
-        if partition != 0 {
-            return Err(DataFusionError::Execution(format!(
-                "InsertExec only exposes one partition, got {partition}"
-            )));
-        }
-
-        let input = Arc::clone(&self.input);
-        let sink = Arc::clone(&self.sink);
-        let stream_schema = Arc::clone(&self.result_schema);
-        let result_schema = Arc::clone(&self.result_schema);
-        let stream = stream::once(async move {
-            let batches = runtime::collect_input_plan(input, Arc::clone(&context)).await?;
-            let count = sink.write_batches(batches, &context).await?;
-            dml_count_batch(stream_schema, count)
-        });
-
-        Ok(Box::pin(RecordBatchStreamAdapter::new(
-            result_schema,
-            stream,
-        )))
-    }
-}
-
-fn dml_count_schema() -> SchemaRef {
-    Arc::new(Schema::new(vec![Field::new(
-        "count",
-        DataType::UInt64,
-        false,
-    )]))
-}
-
-fn dml_count_batch(schema: SchemaRef, count: u64) -> Result<RecordBatch> {
-    RecordBatch::try_new(
-        schema,
-        vec![Arc::new(UInt64Array::from(vec![count])) as ArrayRef],
-    )
-    .map_err(DataFusionError::from)
-}

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

@@ -1,440 +0,0 @@
-use std::any::Any;
-use std::sync::Arc;
-
-use async_trait::async_trait;
-use datafusion::arrow::array::{ArrayRef, BooleanArray, Float64Array, Int64Array, StringArray};
-use datafusion::arrow::datatypes::SchemaRef;
-use datafusion::arrow::record_batch::{RecordBatch, RecordBatchOptions};
-use datafusion::catalog::{Session, TableProvider};
-use datafusion::common::{DataFusionError, Result};
-use datafusion::datasource::TableType;
-use datafusion::execution::TaskContext;
-use datafusion::logical_expr::{Expr, TableProviderFilterPushDown};
-use datafusion::physical_expr::EquivalenceProperties;
-use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType, PlanProperties};
-use datafusion::physical_plan::stream::RecordBatchStreamAdapter;
-use datafusion::physical_plan::{
-    DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, SendableRecordBatchStream,
-};
-use futures_util::stream;
-use serde_json::Value as JsonValue;
-use tokio::sync::Mutex;
-
-use crate::commit_graph::CommitGraphReader;
-use crate::serialize_row_metadata;
-use crate::LixError;
-
-use super::entity_provider::{
-    entity_f64_value, entity_i64_value, entity_json_text_value, entity_surface_schema,
-    parse_snapshot, string_array, EntityColumnType, EntityProviderVariant, EntitySurfaceSpec,
-};
-use super::history_projection::{tombstone_identity_column_value, HistoryIdentityProjection};
-use super::history_route::{
-    load_history_entries, parse_history_filter, HistoryColumnStyle, HistoryRoute,
-    HistoryViewDescriptor, HISTORY_COL_START_COMMIT_ID,
-};
-use super::SqlCommitStoreQuerySource;
-use crate::commit_store::MaterializedChange;
-
-/// Schema-specific history surface backed directly by the commit graph.
-///
-/// The provider does not query `lix_state_history` through SQL. It uses the same
-/// commit graph primitive as the generic history surface, then shapes canonical
-/// changes into the typed entity columns for one registered schema.
-pub(crate) struct EntityHistoryProvider {
-    spec: Arc<EntitySurfaceSpec>,
-    schema: SchemaRef,
-    commit_graph: Arc<Mutex<Box<dyn CommitGraphReader>>>,
-    query_source: SqlCommitStoreQuerySource,
-}
-
-impl std::fmt::Debug for EntityHistoryProvider {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("EntityHistoryProvider")
-            .field("schema_key", &self.spec.schema_key)
-            .finish()
-    }
-}
-
-impl EntityHistoryProvider {
-    pub(crate) fn new(
-        spec: Arc<EntitySurfaceSpec>,
-        commit_graph: Arc<Mutex<Box<dyn CommitGraphReader>>>,
-        query_source: SqlCommitStoreQuerySource,
-    ) -> Self {
-        Self {
-            schema: entity_surface_schema(&spec, EntityProviderVariant::History),
-            spec,
-            commit_graph,
-            query_source,
-        }
-    }
-}
-
-#[async_trait]
-impl TableProvider for EntityHistoryProvider {
-    fn as_any(&self) -> &dyn Any {
-        self
-    }
-
-    fn schema(&self) -> SchemaRef {
-        Arc::clone(&self.schema)
-    }
-
-    fn table_type(&self) -> TableType {
-        TableType::View
-    }
-
-    fn supports_filters_pushdown(
-        &self,
-        filters: &[&Expr],
-    ) -> Result<Vec<TableProviderFilterPushDown>> {
-        Ok(filters
-            .iter()
-            .map(|filter| {
-                if parse_history_filter(filter, HistoryColumnStyle::Prefixed).is_some() {
-                    TableProviderFilterPushDown::Exact
-                } else {
-                    TableProviderFilterPushDown::Unsupported
-                }
-            })
-            .collect())
-    }
-
-    async fn scan(
-        &self,
-        _state: &dyn Session,
-        projection: Option<&Vec<usize>>,
-        filters: &[Expr],
-        limit: Option<usize>,
-    ) -> Result<Arc<dyn ExecutionPlan>> {
-        let route = HistoryRoute::from_filters(filters, HistoryColumnStyle::Prefixed);
-        let schema = projected_schema(&self.schema, projection)?;
-        Ok(Arc::new(EntityHistoryScanExec::new(
-            Arc::clone(&self.spec),
-            Arc::clone(&self.commit_graph),
-            self.query_source.clone(),
-            schema,
-            route,
-            limit,
-        )))
-    }
-}
-
-struct EntityHistoryScanExec {
-    spec: Arc<EntitySurfaceSpec>,
-    commit_graph: Arc<Mutex<Box<dyn CommitGraphReader>>>,
-    query_source: SqlCommitStoreQuerySource,
-    schema: SchemaRef,
-    route: HistoryRoute,
-    limit: Option<usize>,
-    properties: Arc<PlanProperties>,
-}
-
-impl std::fmt::Debug for EntityHistoryScanExec {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        f.debug_struct("EntityHistoryScanExec")
-            .field("schema_key", &self.spec.schema_key)
-            .field("route", &self.route)
-            .field("limit", &self.limit)
-            .finish()
-    }
-}
-
-impl EntityHistoryScanExec {
-    fn new(
-        spec: Arc<EntitySurfaceSpec>,
-        commit_graph: Arc<Mutex<Box<dyn CommitGraphReader>>>,
-        query_source: SqlCommitStoreQuerySource,
-        schema: SchemaRef,
-        route: HistoryRoute,
-        limit: Option<usize>,
-    ) -> Self {
-        let properties = PlanProperties::new(
-            EquivalenceProperties::new(Arc::clone(&schema)),
-            Partitioning::UnknownPartitioning(1),
-            EmissionType::Incremental,
-            Boundedness::Bounded,
-        );
-        Self {
-            spec,
-            commit_graph,
-            query_source,
-            schema,
-            route,
-            limit,
-            properties: Arc::new(properties),
-        }
-    }
-}
-
-impl DisplayAs for EntityHistoryScanExec {
-    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match t {
-            DisplayFormatType::Default | DisplayFormatType::Verbose => write!(
-                f,
-                "EntityHistoryScanExec(schema_key={}, route={:?}, limit={:?})",
-                self.spec.schema_key, self.route, self.limit
-            ),
-            DisplayFormatType::TreeRender => write!(f, "EntityHistoryScanExec"),
-        }
-    }
-}
-
-impl ExecutionPlan for EntityHistoryScanExec {
-    fn name(&self) -> &str {
-        "EntityHistoryScanExec"
-    }
-
-    fn as_any(&self) -> &dyn Any {
-        self
-    }
-
-    fn properties(&self) -> &Arc<PlanProperties> {
-        &self.properties
-    }
-
-    fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
-        Vec::new()
-    }
-
-    fn with_new_children(
-        self: Arc<Self>,
-        children: Vec<Arc<dyn ExecutionPlan>>,
-    ) -> Result<Arc<dyn ExecutionPlan>> {
-        if !children.is_empty() {
-            return Err(DataFusionError::Internal(
-                "EntityHistoryScanExec does not accept children".to_string(),
-            ));
-        }
-        Ok(self)
-    }
-
-    fn execute(
-        &self,
-        partition: usize,
-        _context: Arc<TaskContext>,
-    ) -> Result<SendableRecordBatchStream> {
-        if partition != 0 {
-            return Err(DataFusionError::Execution(format!(
-                "EntityHistoryScanExec only exposes one partition, got {partition}"
-            )));
-        }
-
-        let spec = Arc::clone(&self.spec);
-        let commit_graph = Arc::clone(&self.commit_graph);
-        let query_source = self.query_source.clone();
-        let schema = Arc::clone(&self.schema);
-        let route = self.route.clone();
-        let limit = self.limit;
-        let stream_schema = Arc::clone(&schema);
-        let fut = async move {
-            let rows = load_entity_history_rows(&spec, commit_graph, query_source, &route, limit)
-                .await
-                .map_err(lix_error_to_datafusion_error)?;
-            entity_history_record_batch(&stream_schema, &spec, &rows)
-        };
-
-        Ok(Box::pin(RecordBatchStreamAdapter::new(
-            schema,
-            stream::once(fut),
-        )))
-    }
-}
-
-#[derive(Debug, Clone)]
-struct EntityHistoryRow {
-    change: MaterializedChange,
-    observed_commit_id: String,
-    commit_created_at: String,
-    start_commit_id: String,
-    depth: u32,
-}
-
-async fn load_entity_history_rows(
-    spec: &EntitySurfaceSpec,
-    commit_graph: Arc<Mutex<Box<dyn CommitGraphReader>>>,
-    query_source: SqlCommitStoreQuerySource,
-    route: &HistoryRoute,
-    limit: Option<usize>,
-) -> Result<Vec<EntityHistoryRow>, LixError> {
-    let history_view_name = format!("{}_history", spec.schema_key);
-    let entries = load_history_entries(
-        HistoryViewDescriptor {
-            view_name: history_view_name.as_str(),
-            start_commit_column: HISTORY_COL_START_COMMIT_ID,
-        },
-        commit_graph,
-        query_source.json_reader,
-        route,
-        vec![spec.schema_key.clone()],
-    )
-    .await?;
-    let mut rows = entries
-        .into_iter()
-        .map(|entry| EntityHistoryRow {
-            change: entry.change,
-            observed_commit_id: entry.observed_commit_id,
-            commit_created_at: entry.commit_created_at,
-            start_commit_id: entry.start_commit_id,
-            depth: entry.depth,
-        })
-        .collect::<Vec<_>>();
-    if let Some(limit) = limit {
-        rows.truncate(limit);
-    }
-    Ok(rows)
-}
-
-fn entity_history_record_batch(
-    schema: &SchemaRef,
-    spec: &EntitySurfaceSpec,
-    rows: &[EntityHistoryRow],
-) -> Result<RecordBatch> {
-    let columns = schema
-        .fields()
-        .iter()
-        .map(|field| entity_history_column_array(field.name(), spec, rows))
-        .collect::<Result<Vec<_>>>()?;
-    Ok(RecordBatch::try_new_with_options(
-        Arc::clone(schema),
-        columns,
-        &RecordBatchOptions::new().with_row_count(Some(rows.len())),
-    )?)
-}
-
-fn entity_history_column_array(
-    column_name: &str,
-    spec: &EntitySurfaceSpec,
-    rows: &[EntityHistoryRow],
-) -> Result<ArrayRef> {
-    if let Some(system_column) = column_name.strip_prefix("lixcol_") {
-        return entity_history_system_column_array(system_column, rows);
-    }
-
-    let column_type = spec
-        .visible_column(column_name)
-        .ok_or_else(|| {
-            DataFusionError::Execution(format!(
-                "sql2 entity history provider '{}' does not expose column '{}'",
-                spec.schema_key, column_name
-            ))
-        })?
-        .column_type;
-    let projected_values = rows
-        .iter()
-        .map(|row| entity_history_column_value(row, spec, column_name))
-        .collect::<Result<Vec<_>>>()?;
-
-    Ok(match column_type {
-        EntityColumnType::String | EntityColumnType::Json => Arc::new(StringArray::from(
-            projected_values
-                .iter()
-                .map(|snapshot| entity_json_text_value(snapshot.as_ref(), column_type))
-                .collect::<Result<Vec<_>>>()?,
-        )) as ArrayRef,
-        EntityColumnType::Integer => Arc::new(Int64Array::from(
-            projected_values
-                .iter()
-                .map(|snapshot| entity_i64_value(snapshot.as_ref()))
-                .collect::<Vec<_>>(),
-        )) as ArrayRef,
-        EntityColumnType::Number => Arc::new(Float64Array::from(
-            projected_values
-                .iter()
-                .map(|snapshot| entity_f64_value(snapshot.as_ref()))
-                .collect::<Vec<_>>(),
-        )) as ArrayRef,
-        EntityColumnType::Boolean => Arc::new(BooleanArray::from(
-            projected_values
-                .iter()
-                .map(|snapshot| snapshot.as_ref().and_then(JsonValue::as_bool))
-                .collect::<Vec<_>>(),
-        )) as ArrayRef,
-    })
-}
-
-fn entity_history_column_value(
-    row: &EntityHistoryRow,
-    spec: &EntitySurfaceSpec,
-    column_name: &str,
-) -> Result<Option<JsonValue>> {
-    let snapshot = parse_snapshot(row.change.snapshot_content.as_deref())?;
-    if let Some(snapshot) = snapshot {
-        return Ok(snapshot.get(column_name).cloned());
-    }
-
-    let entity_id = row.change.entity_id.as_json_array_text().map_err(|error| {
-        DataFusionError::Execution(format!(
-            "sql2 entity history provider failed to project entity id: {error}"
-        ))
-    })?;
-    tombstone_identity_column_value(
-        column_name,
-        &entity_id,
-        HistoryIdentityProjection::PrimaryKeyPaths(&spec.primary_key_paths),
-    )
-    .map_err(|error| DataFusionError::Execution(error.to_string()))
-}
-
-fn entity_history_system_column_array(
-    column_name: &str,
-    rows: &[EntityHistoryRow],
-) -> Result<ArrayRef> {
-    Ok(match column_name {
-        "entity_id" => Arc::new(StringArray::from(
-            rows.iter()
-                .map(|row| {
-                    Some(
-                        row.change
-                            .entity_id
-                            .as_json_array_text()
-                            .expect("canonical change entity identity should project"),
-                    )
-                })
-                .collect::<Vec<_>>(),
-        )) as ArrayRef,
-        "schema_key" => string_array(rows.iter().map(|row| Some(row.change.schema_key.as_str()))),
-        "file_id" => string_array(rows.iter().map(|row| row.change.file_id.as_deref())),
-        "snapshot_content" => string_array(
-            rows.iter()
-                .map(|row| row.change.snapshot_content.as_deref()),
-        ),
-        "metadata" => Arc::new(StringArray::from(
-            rows.iter()
-                .map(|row| row.change.metadata.as_ref().map(serialize_row_metadata))
-                .collect::<Vec<_>>(),
-        )) as ArrayRef,
-        "change_id" => string_array(rows.iter().map(|row| Some(row.change.id.as_str()))),
-        "observed_commit_id" => {
-            string_array(rows.iter().map(|row| Some(row.observed_commit_id.as_str())))
-        }
-        "commit_created_at" => {
-            string_array(rows.iter().map(|row| Some(row.commit_created_at.as_str())))
-        }
-        "start_commit_id" => {
-            string_array(rows.iter().map(|row| Some(row.start_commit_id.as_str())))
-        }
-        "depth" => Arc::new(Int64Array::from(
-            rows.iter()
-                .map(|row| i64::from(row.depth))
-                .collect::<Vec<_>>(),
-        )) as ArrayRef,
-        other => {
-            return Err(DataFusionError::Execution(format!(
-                "sql2 entity history provider does not support system column 'lixcol_{other}'"
-            )))
-        }
-    })
-}
-
-fn projected_schema(schema: &SchemaRef, projection: Option<&Vec<usize>>) -> Result<SchemaRef> {
-    let Some(projection) = projection else {
-        return Ok(Arc::clone(schema));
-    };
-    Ok(Arc::new(schema.project(projection)?))
-}
-
-fn lix_error_to_datafusion_error(error: LixError) -> DataFusionError {
-    super::error::lix_error_to_datafusion_error(error)
-}