Skip to content

Commit 59030fe

Browse files
Licht-Twesm
authored andcommitted
ARROW-1730, ARROW-1738: [Python] Fix wrong datetime conversion
This closes [ARROW-1730](https://issues.apache.org/jira/browse/ARROW-1730) and [ARROW-1738](https://issues.apache.org/jira/browse/ARROW-1738). Author: Licht-T <licht-t@outlook.jp> Author: Wes McKinney <wes.mckinney@twosigma.com> Closes apache#1256 from Licht-T/fix-pylong-force-cast-to-date-type and squashes the following commits: b433464 [Wes McKinney] Code review comments, Add test cases for invalid values 23ad1a5 [Licht-T] Fix C++ lint issues d9cd15b [Licht-T] TST: Add tests for datetime converter 560a97e [Licht-T] BUG: Implement time unit resolution converter c1063eb [Licht-T] BUG: Fix PyLong force casting to PyDateTime/PyDate type
1 parent 2e04089 commit 59030fe

3 files changed

Lines changed: 122 additions & 5 deletions

File tree

‎cpp/src/arrow/python/builtin_convert.cc‎

Lines changed: 39 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -522,18 +522,51 @@ class UInt64Converter : public TypedConverterVisitor<UInt64Builder, UInt64Conver
522522
class DateConverter : public TypedConverterVisitor<Date64Builder, DateConverter> {
523523
public:
524524
inline Status AppendItem(const OwnedRef& item) {
525-
auto pydate = reinterpret_cast<PyDateTime_Date*>(item.obj());
526-
return typed_builder_->Append(PyDate_to_ms(pydate));
525+
int64_t t;
526+
if (PyDate_Check(item.obj())) {
527+
auto pydate = reinterpret_cast<PyDateTime_Date*>(item.obj());
528+
t = PyDate_to_ms(pydate);
529+
} else {
530+
t = static_cast<int64_t>(PyLong_AsLongLong(item.obj()));
531+
RETURN_IF_PYERROR();
532+
}
533+
return typed_builder_->Append(t);
527534
}
528535
};
529536

530537
class TimestampConverter
531538
: public TypedConverterVisitor<Date64Builder, TimestampConverter> {
532539
public:
540+
explicit TimestampConverter(TimeUnit::type unit) : unit_(unit) {}
541+
533542
inline Status AppendItem(const OwnedRef& item) {
534-
auto pydatetime = reinterpret_cast<PyDateTime_DateTime*>(item.obj());
535-
return typed_builder_->Append(PyDateTime_to_us(pydatetime));
543+
int64_t t;
544+
if (PyDateTime_Check(item.obj())) {
545+
auto pydatetime = reinterpret_cast<PyDateTime_DateTime*>(item.obj());
546+
547+
switch (unit_) {
548+
case TimeUnit::SECOND:
549+
t = PyDateTime_to_s(pydatetime);
550+
break;
551+
case TimeUnit::MILLI:
552+
t = PyDateTime_to_ms(pydatetime);
553+
break;
554+
case TimeUnit::MICRO:
555+
t = PyDateTime_to_us(pydatetime);
556+
break;
557+
case TimeUnit::NANO:
558+
t = PyDateTime_to_ns(pydatetime);
559+
break;
560+
}
561+
} else {
562+
t = static_cast<int64_t>(PyLong_AsLongLong(item.obj()));
563+
RETURN_IF_PYERROR();
564+
}
565+
return typed_builder_->Append(t);
536566
}
567+
568+
private:
569+
TimeUnit::type unit_;
537570
};
538571

539572
class DoubleConverter : public TypedConverterVisitor<DoubleBuilder, DoubleConverter> {
@@ -687,7 +720,8 @@ std::shared_ptr<SeqConverter> GetConverter(const std::shared_ptr<DataType>& type
687720
case Type::DATE64:
688721
return std::make_shared<DateConverter>();
689722
case Type::TIMESTAMP:
690-
return std::make_shared<TimestampConverter>();
723+
return std::make_shared<TimestampConverter>(
724+
static_cast<const TimestampType&>(*type).unit());
691725
case Type::DOUBLE:
692726
return std::make_shared<DoubleConverter>();
693727
case Type::BINARY:

‎cpp/src/arrow/python/util/datetime.h‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -247,12 +247,26 @@ static inline int64_t PyDate_to_ms(PyDateTime_Date* pydate) {
247247
return total_seconds * 1000;
248248
}
249249

250+
static inline int64_t PyDateTime_to_s(PyDateTime_DateTime* pydatetime) {
251+
return PyDate_to_ms(reinterpret_cast<PyDateTime_Date*>(pydatetime)) / 1000LL;
252+
}
253+
254+
static inline int64_t PyDateTime_to_ms(PyDateTime_DateTime* pydatetime) {
255+
int64_t date_ms = PyDate_to_ms(reinterpret_cast<PyDateTime_Date*>(pydatetime));
256+
int ms = PyDateTime_DATE_GET_MICROSECOND(pydatetime) / 1000;
257+
return date_ms + ms;
258+
}
259+
250260
static inline int64_t PyDateTime_to_us(PyDateTime_DateTime* pydatetime) {
251261
int64_t ms = PyDate_to_ms(reinterpret_cast<PyDateTime_Date*>(pydatetime));
252262
int us = PyDateTime_DATE_GET_MICROSECOND(pydatetime);
253263
return ms * 1000 + us;
254264
}
255265

266+
static inline int64_t PyDateTime_to_ns(PyDateTime_DateTime* pydatetime) {
267+
return PyDateTime_to_us(pydatetime) * 1000;
268+
}
269+
256270
static inline int32_t PyDate_to_days(PyDateTime_Date* pydate) {
257271
return static_cast<int32_t>(PyDate_to_ms(pydate) / 86400000LL);
258272
}

‎python/pyarrow/tests/test_convert_builtin.py‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -197,6 +197,75 @@ def test_timestamp(self):
197197
assert arr[3].as_py() == datetime.datetime(2010, 8, 13, 5,
198198
46, 57, 437699)
199199

200+
def test_timestamp_with_unit(self):
201+
data = [
202+
datetime.datetime(2007, 7, 13, 1, 23, 34, 123456),
203+
]
204+
205+
s = pa.timestamp('s')
206+
ms = pa.timestamp('ms')
207+
us = pa.timestamp('us')
208+
ns = pa.timestamp('ns')
209+
210+
arr_s = pa.array(data, type=s)
211+
assert len(arr_s) == 1
212+
assert arr_s.type == s
213+
assert arr_s[0].as_py() == datetime.datetime(2007, 7, 13, 1,
214+
23, 34, 0)
215+
216+
arr_ms = pa.array(data, type=ms)
217+
assert len(arr_ms) == 1
218+
assert arr_ms.type == ms
219+
assert arr_ms[0].as_py() == datetime.datetime(2007, 7, 13, 1,
220+
23, 34, 123000)
221+
222+
arr_us = pa.array(data, type=us)
223+
assert len(arr_us) == 1
224+
assert arr_us.type == us
225+
assert arr_us[0].as_py() == datetime.datetime(2007, 7, 13, 1,
226+
23, 34, 123456)
227+
228+
arr_ns = pa.array(data, type=ns)
229+
assert len(arr_ns) == 1
230+
assert arr_ns.type == ns
231+
assert arr_ns[0].as_py() == datetime.datetime(2007, 7, 13, 1,
232+
23, 34, 123456)
233+
234+
def test_timestamp_from_int_with_unit(self):
235+
data = [1]
236+
237+
s = pa.timestamp('s')
238+
ms = pa.timestamp('ms')
239+
us = pa.timestamp('us')
240+
ns = pa.timestamp('ns')
241+
242+
arr_s = pa.array(data, type=s)
243+
assert len(arr_s) == 1
244+
assert arr_s.type == s
245+
assert str(arr_s[0]) == "Timestamp('1970-01-01 00:00:01')"
246+
247+
arr_ms = pa.array(data, type=ms)
248+
assert len(arr_ms) == 1
249+
assert arr_ms.type == ms
250+
assert str(arr_ms[0]) == "Timestamp('1970-01-01 00:00:00.001000')"
251+
252+
arr_us = pa.array(data, type=us)
253+
assert len(arr_us) == 1
254+
assert arr_us.type == us
255+
assert str(arr_us[0]) == "Timestamp('1970-01-01 00:00:00.000001')"
256+
257+
arr_ns = pa.array(data, type=ns)
258+
assert len(arr_ns) == 1
259+
assert arr_ns.type == ns
260+
assert str(arr_ns[0]) == "Timestamp('1970-01-01 00:00:00.000000001')"
261+
262+
with pytest.raises(pa.ArrowException):
263+
class CustomClass():
264+
pass
265+
pa.array([1, CustomClass()], type=ns)
266+
pa.array([1, CustomClass()], type=pa.date32())
267+
pa.array([1, CustomClass()], type=pa.date64())
268+
200269
def test_mixed_nesting_levels(self):
201270
pa.array([1, 2, None])
202271
pa.array([[1], [2], None])

0 commit comments

Comments
 (0)