diff --git a/lance_ray/io.py b/lance_ray/io.py
index 270ab287..dd093c6f 100644
--- a/lance_ray/io.py
+++ b/lance_ray/io.py
@@ -799,9 +799,8 @@ def _wrap_transform(batch: pa.Table) -> pa.Table:
             result = transform(batch_dict)
             if isinstance(result, pa.RecordBatch):
                 new_cols = pa.Table.from_batches([result])
-            elif isinstance(result, pa.Table | dict):
-                new_cols = result
             else:
+                # Already a pa.Table or dict (normalized below).
                 new_cols = result
         else:
             reader = pa.RecordBatchReader.from_batches(