feat: serialize Iceberg file scan work units

iceberg/Cargo.toml

@@ -20,6 +20,7 @@ typetag = "0.2"
 datafusion-distributed = { path = "..", features = ["sql"] }
 delegate = "0.13"
 prost = "0.14.1"
+rmp-serde = "1"
 tokio = { version = "1.48", features = ["full"] }
 tokio-stream = "0.1"
 

iceberg/src/codec.rs

@@ -1,11 +1,7 @@
-use crate::work_unit_feed::FileScanTaskMessage;
-use bytes::{Buf, BufMut};
 use datafusion::common::Result;
 use datafusion::execution::TaskContext;
 use datafusion::physical_plan::ExecutionPlan;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
-use prost::encoding::{DecodeContext, WireType};
-use prost::{DecodeError, Message};
 use std::sync::Arc;
 
 #[derive(Debug)]
@@ -26,40 +22,3 @@ impl PhysicalExtensionCodec for IcebergCodec {
         unimplemented!()
     }
 }
-
-// TODO: Implement serde for FileScanTaskMessage
-//  This message is an individual WorkUnit, but it cannot be serialized yet. During distributed
-//  execution, this will be streamed over the wire from coordinator to workers, but for that to
-//  happen, it will need to be represented as a prost::Message.
-//
-// WARNING: for the ones who end up implementing this. I have no idea if implementing Message here
-//  is really the best option for serialization, it might not be.
-impl Message for FileScanTaskMessage {
-    fn encode_raw(&self, _buf: &mut impl BufMut)
-    where
-        Self: Sized,
-    {
-        unimplemented!()
-    }
-
-    fn merge_field(
-        &mut self,
-        _tag: u32,
-        _wire_type: WireType,
-        _buf: &mut impl Buf,
-        _ctx: DecodeContext,
-    ) -> std::result::Result<(), DecodeError>
-    where
-        Self: Sized,
-    {
-        unimplemented!()
-    }
-
-    fn encoded_len(&self) -> usize {
-        unimplemented!()
-    }
-
-    fn clear(&mut self) {
-        unimplemented!()
-    }
-}

iceberg/src/data_source.rs

@@ -2,7 +2,7 @@ use std::sync::Arc;
 
 use datafusion::arrow::datatypes::SchemaRef;
 use datafusion::common::stats::Precision;
-use datafusion::common::{ColumnStatistics, Statistics, exec_datafusion_err};
+use datafusion::common::{ColumnStatistics, Statistics};
 use datafusion::config::ConfigOptions;
 use datafusion::datasource::source::DataSource;
 use datafusion::error::Result;
@@ -22,6 +22,7 @@ use iceberg::arrow::ArrowReaderBuilder;
 use iceberg::spec::SnapshotRef;
 
 use crate::common::{convert_filters_to_predicate, df_err, iceberg_err};
+use crate::work_unit_wire::FileScanTaskDecoder;
 use crate::{IcebergConfig, IcebergWorkUnitFeed};
 
 /// Snapshot summary keys defined by the Iceberg table spec:
@@ -209,14 +210,12 @@ impl DataSource for IcebergDataSource {
                 .with_row_selection_enabled(config.row_selection_enabled)
                 .build();
 
+        let mut decoder = FileScanTaskDecoder::default();
         let feed = self
             .feed
             .feed(partition, context)?
-            .map(|msg_or_err| match msg_or_err {
-                Ok(msg) => match msg.inner {
-                    Some(msg) => Ok(msg),
-                    None => Err(iceberg_err(exec_datafusion_err!("Missing inner"))),
-                },
+            .map(move |msg_or_err| match msg_or_err {
+                Ok(msg) => decoder.decode(msg).map_err(iceberg_err),
                 Err(err) => Err(iceberg_err(err)),
             })
             .boxed();

iceberg/src/lib.rs

@@ -13,6 +13,7 @@ mod distributed_desired_task_count_handler;
 mod iceberg_ext;
 mod table_provider;
 mod work_unit_feed;
+mod work_unit_wire;
 
 mod codec;
 #[doc(hidden)]

iceberg/src/test_utils/harness.rs

@@ -62,6 +62,11 @@ impl IcebergTestHarness {
     pub async fn physical_plan(&self, sql: &str) -> Result<Arc<dyn ExecutionPlan>> {
         self.ctx.sql(sql).await?.create_physical_plan().await
     }
+
+    #[cfg(test)]
+    pub(crate) fn task_context(&self) -> Arc<datafusion::execution::TaskContext> {
+        self.ctx.task_ctx()
+    }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

iceberg/src/work_unit_feed.rs

@@ -6,14 +6,14 @@ use datafusion::error::DataFusionError;
 use datafusion::execution::TaskContext;
 use datafusion::physical_expr::Partitioning;
 use datafusion_distributed::{DistributedWorkUnitFeedContext, WorkUnitFeedProvider};
+use futures::StreamExt;
 use futures::stream::BoxStream;
-use futures::{StreamExt, TryStreamExt};
 use iceberg::expr::Predicate;
-use iceberg::scan::FileScanTask;
 use tokio::sync::mpsc::UnboundedReceiver;
 use tokio_stream::wrappers::UnboundedReceiverStream;
 
 use crate::common::df_err;
+use crate::work_unit_wire::{FileScanTaskEncoder, FileScanTaskMessage};
 
 /// Work unit feed implementation that yields [FileScanTask] messages at execution time.
 ///
@@ -116,17 +116,6 @@ pub(crate) struct SyncManager {
     feeds: TakeableVec<UnboundedReceiver<Result<FileScanTaskMessage>>>,
 }
 
-#[derive(Debug, Clone, Default)]
-pub struct FileScanTaskMessage {
-    pub(crate) inner: Option<FileScanTask>,
-}
-
-impl FileScanTaskMessage {
-    fn new(inner: FileScanTask) -> Self {
-        Self { inner: Some(inner) }
-    }
-}
-
 impl WorkUnitFeedProvider for IcebergWorkUnitFeed {
     type WorkUnit = FileScanTaskMessage;
 
@@ -173,7 +162,7 @@ impl WorkUnitFeedProvider for IcebergWorkUnitFeed {
             // the return streams of the `feed()` method.
             let task = SpawnedTask::spawn(async move {
                 let mut stream = match table_scan.plan_files().await {
-                    Ok(stream) => stream.map_ok(FileScanTaskMessage::new),
+                    Ok(stream) => stream,
                     Err(err) => {
                         let _ = txs[0].send(Err(df_err(err)));
                         return;
@@ -183,9 +172,16 @@ impl WorkUnitFeedProvider for IcebergWorkUnitFeed {
                 // Round robing across output partitions.
                 // TODO: this is fine for Partitioning::UnknownPartitioning, but any other
                 //  partitioning will require smarter routing across output channels.
+                let mut encoders = (0..txs.len())
+                    .map(|_| FileScanTaskEncoder::default())
+                    .collect::<Vec<_>>();
                 let mut i = 0;
                 while let Some(scan_task_or_err) = stream.next().await {
-                    let _ = txs[i % txs.len()].send(scan_task_or_err.map_err(df_err));
+                    let partition = i % txs.len();
+                    let message = scan_task_or_err
+                        .map_err(df_err)
+                        .and_then(|task| encoders[partition].encode(task));
+                    let _ = txs[partition].send(message);
                     i += 1;
                 }
             });