Add ObjectStore-backed TempFileFactor / spill example

datafusion-examples/Cargo.toml

@@ -60,7 +60,7 @@ futures = { workspace = true }
 insta = { workspace = true }
 log = { workspace = true }
 mimalloc = { version = "0.1", default-features = false }
-object_store = { workspace = true, features = ["aws", "http"] }
+object_store = { workspace = true, features = ["aws", "fs", "http"] }
 prost = { workspace = true }
 rand = { workspace = true }
 serde = { version = "1", features = ["derive"] }

datafusion-examples/README.md

@@ -93,6 +93,7 @@ cargo run --example dataframe -- dataframe
 | catalog                 | [`data_io/catalog.rs`](examples/data_io/catalog.rs)                                       | Register tables into a custom catalog                                     |
 | in_memory_object_store  | [`data_io/in_memory_object_store.rs`](examples/data_io/in_memory_object_store.rs)         | Read CSV from an in-memory object store (pattern applies to JSON/Parquet) |
 | json_shredding          | [`data_io/json_shredding.rs`](examples/data_io/json_shredding.rs)                         | Implement filter rewriting for JSON shredding                             |
+| object_store_spill      | [`data_io/object_store_spill.rs`](examples/data_io/object_store_spill.rs)                 | Use ObjectStore-backed spill files                                        |
 | parquet_adv_idx         | [`data_io/parquet_advanced_index.rs`](examples/data_io/parquet_advanced_index.rs)         | Create a secondary index across multiple parquet files                    |
 | parquet_emb_idx         | [`data_io/parquet_embedded_index.rs`](examples/data_io/parquet_embedded_index.rs)         | Store a custom index inside Parquet files                                 |
 | parquet_enc             | [`data_io/parquet_encrypted.rs`](examples/data_io/parquet_encrypted.rs)                   | Read & write encrypted Parquet files                                      |

datafusion-examples/examples/data_io/main.rs

@@ -21,7 +21,7 @@
 //!
 //! ## Usage
 //! ```bash
-//! cargo run --example data_io -- [all|catalog|in_memory_object_store|json_shredding|parquet_adv_idx|parquet_emb_idx|parquet_enc_with_kms|parquet_enc|parquet_exec_visitor|parquet_idx|query_http_csv|remote_catalog]
+//! cargo run --example data_io -- [all|catalog|in_memory_object_store|json_shredding|object_store_spill|parquet_adv_idx|parquet_emb_idx|parquet_enc_with_kms|parquet_enc|parquet_exec_visitor|parquet_idx|query_http_csv|remote_catalog]
 //! ```
 //!
 //! Each subcommand runs a corresponding example:
@@ -36,6 +36,9 @@
 //! - `json_shredding`
 //!   (file: json_shredding.rs, desc: Implement filter rewriting for JSON shredding)
 //!
+//! - `object_store_spill`
+//!   (file: object_store_spill.rs, desc: Use ObjectStore-backed spill files)
+//!
 //! - `parquet_adv_idx`
 //!   (file: parquet_advanced_index.rs, desc: Create a secondary index across multiple parquet files)
 //!
@@ -66,6 +69,7 @@
 mod catalog;
 mod in_memory_object_store;
 mod json_shredding;
+mod object_store_spill;
 mod parquet_advanced_index;
 mod parquet_embedded_index;
 mod parquet_encrypted;
@@ -87,6 +91,7 @@ enum ExampleKind {
     Catalog,
     InMemoryObjectStore,
     JsonShredding,
+    ObjectStoreSpill,
     ParquetAdvIdx,
     ParquetEmbIdx,
     ParquetEnc,
@@ -118,6 +123,9 @@ impl ExampleKind {
                 in_memory_object_store::in_memory_object_store().await?
             }
             ExampleKind::JsonShredding => json_shredding::json_shredding().await?,
+            ExampleKind::ObjectStoreSpill => {
+                object_store_spill::object_store_spill().await?
+            }
             ExampleKind::ParquetAdvIdx => {
                 parquet_advanced_index::parquet_advanced_index().await?
             }

datafusion-examples/examples/data_io/object_store_spill.rs

@@ -0,0 +1,273 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! See `main.rs` for how to run it.
+//!
+//! [`object_store_spill`] demonstrates how to use the [`TempFileFactory`] API to configure
+//! DataFusion to spill intermediate results to remote storage when it exceeds
+//! the configured memory limits.
+//!
+//! See [`datafusion::execution::memory_pool`] for more information on how
+//! DataFusion decides when operators should spill, and [`SpillFile`] for the
+//! spill file abstraction this example implements.
+use std::future::Future;
+use std::io::Write;
+use std::path::Path as StdPath;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+
+use bytes::Bytes;
+use datafusion::common::Result;
+use datafusion::execution::disk_manager::DiskManagerBuilder;
+use datafusion::execution::runtime_env::RuntimeEnvBuilder;
+use datafusion::execution::{SpillFile, SpillWriter, TempFileFactory};
+use datafusion::prelude::{SessionConfig, SessionContext};
+use datafusion_common::exec_err;
+use futures::{Stream, StreamExt, TryStreamExt, stream};
+use object_store::local::LocalFileSystem;
+use object_store::path::Path;
+use object_store::{ObjectStore, ObjectStoreExt, PutPayload};
+use tempfile::tempdir;
+
+/// Demonstrates configuring DataFusion with spill files backed by an ObjectStore.
+pub async fn object_store_spill() -> Result<()> {
+    // A real system would use S3, GCS, Azure, or some other ObjectStore for
+    // remote spills. This example uses a local-file-backed ObjectStore for
+    // simplicity.
+    let tmp_dir = tempdir()?;
+    let store: Arc<dyn ObjectStore> =
+        Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path())?);
+
+    // Create the custom TempFileFactory that creates spill files in the ObjectStore.
+    let temp_file_factory = Arc::new(ObjectStoreTempFileFactory::new(store));
+    let disk_manager_builder =
+        DiskManagerBuilder::default().with_temp_file_factory(temp_file_factory.clone());
+    let runtime = RuntimeEnvBuilder::new()
+        .with_disk_manager_builder(disk_manager_builder) // use the factory
+        // and set a small memory limit so the example spills
+        .with_memory_limit(1024 * 1024, 1.0)
+        .build_arc()?;
+
+    // Configure a SessionContext for running queries; use a single partition
+    // and no sort spill reservation to make the example deterministic and keep
+    // the spill behavior easy to observe.
+    let config = SessionConfig::new()
+        .with_sort_spill_reservation_bytes(0)
+        .with_sort_in_place_threshold_bytes(0)
+        .with_target_partitions(1);
+    let ctx = SessionContext::new_with_config_rt(config, Arc::clone(&runtime));
+
+    // Run an SQL query that sorts a "large" amount of data. Given the
+    // SessionContext's low memory limit, the sort will spill.
+    let row_count = 10_000_000;
+    let mut stream = ctx
+        .sql(&format!(
+            "SELECT * FROM generate_series(1, {row_count}) AS t(v) ORDER BY v DESC"
+        ))
+        .await?
+        .execute_stream()
+        .await?;
+
+    // Drive the query to completion, and verify output
+    let mut output_rows = 0;
+    while let Some(batch) = stream.next().await {
+        output_rows += batch?.num_rows();
+    }
+
+    assert_eq!(output_rows, row_count as usize);
+    assert!(
+        temp_file_factory.created_files() > 0,
+        "expected the custom TempFileFactory to be used for spilling"
+    );
+
+    Ok(())
+}
+
+/// Creates spill files backed by an [`ObjectStore`].
+///
+/// DataFusion calls this factory whenever an operator needs a new temporary
+/// file for spilling. A remote deployment would use the same pattern with an
+/// S3, GCS, Azure, or other remote ObjectStore implementation.
+struct ObjectStoreTempFileFactory {
+    /// ObjectStore used for spill file reads and writes.
+    store: Arc<dyn ObjectStore>,
+    /// Monotonic counter used to create unique object paths.
+    counter: AtomicU64,
+    /// Counts how many spill files DataFusion requested from this factory.
+    created_files: AtomicU64,
+}
+
+impl ObjectStoreTempFileFactory {
+    /// Create a new spill file factory that stores spill data in `store`.
+    fn new(store: Arc<dyn ObjectStore>) -> Self {
+        Self {
+            store,
+            counter: AtomicU64::new(0),
+            created_files: AtomicU64::new(0),
+        }
+    }
+
+    /// Return the number of spill files created through this factory.
+    fn created_files(&self) -> u64 {
+        self.created_files.load(Ordering::Relaxed)
+    }
+}
+
+impl TempFileFactory for ObjectStoreTempFileFactory {
+    /// Create one logical spill file backed by an ObjectStore path.
+    fn create_temp_file(&self, description: &str) -> Result<Arc<dyn SpillFile>> {
+        let id = self.counter.fetch_add(1, Ordering::Relaxed);
+        self.created_files.fetch_add(1, Ordering::Relaxed);
+
+        // Convert a query-provided spill description into an ObjectStore-safe path component.
+        //
+        // For example, `"Sort Spill: partition 0"` becomes `"Sort_Spill__partition_0"`.
+        let cleaned_description: String = description
+            .chars()
+            .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
+            .collect();
+        let location = Path::from(format!("spill/{cleaned_description}-{id}.bin"));
+
+        // Return a SpillFile implementation that reads and writes this ObjectStore path.
+        Ok(Arc::new(ObjectStoreSpillFile {
+            store: Arc::clone(&self.store),
+            location,
+            size: Arc::new(AtomicU64::new(0)),
+        }))
+    }
+}
+
+/// Logical spill file stored at an ObjectStore path.
+///
+/// DataFusion writes spill data by calling [`SpillFile::open_writer`] and reads
+/// it back by calling [`SpillFile::read_stream`].
+struct ObjectStoreSpillFile {
+    /// ObjectStore containing the spill object.
+    store: Arc<dyn ObjectStore>,
+    /// ObjectStore path for this spill object.
+    location: Path,
+    /// Last committed object size, updated when the writer finishes.
+    size: Arc<AtomicU64>,
+}
+
+impl SpillFile for ObjectStoreSpillFile {
+    /// Return no local filesystem path because the spill file is accessed through ObjectStore.
+    fn path(&self) -> Option<&StdPath> {
+        None // Remote ObjectStores do not have a local OS path.
+    }
+
+    /// Return the size of the uploaded object
+    fn size(&self) -> Option<u64> {
+        // Return the last committed size, which this example tracks after upload.
+        Some(self.size.load(Ordering::Relaxed))
+    }
+
+    /// Read the spill file contents as a byte stream.
+    fn read_stream(&self) -> Result<Pin<Box<dyn Stream<Item = Result<Bytes>> + Send>>> {
+        let store = Arc::clone(&self.store);
+        let location = self.location.clone();
+
+        // Use `stream::once` to defer the ObjectStore read until DataFusion
+        // polls the returned stream.
+        let result_stream =
+            async move { store.get(&location).await.map(|r| r.into_stream()) };
+        let stream = stream::once(result_stream)
+            .try_flatten()
+            .map_err(Into::into);
+
+        Ok(Box::pin(stream))
+    }
+
+    /// Open a synchronous writer for this spill file.
+    fn open_writer(&self) -> Result<Box<dyn SpillWriter>> {
+        // Create a writer that buffers bytes and uploads them on finish.
+        Ok(Box::new(ObjectStoreSpillWriter {
+            store: Arc::clone(&self.store),
+            location: self.location.clone(),
+            size: Arc::clone(&self.size),
+            buffer: Vec::new(),
+        }))
+    }
+}
+
+/// Adapts DataFusion's [`SpillWriter`] API to ObjectStore.
+///
+/// This simple example buffers bytes in memory and uploads them in
+/// [`SpillWriter::finish`]. A production remote implementation should consider
+/// multipart or streaming uploads.
+struct ObjectStoreSpillWriter {
+    /// ObjectStore to read/write bytes to.
+    store: Arc<dyn ObjectStore>,
+    /// ObjectStore path to upload to.
+    location: Path,
+    /// Shared size field on the corresponding [`ObjectStoreSpillFile`].
+    size: Arc<AtomicU64>,
+    /// Buffered spill bytes waiting to be uploaded.
+    ///
+    /// This simple example buffers the spill and uploads it on finish.
+    /// Production remote stores should consider multipart or streaming uploads.
+    buffer: Vec<u8>,
+}
+
+impl Write for ObjectStoreSpillWriter {
+    /// Append bytes to the in-memory buffer.
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        // Buffer bytes written through the synchronous Write API.
+        self.buffer.extend_from_slice(buf);
+        Ok(buf.len())
+    }
+
+    /// No-op because data is committed in [`SpillWriter::finish`].
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
+    }
+}
+
+impl SpillWriter for ObjectStoreSpillWriter {
+    /// Upload buffered bytes to ObjectStore and mark the spill file complete.
+    fn finish(&mut self) -> Result<()> {
+        // Move the buffered bytes into the upload future.
+        let store = Arc::clone(&self.store);
+        let location = self.location.clone();
+        let data = std::mem::take(&mut self.buffer);
+        let size = data.len() as u64;
+
+        // This simple example buffers the spill and uploads it on finish.
+        // Production remote stores should consider multipart or streaming uploads.
+        block_on_object_store(async move {
+            store
+                .put(&location, PutPayload::from_bytes(data.into()))
+                .await?;
+            Ok(())
+        })?;
+
+        self.size.store(size, Ordering::Relaxed);
+        Ok(())
+    }
+}
+
+/// Run an async ObjectStore operation.
+///
+/// Adding a native async API is tracked in <https://github.com/apache/datafusion/issues/23247>
+fn block_on_object_store<T>(future: impl Future<Output = Result<T>>) -> Result<T> {
+    if let Ok(handle) = tokio::runtime::Handle::try_current() {
+        tokio::task::block_in_place(|| handle.block_on(future))
+    } else {
+        exec_err!("No current Tokio runtime available")
+    }
+}

datafusion/execution/src/disk_manager.rs

@@ -69,6 +69,26 @@ impl DiskManagerBuilder {
         self
     }
 
+    /// Configure a custom factory for creating temporary spill files.
+    ///
+    /// This sets the disk manager mode to [`DiskManagerMode::Custom`], so
+    /// operators that spill during query execution create files through the
+    /// provided [`TempFileFactory`] instead of using local temporary files.
+    pub fn set_temp_file_factory(&mut self, temp_file_factory: Arc<dyn TempFileFactory>) {
+        self.mode = DiskManagerMode::Custom(temp_file_factory);
+    }
+
+    /// Configure a custom factory for creating temporary spill files.
+    ///
+    /// See details on [`Self::set_temp_file_factory`].
+    pub fn with_temp_file_factory(
+        mut self,
+        temp_file_factory: Arc<dyn TempFileFactory>,
+    ) -> Self {
+        self.set_temp_file_factory(temp_file_factory);
+        self
+    }
+
     pub fn set_max_temp_directory_size(&mut self, value: u64) {
         self.max_temp_directory_size = value;
     }