WIP: Add support for remote Parquet writer with openDAL

Author: comphead

native/core/src/execution/operators/parquet_writer.rs

@@ -22,9 +22,13 @@ use std::{
     fmt,
     fmt::{Debug, Formatter},
     fs::File,
+    io::Cursor,
     sync::Arc,
 };
 
+use bytes::Bytes;
+use url::Url;
+
 use arrow::datatypes::{Schema, SchemaRef};
 use arrow::record_batch::RecordBatch;
 use async_trait::async_trait;
@@ -48,6 +52,71 @@ use parquet::{
 };
 
 use crate::execution::shuffle::CompressionCodec;
+use crate::parquet::parquet_support::write_to_hdfs_with_opendal_async;
+
+/// Enum representing different types of Arrow writers based on storage backend
+enum ParquetWriter {
+    /// Writer for local file system
+    LocalFile(ArrowWriter<File>),
+    /// Writer for HDFS or other remote storage (writes to in-memory buffer)
+    /// Contains the writer and the destination HDFS path
+    Remote(ArrowWriter<Cursor<Vec<u8>>>, String),
+}
+
+impl ParquetWriter {
+    /// Write a RecordBatch to the underlying writer
+    async fn write(&mut self, batch: &RecordBatch) -> std::result::Result<(), parquet::errors::ParquetError> {
+        match self {
+            ParquetWriter::LocalFile(writer) => writer.write(batch),
+            ParquetWriter::Remote(writer, output_path) => {
+                // Write batch to in-memory buffer
+                writer.write(batch)?;
+
+                // Flush and get the current buffer content
+                writer.flush()?;
+                let cursor = writer.inner_mut();
+                let buffer = cursor.get_ref().clone();
+
+                // Upload
+                let url = Url::parse(output_path).map_err(|e| {
+                    parquet::errors::ParquetError::General(format!(
+                        "Failed to parse URL '{}': {}",
+                        output_path, e
+                    ))
+                })?;
+                
+                write_to_hdfs_with_opendal_async(&url, Bytes::from(buffer))
+                    .await
+                    .map_err(|e| {
+                        parquet::errors::ParquetError::General(format!(
+                            "Failed to upload to '{}': {}",
+                            output_path, e
+                        ))
+                    })?;
+
+                // Clear the buffer after upload
+                cursor.get_mut().clear();
+                cursor.set_position(0);
+
+                Ok(())
+            },
+        }
+    }
+
+    /// Close the writer and return the buffer for remote writers
+    fn close(self) -> std::result::Result<Option<(Vec<u8>, String)>, parquet::errors::ParquetError> {
+        match self {
+            ParquetWriter::LocalFile(writer) => {
+                writer.close()?;
+                Ok(None)
+            }
+            ParquetWriter::Remote(writer, path) => {
+                let cursor = writer.into_inner()?;
+                Ok(Some((cursor.into_inner(), path)))
+            }
+        }
+    }
+}
 
 /// Parquet writer operator that writes input batches to a Parquet file
 #[derive(Debug)]
@@ -119,6 +188,59 @@ impl ParquetWriterExec {
             CompressionCodec::Snappy => Ok(Compression::SNAPPY),
         }
     }
+
+    /// Create an Arrow writer based on the storage scheme
+    ///
+    /// # Arguments
+    /// * `storage_scheme` - The storage backend ("hdfs", "s3", or "local")
+    /// * `output_file_path` - The full path to the output file
+    /// * `schema` - The Arrow schema for the Parquet file
+    /// * `props` - Writer properties including compression
+    ///
+    /// # Returns
+    /// * `Ok(ParquetWriter)` - A writer appropriate for the storage scheme
+    /// * `Err(DataFusionError)` - If writer creation fails
+    fn create_arrow_writer(
+        storage_scheme: &str,
+        output_file_path: &str,
+        schema: SchemaRef,
+        props: WriterProperties,
+    ) -> Result<ParquetWriter> {
+        match storage_scheme {
+            "hdfs" | "s3" => {
+                // For remote storage (HDFS, S3), write to an in-memory buffer
+                let buffer = Vec::new();
+                let cursor = Cursor::new(buffer);
+                let writer = ArrowWriter::try_new(cursor, schema, Some(props))
+                    .map_err(|e| DataFusionError::Execution(format!(
+                        "Failed to create {} writer: {}", storage_scheme, e
+                    )))?;
+                Ok(ParquetWriter::Remote(writer, output_file_path.to_string()))
+            }
+            "local" => {
+                // For local file system, write directly to file
+                // Strip file:// or file: prefix if present
+                let local_path = output_file_path
+                    .strip_prefix("file://")
+                    .or_else(|| output_file_path.strip_prefix("file:"))
+                    .unwrap_or(output_file_path);
+
+                let file = File::create(local_path).map_err(|e| {
+                    DataFusionError::Execution(format!(
+                        "Failed to create output file '{}': {}",
+                        local_path, e
+                    ))
+                })?;
+
+                let writer = ArrowWriter::try_new(file, schema, Some(props))
+                    .map_err(|e| DataFusionError::Execution(format!("Failed to create local file writer: {}", e)))?;
+                Ok(ParquetWriter::LocalFile(writer))
+            }
+            _ => Err(DataFusionError::Execution(format!(
+                "Unsupported storage scheme: {}", storage_scheme
+            ))),
+        }
+    }
 }
 
 impl DisplayAs for ParquetWriterExec {
@@ -217,6 +339,15 @@ impl ExecutionPlan for ParquetWriterExec {
             .collect();
         let output_schema = Arc::new(arrow::datatypes::Schema::new(fields));
 
+        // Determine storage scheme from work_dir
+        let storage_scheme = if work_dir.starts_with("hdfs://") {
+            "hdfs"
+        } else if work_dir.starts_with("s3://") || work_dir.starts_with("s3a://") {
+            "s3"
+        } else {
+            "local"
+        };
+
         // Strip file:// or file: prefix if present
         let local_path = work_dir
             .strip_prefix("file://")
@@ -243,21 +374,12 @@ impl ExecutionPlan for ParquetWriterExec {
             format!("{}/part-{:05}.parquet", local_path, self.partition_id)
         };
 
-        // Create the Parquet file
-        let file = File::create(&part_file).map_err(|e| {
-            DataFusionError::Execution(format!(
-                "Failed to create output file '{}': {}",
-                part_file, e
-            ))
-        })?;
-
         // Configure writer properties
         let props = WriterProperties::builder()
             .set_compression(compression)
             .build();
 
-        let mut writer = ArrowWriter::try_new(file, Arc::clone(&output_schema), Some(props))
-            .map_err(|e| DataFusionError::Execution(format!("Failed to create writer: {}", e)))?;
+        let mut writer = Self::create_arrow_writer(storage_scheme, &part_file, Arc::clone(&output_schema), props)?;
 
         // Clone schema for use in async closure
         let schema_for_write = Arc::clone(&output_schema);
@@ -286,7 +408,7 @@ impl ExecutionPlan for ParquetWriterExec {
                     batch
                 };
 
-                writer.write(&renamed_batch).map_err(|e| {
+                writer.write(&renamed_batch).await.map_err(|e| {
                     DataFusionError::Execution(format!("Failed to write batch: {}", e))
                 })?;
             }

native/core/src/parquet/parquet_support.rs

@@ -369,9 +369,11 @@ fn is_hdfs_scheme(url: &Url, object_store_configs: &HashMap<String, String>) ->
     }
 }
 
-// Mirrors object_store::parse::parse_url for the hdfs object store
+// Creates an HDFS object store from a URL using the native HDFS implementation
 #[cfg(feature = "hdfs")]
-fn parse_hdfs_url(url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> {
+fn create_hdfs_object_store(
+    url: &Url,
+) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> {
     match datafusion_comet_objectstore_hdfs::object_store::hdfs::HadoopFileSystem::new(url.as_ref())
     {
         Some(object_store) => {
@@ -385,8 +387,11 @@ fn parse_hdfs_url(url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_stor
     }
 }
 
+// Creates an HDFS object store from a URL using OpenDAL
 #[cfg(feature = "hdfs-opendal")]
-fn parse_hdfs_url(url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> {
+fn create_hdfs_object_store(
+    url: &Url,
+) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> {
     let name_node = get_name_node_uri(url)?;
     let builder = opendal::services::Hdfs::default().name_node(&name_node);
 
@@ -401,6 +406,51 @@ fn parse_hdfs_url(url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_stor
     Ok((Box::new(store), path))
 }
 
+/// Writes data to HDFS using OpenDAL via ObjectStore trait (asynchronous version)
+///
+/// # Arguments
+/// * `url` - The HDFS URL (e.g., hdfs://namenode:port/path/to/file)
+/// * `data` - The bytes to write to the file
+///
+/// # Returns
+/// * `Ok(())` on success
+/// * `Err(object_store::Error)` on failure
+///
+/// # Example
+/// ```ignore
+/// use url::Url;
+/// use bytes::Bytes;
+///
+/// let url = Url::parse("hdfs://namenode:9000/path/to/file.parquet")?;
+/// let data = Bytes::from("file contents");
+/// write_to_hdfs_with_opendal_async(&url, data).await?;
+/// ```
+#[cfg(feature = "hdfs-opendal")]
+pub async fn write_to_hdfs_with_opendal_async(
+    url: &Url,
+    data: bytes::Bytes,
+) -> Result<(), object_store::Error> {
+    // Create the HDFS object store using OpenDAL
+    let (object_store, path) = create_hdfs_object_store(url)?;
+
+    // Use the ObjectStore trait's put method to write the data
+    object_store.put(&path, data.into()).await?;
+
+    Ok(())
+}
+
+/// Stub implementation when hdfs-opendal feature is not enabled
+#[cfg(not(feature = "hdfs-opendal"))]
+pub async fn write_to_hdfs_with_opendal_async(
+    _url: &Url,
+    _data: bytes::Bytes,
+) -> Result<(), object_store::Error> {
+    Err(object_store::Error::Generic {
+        store: "hdfs-opendal",
+        source: "HDFS OpenDAL support is not enabled in this build".into(),
+    })
+}
+
 #[cfg(feature = "hdfs-opendal")]
 fn get_name_node_uri(url: &Url) -> Result<String, object_store::Error> {
     use std::fmt::Write;
@@ -422,8 +472,11 @@ fn get_name_node_uri(url: &Url) -> Result<String, object_store::Error> {
     }
 }
 
+// Stub implementation when HDFS support is not enabled
 #[cfg(all(not(feature = "hdfs"), not(feature = "hdfs-opendal")))]
-fn parse_hdfs_url(_url: &Url) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> {
+fn create_hdfs_object_store(
+    _url: &Url,
+) -> Result<(Box<dyn ObjectStore>, Path), object_store::Error> {
     Err(object_store::Error::Generic {
         store: "HadoopFileSystem",
         source: "Hdfs support is not enabled in this build".into(),
@@ -454,7 +507,7 @@ pub(crate) fn prepare_object_store_with_configs(
     );
 
     let (object_store, object_store_path): (Box<dyn ObjectStore>, Path) = if is_hdfs_scheme {
-        parse_hdfs_url(&url)
+        create_hdfs_object_store(&url)
     } else if scheme == "s3" {
         objectstore::s3::create_store(&url, object_store_configs, Duration::from_secs(300))
     } else {