Revert "revert branch UNPICK"

This reverts commit a63d083.

Author: kosiew

AGENTS.md

@@ -48,6 +48,7 @@ uuid = { version = "1.18", features = ["v4"] }
 mimalloc = { version = "0.1", optional = true, default-features = false, features = ["local_dynamic_tls"] }
 async-trait = "0.1.89"
 futures = "0.3"
+rayon = "1.10"
 object_store = { version = "0.12.3", features = ["aws", "gcp", "azure", "http"] }
 url = "2"
 log = "0.4.27"

Cargo.lock

@@ -0,0 +1,24 @@
+import time
+
+import pyarrow as pa
+from datafusion import SessionContext
+
+
+def run(n_batches: int = 8, batch_size: int = 1_000_000) -> None:
+    ctx = SessionContext()
+    batches = []
+    for i in range(n_batches):
+        start = i * batch_size
+        arr = pa.array(range(start, start + batch_size))
+        batches.append(pa.record_batch([arr], names=["a"]))
+
+    df = ctx.create_dataframe([batches])
+
+    start = time.perf_counter()
+    df.collect()
+    duration = time.perf_counter() - start
+    print(f"{n_batches} batches collected in {duration:.3f}s")
+
+
+if __name__ == "__main__":
+    run()

Cargo.toml

@@ -47,5 +47,26 @@ a :py:class:`~datafusion.context.SessionConfig` and :py:class:`~datafusion.conte
     print(ctx)
 
 
+.. _target_partitions:
+
+Target partitions and threads
+-----------------------------
+
+The :py:meth:`~datafusion.context.SessionConfig.with_target_partitions` method
+controls how many partitions DataFusion uses when executing a query. Each
+partition is processed on its own thread, so this setting effectively limits
+the number of threads that will be scheduled.
+
+For most workloads a good starting value is the number of logical CPU cores on
+your machine. You can use :func:`os.cpu_count` to automatically configure this::
+
+    import os
+    config = SessionConfig().with_target_partitions(os.cpu_count())
+
+Choosing a value significantly higher than the available cores can lead to
+excessive context switching without performance gains, while a much lower value
+may underutilize the machine.
+
+
 You can read more about available :py:class:`~datafusion.context.SessionConfig` options in the `rust DataFusion Configuration guide <https://arrow.apache.org/datafusion/user-guide/configs.html>`_,
 and about :code:`RuntimeEnvBuilder` options in the rust `online API documentation <https://docs.rs/datafusion/latest/datafusion/execution/runtime_env/struct.RuntimeEnvBuilder.html>`_.

benchmarks/collect_gil_bench.py

@@ -0,0 +1,26 @@
+# RecordBatch conversion and the GIL
+
+Profiling `DataFrame.collect` showed that converting each `RecordBatch` to
+PyArrow via `rb.to_pyarrow(py)` spent considerable time holding the Python GIL.
+Using `py-spy` on a query returning many batches indicated that more than
+95 % of the conversion executed while the GIL was held, meaning the work was
+effectively serialised.
+For queries that return many batches this limited CPU utilisation because only
+one conversion could run at a time.
+
+The implementation now converts each batch to Arrow's C data (schema/array)
+while the GIL is released, acquiring the GIL only to wrap those pointers into
+PyArrow objects. This allows the CPU intensive portions of the conversion to
+run fully in parallel.
+
+A simple benchmark is provided in `benchmarks/collect_gil_bench.py`.
+Run it twice to compare serial and parallel conversions:
+
+```bash
+RAYON_NUM_THREADS=1 python benchmarks/collect_gil_bench.py   # serial
+python benchmarks/collect_gil_bench.py                      # parallel
+```
+
+On this container, collecting 128 1 M‑row batches took around 1.5 s with
+`RAYON_NUM_THREADS=1` and 0.8 s with the default thread pool, demonstrating
+that releasing the GIL allows conversions to run in parallel.

docs/source/user-guide/configuration.rst

@@ -25,8 +25,10 @@ The ``DataFrame`` class is the core abstraction in DataFusion that represents ta
 on that data. DataFrames provide a flexible API for transforming data through various operations such as
 filtering, projection, aggregation, joining, and more.
 
-A DataFrame represents a logical plan that is lazily evaluated. The actual execution occurs only when 
-terminal operations like ``collect()``, ``show()``, or ``to_pandas()`` are called.
+A DataFrame represents a logical plan that is lazily evaluated. The actual execution occurs only when
+terminal operations like ``collect()``, ``show()``, or ``to_pandas()`` are called. ``collect()`` loads
+all record batches into Python memory; for large results you may want to stream data instead using
+``execute_stream()`` or ``__arrow_c_stream__()``.
 
 Creating DataFrames
 -------------------
@@ -128,27 +130,47 @@ DataFusion's DataFrame API offers a wide range of operations:
 
 Terminal Operations
 -------------------
-
-To materialize the results of your DataFrame operations:
+``collect()`` materializes every record batch in Python. While convenient, this
+eagerly loads the full result set into memory and can overwhelm the Python
+process for large queries. Alternatives that stream data from Rust avoid this
+memory growth:
 
 .. code-block:: python
 
-    # Collect all data as PyArrow RecordBatches
+    # Collect all data as PyArrow RecordBatches (loads entire result set)
     result_batches = df.collect()
-    
-    # Convert to various formats
+
+    # Stream batches using the native API
+    stream = df.execute_stream()
+    for batch in stream:
+        ...  # process each RecordBatch
+
+    # Stream via the Arrow C Data Interface
+    import pyarrow as pa
+    reader = pa.ipc.RecordBatchStreamReader._import_from_c(df.__arrow_c_stream__())
+    for batch in reader:
+        ...
+
+    # Convert to various formats (also load all data into memory)
     pandas_df = df.to_pandas()        # Pandas DataFrame
     polars_df = df.to_polars()        # Polars DataFrame
     arrow_table = df.to_arrow_table() # PyArrow Table
     py_dict = df.to_pydict()          # Python dictionary
     py_list = df.to_pylist()          # Python list of dictionaries
-    
+
     # Display results
     df.show()                         # Print tabular format to console
-    
+
     # Count rows
     count = df.count()
 
+For large outputs, prefer engine-level writers such as ``df.write_parquet()``
+or other DataFusion writers. These stream data directly to the destination and
+avoid buffering the entire dataset in Python.
+
+For more on parallel record batch conversion and the Python GIL, see
+:doc:`collect-gil`.
+
 HTML Rendering
 --------------
 
@@ -207,3 +229,4 @@ For a complete list of available functions, see the :py:mod:`datafusion.function
    :maxdepth: 1
 
    rendering
+   collect-gil

docs/source/user-guide/dataframe/collect-gil.md

@@ -57,17 +57,34 @@ and returns a ``StructArray``. Common pyarrow sources you can use are:
 Exporting from DataFusion
 -------------------------
 
-DataFusion DataFrames implement ``__arrow_c_stream__`` PyCapsule interface, so any
-Python library that accepts these can import a DataFusion DataFrame directly.
+DataFusion DataFrames implement ``__arrow_c_stream__`` so any Python library
+that accepts this interface can import a DataFusion ``DataFrame`` directly.
 
-.. warning::
-    It is important to note that this will cause the DataFrame execution to happen, which may be
-    a time consuming task. That is, you will cause a
-    :py:func:`datafusion.dataframe.DataFrame.collect` operation call to occur.
+``collect()`` or ``pa.table(df)`` will materialize every record batch in
+Python. For large results this can quickly exhaust memory. Instead, stream the
+output incrementally:
 
+.. ipython:: python
+
+    # Stream batches with DataFusion's native API
+    stream = df.execute_stream()
+    for batch in stream:
+        ...  # process each RecordBatch as it arrives
+
+.. ipython:: python
+
+    # Expose a C stream that PyArrow can consume lazily
+    import pyarrow as pa
+    reader = pa.ipc.RecordBatchStreamReader._import_from_c(df.__arrow_c_stream__())
+    for batch in reader:
+        ...  # process each batch without buffering the entire table
+
+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.
 
 .. ipython:: python
 
     df = df.select((col("a") * lit(1.5)).alias("c"), lit("df").alias("d"))
-    pa.table(df)
+    pa.table(df)  # loads all batches into memory
 

docs/source/user-guide/dataframe/index.rst

@@ -161,7 +161,13 @@ def with_batch_size(self, batch_size: int) -> SessionConfig:
     def with_target_partitions(self, target_partitions: int) -> SessionConfig:
         """Customize the number of target partitions for query execution.
 
-        Increasing partitions can increase concurrency.
+        Each partition is processed on its own thread, so this value controls
+        the degree of parallelism. A good starting point is the number of
+        logical CPU cores on your machine, for example
+        ``SessionConfig().with_target_partitions(os.cpu_count())``.
+
+        See the :ref:`configuration guide <target_partitions>` for more
+        discussion on choosing a value.
 
         Args:
             target_partitions: Number of target partitions.