Add datetime.date conversions

Change-Id: I6627e7be236b669600301ff6f073d79cef0f3ed0

cpp/src/arrow/type.cc

@@ -88,6 +88,10 @@ std::string StructType::ToString() const {
   return s.str();
 }
 
+std::string DateType::ToString() const {
+  return std::string("date");
+}
+
 std::string UnionType::ToString() const {
   std::stringstream s;
 

cpp/src/arrow/type.h

@@ -420,7 +420,7 @@ struct ARROW_EXPORT DateType : public FixedWidthType {
   int bit_width() const override { return sizeof(c_type) * 8; }
 
   Status Accept(TypeVisitor* visitor) const override;
-  std::string ToString() const override { return name(); }
+  std::string ToString() const override;
   static std::string name() { return "date"; }
 };
 

python/pyarrow/__init__.py

@@ -45,6 +45,7 @@ from pyarrow.scalar import (ArrayValue, Scalar, NA, NAType,
 from pyarrow.schema import (null, bool_,
                             int8, int16, int32, int64,
                             uint8, uint16, uint32, uint64,
+                            timestamp, date,
                             float_, double, string,
                             list_, struct, field,
                             DataType, Field, Schema, schema)

python/pyarrow/array.pyx

@@ -212,6 +212,10 @@ cdef class UInt64Array(NumericArray):
     pass
 
 
+cdef class DateArray(NumericArray):
+    pass
+
+
 cdef class FloatArray(NumericArray):
     pass
 
@@ -239,6 +243,7 @@ cdef dict _array_classes = {
     Type_INT16: Int16Array,
     Type_INT32: Int32Array,
     Type_INT64: Int64Array,
+    Type_DATE: DateArray,
     Type_FLOAT: FloatArray,
     Type_DOUBLE: DoubleArray,
     Type_LIST: ListArray,
@@ -278,7 +283,7 @@ def from_pylist(object list_obj, DataType type=None):
     if type is None:
         check_status(pyarrow.ConvertPySequence(list_obj, &sp_array))
     else:
-        raise NotImplementedError
+        raise NotImplementedError()
 
     return box_arrow_array(sp_array)
 

python/pyarrow/includes/libarrow.pxd

@@ -39,6 +39,7 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         Type_DOUBLE" arrow::Type::DOUBLE"
 
         Type_TIMESTAMP" arrow::Type::TIMESTAMP"
+        Type_DATE" arrow::Type::DATE"
         Type_STRING" arrow::Type::STRING"
 
         Type_LIST" arrow::Type::LIST"
@@ -132,6 +133,9 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
     cdef cppclass CInt64Array" arrow::Int64Array"(CArray):
         int64_t Value(int i)
 
+    cdef cppclass CDateArray" arrow::DateArray"(CArray):
+        int64_t Value(int i)
+
     cdef cppclass CFloatArray" arrow::FloatArray"(CArray):
         float Value(int i)
 

python/pyarrow/scalar.pyx

@@ -20,6 +20,9 @@ from pyarrow.schema cimport DataType, box_data_type
 from pyarrow.compat import frombytes
 import pyarrow.schema as schema
 
+import datetime
+
+
 NA = None
 
 cdef class NAType(Scalar):
@@ -120,6 +123,13 @@ cdef class UInt64Value(ArrayValue):
         return ap.Value(self.index)
 
 
+cdef class DateValue(ArrayValue):
+
+    def as_py(self):
+        cdef CDateArray* ap = <CDateArray*> self.sp_array.get()
+        return datetime.date.fromtimestamp(ap.Value(self.index) / 1000)
+
+
 cdef class FloatValue(ArrayValue):
 
     def as_py(self):
@@ -184,6 +194,7 @@ cdef dict _scalar_classes = {
     Type_INT16: Int16Value,
     Type_INT32: Int32Value,
     Type_INT64: Int64Value,
+    Type_DATE: DateValue,
     Type_FLOAT: FloatValue,
     Type_DOUBLE: DoubleValue,
     Type_LIST: ListValue,

python/pyarrow/schema.pyx

@@ -164,6 +164,7 @@ cdef set PRIMITIVE_TYPES = set([
     Type_UINT16, Type_INT16,
     Type_UINT32, Type_INT32,
     Type_UINT64, Type_INT64,
+    Type_TIMESTAMP, Type_DATE,
     Type_FLOAT, Type_DOUBLE])
 
 def null():
@@ -196,6 +197,12 @@ def uint64():
 def int64():
     return primitive_type(Type_INT64)
 
+def timestamp():
+    return primitive_type(Type_TIMESTAMP)
+
+def date():
+    return primitive_type(Type_DATE)
+
 def float_():
     return primitive_type(Type_FLOAT)
 

python/pyarrow/tests/test_convert_builtin.py

@@ -18,6 +18,7 @@
 from pyarrow.compat import unittest
 import pyarrow
 
+import datetime
 
 class TestConvertList(unittest.TestCase):
 
@@ -70,6 +71,17 @@ class TestConvertList(unittest.TestCase):
         assert arr.null_count == 1
         assert arr.type == pyarrow.string()
 
+    def test_date(self):
+        data = [datetime.date(2000, 1, 1), None, datetime.date(1970, 1, 1), datetime.date(2040, 2, 26)]
+        arr = pyarrow.from_pylist(data)
+        assert len(arr) == 4
+        assert arr.type == pyarrow.date()
+        assert arr.null_count == 1
+        assert arr[0].as_py() == datetime.date(2000, 1, 1)
+        assert arr[1].as_py() is None
+        assert arr[2].as_py() == datetime.date(1970, 1, 1)
+        assert arr[3].as_py() == datetime.date(2040, 2, 26)
+
     def test_mixed_nesting_levels(self):
         pyarrow.from_pylist([1, 2, None])
         pyarrow.from_pylist([[1], [2], None])

python/src/pyarrow/adapters/builtin.cc

@@ -16,6 +16,7 @@
 // under the License.
 
 #include <Python.h>
+#include <datetime.h>
 #include <sstream>
 
 #include "pyarrow/adapters/builtin.h"
@@ -55,6 +56,7 @@ class ScalarVisitor {
       none_count_(0),
       bool_count_(0),
       int_count_(0),
+      date_count_(0),
       float_count_(0),
       string_count_(0) {}
 
@@ -68,6 +70,8 @@ class ScalarVisitor {
       ++float_count_;
     } else if (IsPyInteger(obj)) {
       ++int_count_;
+    } else if (PyDate_Check(obj)) {
+      ++date_count_;
     } else if (IsPyBaseString(obj)) {
       ++string_count_;
     } else {
@@ -82,6 +86,8 @@ class ScalarVisitor {
     } else if (int_count_) {
       // TODO(wesm): tighter type later
       return INT64;
+    } else if (date_count_) {
+      return DATE;
     } else if (bool_count_) {
       return BOOL;
     } else if (string_count_) {
@@ -100,6 +106,7 @@ class ScalarVisitor {
   int64_t none_count_;
   int64_t bool_count_;
   int64_t int_count_;
+  int64_t date_count_;
   int64_t float_count_;
   int64_t string_count_;
 
@@ -297,6 +304,35 @@ class Int64Converter : public TypedConverter<arrow::Int64Builder> {
   }
 };
 
+class DateConverter : public TypedConverter<arrow::DateBuilder> {
+ 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_Date* pydate = reinterpret_cast<PyDateTime_Date*>(item.obj());
+        struct tm date = {0};
+        date.tm_year = PyDateTime_GET_YEAR(pydate) - 1900;
+        date.tm_mon = PyDateTime_GET_MONTH(pydate) - 1;
+        date.tm_mday = PyDateTime_GET_DAY(pydate);
+        RETURN_IF_PYERROR();
+        struct tm epoch = {0};
+        epoch.tm_year = 70;
+        epoch.tm_mday = 1;
+        // Milliseconds since the epoch
+        int64_t val = lrint(difftime(mktime(&date), mktime(&epoch)) * 1000);
+        typed_builder_->Append(val);
+      }
+    }
+    return Status::OK();
+  }
+};
+
 class DoubleConverter : public TypedConverter<arrow::DoubleBuilder> {
  public:
   Status AppendData(PyObject* seq) override {
@@ -379,6 +415,8 @@ std::shared_ptr<SeqConverter> GetConverter(const std::shared_ptr<DataType>& type
       return std::make_shared<BoolConverter>();
     case Type::INT64:
       return std::make_shared<Int64Converter>();
+    case Type::DATE:
+      return std::make_shared<DateConverter>();
     case Type::DOUBLE:
       return std::make_shared<DoubleConverter>();
     case Type::STRING:
@@ -409,6 +447,7 @@ Status ListConverter::Init(const std::shared_ptr<ArrayBuilder>& builder) {
 Status ConvertPySequence(PyObject* obj, std::shared_ptr<arrow::Array>* out) {
   std::shared_ptr<DataType> type;
   int64_t size;
+  PyDateTime_IMPORT;
   RETURN_NOT_OK(InferArrowType(obj, &size, &type));
 
   // Handle NA / NullType case

python/src/pyarrow/helpers.cc

@@ -33,6 +33,8 @@ const std::shared_ptr<Int8Type> INT8 = std::make_shared<Int8Type>();
 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<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>();
@@ -54,6 +56,8 @@ std::shared_ptr<DataType> GetPrimitiveType(Type::type type) {
     GET_PRIMITIVE_TYPE(INT32, Int32Type);
     GET_PRIMITIVE_TYPE(UINT64, UInt64Type);
     GET_PRIMITIVE_TYPE(INT64, Int64Type);
+    GET_PRIMITIVE_TYPE(DATE, DateType);
+    GET_PRIMITIVE_TYPE(TIMESTAMP, TimestampType);
     GET_PRIMITIVE_TYPE(BOOL, BooleanType);
     GET_PRIMITIVE_TYPE(FLOAT, FloatType);
     GET_PRIMITIVE_TYPE(DOUBLE, DoubleType);

python/src/pyarrow/helpers.h

@@ -38,6 +38,8 @@ extern const std::shared_ptr<arrow::Int8Type> INT8;
 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::FloatType> FLOAT;
 extern const std::shared_ptr<arrow::DoubleType> DOUBLE;
 extern const std::shared_ptr<arrow::StringType> STRING;