add timestamp_tz_ns support

Author: evertlammerts

src/duckdb_py/arrow/pyarrow_filter_pushdown.cpp

@@ -112,6 +112,11 @@ py::object MakePyArrowScalar(const Value &constant, const string &timezone_confi
 		py::handle date_type = import_cache.pyarrow.timestamp();
 		return dataset_scalar(scalar(converted_value, date_type(time_unit_string, py::arg("tz") = timezone_config)));
 	}
+	case LogicalTypeId::TIMESTAMP_TZ_NS: {
+		py::handle date_type = import_cache.pyarrow.timestamp();
+		auto converted_value = Timestamp::GetEpochNanoSeconds(timestamp_t(constant.GetValue<int64_t>()));
+		return dataset_scalar(scalar(converted_value, date_type("ns", py::arg("tz") = timezone_config)));
+	}
 	case LogicalTypeId::UTINYINT: {
 		py::handle integer_type = import_cache.pyarrow.uint8();
 		return dataset_scalar(scalar(constant.GetValue<uint8_t>(), integer_type()));

src/duckdb_py/native/python_objects.cpp

@@ -437,6 +437,7 @@ static bool KeyIsHashable(const LogicalType &type) {
 	case LogicalTypeId::TIMESTAMP_NS:
 	case LogicalTypeId::TIMESTAMP_SEC:
 	case LogicalTypeId::TIMESTAMP_TZ:
+	case LogicalTypeId::TIMESTAMP_TZ_NS:
 	case LogicalTypeId::TIME_TZ:
 	case LogicalTypeId::TIME:
 	case LogicalTypeId::DATE:
@@ -518,7 +519,8 @@ py::object PythonObject::FromValue(const Value &val, const LogicalType &type,
 	case LogicalTypeId::TIMESTAMP_MS:
 	case LogicalTypeId::TIMESTAMP_NS:
 	case LogicalTypeId::TIMESTAMP_SEC:
-	case LogicalTypeId::TIMESTAMP_TZ: {
+	case LogicalTypeId::TIMESTAMP_TZ:
+	case LogicalTypeId::TIMESTAMP_TZ_NS: {
 		D_ASSERT(type.InternalType() == PhysicalType::INT64);
 		auto timestamp = val.GetValueUnsafe<timestamp_t>();
 
@@ -532,7 +534,7 @@ py::object PythonObject::FromValue(const Value &val, const LogicalType &type,
 
 		if (type.id() == LogicalTypeId::TIMESTAMP_MS) {
 			timestamp = Timestamp::FromEpochMs(timestamp.value);
-		} else if (type.id() == LogicalTypeId::TIMESTAMP_NS) {
+		} else if (type.id() == LogicalTypeId::TIMESTAMP_NS || type.id() == LogicalTypeId::TIMESTAMP_TZ_NS) {
 			timestamp = Timestamp::FromEpochNanoSeconds(timestamp.value);
 		} else if (type.id() == LogicalTypeId::TIMESTAMP_SEC) {
 			timestamp = Timestamp::FromEpochSeconds(timestamp.value);
@@ -555,7 +557,7 @@ py::object PythonObject::FromValue(const Value &val, const LogicalType &type,
 			// Failed to convert, fall back to str
 			return py::str(val.ToString());
 		}
-		if (type.id() == LogicalTypeId::TIMESTAMP_TZ) {
+		if (type.id() == LogicalTypeId::TIMESTAMP_TZ || type.id() == LogicalTypeId::TIMESTAMP_TZ_NS) {
 			// We have to add the timezone info
 			auto tz_utc = import_cache.pytz.timezone()("UTC");
 			auto timestamp_utc = tz_utc.attr("localize")(py_timestamp);

src/duckdb_py/numpy/array_wrapper.cpp

@@ -654,6 +654,7 @@ void ArrayWrapper::Append(idx_t current_offset, Vector &input, idx_t source_size
 		break;
 	case LogicalTypeId::TIMESTAMP:
 	case LogicalTypeId::TIMESTAMP_TZ:
+	case LogicalTypeId::TIMESTAMP_TZ_NS:
 	case LogicalTypeId::TIMESTAMP_SEC:
 	case LogicalTypeId::TIMESTAMP_MS:
 	case LogicalTypeId::TIMESTAMP_NS:

src/duckdb_py/numpy/raw_array_wrapper.cpp

@@ -46,6 +46,7 @@ static idx_t GetNumpyTypeWidth(const LogicalType &type) {
 	case LogicalTypeId::DATE:
 	case LogicalTypeId::INTERVAL:
 	case LogicalTypeId::TIMESTAMP_TZ:
+	case LogicalTypeId::TIMESTAMP_TZ_NS:
 		return sizeof(int64_t);
 	case LogicalTypeId::TIME:
 	case LogicalTypeId::TIME_NS:
@@ -102,6 +103,7 @@ string RawArrayWrapper::DuckDBToNumpyDtype(const LogicalType &type) {
 		return "datetime64[us]";
 	case LogicalTypeId::TIMESTAMP_TZ:
 		return "datetime64[us]";
+	case LogicalTypeId::TIMESTAMP_TZ_NS:
 	case LogicalTypeId::TIMESTAMP_NS:
 		return "datetime64[ns]";
 	case LogicalTypeId::TIMESTAMP_MS:

tests/fast/arrow/test_timestamp_timezone.py

@@ -37,17 +37,17 @@ def test_timestamp_timezone_overflow(self, duckdb_cursor):
             with pytest.raises(duckdb.ConversionException, match="Could not convert"):
                 duckdb.from_arrow(arrow_table).execute().fetchall()
 
-    def test_timestamp_tz_to_arrow(self, duckdb_cursor):
-        precisions = ["us", "s", "ns", "ms"]
+    @pytest.mark.parametrize("precision", ["us", "s", "ns", "ms"])
+    def test_timestamp_tz_to_arrow(self, duckdb_cursor, precision):
         current_time = datetime.datetime(2017, 11, 28, 23, 55, 59)
         con = duckdb.connect()
-        for precision in precisions:
-            for timezone in timezones:
-                con.execute("SET TimeZone = '" + timezone + "'")
-                arrow_table = generate_table(current_time, precision, timezone)
-                res = con.from_arrow(arrow_table).to_arrow_table()
-                assert res[0].type == pa.timestamp("us", tz=timezone)
-                assert res == generate_table(current_time, "us", timezone)
+        expected_precision = "ns" if precision == "ns" else "us"
+        for timezone in timezones:
+            con.execute("SET TimeZone = '" + timezone + "'")
+            arrow_table = generate_table(current_time, precision, timezone)
+            res = con.from_arrow(arrow_table).to_arrow_table()
+            assert res[0].type == pa.timestamp(expected_precision, tz=timezone)
+            assert res == generate_table(current_time, expected_precision, timezone)
 
     def test_timestamp_tz_with_null(self, duckdb_cursor):
         con = duckdb.connect()