Add MapType as JSON-compatible

mosaicml · Sep 6, 2024 · e39541f · e39541f
1 parent 893d9bd
commit e39541f
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/streaming/base/converters/dataframe_to_mds.py b/streaming/base/converters/dataframe_to_mds.py
@@ -19,8 +19,9 @@
     from pyspark.sql.dataframe import DataFrame
     from pyspark.sql.types import (ArrayType, BinaryType, BooleanType, ByteType, DateType,
                                    DayTimeIntervalType, DecimalType, DoubleType, FloatType,
-                                   IntegerType, LongType, NullType, ShortType, StringType,
-                                   StructField, StructType, TimestampNTZType, TimestampType)
+                                   IntegerType, LongType, MapType, NullType, ShortType,
+                                   StringType, StructField, StructType, TimestampNTZType,
+                                   TimestampType)
 except ImportError as e:
     e.msg = get_import_exception_message(e.name, extra_deps='spark')  # pyright: ignore
     raise e
@@ -70,6 +71,8 @@ def is_json_compatible(data_type: Any):
         return all(is_json_compatible(field.dataType) for field in data_type.fields)
     elif isinstance(data_type, ArrayType):
         return is_json_compatible(data_type.elementType)
+    elif isinstance(data_type, MapType):
+        return is_json_compatible(data_type.keyType) and is_json_compatible(data_type.valueType)
     elif isinstance(data_type, (StringType, IntegerType, FloatType, BooleanType, NullType)):
         return True
     else: