From 4beee15c1afd9ea28861f7a51fd2e183e55c389c Mon Sep 17 00:00:00 2001
From: Jan Heinrich Reimer <heinrich.reimer@uni-jena.de>
Date: Thu, 4 Jul 2024 01:46:22 +0200
Subject: [PATCH] Fix schema usage

---
 ray_elasticsearch/__init__.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/ray_elasticsearch/__init__.py b/ray_elasticsearch/__init__.py
index 1d85c0c..b6343c4 100644
--- 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"]