Fix DataFrame loading doesn't support int16/uint16 types

aloneguid · Apr 16, 2024 · db69aa5 · db69aa5
1 parent 3feaac2
commit db69aa5
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 0 deletions.
diff --git a/src/Parquet.Test/DataAnalysis/DataFrameReaderTest.cs b/src/Parquet.Test/DataAnalysis/DataFrameReaderTest.cs
@@ -12,6 +12,8 @@ namespace Parquet.Test.DataAnalysis {
     public class DataFrameReaderTest : TestBase {
 
         [Theory]
+        [InlineData(typeof(short), (short)1, (short)2)]
+        [InlineData(typeof(short?), null, (short)2)]
         [InlineData(typeof(int), 1, 2)]
         [InlineData(typeof(int?), null, 2)]
         [InlineData(typeof(bool), true, false)]

diff --git a/src/Parquet/Data/Analysis/DataFrameMapper.cs b/src/Parquet/Data/Analysis/DataFrameMapper.cs
@@ -213,6 +213,32 @@ public static void AppendValues(DataFrameColumn dfc, DataColumn dc) {
                 }
                 return;
             }
+            if(dc.Field.ClrType == typeof(short)) {
+                var tdfc = (PrimitiveDataFrameColumn<short>)dfc;
+                if(dc.Field.ClrType == dc.Field.ClrNullableIfHasNullsType) {
+                    foreach(short el in (short[])dc.Data) {
+                        tdfc.Append(el);
+                    }
+                } else {
+                    foreach(short? el in (short?[])dc.Data) {
+                        tdfc.Append(el);
+                    }
+                }
+                return;
+            }
+            if(dc.Field.ClrType == typeof(ushort)) {
+                var tdfc = (PrimitiveDataFrameColumn<ushort>)dfc;
+                if(dc.Field.ClrType == dc.Field.ClrNullableIfHasNullsType) {
+                    foreach(ushort el in (ushort[])dc.Data) {
+                        tdfc.Append(el);
+                    }
+                } else {
+                    foreach(ushort? el in (ushort?[])dc.Data) {
+                        tdfc.Append(el);
+                    }
+                }
+                return;
+            }
             if(dc.Field.ClrType == typeof(DateTime)) {
                 var tdfc = (PrimitiveDataFrameColumn<DateTime>)dfc;
                 if(dc.Field.ClrType == dc.Field.ClrNullableIfHasNullsType) {
@@ -313,6 +339,12 @@ public static Array GetTypedDataFast(DataFrameColumn col) {
             if(col.DataType == typeof(sbyte)) {
                 return ((PrimitiveDataFrameColumn<sbyte>)col).ToArray();
             }
+            if(col.DataType == typeof(short)) {
+                return ((PrimitiveDataFrameColumn<short>)col).ToArray();
+            }
+            if(col.DataType == typeof(ushort)) {
+                return ((PrimitiveDataFrameColumn<ushort>)col).ToArray();
+            }
             if(col.DataType == typeof(DateTime)) {
                 return ((PrimitiveDataFrameColumn<DateTime>)col).ToArray();
             }

diff --git a/src/Parquet/Data/Analysis/DataFrameMapper.tt b/src/Parquet/Data/Analysis/DataFrameMapper.tt
@@ -10,6 +10,8 @@
         "ulong",
         "byte",
         "sbyte",
+        "short",
+        "ushort",
         "DateTime",
         "TimeSpan",
         "decimal",