From 81e34c512e12d3113d5a8edb098b79695ebda116 Mon Sep 17 00:00:00 2001
From: Aleksei Smirnov <tlalok@inbox.ru>
Date: Mon, 8 Apr 2024 11:41:20 +0300
Subject: [PATCH] Fix DataFrame loading doesn't support int16/uint16 types

---
 .../DataAnalysis/DataFrameReaderTest.cs       |  2 ++
 src/Parquet/Data/Analysis/DataFrameMapper.cs  | 32 +++++++++++++++++++
 src/Parquet/Data/Analysis/DataFrameMapper.tt  |  2 ++
 3 files changed, 36 insertions(+)

diff --git a/src/Parquet.Test/DataAnalysis/DataFrameReaderTest.cs b/src/Parquet.Test/DataAnalysis/DataFrameReaderTest.cs
index 4b376a90..befc25c4 100644
--- 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
index 6e92831a..8fedd9e9 100644
--- 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
index 71e94cce..57fbb484 100644
--- 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",