fix

Signed-off-by: Zhi Lin <[email protected]>
oap-project · Apr 9, 2024 · 2713e96 · 2713e96
1 parent 40d192c
commit 2713e96
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/python/raydp/spark/dataset.py b/python/raydp/spark/dataset.py
@@ -286,12 +286,11 @@ def ray_dataset_to_spark_dataframe(spark: sql.SparkSession,
                                    blocks: List[ObjectRef],
                                    locations = None) -> DataFrame:
     locations = get_locations(blocks)
+    if hasattr(arrow_schema, "base_schema"):
+        arrow_schema = arrow_schema.base_schema
     if not isinstance(arrow_schema, pa.lib.Schema):
-        if hasattr(arrow_schema, "base_schema"):
-            arrow_schema = arrow_schema.base_schema
-        if isinstance(arrow_schema, pa.lib.Schema):
-            raise RuntimeError(f"Schema is {type(arrow_schema)}, required pyarrow.lib.Schema. \n" \
-                               f"to_spark does not support converting non-arrow ray datasets.")
+        raise RuntimeError(f"Schema is {type(arrow_schema)}, required pyarrow.lib.Schema. \n" \
+                            f"to_spark does not support converting non-arrow ray datasets.")
     schema = StructType()
     for field in arrow_schema:
         schema.add(field.name, from_arrow_type(field.type), nullable=field.nullable)