Merge pull request JanKaul#210 from JanKaul/add-tracing

Add tracing

Author: JanKaul

Cargo.lock

@@ -1479,6 +1479,197 @@ mod tests {
         };
     }
 
+    #[tokio::test]
+    pub async fn test_datafusion_table_insert_truncate_partitioned() {
+        let object_store = ObjectStoreBuilder::memory();
+
+        let catalog: Arc<dyn Catalog> = Arc::new(
+            SqlCatalog::new("sqlite://", "test", object_store)
+                .await
+                .unwrap(),
+        );
+
+        let schema = Schema::builder()
+            .with_struct_field(StructField {
+                id: 1,
+                name: "id".to_string(),
+                required: true,
+                field_type: Type::Primitive(PrimitiveType::Long),
+                doc: None,
+            })
+            .with_struct_field(StructField {
+                id: 2,
+                name: "customer_id".to_string(),
+                required: true,
+                field_type: Type::Primitive(PrimitiveType::Long),
+                doc: None,
+            })
+            .with_struct_field(StructField {
+                id: 3,
+                name: "product_id".to_string(),
+                required: true,
+                field_type: Type::Primitive(PrimitiveType::Long),
+                doc: None,
+            })
+            .with_struct_field(StructField {
+                id: 4,
+                name: "date".to_string(),
+                required: true,
+                field_type: Type::Primitive(PrimitiveType::Date),
+                doc: None,
+            })
+            .with_struct_field(StructField {
+                id: 5,
+                name: "amount".to_string(),
+                required: true,
+                field_type: Type::Primitive(PrimitiveType::Int),
+                doc: None,
+            })
+            .build()
+            .unwrap();
+
+        let partition_spec = PartitionSpec::builder()
+            .with_partition_field(PartitionField::new(
+                2,
+                1000,
+                "customer_id_truncate",
+                Transform::Truncate(2),
+            ))
+            .build()
+            .expect("Failed to create partition spec");
+
+        let table = Table::builder()
+            .with_name("orders")
+            .with_location("/test/orders")
+            .with_schema(schema)
+            .with_partition_spec(partition_spec)
+            .build(&["test".to_owned()], catalog)
+            .await
+            .expect("Failed to create table");
+
+        let table = Arc::new(DataFusionTable::from(table));
+
+        let ctx = SessionContext::new();
+
+        ctx.register_table("orders", table.clone()).unwrap();
+
+        ctx.sql(
+            "INSERT INTO orders (id, customer_id, product_id, date, amount) VALUES
+                (1, 123, 1, '2020-01-01', 1),
+                (2, 234, 1, '2020-01-01', 1),
+                (3, 345, 1, '2020-01-01', 3),
+                (4, 123, 2, '2020-02-02', 1),
+                (5, 123, 1, '2020-02-02', 2),
+                (6, 345, 3, '2020-02-02', 3);",
+        )
+        .await
+        .expect("Failed to create query plan for insert")
+        .collect()
+        .await
+        .expect("Failed to insert values into table");
+
+        let batches = ctx
+            .sql("select product_id, sum(amount) from orders where customer_id = 123 group by product_id;")
+            .await
+            .expect("Failed to create plan for select")
+            .collect()
+            .await
+            .expect("Failed to execute select query");
+
+        for batch in batches {
+            if batch.num_rows() != 0 {
+                let (product_ids, amounts) = (
+                    batch
+                        .column(0)
+                        .as_any()
+                        .downcast_ref::<Int64Array>()
+                        .unwrap(),
+                    batch
+                        .column(1)
+                        .as_any()
+                        .downcast_ref::<Int64Array>()
+                        .unwrap(),
+                );
+                for (product_id, amount) in product_ids.iter().zip(amounts) {
+                    if product_id.unwrap() == 1 {
+                        assert_eq!(amount.unwrap(), 3)
+                    } else if product_id.unwrap() == 2 {
+                        assert_eq!(amount.unwrap(), 1)
+                    } else if product_id.unwrap() == 3 {
+                        assert_eq!(amount.unwrap(), 0)
+                    } else {
+                        panic!("Unexpected order id")
+                    }
+                }
+            }
+        }
+
+        ctx.sql(
+            "INSERT INTO orders (id, customer_id, product_id, date, amount) VALUES
+                (7, 123, 3, '2020-01-03', 1),
+                (8, 234, 1, '2020-01-03', 2),
+                (9, 234, 2, '2020-01-03', 1),
+                (10, 123, 2, '2020-01-04', 3),
+                (11, 345, 1, '2020-01-04', 2),
+                (12, 234, 3, '2020-01-04', 1),
+                (13, 123, 1, '2020-01-05', 4),
+                (14, 345, 2, '2020-01-05', 2),
+                (15, 234, 3, '2020-01-05', 3),
+                (16, 234, 3, '2020-01-05', 3),
+                (17, 123, 3, '2020-01-06', 1),
+                (18, 234, 1, '2020-01-06', 2),
+                (19, 234, 2, '2020-01-06', 1),
+                (20, 123, 2, '2020-01-07', 3),
+                (21, 345, 1, '2020-01-07', 2),
+                (22, 234, 3, '2020-01-07', 1),
+                (23, 123, 1, '2020-01-08', 4),
+                (24, 345, 2, '2020-01-08', 2),
+                (25, 234, 3, '2020-01-08', 3);",
+        )
+        .await
+        .expect("Failed to create query plan for insert")
+        .collect()
+        .await
+        .expect("Failed to insert values into table");
+
+        let batches = ctx
+            .sql("select product_id, sum(amount) from orders where customer_id = 123 group by product_id;")
+            .await
+            .expect("Failed to create plan for select")
+            .collect()
+            .await
+            .expect("Failed to execute select query");
+
+        for batch in batches {
+            if batch.num_rows() != 0 {
+                let (product_ids, amounts) = (
+                    batch
+                        .column(0)
+                        .as_any()
+                        .downcast_ref::<Int64Array>()
+                        .unwrap(),
+                    batch
+                        .column(1)
+                        .as_any()
+                        .downcast_ref::<Int64Array>()
+                        .unwrap(),
+                );
+                for (product_id, amount) in product_ids.iter().zip(amounts) {
+                    match product_id.unwrap() {
+                        1 => assert_eq!(amount.unwrap(), 11),
+                        2 => assert_eq!(amount.unwrap(), 7),
+                        3 => assert_eq!(amount.unwrap(), 2),
+                        _ => panic!("Unexpected order id"),
+                    }
+                }
+            }
+        }
+
+        if let Tabular::Table(table) = table.tabular.read().await.deref() {
+            assert_eq!(table.manifests(None, None).await.unwrap().len(), 2);
+        };
+    }
+
     #[tokio::test]
     pub async fn test_datafusion_table_branch_insert() {
         let object_store = ObjectStoreBuilder::memory();

datafusion_iceberg/src/table.rs

@@ -32,6 +32,7 @@ smallvec = { version = "1.14.0", features = ["const_generics"] }
 sqlparser = { workspace = true }
 thiserror = { workspace = true }
 thrift = { version = "0.17.0", default-features = false }
+tracing = { workspace = true }
 tokio = { version = "1.43", features = ["sync"] }
 url = { workspace = true }
 uuid = { workspace = true }

iceberg-rust/Cargo.toml

@@ -40,6 +40,7 @@ use object_store::{buffered::BufWriter, ObjectStore};
 use std::fmt::Write;
 use std::sync::Arc;
 use tokio::task::JoinSet;
+use tracing::instrument;
 
 use arrow::{datatypes::Schema as ArrowSchema, error::ArrowError, record_batch::RecordBatch};
 use futures::Stream;
@@ -65,7 +66,7 @@ use super::partition::PartitionStream;
 
 const MAX_PARQUET_SIZE: usize = 512_000_000;
 
-#[inline]
+#[instrument(skip(table, batches), fields(table_name = %table.identifier().name()))]
 /// Writes Arrow record batches as partitioned Parquet files.
 ///
 /// This function writes Arrow record batches to Parquet files, partitioning them according
@@ -94,7 +95,7 @@ pub async fn write_parquet_partitioned(
     store_parquet_partitioned(table, batches, branch, None).await
 }
 
-#[inline]
+#[instrument(skip(table, batches), fields(table_name = %table.identifier().name(), equality_ids = ?equality_ids))]
 /// Writes equality delete records as partitioned Parquet files.
 ///
 /// This function writes Arrow record batches containing equality delete records to Parquet files,
@@ -125,6 +126,7 @@ pub async fn write_equality_deletes_parquet_partitioned(
     store_parquet_partitioned(table, batches, branch, Some(equality_ids)).await
 }
 
+#[instrument(skip(table, batches), fields(table_name = %table.identifier().name(), equality_ids = ?equality_ids))]
 /// Stores Arrow record batches as partitioned Parquet files.
 ///
 /// This is an internal function that handles the core storage logic for both regular data files
@@ -268,6 +270,7 @@ async fn store_parquet_partitioned(
 type ArrowSender = Sender<(String, FileMetaData)>;
 type ArrowReciever = Receiver<(String, FileMetaData)>;
 
+#[instrument(skip(batches, object_store), fields(data_location, equality_ids = ?equality_ids))]
 /// Writes a stream of Arrow record batches to multiple Parquet files.
 ///
 /// This internal function handles the low-level details of writing record batches to Parquet files,
@@ -438,6 +441,7 @@ fn generate_partition_path(
         .collect::<Result<String, ArrowError>>()
 }
 
+#[instrument(skip(schema, object_store), fields(data_location))]
 /// Creates a new Arrow writer for writing record batches to a Parquet file.
 ///
 /// This internal function creates a new buffered writer and configures it with