Skip to content

Commit 255d968

Browse files
fangchenliclaude
andcommitted
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>
1 parent 5ce35f8 commit 255d968

2 files changed

Lines changed: 26 additions & 1 deletion

File tree

crates/core/src/context.rs

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -832,7 +832,20 @@ impl PySessionContext {
832832
name: &str,
833833
partitions: PyArrowType<Vec<Vec<RecordBatch>>>,
834834
) -> PyDataFusionResult<()> {
835-
let schema = partitions.0[0][0].schema();
835+
// Take the schema from the first available batch; error instead of
836+
// panicking when the partitions hold no batches.
837+
let schema = partitions
838+
.0
839+
.iter()
840+
.flatten()
841+
.next()
842+
.ok_or_else(|| {
843+
PyValueError::new_err(
844+
"Cannot register record batches without a schema: the \
845+
provided partitions contain no record batches.",
846+
)
847+
})?
848+
.schema();
836849
let table = MemTable::try_new(schema, partitions.0)?;
837850
self.ctx.register_table(name, Arc::new(table))?;
838851
Ok(())

python/tests/test_context.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,18 @@ def test_register_record_batches(ctx):
116116
assert result[0].column(1) == pa.array([-3, -3, -3])
117117

118118

119+
def test_register_record_batches_empty(ctx):
120+
# A partition list with no record batches carries no schema, so this used to
121+
# panic on unchecked `[0][0]` indexing. It should now raise a clear error.
122+
with pytest.raises(ValueError, match="no record batches"):
123+
ctx.register_record_batches("t", [[]])
124+
125+
# The schema is still recovered from a later non-empty partition.
126+
batch = pa.RecordBatch.from_arrays([pa.array([1, 2, 3])], names=["a"])
127+
ctx.register_record_batches("t2", [[], [batch]])
128+
assert ctx.sql("SELECT a FROM t2").collect()[0].column(0) == pa.array([1, 2, 3])
129+
130+
119131
def test_create_dataframe_registers_unique_table_name(ctx):
120132
# create a RecordBatch and register it as memtable
121133
batch = pa.RecordBatch.from_arrays(

0 commit comments

Comments
 (0)