Add metrics broadcasting to all write functions

Eventual-Inc · Feb 5, 2025 · 0517f7c · 0517f7c
1 parent 568f5ce
commit 0517f7c
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/daft/dataframe/dataframe.py b/daft/dataframe/dataframe.py
@@ -629,6 +629,7 @@ def _from_schema(cls, schema: Schema) -> "DataFrame":
     # Write methods
     ###
 
+    @BroadcastMetrics
     @DataframePublicAPI
     def write_parquet(
         self,
@@ -708,6 +709,7 @@ def write_parquet(
                 }
             )
 
+    @BroadcastMetrics
     @DataframePublicAPI
     def write_csv(
         self,
@@ -779,6 +781,7 @@ def write_csv(
                 }
             )
 
+    @BroadcastMetrics
     @DataframePublicAPI
     def write_iceberg(
         self, table: "pyiceberg.table.Table", mode: str = "append", io_config: Optional[IOConfig] = None
@@ -929,6 +932,7 @@ def write_iceberg(
         # This is due to the fact that the logical plan of the write_iceberg returns datafiles but we want to return the above data
         return from_pydict(with_operations)
 
+    @BroadcastMetrics
     @DataframePublicAPI
     def write_deltalake(
         self,
@@ -1140,6 +1144,7 @@ def write_deltalake(
 
         return with_operations
 
+    @BroadcastMetrics
     @DataframePublicAPI
     def write_lance(
         self,