Mostly finish the refactor

Signed-off-by: Abhi Agarwal <[email protected]>

Author: abhi-airspace-intelligence

Cargo.toml

@@ -55,7 +55,7 @@ k8s-openapi = { version = "0.25.0", default-features = false }
 kube = { version = "1.1.0", default-features = false }
 metrics = { version = "0.24.2", default-features = false }
 metrics-exporter-prometheus = { version = "0.17.2", default-features = false }
-parquet = { version = "56.2.0", default-features = false }
+parquet = { version = "55.0.0", default-features = false }
 pg_escape = { version = "0.1.1", default-features = false }
 pin-project-lite = { version = "0.2.16", default-features = false }
 postgres-replication = { git = "https://github.com/MaterializeInc/rust-postgres", default-features = false, rev = "c4b473b478b3adfbf8667d2fbe895d8423f1290b" }

etl-destinations/Cargo.toml

@@ -8,7 +8,6 @@ repository.workspace = true
 homepage.workspace = true
 
 [features]
-arrow = ["dep:arrow"]
 bigquery = [
     "dep:gcp-bigquery-client",
     "dep:prost",
@@ -21,15 +20,15 @@ iceberg = [
     "dep:iceberg-catalog-rest",
     "dep:parquet",
     "dep:uuid",
-    "arrow",
+    "dep:arrow",
 ]
 deltalake = [
     "dep:dashmap",
     "dep:deltalake",
     "dep:futures",
     "dep:tokio",
     "dep:tracing",
-    "arrow",
+    "dep:arrow",
 ]
 
 [dependencies]

etl-destinations/src/arrow/encoding.rs

@@ -7,15 +7,13 @@ use arrow::{
         TimestampMicrosecondBuilder,
     },
     datatypes::{
-        DataType, Date32Type, FieldRef, Float32Type, Float64Type, Int16Type, Int32Type,
-        Int64Type, Schema, Time64MicrosecondType, TimeUnit, TimestampMicrosecondType, UInt32Type,
+        DataType, Date32Type, FieldRef, Float32Type, Float64Type, Int16Type, Int32Type, Int64Type,
+        Schema, Time64MicrosecondType, TimeUnit, TimestampMicrosecondType, UInt32Type,
     },
     error::ArrowError,
 };
 use chrono::{NaiveDate, NaiveTime};
-use etl::types::{
-    ArrayCell, Cell, DATE_FORMAT, TIME_FORMAT, TIMESTAMP_FORMAT, TableRow,
-};
+use etl::types::{ArrayCell, Cell, DATE_FORMAT, TIME_FORMAT, TIMESTAMP_FORMAT, TableRow};
 
 pub const UNIX_EPOCH: NaiveDate =
     NaiveDate::from_ymd_opt(1970, 1, 1).expect("unix epoch is a valid date");

etl-destinations/src/deltalake/core.rs

@@ -20,7 +20,7 @@ use crate::deltalake::config::DeltaTableConfig;
 use crate::deltalake::events::{materialize_events, materialize_events_append_only};
 use crate::deltalake::maintenance::TableMaintenanceState;
 use crate::deltalake::operations::{append_to_table, delete_from_table, merge_to_table};
-use crate::deltalake::schema::postgres_to_delta_schema;
+use crate::deltalake::schema::{postgres_to_arrow_schema, postgres_to_delta_schema};
 
 /// Configuration for Delta Lake destination
 #[derive(Debug, Clone)]
@@ -111,10 +111,12 @@ where
             })?;
 
         let table_name = &table_schema.name.name;
-        let table_path = parse_table_uri(format!("{}/{}", self.config.base_uri, table_name))
-            .map_err(|e| {
-                etl_error!(ErrorKind::DestinationError, "Failed to parse table path", e)
-            })?;
+        let pg_table_schema = &table_schema.name.schema;
+        let table_path = parse_table_uri(format!(
+            "{}/{}/{}",
+            self.config.base_uri, pg_table_schema, table_name
+        ))
+        .map_err(|e| etl_error!(ErrorKind::DestinationError, "Failed to parse table path", e))?;
 
         let mut table_builder = DeltaTableBuilder::from_uri(table_path).map_err(|e| {
             etl_error!(
@@ -245,14 +247,16 @@ where
             .append_only;
 
         if is_append_only {
-            let rows = materialize_events_append_only(&events, &table_schema)?;
-            self.write_table_rows_internal(&table_id, rows).await?;
+            let row_refs = materialize_events_append_only(&events, &table_schema)?;
+            let rows: Vec<PgTableRow> = row_refs.into_iter().cloned().collect();
+            self.write_table_rows_internal(&table_id, &rows).await?;
         } else {
-            let (delete_predicates, rows) = materialize_events(&events, &table_schema)?;
+            let (delete_predicates, row_refs) = materialize_events(&events, &table_schema)?;
+            let rows: Vec<PgTableRow> = row_refs.into_iter().cloned().collect();
             self.execute_delete_append_transaction_expr(
                 table_id,
                 &table_schema,
-                rows,
+                &rows,
                 delete_predicates,
             )
             .await?;
@@ -266,7 +270,7 @@ where
         &self,
         table_id: TableId,
         table_schema: &PgTableSchema,
-        upsert_rows: Vec<&PgTableRow>,
+        upsert_rows: &[PgTableRow],
         delete_predicates: Vec<Expr>,
     ) -> EtlResult<()> {
         let combined_predicate = delete_predicates.into_iter().reduce(|acc, e| acc.or(e));
@@ -347,29 +351,47 @@ where
     async fn write_table_rows_internal(
         &self,
         table_id: &TableId,
-        table_rows: Vec<&PgTableRow>,
+        table_rows: &[PgTableRow],
     ) -> EtlResult<()> {
         if table_rows.is_empty() {
             return Ok(());
         }
 
         let table = self.table_handle(table_id).await?;
 
+        let table_schema = self
+            .store
+            .get_table_schema(table_id)
+            .await?
+            .ok_or_else(|| {
+                etl_error!(
+                    ErrorKind::MissingTableSchema,
+                    "Table schema not found",
+                    format!("Schema for table {} not found in store", table_id.0)
+                )
+            })?;
+
         let row_length = table_rows.len();
         trace!("Writing {} rows to Delta table", row_length);
-            
-        let config = self.config_for_table_name(&table_schema.name.name);
+
         let mut table_guard = table.lock().await;
-        let schema = table_guard.snapshot().schema();
+        let arrow_schema = postgres_to_arrow_schema(&table_schema).map_err(|e| {
+            etl_error!(
+                ErrorKind::ConversionError,
+                "Failed to convert table schema to Arrow schema",
+                e
+            )
+        })?;
 
-        let record_batch =
-            rows_to_record_batch(table_rows.iter(), table_schema.clone()).map_err(|e| {
-                etl_error!(
-                    ErrorKind::ConversionError,
-                    "Failed to encode table rows",
-                    format!("Error converting to Arrow: {}", e)
-                )
-            })?;
+        let config = self.config_for_table_name(&table_schema.name.name);
+
+        let record_batch = rows_to_record_batch(table_rows, arrow_schema.clone()).map_err(|e| {
+            etl_error!(
+                ErrorKind::ConversionError,
+                "Failed to encode table rows",
+                format!("Error converting to Arrow: {}", e)
+            )
+        })?;
         append_to_table(&mut table_guard, config.as_ref(), record_batch)
             .await
             .map_err(|e| {
@@ -450,8 +472,7 @@ where
         table_id: TableId,
         table_rows: Vec<PgTableRow>,
     ) -> EtlResult<()> {
-        self.write_table_rows_internal(&table_id, table_rows.iter().collect())
-            .await
+        self.write_table_rows_internal(&table_id, &table_rows).await
     }
 
     async fn write_events(&self, events: Vec<Event>) -> EtlResult<()> {

etl-destinations/src/deltalake/events.rs

@@ -24,7 +24,7 @@ pub fn materialize_events_append_only<'a>(
         match event {
             Event::Insert(e) => {
                 let marker = (e.commit_lsn, e.start_lsn);
-                let pk_expr = build_pk_expr(table_schema, &e.table_row);
+                let pk_expr = build_pk_expr(table_schema, &e.table_row)?;
                 let entry = crdt_by_key.entry(pk_expr).or_insert_with(|| LWWReg {
                     val: RowOp::Upsert(&e.table_row),
                     marker,
@@ -68,7 +68,7 @@ pub fn materialize_events<'a>(
         match event {
             Event::Insert(e) => {
                 let marker = (e.commit_lsn, e.start_lsn);
-                let pk_expr = build_pk_expr(table_schema, &e.table_row);
+                let pk_expr = build_pk_expr(table_schema, &e.table_row)?;
                 let entry = crdt_by_key.entry(pk_expr).or_insert_with(|| LWWReg {
                     val: RowOp::Upsert(&e.table_row),
                     marker,
@@ -77,7 +77,7 @@ pub fn materialize_events<'a>(
             }
             Event::Update(e) => {
                 let marker = (e.commit_lsn, e.start_lsn);
-                let pk_expr = build_pk_expr(table_schema, &e.table_row);
+                let pk_expr = build_pk_expr(table_schema, &e.table_row)?;
                 let entry = crdt_by_key.entry(pk_expr).or_insert_with(|| LWWReg {
                     val: RowOp::Upsert(&e.table_row),
                     marker,
@@ -87,7 +87,7 @@ pub fn materialize_events<'a>(
             Event::Delete(e) => {
                 if let Some((_, ref old_row)) = e.old_table_row {
                     let marker = (e.commit_lsn, e.start_lsn);
-                    let pk_expr = build_pk_expr(table_schema, old_row);
+                    let pk_expr = build_pk_expr(table_schema, old_row)?;
                     let entry = crdt_by_key.entry(pk_expr).or_insert_with(|| LWWReg {
                         val: RowOp::Delete,
                         marker,
@@ -127,6 +127,7 @@ mod tests {
         Cell as PgCell, ColumnSchema as PgColumnSchema, DeleteEvent, InsertEvent, TableId,
         TableName, Type as PgType, UpdateEvent,
     };
+    use insta::assert_debug_snapshot;
 
     fn schema_single_pk(table_id: TableId) -> PgTableSchema {
         PgTableSchema::new(
@@ -221,8 +222,20 @@ mod tests {
 
         let (deletes, upserts) = materialize_events(&events, &schema).unwrap();
         assert!(deletes.is_empty());
-        assert_eq!(upserts.len(), 1);
-        assert_eq!(upserts[0].values[1], PgCell::String("b".to_string()));
+        assert_debug_snapshot!(upserts, @r#"
+        [
+            TableRow {
+                values: [
+                    I64(
+                        1,
+                    ),
+                    String(
+                        "b",
+                    ),
+                ],
+            },
+        ]
+        "#);
     }
 
     #[test]
@@ -246,8 +259,26 @@ mod tests {
         let events = vec![ins, del];
 
         let (deletes, upserts) = materialize_events(&events, &schema).unwrap();
+        assert_debug_snapshot!(deletes, @r#"
+        [
+            BinaryExpr(
+                BinaryExpr {
+                    left: Column(
+                        Column {
+                            relation: None,
+                            name: "id",
+                        },
+                    ),
+                    op: Eq,
+                    right: Literal(
+                        Int64(1),
+                        None,
+                    ),
+                },
+            ),
+        ]
+        "#);
         assert!(upserts.is_empty());
-        assert_eq!(deletes.len(), 1);
     }
 
     #[test]
@@ -271,8 +302,20 @@ mod tests {
         let events = vec![ins, upd];
 
         let upserts = materialize_events_append_only(&events, &schema).unwrap();
-        assert_eq!(upserts.len(), 1);
-        assert_eq!(upserts[0].values[1], PgCell::String("a".to_string()));
+        assert_debug_snapshot!(upserts, @r#"
+        [
+            TableRow {
+                values: [
+                    I64(
+                        1,
+                    ),
+                    String(
+                        "a",
+                    ),
+                ],
+            },
+        ]
+        "#);
     }
 
     #[test]
@@ -317,8 +360,61 @@ mod tests {
 
         // We expect one delete predicate (for tenant_id=10 AND user_id=101)
         // and one upsert (tenant_id=10 AND user_id=100 with name=a2)
-        assert_eq!(deletes.len(), 1);
-        assert_eq!(upserts.len(), 1);
-        assert_eq!(upserts[0].values[2], PgCell::String("a2".to_string()));
+        assert_debug_snapshot!(deletes, @r#"
+        [
+            BinaryExpr(
+                BinaryExpr {
+                    left: BinaryExpr(
+                        BinaryExpr {
+                            left: Column(
+                                Column {
+                                    relation: None,
+                                    name: "tenant_id",
+                                },
+                            ),
+                            op: Eq,
+                            right: Literal(
+                                Int32(10),
+                                None,
+                            ),
+                        },
+                    ),
+                    op: And,
+                    right: BinaryExpr(
+                        BinaryExpr {
+                            left: Column(
+                                Column {
+                                    relation: None,
+                                    name: "user_id",
+                                },
+                            ),
+                            op: Eq,
+                            right: Literal(
+                                Int64(101),
+                                None,
+                            ),
+                        },
+                    ),
+                },
+            ),
+        ]
+        "#);
+        assert_debug_snapshot!(upserts, @r#"
+        [
+            TableRow {
+                values: [
+                    I32(
+                        10,
+                    ),
+                    I64(
+                        100,
+                    ),
+                    String(
+                        "a2",
+                    ),
+                ],
+            },
+        ]
+        "#);
     }
 }