Support writing datetime columns

CHANGES.md

@@ -22,6 +22,7 @@
 -   Use certifi to set `GDAL_CURL_CA_BUNDLE` / `PROJ_CURL_CA_BUNDLE` defaults (#97)
 -   automatically detect driver for `.geojson`, `.geojsonl` and `.geojsons` files (#101)
 -   read DateTime fields with millisecond accuracy (#111)
+-   support writing datetime columns in write_dataframe (#120)
 
 ### Breaking changes
 

pyogrio/_io.pyx

@@ -76,9 +76,11 @@ DTYPE_OGR_FIELD_TYPES = {
 
     'float32': (OFTReal,OFSTFloat32),
     'float': (OFTReal, OFSTNone),
-    'float64': (OFTReal, OFSTNone)
-}
+    'float64': (OFTReal, OFSTNone),
 
+    'datetime64[D]': (OFTDate, OFSTNone),
+    'datetime64': (OFTDateTime, OFSTNone),
+}
 
 
 cdef int start_transaction(OGRDataSourceH ogr_dataset, int force) except 1:
@@ -1121,7 +1123,12 @@ cdef infer_field_types(list dtypes):
             field_types_view[i, 0] = OFTString
             field_types_view[i, 2] = int(dtype.itemsize // 4)
 
-        # TODO: datetime types
+        elif dtype.name.startswith("datetime64"):
+            # datetime dtype precision is specified with eg. [ms], but this isn't
+            # usefull when writing to gdal.
+            field_type, field_subtype = DTYPE_OGR_FIELD_TYPES["datetime64"]
+            field_types_view[i, 0] = field_type
+            field_types_view[i, 1] = field_subtype
 
         else:
             raise NotImplementedError(f"field type is not supported {dtype.name} (field index: {i})")
@@ -1385,6 +1392,35 @@ def ogr_write(str path, str layer, str driver, geometry, field_data, fields,
                 elif field_type == OFTReal:
                     OGR_F_SetFieldDouble(ogr_feature, field_idx, field_value)
 
+                elif field_type == OFTDate:
+                    datetime = field_value.astype(object)
+                    OGR_F_SetFieldDateTimeEx(
+                        ogr_feature,
+                        field_idx,
+                        datetime.year,
+                        datetime.month,
+                        datetime.day,
+                        0,
+                        0,
+                        0.0,
+                        0
+                    )
+
+                elif field_type == OFTDateTime:
+                    # TODO: what to do with tzinfo?
+                    datetime = field_value.astype("datetime64[ms]").astype(object)
+                    OGR_F_SetFieldDateTimeEx(
+                        ogr_feature,
+                        field_idx,
+                        datetime.year,
+                        datetime.month,
+                        datetime.day,
+                        datetime.hour,
+                        datetime.minute,
+                        float(f"{datetime.second}.{datetime.microsecond:06d}"),
+                        0
+                    )
+
                 else:
                     raise NotImplementedError(f"OGR field type is not supported for writing: {field_type}")
 

pyogrio/_ogr.pxd

@@ -210,6 +210,16 @@ cdef extern from "ogr_api.h":
     void            OGR_F_SetFieldString(OGRFeatureH feature, int n, char *value)
     void            OGR_F_SetFieldBinary(OGRFeatureH feature, int n, int l, unsigned char *value)
     void            OGR_F_SetFieldNull(OGRFeatureH feature, int n)  # new in GDAL 2.2
+    void            OGR_F_SetFieldDateTimeEx(
+                            OGRFeatureH hFeat,
+                            int iField,
+                            int nYear,
+                            int nMonth,
+                            int nDay,
+                            int nHour,
+                            int nMinute,
+                            float fSecond,
+                            int nTZFlag)
     OGRErr          OGR_F_SetGeometryDirectly(OGRFeatureH feature, OGRGeometryH geometry)
 
     OGRFeatureDefnH     OGR_FD_Create(const char *name)

pyogrio/tests/test_raw_io.py

@@ -422,6 +422,31 @@ def test_read_write_data_types_numeric(tmp_path, ext):
         assert result.dtype == result_dtype
 
 
+def test_read_write_datetime(tmp_path):
+    field_data = [
+        np.array(["2005-02-01", "2005-02-02"], dtype="datetime64[D]"),
+        np.array(["2001-01-01T12:00", "2002-02-03T13:56:03"], dtype="datetime64[s]"),
+        np.array(
+            ["2001-01-01T12:00", "2002-02-03T13:56:03.072"], dtype="datetime64[ms]"
+        ),
+        np.array(
+            ["2001-01-01T12:00", "2002-02-03T13:56:03.072"], dtype="datetime64[ns]"
+        ),
+    ]
+    fields = ["datetime64_d", "datetime64_s", "datetime64_ms", "datetime64_ns"]
+
+    # Point(0, 0)
+    geometry = np.array(
+        [bytes.fromhex("010100000000000000000000000000000000000000")] * 2, dtype=object
+    )
+    meta = dict(geometry_type="Point", crs="EPSG:4326", spatial_index=False)
+
+    filename = tmp_path / f"test.gpkg"
+    write(filename, geometry, field_data, fields, **meta)
+    result = read(filename)[3]
+    assert all([np.array_equal(f1, f2) for f1, f2 in zip(result, field_data)])
+
+
 def test_read_data_types_numeric_with_null(test_gpkg_nulls):
     fields = read(test_gpkg_nulls)[3]