From 255d9680249f5bc0f2131847b29f16f80ffbf68c Mon Sep 17 00:00:00 2001 From: Fangchen Li Date: Fri, 10 Jul 2026 14:37:42 -0700 Subject: [PATCH 1/2] 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) --- 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>>, ) -> 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( From f3a3fcc5030017f755f53e2724d500517a4e89fc Mon Sep 17 00:00:00 2001 From: Fangchen Li Date: Mon, 13 Jul 2026 12:03:13 -0700 Subject: [PATCH 2/2] refactor: address review feedback - Use find_map(|p| p.first()) instead of flatten().next() to make the "first batch in any partition" intent clearer (per @kosiew). - Extend regression test to also assert the empty outer partition list case ([]) raises the same ValueError. Co-Authored-By: Claude Opus 4.8 (1M context) --- crates/core/src/context.rs | 3 +-- python/tests/test_context.py | 4 ++++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/crates/core/src/context.rs b/crates/core/src/context.rs index c484f0d57..6184c7c74 100644 --- a/crates/core/src/context.rs +++ b/crates/core/src/context.rs @@ -837,8 +837,7 @@ impl PySessionContext { let schema = partitions .0 .iter() - .flatten() - .next() + .find_map(|partition| partition.first()) .ok_or_else(|| { PyValueError::new_err( "Cannot register record batches without a schema: the \ diff --git a/python/tests/test_context.py b/python/tests/test_context.py index d9887680a..2f068b05d 100644 --- a/python/tests/test_context.py +++ b/python/tests/test_context.py @@ -122,6 +122,10 @@ def test_register_record_batches_empty(ctx): with pytest.raises(ValueError, match="no record batches"): ctx.register_record_batches("t", [[]]) + # An empty outer partition list carries no schema either, and raises the same 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]])