feat: Datafusion file sink

vortex-datafusion/src/persistent/format.rs

@@ -12,8 +12,10 @@ use datafusion_common::stats::Precision;
 use datafusion_common::{
     not_impl_err, ColumnStatistics, DataFusionError, Result as DFResult, ScalarValue, Statistics,
 };
+use datafusion_expr::dml::InsertOp;
 use datafusion_expr::Expr;
 use datafusion_physical_expr::{LexRequirement, PhysicalExpr};
+use datafusion_physical_plan::insert::DataSinkExec;
 use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet;
 use datafusion_physical_plan::ExecutionPlan;
 use futures::{stream, StreamExt as _, TryStreamExt as _};
@@ -29,6 +31,7 @@ use vortex_scalar::Scalar;
 
 use super::cache::FileLayoutCache;
 use super::execution::VortexExec;
+use super::sink::VortexSink;
 use crate::can_be_pushed_down;
 
 /// Vortex implementation of a DataFusion [`FileFormat`].
@@ -237,12 +240,24 @@ impl FileFormat for VortexFormat {
 
     async fn create_writer_physical_plan(
         &self,
-        _input: Arc<dyn ExecutionPlan>,
+        input: Arc<dyn ExecutionPlan>,
         _state: &SessionState,
-        _conf: FileSinkConfig,
-        _order_requirements: Option<LexRequirement>,
+        conf: FileSinkConfig,
+        order_requirements: Option<LexRequirement>,
     ) -> DFResult<Arc<dyn ExecutionPlan>> {
-        not_impl_err!("Writer not implemented for this format")
+        if conf.insert_op != InsertOp::Append {
+            return not_impl_err!("Overwrites are not implemented yet for Parquet");
+        }
+
+        let sink = Arc::new(VortexSink::new(conf));
+        let sink_schema = input.schema();
+
+        Ok(Arc::new(DataSinkExec::new(
+            input,
+            sink,
+            sink_schema,
+            order_requirements,
+        )) as _)
     }
 
     fn supports_filters_pushdown(

vortex-datafusion/src/persistent/mod.rs

@@ -4,5 +4,6 @@ mod config;
 mod execution;
 mod format;
 mod opener;
+mod sink;
 
 pub use format::{VortexFormat, VortexFormatOptions};

vortex-datafusion/src/persistent/sink.rs

@@ -0,0 +1,60 @@
+use std::any::Any;
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use datafusion::datasource::physical_plan::FileSinkConfig;
+use datafusion_execution::{SendableRecordBatchStream, TaskContext};
+use datafusion_physical_plan::insert::DataSink;
+use datafusion_physical_plan::metrics::MetricsSet;
+use datafusion_physical_plan::{DisplayAs, DisplayFormatType};
+
+pub struct VortexSink {
+    config: FileSinkConfig,
+}
+
+impl VortexSink {
+    pub fn new(config: FileSinkConfig) -> Self {
+        Self { config }
+    }
+}
+
+impl std::fmt::Debug for VortexSink {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("VortexSink").finish()
+    }
+}
+
+impl DisplayAs for VortexSink {
+    fn fmt_as(&self, t: DisplayFormatType, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match t {
+            DisplayFormatType::Default | DisplayFormatType::Verbose => {
+                write!(f, "VortexSink")
+            }
+        }
+    }
+}
+
+#[async_trait]
+impl DataSink for VortexSink {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn metrics(&self) -> Option<MetricsSet> {
+        None
+    }
+
+    async fn write_all(
+        &self,
+        mut data: SendableRecordBatchStream,
+        context: &Arc<TaskContext>,
+    ) -> datafusion_common::error::Result<u64> {
+        let object_store = context
+            .runtime_env()
+            .object_store(&self.config.object_store_url)?;
+
+        let base_output_path = &self.config.table_paths[0];
+
+        todo!()
+    }
+}