kosiew commented on code in PR #1627:
URL:
https://github.com/apache/datafusion-python/pull/1627#discussion_r3569629440
##########
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"):
Review Comment:
Could we also assert that `ctx.register_record_batches("t", [])` raises the
same clear `ValueError`? The fix now covers both an empty partition (`[[]]`)
and an empty outer partition list (`[]`), while the regression test currently
locks in only the first case.
##########
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
Review Comment:
Could we use `partitions.0.iter().find_map(|partition| partition.first())`
here instead of `iter().flatten().next()`? The behavior is the same, but I
think it makes the intent a little clearer: find the first record batch in any
partition, then use its schema.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
