Add support for datetime.datetime

Change-Id: Id38076d044006da7bfcf939832f82b345c7ca03b

python/pyarrow/includes/libarrow.pxd

@@ -45,6 +45,12 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         Type_LIST" arrow::Type::LIST"
         Type_STRUCT" arrow::Type::STRUCT"
 
+    enum TimeUnit" arrow::TimeUnit":
+        TimeUnit_SECOND" arrow::TimeUnit::SECOND"
+        TimeUnit_MILLI" arrow::TimeUnit::MILLI"
+        TimeUnit_MICRO" arrow::TimeUnit::MICRO"
+        TimeUnit_NANO" arrow::TimeUnit::NANO"
+
     cdef cppclass CDataType" arrow::DataType":
         Type type
 
@@ -75,6 +81,9 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
     cdef cppclass CStringType" arrow::StringType"(CDataType):
         pass
 
+    cdef cppclass CTimestampType" arrow::TimestampType"(CDataType):
+        TimeUnit unit
+
     cdef cppclass CField" arrow::Field":
         c_string name
         shared_ptr[CDataType] type
@@ -136,6 +145,9 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
     cdef cppclass CDateArray" arrow::DateArray"(CArray):
         int64_t Value(int i)
 
+    cdef cppclass CTimestampArray" arrow::TimestampArray"(CArray):
+        int64_t Value(int i)
+
     cdef cppclass CFloatArray" arrow::FloatArray"(CArray):
         float Value(int i)
 

python/pyarrow/scalar.pyx

@@ -130,6 +130,25 @@ cdef class DateValue(ArrayValue):
         return datetime.date.fromtimestamp(ap.Value(self.index) / 1000)
 
 
+cdef class TimestampValue(ArrayValue):
+
+    def as_py(self):
+        cdef:
+            CTimestampArray* ap = <CTimestampArray*> self.sp_array.get()
+            CTimestampType* dtype = <CTimestampType*>ap.type().get()
+            int64_t val = ap.Value(self.index)
+
+        if dtype.unit == TimeUnit_SECOND:
+            return datetime.datetime.utcfromtimestamp(val)
+        elif dtype.unit == TimeUnit_MILLI:
+            return datetime.datetime.utcfromtimestamp(float(val) / 1000)
+        elif dtype.unit == TimeUnit_MICRO:
+            return datetime.datetime.utcfromtimestamp(float(val) / 1000000)
+        else:
+            # TimeUnit_NANO
+            raise NotImplementedError("Cannot convert nanosecond timestamps to datetime.datetime")
+
+
 cdef class FloatValue(ArrayValue):
 
     def as_py(self):
@@ -195,6 +214,7 @@ cdef dict _scalar_classes = {
     Type_INT32: Int32Value,
     Type_INT64: Int64Value,
     Type_DATE: DateValue,
+    Type_TIMESTAMP: TimestampValue,
     Type_FLOAT: FloatValue,
     Type_DOUBLE: DoubleValue,
     Type_LIST: ListValue,

python/pyarrow/tests/test_convert_builtin.py

@@ -82,6 +82,22 @@ class TestConvertList(unittest.TestCase):
         assert arr[2].as_py() == datetime.date(1970, 1, 1)
         assert arr[3].as_py() == datetime.date(2040, 2, 26)
 
+    def test_timestamp(self):
+        data = [
+            datetime.datetime(2007, 7, 13, 1, 23, 34, 123456),
+            None,
+            datetime.datetime(2006, 1, 13, 12, 34, 56, 432539),
+            datetime.datetime(2010, 8, 13, 5, 46, 57, 437699)
+        ]
+        arr = pyarrow.from_pylist(data)
+        assert len(arr) == 4
+        assert arr.type == pyarrow.timestamp()
+        assert arr.null_count == 1
+        assert arr[0].as_py() == datetime.datetime(2007, 7, 13, 1, 23, 34, 123456)
+        assert arr[1].as_py() is None
+        assert arr[2].as_py() == datetime.datetime(2006, 1, 13, 12, 34, 56, 432539)
+        assert arr[3].as_py() == datetime.datetime(2010, 8, 13, 5, 46, 57, 437699)
+
     def test_mixed_nesting_levels(self):
         pyarrow.from_pylist([1, 2, None])
         pyarrow.from_pylist([[1], [2], None])

python/src/pyarrow/adapters/builtin.cc

@@ -57,6 +57,7 @@ class ScalarVisitor {
       bool_count_(0),
       int_count_(0),
       date_count_(0),
+      timestamp_count_(0),
       float_count_(0),
       string_count_(0) {}
 
@@ -70,8 +71,10 @@ class ScalarVisitor {
       ++float_count_;
     } else if (IsPyInteger(obj)) {
       ++int_count_;
-    } else if (PyDate_Check(obj)) {
+    } else if (PyDate_CheckExact(obj)) {
       ++date_count_;
+    } else if (PyDateTime_CheckExact(obj)) {
+      ++timestamp_count_;
     } else if (IsPyBaseString(obj)) {
       ++string_count_;
     } else {
@@ -88,6 +91,8 @@ class ScalarVisitor {
       return INT64;
     } else if (date_count_) {
       return DATE;
+    } else if (timestamp_count_) {
+      return TIMESTAMP_US;
     } else if (bool_count_) {
       return BOOL;
     } else if (string_count_) {
@@ -107,6 +112,7 @@ class ScalarVisitor {
   int64_t bool_count_;
   int64_t int_count_;
   int64_t date_count_;
+  int64_t timestamp_count_;
   int64_t float_count_;
   int64_t string_count_;
 
@@ -333,6 +339,39 @@ class DateConverter : public TypedConverter<arrow::DateBuilder> {
   }
 };
 
+class TimestampConverter : public TypedConverter<arrow::TimestampBuilder> {
+ public:
+  Status AppendData(PyObject* seq) override {
+    int64_t val;
+    Py_ssize_t size = PySequence_Size(seq);
+    RETURN_NOT_OK(typed_builder_->Reserve(size));
+    for (int64_t i = 0; i < size; ++i) {
+      OwnedRef item(PySequence_GetItem(seq, i));
+      if (item.obj() == Py_None) {
+        typed_builder_->AppendNull();
+      } else {
+        PyDateTime_DateTime* pydatetime = reinterpret_cast<PyDateTime_DateTime*>(item.obj());
+        struct tm datetime = {0};
+        datetime.tm_year = PyDateTime_GET_YEAR(pydatetime) - 1900;
+        datetime.tm_mon = PyDateTime_GET_MONTH(pydatetime) - 1;
+        datetime.tm_mday = PyDateTime_GET_DAY(pydatetime);
+        datetime.tm_hour = PyDateTime_DATE_GET_HOUR(pydatetime);
+        datetime.tm_min = PyDateTime_DATE_GET_MINUTE(pydatetime);
+        datetime.tm_sec = PyDateTime_DATE_GET_SECOND(pydatetime);
+        int us = PyDateTime_DATE_GET_MICROSECOND(pydatetime);
+        RETURN_IF_PYERROR();
+        struct tm epoch = {0};
+        epoch.tm_year = 70;
+        epoch.tm_mday = 1;
+        // Microseconds since the epoch
+        int64_t val = lrint(difftime(mktime(&datetime), mktime(&epoch))) * 1000000 + us;
+        typed_builder_->Append(val);
+      }
+    }
+    return Status::OK();
+  }
+};
+
 class DoubleConverter : public TypedConverter<arrow::DoubleBuilder> {
  public:
   Status AppendData(PyObject* seq) override {
@@ -417,6 +456,8 @@ std::shared_ptr<SeqConverter> GetConverter(const std::shared_ptr<DataType>& type
       return std::make_shared<Int64Converter>();
     case Type::DATE:
       return std::make_shared<DateConverter>();
+    case Type::TIMESTAMP:
+      return std::make_shared<TimestampConverter>();
     case Type::DOUBLE:
       return std::make_shared<DoubleConverter>();
     case Type::STRING:

python/src/pyarrow/helpers.cc

@@ -34,7 +34,7 @@ const std::shared_ptr<Int16Type> INT16 = std::make_shared<Int16Type>();
 const std::shared_ptr<Int32Type> INT32 = std::make_shared<Int32Type>();
 const std::shared_ptr<Int64Type> INT64 = std::make_shared<Int64Type>();
 const std::shared_ptr<DateType> DATE = std::make_shared<DateType>();
-const std::shared_ptr<TimestampType> TIMESTAMP = std::make_shared<TimestampType>();
+const std::shared_ptr<TimestampType> TIMESTAMP_US = std::make_shared<TimestampType>(TimeUnit::MICRO);
 const std::shared_ptr<FloatType> FLOAT = std::make_shared<FloatType>();
 const std::shared_ptr<DoubleType> DOUBLE = std::make_shared<DoubleType>();
 const std::shared_ptr<StringType> STRING = std::make_shared<StringType>();
@@ -57,7 +57,9 @@ std::shared_ptr<DataType> GetPrimitiveType(Type::type type) {
     GET_PRIMITIVE_TYPE(UINT64, UInt64Type);
     GET_PRIMITIVE_TYPE(INT64, Int64Type);
     GET_PRIMITIVE_TYPE(DATE, DateType);
-    GET_PRIMITIVE_TYPE(TIMESTAMP, TimestampType);
+    case Type::TIMESTAMP:
+      return TIMESTAMP_US;
+      break;
     GET_PRIMITIVE_TYPE(BOOL, BooleanType);
     GET_PRIMITIVE_TYPE(FLOAT, FloatType);
     GET_PRIMITIVE_TYPE(DOUBLE, DoubleType);

python/src/pyarrow/helpers.h

@@ -39,7 +39,7 @@ extern const std::shared_ptr<arrow::Int16Type> INT16;
 extern const std::shared_ptr<arrow::Int32Type> INT32;
 extern const std::shared_ptr<arrow::Int64Type> INT64;
 extern const std::shared_ptr<arrow::DateType> DATE;
-extern const std::shared_ptr<arrow::TimestampType> TIMESTAMP;
+extern const std::shared_ptr<arrow::TimestampType> TIMESTAMP_US;
 extern const std::shared_ptr<arrow::FloatType> FLOAT;
 extern const std::shared_ptr<arrow::DoubleType> DOUBLE;
 extern const std::shared_ptr<arrow::StringType> STRING;