diff --git a/benchmarks/requirements-bench.txt b/benchmarks/requirements-bench.txt index 8a9f49e68..50a6edee9 100644 --- a/benchmarks/requirements-bench.txt +++ b/benchmarks/requirements-bench.txt @@ -18,6 +18,5 @@ pygments==2.20.0 pytest==9.1.1 pytest-codspeed==5.0.3 python-dateutil==2.9.0.post0 -pytz==2026.2 rich==15.0.0 six==1.17.0 diff --git a/benchmarks/test_fetch_perf.py b/benchmarks/test_fetch_perf.py index 1b07dfb74..afb60ec1b 100644 --- a/benchmarks/test_fetch_perf.py +++ b/benchmarks/test_fetch_perf.py @@ -67,7 +67,7 @@ def run() -> None: # small-N *_gate variants: at range(2048) the measured region is ~60% SQL compile + engine, ~40% fetch, so these -# catch a fixed-cost regression (not a pure per-row one). Plus expensive scalar types (timestamptz pytz-per-row, +# catch a fixed-cost regression (not a pure per-row one). Plus expensive scalar types (timestamptz zoneinfo per row, # blob, null-heavy), a heterogeneous per-cell-dispatch row, and the batched fetchmany loop. diff --git a/duckdb/_relation.py b/duckdb/_relation.py index 7088f6eca..4e11a02ae 100644 --- a/duckdb/_relation.py +++ b/duckdb/_relation.py @@ -340,8 +340,8 @@ def _assemble(ctes: list[tuple[str, str]], body: str) -> str: def _stub(names: list[str], types: list[DuckDBPyType]) -> str | None: """An empty SELECT of typed NULLs standing in for these columns, or None where one has a type it cannot write. - A type's text leaves out its collation, which a column of its own carries as COLLATE; a collation on a type inside - a list or struct has no such spelling. + A 2.0 engine's type text leaves a collation out, so a column of its own carries it as COLLATE, and a collation on a + type inside a list or struct cannot be written from the type's text there. """ columns: list[str] = [] for name, type_ in zip(names, types, strict=True): diff --git a/duckdb/_sql_keywords.py b/duckdb/_sql_keywords.py index abbba9bdc..e7429e886 100644 --- a/duckdb/_sql_keywords.py +++ b/duckdb/_sql_keywords.py @@ -33,6 +33,7 @@ "attach": "unreserved", "attribute": "unreserved", "authorization": "type_function", + "auto": "unreserved", "backward": "unreserved", "before": "unreserved", "begin": "unreserved", @@ -106,6 +107,7 @@ "deferrable": "reserved", "deferred": "unreserved", "definer": "unreserved", + "defines": "unreserved", "delete": "unreserved", "delimiter": "unreserved", "delimiters": "unreserved", @@ -240,6 +242,7 @@ "map": "column_name/type_function", "mapping": "unreserved", "match": "unreserved", + "match_recognize": "reserved", "matched": "unreserved", "materialized": "unreserved", "maxvalue": "unreserved", @@ -415,6 +418,7 @@ "strip": "unreserved", "struct": "column_name/type_function", "subscription": "unreserved", + "subset": "unreserved", "substring": "column_name", "summarize": "reserved", "symmetric": "reserved", diff --git a/duckdb/experimental/spark/sql/column.py b/duckdb/experimental/spark/sql/column.py index e1248482e..3f8eb0434 100644 --- a/duckdb/experimental/spark/sql/column.py +++ b/duckdb/experimental/spark/sql/column.py @@ -420,7 +420,7 @@ def _truncate(value: Expression, target: DuckDBPyType) -> Expression: Spark in ANSI mode does. """ source_type = FunctionExpression("typeof", value) - is_string = source_type == ConstantExpression("VARCHAR") + is_string = _is_string(value) text = value.cast(VARCHAR) is_number = FunctionExpression( "regexp_full_match", text, ConstantExpression(r"\s*[+-]?([0-9]+\.?[0-9]*|\.[0-9]+)\s*") @@ -451,10 +451,18 @@ def _truncate(value: Expression, target: DuckDBPyType) -> Expression: ) +def _is_string(value: Expression) -> Expression: + """Whether ``value`` is a string at run time; a collated string's type reads VARCHAR('NOCASE') on 2.1 engines.""" + return FunctionExpression( + "regexp_full_match", FunctionExpression("typeof", value), ConstantExpression(r"VARCHAR(\(.*\))?") + ) + + def _to_binary(value: Expression) -> Expression: """Cast to binary like Spark, which takes a string's UTF-8 bytes where DuckDB reads escape sequences.""" - is_string = FunctionExpression("typeof", value) == ConstantExpression("VARCHAR") - return CaseExpression(is_string, FunctionExpression("encode", value.cast(VARCHAR))).otherwise(value.cast(BLOB)) + return CaseExpression(_is_string(value), FunctionExpression("encode", value.cast(VARCHAR))).otherwise( + value.cast(BLOB) + ) def _children(dtype: DuckDBPyType) -> list[DuckDBPyType]: diff --git a/duckdb/experimental/spark/sql/functions.py b/duckdb/experimental/spark/sql/functions.py index a2fcc70aa..96509410c 100644 --- a/duckdb/experimental/spark/sql/functions.py +++ b/duckdb/experimental/spark/sql/functions.py @@ -3094,19 +3094,11 @@ def first(col: "ColumnOrName", ignorenulls: bool = False) -> Column: Examples: -------- - >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5), ("Alice", None)], ("name", "age")) - >>> df = df.orderBy(df.age) - >>> df.groupby("name").agg(first("age")).orderBy("name").show() - +-----+----------+ - | name|first(age)| - +-----+----------+ - |Alice| NULL| - | Bob| 5| - +-----+----------+ - - Now, to ignore any nulls we needs to set ``ignorenulls`` to `True` + The row a group's first value comes from is not set by the frame's order, so a group with a null and a value can + give either. ``ignorenulls`` is not supported here and raises. - >>> df.groupby("name").agg(first("age", ignorenulls=True)).orderBy("name").show() + >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5)], ("name", "age")) + >>> df.groupby("name").agg(first("age")).orderBy("name").show() +-----+----------+ | name|first(age)| +-----+----------+ @@ -3149,19 +3141,11 @@ def last(col: "ColumnOrName", ignorenulls: bool = False) -> Column: Examples: -------- - >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5), ("Alice", None)], ("name", "age")) - >>> df = df.orderBy(df.age.desc()) - >>> df.groupby("name").agg(last("age")).orderBy("name").show() - +-----+---------+ - | name|last(age)| - +-----+---------+ - |Alice| NULL| - | Bob| 5| - +-----+---------+ - - Now, to ignore any nulls we needs to set ``ignorenulls`` to `True` + The row a group's last value comes from is not set by the frame's order, so a group with a null and a value can + give either. ``ignorenulls`` is not supported here and raises. - >>> df.groupby("name").agg(last("age", ignorenulls=True)).orderBy("name").show() + >>> df = spark.createDataFrame([("Alice", 2), ("Bob", 5)], ("name", "age")) + >>> df.groupby("name").agg(last("age")).orderBy("name").show() +-----+---------+ | name|last(age)| +-----+---------+ diff --git a/external/duckdb b/external/duckdb index a2353736a..ad2d8e9ac 160000 --- a/external/duckdb +++ b/external/duckdb @@ -1 +1 @@ -Subproject commit a2353736a58acd17f0ba8fae400068ffbf9d465e +Subproject commit ad2d8e9acc9e570855c1a5a3143a6d843160a795 diff --git a/pyproject.toml b/pyproject.toml index f4d9bd504..eb7caf963 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -291,7 +291,6 @@ test = [ # dependencies used for running tests "py4j", "pyotp", "pyspark", - "pytz", "requests", "urllib3", "fsspec>=2022.11.0; sys_platform != 'win32' or platform_machine != 'ARM64'", @@ -316,7 +315,6 @@ bench = [ # Pinned deps for the benchmark suite (see benchmarks/README.md). Mini "pytest", "pytest_codspeed", "polars>=1.33.0", - "pytz", "numpy<2; ( sys_platform != 'win32' or platform_machine != 'ARM64' ) and python_version < '3.12'", "numpy>=2; ( sys_platform != 'win32' or platform_machine != 'ARM64' ) and python_version >= '3.12'", "numpy>=2.3; sys_platform == 'win32' and platform_machine == 'ARM64'", @@ -331,7 +329,6 @@ scripts = [ # dependencies used for running scripts "pandas", "polars>=1.33.0", "pyarrow; sys_platform != 'win32' or platform_machine != 'ARM64'", - "pytz" ] pypi = [ # dependencies used by the pypi cleanup script "pyotp>=2.9.0", diff --git a/scripts/cache_data.json b/scripts/cache_data.json index a822fa81d..97e26642f 100644 --- a/scripts/cache_data.json +++ b/scripts/cache_data.json @@ -258,7 +258,9 @@ "type": "attribute", "full_path": "datetime.timezone", "name": "timezone", - "children": [] + "children": [ + "datetime.timezone.utc" + ] }, "decimal": { "type": "module", @@ -582,20 +584,6 @@ "name": "Value", "children": [] }, - "pytz": { - "type": "module", - "full_path": "pytz", - "name": "pytz", - "children": [ - "pytz.timezone" - ] - }, - "pytz.timezone": { - "type": "attribute", - "full_path": "pytz.timezone", - "name": "timezone", - "children": [] - }, "types": { "type": "module", "full_path": "types", @@ -858,5 +846,25 @@ "full_path": "pandas.arrays.StringArray", "name": "StringArray", "children": [] + }, + "datetime.timezone.utc": { + "type": "attribute", + "full_path": "datetime.timezone.utc", + "name": "utc", + "children": [] + }, + "zoneinfo": { + "type": "module", + "full_path": "zoneinfo", + "name": "zoneinfo", + "children": [ + "zoneinfo.ZoneInfo" + ] + }, + "zoneinfo.ZoneInfo": { + "type": "attribute", + "full_path": "zoneinfo.ZoneInfo", + "name": "ZoneInfo", + "children": [] } } \ No newline at end of file diff --git a/scripts/imports.py b/scripts/imports.py index 477e8a8f9..b21ab9023 100644 --- a/scripts/imports.py +++ b/scripts/imports.py @@ -60,6 +60,7 @@ datetime.datetime.max datetime.datetime.combine datetime.timezone +datetime.timezone.utc import decimal @@ -122,9 +123,9 @@ duckdb.filesystem.ModifiedMemoryFileSystem duckdb.Value -import pytz +import zoneinfo -pytz.timezone +zoneinfo.ZoneInfo import types diff --git a/src/include/duckdb_python/import_cache/modules/datetime_module.hpp b/src/include/duckdb_python/import_cache/modules/datetime_module.hpp index c05f77d4f..c3e4426a9 100644 --- a/src/include/duckdb_python/import_cache/modules/datetime_module.hpp +++ b/src/include/duckdb_python/import_cache/modules/datetime_module.hpp @@ -20,6 +20,18 @@ namespace duckdb { +struct DatetimeTimezoneCacheItem : public PythonImportCacheItem { + +public: + DatetimeTimezoneCacheItem(optional_ptr parent) + : PythonImportCacheItem("timezone", parent), utc("utc", this) { + } + ~DatetimeTimezoneCacheItem() override { + } + + PythonImportCacheItem utc; +}; + struct DatetimeDatetimeCacheItem : public PythonImportCacheItem { public: @@ -55,7 +67,7 @@ struct DatetimeCacheItem : public PythonImportCacheItem { public: DatetimeCacheItem() : PythonImportCacheItem("datetime"), date(this), time("time", this), timedelta("timedelta", this), - datetime(this), timezone("timezone", this) { + datetime(this), timezone(this) { } ~DatetimeCacheItem() override { } @@ -64,7 +76,7 @@ struct DatetimeCacheItem : public PythonImportCacheItem { PythonImportCacheItem time; PythonImportCacheItem timedelta; DatetimeDatetimeCacheItem datetime; - PythonImportCacheItem timezone; + DatetimeTimezoneCacheItem timezone; }; } // namespace duckdb diff --git a/src/include/duckdb_python/import_cache/modules/pytz_module.hpp b/src/include/duckdb_python/import_cache/modules/zoneinfo_module.hpp similarity index 67% rename from src/include/duckdb_python/import_cache/modules/pytz_module.hpp rename to src/include/duckdb_python/import_cache/modules/zoneinfo_module.hpp index ccdde0c89..5671433dd 100644 --- a/src/include/duckdb_python/import_cache/modules/pytz_module.hpp +++ b/src/include/duckdb_python/import_cache/modules/zoneinfo_module.hpp @@ -2,7 +2,7 @@ //===----------------------------------------------------------------------===// // DuckDB // -// duckdb_python/import_cache/modules/pytz_module.hpp +// duckdb_python/import_cache/modules/zoneinfo_module.hpp // // //===----------------------------------------------------------------------===// @@ -20,18 +20,18 @@ namespace duckdb { -struct PytzCacheItem : public PythonImportCacheItem { +struct ZoneinfoCacheItem : public PythonImportCacheItem { public: - static constexpr const char *Name = "pytz"; + static constexpr const char *Name = "zoneinfo"; public: - PytzCacheItem() : PythonImportCacheItem("pytz"), timezone("timezone", this) { + ZoneinfoCacheItem() : PythonImportCacheItem("zoneinfo"), ZoneInfo("ZoneInfo", this) { } - ~PytzCacheItem() override { + ~ZoneinfoCacheItem() override { } - PythonImportCacheItem timezone; + PythonImportCacheItem ZoneInfo; }; } // namespace duckdb diff --git a/src/include/duckdb_python/import_cache/python_import_cache.hpp b/src/include/duckdb_python/import_cache/python_import_cache.hpp index 6438332f6..bbe67d00b 100644 --- a/src/include/duckdb_python/import_cache/python_import_cache.hpp +++ b/src/include/duckdb_python/import_cache/python_import_cache.hpp @@ -33,11 +33,11 @@ struct PythonImportCache { PathlibCacheItem pathlib; PolarsCacheItem polars; DuckdbCacheItem duckdb; - PytzCacheItem pytz; TypesCacheItem types; TypingCacheItem typing; UuidCacheItem uuid; CollectionsCacheItem collections; + ZoneinfoCacheItem zoneinfo; public: nb::handle AddCache(nb::object item); diff --git a/src/include/duckdb_python/import_cache/python_import_cache_modules.hpp b/src/include/duckdb_python/import_cache/python_import_cache_modules.hpp index e82cd2dcc..f500b8995 100644 --- a/src/include/duckdb_python/import_cache/python_import_cache_modules.hpp +++ b/src/include/duckdb_python/import_cache/python_import_cache_modules.hpp @@ -8,8 +8,8 @@ #include "duckdb_python/import_cache/modules/pathlib_module.hpp" #include "duckdb_python/import_cache/modules/polars_module.hpp" #include "duckdb_python/import_cache/modules/duckdb_module.hpp" -#include "duckdb_python/import_cache/modules/pytz_module.hpp" #include "duckdb_python/import_cache/modules/types_module.hpp" #include "duckdb_python/import_cache/modules/typing_module.hpp" #include "duckdb_python/import_cache/modules/uuid_module.hpp" -#include "duckdb_python/import_cache/modules/collections_module.hpp" \ No newline at end of file +#include "duckdb_python/import_cache/modules/collections_module.hpp" +#include "duckdb_python/import_cache/modules/zoneinfo_module.hpp" \ No newline at end of file diff --git a/src/include/duckdb_python/pandas/column/pandas_numpy_column.hpp b/src/include/duckdb_python/pandas/column/pandas_numpy_column.hpp index 2e28dfa6e..deba2aa83 100644 --- a/src/include/duckdb_python/pandas/column/pandas_numpy_column.hpp +++ b/src/include/duckdb_python/pandas/column/pandas_numpy_column.hpp @@ -13,11 +13,14 @@ class PandasNumpyColumn : public PandasColumn { D_ASSERT(nb::hasattr(arr, "strides")); auto first_stride = nb::cast(arr.attr("strides").attr("__getitem__")(0)); auto item_size = nb::cast(arr.attr("dtype").attr("itemsize")); - // The scan steps through whole items at aligned addresses, which a field of a packed structured array lacks - if (item_size <= 0 || first_stride % item_size != 0 || !nb::cast(arr.attr("flags").attr("aligned"))) { + // The scan steps forward through whole aligned items, which a reversed view or a packed structured field lacks + if (item_size <= 0 || first_stride < 0 || first_stride % item_size != 0 || + !nb::cast(arr.attr("flags").attr("aligned"))) { array = NumpyArray(nb::module_::import_("numpy").attr("require")(arr, nb::arg("requirements") = "CA")); } - stride = nb::cast(array.GetArray().attr("strides").attr("__getitem__")(0)); + auto copied_stride = nb::cast(array.GetArray().attr("strides").attr("__getitem__")(0)); + // numpy calls an array of one item contiguous whatever its stride, so the copy can leave it reversed + stride = UnsafeNumericCast(copied_stride < 0 ? item_size : copied_stride); } public: diff --git a/src/native/python_conversion.cpp b/src/native/python_conversion.cpp index 203d68576..5a86ae445 100644 --- a/src/native/python_conversion.cpp +++ b/src/native/python_conversion.cpp @@ -762,7 +762,9 @@ struct PythonVectorConversion { if (result.GetType().id() != LogicalTypeId::BOOLEAN) { throw TypeMismatchException( LogicalType::BOOLEAN, result.GetType(), - "Python Conversion Failure: Expected a value of type %s, but got a value of type boolean"); + StringUtil::Format( + "Python Conversion Failure: Expected a value of type %s, but got a value of type boolean", + result.GetType().ToString())); } FlatVector::GetDataMutable(result)[result_offset] = val; } diff --git a/src/native/python_objects.cpp b/src/native/python_objects.cpp index 2b476ddd6..f147cb988 100644 --- a/src/native/python_objects.cpp +++ b/src/native/python_objects.cpp @@ -661,11 +661,29 @@ nb::object PythonObject::FromValue(const Value &val, const LogicalType &type, return nb::str(fallback_str.c_str(), fallback_str.size()); } 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); - auto tz_info = import_cache.pytz.timezone()(client_properties.time_zone); - return timestamp_utc.attr("astimezone")(tz_info); + nb::object zone; + try { + zone = import_cache.zoneinfo.ZoneInfo()(client_properties.time_zone); + } catch (nb::python_error &e) { + if (!e.matches(PyExc_KeyError)) { + throw; + } + throw InvalidInputException( + "Time zone '%s' is not available to Python's zoneinfo. Install the 'tzdata' " + "package, or SET TimeZone to a name the IANA database has.", + client_properties.time_zone); + } + auto timestamp_utc = py_timestamp.attr("replace")(nb::arg("tzinfo") = import_cache.datetime.timezone.utc()); + try { + return timestamp_utc.attr("astimezone")(zone); + } catch (nb::python_error &e) { + // The local wall time can leave Python's year range while the UTC instant is inside it. + if (!e.matches(PyExc_OverflowError)) { + throw; + } + auto fallback_str = val.ToString(); + return nb::str(fallback_str.c_str(), fallback_str.size()); + } } return py_timestamp; } diff --git a/src/numpy/numpy_scan.cpp b/src/numpy/numpy_scan.cpp index 68dc4d1a0..aac3fdb29 100644 --- a/src/numpy/numpy_scan.cpp +++ b/src/numpy/numpy_scan.cpp @@ -26,7 +26,7 @@ void ScanNumpyColumn(NumpyArray &numpy_col, idx_t stride, idx_t offset, Vector & if (stride == sizeof(T)) { FlatVector::SetData(out, data_ptr_cast(src_ptr + offset), count_t(count)); } else { - auto tgt_ptr = (T *)FlatVector::GetData(out); + auto tgt_ptr = FlatVector::GetDataMutable(out); for (idx_t i = 0; i < count; i++) { tgt_ptr[i] = src_ptr[stride / sizeof(T) * (i + offset)]; } @@ -34,30 +34,31 @@ void ScanNumpyColumn(NumpyArray &numpy_col, idx_t stride, idx_t offset, Vector & } template -void ScanNumpyCategoryTemplated(NumpyArray &column, idx_t offset, Vector &out, idx_t count) { +void ScanNumpyCategoryTemplated(NumpyArray &column, idx_t stride, idx_t offset, Vector &out, idx_t count) { auto src_ptr = (T *)column.Data(); - auto tgt_ptr = (V *)FlatVector::GetData(out); + auto tgt_ptr = FlatVector::GetDataMutable(out); auto &tgt_mask = FlatVector::ValidityMutable(out); for (idx_t i = 0; i < count; i++) { - if (src_ptr[i + offset] == -1) { + auto code = src_ptr[stride / sizeof(T) * (i + offset)]; + if (code == -1) { // Null value tgt_mask.SetInvalid(i); } else { - tgt_ptr[i] = src_ptr[i + offset]; + tgt_ptr[i] = code; } } } template -void ScanNumpyCategory(NumpyArray &column, idx_t count, idx_t offset, Vector &out, string &src_type) { +void ScanNumpyCategory(NumpyArray &column, idx_t stride, idx_t count, idx_t offset, Vector &out, string &src_type) { if (src_type == "int8") { - ScanNumpyCategoryTemplated(column, offset, out, count); + ScanNumpyCategoryTemplated(column, stride, offset, out, count); } else if (src_type == "int16") { - ScanNumpyCategoryTemplated(column, offset, out, count); + ScanNumpyCategoryTemplated(column, stride, offset, out, count); } else if (src_type == "int32") { - ScanNumpyCategoryTemplated(column, offset, out, count); + ScanNumpyCategoryTemplated(column, stride, offset, out, count); } else if (src_type == "int64") { - ScanNumpyCategoryTemplated(column, offset, out, count); + ScanNumpyCategoryTemplated(column, stride, offset, out, count); } else { throw NotImplementedException("The Pandas type " + src_type + " for categorical types is not implemented yet"); } @@ -537,13 +538,13 @@ void NumpyScan::Scan(ClientContext &context, PandasColumnBindData &bind_data, id case NumpyNullableType::CATEGORY: { switch (out.GetType().InternalType()) { case PhysicalType::UINT8: - ScanNumpyCategory(array, count, offset, out, bind_data.internal_categorical_type); + ScanNumpyCategory(array, stride, count, offset, out, bind_data.internal_categorical_type); break; case PhysicalType::UINT16: - ScanNumpyCategory(array, count, offset, out, bind_data.internal_categorical_type); + ScanNumpyCategory(array, stride, count, offset, out, bind_data.internal_categorical_type); break; case PhysicalType::UINT32: - ScanNumpyCategory(array, count, offset, out, bind_data.internal_categorical_type); + ScanNumpyCategory(array, stride, count, offset, out, bind_data.internal_categorical_type); break; default: throw InternalException("Invalid Physical Type for ENUMs"); diff --git a/src/pyconnection.cpp b/src/pyconnection.cpp index 18096b1fb..6ff11149f 100644 --- a/src/pyconnection.cpp +++ b/src/pyconnection.cpp @@ -1816,7 +1816,7 @@ NumpyObjectType DuckDBPyConnection::IsAcceptedNumpyObject(const nb::object &obje default: return NumpyObjectType::INVALID; } - } else if (duckdb::PyUtil::IsDictLike(object)) { + } else if (PyDict_Check(object.ptr())) { int dim = -1; for (auto item : nb::cast(object)) { if (!IsValidNumpyDimensions(item.second, dim)) { @@ -1824,9 +1824,9 @@ NumpyObjectType DuckDBPyConnection::IsAcceptedNumpyObject(const nb::object &obje } } return NumpyObjectType::DICT; - } else if (duckdb::PyUtil::IsListLike(object)) { + } else if (PyList_Check(object.ptr()) || PyTuple_Check(object.ptr())) { int dim = -1; - for (auto item : nb::cast(object)) { + for (auto item : object) { if (!IsValidNumpyDimensions(item, dim)) { return NumpyObjectType::INVALID; } diff --git a/src/python_replacement_scan.cpp b/src/python_replacement_scan.cpp index 335329517..10db51222 100644 --- a/src/python_replacement_scan.cpp +++ b/src/python_replacement_scan.cpp @@ -71,15 +71,20 @@ ArrowObjectKind PythonReplacementScan::ClassifyArrowObject(const nb::handle &ent return ArrowObjectKind::NONE; } -//! Whether the object is a dict, a list or a numpy array, the containers pandas_scan reads as columns +//! Whether the object is a dict, a list, a tuple or a numpy array, the containers pandas_scan reads as columns static bool IsNumpyContainer(const nb::handle &entry) { - if (PyDict_Check(entry.ptr()) || PyList_Check(entry.ptr())) { + if (PyDict_Check(entry.ptr()) || PyList_Check(entry.ptr()) || PyTuple_Check(entry.ptr())) { return true; } return ModuleIsLoaded() && duckdb::PyUtil::IsInstance(entry, DuckDBPyConnection::ImportCache()->numpy.ndarray()); } +static bool IsPandasSeries(const nb::handle &entry) { + return ModuleIsLoaded() && + duckdb::PyUtil::IsInstance(entry, DuckDBPyConnection::ImportCache()->pandas.Series()); +} + static unique_ptr CreateArrowScan(const nb::handle &entry, const string &name, ArrowSourceKind kind) { nb::object source; try { @@ -260,6 +265,11 @@ static unique_ptr ReplacementFor(const nb::object &entry, const string break; } } + if (IsPandasSeries(entry)) { + // A pandas 2 Series exports no Arrow stream, so it reads as the one column pandas 3 exports, without its index + auto frame = entry.attr("to_frame")(nb::arg("name") = "value").attr("reset_index")(nb::arg("drop") = true); + return ReplacementFor(frame, name, context, kind); + } if (DuckDBPyConnection::IsAcceptedNumpyObject(entry) != NumpyObjectType::INVALID) { numpytype = DuckDBPyConnection::IsAcceptedNumpyObject(entry); string np_name = "np_" + StringUtil::GenerateRandomName(); @@ -280,7 +290,7 @@ static unique_ptr ReplacementFor(const nb::object &entry, const string } case NumpyObjectType::LIST: idx = 0; - for (auto item : nb::cast(entry)) { + for (auto item : entry) { data[("column" + std::to_string(idx)).c_str()] = item; idx++; } diff --git a/tests/conftest.py b/tests/conftest.py index 29ddf927b..c33319862 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -31,6 +31,16 @@ def _get_pandas_ge_3(): PANDAS_GE_3 = _get_pandas_ge_3() +def _get_engine_ge_2_1(): + from packaging.version import Version + + return Version(duckdb.__version__) >= Version("2.1.0.dev0") + + +# A client branch builds against the engine branch of its own version, so the version names the engine's behaviour. +ENGINE_GE_2_1 = _get_engine_ge_2_1() + + def is_string_dtype(dtype): """Check if a dtype is a string dtype (works across Pandas 2.x and 3.0). diff --git a/tests/fast/api/test_cursor.py b/tests/fast/api/test_cursor.py index 044e736dd..351b69a13 100644 --- a/tests/fast/api/test_cursor.py +++ b/tests/fast/api/test_cursor.py @@ -1,5 +1,8 @@ # simple DB API testcase +import datetime +import zoneinfo + import pytest import duckdb @@ -93,7 +96,8 @@ def test_cursor_timezone(self): # Because the 'timezone' setting was not explicitly set for the connection # the setting of the DBConfig is used instead res = con1.execute("SELECT make_timestamptz(2000,01,20,03,30,59)").fetchone() - assert str(res) == "(datetime.datetime(2000, 1, 20, 3, 30, 59, tzinfo=),)" + assert res == (datetime.datetime(2000, 1, 20, 3, 30, 59, tzinfo=zoneinfo.ZoneInfo("UTC")),) + assert res[0].tzinfo.key == "UTC" def test_cursor_closed(self): con = duckdb.connect(":memory:") diff --git a/tests/fast/api/test_fetch_conversion_errors.py b/tests/fast/api/test_fetch_conversion_errors.py index c9f4965b6..b54967acd 100644 --- a/tests/fast/api/test_fetch_conversion_errors.py +++ b/tests/fast/api/test_fetch_conversion_errors.py @@ -88,11 +88,9 @@ def test_a_later_fetch_raises_the_same_error(self, con, fetch): assert str(again.value) == str(first.value) assert con.sql("SELECT 42").fetchall() == [(42,)] - @needs_icu def test_a_python_error_is_raised_again_as_its_text(self, con): - # The instant fits a datetime in UTC, but its wall clock in Tokyo is past datetime.max. - con.execute("SET TimeZone = 'Asia/Tokyo'") - con.execute("SELECT TIMESTAMPTZ '9999-12-31 23:59:59+00'") + # The day count fits the engine's interval, but datetime.timedelta caps days at 999999999. + con.execute("SELECT INTERVAL 1000000000 DAYS") with pytest.raises(OverflowError) as first: con.fetchone() with pytest.raises(duckdb.InvalidInputException, match=re.escape(str(first.value))): diff --git a/tests/fast/api/test_native_tz.py b/tests/fast/api/test_native_tz.py index 0545182ee..d869bffa1 100644 --- a/tests/fast/api/test_native_tz.py +++ b/tests/fast/api/test_native_tz.py @@ -30,23 +30,23 @@ def test_native_python_timestamp_timezone(self, duckdb_cursor): res = duckdb_cursor.execute(f"select TimeRecStart as tz from '{filename}'").fetchone() assert res[0].hour == 14 assert res[0].minute == 52 - assert res[0].tzinfo.zone == "America/Los_Angeles" + assert res[0].tzinfo.key == "America/Los_Angeles" res = duckdb_cursor.execute(f"select TimeRecStart as tz from '{filename}'").fetchall()[0] assert res[0].hour == 14 assert res[0].minute == 52 - assert res[0].tzinfo.zone == "America/Los_Angeles" + assert res[0].tzinfo.key == "America/Los_Angeles" res = duckdb_cursor.execute(f"select TimeRecStart as tz from '{filename}'").fetchmany(1)[0] assert res[0].hour == 14 assert res[0].minute == 52 - assert res[0].tzinfo.zone == "America/Los_Angeles" + assert res[0].tzinfo.key == "America/Los_Angeles" duckdb_cursor.execute("SET timezone='UTC';") res = duckdb_cursor.execute(f"select TimeRecStart as tz from '{filename}'").fetchone() assert res[0].hour == 21 assert res[0].minute == 52 - assert res[0].tzinfo.zone == "UTC" + assert res[0].tzinfo.key == "UTC" def test_native_python_time_timezone(self, duckdb_cursor): res = duckdb_cursor.execute(f"select TimeRecStart::TIMETZ as tz from '{filename}'").fetchone() diff --git a/tests/fast/arrow/scan/test_arrow_scan_row_ranges.py b/tests/fast/arrow/scan/test_arrow_scan_row_ranges.py index abf5588bd..736b9449c 100644 --- a/tests/fast/arrow/scan/test_arrow_scan_row_ranges.py +++ b/tests/fast/arrow/scan/test_arrow_scan_row_ranges.py @@ -919,14 +919,6 @@ def test_a_struct_dictionary_without_a_null_value_reads_its_values(self, indices con.register("src", pa.table({"d": column})) assert con.execute("SELECT d FROM src").fetchall() == expected - @pytest.mark.xfail( - strict=True, - raises=duckdb.InternalException, - reason=( - "the engine fails to read a list of dictionaries whose first vector holds only empty lists: it " - "dereferences a dictionary it never built" - ), - ) @pytest.mark.timeout(60) @pytest.mark.parametrize("threads", [1, 8]) @pytest.mark.parametrize(("empty", "valued"), [(2048, 1), (300_000, 1000)], ids=["one_vector", "several_ranges"]) diff --git a/tests/fast/arrow/scan/test_arrow_scan_rules.py b/tests/fast/arrow/scan/test_arrow_scan_rules.py index dc80e3831..49d3b50fc 100644 --- a/tests/fast/arrow/scan/test_arrow_scan_rules.py +++ b/tests/fast/arrow/scan/test_arrow_scan_rules.py @@ -1104,11 +1104,16 @@ def test_each_node_renders_its_expected_text(self): SUBINTERPRETER_TESTS = Path(__file__).resolve().parents[2] / "test_subinterpreter_gil.py" -# Scanned after the process created a subinterpreter, a Python generator backed reader must still be pulled safely. -GENERATOR_INPUT = """ +# pyarrow imports pandas when it first builds an array from a Python object, an import the spinner slows to a minute +# or more, as every file access then waits a switch interval for the GIL. +GENERATOR_IMPORTS = """ import pyarrow as pa +pa.array([0]) +""" +# Scanned after the process created a subinterpreter, a Python generator backed reader must still be pulled safely. +GENERATOR_INPUT = """ def produce(): for index in range(200): yield pa.record_batch([pa.array(range(index * 100, index * 100 + 100))], names=["n"]) @@ -1126,10 +1131,6 @@ def produce(): class TestAfterASubinterpreter: - @pytest.mark.skipif( - sys.platform == "win32", - reason="the child process hangs in the scan on Windows once a subinterpreter exists; Linux and macOS pass", - ) @pytest.mark.timeout(120) def test_a_generator_backed_reader_is_scanned_at_one_and_four_threads(self): import importlib.util @@ -1139,7 +1140,7 @@ def test_a_generator_backed_reader_is_scanned_at_one_and_four_threads(self): spec.loader.exec_module(harness) if importlib.util.find_spec(harness.INTERPRETERS) is None: pytest.skip(f"{harness.INTERPRETERS} is not available") - proc = harness._run("", GENERATOR_INPUT) + proc = harness._run(GENERATOR_IMPORTS, GENERATOR_INPUT) report = f"exit status {proc.returncode}\nstdout:\n{proc.stdout[-2000:]}\nstderr:\n{proc.stderr[-2000:]}" answers = [line.removeprefix(harness.PRECONDITION) for line in proc.stdout.splitlines()] answers = [line for line in answers if line in ("0", "1")] diff --git a/tests/fast/arrow/scan/test_arrow_scan_streams.py b/tests/fast/arrow/scan/test_arrow_scan_streams.py index f02effebc..4a49b7aa0 100644 --- a/tests/fast/arrow/scan/test_arrow_scan_streams.py +++ b/tests/fast/arrow/scan/test_arrow_scan_streams.py @@ -17,6 +17,7 @@ import _pushdown_helpers import pytest from _pushdown_helpers import eventually, scan_of +from conftest import ENGINE_GE_2_1 import duckdb @@ -1109,13 +1110,16 @@ def test_a_coarse_timestamp_beyond_microseconds_try_casts_to_null(self, con, uni assert rows(con, "SELECT 42") == [(42,)] @pytest.mark.xfail( + not ENGINE_GE_2_1, strict=True, raises=duckdb.InternalException, - reason="the engine's text form of a timestamp before its date range throws an internal error", + reason=( + "the 2.0 engine's text form of a timestamp before the first midnight of its range throws an internal error" + ), ) - def test_a_timestamp_before_the_date_range_try_casts_to_null(self, con): + def test_a_timestamp_before_the_first_midnight_of_the_range_reads_as_text(self, con): con.register("early", pa.table({"t": pa.array([-106_751_991 * 86_400_000_000 - 1], pa.timestamp("us"))})) - assert rows(con, "SELECT TRY_CAST(t AS VARCHAR) FROM early") == [(None,)] + assert rows(con, "SELECT TRY_CAST(t AS VARCHAR) FROM early") == [("290309-12-21 (BC) 23:59:59.999999",)] assert rows(con, "SELECT 42") == [(42,)] def test_the_first_day_of_nanoseconds_reads_as_text(self, con): diff --git a/tests/fast/arrow/test_arrow_offsets.py b/tests/fast/arrow/test_arrow_offsets.py index 304b87958..85949cb99 100644 --- a/tests/fast/arrow/test_arrow_offsets.py +++ b/tests/fast/arrow/test_arrow_offsets.py @@ -3,7 +3,6 @@ import decimal import pytest -import pytz pa = pytest.importorskip("pyarrow") @@ -324,10 +323,10 @@ def test_struct_of_duration(self, duckdb_cursor, constructor, unit, expected, co @pytest.mark.parametrize( ("constructor", "unit", "expected"), [ - (pa_timestamp(), "ms", datetime.datetime(1970, 1, 1, 0, 2, 11, 72000, tzinfo=pytz.utc)), - (pa_timestamp(), "s", datetime.datetime(1970, 1, 2, 12, 24, 32, 0, tzinfo=pytz.utc)), - (pa_timestamp(), "ns", datetime.datetime(1970, 1, 1, 0, 0, 0, 131, tzinfo=pytz.utc)), - (pa_timestamp(), "us", datetime.datetime(1970, 1, 1, 0, 0, 0, 131072, tzinfo=pytz.utc)), + (pa_timestamp(), "ms", datetime.datetime(1970, 1, 1, 0, 2, 11, 72000, tzinfo=datetime.timezone.utc)), + (pa_timestamp(), "s", datetime.datetime(1970, 1, 2, 12, 24, 32, 0, tzinfo=datetime.timezone.utc)), + (pa_timestamp(), "ns", datetime.datetime(1970, 1, 1, 0, 0, 0, 131, tzinfo=datetime.timezone.utc)), + (pa_timestamp(), "us", datetime.datetime(1970, 1, 1, 0, 0, 0, 131072, tzinfo=datetime.timezone.utc)), ], ) def test_struct_of_timestamp_tz(self, duckdb_cursor, constructor, unit, expected, col1_null, col2_null): diff --git a/tests/fast/arrow/test_timestamp_timezone.py b/tests/fast/arrow/test_timestamp_timezone.py index c847c28d6..f14edfc2b 100644 --- a/tests/fast/arrow/test_timestamp_timezone.py +++ b/tests/fast/arrow/test_timestamp_timezone.py @@ -1,7 +1,6 @@ import datetime import pytest -import pytz import duckdb @@ -21,7 +20,7 @@ def generate_table(current_time, precision, timezone): class TestArrowTimestampsTimezone: def test_timestamp_timezone(self, duckdb_cursor): precisions = ["us", "s", "ns", "ms"] - current_time = datetime.datetime(2017, 11, 28, 23, 55, 59, tzinfo=pytz.UTC) + current_time = datetime.datetime(2017, 11, 28, 23, 55, 59, tzinfo=datetime.timezone.utc) con = duckdb.connect() con.execute("SET TimeZone = 'UTC'") for precision in precisions: diff --git a/tests/fast/pandas/test_append_df.py b/tests/fast/pandas/test_append_df.py index 839fe4085..1a98abca0 100644 --- a/tests/fast/pandas/test_append_df.py +++ b/tests/fast/pandas/test_append_df.py @@ -2,10 +2,10 @@ import enum import math import pathlib +import zoneinfo from decimal import Decimal import pytest -import pytz import duckdb import numpy as np @@ -416,7 +416,7 @@ def test_naive_datetime_into_timestamptz_reads_in_the_session_time_zone(self, du duckdb_cursor.execute("CREATE TABLE t (a TIMESTAMPTZ)") duckdb_cursor.append("t", pd.DataFrame({"a": _outside_the_sample(AWARE, naive)})) values = _column(duckdb_cursor) - assert (values[0], values[-1]) == (AWARE, pytz.timezone(zone).localize(NAIVE)) + assert (values[0], values[-1]) == (AWARE, NAIVE.replace(tzinfo=zoneinfo.ZoneInfo(zone))) @pytest.mark.parametrize( "zone", @@ -439,7 +439,7 @@ def test_date_into_timestamptz_reads_as_midnight_in_the_session_time_zone(self, duckdb_cursor.execute("CREATE TABLE t (a TIMESTAMPTZ)") duckdb_cursor.append("t", pd.DataFrame({"a": _outside_the_sample(AWARE, datetime.date(2024, 5, 6))})) values = _column(duckdb_cursor) - assert (values[0], values[-1]) == (AWARE, pytz.timezone(zone).localize(datetime.datetime(2024, 5, 6))) + assert (values[0], values[-1]) == (AWARE, datetime.datetime(2024, 5, 6, tzinfo=zoneinfo.ZoneInfo(zone))) @pytest.mark.xfail(strict=True, reason=NOT_YET_DECLARED) def test_by_name_declares_each_column_from_its_name(self, duckdb_cursor): diff --git a/tests/fast/pandas/test_object_column_conversion.py b/tests/fast/pandas/test_object_column_conversion.py index c6a8bc4d9..bd9fafc96 100644 --- a/tests/fast/pandas/test_object_column_conversion.py +++ b/tests/fast/pandas/test_object_column_conversion.py @@ -1417,6 +1417,15 @@ def test_other_conversions_keep_their_casts(duckdb_cursor, prepare, expected): convert() +def test_bool_refused_by_a_number_type_names_the_expected_type(duckdb_cursor): + duckdb_cursor.create_function("produce", lambda: True, [], duckdb.sqltype("INTEGER"), side_effects=True) + with pytest.raises(duckdb.TypeMismatchException) as error: + duckdb_cursor.execute("SELECT produce()").fetchall() + message = str(error.value) + assert "%s" not in message + assert "Expected a value of type INTEGER, but got a value of type boolean" in message + + @not_yet("the error for a long string outside the sample repeats the whole string") def test_error_preview_of_a_long_string_is_bounded(duckdb_cursor): duckdb_cursor.register("frame", column(1, "é" * 100000, at=2000)) diff --git a/tests/fast/pandas/test_pandas_scan_bind_reads.py b/tests/fast/pandas/test_pandas_scan_bind_reads.py index b687caf37..dd6dbabb8 100644 --- a/tests/fast/pandas/test_pandas_scan_bind_reads.py +++ b/tests/fast/pandas/test_pandas_scan_bind_reads.py @@ -10,6 +10,7 @@ import sys import textwrap import tracemalloc +import types from collections.abc import Callable from typing import Any, NamedTuple @@ -958,6 +959,20 @@ def test_nullable_column_with_strided_data_keeps_its_mask(duckdb_cursor, path, d assert _length_and_first_differences(result, expected) == (len(expected), []) +@pytest.mark.parametrize("path", PATHS) +@pytest.mark.parametrize("dtype", ["bool", "int8", "int16", "int32", "int64", "uint8", "uint16", "uint32", "uint64"]) +def test_sliced_numpy_column_reads_its_rows(duckdb_cursor, path, dtype): + positions = np.arange(_SLICED_ROWS, dtype="int64") + values = positions % 5 < 2 if dtype == "bool" else positions.astype(dtype) + source = pd.DataFrame({"i": positions, "x": values}).iloc[1::3] + column = source["x"].to_numpy() + assert column.strides[0] == 3 * column.itemsize + expected = column.tolist() + assert len(expected) > PANDAS_PARTITION_ROWS + result = _scan_in_order(duckdb_cursor, source, path) + assert _length_and_first_differences(result, expected) == (len(expected), []) + + BIND_HEAP_ROWS = 200_000 @@ -1031,3 +1046,210 @@ def test_dict_array_bind_allocates_nothing_per_row(duckdb_cursor, path, make_arr positions = np.arange(BIND_HEAP_ROWS, dtype="int64") source = {"i": positions, "x": make_array(positions)} assert _bind_heap_peak(duckdb_cursor, source, path) < BIND_HEAP_ROWS // 4 + + +def _categorical_codes_as_a_view(monkeypatch): + # pandas 2 hands out cat.codes as a view of the column's codes, which pandas 3 copies; a view reproduces pandas 2. + def codes(accessor): + return pd.Series(accessor._parent.codes, index=accessor._index, copy=False) + + monkeypatch.setattr(pd.core.arrays.categorical.CategoricalAccessor, "codes", property(codes)) + + +@pytest.mark.parametrize("path", PATHS) +@pytest.mark.parametrize( + "rows", + [slice(None, None, 2), slice(1, None, 2), slice(None, None, -2)], + ids=["every_second", "every_second_from_1", "every_second_reversed"], +) +@pytest.mark.parametrize( + ("categories", "codes"), + [(3, "int8"), (200, "int16"), (300, "int16"), (40000, "int32")], + ids=["int8_codes", "int16_codes_of_a_one_byte_enum", "int16_codes", "int32_codes"], +) +def test_strided_view_of_a_text_categorical_reads_its_rows(duckdb_cursor, monkeypatch, path, rows, categories, codes): + _categorical_codes_as_a_view(monkeypatch) + positions = np.arange(_SLICED_ROWS, dtype="int64") + kept = positions[rows] + category_codes = (positions % categories).astype(codes) + category_codes[_missing_at(positions, kept[_edge_rows(len(kept))])] = -1 + labels = pd.Index([f"k{n}" for n in range(categories)], dtype=object) + column = pd.Categorical.from_codes(category_codes, categories=labels) + source = pd.DataFrame({"i": positions, "x": column}).iloc[rows] + read_codes = source["x"].cat.codes.to_numpy() + assert (read_codes.dtype, read_codes.strides[0]) == (np.dtype(codes), rows.step * read_codes.itemsize) + expected = [None if pd.isna(value) else value for value in source.sort_values("i")["x"].tolist()] + assert len(expected) > PANDAS_PARTITION_ROWS + result = _scan_in_order(duckdb_cursor, source, path) + assert _length_and_first_differences(result, expected) == (len(expected), []) + + +_REVERSED_ROWS = [slice(None, None, -1), slice(None, None, -2), slice(0, None, -1)] +_REVERSED_IDS = ["reversed", "every_second_reversed", "first_row_reversed"] + + +@pytest.mark.parametrize("path", PATHS) +@pytest.mark.parametrize("rows", _REVERSED_ROWS, ids=_REVERSED_IDS) +@pytest.mark.parametrize( + "name", + [pytest.param(name, marks=_NEEDS_PYARROW) if name in _COLUMNS_NEEDING_PYARROW else name for name in FRAME_COLUMNS], +) +def test_reversed_frame_column_reads_its_rows_in_order(duckdb_cursor, path, rows, name): + source = pd.DataFrame({name: FRAME_COLUMNS[name].build()}).iloc[rows] + sql = f'SELECT typeof("{name}"), "{name}" FROM {{source}}' + # The type can follow the view: a category column whose view skips its only missing value reads as BIGINT. + contiguous = _execute(duckdb_cursor, source.copy(), path, sql).fetchall() + assert [value for _, value in contiguous] == FRAME_COLUMNS[name].values[rows] + assert _execute(duckdb_cursor, source, path, sql).fetchall() == contiguous + + +_LONG_REVERSED_ROWS = 2 * VECTOR_ROWS + 5 + + +def _reversed_dict_column(name): + build, sql_type, values = DICT_COLUMNS[name] + return pytest.param( + lambda: {name: build()[::-1]}, + f"SELECT typeof({name}), {name} FROM {{source}}", + ([f"typeof({name})", name], [(sql_type, value) for value in values[::-1]]), + id=f"dict_{name}", + ) + + +@pytest.mark.parametrize("path", PATHS) +@pytest.mark.parametrize( + ("make_source", "sql", "expected"), + [ + *[_reversed_dict_column(name) for name in DICT_COLUMNS], + pytest.param( + lambda: {"a": np.arange(_LONG_REVERSED_ROWS, dtype="int64")[::-1]}, + "SELECT * FROM {source}", + (["a"], [(value,) for value in range(_LONG_REVERSED_ROWS - 1, -1, -1)]), + id="dict_across_vectors", + ), + pytest.param( + lambda: np.arange(5, dtype="int64")[::-1], + "SELECT * FROM {source}", + (["column0"], [(4,), (3,), (2,), (1,), (0,)]), + id="ndarray", + ), + pytest.param( + lambda: np.arange(6, dtype="int64").reshape(2, 3)[:, ::-1], + "SELECT * FROM {source}", + (["column0", "column1"], [(2, 5), (1, 4), (0, 3)]), + id="two_dimensional_ndarray", + ), + pytest.param( + lambda: [np.arange(3, dtype="int64")[::-1]], + "SELECT * FROM {source}", + (["column0"], [(2,), (1,), (0,)]), + id="list_of_arrays", + ), + pytest.param( + lambda: {"a": np.array([42], dtype="int64")[::-1]}, + "SELECT * FROM {source}", + (["a"], [(42,)]), + id="dict_of_one_reversed_item", + ), + ], +) +def test_reversed_numpy_array_reads_its_rows_in_order(duckdb_cursor, path, make_source, sql, expected): + assert _names_and_rows(duckdb_cursor, make_source(), path, sql) == expected + + +@pytest.mark.parametrize("path", PATHS) +@pytest.mark.parametrize("container", [tuple, list]) +@pytest.mark.parametrize( + ("arrays", "expected"), + [ + pytest.param( + [np.array([1, 2], dtype="int64"), np.array(["a", "b"], dtype=object)], + (["column0", "column1"], [(1, "a"), (2, "b")]), + id="two_arrays", + ), + pytest.param([np.array([1.5, 2.5], dtype="float64")], (["column0"], [(1.5,), (2.5,)]), id="one_array"), + pytest.param([[1, 2], [3, 4]], "not suitable for replacement scans", id="lists_that_are_not_arrays"), + pytest.param( + [np.arange(2, dtype="int64"), np.arange(3, dtype="int64")], + "not suitable for replacement scans", + id="arrays_of_unequal_lengths", + ), + pytest.param([], "at least one column", id="empty"), + ], +) +def test_tuple_of_arrays_reads_like_a_list_of_arrays(duckdb_cursor, path, container, arrays, expected): + source = container(arrays) + if isinstance(expected, str): + with pytest.raises(duckdb.InvalidInputException, match=expected): + _execute(duckdb_cursor, source, path, "SELECT * FROM {source}").fetchall() + else: + assert _names_and_rows(duckdb_cursor, source, path, "SELECT * FROM {source}") == expected + + +@pytest.mark.parametrize("path", PATHS) +@pytest.mark.parametrize( + "arrow_stream", + [ + pytest.param(True, id="as_pandas_gives_it", marks=_NEEDS_PYARROW), + pytest.param(False, id="without_an_arrow_stream"), + ], +) +@pytest.mark.parametrize( + ("make_series", "sql_type", "values"), + [ + pytest.param(lambda: pd.Series(np.array([1, 2, 3], dtype="int64")), "BIGINT", [1, 2, 3], id="int64"), + pytest.param(lambda: pd.Series(np.array([1, 2], dtype="int64"), name="x"), "BIGINT", [1, 2], id="named_int64"), + pytest.param( + lambda: pd.Series(np.array([1.5, np.nan], dtype="float64")), "DOUBLE", [1.5, None], id="float64_nan" + ), + pytest.param(lambda: pd.Series(pd.array([1, None, 3], dtype="Int64")), "BIGINT", [1, None, 3], id="Int64"), + pytest.param(lambda: pd.Series(["a", None, "c"], dtype=object), "VARCHAR", ["a", None, "c"], id="object_text"), + pytest.param( + lambda: pd.Series(_NAIVE_TIMES), + _NAIVE_TIMESTAMP, + [timestamp.to_pydatetime() for timestamp in _NAIVE_TIMES], + id="datetime", + ), + pytest.param( + lambda: pd.Series(np.array([1, 2], dtype="int64"), index=["a", "b"]), "BIGINT", [1, 2], id="text_index" + ), + pytest.param( + lambda: pd.Series([1, 2], index=["a", "b"], dtype="int64[pyarrow]"), + "BIGINT", + [1, 2], + id="arrow_backed_with_a_text_index", + marks=_NEEDS_PYARROW, + ), + ], +) +def test_series_reads_as_one_column_named_value( + duckdb_cursor, monkeypatch, path, arrow_stream, make_series, sql_type, values +): + if not arrow_stream: + # A pandas 2 Series has no Arrow stream; taking pandas 3's away reproduces it. + monkeypatch.delattr(pd.Series, "__arrow_c_stream__", raising=False) + series = make_series() + assert _names_and_rows(duckdb_cursor, series, path, "SELECT * FROM {source}") == ( + ["value"], + [(value,) for value in values], + ) + types_read = _execute(duckdb_cursor, series, path, "SELECT DISTINCT typeof(value) FROM {source}").fetchall() + assert types_read == [(sql_type,)] + + +@pytest.mark.parametrize("path", PATHS) +@pytest.mark.parametrize( + ("make_source", "type_name"), + [ + pytest.param(lambda: range(3), "range", id="range"), + pytest.param(lambda: {1, 2}, "set", id="set"), + pytest.param(lambda: (array for array in [np.arange(2)]), "generator", id="generator_of_arrays"), + pytest.param(lambda: iter([np.arange(2)]), "list_iterator", id="iterator_of_arrays"), + pytest.param(lambda: collections.deque([np.arange(2)]), "deque", id="deque_of_arrays"), + pytest.param(lambda: {"a": np.arange(2)}.values(), "dict_values", id="dict_values_of_arrays"), + pytest.param(lambda: types.MappingProxyType({"a": np.arange(2)}), "mappingproxy", id="mapping_of_arrays"), + ], +) +def test_iterable_that_is_not_a_supported_container_is_refused(duckdb_cursor, path, make_source, type_name): + with pytest.raises(duckdb.InvalidInputException, match=rf'(?s)of type "{type_name}".*not suitable for replacement'): + _execute(duckdb_cursor, make_source(), path, "SELECT * FROM {source}").fetchall() diff --git a/tests/fast/pandas/test_stride.py b/tests/fast/pandas/test_stride.py index 9434672b7..a17dbd291 100644 --- a/tests/fast/pandas/test_stride.py +++ b/tests/fast/pandas/test_stride.py @@ -7,14 +7,18 @@ class TestPandasStride: def test_stride(self, duckdb_cursor): - expected_df = pd.DataFrame(np.arange(20).reshape(5, 4), columns=["a", "b", "c", "d"]) + expected_df = pd.DataFrame(np.arange(20).reshape(5, 4), columns=["a", "b", "c", "d"], copy=False) + assert expected_df["a"].to_numpy().strides[0] == 4 * expected_df["a"].to_numpy().itemsize con = duckdb.connect() con.register("df_view", expected_df) output_df = con.execute("SELECT * FROM df_view;").fetchdf() pd.testing.assert_frame_equal(expected_df, output_df) def test_stride_fp32(self, duckdb_cursor): - expected_df = pd.DataFrame(np.arange(20, dtype="float32").reshape(5, 4), columns=["a", "b", "c", "d"]) + expected_df = pd.DataFrame( + np.arange(20, dtype="float32").reshape(5, 4), columns=["a", "b", "c", "d"], copy=False + ) + assert expected_df["a"].to_numpy().strides[0] == 4 * expected_df["a"].to_numpy().itemsize con = duckdb.connect() con.register("df_view", expected_df) output_df = con.execute("SELECT * FROM df_view;").fetchdf() @@ -61,7 +65,10 @@ def test_stride_timedelta(self, duckdb_cursor): pd.testing.assert_frame_equal(roundtrip, expected, check_dtype=False) def test_stride_fp64(self, duckdb_cursor): - expected_df = pd.DataFrame(np.arange(20, dtype="float64").reshape(5, 4), columns=["a", "b", "c", "d"]) + expected_df = pd.DataFrame( + np.arange(20, dtype="float64").reshape(5, 4), columns=["a", "b", "c", "d"], copy=False + ) + assert expected_df["a"].to_numpy().strides[0] == 4 * expected_df["a"].to_numpy().itemsize con = duckdb.connect() con.register("df_view", expected_df) output_df = con.execute("SELECT * FROM df_view;").fetchdf() diff --git a/tests/fast/relational_api/test_rapi_aggregations.py b/tests/fast/relational_api/test_rapi_aggregations.py index aeec5d854..6d0ca7b93 100644 --- a/tests/fast/relational_api/test_rapi_aggregations.py +++ b/tests/fast/relational_api/test_rapi_aggregations.py @@ -36,10 +36,12 @@ def test_any_value(self, table): expected = [(1,)] assert len(result) == len(expected) assert all(r == e for r, e in zip(result, expected, strict=False)) - result = table.order("id, t").any_value("v", groups="id", projected_columns="id").order("id").fetchall() - expected = [(1, 1), (2, 11), (3, 5)] - assert len(result) == len(expected) - assert all(r == e for r, e in zip(result, expected, strict=False)) + # A grouped aggregate reads its rows in no set order, so any non-NULL value of the group is an answer. + result = table.any_value("v", groups="id", projected_columns="id").order("id").fetchall() + assert [group for group, _ in result] == [1, 2, 3] + assert result[0][1] in {1, 2} + assert result[1][1] in {10, 11} + assert result[2][1] in {-1, 5} def test_arg_max(self, table): result = table.arg_max("t", "v").fetchall() diff --git a/tests/fast/spark/test_spark_column.py b/tests/fast/spark/test_spark_column.py index 290b2e26e..fbd4bf9b0 100644 --- a/tests/fast/spark/test_spark_column.py +++ b/tests/fast/spark/test_spark_column.py @@ -518,6 +518,16 @@ def test_string_to_binary_takes_utf8_bytes(self, spark): Row(s=b"a\\x41", b=None, a=[b"a\\x41"]), ] + @pytest.mark.skipif(USE_ACTUAL_SPARK, reason="COLLATE NOCASE is DuckDB's collation syntax") + def test_collated_string_casts_as_a_string(self, spark): + df = spark.sql( + "SELECT 'é' COLLATE NOCASE AS s, '9007199254740992.9' COLLATE NOCASE AS t, '1e2' COLLATE NOCASE AS u" + ) + res = df.select(F.col("s").cast("binary"), F.col("t").cast("bigint")) + assert res.collect() == [Row(s=b"\xc3\xa9", t=9007199254740992)] + with pytest.raises(Exception, match="CAST_INVALID_INPUT"): + df.select(F.col("u").cast("bigint")).collect() + def test_integral_to_binary(self, spark): df = spark.range(1) if USE_ACTUAL_SPARK and pyspark.__version__.startswith("3."): diff --git a/tests/fast/spark/test_spark_group_by.py b/tests/fast/spark/test_spark_group_by.py index 3c25ac35a..ba05cc344 100644 --- a/tests/fast/spark/test_spark_group_by.py +++ b/tests/fast/spark/test_spark_group_by.py @@ -184,7 +184,6 @@ def test_group_by_empty(self, spark): def test_group_by_first_and_last(self, spark): df = spark.createDataFrame([("Alice", 2), ("Bob", 5), ("Alice", None)], ("name", "age")) - df = df.orderBy(df.age) res = ( df.groupBy("name") .agg(first("age").alias("first_age"), last("age").alias("last_age")) @@ -192,7 +191,12 @@ def test_group_by_first_and_last(self, spark): .collect() ) - assert res == [Row(name="Alice", first_age=None, last_age=2), Row(name="Bob", first_age=5, last_age=5)] + # A grouped first and last read the rows in no set order, in Spark as in DuckDB. With three rows both still + # come from different rows of Alice's, so the pair is her two values in either order. + alice, bob = res + assert alice.name == "Alice" + assert {alice.first_age, alice.last_age} == {None, 2} + assert bob == Row(name="Bob", first_age=5, last_age=5) def test_standard_deviations(self, spark): df = spark.createDataFrame( diff --git a/tests/fast/spark/test_spark_readjson.py b/tests/fast/spark/test_spark_readjson.py index cb677fd46..991932db1 100644 --- a/tests/fast/spark/test_spark_readjson.py +++ b/tests/fast/spark/test_spark_readjson.py @@ -6,6 +6,7 @@ _ = pytest.importorskip("duckdb.experimental.spark") +from conftest import ENGINE_GE_2_1 from spark_namespace import USE_ACTUAL_SPARK from spark_namespace.errors import AnalysisException, error_condition from spark_namespace.sql.types import Row @@ -131,9 +132,12 @@ def test_decimals(self, spark, tmp_path): thousand, one = (Decimal("1000.00"), None) if USE_ACTUAL_SPARK else (None, Decimal("1.00")) expected = [Decimal("1.00"), Decimal("1.70"), Decimal("1.00"), thousand, one, None, None] assert _field(spark, tmp_path, "decimal(10,2)", *values) == expected - # Spark rounds the decimal text, DuckDB the double it reads. + # Spark and the 2.1 engine round the decimal text, the 2.0 engine the double it reads. values = _field(spark, tmp_path, "decimal(10,2)", "1.005", "true") - assert values == ([Decimal("1.01"), None] if USE_ACTUAL_SPARK else [Decimal("1.00"), Decimal("1.00")]) + if USE_ACTUAL_SPARK: + assert values == [Decimal("1.01"), None] + else: + assert values == [Decimal("1.01") if ENGINE_GE_2_1 else Decimal("1.00"), Decimal("1.00")] def test_booleans(self, spark, tmp_path): assert _field(spark, tmp_path, "boolean", "true", "false") == [True, False] diff --git a/tests/fast/test_all_types.py b/tests/fast/test_all_types.py index f346fa903..df04b9e80 100644 --- a/tests/fast/test_all_types.py +++ b/tests/fast/test_all_types.py @@ -6,7 +6,6 @@ from uuid import UUID import pytest -import pytz import duckdb import numpy as np @@ -133,9 +132,9 @@ def test_fetchall(self, cur_type): "time_tz": """CASE WHEN time_tz = '24:00:00-1559'::TIMETZ THEN '23:59:59.999999-1559'::TIMETZ ELSE time_tz END AS "time_tz" """, # noqa: E501 } min_datetime = datetime.datetime.min - min_datetime_with_utc = min_datetime.replace(tzinfo=pytz.UTC) + min_datetime_with_utc = min_datetime.replace(tzinfo=datetime.timezone.utc) max_datetime = datetime.datetime.max - max_datetime_with_utc = max_datetime.replace(tzinfo=pytz.UTC) + max_datetime_with_utc = max_datetime.replace(tzinfo=datetime.timezone.utc) correct_answer_map = { "bool": [(False,), (True,), (None,)], "tinyint": [(-128,), (127,), (None,)], @@ -196,7 +195,7 @@ def test_fetchall(self, cur_type): ( [], [ - datetime.datetime(1970, 1, 1, tzinfo=pytz.UTC), + datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc), None, min_datetime_with_utc, max_datetime_with_utc, @@ -237,7 +236,7 @@ def test_fetchall(self, cur_type): "timestamp_s": [(datetime.datetime(1990, 1, 1, 0, 0),)], "timestamp_ns": [(datetime.datetime(1990, 1, 1, 0, 0),)], "timestamp_ms": [(datetime.datetime(1990, 1, 1, 0, 0),)], - "timestamp_tz": [(datetime.datetime(1990, 1, 1, 0, 0, tzinfo=pytz.UTC),)], + "timestamp_tz": [(datetime.datetime(1990, 1, 1, 0, 0, tzinfo=datetime.timezone.utc),)], "union": [("Frank",), (5,), (None,)], "fixed_int_array": [((None, 2, 3),), ((4, 5, 6),), (None,)], "fixed_varchar_array": [(("a", None, "c"),), (("d", "e", "f"),), (None,)], diff --git a/tests/fast/test_optional_dependencies.py b/tests/fast/test_optional_dependencies.py index 2cacc71f3..706e68a6b 100644 --- a/tests/fast/test_optional_dependencies.py +++ b/tests/fast/test_optional_dependencies.py @@ -1,4 +1,5 @@ import importlib.util +import pathlib import subprocess import sys import textwrap @@ -11,7 +12,6 @@ "pyarrow", "polars", "fsspec", - "pytz", "typing_extensions", "pyspark", "IPython", @@ -41,8 +41,11 @@ def find_spec(self, fullname, path=None, target=None): """ -def run_blocked(cwd, script, blocked=OPTIONAL_PACKAGES): +def run_blocked(cwd, script, blocked=OPTIONAL_PACKAGES, pytz_tripwire=False): # An empty working directory keeps modules next to the caller's off the path that -c puts first. + if pytz_tripwire: + # pandas and pyarrow import pytz inside a try block, so the stand-in only serves scripts without them. + (pathlib.Path(cwd) / "pytz.py").write_text('raise AssertionError("pytz imported")\n') proc = subprocess.run( [sys.executable, "-S", "-c", BLOCK_IMPORTS + textwrap.dedent(script), ",".join(blocked)], capture_output=True, @@ -162,7 +165,6 @@ def test_spark_timestamps(tmp_path): def test_spark_timestamps_to_pandas(tmp_path): - # pandas before 3.0 cannot be imported without pytz. pytest.importorskip("pandas", minversion="3.0.0") run_blocked( tmp_path, @@ -201,14 +203,15 @@ def test_spark_udf_needs_numpy(tmp_path, blocked): ) -def test_spark_udf_with_a_timestamp_needs_pytz(tmp_path): - """The binding's Python UDF path needs pytz to pass a timestamp to a function.""" +def test_spark_udf_with_a_timestamp_needs_no_package(tmp_path): + """A timestamp reaches a Python UDF through the standard library's zoneinfo, so pytz is not needed.""" if importlib.util.find_spec("numpy") is None: pytest.skip("a UDF needs numpy") run_blocked( tmp_path, """ import datetime + import sys import duckdb from duckdb.experimental.spark.sql import SparkSession @@ -218,12 +221,30 @@ def test_spark_udf_with_a_timestamp_needs_pytz(tmp_path): assert spark.createDataFrame([(1,)], "n int").select(plus_one("n")).collect()[0][0] == 2 year = spark.udf.register("year_of", lambda t: t.year, "int") df = spark.createDataFrame([(datetime.datetime(2024, 1, 2),)], "t timestamp") - try: - df.select(year("t")).collect() - except duckdb.InvalidInputException as error: - assert "'pytz' failed to import" in str(error), repr(error) - else: - raise AssertionError("a UDF of a timestamp without pytz did not raise") + assert df.select(year("t")).collect()[0][0] == 2024 + assert "pytz" not in sys.modules + """, + blocked=tuple(package for package in OPTIONAL_PACKAGES if package != "numpy"), + pytz_tripwire=True, + ) + + +def test_timestamptz_fetch_needs_no_package(tmp_path): + run_blocked( + tmp_path, + """ + import datetime + import sys + import zoneinfo + + import duckdb + + con = duckdb.connect() + con.execute("SET TimeZone = 'Europe/Amsterdam'") + value = con.execute("SELECT TIMESTAMPTZ '2024-07-01 00:00:00+00'").fetchone()[0] + assert value == datetime.datetime(2024, 7, 1, 2, tzinfo=zoneinfo.ZoneInfo("Europe/Amsterdam")) + assert value.tzinfo is zoneinfo.ZoneInfo("Europe/Amsterdam") + assert "pytz" not in sys.modules """, - ("pytz",), + pytz_tripwire=True, ) diff --git a/tests/fast/test_subinterpreter_gil.py b/tests/fast/test_subinterpreter_gil.py index aa23c5efb..c6cd35081 100644 --- a/tests/fast/test_subinterpreter_gil.py +++ b/tests/fast/test_subinterpreter_gil.py @@ -229,7 +229,9 @@ def _native_stacks(pid: int) -> str: sample = subprocess.run(["sample", str(pid), "2"], capture_output=True, text=True, timeout=30, check=False) except (OSError, subprocess.TimeoutExpired): return "" - return f"\nnative stacks of the hung child:\n{sample.stdout[-6000:]}" + # The report ends in a list of loaded images longer than any cap, so only the call graph before it is kept. + call_graph = sample.stdout.partition("Binary Images:")[0] + return f"\nnative stacks of the hung child:\n{call_graph[:40000]}" @_needs(INTERPRETERS) diff --git a/tests/fast/test_type.py b/tests/fast/test_type.py index 9d2b5b8cb..9307e23f9 100644 --- a/tests/fast/test_type.py +++ b/tests/fast/test_type.py @@ -1,6 +1,7 @@ import sys import pytest +from conftest import ENGINE_GE_2_1 import duckdb from duckdb.sqltypes import ( @@ -112,8 +113,7 @@ def test_string_type(self): def test_string_type_collation(self): type = duckdb.string_type("NOCASE") - # collation does not show up in the string representation.. - assert str(type) == "VARCHAR" + assert str(type) == ("VARCHAR('NOCASE')" if ENGINE_GE_2_1 else "VARCHAR") def test_union_type(self): type = duckdb.union_type([BIGINT, VARCHAR, TINYINT]) diff --git a/tests/fast/types/test_time_ns.py b/tests/fast/types/test_time_ns.py index 0bcb29938..6eb0d0fb7 100644 --- a/tests/fast/types/test_time_ns.py +++ b/tests/fast/types/test_time_ns.py @@ -2,7 +2,7 @@ import pytest -from duckdb import ConversionException, sqltypes +from duckdb import sqltypes def test_time_ns_select(duckdb_cursor): @@ -12,10 +12,6 @@ def test_time_ns_select(duckdb_cursor): assert isinstance(result, datetime.time) -@pytest.mark.xfail( - raises=ConversionException, - reason="Conversion Error: Unimplemented type for cast (TIME -> TIME_NS)", -) def test_time_ns_insert(duckdb_cursor): """This tests that datetime.time values can be inserted as TIME_NS.""" duckdb_cursor.execute("SELECT TIME_NS '1992-09-20 11:30:00.123456'") @@ -66,7 +62,7 @@ def test_time_ns_pandas_roundtrip(duckdb_cursor): def test_time_pandas_roundtrip(duckdb_cursor): - """For Pandas, creating a table using CREATE .... AS SELECT FROM df, will create TIME_NS cols by default.""" + """A pandas column of datetime.time values comes back as TIME, not TIME_NS.""" pytest.importorskip("pandas") df = duckdb_cursor.execute("SELECT TIME '12:34:56.123456789' AS time_col").df() assert df["time_col"].dtype == "object" diff --git a/tests/fast/types/test_timestamp_tz.py b/tests/fast/types/test_timestamp_tz.py new file mode 100644 index 000000000..87f580747 --- /dev/null +++ b/tests/fast/types/test_timestamp_tz.py @@ -0,0 +1,160 @@ +import datetime +import os +import subprocess +import sys +import textwrap +import zoneinfo + +import pytest + +import duckdb + +UTC = datetime.timezone.utc + + +def fetch(con, literal): + return con.execute(f"SELECT TIMESTAMPTZ '{literal}'").fetchone()[0] + + +@pytest.fixture +def con(): + with duckdb.connect() as con: + yield con + + +@pytest.mark.parametrize("zone", ["UTC", "Europe/Amsterdam", "America/Los_Angeles", "Asia/Kolkata"]) +def test_fetch_in_the_session_time_zone(con, zone): + con.execute(f"SET TimeZone = '{zone}'") + value = fetch(con, "2024-07-01 12:00:00+00") + expected = datetime.datetime(2024, 7, 1, 12, tzinfo=UTC).astimezone(zoneinfo.ZoneInfo(zone)) + assert value == expected + assert value.utcoffset() == expected.utcoffset() + assert value.tzinfo is zoneinfo.ZoneInfo(zone) + assert value.tzinfo.key == zone + + +@pytest.mark.parametrize("literal", ["2024-10-27 00:30:00+00", "2024-10-27 01:30:00+00"]) +def test_repeated_hour_keeps_the_instant(con, literal): + con.execute("SET TimeZone = 'Europe/Amsterdam'") + value = fetch(con, literal) + instant = datetime.datetime.fromisoformat(literal.replace(" ", "T") + ":00") + assert value.replace(tzinfo=None) == datetime.datetime(2024, 10, 27, 2, 30) + assert value.fold == (0 if literal.startswith("2024-10-27 00") else 1) + assert value.utcoffset() == datetime.timedelta(hours=2 if value.fold == 0 else 1) + assert value.timestamp() == instant.timestamp() + assert value.astimezone(UTC) == instant + + +def test_values_of_one_session_share_the_zone_and_compare_by_wall_clock(con): + con.execute("SET TimeZone = 'Europe/Amsterdam'") + first, second = con.execute( + "SELECT TIMESTAMPTZ '2024-10-27 00:30:00+00', TIMESTAMPTZ '2024-10-27 01:30:00+00'" + ).fetchone() + assert first.tzinfo is second.tzinfo + assert first == second + assert first.fold == 0 + assert second.fold == 1 + assert first.astimezone(UTC) != second.astimezone(UTC) + assert second.timestamp() - first.timestamp() == 3600 + keyed = con.execute( + "SELECT MAP {TIMESTAMPTZ '2024-10-27 00:30:00+00': 1, TIMESTAMPTZ '2024-10-27 01:30:00+00': 2}" + ).fetchone()[0] + assert keyed == {second: 2} + + +def test_timestamptz_ns_takes_the_same_path(con): + con.execute("SET TimeZone = 'Asia/Tokyo'") + value = con.execute("SELECT '2024-01-01 00:00:00.123456789+00'::TIMESTAMPTZ_NS").fetchone()[0] + assert value.tzinfo is zoneinfo.ZoneInfo("Asia/Tokyo") + assert value.replace(tzinfo=None) == datetime.datetime(2024, 1, 1, 9, 0, 0, 123456) + assert value == datetime.datetime(2024, 1, 1, 0, 0, 0, 123456, tzinfo=UTC) + + +def test_dst_offset_after_2037(con): + con.execute("SET TimeZone = 'America/Chicago'") + assert fetch(con, "2037-08-01 08:00:00+00").isoformat() == "2037-08-01T03:00:00-05:00" + assert fetch(con, "2038-08-01 08:00:00+00").isoformat() == "2038-08-01T03:00:00-05:00" + assert fetch(con, "2050-01-15 08:00:00+00").isoformat() == "2050-01-15T02:00:00-06:00" + + +def test_nested_values_carry_the_zone(con): + con.execute("SET TimeZone = 'Europe/Amsterdam'") + literal = "TIMESTAMPTZ '2024-07-01 00:00:00+00'" + row = con.execute(f"SELECT [{literal}], {{'t': {literal}}}, MAP {{'k': {literal}}}").fetchone() + expected = datetime.datetime(2024, 7, 1, 2, tzinfo=zoneinfo.ZoneInfo("Europe/Amsterdam")) + assert row == ([expected], {"t": expected}, {"k": expected}) + assert row[0][0].tzinfo is expected.tzinfo + + +def test_udf_argument_is_aware(con): + pytest.importorskip("numpy") + con.execute("SET TimeZone = 'Asia/Tokyo'") + con.create_function( + "describe_tz", lambda t: f"{t.isoformat()} {type(t.tzinfo).__name__} {t.tzinfo.key}", ["TIMESTAMPTZ"], "VARCHAR" + ) + assert con.execute("SELECT describe_tz(TIMESTAMPTZ '2024-01-01 00:00:00+00')").fetchone()[0] == ( + "2024-01-01T09:00:00+09:00 ZoneInfo Asia/Tokyo" + ) + + +def test_icu_only_zone_name_raises(con): + con.execute("SET TimeZone = 'ACT'") + with pytest.raises(duckdb.InvalidInputException, match=r"'ACT'.*tzdata"): + fetch(con, "2024-01-01 00:00:00+00") + + +def test_without_a_tz_database_raises(tmp_path): + script = textwrap.dedent( + """ + import sys + + sys.modules["tzdata"] = None + + import duckdb + con = duckdb.connect() + con.execute("SET TimeZone = 'Europe/Amsterdam'") + try: + con.execute("SELECT TIMESTAMPTZ '2024-01-01 00:00:00+00'").fetchone() + except duckdb.InvalidInputException as error: + assert "'Europe/Amsterdam'" in str(error) and "tzdata" in str(error), repr(error) + else: + raise AssertionError("a fetch without a tz database did not raise") + """ + ) + env = {**os.environ, "PYTHONTZPATH": ""} + proc = subprocess.run( + [sys.executable, "-c", script], env=env, cwd=tmp_path, capture_output=True, text=True, timeout=60 + ) + assert proc.returncode == 0, proc.stdout + proc.stderr + + +@pytest.mark.parametrize( + ("zone", "literal"), + [ + ("Asia/Tokyo", "9999-12-31 23:00:00+00"), + ("America/Los_Angeles", "0001-01-01 00:00:00+00"), + ], +) +def test_local_wall_time_outside_python_range_falls_back_to_str(con, zone, literal): + con.execute(f"SET TimeZone = '{zone}'") + assert fetch(con, literal) == literal + + +def test_utc_wall_time_outside_python_range_is_already_str(con): + con.execute("SET TimeZone = 'America/Los_Angeles'") + assert fetch(con, "9999-12-31 23:59:59.999999-08") == "10000-01-01 07:59:59.999999+00" + + +def test_local_mean_time_offset_keeps_seconds(con): + con.execute("SET TimeZone = 'Asia/Tokyo'") + value = fetch(con, "0001-01-01 00:00:00+00") + assert value == datetime.datetime(1, 1, 1, tzinfo=UTC) + assert value.utcoffset() == datetime.timedelta(hours=9, minutes=18, seconds=59) + + +def test_infinity_and_null_are_unchanged(con): + con.execute("SET TimeZone = 'Europe/Amsterdam'") + row = con.execute("SELECT TIMESTAMPTZ '-infinity', TIMESTAMPTZ 'infinity', NULL::TIMESTAMPTZ").fetchone() + assert row == (datetime.datetime.min, datetime.datetime.max, None) + assert row[0].tzinfo is None + assert row[1].tzinfo is None