Reorder to allow type coercion for Unions

src/pydantic_spark/base.py

@@ -64,13 +64,18 @@ def get_type(value: dict) -> Tuple[str, dict]:
             a = value.get("additionalProperties")
             ft = value.get("coerce_type")
             metadata = {}
+
+            if ft is not None:
+                return ft, metadata
+
             if ao is not None:
                 if len(ao) == 2 and (ao[0].get("type") == "null" or ao[1].get("type") == "null"):
                     # this is an optional column. We will remove the null type
                     t = ao[0].get("type") if ao[0].get("type") != "null" else ao[1].get("type")
                     f = ao[0].get("format") if ao[0].get("type") != "null" else ao[1].get("format")
                 else:
-                    NotImplementedError(f"Union type {ao} is not supported yet")
+                    NotImplementedError(f"Union type {ao} is not supported yet. Use coerce_type option to specify type")
+
             if "default" in value:
                 metadata["default"] = value.get("default")
             if r is not None:
@@ -80,8 +85,6 @@ def get_type(value: dict) -> Tuple[str, dict]:
                 else:
                     spark_type = get_type_of_definition(r, schema)
                     classes_seen[class_name] = spark_type
-            elif ft is not None:
-                spark_type = ft
             elif t == "array":
                 items = value.get("items")
                 tn, metadata = get_type(items)

tests/test_to_spark.py

@@ -1,7 +1,9 @@
+from __future__ import annotations
+
 import json
 from datetime import date, datetime
 from enum import Enum
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 from uuid import UUID
 
 from pydantic import Field
@@ -290,9 +292,11 @@ def test_enum():
 def test_coerce_type():
     class TestCoerceType(SparkBase):
         c1: int = Field(json_schema_extra={"coerce_type": CoerceType.integer})
+        c2: Union[str, int] = Field(json_schema_extra={"coerce_type": CoerceType.string})
 
     result = TestCoerceType.spark_schema()
     assert result["fields"][0]["type"] == "integer"
+    assert result["fields"][1]["type"] == "string"
 
 
 class Nested2ModelCoerceType(SparkBase):