Support for reading/writing DATE32 (#220)

innobi · Jan 14, 2024 · bbdb71e · bbdb71e
1 parent 6059ad7
commit bbdb71e
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/pantab/src/pantab.cpp b/pantab/src/pantab.cpp
@@ -49,6 +49,8 @@ static hyperapi::SqlType hyperTypeFromArrowSchema(struct ArrowSchema *schema,
   case NANOARROW_TYPE_STRING:
   case NANOARROW_TYPE_LARGE_STRING:
     return hyperapi::SqlType::text();
+  case NANOARROW_TYPE_DATE32:
+    return hyperapi::SqlType::date();
   case NANOARROW_TYPE_TIMESTAMP:
     if (std::strcmp("", schema_view.timezone)) {
       return hyperapi::SqlType::timestampTZ();
@@ -137,6 +139,40 @@ template <typename OffsetT> class Utf8InsertHelper : public InsertHelper {
   }
 };
 
+class Date32InsertHelper : public InsertHelper {
+public:
+  using InsertHelper::InsertHelper;
+
+  void insertValueAtIndex(size_t idx) override {
+    constexpr size_t elem_size = sizeof(int32_t);
+    int32_t value;
+    if (ArrowArrayViewIsNull(&array_view_, idx)) {
+      // MSVC on cibuildwheel doesn't like this templated optional
+      // inserter_->add(std::optional<timestamp_t>{std::nullopt});
+      hyperapi::internal::ValueInserter{*inserter_}.addNull();
+      return;
+    }
+
+    memcpy(&value,
+           array_view_.buffer_views[1].data.as_uint8 + (idx * elem_size),
+           elem_size);
+
+    const std::chrono::duration<int32_t, std::ratio<86400>> dur{value};
+    const std::chrono::time_point<
+        std::chrono::system_clock,
+        std::chrono::duration<int32_t, std::ratio<86400>>>
+        tp{dur};
+    const auto tt = std::chrono::system_clock::to_time_t(tp);
+
+    const struct tm utc_tm = *std::gmtime(&tt);
+    hyperapi::Date dt{1900 + utc_tm.tm_year,
+                      static_cast<int16_t>(1 + utc_tm.tm_mon),
+                      static_cast<int16_t>(1 + utc_tm.tm_yday)};
+
+    inserter_->add(dt);
+  }
+};
+
 template <enum TimeUnit TU, bool TZAware>
 class TimestampInsertHelper : public InsertHelper {
 public:
@@ -236,6 +272,9 @@ makeInsertHelper(std::shared_ptr<hyperapi::Inserter> inserter,
   case NANOARROW_TYPE_LARGE_STRING:
     return std::unique_ptr<InsertHelper>(new Utf8InsertHelper<int64_t>(
         inserter, chunk, schema, error, column_position));
+  case NANOARROW_TYPE_DATE32:
+    return std::unique_ptr<InsertHelper>(new Date32InsertHelper(
+        inserter, chunk, schema, error, column_position));
   case NANOARROW_TYPE_TIMESTAMP:
     switch (schema_view.time_unit) {
     case NANOARROW_TIME_UNIT_SECOND:

diff --git a/pantab/tests/conftest.py b/pantab/tests/conftest.py
@@ -22,6 +22,7 @@ def get_basic_dataframe():
                 2.0,
                 True,
                 True,
+                pd.to_datetime("2024-01-01"),
                 pd.to_datetime("2018-01-01"),
                 pd.to_datetime("2018-01-01", utc=True),
                 "foo",
@@ -46,6 +47,7 @@ def get_basic_dataframe():
                 2.0,
                 False,
                 False,
+                pd.to_datetime("2024-01-01"),
                 pd.to_datetime("1/1/19"),
                 pd.to_datetime("2019-01-01", utc=True),
                 "bar",
@@ -72,6 +74,7 @@ def get_basic_dataframe():
                 pd.NA,
                 pd.NaT,
                 pd.NaT,
+                pd.NaT,
                 np.nan,
                 pd.NA,
                 0,
@@ -95,6 +98,7 @@ def get_basic_dataframe():
             "Float64",
             "bool",
             "boolean",
+            "date32",
             "datetime64",
             "datetime64_utc",
             "object",
@@ -122,6 +126,7 @@ def get_basic_dataframe():
             "Float64": "Float64",
             "bool": bool,
             "boolean": "boolean",
+            "date32": "date32[pyarrow]",
             "datetime64": "datetime64[ns]",
             "datetime64_utc": "datetime64[ns, UTC]",
             "object": "object",