feat: Add write support with streaming API for DuckLake catalogs (#46)

Author: shefeek-jinnah

Cargo.toml

@@ -32,6 +32,9 @@ thiserror = "2.0"
 base64 = { version = "0.22", optional = true }
 hex = { version = "0.4", optional = true }
 
+# Write support (optional)
+uuid = { version = "1.0", features = ["v4"], optional = true }
+
 # Metadata providers (optional)
 duckdb = { version = "1.4.1", features = ["bundled"], optional = true }
 sqlx = { version = "0.8", features = ["runtime-tokio"], optional = true }
@@ -58,4 +61,8 @@ metadata-mysql = ["dep:sqlx", "sqlx/mysql", "sqlx/chrono"]
 metadata-sqlite = ["dep:sqlx", "sqlx/sqlite", "sqlx/chrono"]
 
 # Encryption support for Parquet files
-encryption = ["parquet/encryption", "datafusion/parquet_encryption", "dep:base64", "dep:hex"]
+encryption = ["parquet/encryption", "datafusion/parquet_encryption", "dep:base64", "dep:hex"]
+
+# Write support
+write = ["dep:uuid"]
+write-sqlite = ["write", "metadata-sqlite"]

src/catalog.rs

@@ -6,7 +6,7 @@ use std::sync::Arc;
 use crate::Result;
 use crate::information_schema::InformationSchemaProvider;
 use crate::metadata_provider::MetadataProvider;
-use crate::path_resolver::parse_object_store_url;
+use crate::path_resolver::{parse_object_store_url, resolve_path};
 use crate::schema::DuckLakeSchema;
 use datafusion::catalog::{CatalogProvider, SchemaProvider};
 use datafusion::datasource::object_store::ObjectStoreUrl;
@@ -116,14 +116,9 @@ impl CatalogProvider for DuckLakeCatalog {
         // Query database with the pinned snapshot_id for data schemas
         match self.provider.get_schema_by_name(name, self.snapshot_id) {
             Ok(Some(meta)) => {
-                // Resolve schema path hierarchically
-                let schema_path = if meta.path_is_relative {
-                    // Schema path is relative to catalog path
-                    format!("{}{}", self.catalog_path, meta.path)
-                } else {
-                    // Schema path is absolute
-                    meta.path
-                };
+                // Resolve schema path hierarchically using path_resolver utility
+                let schema_path =
+                    resolve_path(&self.catalog_path, &meta.path, meta.path_is_relative);
 
                 // Pass the pinned snapshot_id to schema
                 Some(Arc::new(DuckLakeSchema::new(

src/encryption.rs

@@ -126,7 +126,7 @@ impl DuckLakeEncryptionFactory {
     /// 2. Hex (if decodes to valid AES length)
     /// 3. Raw bytes (if exactly 16, 24, or 32 chars)
     #[cfg(feature = "encryption")]
-    fn decode_key(key: &str) -> Result<Vec<u8>> {
+    pub fn decode_key(key: &str) -> Result<Vec<u8>> {
         use base64::Engine;
         use datafusion::error::DataFusionError;
 

src/error.rs

@@ -59,6 +59,10 @@ pub enum DuckLakeError {
     #[error("IO error: {0}")]
     Io(#[from] std::io::Error),
 
+    /// Parquet error
+    #[error("Parquet error: {0}")]
+    Parquet(#[from] parquet::errors::ParquetError),
+
     /// Generic error
     #[error("Internal error: {0}")]
     Internal(String),

src/lib.rs

@@ -60,6 +60,14 @@ pub mod metadata_provider_postgres;
 #[cfg(feature = "metadata-sqlite")]
 pub mod metadata_provider_sqlite;
 
+// Write support (feature-gated)
+#[cfg(feature = "write")]
+pub mod metadata_writer;
+#[cfg(feature = "write-sqlite")]
+pub mod metadata_writer_sqlite;
+#[cfg(feature = "write")]
+pub mod table_writer;
+
 // Result type for DuckLake operations
 pub type Result<T> = std::result::Result<T, DuckLakeError>;
 
@@ -80,3 +88,13 @@ pub use metadata_provider_mysql::MySqlMetadataProvider;
 pub use metadata_provider_postgres::PostgresMetadataProvider;
 #[cfg(feature = "metadata-sqlite")]
 pub use metadata_provider_sqlite::SqliteMetadataProvider;
+
+// Re-export write types (feature-gated)
+#[cfg(feature = "write")]
+pub use metadata_writer::{
+    ColumnDef, DataFileInfo, MetadataWriter, WriteMode, WriteResult, WriteSetupResult,
+};
+#[cfg(feature = "write-sqlite")]
+pub use metadata_writer_sqlite::SqliteMetadataWriter;
+#[cfg(feature = "write")]
+pub use table_writer::{DuckLakeTableWriter, TableWriteSession};

src/metadata_provider.rs

@@ -19,7 +19,7 @@ pub const SQL_LIST_TABLES: &str =
 
 pub const SQL_GET_TABLE_COLUMNS: &str = "SELECT column_id, column_name, column_type, nulls_allowed
      FROM ducklake_column
-     WHERE table_id = ?
+     WHERE table_id = ? AND end_snapshot IS NULL
      ORDER BY column_order";
 
 pub const SQL_GET_DATA_FILES: &str = "

src/metadata_provider_sqlite.rs

@@ -142,7 +142,7 @@ impl MetadataProvider for SqliteMetadataProvider {
             let rows = sqlx::query(
                 "SELECT column_id, column_name, column_type, nulls_allowed
                  FROM ducklake_column
-                 WHERE table_id = ?
+                 WHERE table_id = ? AND end_snapshot IS NULL
                  ORDER BY column_order",
             )
             .bind(table_id)

src/metadata_writer.rs

@@ -0,0 +1,236 @@
+//! Metadata writer trait and common types for DuckLake catalog writes.
+//!
+//! This module provides the `MetadataWriter` trait for writing metadata to DuckLake catalogs,
+//! along with helper types for column definitions and data file registration.
+
+use crate::Result;
+
+/// Write mode for table operations.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum WriteMode {
+    /// Drop existing data and replace with new data
+    Replace,
+    /// Keep existing data and append new records
+    Append,
+}
+use crate::types::arrow_to_ducklake_type;
+use arrow::datatypes::DataType;
+
+/// Column definition for creating or updating a table's schema.
+///
+/// Unlike `DuckLakeTableColumn` (used for reading), this struct doesn't have a `column_id`
+/// field since IDs are assigned by the catalog during write operations.
+#[derive(Debug, Clone)]
+pub struct ColumnDef {
+    /// Column name
+    pub name: String,
+    /// DuckLake type string (e.g., "varchar", "int64", "decimal(10,2)")
+    pub ducklake_type: String,
+    /// Whether this column allows NULL values
+    pub is_nullable: bool,
+}
+
+impl ColumnDef {
+    /// Create a new column definition.
+    pub fn new(
+        name: impl Into<String>,
+        ducklake_type: impl Into<String>,
+        is_nullable: bool,
+    ) -> Self {
+        Self {
+            name: name.into(),
+            ducklake_type: ducklake_type.into(),
+            is_nullable,
+        }
+    }
+
+    /// Create a column definition from an Arrow DataType.
+    ///
+    /// This is a convenience constructor that converts the Arrow type to a DuckLake type string.
+    pub fn from_arrow(
+        name: impl Into<String>,
+        data_type: &DataType,
+        is_nullable: bool,
+    ) -> Result<Self> {
+        let ducklake_type = arrow_to_ducklake_type(data_type)?;
+        Ok(Self::new(name, ducklake_type, is_nullable))
+    }
+}
+
+/// Information about a data file to register in the catalog.
+///
+/// This struct contains the metadata needed to register a Parquet file in the DuckLake catalog.
+#[derive(Debug, Clone)]
+pub struct DataFileInfo {
+    /// Path to the file (relative to table path or absolute)
+    pub path: String,
+    /// Whether the path is relative to the table's path
+    pub path_is_relative: bool,
+    /// Size of the file in bytes
+    pub file_size_bytes: i64,
+    /// Size of the Parquet footer in bytes (optimization hint for reads)
+    pub footer_size: Option<i64>,
+    /// Number of records in the file
+    pub record_count: i64,
+}
+
+impl DataFileInfo {
+    /// Create a new data file info with relative path.
+    pub fn new(path: impl Into<String>, file_size_bytes: i64, record_count: i64) -> Self {
+        Self {
+            path: path.into(),
+            path_is_relative: true,
+            file_size_bytes,
+            footer_size: None,
+            record_count,
+        }
+    }
+
+    /// Set the footer size for read optimization.
+    pub fn with_footer_size(mut self, footer_size: i64) -> Self {
+        self.footer_size = Some(footer_size);
+        self
+    }
+
+    /// Mark this file as having an absolute path.
+    pub fn with_absolute_path(mut self) -> Self {
+        self.path_is_relative = false;
+        self
+    }
+}
+
+/// Result of a write operation.
+#[derive(Debug)]
+pub struct WriteResult {
+    /// Snapshot ID of the write operation
+    pub snapshot_id: i64,
+    /// Table ID (may be newly created)
+    pub table_id: i64,
+    /// Schema ID (may be newly created)
+    pub schema_id: i64,
+    /// Number of files written
+    pub files_written: usize,
+    /// Total records written
+    pub records_written: i64,
+}
+
+/// Result of a transactional write setup operation.
+#[derive(Debug)]
+pub struct WriteSetupResult {
+    /// Snapshot ID created for this write
+    pub snapshot_id: i64,
+    /// Schema ID (may be newly created)
+    pub schema_id: i64,
+    /// Table ID (may be newly created)
+    pub table_id: i64,
+    /// Column IDs in order
+    pub column_ids: Vec<i64>,
+}
+
+/// Trait for writing metadata to DuckLake catalogs.
+///
+/// Implementations must be thread-safe (`Send + Sync`).
+pub trait MetadataWriter: Send + Sync + std::fmt::Debug {
+    /// Create a new snapshot and return its ID.
+    fn create_snapshot(&self) -> Result<i64>;
+
+    /// Get or create a schema, returning `(schema_id, was_created)`.
+    fn get_or_create_schema(
+        &self,
+        name: &str,
+        path: Option<&str>,
+        snapshot_id: i64,
+    ) -> Result<(i64, bool)>;
+
+    /// Get or create a table, returning `(table_id, was_created)`.
+    fn get_or_create_table(
+        &self,
+        schema_id: i64,
+        name: &str,
+        path: Option<&str>,
+        snapshot_id: i64,
+    ) -> Result<(i64, bool)>;
+
+    /// Set columns for a table, returning assigned column IDs.
+    /// Ends existing columns using end_snapshot pattern for time travel.
+    fn set_columns(
+        &self,
+        table_id: i64,
+        columns: &[ColumnDef],
+        snapshot_id: i64,
+    ) -> Result<Vec<i64>>;
+
+    /// Register a new data file. Returns the assigned data_file_id.
+    fn register_data_file(
+        &self,
+        table_id: i64,
+        snapshot_id: i64,
+        file: &DataFileInfo,
+    ) -> Result<i64>;
+
+    /// End all existing data files for a table. Returns count of files ended.
+    fn end_table_files(&self, table_id: i64, snapshot_id: i64) -> Result<u64>;
+
+    /// Get the data path from catalog metadata.
+    fn get_data_path(&self) -> Result<String>;
+
+    /// Set the data path in catalog metadata.
+    fn set_data_path(&self, path: &str) -> Result<()>;
+
+    /// Initialize DuckLake schema tables if they don't exist.
+    fn initialize_schema(&self) -> Result<()>;
+
+    /// Atomically set up catalog metadata for a write operation.
+    /// Creates snapshot, schema, table, columns in a single transaction.
+    /// If mode is `WriteMode::Replace`, ends existing data files.
+    fn begin_write_transaction(
+        &self,
+        schema_name: &str,
+        table_name: &str,
+        columns: &[ColumnDef],
+        mode: WriteMode,
+    ) -> Result<WriteSetupResult>;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_column_def_new() {
+        let col = ColumnDef::new("test_col", "int32", true);
+        assert_eq!(col.name, "test_col");
+        assert_eq!(col.ducklake_type, "int32");
+        assert!(col.is_nullable);
+    }
+
+    #[test]
+    fn test_column_def_from_arrow() {
+        let col = ColumnDef::from_arrow("id", &DataType::Int64, false).unwrap();
+        assert_eq!(col.name, "id");
+        assert_eq!(col.ducklake_type, "int64");
+        assert!(!col.is_nullable);
+    }
+
+    #[test]
+    fn test_data_file_info_new() {
+        let file = DataFileInfo::new("test.parquet", 1024, 100);
+        assert_eq!(file.path, "test.parquet");
+        assert!(file.path_is_relative);
+        assert_eq!(file.file_size_bytes, 1024);
+        assert_eq!(file.record_count, 100);
+        assert!(file.footer_size.is_none());
+    }
+
+    #[test]
+    fn test_data_file_info_with_footer_size() {
+        let file = DataFileInfo::new("test.parquet", 1024, 100).with_footer_size(256);
+        assert_eq!(file.footer_size, Some(256));
+    }
+
+    #[test]
+    fn test_data_file_info_with_absolute_path() {
+        let file = DataFileInfo::new("/absolute/path.parquet", 1024, 100).with_absolute_path();
+        assert!(!file.path_is_relative);
+    }
+}