Fix schema usage

janheinrichmerker · Jul 3, 2024 · 4beee15 · 4beee15
1 parent 5d6118e
commit 4beee15
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/ray_elasticsearch/__init__.py b/ray_elasticsearch/__init__.py
@@ -132,7 +132,11 @@ def iter_blocks() -> Iterator[Table]:
                     break
                 yield Table.from_pylist(
                     mapping=hits,
-                    schema=schema,
+                    schema=(
+                        schema
+                        if schema is not None and isinstance(schema, Schema)
+                        else None
+                    ),
                 )
                 search_after = max(
                     hit["sort"]