From 255d9680249f5bc0f2131847b29f16f80ffbf68c Mon Sep 17 00:00:00 2001
From: Fangchen Li <fangchen.li@outlook.com>
Date: Fri, 10 Jul 2026 14:37:42 -0700
Subject: [PATCH] fix: no panic on empty partition in register_record_batches

`register_record_batches` read the schema via unchecked `partitions.0[0][0]`,
which panics when a partition contains no record batches (e.g. the batches of
an empty pyarrow table, whose `to_batches()` returns `[]`).

Take the schema from the first available batch across all partitions and return
a clear error when there is none, mirroring the empty-table handling added for
`from_arrow_table` in #613. A non-empty later partition now also works instead
of panicking on an empty leading one.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 crates/core/src/context.rs   | 15 ++++++++++++++-
 python/tests/test_context.py | 12 ++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs
index 0613a96dc..c484f0d57 100644
--- a/crates/core/src/context.rs
+++ b/crates/core/src/context.rs
@@ -832,7 +832,20 @@ impl PySessionContext {
         name: &str,
         partitions: PyArrowType<Vec<Vec<RecordBatch>>>,
     ) -> PyDataFusionResult<()> {
-        let schema = partitions.0[0][0].schema();
+        // Take the schema from the first available batch; error instead of
+        // panicking when the partitions hold no batches.
+        let schema = partitions
+            .0
+            .iter()
+            .flatten()
+            .next()
+            .ok_or_else(|| {
+                PyValueError::new_err(
+                    "Cannot register record batches without a schema: the \
+                     provided partitions contain no record batches.",
+                )
+            })?
+            .schema();
         let table = MemTable::try_new(schema, partitions.0)?;
         self.ctx.register_table(name, Arc::new(table))?;
         Ok(())
diff --git a/python/tests/test_context.py b/python/tests/test_context.py
index 112a6fd7b..d9887680a 100644
--- a/python/tests/test_context.py
+++ b/python/tests/test_context.py
@@ -116,6 +116,18 @@ def test_register_record_batches(ctx):
     assert result[0].column(1) == pa.array([-3, -3, -3])
 
 
+def test_register_record_batches_empty(ctx):
+    # A partition list with no record batches carries no schema, so this used to
+    # panic on unchecked `[0][0]` indexing. It should now raise a clear error.
+    with pytest.raises(ValueError, match="no record batches"):
+        ctx.register_record_batches("t", [[]])
+
+    # The schema is still recovered from a later non-empty partition.
+    batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], names=["a"])
+    ctx.register_record_batches("t2", [[], [batch]])
+    assert ctx.sql("SELECT a FROM t2").collect()[0].column(0) == pa.array([1, 2, 3])
+
+
 def test_create_dataframe_registers_unique_table_name(ctx):
     # create a RecordBatch and register it as memtable
     batch = pa.RecordBatch.from_arrays(