fix: Handle nullable enums

aloneguid · Sep 30, 2024 · 8c3634e · 8c3634e
1 parent a12d5a9
commit 8c3634e
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 8 deletions.
diff --git a/src/Parquet.Test/Serialisation/ParquetSerializerTest.cs b/src/Parquet.Test/Serialisation/ParquetSerializerTest.cs
@@ -838,18 +838,22 @@ enum ShortEnum {
 
         class EnumProps {
             public int Id { get; set; }
-
+            
             public DefaultEnum DE { get; set; }
-
+            
             public ShortEnum SE { get; set; }
+
+            // Nullable Enum
+            public DefaultEnum? NE { get; set; }
         }
 
         [Fact]
         public async Task Enums_Serde() {
             var data = Enumerable.Range(0, 1_000).Select(i => new EnumProps {
                 Id = i,
                 DE = i % 2 == 0 ? DefaultEnum.One : DefaultEnum.Two,
-                SE = i % 2 == 0 ? ShortEnum.One : ShortEnum.Two
+                SE = i % 2 == 0 ? ShortEnum.One : ShortEnum.Two,
+                NE = i % 2 == 0 ? null : DefaultEnum.One,
             }).ToList();
 
             await Compare(data);

diff --git a/src/Parquet.Test/Serialisation/SchemaReflectorTest.cs b/src/Parquet.Test/Serialisation/SchemaReflectorTest.cs
@@ -633,20 +633,28 @@ public class EnumsInClasses {
             public int Id { get; set; }
 
             public DefaultEnum DE { get; set; }
+
+            // Nullable Enum
+            public DefaultEnum? NE { get; set; }
 
             public ShortEnum SE { get; set; }
         }
 
         [Fact]
         public void Enums() {
             ParquetSchema schema = typeof(EnumsInClasses).GetParquetSchema(true);
-            Assert.Equal(3, schema.Fields.Count);
+            Assert.Equal(4, schema.Fields.Count);
 
             DataField dedf = schema.FindDataField("DE");
+            DataField nedf = schema.FindDataField("NE");
             DataField sedf = schema.FindDataField("SE");
 
             Assert.Equal(typeof(int), dedf.ClrType);
+            Assert.False(dedf.IsNullable);
+            Assert.Equal(typeof(int), nedf.ClrType);
+            Assert.True(nedf.IsNullable);
             Assert.Equal(typeof(short), sedf.ClrType);
+            Assert.False(dedf.IsNullable);
         }
     }
 }
diff --git a/src/Parquet/Serialization/TypeExtensions.cs b/src/Parquet/Serialization/TypeExtensions.cs
@@ -219,10 +219,6 @@ private static Field ConstructDataField(string name, string propertyName, Type t
                     : new DecimalDataField(name, ps.Precision, ps.Scale,
                         isNullable: isTypeNullable, propertyName: propertyName);
             } else {
-
-                if(t.IsEnum) {
-                    t = t.GetEnumUnderlyingType();
-                }
                 bool? isMemberNullable = null;
                 if (isCompiledWithNullable) {
                     isMemberNullable = member?.IsNullable(t);
@@ -231,6 +227,17 @@ private static Field ConstructDataField(string name, string propertyName, Type t
                 if(isMemberNullable is not null) {
                     isNullable = isMemberNullable.Value;
                 }
+
+                Type? nt = Nullable.GetUnderlyingType(t);
+
+                if(nt is { IsEnum: true }) {
+                    isNullable = true;
+                    t = nt.GetEnumUnderlyingType();
+                }
+                if(t.IsEnum) {
+                    t = t.GetEnumUnderlyingType();
+                }
+
                 r = new DataField(name, t, isNullable, null, propertyName, isCompiledWithNullable);
             }