Skip to content

Commit e7379e7

Browse files
AEsir777AEsir777
committed
finish creating datatype timestamp_ns/ms/sec/tz and adding corresponding apis
1 parent 996429d commit e7379e7

10 files changed

Lines changed: 343 additions & 127 deletions

File tree

‎src_cpp/include/numpy/numpy_type.h‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,13 +23,16 @@ enum class NumpyNullableType : uint8_t {
2323
FLOAT_16, //! float16, half
2424
FLOAT_32, //! float32, single
2525
FLOAT_64, //! float64, float_, double
26-
DATETIME_US, //! datetime64[us], <M8[us]
26+
DATETIME_S, //! datetime64[s], <M8[s]
27+
DATETIME_MS, //! datetime64[ms], <M8[ms]
2728
DATETIME_NS, //! datetime64[ns], <M8[ns]
29+
DATETIME_US, //! datetime64[us], <M8[us]
2830
TIMEDELTA, //! timedelta64[D], timedelta64
2931
};
3032

3133
struct NumpyType {
3234
NumpyNullableType type;
35+
bool hasTimezone = false;
3336
};
3437

3538
struct NumpyTypeUtils {

‎src_cpp/numpy/numpy_scan.cpp‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -83,16 +83,15 @@ void NumpyScan::scan(PandasColumnBindData* bindData, uint64_t count, uint64_t of
8383
ScanNumpyFpColumn<double>(
8484
reinterpret_cast<const double*>(array.data()), count, offset, outputVector);
8585
break;
86+
case NumpyNullableType::DATETIME_S:
87+
case NumpyNullableType::DATETIME_MS:
8688
case NumpyNullableType::DATETIME_NS:
8789
case NumpyNullableType::DATETIME_US: {
8890
auto sourceData = reinterpret_cast<const int64_t*>(array.data());
8991
auto dstData = reinterpret_cast<timestamp_t*>(outputVector->getData());
90-
auto timestampCastFunc = bindData->npType.type == NumpyNullableType::DATETIME_NS ?
91-
Timestamp::fromEpochNanoSeconds :
92-
Timestamp::fromEpochMicroSeconds;
9392
for (auto i = 0u; i < count; i++) {
9493
auto pos = offset + i;
95-
dstData[i] = timestampCastFunc(sourceData[pos]);
94+
dstData[i].value = sourceData[pos];
9695
outputVector->setNull(i, false /* isNull */);
9796
}
9897
break;

‎src_cpp/numpy/numpy_type.cpp‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,18 @@ namespace kuzu {
66

77
using namespace kuzu::common;
88

9+
static bool isDateTime(NumpyNullableType type) {
10+
switch(type) {
11+
case NumpyNullableType::DATETIME_US:
12+
case NumpyNullableType::DATETIME_S:
13+
case NumpyNullableType::DATETIME_NS:
14+
case NumpyNullableType::DATETIME_MS:
15+
return true;
16+
default:
17+
return false;
18+
}
19+
}
20+
921
static NumpyNullableType convertNumpyTypeInternal(const std::string& colTypeStr) {
1022
if (colTypeStr == "bool" || colTypeStr == "boolean") {
1123
return NumpyNullableType::BOOL;
@@ -53,6 +65,12 @@ static NumpyNullableType convertNumpyTypeInternal(const std::string& colTypeStr)
5365
if (colTypeStr.starts_with("datetime64[ns")) {
5466
return NumpyNullableType::DATETIME_NS;
5567
}
68+
if (colTypeStr.starts_with("datetime64[ms")) {
69+
return NumpyNullableType::DATETIME_MS;
70+
}
71+
if (colTypeStr.starts_with("datetime64[s")) {
72+
return NumpyNullableType::DATETIME_S;
73+
}
5674
// LCOV_EXCL_START
5775
// Legacy datetime type indicators
5876
if (colTypeStr.starts_with("<M8[us")) {
@@ -69,6 +87,12 @@ NumpyType NumpyTypeUtils::convertNumpyType(const py::handle& colType) {
6987
auto colTypeStr = std::string(py::str(colType));
7088
NumpyType resultNPType;
7189
resultNPType.type = convertNumpyTypeInternal(colTypeStr);
90+
if (isDateTime(resultNPType.type)) {
91+
if (hasattr(colType, "tz")) {
92+
resultNPType.hasTimezone = true;
93+
}
94+
}
95+
7296
return resultNPType;
7397
}
7498

@@ -101,8 +125,17 @@ std::unique_ptr<LogicalType> NumpyTypeUtils::numpyToLogicalType(const NumpyType&
101125
case NumpyNullableType::TIMEDELTA:
102126
return std::make_unique<LogicalType>(LogicalTypeID::INTERVAL);
103127
case NumpyNullableType::DATETIME_US:
104-
case NumpyNullableType::DATETIME_NS:
128+
// because we currently don't support object, only US with timezone is allowed
129+
if (col_type.hasTimezone) {
130+
return std::make_unique<LogicalType>(LogicalTypeID::TIMESTAMP_TZ);
131+
}
105132
return std::make_unique<LogicalType>(LogicalTypeID::TIMESTAMP);
133+
case NumpyNullableType::DATETIME_NS:
134+
return std::make_unique<LogicalType>(LogicalTypeID::TIMESTAMP_NS);
135+
case NumpyNullableType::DATETIME_MS:
136+
return std::make_unique<LogicalType>(LogicalTypeID::TIMESTAMP_MS);
137+
case NumpyNullableType::DATETIME_S:
138+
return std::make_unique<LogicalType>(LogicalTypeID::TIMESTAMP_SEC);
106139
default:
107140
KU_UNREACHABLE;
108141
}

‎src_cpp/py_query_result.cpp‎

Lines changed: 34 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,10 @@
1212

1313
using namespace kuzu::common;
1414

15+
#define PyDateTimeTZ_FromDateAndTime(year, month, day, hour, min, sec, usec, timezone) \
16+
PyDateTimeAPI->DateTime_FromDateAndTime(year, month, day, hour, \
17+
min, sec, usec, timezone, PyDateTimeAPI->DateTimeType)
18+
1519
void PyQueryResult::initialize(py::handle& m) {
1620
py::class_<PyQueryResult>(m, "result")
1721
.def("hasNext", &PyQueryResult::hasNext)
@@ -67,6 +71,17 @@ void PyQueryResult::close() {
6771
queryResult.reset();
6872
}
6973

74+
static py::object converTimestampToPyObject(timestamp_t& timestamp) {
75+
int32_t year, month, day, hour, min, sec, micros;
76+
date_t date;
77+
dtime_t time;
78+
Timestamp::convert(timestamp, date, time);
79+
Date::convert(date, year, month, day);
80+
Time::Convert(time, hour, min, sec, micros);
81+
return py::cast<py::object>(
82+
PyDateTime_FromDateAndTime(year, month, day, hour, min, sec, micros));
83+
}
84+
7085
py::object PyQueryResult::convertValueToPyObject(const Value& value) {
7186
if (value.isNull()) {
7287
return py::none();
@@ -130,14 +145,31 @@ py::object PyQueryResult::convertValueToPyObject(const Value& value) {
130145
}
131146
case LogicalTypeID::TIMESTAMP: {
132147
auto timestampVal = value.getValue<timestamp_t>();
148+
return converTimestampToPyObject(timestampVal);
149+
}
150+
case LogicalTypeID::TIMESTAMP_TZ: {
151+
auto timestampVal = value.getValue<timestamp_tz_t>();
133152
int32_t year, month, day, hour, min, sec, micros;
134153
date_t date;
135154
dtime_t time;
136155
Timestamp::convert(timestampVal, date, time);
137156
Date::convert(date, year, month, day);
138157
Time::Convert(time, hour, min, sec, micros);
139-
return py::cast<py::object>(
140-
PyDateTime_FromDateAndTime(year, month, day, hour, min, sec, micros));
158+
159+
return py::cast<py::object>(PyDateTimeTZ_FromDateAndTime(year, month, day, hour, min, sec,
160+
micros, PyDateTime_TimeZone_UTC));
161+
}
162+
case LogicalTypeID::TIMESTAMP_NS: {
163+
timestamp_t timestampVal = Timestamp::fromEpochNanoSeconds(value.getValue<timestamp_ns_t>().value);
164+
return converTimestampToPyObject(timestampVal);
165+
}
166+
case LogicalTypeID::TIMESTAMP_MS: {
167+
timestamp_t timestampVal = Timestamp::fromEpochMilliSeconds(value.getValue<timestamp_ms_t>().value);
168+
return converTimestampToPyObject(timestampVal);
169+
}
170+
case LogicalTypeID::TIMESTAMP_SEC: {
171+
timestamp_t timestampVal = Timestamp::fromEpochSeconds(value.getValue<timestamp_sec_t>().value);
172+
return converTimestampToPyObject(timestampVal);
141173
}
142174
case LogicalTypeID::INTERVAL: {
143175
auto intervalVal = value.getValue<interval_t>();

‎src_cpp/py_query_result_converter.cpp‎

Lines changed: 28 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -54,11 +54,22 @@ void NPArrayWrapper::appendElement(Value* value) {
5454
} break;
5555
case LogicalTypeID::DATE: {
5656
((int64_t*)dataBuffer)[numElements] =
57-
Date::getEpochNanoSeconds(value->getValue<date_t>());
57+
Date::getEpochNanoSeconds(value->getValue<date_t>()) / Interval::NANOS_PER_MICRO;
5858
} break;
5959
case LogicalTypeID::TIMESTAMP: {
60-
((int64_t*)dataBuffer)[numElements] =
61-
Timestamp::getEpochNanoSeconds(value->getValue<timestamp_t>());
60+
((int64_t*)dataBuffer)[numElements] = value->getValue<timestamp_t>().value;
61+
} break;
62+
case LogicalTypeID::TIMESTAMP_TZ: {
63+
((int64_t*)dataBuffer)[numElements] = value->getValue<timestamp_tz_t>().value;
64+
} break;
65+
case LogicalTypeID::TIMESTAMP_NS: {
66+
((int64_t*)dataBuffer)[numElements] = value->getValue<timestamp_ns_t>().value;
67+
} break;
68+
case LogicalTypeID::TIMESTAMP_MS: {
69+
((int64_t*)dataBuffer)[numElements] = value->getValue<timestamp_ms_t>().value;
70+
} break;
71+
case LogicalTypeID::TIMESTAMP_SEC: {
72+
((int64_t*)dataBuffer)[numElements] = value->getValue<timestamp_sec_t>().value;
6273
} break;
6374
case LogicalTypeID::INTERVAL: {
6475
((int64_t*)dataBuffer)[numElements] =
@@ -134,17 +145,24 @@ py::dtype NPArrayWrapper::convertToArrayType(const LogicalType& type) {
134145
} break;
135146
case LogicalTypeID::BOOL: {
136147
dtype = "bool";
137-
break;
138-
}
148+
} break;
139149
case LogicalTypeID::DATE:
150+
case LogicalTypeID::TIMESTAMP_TZ:
140151
case LogicalTypeID::TIMESTAMP: {
152+
dtype = "datetime64[us]";
153+
} break;
154+
case LogicalTypeID::TIMESTAMP_NS: {
141155
dtype = "datetime64[ns]";
142-
break;
143-
}
156+
} break;
157+
case LogicalTypeID::TIMESTAMP_MS: {
158+
dtype = "datetime64[ms]";
159+
} break;
160+
case LogicalTypeID::TIMESTAMP_SEC: {
161+
dtype = "datetime64[s]";
162+
} break;
144163
case LogicalTypeID::INTERVAL: {
145164
dtype = "timedelta64[ns]";
146-
break;
147-
}
165+
} break;
148166
case LogicalTypeID::UNION:
149167
case LogicalTypeID::BLOB:
150168
case LogicalTypeID::STRUCT:
@@ -156,8 +174,7 @@ py::dtype NPArrayWrapper::convertToArrayType(const LogicalType& type) {
156174
case LogicalTypeID::MAP:
157175
case LogicalTypeID::RECURSIVE_REL: {
158176
dtype = "object";
159-
break;
160-
}
177+
} break;
161178
default: {
162179
KU_UNREACHABLE;
163180
}

‎src_py/types.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,10 @@ class Type(Enum):
2222
FLOAT = "FLOAT"
2323
DATE = "DATE"
2424
TIMESTAMP = "TIMESTAMP"
25+
TIMSTAMP_TZ = "TIMESTAMP_TZ"
26+
TIMESTAMP_NS = "TIMESTAMP_NS"
27+
TIMESTAMP_MS = "TIMESTAMP_MS"
28+
TIMESTAMP_SEC = "TIMESTAMP_SEC"
2529
INTERVAL = "INTERVAL"
2630
FIXED_LIST = "FIXED_LIST"
2731
INTERNAL_ID = "INTERNAL_ID"

0 commit comments

Comments
 (0)