fix: no panic on empty partition in register_record_batches

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(())

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(