Add memory limit configuration for DataFusion runtime to prevent out-of-memory errors

Author: kosiew

docs/source/user-guide/configuration.rst

@@ -46,6 +46,26 @@ a :py:class:`~datafusion.context.SessionConfig` and :py:class:`~datafusion.conte
     ctx = SessionContext(config, runtime)
     print(ctx)
 
+.. _memory_limit:
+
+Limiting memory usage
+---------------------
+
+For long-running queries or streaming results it can be helpful to cap
+DataFusion's memory consumption. Use
+``RuntimeEnvBuilder.with_memory_limit`` to set a maximum in bytes. When this
+limit is reached DataFusion will spill intermediate data to disk using a fair
+spill pool, preventing the process from exhausting system memory.
+
+.. code-block:: python
+
+    from datafusion import RuntimeEnvBuilder, SessionContext
+
+    runtime = RuntimeEnvBuilder().with_memory_limit(1_000_000)
+    ctx = SessionContext(runtime=runtime)
+
+    # queries now respect the 1 MB cap and may spill to disk
+
 
 .. _target_partitions:
 

docs/source/user-guide/io/arrow.rst

@@ -86,6 +86,9 @@ If the goal is simply to persist results, prefer engine-level writers such as
 ``df.write_parquet()``. These writers stream data from Rust directly to the
 destination and avoid Python-side memory growth.
 
+See :ref:`memory_limit` for how to configure a runtime memory cap to further
+guard against out-of-memory conditions when streaming large datasets.
+
 .. ipython:: python
 
     df = df.select((col("a") * lit(1.5)).alias("c"), lit("df").alias("d"))

python/datafusion/context.py

@@ -339,6 +339,23 @@ def with_unbounded_memory_pool(self) -> RuntimeEnvBuilder:
         self.config_internal = self.config_internal.with_unbounded_memory_pool()
         return self
 
+    def with_memory_limit(self, size: int) -> RuntimeEnvBuilder:
+        """Limit runtime memory using a fair spill pool.
+
+        Args:
+            size: Maximum number of bytes that the runtime may allocate before spilling.
+
+        Returns:
+            A new :py:class:`RuntimeEnvBuilder` object with the updated setting.
+
+        Example usage::
+
+            RuntimeEnvBuilder().with_memory_limit(1024)
+        """
+
+        self.config_internal = self.config_internal.with_memory_limit(size)
+        return self
+
     def with_fair_spill_pool(self, size: int) -> RuntimeEnvBuilder:
         """Use a fair spill pool with the specified size.
 

python/tests/test_context.py

@@ -634,6 +634,21 @@ def test_sql_with_options_no_statements(ctx):
         ctx.sql_with_options(sql, options=options)
 
 
+def test_memory_limit_streaming():
+    runtime = RuntimeEnvBuilder().with_memory_limit(1024)
+    ctx = SessionContext(runtime=runtime)
+
+    # Create a batch larger than the 1KB limit
+    array = pa.array(range(10000))
+    batch = pa.RecordBatch.from_arrays([array], names=["a"])
+    df = ctx.create_dataframe([[batch]])
+    stream = df.select(column("a") + literal(1)).execute_stream()
+
+    with pytest.raises(Exception):
+        for _ in stream:
+            pass
+
+
 @pytest.fixture
 def batch():
     return pa.RecordBatch.from_arrays(

src/context.rs

@@ -250,6 +250,10 @@ impl PyRuntimeEnvBuilder {
         Self { builder }
     }
 
+    fn with_memory_limit(&self, size: usize) -> Self {
+        self.with_fair_spill_pool(size)
+    }
+
     fn with_temp_file_path(&self, path: &str) -> Self {
         let builder = self.builder.clone();
         let builder = builder.with_temp_file_path(path);
@@ -322,7 +326,7 @@ impl PySessionContext {
         } else {
             RuntimeEnvBuilder::default()
         };
-        let runtime = Arc::new(runtime_env_builder.build()?);
+        let runtime = runtime_env_builder.build_arc()?;
         let session_state = SessionStateBuilder::new()
             .with_config(config)
             .with_runtime_env(runtime)