From c1802141e86aa8051b6d0108dae2890c80078c38 Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Fri, 6 Nov 2026 01:07:41 +0000 Subject: [PATCH 1/9] Make map() keep its column types, ENUMs included --- .../numpy/numpy_result_conversion.hpp | 3 + .../duckdb_python/pandas/pandas_bind.hpp | 12 +- src/include/duckdb_python/pyresult.hpp | 2 - src/map.cpp | 162 +++++-- src/numpy/numpy_result_conversion.cpp | 15 + src/numpy/numpy_scan.cpp | 27 +- src/pandas/bind.cpp | 22 +- src/pyresult.cpp | 25 +- tests/fast/pandas/test_append_df.py | 21 + tests/fast/test_map_types.py | 416 ++++++++++++++++++ 10 files changed, 630 insertions(+), 75 deletions(-) create mode 100644 tests/fast/test_map_types.py diff --git a/src/include/duckdb_python/numpy/numpy_result_conversion.hpp b/src/include/duckdb_python/numpy/numpy_result_conversion.hpp index f068dc79..6425d9e6 100644 --- a/src/include/duckdb_python/numpy/numpy_result_conversion.hpp +++ b/src/include/duckdb_python/numpy/numpy_result_conversion.hpp @@ -38,4 +38,7 @@ class NumpyResultConversion { bool pandas; }; +//! The dtype df() gives an ENUM column: a pandas CategoricalDtype over the enum's values in order, ordered +nb::object EnumCategoricalDtype(const LogicalType &enum_type); + } // namespace duckdb diff --git a/src/include/duckdb_python/pandas/pandas_bind.hpp b/src/include/duckdb_python/pandas/pandas_bind.hpp index 7a9fce9c..ee5a6c3c 100644 --- a/src/include/duckdb_python/pandas/pandas_bind.hpp +++ b/src/include/duckdb_python/pandas/pandas_bind.hpp @@ -42,11 +42,19 @@ struct PandasColumnBindData { bool declared = false; }; +//! The categories of a CategoricalDtype the caller built from an ENUM's values, in order +struct PandasEnumDtype { + nb::handle categories; + LogicalType type; +}; + struct Pandas { - //! Returns the frame's row count; declared_types, when given, hold one type per column, INVALID where none + //! Returns the frame's row count; declared_types, when given, hold one type per column, INVALID where none; a + //! categorical column whose categories are those of one of enum_dtypes, that very object, reads its codes as that + //! ENUM static idx_t Bind(ClientContext &config, nb::handle df, vector &out, vector &return_types, vector &names, - const vector &declared_types = {}); + const vector &declared_types = {}, const vector &enum_dtypes = {}); }; } // namespace duckdb diff --git a/src/include/duckdb_python/pyresult.hpp b/src/include/duckdb_python/pyresult.hpp index abad6e35..50157de2 100644 --- a/src/include/duckdb_python/pyresult.hpp +++ b/src/include/duckdb_python/pyresult.hpp @@ -136,8 +136,6 @@ struct DuckDBPyResult { //! and whose originals the caller carries here. Empty means the result's own names. vector names_override; unique_ptr current_chunk; - // Holds the categories of Categorical/ENUM types - unordered_map categories; // Holds the categorical type of Categorical/ENUM types unordered_map categories_type; }; diff --git a/src/map.cpp b/src/map.cpp index e9724243..e3746f16 100644 --- a/src/map.cpp +++ b/src/map.cpp @@ -26,15 +26,30 @@ struct MapFunctionData : public TableFunctionData { shared_ptr info; vector in_types, out_types; vector in_names, out_names; + //! Per output column, whether its type is only a guess from the typed input column of its name, which each chunk + //! must confirm, rather than imposed by the schema or an object input + vector guessed; + //! Per input column, the CategoricalDtype an ENUM column is handed over in, built once since it lists every value + vector enum_dtypes; + //! The categories of those dtypes with their ENUM types, by which a returned column of them reads its codes + vector enum_categories; + vector known_enums; }; -static nb::object FunctionCall(NumpyResultConversion &conversion, const vector &names, PyObject *function) { +static nb::object FunctionCall(const MapFunctionData &data, NumpyResultConversion &conversion) { + auto &import_cache = *DuckDBPyConnection::ImportCache(); nb::dict in_numpy_dict; - for (idx_t col_idx = 0; col_idx < names.size(); col_idx++) { - in_numpy_dict[names[col_idx].c_str()] = conversion.ToArray(col_idx); + for (idx_t col_idx = 0; col_idx < data.in_names.size(); col_idx++) { + auto column = conversion.ToArray(col_idx); + // As df() gives it: the codes alone bind back as TINYINT, a NULL as the code -1. + if (data.in_types[col_idx].id() == LogicalTypeId::ENUM) { + column = import_cache.pandas.Categorical().attr("from_codes")(column, nb::arg("dtype") = + data.enum_dtypes[col_idx].obj); + } + in_numpy_dict[data.in_names[col_idx].c_str()] = column; } + auto function = data.info->function.obj.ptr(); - auto &import_cache = *DuckDBPyConnection::ImportCache(); auto pandas_df = import_cache.pandas.DataFrame(); auto in_df = pandas_df(in_numpy_dict); D_ASSERT(in_df.ptr()); @@ -65,39 +80,37 @@ static nb::object FunctionCall(NumpyResultConversion &conversion, const vector &types) { - for (auto &type : types) { - if (type.id() == LogicalTypeId::SQLNULL) { - return true; +//! The input column of that name, the one at the same position first; the input names are deduplicated at bind +static optional_idx InputColumnNamed(const vector &input_names, const Identifier &name, idx_t position) { + if (position < input_names.size() && input_names[position] == name) { + return position; + } + for (idx_t col_idx = 0; col_idx < input_names.size(); col_idx++) { + if (input_names[col_idx] == name) { + return col_idx; } } - return false; + return optional_idx(); } +//! An object column of the empty bind frame has no values to type it, so it takes the type of the input column of its +//! name. That type is imposed when the function got the input as an object column too; from a typed input it is only +//! a guess, since the function may have moved the column to another type, and casting back would lose values. static void OverrideNullType(vector &return_types, const vector &return_names, - const vector &original_types, const vector &original_names) { - if (!ContainsNullType(return_types)) { - // Nothing to override, none of the returned types are NULL - return; - } - if (return_types.size() != original_types.size()) { - // FIXME: we can probably infer from the names in this case - // Cant infer what the type should be - return; - } - for (idx_t i = 0; i < return_types.size(); i++) { - auto &return_type = return_types[i]; - auto &original_type = original_types[i]; - - if (return_type != LogicalTypeId::SQLNULL) { + const vector &input_types, const vector &input_names, + vector &guessed) { + guessed.assign(return_types.size(), false); + for (idx_t col_idx = 0; col_idx < return_types.size(); col_idx++) { + if (return_types[col_idx].id() != LogicalTypeId::SQLNULL) { continue; } - if (return_names[i] != original_names[i]) { - throw InvalidInputException( - "Returned dataframe contains NULL type, and we could not infer the desired type"); + auto input = InputColumnNamed(input_names, return_names[col_idx], col_idx); + if (!input.IsValid()) { + continue; } - // Override the NULL with the original type - return_type = original_type; + auto &input_type = input_types[input.GetIndex()]; + return_types[col_idx] = input_type; + guessed[col_idx] = RawArrayWrapper::DuckDBToNumpyDtype(input_type) != "object"; } } @@ -135,6 +148,7 @@ unique_ptr BindExplicitSchema(unique_ptr function function_data->out_names = names; function_data->out_types = types; + function_data->guessed.assign(types.size(), false); return std::move(function_data); } @@ -146,22 +160,31 @@ static unique_ptr BindMap(ClientContext &context, shared_ptr &names) { auto data_uptr = make_uniq(std::move(info)); auto &data = *data_uptr; - auto function = data.info->function.obj.ptr(); auto explicit_schema = data.info->schema.obj.ptr(); data.in_names = input.input_table_names; + // The frame is keyed by name, so two input columns of one name would leave the function one of them + QueryResult::DeduplicateColumns(data.in_names); data.in_types = input.input_table_types; + for (auto &type : data.in_types) { + if (type.id() != LogicalTypeId::ENUM) { + data.enum_dtypes.emplace_back(); + continue; + } + data.enum_dtypes.emplace_back(EnumCategoricalDtype(type)); + data.enum_categories.emplace_back(data.enum_dtypes.back().obj.attr("categories")); + data.known_enums.push_back({data.enum_categories.back().obj, type}); + } if (explicit_schema != Py_None) { return BindExplicitSchema(std::move(data_uptr), explicit_schema, return_types, names); } NumpyResultConversion conversion(data.in_types, 0, context.GetClientProperties()); - auto df = PandasScanFunction::PandasWithBoundNames(FunctionCall(conversion, data.in_names, function)); + auto df = PandasScanFunction::PandasWithBoundNames(FunctionCall(data, conversion)); vector pandas_bind_data; // unused - Pandas::Bind(context, df, pandas_bind_data, return_types, names); + Pandas::Bind(context, df, pandas_bind_data, return_types, names, {}, data.known_enums); - // output types are potentially NULL, this happens for types that map to 'object' dtype - OverrideNullType(return_types, names, data.in_types, data.in_names); + OverrideNullType(return_types, names, data.in_types, data.in_names, data.guessed); data.out_names = names; data.out_types = return_types; @@ -184,6 +207,12 @@ unique_ptr MapFunction::MapFunctionBind(ClientContext &context, Ta } } +static InvalidInputException UntypedColumnError(const Identifier &name) { + return InvalidInputException( + "UDF column '%s' has no type: the function returned no values for it at bind time; pass its type with schema=", + name.c_str()); +} + static string TypeVectorToString(const vector &types) { return StringUtil::Join(types, types.size(), ", ", [](const LogicalType &argument) { return argument.ToString(); }); } @@ -200,18 +229,55 @@ static OperatorResultType ExecuteMap(ExecutionContext &context, TableFunctionInp NumpyResultConversion conversion(data.in_types, input.size(), context.client.GetClientProperties()); conversion.Append(input); - auto df = PandasScanFunction::PandasWithBoundNames( - FunctionCall(conversion, data.in_names, data.info->function.obj.ptr())); + auto df = PandasScanFunction::PandasWithBoundNames(FunctionCall(data, conversion)); + auto column_count = nb::len(nb::object(df.attr("columns"))); + if (column_count != output.ColumnCount()) { + throw InvalidInputException("Expected %llu columns from UDF, got %llu", output.ColumnCount(), column_count); + } + // Before anything types the frame, which could fail on a value it cannot read + for (idx_t col_idx = 0; col_idx < output.ColumnCount(); col_idx++) { + if (data.out_types[col_idx].id() != LogicalTypeId::SQLNULL) { + continue; + } + auto column = df.attr("__getitem__")(nb::object(df.attr("columns"))[nb::int_(col_idx)]); + auto has_value = PyObject_IsTrue(nb::object(column.attr("notna")().attr("any")()).ptr()); + if (has_value < 0) { + throw nb::python_error(); + } + if (has_value) { + throw UntypedColumnError(data.out_names[col_idx]); + } + } + // An object column has no type of its own, only values, so it is read into the type the schema or an object input + // imposed, whatever its values. A guessed type and a column bind could not type, which holds only missing values + // here, are left to the analyzer. + vector declared_types; + for (idx_t col_idx = 0; col_idx < output.ColumnCount(); col_idx++) { + auto &type = data.out_types[col_idx]; + auto imposed = type.id() != LogicalTypeId::SQLNULL && !data.guessed[col_idx]; + declared_types.push_back(imposed ? type : LogicalType::INVALID); + } vector pandas_bind_data; vector pandas_return_types; vector pandas_names; + auto row_count = Pandas::Bind(context.client, df, pandas_bind_data, pandas_return_types, pandas_names, + declared_types, data.known_enums); - auto row_count = Pandas::Bind(context.client, df, pandas_bind_data, pandas_return_types, pandas_names); - if (pandas_return_types.size() != output.ColumnCount()) { - throw InvalidInputException("Expected %llu columns from UDF, got %llu", output.ColumnCount(), - pandas_return_types.size()); + // A guess is only checked: a chunk of only missing values confirms nothing and reads NULL, any other type refutes + // it, and the analyzer's sample decides which, as it did before for every column. + vector missing_values(output.ColumnCount(), false); + for (idx_t col_idx = 0; col_idx < output.ColumnCount(); col_idx++) { + if (!data.guessed[col_idx] || pandas_return_types[col_idx] == data.out_types[col_idx]) { + continue; + } + if (pandas_return_types[col_idx].id() != LogicalTypeId::SQLNULL) { + throw UntypedColumnError(data.out_names[col_idx]); + } + pandas_return_types[col_idx] = data.out_types[col_idx]; + missing_values[col_idx] = true; } + D_ASSERT(output.GetTypes() == data.out_types); if (pandas_return_types != output.GetTypes()) { throw InvalidInputException("UDF column type mismatch, expected [%s], got [%s]", @@ -227,9 +293,23 @@ static OperatorResultType ExecuteMap(ExecutionContext &context, TableFunctionInp } for (idx_t col_idx = 0; col_idx < output.ColumnCount(); col_idx++) { + if (missing_values[col_idx]) { + ConstantVector::SetNull(output.data[col_idx], count_t(row_count)); + continue; + } auto &bind_data = pandas_bind_data[col_idx]; Vector temp(output.data[col_idx].GetType(), row_count); - PandasScanFunction::PandasBackendScanSwitch(context.client, bind_data, row_count, 0, temp); + try { + PandasScanFunction::PandasBackendScanSwitch(context.client, bind_data, row_count, 0, temp); + } catch (std::exception &ex) { + ErrorData error(ex); + // The cast's own message names no column; a Python error keeps the traceback its message starts with + if (!bind_data.declared || error.Type() == ExceptionType::INVALID || + error.Type() == ExceptionType::INTERRUPT) { + throw; + } + error.Throw(StringUtil::Format("UDF column '%s': ", data.out_names[col_idx].c_str())); + } VectorOperations::Copy(temp, output.data[col_idx], row_count, 0, 0); } output.SetChildCardinality(row_count); diff --git a/src/numpy/numpy_result_conversion.cpp b/src/numpy/numpy_result_conversion.cpp index f511baed..db331516 100644 --- a/src/numpy/numpy_result_conversion.cpp +++ b/src/numpy/numpy_result_conversion.cpp @@ -1,5 +1,6 @@ #include "duckdb_python/numpy/array_wrapper.hpp" #include "duckdb_python/numpy/numpy_result_conversion.hpp" +#include "duckdb_python/pyconnection/pyconnection.hpp" namespace duckdb { @@ -46,4 +47,18 @@ void NumpyResultConversion::Append(DataChunk &chunk) { #endif } +nb::object EnumCategoricalDtype(const LogicalType &enum_type) { + DUCKDB_PY_ASSERT_GIL_HELD(); + auto categorical_dtype = DuckDBPyConnection::ImportCache()->pandas.CategoricalDtype(); + if (!categorical_dtype) { + throw InvalidInputException("'pandas' is required for this operation but it was not installed"); + } + auto &values = EnumType::GetValuesInsertOrder(enum_type); + nb::list categories; + for (idx_t i = 0; i < EnumType::GetSize(enum_type); i++) { + categories.append(nb::cast(values.GetValue(i).ToString())); + } + return categorical_dtype(categories, true); +} + } // namespace duckdb diff --git a/src/numpy/numpy_scan.cpp b/src/numpy/numpy_scan.cpp index ec3cc615..d9ff6fe9 100644 --- a/src/numpy/numpy_scan.cpp +++ b/src/numpy/numpy_scan.cpp @@ -352,7 +352,12 @@ static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb //! shape of its own (a UNION member that is NULL), and the output owns every string the cast produced. class DeclaredVector { public: - explicit DeclaredVector(idx_t count) : values(count), source(LogicalTypeId::SQLNULL) { + explicit DeclaredVector(idx_t count) : values(count), written(count, false), source(LogicalTypeId::SQLNULL) { + } + + //! A row the scan wrote into the output itself, which Write leaves as it is + void MarkWritten(idx_t row) { + written[row] = true; } void Set(idx_t row, Value value) { @@ -389,6 +394,10 @@ class DeclaredVector { VectorOperations::TryCast(context, staged, casted, present, nullptr); idx_t copied = 0; for (idx_t row = 0; row < count;) { + if (written[row]) { + row++; + continue; + } if (values[row].IsNull()) { FlatVector::SetNull(out, row, true); row++; @@ -406,6 +415,9 @@ class DeclaredVector { } DeclaredCasts casts(context, target); for (idx_t row = 0; row < count; row++) { + if (written[row]) { + continue; + } if (values[row].IsNull()) { FlatVector::SetNull(out, row, true); } else { @@ -416,6 +428,7 @@ class DeclaredVector { private: vector values; + vector written; LogicalType source; bool uniform = true; }; @@ -461,8 +474,16 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t idx_t count, idx_t offset, Vector &out) { DeclaredVector staged(count); auto &target = out.GetType(); - ScanObjectCells(context, col, stride, masked, count, offset, out, - [&](nb::handle cell, idx_t row) { staged.Set(row, DeclaredValue(context, target, cell)); }); + // A str into plain VARCHAR reads as TransformPythonValue would read it, without building a Value per cell + const bool plain_text = target == LogicalType::VARCHAR; + ScanObjectCells(context, col, stride, masked, count, offset, out, [&](nb::handle cell, idx_t row) { + if (plain_text && PyUnicode_CheckExact(cell.ptr())) { + FlatVector::GetDataMutable(out)[row] = StringVector::AddString(out, TransformPythonString(cell)); + staged.MarkWritten(row); + return; + } + staged.Set(row, DeclaredValue(context, target, cell)); + }); staged.Write(context, out, count); VerifyTypeConstraints(out, count); } diff --git a/src/pandas/bind.cpp b/src/pandas/bind.cpp index 73b20205..e998558b 100644 --- a/src/pandas/bind.cpp +++ b/src/pandas/bind.cpp @@ -120,7 +120,8 @@ std::unique_ptr RegisterMask(nb::handle mask) { } static LogicalType BindColumn(ClientContext &context, PandasBindColumn &column_p, PandasColumnBindData &bind_data, - idx_t row_count, const LogicalType &declared_type) { + idx_t row_count, const LogicalType &declared_type, + const vector &enum_dtypes) { const string counted_from = "the frame"; LogicalType column_type; auto &column = column_p.handle; @@ -137,6 +138,18 @@ static LogicalType BindColumn(ClientContext &context, PandasBindColumn &column_p nb::object data; if (bind_data.numpy_type.type == NumpyNullableType::CATEGORY) { + auto column_categories = enum_dtypes.empty() ? nb::object() : nb::object(column_p.type.attr("categories")); + for (auto &known : enum_dtypes) { + // A categories index is immutable, so this object still holds the enum's values, and reading them again + // would cost a string per category on every call; pandas 2 builds a new dtype around the same index + if (column_categories.is(known.categories)) { + data = column.attr("cat").attr("codes").attr("to_numpy")(); + bind_data.internal_categorical_type = nb::cast(nb::str(nb::object(data.attr("dtype")))); + CheckColumnArray(bind_data.name, data, row_count, counted_from); + bind_data.pandas_col = std::make_unique(NumpyArray(std::move(data))); + return known.type; + } + } // for category types, we create an ENUM type for string or use the converted numpy type for the rest D_ASSERT(nb::hasattr(column, "cat")); D_ASSERT(nb::hasattr(column.attr("cat"), "categories")); @@ -223,7 +236,7 @@ static LogicalType BindColumn(ClientContext &context, PandasBindColumn &column_p idx_t Pandas::Bind(ClientContext &context, nb::handle df_p, vector &bind_columns, vector &return_types, vector &names, - const vector &declared_types) { + const vector &declared_types, const vector &enum_dtypes) { PandasDataFrameBind df(df_p); idx_t column_count = nb::len(df.names); @@ -256,8 +269,9 @@ idx_t Pandas::Bind(ClientContext &context, nb::handle df_p, vector &types, unordered_map &categories) { - for (idx_t col_idx = 0; col_idx < types.size(); col_idx++) { - auto &type = types[col_idx]; - if (type.id() == LogicalTypeId::ENUM) { - // It's an ENUM type, in addition to converting the codes we must convert the categories - if (categories.find(col_idx) == categories.end()) { - auto &categories_list = EnumType::GetValuesInsertOrder(type); - auto categories_size = EnumType::GetSize(type); - for (idx_t i = 0; i < categories_size; i++) { - categories[col_idx].append(nb::cast(categories_list.GetValue(i).ToString())); - } - } - } - } -} - std::unique_ptr DuckDBPyResult::InitializeNumpyConversion(bool pandas) { if (Empty()) { throw InvalidInputException("result closed"); @@ -313,7 +293,6 @@ nb::dict DuckDBPyResult::FetchNumpyInternal(bool chunked, idx_t vectors_per_chun } conversion.Append(*chunk); } - InsertCategory(GetTypes(), categories); // now that we have materialized the result in contiguous arrays, construct the actual NumPy arrays or // categorical types diff --git a/tests/fast/pandas/test_append_df.py b/tests/fast/pandas/test_append_df.py index 4360a8bc..0f3099f8 100644 --- a/tests/fast/pandas/test_append_df.py +++ b/tests/fast/pandas/test_append_df.py @@ -428,6 +428,27 @@ def test_text_with_a_lone_surrogate_names_it(self, duckdb_cursor): with pytest.raises(duckdb.InvalidInputException, match="surrogate"): duckdb_cursor.append("t", frame) + def test_text_column_mixing_str_with_other_cells_appends_each_as_text(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (a VARCHAR)") + cells = ["a", 5, 6, "b", None, 7, np.str_("c"), "\u00e9 \U0001f600"] + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series(cells, dtype=object)})) + assert duckdb_cursor.table("t").fetchall() == [ + ("a",), + ("5",), + ("6",), + ("b",), + (None,), + ("7",), + ("c",), + ("\u00e9 \U0001f600",), + ] + + def test_text_with_a_lone_surrogate_in_a_str_names_it(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (a VARCHAR)") + frame = pd.DataFrame({"a": pd.Series(["a", "b\udcff"], dtype=object)}) + with pytest.raises(duckdb.InvalidInputException, match="surrogate"): + duckdb_cursor.append("t", frame) + def test_null_row_among_struct_values_the_cast_converts(self, duckdb_cursor): duckdb_cursor.execute("CREATE TABLE t (a STRUCT(x INTEGER))") # Text fields cast into INTEGER, so the whole vector goes through one cast with the NULL row inside it diff --git a/tests/fast/test_map_types.py b/tests/fast/test_map_types.py new file mode 100644 index 00000000..8672d4ca --- /dev/null +++ b/tests/fast/test_map_types.py @@ -0,0 +1,416 @@ +"""map() fixes one type per output column at bind and reads each chunk's returned column into it.""" + +import re + +import pytest +from conftest import PANDAS_GE_3 + +import duckdb +import numpy as np +import pandas as pd + +ENUM = "ENUM('a', 'b')" +SCHEMAS = pytest.mark.parametrize("explicit_schema", [False, True], ids=["inferred", "explicit_schema"]) + + +def _schema_of(relation, explicit_schema): + return dict(zip(relation.columns, relation.types, strict=True)) if explicit_schema else None + + +@pytest.mark.parametrize("with_null", [False, True], ids=["no_null", "a_null"]) +def test_map_enum_column_reaches_the_function_as_df_gives_it(duckdb_cursor, with_null): + second = "NULL" if with_null else f"'a'::{ENUM}" + source = duckdb_cursor.sql(f"select * from (values ('b'::{ENUM}), ({second})) t(e)") + seen = [] + + def record(df): + seen.append((len(df), df["e"].dtype, [None if pd.isna(v) else v for v in df["e"].tolist()])) + return df + + mapped = source.map(record) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + categorical = pd.CategoricalDtype(["a", "b"], ordered=True) + assert source.df()["e"].dtype == categorical + assert all(dtype == categorical for _, dtype, _ in seen), seen + assert [values for rows, _, values in seen if rows] == [["b", None if with_null else "a"]] + + +def test_map_enum_column_reaches_every_call_with_the_dtype_built_at_bind(duckdb_cursor): + duckdb_cursor.execute("set threads = 1") + seen = [] + + def record(df): + if len(df): + seen.append(df["e"].dtype) + return df + + source = duckdb_cursor.sql(f"select (case when i % 3 = 0 then 'a' else 'b' end)::{ENUM} as e from range(4096) t(i)") + assert source.map(record).fetchall() == source.fetchall() + assert len(seen) == 2 + # pandas 3 keeps the dtype object; pandas 2 builds a new dtype around the same categories index + assert all(dtype.categories is seen[0].categories for dtype in seen) + if PANDAS_GE_3: + assert all(dtype is seen[0] for dtype in seen) + + +@pytest.mark.parametrize( + ("categories", "ordered", "result_type"), + [ + pytest.param(["a", "b"], True, ENUM, id="equal_categories"), + pytest.param(["a", "b"], False, ENUM, id="equal_categories_unordered"), + pytest.param(["b", "a"], True, "ENUM('b', 'a')", id="reordered_categories"), + ], +) +def test_map_enum_column_rebuilt_as_a_new_categorical_reads_by_its_categories( + duckdb_cursor, categories, ordered, result_type +): + def rebuild(df): + return df.assign(e=pd.Categorical(df["e"].astype(object), categories=categories, ordered=ordered)) + + source = duckdb_cursor.sql(f"select * from (values ('b'::{ENUM}), (NULL), ('a'::{ENUM})) t(e)") + mapped = source.map(rebuild) + assert mapped.types == [duckdb.sqltype(result_type)] + assert mapped.fetchall() == [("b",), (None,), ("a",)] + + +def test_map_enum_column_compares_with_its_text(duckdb_cursor): + source = duckdb_cursor.sql(f"select * from (values ('b'::{ENUM}), (NULL), ('a'::{ENUM})) t(e)") + mapped = source.map(lambda df: df.assign(is_b=df["e"] == "b")) + assert mapped.types == [source.types[0], duckdb.sqltypes.BOOLEAN] + assert mapped.fetchall() == [("b", True), (None, False), ("a", False)] + + +@SCHEMAS +@pytest.mark.parametrize( + ("sql_type", "value"), + [ + pytest.param("STRUCT(x BIGINT)", "{'x': i::BIGINT}", id="STRUCT_BIGINT_field"), + pytest.param("MAP(INTEGER, INTEGER)", "map([i::INTEGER], [i::INTEGER])", id="MAP_INTEGER_keys"), + pytest.param("MAP(VARCHAR, BIGINT)", "map([i::VARCHAR], [i::BIGINT])", id="MAP_BIGINT_values"), + pytest.param("INTEGER[2]", "[i, i]::INTEGER[2]", id="ARRAY"), + pytest.param("UNION(a INTEGER, b VARCHAR)", "i::INTEGER::UNION(a INTEGER, b VARCHAR)", id="UNION"), + pytest.param("TIME_NS", "'01:00:00'::TIME_NS", id="TIME_NS"), + pytest.param( + "STRUCT(a INTEGER[], b MAP(VARCHAR, INTEGER))", + "{'a': [i::INTEGER], 'b': map(['k'], [i::INTEGER])}", + id="nested", + ), + ], +) +def test_map_identity_returns_an_object_column_in_its_own_type(duckdb_cursor, sql_type, value, explicit_schema): + source = duckdb_cursor.sql( + f"select (case when i = 1 then NULL else {value} end)::{sql_type} as a from range(3) t(i)" + ) + mapped = source.map(lambda df: df, schema=_schema_of(source, explicit_schema)) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + + +@SCHEMAS +def test_map_identity_returns_a_struct_whose_field_is_null_across_a_chunk(duckdb_cursor, explicit_schema): + duckdb_cursor.execute("set threads = 1") + # The field holds values in the first 2048 row chunk and only NULLs in the second. + source = duckdb_cursor.sql("select {'a': case when i < 2048 then i::INTEGER end} as s from range(4096) t(i)") + mapped = source.map(lambda df: df, schema=_schema_of(source, explicit_schema)) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + + +@SCHEMAS +def test_map_identity_returns_an_enum_column_in_its_own_type(duckdb_cursor, explicit_schema): + source = duckdb_cursor.sql(f"select * from (values ('b'::{ENUM}), (NULL), ('a'::{ENUM})) t(e)") + mapped = source.map(lambda df: df, schema=_schema_of(source, explicit_schema)) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + + +def test_map_column_takes_the_input_type_of_its_name_when_columns_are_added_or_reordered(duckdb_cursor): + source = duckdb_cursor.sql("select i as n, {'x': i::BIGINT} as s from range(3) t(i)") + mapped = source.map(lambda df: df[["s", "n"]].assign(k=1)) + assert mapped.types == [source.types[1], duckdb.sqltypes.BIGINT, duckdb.sqltypes.BIGINT] + assert mapped.fetchall() == [(s, n, 1) for n, s in source.fetchall()] + + +def _text_cells(values): + def replace(df): + return df.assign(a=pd.Series(values[: len(df)], dtype=object, index=df.index)) if len(df) else df + + return replace + + +MISSING_VALUES = [ + pytest.param(lambda: None, id="None"), + pytest.param(duckdb.NullValue, id="NullValue"), + pytest.param(lambda: np.datetime64("NaT", "us"), id="NaT"), + pytest.param(lambda: pd.NA, id="pd.NA"), + pytest.param(lambda: float("nan"), id="NaN"), +] + + +def _missing_cells(make_missing): + def missing(df): + return df.assign(a=pd.Series([make_missing() for _ in range(len(df))], dtype=object)) if len(df) else df + + return missing + + +@pytest.mark.parametrize("sql_type", ["BIGINT", "DOUBLE", "VARCHAR"]) +@pytest.mark.parametrize("make_missing", MISSING_VALUES) +def test_map_object_column_of_missing_values_reads_null_in_its_schema_type(duckdb_cursor, make_missing, sql_type): + source = duckdb_cursor.sql(f"select i::{sql_type} as a from range(3) t(i)") + mapped = source.map(_missing_cells(make_missing), schema=_schema_of(source, explicit_schema=True)) + assert mapped.types == source.types + assert mapped.fetchall() == [(None,), (None,), (None,)] + + +@pytest.mark.parametrize("make_missing", MISSING_VALUES) +def test_map_object_column_of_missing_values_keeps_an_object_input_type(duckdb_cursor, make_missing): + source = duckdb_cursor.sql("select i::VARCHAR as a from range(3) t(i)") + mapped = source.map(_missing_cells(make_missing)) + assert mapped.types == source.types + assert mapped.fetchall() == [(None,), (None,), (None,)] + + +def _odd_cell_at(row, value): + """Every cell of column a as its Python value, and value at row.""" + + def replace(df): + if not len(df): + return df + cells = df["a"].astype(object).tolist() + cells[row] = value + return df.assign(a=pd.Series(cells, dtype=object, index=df.index)) + + return replace + + +# In a 2048 row chunk the analyzer samples the even rows only, so the last row is outside its sample. +ODD_ROWS = pytest.mark.parametrize("row", [0, 2047], ids=["first_row", "row_outside_the_sample"]) +BIGINT_A = {"a": duckdb.sqltypes.BIGINT} + + +@ODD_ROWS +def test_map_object_column_cell_is_cast_into_the_bound_type(duckdb_cursor, row): + duckdb_cursor.execute("set threads = 1") + mapped = duckdb_cursor.sql("select i as a from range(2048) t(i)").map(_odd_cell_at(row, "5"), schema=BIGINT_A) + assert mapped.types == [duckdb.sqltypes.BIGINT] + assert mapped.fetchall() == [(5 if i == row else i,) for i in range(2048)] + + +@ODD_ROWS +def test_map_object_column_cell_the_cast_refuses_raises_the_casts_error(duckdb_cursor, row): + duckdb_cursor.execute("set threads = 1") + mapped = duckdb_cursor.sql("select i as a from range(2048) t(i)").map(_odd_cell_at(row, "x"), schema=BIGINT_A) + expected = "Conversion Error: UDF column 'a': Could not convert string 'x' to INT64" + with pytest.raises(duckdb.ConversionException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +@ODD_ROWS +def test_map_struct_cell_without_a_field_reads_it_as_null(duckdb_cursor, row): + duckdb_cursor.execute("set threads = 1") + source = duckdb_cursor.sql("select {'x': i::INTEGER, 'y': i::INTEGER} as a from range(2048) t(i)") + mapped = source.map(_odd_cell_at(row, {"x": 3})) + expected = [({"x": 3, "y": None} if i == row else {"x": i, "y": i},) for i in range(2048)] + assert mapped.types == source.types + assert mapped.fetchall() == expected + + +def test_map_object_column_of_zero_dimensional_arrays_reads_their_values(duckdb_cursor): + def arrays(df): + cells = np.empty(len(df), dtype=object) + for i, v in enumerate(df["a"]): + cells[i] = np.array(int(v), dtype=object) + return df.assign(a=pd.Series(cells, dtype=object, index=df.index)) + + mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map(arrays, schema={"a": duckdb.sqltypes.BIGINT}) + assert mapped.fetchall() == [(0,), (1,), (2,)] + + +def test_map_declared_column_read_error_names_the_column(duckdb_cursor): + source = duckdb_cursor.sql("select i as n, {'x': i::INTEGER} as s from range(3) t(i)") + mapped = source.map(lambda df: df.assign(n=df["s"]), schema=_schema_of(source, explicit_schema=True)) + expected = "Conversion Error: UDF column 'n': Unimplemented type for cast (STRUCT(x INTEGER) -> BIGINT)" + with pytest.raises(duckdb.ConversionException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +@pytest.mark.parametrize( + ("cells", "expected"), + [ + # Every staged value is an int, so one cast reads them; the str rows between them are written directly. + pytest.param(["a", 5, 6, "b", None, 7, "c"], ["a", "5", "6", "b", None, "7", "c"], id="str_int_and_null"), + pytest.param(["a", 5, np.str_("c"), None], ["a", "5", "c", None], id="with_a_str_subclass"), + ], +) +def test_map_text_column_mixing_str_with_other_cells_reads_each_as_text(duckdb_cursor, cells, expected): + source = duckdb_cursor.sql(f"select i::VARCHAR as a from range({len(cells)}) t(i)") + mapped = source.map(_text_cells(cells)) + assert mapped.fetchall() == [(value,) for value in expected] + + +def test_map_json_column_still_validates_its_text(duckdb_cursor): + schema = {"a": duckdb.sqltype("JSON")} + good = duckdb_cursor.sql("select 1 as a").map(_text_cells(['{"k": 1}']), schema=schema) + assert good.fetchall() == [('{"k": 1}',)] + bad = duckdb_cursor.sql("select 1 as a").map(_text_cells(["not json"]), schema=schema) + expected = ( + "Conversion Error: UDF column 'a': Malformed JSON at byte 0 of input: invalid literal. Input: \"not json\"" + ) + with pytest.raises(duckdb.ConversionException, match=f"^{re.escape(expected)}$"): + bad.fetchall() + + +@pytest.mark.parametrize( + ("projection", "names"), + [ + pytest.param("col0 as a, col1 as a", ["a", "a_1"], id="one_name"), + pytest.param('col0 as a, col1 as "A"', ["a", "A_1"], id="names_differing_in_case"), + ], +) +def test_map_identity_keeps_both_columns_of_one_name(duckdb_cursor, projection, names): + source = duckdb_cursor.values([1, "x"]).project(projection) + seen = [] + + def record(df): + seen.append(list(df.columns)) + return df + + mapped = source.map(record) + assert seen[0] == names == list(source.df().columns) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + + +def _with_list_column(df): + return df.assign(s=pd.Series([[v] for v in df["a"]], dtype=object, index=df.index)) + + +def _with_text_column(df): + # An object column, as pandas 2 gives this expression on every frame and pandas 3 without pyarrow on the empty one. + return df.assign(s=pd.Series(("prefix " + df["a"].astype(str)).tolist(), dtype=object, index=df.index)) + + +NEW_COLUMNS = pytest.mark.parametrize( + ("function", "sql_type", "values"), + [ + pytest.param(_with_list_column, "BIGINT[]", [[0], [1], [2]], id="list"), + pytest.param(_with_text_column, "VARCHAR", ["prefix 0", "prefix 1", "prefix 2"], id="text"), + ], +) + + +@NEW_COLUMNS +def test_map_new_object_column_reads_into_the_schema_type(duckdb_cursor, function, sql_type, values): + schema = {"a": duckdb.sqltypes.BIGINT, "s": duckdb.sqltype(sql_type)} + mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map(function, schema=schema) + assert mapped.types == [duckdb.sqltypes.BIGINT, duckdb.sqltype(sql_type)] + assert mapped.fetchall() == list(enumerate(values)) + + +@NEW_COLUMNS +def test_map_new_object_column_the_bind_frame_cannot_type_names_the_column_and_schema( + duckdb_cursor, function, sql_type, values +): + mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map(function) + with pytest.raises(duckdb.InvalidInputException, match=_no_type("s")): + mapped.fetchall() + + +def _no_type(column): + message = f"UDF column '{column}' has no type: the function returned no values for it at bind time; pass its type" + return f"^{re.escape(f'Invalid Input Error: {message} with schema=')}$" + + +def test_map_new_object_column_with_a_cell_the_analyzer_cannot_read_names_the_column_and_schema(duckdb_cursor): + def array_cell(df): + cells = np.empty(len(df), dtype=object) + if len(df): + cells[-1] = np.array(1, dtype=object) + return df.assign(s=pd.Series(cells, dtype=object, index=df.index)) + + mapped = duckdb_cursor.sql("select i as a from range(2048) t(i)").map(array_cell) + with pytest.raises(duckdb.InvalidInputException, match=_no_type("s")): + mapped.fetchall() + + +@pytest.mark.parametrize( + ("function", "column"), + [ + pytest.param( + lambda df: df.assign(a=pd.Series(df["a"].astype(str).str.zfill(5).tolist(), dtype=object, index=df.index)), + "a", + id="typed_input_turned_into_text", + ), + pytest.param(lambda df: df.rename(columns={"a": "b", "b": "a"}), "a", id="names_swapped"), + ], +) +def test_map_typed_input_column_the_function_moved_to_another_type_is_refused(duckdb_cursor, function, column): + mapped = duckdb_cursor.sql("select i as a, i::VARCHAR as b from range(3) t(i)").map(function) + with pytest.raises(duckdb.InvalidInputException, match=_no_type(column)): + mapped.fetchall() + + +@pytest.mark.parametrize( + ("make_cell", "expected"), + [ + pytest.param(int, [(0,), (1,), (2,)], id="python_ints"), + pytest.param(lambda v: None, [(None,), (None,), (None,)], id="only_none"), + ], +) +def test_map_typed_input_column_the_function_turned_into_objects_keeps_a_confirmed_type( + duckdb_cursor, make_cell, expected +): + def objects(df): + return df.assign(a=pd.Series([make_cell(v) for v in df["a"]], dtype=object, index=df.index)) + + source = duckdb_cursor.sql("select i::INTEGER as a from range(3) t(i)") + mapped = source.map(objects) + assert mapped.types == source.types + assert mapped.fetchall() == expected + + +def test_map_new_object_column_fails_on_the_first_chunk_with_a_value(duckdb_cursor): + duckdb_cursor.execute("set threads = 1") + rows_per_call = [] + + def late_values(df): + rows_per_call.append(len(df)) + values = [None if i < 2048 else str(i) for i in df["a"]] + return df.assign(s=pd.Series(values, dtype=object, index=df.index)) + + mapped = duckdb_cursor.sql("select i as a from range(4096) t(i)").map(late_values) + with pytest.raises(duckdb.InvalidInputException, match=_no_type("s")): + mapped.fetchall() + # The first chunk, of only None, was read before the second one failed. + assert [rows for rows in rows_per_call if rows] == [2048, 2048] + + +@pytest.mark.parametrize( + "missing", + [ + pytest.param(None, id="None"), + pytest.param(float("nan"), id="NaN"), + pytest.param(pd.NA, id="pd.NA"), + pytest.param(pd.NaT, id="pandas_NaT"), + pytest.param(np.datetime64("NaT", "us"), id="numpy_NaT"), + ], +) +def test_map_new_object_column_of_only_missing_values_stays_null(duckdb_cursor, missing): + def note(df): + return df.assign(note=pd.Series([missing] * len(df), dtype=object, index=df.index)) + + mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map(note) + assert mapped.types == [duckdb.sqltypes.BIGINT, duckdb.sqltypes.SQLNULL] + assert mapped.fetchall() == [(0, None), (1, None), (2, None)] + + +def test_map_new_object_column_of_null_values_is_refused_as_untyped(duckdb_cursor): + # pandas does not read a duckdb.NullValue() as missing, so the column holds values bind could not type + def nulls(df): + return df.assign(s=pd.Series([duckdb.NullValue() for _ in range(len(df))], dtype=object, index=df.index)) + + mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map(nulls) + with pytest.raises(duckdb.InvalidInputException, match=_no_type("s")): + mapped.fetchall() From b504c5e38445d125bfe447d34bbc8123885ed5cd Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Fri, 6 Nov 2026 15:12:55 +0000 Subject: [PATCH 2/9] Read a map() cell into an inferred type only when its kind fits, keep UNION a guess, confirm a guess by lossless widening, and cast MAP keys as values --- .../duckdb_python/numpy/numpy_scan.hpp | 18 +- .../duckdb_python/pandas/pandas_bind.hpp | 6 + src/map.cpp | 216 ++++-- src/native/python_conversion.cpp | 25 +- src/numpy/numpy_scan.cpp | 443 +++++++++++- tests/fast/pandas/test_append_df.py | 98 +++ tests/fast/test_map.py | 28 +- tests/fast/test_map_types.py | 677 +++++++++++++++++- 8 files changed, 1401 insertions(+), 110 deletions(-) diff --git a/src/include/duckdb_python/numpy/numpy_scan.hpp b/src/include/duckdb_python/numpy/numpy_scan.hpp index 2fdd8410..940688b9 100644 --- a/src/include/duckdb_python/numpy/numpy_scan.hpp +++ b/src/include/duckdb_python/numpy/numpy_scan.hpp @@ -2,19 +2,33 @@ #include "duckdb_python/nb/casters.hpp" #include "duckdb/common/common.hpp" +#include "duckdb/common/types.hpp" namespace duckdb { struct PandasColumnBindData; +//! A cell of a declared column whose type was implied is not of that type's kind; source is the type the cell +//! converted to, and the caller reports it in the words of its own type check +struct ImpliedTypeRefused : public std::exception { + explicit ImpliedTypeRefused(LogicalType source_p) : source(std::move(source_p)) { + } + const char *what() const noexcept override { + return "a value is not of the kind of its column's implied type"; + } + LogicalType source; +}; + struct NumpyScan { static void Scan(ClientContext &context, PandasColumnBindData &bind_data, idx_t count, idx_t offset, Vector &out); //! masked holds one numpy bool per row of the whole column, read as a byte, or is null when the column has no mask static void ScanObjectColumn(ClientContext &context, PyObject **col, idx_t stride, const uint8_t *masked, idx_t count, idx_t offset, Vector &out); - //! Converts each cell into the declared type of the column as the appender converts a value + //! Converts each cell into the declared type of the column as the appender converts a value; with implied_type, + //! a cell that is not of the type's kind raises ImpliedTypeRefused, and with bit_storage_refused so does one + //! holding bytes where the type holds a BIT static void ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t stride, const uint8_t *masked, - idx_t count, idx_t offset, Vector &out); + idx_t count, idx_t offset, Vector &out, bool implied_type, bool bit_storage_refused); }; } // namespace duckdb diff --git a/src/include/duckdb_python/pandas/pandas_bind.hpp b/src/include/duckdb_python/pandas/pandas_bind.hpp index ee5a6c3c..b201f044 100644 --- a/src/include/duckdb_python/pandas/pandas_bind.hpp +++ b/src/include/duckdb_python/pandas/pandas_bind.hpp @@ -40,6 +40,12 @@ struct PandasColumnBindData { PythonObjectContainer object_str_val; //! The caller declared the column's type: each value converts into it as the appender converts a row bool declared = false; + //! Of a declared column: the caller implied the type rather than stating it, so a cell is read only when it is of + //! the type's kind, since a cast could change what it was + bool implied_type = false; + //! Of a declared column map() reads: a cell holding bytes where the type holds a BIT is refused, since df() hands + //! a BIT value to the function as its storage bytes + bool bit_storage_refused = false; }; //! The categories of a CategoricalDtype the caller built from an ENUM's values, in order diff --git a/src/map.cpp b/src/map.cpp index 4c8378d4..b73d286e 100644 --- a/src/map.cpp +++ b/src/map.cpp @@ -13,6 +13,8 @@ #include "duckdb_python/dataframe.hpp" #include "duckdb_python/pyconnection/pyconnection.hpp" #include "duckdb/common/vector_operations/vector_operations.hpp" +#include "duckdb/common/type_visitor.hpp" +#include "duckdb_python/python_conversion.hpp" namespace duckdb { @@ -20,15 +22,31 @@ MapFunction::MapFunction() : TableFunction("python_map_function", {LogicalType:: in_out_function = MapFunctionExec; } +//! Where an output column's type came from, which decides how a chunk's column is read into it +enum class TypeSource : uint8_t { + //! The schema stated it, so an object column's values are cast into it whatever they hold + SCHEMA, + //! The empty bind frame typed the column, or the object input column of its name lent its type: an object + //! column's cells are read into it only when they are of its kind, since a function that changed what they hold + //! would otherwise see them cast back + INFERRED, + //! The input column of its name lent its type, but the function got that input typed, or as the member values of + //! a UNION, whose tags cannot travel: an object column is read into it only from cells of its kind, as an + //! inferred one is, and a typed column must have the type, or one that widens into it losslessly + GUESSED, + //! Nothing typed the column: it may hold only missing values + UNTYPED +}; + struct MapFunctionData : public TableFunctionData { explicit MapFunctionData(shared_ptr info_p) : info(std::move(info_p)) { } shared_ptr info; vector in_types, out_types; vector in_names, out_names; - //! Per output column, whether its type is only a guess from the typed input column of its name, which each chunk - //! must confirm, rather than imposed by the schema or an object input - vector guessed; + vector sources; + //! Per output column, the type a returned object column is read into, INVALID where the analyzer types it + vector declared_types; //! Per input column, the CategoricalDtype an ENUM column is handed over in, built once since it lists every value vector enum_dtypes; //! The categories of those dtypes with their ENUM types, by which a returned column of them reads its codes @@ -95,23 +113,44 @@ static optional_idx InputColumnNamed(const vector &input_names, cons } //! An object column of the empty bind frame has no values to type it, so it takes the type of the input column of its -//! name. That type is imposed when the function got the input as an object column too; from a typed input it is only -//! a guess, since the function may have moved the column to another type, and casting back would lose values. +//! name: inferred when the function got that input as an object column too, a guess otherwise (TypeSource). static void OverrideNullType(vector &return_types, const vector &return_names, const vector &input_types, const vector &input_names, - vector &guessed) { - guessed.assign(return_types.size(), false); + vector &sources) { + sources.clear(); for (idx_t col_idx = 0; col_idx < return_types.size(); col_idx++) { if (return_types[col_idx].id() != LogicalTypeId::SQLNULL) { + sources.push_back(TypeSource::INFERRED); continue; } auto input = InputColumnNamed(input_names, return_names[col_idx], col_idx); if (!input.IsValid()) { + sources.push_back(TypeSource::UNTYPED); continue; } auto &input_type = input_types[input.GetIndex()]; return_types[col_idx] = input_type; - guessed[col_idx] = RawArrayWrapper::DuckDBToNumpyDtype(input_type) != "object"; + auto object_input = RawArrayWrapper::DuckDBToNumpyDtype(input_type) == "object"; + auto lends_type = object_input && !TypeVisitor::Contains(input_type, LogicalTypeId::UNION); + sources.push_back(lends_type ? TypeSource::INFERRED : TypeSource::GUESSED); + } +} + +//! Every object column is declared, so the analyzer never reads one and no sample decides how a cell reads; an untyped +//! column is declared NULL, its check that it holds no value comes after the frame's own checks +static void DeclareTypes(MapFunctionData &data) { + data.declared_types.clear(); + for (idx_t col_idx = 0; col_idx < data.out_types.size(); col_idx++) { + switch (data.sources[col_idx]) { + case TypeSource::SCHEMA: + case TypeSource::INFERRED: + case TypeSource::GUESSED: + data.declared_types.push_back(data.out_types[col_idx]); + break; + case TypeSource::UNTYPED: + data.declared_types.push_back(LogicalType::SQLNULL); + break; + } } } @@ -149,7 +188,8 @@ unique_ptr BindExplicitSchema(unique_ptr function function_data->out_names = names; function_data->out_types = types; - function_data->guessed.assign(types.size(), false); + function_data->sources.assign(types.size(), TypeSource::SCHEMA); + DeclareTypes(*function_data); return std::move(function_data); } @@ -185,10 +225,11 @@ static unique_ptr BindMap(ClientContext &context, shared_ptr pandas_bind_data; // unused Pandas::Bind(context, df, pandas_bind_data, return_types, names, {}, data.known_enums); - OverrideNullType(return_types, names, data.in_types, data.in_names, data.guessed); + OverrideNullType(return_types, names, data.in_types, data.in_names, data.sources); data.out_names = names; data.out_types = return_types; + DeclareTypes(data); return std::move(data_uptr); } @@ -218,18 +259,79 @@ static string TypeVectorToString(const vector &types) { return StringUtil::Join(types, types.size(), ", ", [](const LogicalType &argument) { return argument.ToString(); }); } -//! A returned column the analyzer reads as all NULL binds as NULL, which holds for any type, as OverrideNullType -//! assumes at bind -static bool MatchesBoundTypes(const vector &returned, const vector &bound) { - D_ASSERT(returned.size() == bound.size()); - for (idx_t col_idx = 0; col_idx < returned.size(); col_idx++) { - if (returned[col_idx] != bound[col_idx] && returned[col_idx].id() != LogicalTypeId::SQLNULL) { - return false; +static InvalidInputException TypeMismatch(const MapFunctionData &data, const vector &returned) { + return InvalidInputException("UDF column type mismatch, expected [%s], got [%s]", + TypeVectorToString(data.out_types), TypeVectorToString(returned)); +} + +static bool IsTrue(const nb::object &object) { + auto truth = PyObject_IsTrue(object.ptr()); + if (truth < 0) { + throw nb::python_error(); + } + return truth; +} + +//! Whether a column bind could not type holds a value: pandas' missing values, numpy's masked constant and DuckDB's +//! NULL Values are none, as the conversion reads them all as NULL +static bool HoldsValue(const nb::object &column) { + auto present = column.attr("notna")(); + if (!IsTrue(present.attr("any")())) { + return false; + } + auto &import_cache = *DuckDBPyConnection::ImportCache(); + auto cells = nb::list(column.attr("__getitem__")(present).attr("tolist")()); + for (auto cell : cells) { + if (cell.is(import_cache.numpy.ma.masked())) { + continue; + } + if (!duckdb::PyUtil::IsInstance(cell, import_cache.duckdb.Value()) || + !TransformPythonValue(nullptr, cell).IsNull()) { + return true; } } + return false; +} + +static bool IntegerRange(const LogicalType &type, bool &is_signed, idx_t &width) { + switch (type.id()) { + case LogicalTypeId::TINYINT: + case LogicalTypeId::SMALLINT: + case LogicalTypeId::INTEGER: + case LogicalTypeId::BIGINT: + case LogicalTypeId::HUGEINT: + is_signed = true; + break; + case LogicalTypeId::UTINYINT: + case LogicalTypeId::USMALLINT: + case LogicalTypeId::UINTEGER: + case LogicalTypeId::UBIGINT: + case LogicalTypeId::UHUGEINT: + is_signed = false; + break; + default: + return false; + } + width = GetTypeIdSize(type.InternalType()); return true; } +//! Whether every value of one type is a value of the other: an integer into a wider one, a FLOAT into a DOUBLE +static bool WidensLosslessly(const LogicalType &from, const LogicalType &to) { + if (from.id() == LogicalTypeId::FLOAT && to.id() == LogicalTypeId::DOUBLE) { + return true; + } + bool from_signed, to_signed; + idx_t from_width, to_width; + if (!IntegerRange(from, from_signed, from_width) || !IntegerRange(to, to_signed, to_width)) { + return false; + } + if (from_signed) { + return to_signed && to_width >= from_width; + } + return to_signed ? to_width > from_width : to_width >= from_width; +} + static OperatorResultType ExecuteMap(ExecutionContext &context, TableFunctionInput &data_p, DataChunk &input, DataChunk &output) { if (input.size() == 0) { @@ -244,57 +346,38 @@ static OperatorResultType ExecuteMap(ExecutionContext &context, TableFunctionInp auto df = PandasScanFunction::PandasWithBoundNames(FunctionCall(data, conversion)); - auto column_count = nb::len(nb::object(df.attr("columns"))); - if (column_count != output.ColumnCount()) { - throw InvalidInputException("Expected %llu columns from UDF, got %llu", output.ColumnCount(), column_count); - } - // Before anything types the frame, which could fail on a value it cannot read - for (idx_t col_idx = 0; col_idx < output.ColumnCount(); col_idx++) { - if (data.out_types[col_idx].id() != LogicalTypeId::SQLNULL) { - continue; - } - auto column = df.attr("__getitem__")(nb::object(df.attr("columns"))[nb::int_(col_idx)]); - auto has_value = PyObject_IsTrue(nb::object(column.attr("notna")().attr("any")()).ptr()); - if (has_value < 0) { - throw nb::python_error(); - } - if (has_value) { - throw UntypedColumnError(data.out_names[col_idx]); - } - } - // An object column has no type of its own, only values, so it is read into the type the schema or an object input - // imposed, whatever its values. A guessed type and a column bind could not type, which holds only missing values - // here, are left to the analyzer. - vector declared_types; - for (idx_t col_idx = 0; col_idx < output.ColumnCount(); col_idx++) { - auto &type = data.out_types[col_idx]; - auto imposed = type.id() != LogicalTypeId::SQLNULL && !data.guessed[col_idx]; - declared_types.push_back(imposed ? type : LogicalType::INVALID); + auto labels = nb::list(nb::object(df.attr("columns"))); + if (labels.size() != output.ColumnCount()) { + throw InvalidInputException("Expected %llu columns from UDF, got %llu", output.ColumnCount(), labels.size()); } vector pandas_bind_data; vector pandas_return_types; vector pandas_names; auto row_count = Pandas::Bind(context.client, df, pandas_bind_data, pandas_return_types, pandas_names, - declared_types, data.known_enums); + data.declared_types, data.known_enums); + for (idx_t col_idx = 0; col_idx < output.ColumnCount(); col_idx++) { + pandas_bind_data[col_idx].implied_type = + data.sources[col_idx] == TypeSource::INFERRED || data.sources[col_idx] == TypeSource::GUESSED; + pandas_bind_data[col_idx].bit_storage_refused = true; + } - // A guess is only checked: a chunk of only missing values confirms nothing and reads NULL, any other type refutes - // it, and the analyzer's sample decides which, as it did before for every column. - vector missing_values(output.ColumnCount(), false); + // A typed column for a guess must have the guessed type, or one that widens losslessly into it and is cast after + // the read; an object column was declared the guess and reads only cells of its kind + vector widened(output.ColumnCount(), false); + auto checked_types = pandas_return_types; for (idx_t col_idx = 0; col_idx < output.ColumnCount(); col_idx++) { - if (!data.guessed[col_idx] || pandas_return_types[col_idx] == data.out_types[col_idx]) { + auto &returned = pandas_return_types[col_idx]; + auto &bound = data.out_types[col_idx]; + if (data.sources[col_idx] != TypeSource::GUESSED || returned == bound || !WidensLosslessly(returned, bound)) { continue; } - if (pandas_return_types[col_idx].id() != LogicalTypeId::SQLNULL) { - throw UntypedColumnError(data.out_names[col_idx]); - } - pandas_return_types[col_idx] = data.out_types[col_idx]; - missing_values[col_idx] = true; + widened[col_idx] = true; + checked_types[col_idx] = bound; } D_ASSERT(output.GetTypes() == data.out_types); - if (!MatchesBoundTypes(pandas_return_types, output.GetTypes())) { - throw InvalidInputException("UDF column type mismatch, expected [%s], got [%s]", - TypeVectorToString(data.out_types), TypeVectorToString(pandas_return_types)); + if (checked_types != output.GetTypes()) { + throw TypeMismatch(data, pandas_return_types); } if (pandas_names != data.out_names) { throw InvalidInputException("UDF column name mismatch, expected [%s], got [%s]", @@ -304,18 +387,25 @@ static OperatorResultType ExecuteMap(ExecutionContext &context, TableFunctionInp if (row_count > STANDARD_VECTOR_SIZE) { throw InvalidInputException("UDF returned more than %llu rows, which is not allowed.", STANDARD_VECTOR_SIZE); } + for (idx_t col_idx = 0; col_idx < output.ColumnCount(); col_idx++) { + if (data.sources[col_idx] == TypeSource::UNTYPED && HoldsValue(df.attr("__getitem__")(labels[col_idx]))) { + throw UntypedColumnError(data.out_names[col_idx]); + } + } for (idx_t col_idx = 0; col_idx < output.ColumnCount(); col_idx++) { - // A guessed column whose chunk holds only missing values, and a column bind left untyped, the only one still - // typed NULL here; the bound type's scanner would read some missing values as data, such as NaT as 'NaT'. - if (missing_values[col_idx] || pandas_return_types[col_idx].id() == LogicalTypeId::SQLNULL) { + if (pandas_return_types[col_idx].id() == LogicalTypeId::SQLNULL) { ConstantVector::SetNull(output.data[col_idx], count_t(row_count)); continue; } auto &bind_data = pandas_bind_data[col_idx]; - Vector temp(output.data[col_idx].GetType(), row_count); + Vector temp(pandas_return_types[col_idx], row_count); try { PandasScanFunction::PandasBackendScanSwitch(context.client, bind_data, row_count, 0, temp); + } catch (ImpliedTypeRefused &refused) { + auto returned = pandas_return_types; + returned[col_idx] = refused.source; + throw TypeMismatch(data, returned); } catch (std::exception &ex) { ErrorData error(ex); // The cast's own message names no column; a Python error keeps the traceback its message starts with @@ -325,7 +415,11 @@ static OperatorResultType ExecuteMap(ExecutionContext &context, TableFunctionInp } error.Throw(StringUtil::Format("UDF column '%s': ", data.out_names[col_idx].c_str())); } - VectorOperations::Copy(temp, output.data[col_idx], row_count, 0, 0); + if (widened[col_idx]) { + VectorOperations::Cast(context.client, temp, output.data[col_idx], row_count); + } else { + VectorOperations::Copy(temp, output.data[col_idx], row_count, 0, 0); + } } output.SetChildCardinality(row_count); return OperatorResultType::NEED_MORE_INPUT; diff --git a/src/native/python_conversion.cpp b/src/native/python_conversion.cpp index 5a86ae44..cd3be740 100644 --- a/src/native/python_conversion.cpp +++ b/src/native/python_conversion.cpp @@ -5,6 +5,7 @@ #include "duckdb_python/pyresult.hpp" #include "duckdb_python/pyconnection/pyconnection.hpp" #include "duckdb/common/types.hpp" +#include "duckdb/common/type_visitor.hpp" #include "duckdb/common/vector/flat_vector.hpp" #include "duckdb/common/vector/list_vector.hpp" #include "duckdb/common/vector/array_vector.hpp" @@ -151,12 +152,15 @@ Value TransformDictionaryToStruct(optional_ptr context, const PyD dict.ToString(), target_type.ToString()); }; bool struct_target = StructType::IsStruct(target_type); - if (struct_target && dict.len != StructType::GetChildCount(target_type)) { + // A target holding UNKNOWN gives only the shape of the values below it: the dict's own keys name the fields, as + // without a target, and each value converts with the shape of the field of its name + bool shape_only = struct_target && TypeVisitor::Contains(target_type, LogicalTypeId::UNKNOWN); + if (struct_target && !shape_only && dict.len != StructType::GetChildCount(target_type)) { throw conversion_error(); } child_list_t struct_values; - if (StructType::IsUnnamed(target_type)) { + if (!shape_only && StructType::IsUnnamed(target_type)) { // A TUPLE has no member names, so its members take the dict's values in dict order. for (idx_t i = 0; i < dict.len; i++) { auto val = TransformPythonValue(context, dict.values.attr("__getitem__")(i), @@ -167,10 +171,16 @@ Value TransformDictionaryToStruct(optional_ptr context, const PyD } auto struct_keys = TransformStructKeys(dict.keys, dict.len); - if (!struct_target) { + if (!struct_target || shape_only) { for (idx_t i = 0; i < dict.len; i++) { - auto val = - TransformPythonValue(context, dict.values.attr("__getitem__")(i), LogicalType::UNKNOWN, nan_as_null); + LogicalType child_type = LogicalType::UNKNOWN; + for (idx_t child = 0; shape_only && child < StructType::GetChildCount(target_type); child++) { + if (StructType::GetChildName(target_type, child) == struct_keys[i]) { + child_type = StructType::GetChildType(target_type, child); + break; + } + } + auto val = TransformPythonValue(context, dict.values.attr("__getitem__")(i), child_type, nan_as_null); struct_values.emplace_back(make_pair(std::move(struct_keys[i]), std::move(val))); } return Value::STRUCT(std::move(struct_values)); @@ -241,10 +251,11 @@ Value TransformStructFormatDictionaryToMap(optional_ptr context, elements.push_back(Value::STRUCT(std::move(struct_values))); } - if (key_type.id() == LogicalTypeId::SQLNULL) { + // A target holding UNKNOWN asks for no type there, and a MAP of such a type is no value the engine can hold + if (key_type.id() == LogicalTypeId::SQLNULL && !TypeVisitor::Contains(key_target, LogicalTypeId::UNKNOWN)) { key_type = key_target; } - if (value_type.id() == LogicalTypeId::SQLNULL) { + if (value_type.id() == LogicalTypeId::SQLNULL && !TypeVisitor::Contains(value_target, LogicalTypeId::UNKNOWN)) { value_type = value_target; } diff --git a/src/numpy/numpy_scan.cpp b/src/numpy/numpy_scan.cpp index d9ff6fe9..9242efad 100644 --- a/src/numpy/numpy_scan.cpp +++ b/src/numpy/numpy_scan.cpp @@ -4,6 +4,7 @@ #include "duckdb_python/python_conversion.hpp" #include "duckdb/common/string_util.hpp" #include "duckdb/common/swar.hpp" +#include "duckdb/common/type_visitor.hpp" #include "duckdb/common/types/timestamp.hpp" #include "duckdb/common/vector/struct_vector.hpp" #include "duckdb/common/vector/list_vector.hpp" @@ -329,7 +330,11 @@ static nb::object NumpyScalarValue(nb::handle scalar) { return scalar.attr("item")(); } -static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb::handle cell) { +//! With a shape (KeyedShape) the cell converts with it as its target, which a type holding a MAP needs: without it a +//! dict reads as a STRUCT named by its keys' text, and a key whose text does not cast back, such as bytes, is lost. +//! An object without a DuckDB type reads as its text into a stated VARCHAR, and has no kind for an implied one. +static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb::handle cell, + optional_ptr shape, bool implied) { if (cell.ptr() == Py_None) { return Value(); } @@ -337,13 +342,36 @@ static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb auto value = NumpyScalarValue(cell); // A long double's item() is a numpy scalar again; one step, then the conversion refuses what is left if (!IsNumpyScalar(value)) { - return DeclaredValue(context, target, value); + return DeclaredValue(context, target, value, shape, implied); } } - if (target.id() == LogicalTypeId::VARCHAR && GetPythonObjectType(cell) == PythonObjectType::Other) { + if (!implied && target.id() == LogicalTypeId::VARCHAR && GetPythonObjectType(cell) == PythonObjectType::Other) { return Value(TransformPythonString(nb::str(cell))); } - return TransformPythonValue(context, cell, LogicalType::UNKNOWN, true); + return TransformPythonValue(context, cell, shape ? *shape : LogicalType(LogicalTypeId::UNKNOWN), true); +} + +//! The type a cell of a type holding a MAP converts with: its containers and no other type, so a dict builds its MAP +//! from its keys and values in the types they have in Python, and every key and value meets the engine's cast with +//! the connection's settings, as a TIMESTAMPTZ in the session time zone +static LogicalType KeyedShape(const LogicalType &type) { + switch (type.id()) { + case LogicalTypeId::STRUCT: { + child_list_t children; + for (idx_t i = 0; i < StructType::GetChildCount(type); i++) { + children.emplace_back(StructType::GetChildName(type, i), KeyedShape(StructType::GetChildType(type, i))); + } + return LogicalType::STRUCT(std::move(children)); + } + case LogicalTypeId::LIST: + return LogicalType::LIST(KeyedShape(ListType::GetChildType(type))); + case LogicalTypeId::ARRAY: + return LogicalType::ARRAY(KeyedShape(ArrayType::GetChildType(type)), ArrayType::GetSize(type)); + case LogicalTypeId::MAP: + return LogicalType::MAP(KeyedShape(MapType::KeyType(type)), KeyedShape(MapType::ValueType(type))); + default: + return LogicalType::UNKNOWN; + } } //! The values of one vector of a declared column, written once the cells are read: a vector whose values share one @@ -358,6 +386,7 @@ class DeclaredVector { //! A row the scan wrote into the output itself, which Write leaves as it is void MarkWritten(idx_t row) { written[row] = true; + written_count++; } void Set(idx_t row, Value value) { @@ -373,6 +402,9 @@ class DeclaredVector { } void Write(ClientContext &context, Vector &out, idx_t count) { + if (written_count == count) { + return; + } auto &target = out.GetType(); if (uniform && source.id() != LogicalTypeId::SQLNULL && source != target) { idx_t present = 0; @@ -429,10 +461,370 @@ class DeclaredVector { private: vector values; vector written; + idx_t written_count = 0; LogicalType source; bool uniform = true; }; +//! Whether the value a cell of a column whose type was implied, not stated, converted to is of that type's kind: it +//! is the type, or each of its parts is of the kind of the type's part in the same place, by the rules below, so a +//! function that changed what a column holds is refused, not cast back. The check reads the converted value, so it +//! judges exactly what the write would read. +class ImpliedKinds { +public: + explicit ImpliedKinds(ClientContext &context) : context(context) { + } + + bool Fits(const Value &value, const LogicalType &type) { + if (value.IsNull()) { + return true; + } + auto &source = value.type(); + if (source == type) { + return true; + } + for (auto &entry : seen) { + if (entry.source == source) { + return entry.verdict == Verdict::FITS || + (entry.verdict == Verdict::BY_VALUE && FitsByValue(value, type)); + } + } + auto verdict = Judge(source, type); + if (seen.size() < CACHED_SOURCE_TYPES) { + seen.push_back({source, verdict}); + } + return verdict == Verdict::FITS || (verdict == Verdict::BY_VALUE && FitsByValue(value, type)); + } + +private: + //! BY_VALUE: the kinds agree, but whether the type holds the value depends on the value (an integer, a DECIMAL) + enum class Verdict : uint8_t { FITS, REFUSED, BY_VALUE }; + + static Verdict Combine(Verdict a, Verdict b) { + if (a == Verdict::REFUSED || b == Verdict::REFUSED) { + return Verdict::REFUSED; + } + return a == Verdict::BY_VALUE || b == Verdict::BY_VALUE ? Verdict::BY_VALUE : Verdict::FITS; + } + + //! The field of a STRUCT type a key names, or INVALID; the engine's STRUCT cast matches names ignoring case + static optional_idx FieldNamed(const LogicalType &type, const Identifier &name) { + for (idx_t i = 0; i < StructType::GetChildCount(type); i++) { + if (StructType::GetChildName(type, i) == name) { + return i; + } + } + return optional_idx(); + } + + //! The unit of a naive time or timestamp type, finer the higher, or 0 for any other type + static idx_t NaiveUnit(LogicalTypeId id) { + switch (id) { + case LogicalTypeId::TIMESTAMP_SEC: + return 1; + case LogicalTypeId::TIMESTAMP_MS: + return 2; + case LogicalTypeId::TIMESTAMP: + return 3; + case LogicalTypeId::TIMESTAMP_NS: + return 4; + case LogicalTypeId::TIME: + return 13; + case LogicalTypeId::TIME_NS: + return 14; + default: + return 0; + } + } + + static bool IsZoned(LogicalTypeId id) { + return id == LogicalTypeId::TIME_TZ || id == LogicalTypeId::TIMESTAMP_TZ || + id == LogicalTypeId::TIMESTAMP_TZ_NS; + } + + static bool IsNested(const LogicalType &type) { + return type.IsNested() || type.id() == LogicalTypeId::UNION; + } + + Verdict Judge(const LogicalType &source, const LogicalType &type) { + if (source == type || source.id() == LogicalTypeId::SQLNULL) { + return Verdict::FITS; + } + switch (type.id()) { + case LogicalTypeId::STRUCT: { + if (source.id() != LogicalTypeId::STRUCT) { + return Verdict::REFUSED; + } + // A key the STRUCT does not name, or a second key for one field, would be dropped by the cast; a field + // without a key reads NULL + vector claimed(StructType::GetChildCount(type), false); + auto verdict = Verdict::FITS; + for (idx_t i = 0; i < StructType::GetChildCount(source); i++) { + auto field = FieldNamed(type, StructType::GetChildName(source, i)); + if (!field.IsValid() || claimed[field.GetIndex()]) { + return Verdict::REFUSED; + } + claimed[field.GetIndex()] = true; + verdict = Combine(verdict, Judge(StructType::GetChildType(source, i), + StructType::GetChildType(type, field.GetIndex()))); + } + return verdict; + } + // The engine casts any MAP into any MAP implicitly without looking at its key and value types + case LogicalTypeId::MAP: + if (source.id() != LogicalTypeId::MAP) { + return Verdict::REFUSED; + } + return Combine(Judge(MapType::KeyType(source), MapType::KeyType(type)), + Judge(MapType::ValueType(source), MapType::ValueType(type))); + case LogicalTypeId::LIST: + if (source.id() != LogicalTypeId::LIST) { + return Verdict::REFUSED; + } + return Judge(ListType::GetChildType(source), ListType::GetChildType(type)); + case LogicalTypeId::ARRAY: + // A Python list reads as a LIST, which the cast takes into an ARRAY of its length or refuses + if (source.id() == LogicalTypeId::LIST) { + return Judge(ListType::GetChildType(source), ArrayType::GetChildType(type)); + } + if (source.id() != LogicalTypeId::ARRAY) { + return Verdict::REFUSED; + } + return Judge(ArrayType::GetChildType(source), ArrayType::GetChildType(type)); + default: + break; + } + if (IsNested(source) || IsNested(type)) { + return Verdict::REFUSED; + } + // A naive value and one with a time zone are of different kinds: the cast between them reads the wall time in + // the session's time zone + if (IsZoned(source.id()) != IsZoned(type.id())) { + return Verdict::REFUSED; + } + // A time or timestamp of another unit holds only the values the cast keeps unchanged: a coarser unit rounds, + // and a finer one covers fewer years + auto source_unit = NaiveUnit(source.id()); + auto unit = NaiveUnit(type.id()); + if (source_unit && unit && source_unit / 10 == unit / 10) { + return source_unit == unit ? Verdict::FITS : Verdict::BY_VALUE; + } + auto zoned_timestamps = + (source.id() == LogicalTypeId::TIMESTAMP_TZ || source.id() == LogicalTypeId::TIMESTAMP_TZ_NS) && + (type.id() == LogicalTypeId::TIMESTAMP_TZ || type.id() == LogicalTypeId::TIMESTAMP_TZ_NS); + if (zoned_timestamps || (source.id() == LogicalTypeId::VARCHAR && type.id() == LogicalTypeId::VARCHAR) || + (source.id() == LogicalTypeId::DOUBLE && type.id() == LogicalTypeId::FLOAT)) { + return Verdict::FITS; + } + // df() hands a BIT value over as its storage bytes, which the cast from BLOB reads as eight bits per byte + if (source.id() == LogicalTypeId::BLOB && type.id() == LogicalTypeId::BIT) { + return Verdict::REFUSED; + } + // A number fits an integer or a DECIMAL when the cast into it and back gives the number again, whatever kind of + // number it was in Python: 2.0 and Decimal(5) are the integer 2 and 5, 2.5 is none + if (source.IsIntegral() && type.IsIntegral()) { + return CastFunctionSet::ImplicitCastCost(context, source, type) >= 0 ? Verdict::FITS : Verdict::BY_VALUE; + } + if (source.id() == LogicalTypeId::DECIMAL && type.id() == LogicalTypeId::DECIMAL) { + auto source_scale = DecimalType::GetScale(source); + auto scale = DecimalType::GetScale(type); + auto fits = source_scale <= scale && + DecimalType::GetWidth(source) - source_scale <= DecimalType::GetWidth(type) - scale; + return fits ? Verdict::FITS : Verdict::BY_VALUE; + } + auto number = source.IsIntegral() || source.id() == LogicalTypeId::DECIMAL || + source.id() == LogicalTypeId::DOUBLE || source.id() == LogicalTypeId::FLOAT; + if (number && (type.IsIntegral() || type.id() == LogicalTypeId::DECIMAL)) { + return Verdict::BY_VALUE; + } + return CastFunctionSet::ImplicitCastCost(context, source, type) >= 0 ? Verdict::FITS : Verdict::REFUSED; + } + + bool FitsByValue(const Value &value, const LogicalType &type) { + if (value.IsNull()) { + return true; + } + auto &source = value.type(); + switch (Judge(source, type)) { + case Verdict::FITS: + return true; + case Verdict::REFUSED: + return false; + case Verdict::BY_VALUE: + break; + } + switch (type.id()) { + case LogicalTypeId::STRUCT: { + auto &children = StructValue::GetChildren(value); + for (idx_t i = 0; i < children.size(); i++) { + auto field = FieldNamed(type, StructType::GetChildName(source, i)); + if (!FitsByValue(children[i], StructType::GetChildType(type, field.GetIndex()))) { + return false; + } + } + return true; + } + case LogicalTypeId::MAP: + for (auto &entry : MapValue::GetChildren(value)) { + auto &pair = StructValue::GetChildren(entry); + if (!FitsByValue(pair[0], MapType::KeyType(type)) || !FitsByValue(pair[1], MapType::ValueType(type))) { + return false; + } + } + return true; + case LogicalTypeId::LIST: + case LogicalTypeId::ARRAY: { + auto &child = + type.id() == LogicalTypeId::LIST ? ListType::GetChildType(type) : ArrayType::GetChildType(type); + auto &children = + source.id() == LogicalTypeId::LIST ? ListValue::GetChildren(value) : ArrayValue::GetChildren(value); + for (auto &element : children) { + if (!FitsByValue(element, child)) { + return false; + } + } + return true; + } + default: { + if ((source.id() == LogicalTypeId::DOUBLE || source.id() == LogicalTypeId::FLOAT) && + !FloatHoldsItsText(value, type)) { + return false; + } + // The type holds the value when the cast into it and back gives the value again + auto cast = value.DefaultTryCastAs(type); + if (!cast) { + return false; + } + auto back = cast->DefaultTryCastAs(source); + return back && Value::NotDistinctFrom(*back, value); + } + } + } + +public: + //! Whether a value holds bytes where the type holds a BIT, at any depth: df() hands a BIT value over as its + //! storage bytes, which no cast reads back, whether or not the type was stated + static bool HoldsBitStorage(const Value &value, const LogicalType &type) { + if (value.IsNull()) { + return false; + } + auto &source = value.type(); + switch (type.id()) { + case LogicalTypeId::BIT: + return source.id() == LogicalTypeId::BLOB; + case LogicalTypeId::STRUCT: { + if (source.id() != LogicalTypeId::STRUCT) { + return false; + } + auto &children = StructValue::GetChildren(value); + for (idx_t i = 0; i < children.size(); i++) { + auto field = FieldNamed(type, StructType::GetChildName(source, i)); + if (field.IsValid() && HoldsBitStorage(children[i], StructType::GetChildType(type, field.GetIndex()))) { + return true; + } + } + return false; + } + case LogicalTypeId::MAP: + if (source.id() != LogicalTypeId::MAP) { + return false; + } + for (auto &entry : MapValue::GetChildren(value)) { + auto &pair = StructValue::GetChildren(entry); + if (HoldsBitStorage(pair[0], MapType::KeyType(type)) || + HoldsBitStorage(pair[1], MapType::ValueType(type))) { + return true; + } + } + return false; + case LogicalTypeId::LIST: + case LogicalTypeId::ARRAY: { + if (source.id() != LogicalTypeId::LIST && source.id() != LogicalTypeId::ARRAY) { + return false; + } + auto &child = + type.id() == LogicalTypeId::LIST ? ListType::GetChildType(type) : ArrayType::GetChildType(type); + auto &children = + source.id() == LogicalTypeId::LIST ? ListValue::GetChildren(value) : ArrayValue::GetChildren(value); + for (auto &element : children) { + if (HoldsBitStorage(element, child)) { + return true; + } + } + return false; + } + default: + return false; + } + } + +private: + //! The significant digits of a number's shortest text ("2.0" has one, "1e+20" twenty one): from its first digit + //! that is not zero to its last, or to the units when its integer part ends in zeros + static idx_t SignificantDigits(const string &text) { + string digits; + int64_t integer_digits = 0; + bool point = false; + idx_t i = 0; + for (; i < text.size() && text[i] != 'e' && text[i] != 'E'; i++) { + if (text[i] == '.') { + point = true; + } else if (StringUtil::CharacterIsDigit(text[i])) { + digits += text[i]; + integer_digits += point ? 0 : 1; + } + } + int64_t exponent = i < text.size() ? std::stoll(text.substr(i + 1)) : 0; + auto units = integer_digits + exponent; + auto end = static_cast(digits.size()); + while (end > MaxValue(units, 0) && digits[NumericCast(end - 1)] == '0') { + end--; + } + int64_t start = 0; + while (start < end && digits[NumericCast(start)] == '0') { + start++; + } + if (start == end) { + return 0; + } + return NumericCast(MaxValue(end, units) - start); + } + + //! A double holds every decimal of at most 15 significant digits exactly enough to give it back as its shortest + //! text, so a float whose text has more may be a rounding already: df() hands a HUGEINT or a DECIMAL over as a + //! float64. And the engine's cast of a float into a DECIMAL rounds the value times 10**scale in double precision, + //! which for a wide scale is no longer that value (2.0 reads 2.000000000000000039769249677312 in a DECIMAL(38,30), + //! exactly in a DECIMAL(38,20)), so the float fits only when the cast gives the value of its text. + static bool FloatHoldsItsText(const Value &value, const LogicalType &type) { + auto number = value.GetValue(); + if (!Value::DoubleIsFinite(number)) { + return false; + } + auto text = Value::DOUBLE(number).ToString(); + if (SignificantDigits(text) > MAX_EXACT_FLOAT_DIGITS) { + return false; + } + if (type.id() != LogicalTypeId::DECIMAL) { + return true; + } + auto cast = value.DefaultTryCastAs(type); + auto exact = Value(text).DefaultTryCastAs(type); + return cast && exact && Value::NotDistinctFrom(*cast, *exact); + } + + struct Seen { + LogicalType source; + Verdict verdict; + }; + //! DBL_DIG: every decimal of this many significant digits survives a round trip through a double + static constexpr idx_t MAX_EXACT_FLOAT_DIGITS = 15; + //! Bounded like DeclaredCasts, since dicts give a STRUCT type per key set + static constexpr idx_t CACHED_SOURCE_TYPES = 16; + + ClientContext &context; + vector seen; +}; + // A column without a mask takes the instantiation that tests no mask bits, so it reads like the plain array template static void ScanObjectRows(PyObject **col, idx_t stride, const uint8_t *masked, idx_t count, idx_t offset, Vector &out, @@ -471,18 +863,51 @@ static void ScanObjectCells(ClientContext &context, PyObject **col, idx_t stride } void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t stride, const uint8_t *masked, - idx_t count, idx_t offset, Vector &out) { + idx_t count, idx_t offset, Vector &out, bool implied_type, + bool bit_storage_refused) { DeclaredVector staged(count); auto &target = out.GetType(); - // A str into plain VARCHAR reads as TransformPythonValue would read it, without building a Value per cell + const auto keyed = TypeVisitor::Contains(target, LogicalTypeId::MAP); + const auto holds_bit = TypeVisitor::Contains(target, LogicalTypeId::BIT); + const auto shape = keyed ? KeyedShape(target) : LogicalType(LogicalTypeId::UNKNOWN); + // A str into VARCHAR without an alias, and bytes or a bytearray into BLOB without one, read as TransformPythonValue + // would read them, without building a Value per cell; a collation governs comparison only, an alias such as JSON + // validates the text in its cast const bool plain_text = target == LogicalType::VARCHAR; + const bool plain_blob = target == LogicalType::BLOB; + ImpliedKinds kinds(context); ScanObjectCells(context, col, stride, masked, count, offset, out, [&](nb::handle cell, idx_t row) { if (plain_text && PyUnicode_CheckExact(cell.ptr())) { - FlatVector::GetDataMutable(out)[row] = StringVector::AddString(out, TransformPythonString(cell)); + Py_ssize_t size; + auto text = PyUnicode_AsUTF8AndSize(cell.ptr(), &size); + if (!text) { + ThrowUnencodableStr(cell); + } + FlatVector::GetDataMutable(out)[row] = + StringVector::AddString(out, text, NumericCast(size)); staged.MarkWritten(row); return; } - staged.Set(row, DeclaredValue(context, target, cell)); + // df() hands a BLOB cell to the function as a bytearray + if (plain_blob && (PyBytes_CheckExact(cell.ptr()) || PyByteArray_CheckExact(cell.ptr()))) { + auto bytes = PyBytes_CheckExact(cell.ptr()); + auto data = bytes ? PyBytes_AS_STRING(cell.ptr()) : PyByteArray_AS_STRING(cell.ptr()); + auto size = bytes ? PyBytes_GET_SIZE(cell.ptr()) : PyByteArray_GET_SIZE(cell.ptr()); + FlatVector::GetDataMutable(out)[row] = + StringVector::AddStringOrBlob(out, data, NumericCast(size)); + staged.MarkWritten(row); + return; + } + auto value = DeclaredValue(context, target, cell, keyed ? &shape : nullptr, implied_type); + // A duckdb.Value states its type, as a schema does, so a cast into the column's type is meant + if (implied_type && !kinds.Fits(value, target) && + !duckdb::PyUtil::IsInstance(cell, DuckDBPyConnection::ImportCache()->duckdb.Value())) { + throw ImpliedTypeRefused(value.type()); + } + if (bit_storage_refused && holds_bit && ImpliedKinds::HoldsBitStorage(value, target)) { + throw ImpliedTypeRefused(value.type()); + } + staged.Set(row, std::move(value)); }); staged.Write(context, out, count); VerifyTypeConstraints(out, count); @@ -761,7 +1186,7 @@ void NumpyScan::Scan(ClientContext &context, PandasColumnBindData &bind_data, id const bool is_object_col = bind_data.numpy_type.type == NumpyNullableType::OBJECT; if (is_object_col && bind_data.declared) { return NumpyScan::ScanDeclaredColumn(context, src_ptr, numpy_col.stride, MaskedRows(bind_data), count, - offset, out); + offset, out, bind_data.implied_type, bind_data.bit_storage_refused); } if (is_object_col && out.GetType().id() != LogicalTypeId::VARCHAR) { //! We have determined the underlying logical type of this object column diff --git a/tests/fast/pandas/test_append_df.py b/tests/fast/pandas/test_append_df.py index 0f3099f8..66007b6b 100644 --- a/tests/fast/pandas/test_append_df.py +++ b/tests/fast/pandas/test_append_df.py @@ -2,6 +2,7 @@ import enum import math import pathlib +import re import zoneinfo from decimal import Decimal @@ -422,6 +423,97 @@ def test_keys_that_collapse_in_a_nested_map_are_refused(self, duckdb_cursor): with pytest.raises(duckdb.InvalidInputException, match=r"(?i)duplicat|unique"): duckdb_cursor.append("t", frame) + @pytest.mark.parametrize( + ("key_type", "key"), + [ + pytest.param("BLOB", b"\x00ab", id="bytes"), + pytest.param("INTERVAL", datetime.timedelta(days=1), id="timedelta"), + ], + ) + def test_dict_key_whose_text_does_not_cast_back_keeps_its_value_in_a_map(self, duckdb_cursor, key_type, key): + duckdb_cursor.execute(f"CREATE TABLE t (a MAP({key_type}, INTEGER), b STRUCT(m MAP({key_type}, INTEGER)))") + frame = pd.DataFrame( + {"a": pd.Series([{key: 1}, None], dtype=object), "b": pd.Series([{"m": {key: 2}}, None], dtype=object)} + ) + duckdb_cursor.append("t", frame) + assert duckdb_cursor.table("t").fetchall() == [({key: 1}, {"m": {key: 2}}), (None, None)] + + @pytest.mark.parametrize( + ("target", "cell", "expected"), + [ + pytest.param("MAP(VARCHAR, INTEGER)", {"k": None}, {"k": None}, id="scalar_values"), + pytest.param("MAP(VARCHAR, STRUCT(x INTEGER))", {"k": None}, {"k": None}, id="struct_values"), + pytest.param("MAP(VARCHAR, MAP(BLOB, INTEGER))", {"k": {b"z": None}}, {"k": {b"z": None}}, id="nested_map"), + ], + ) + def test_map_of_only_null_values_appends_them(self, duckdb_cursor, target, cell, expected): + duckdb_cursor.execute(f"CREATE TABLE t (a {target})") + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) + assert duckdb_cursor.table("t").fetchall() == [(expected,)] + + @pytest.mark.parametrize( + "value", [datetime.datetime(2024, 5, 6, 12, 0), "2024-05-06 12:00:00"], ids=["naive_datetime", "text"] + ) + def test_map_value_into_timestamptz_reads_in_the_session_time_zone(self, duckdb_cursor, value): + duckdb_cursor.execute("SET TimeZone = 'Europe/Amsterdam'") + duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, TIMESTAMPTZ))") + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([{"k": value}], dtype=object)})) + assert duckdb_cursor.sql("SELECT a['k']::VARCHAR FROM t").fetchall() == [("2024-05-06 12:00:00+02",)] + + @pytest.mark.parametrize( + ("target", "key", "expected"), + [ + pytest.param("TIMESTAMP_NS", datetime.datetime(2020, 1, 1), "{'2020-01-01 00:00:00'=1}", id="timestamp_ns"), + pytest.param("TIMESTAMP_MS", datetime.datetime(2020, 1, 1), "{'2020-01-01 00:00:00'=1}", id="timestamp_ms"), + pytest.param("TIMESTAMP_S", datetime.datetime(2020, 1, 1), "{'2020-01-01 00:00:00'=1}", id="timestamp_s"), + pytest.param("DATE", datetime.datetime(2020, 1, 1, 5), "{2020-01-01=1}", id="date_from_datetime"), + pytest.param( + "TIMESTAMPTZ", datetime.datetime(2020, 1, 1), "{'2020-01-01 00:00:00-05'=1}", id="timestamptz_naive" + ), + pytest.param("TIMESTAMPTZ", "2020-01-01 00:00:00", "{'2020-01-01 00:00:00-05'=1}", id="timestamptz_text"), + pytest.param("DECIMAL(9,2)", 1.257, "{1.26=1}", id="decimal_from_float"), + ], + ) + def test_map_key_meets_the_cast_in_the_session_settings(self, duckdb_cursor, target, key, expected): + duckdb_cursor.execute("SET TimeZone = 'America/New_York'") + duckdb_cursor.execute(f"CREATE TABLE t (a MAP({target}, INTEGER))") + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([{key: 1}], dtype=object)})) + assert duckdb_cursor.sql("SELECT a::VARCHAR FROM t").fetchall() == [(expected,)] + + @pytest.mark.parametrize( + ("target", "key", "expected"), + [ + # The engine's cast reads the instant in the session's time zone; the base dropped the offset + pytest.param( + "TIMESTAMP", + datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), + "{'2019-12-31 19:00:00'=1}", + id="aware_datetime_into_timestamp", + ), + # The cast from BLOB gives the bytes' text; the base gave the text of the Python bytes object + pytest.param("VARCHAR", b"x", "{x=1}", id="bytes_into_varchar"), + pytest.param("STRUCT(a INTEGER)", (1,), "{{'a': 1}=1}", id="tuple_into_struct"), + ], + ) + def test_map_key_reads_as_the_cast_from_its_python_type_gives_it(self, duckdb_cursor, target, key, expected): + duckdb_cursor.execute("SET TimeZone = 'America/New_York'") + duckdb_cursor.execute(f"CREATE TABLE t (a MAP({target}, INTEGER))") + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([{key: 1}], dtype=object)})) + assert duckdb_cursor.sql("SELECT a::VARCHAR FROM t").fetchall() == [(expected,)] + + def test_map_keys_a_str_and_its_bytes_are_one_key_and_refused(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, INTEGER))") + frame = pd.DataFrame({"a": pd.Series([{"a": 1, b"a": 2}], dtype=object)}) + with pytest.raises(duckdb.InvalidInputException, match=re.escape("Map keys must be unique")): + duckdb_cursor.append("t", frame) + + def test_map_key_of_none_is_refused_and_the_connection_stays_usable(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, INTEGER))") + with pytest.raises(duckdb.Error) as error: + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([{None: 1}], dtype=object)})) + assert not isinstance(error.value, duckdb.FatalException), error.value + assert duckdb_cursor.sql("SELECT 42").fetchall() == [(42,)] + def test_text_with_a_lone_surrogate_names_it(self, duckdb_cursor): duckdb_cursor.execute("CREATE TABLE t (a VARCHAR)") frame = pd.DataFrame({"a": _outside_the_sample("a", pathlib.PurePosixPath("/\udcff"))}) @@ -443,6 +535,12 @@ def test_text_column_mixing_str_with_other_cells_appends_each_as_text(self, duck ("\u00e9 \U0001f600",), ] + def test_blob_column_mixing_bytes_with_other_cells_appends_each(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (a BLOB)") + cells = [b"a", bytearray(b"b"), None, memoryview(b"c"), b"\x00d", "e"] + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series(cells, dtype=object)})) + assert duckdb_cursor.table("t").fetchall() == [(b"a",), (b"b",), (None,), (b"c",), (b"\x00d",), (b"e",)] + def test_text_with_a_lone_surrogate_in_a_str_names_it(self, duckdb_cursor): duckdb_cursor.execute("CREATE TABLE t (a VARCHAR)") frame = pd.DataFrame({"a": pd.Series(["a", "b\udcff"], dtype=object)}) diff --git a/tests/fast/test_map.py b/tests/fast/test_map.py index dec9c540..ea5142ff 100644 --- a/tests/fast/test_map.py +++ b/tests/fast/test_map.py @@ -393,9 +393,25 @@ def replace(df): ) -@MAP_SCHEMAS -@pytest.mark.xfail( - strict=True, reason="not built: a fractional value outside the chunk's sample is converted with loss into INTEGER" +@pytest.mark.parametrize( + "schema", + [ + pytest.param( + None, + id="inferred", + marks=pytest.mark.xfail( + strict=True, reason="not built: the refusal of a fractional value names neither its row nor the float" + ), + ), + pytest.param( + {"a": duckdb.sqltypes.INTEGER}, + id="explicit_schema", + marks=pytest.mark.xfail( + strict=True, + reason="not built: a fractional value outside the chunk's sample is converted with loss into INTEGER", + ), + ), + ], ) def test_map_fractional_value_outside_the_chunk_sample_raises(duckdb_cursor, schema): duckdb_cursor.execute("set threads = 1") @@ -652,12 +668,14 @@ def missing_on_rows(df): assert mapped.fetchall() == [(None,), (None,), (None,)] -def test_map_returned_object_column_with_a_value_of_another_type_raises_the_casts_error(duckdb_cursor): +def test_map_returned_object_column_with_a_value_of_another_type_still_mismatches(duckdb_cursor): def text_on_rows(df): return df.assign(a=pd.Series(["x"] + [None] * (len(df) - 1), dtype=object)) if len(df) else df mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map(text_on_rows) - with pytest.raises(duckdb.ConversionException, match=re.escape("Could not convert string 'x' to INT64")): + with pytest.raises( + duckdb.InvalidInputException, match=re.escape("UDF column type mismatch, expected [BIGINT], got [VARCHAR]") + ): mapped.fetchall() diff --git a/tests/fast/test_map_types.py b/tests/fast/test_map_types.py index 8672d4ca..59c45892 100644 --- a/tests/fast/test_map_types.py +++ b/tests/fast/test_map_types.py @@ -1,6 +1,10 @@ """map() fixes one type per output column at bind and reads each chunk's returned column into it.""" +import datetime +import json import re +from decimal import Decimal +from fractions import Fraction import pytest from conftest import PANDAS_GE_3 @@ -89,7 +93,9 @@ def test_map_enum_column_compares_with_its_text(duckdb_cursor): pytest.param("MAP(INTEGER, INTEGER)", "map([i::INTEGER], [i::INTEGER])", id="MAP_INTEGER_keys"), pytest.param("MAP(VARCHAR, BIGINT)", "map([i::VARCHAR], [i::BIGINT])", id="MAP_BIGINT_values"), pytest.param("INTEGER[2]", "[i, i]::INTEGER[2]", id="ARRAY"), - pytest.param("UNION(a INTEGER, b VARCHAR)", "i::INTEGER::UNION(a INTEGER, b VARCHAR)", id="UNION"), + pytest.param("MAP(BLOB, INTEGER)", "map([i::VARCHAR::BLOB], [i::INTEGER])", id="MAP_BLOB_keys"), + pytest.param("MAP(INTERVAL, INTEGER)", "map([to_days(i::INTEGER)], [i::INTEGER])", id="MAP_INTERVAL_keys"), + pytest.param("MAP(VARCHAR, STRUCT(x INTEGER))", "map(['k'], [NULL::STRUCT(x INTEGER)])", id="MAP_NULL_values"), pytest.param("TIME_NS", "'01:00:00'::TIME_NS", id="TIME_NS"), pytest.param( "STRUCT(a INTEGER[], b MAP(VARCHAR, INTEGER))", @@ -107,6 +113,54 @@ def test_map_identity_returns_an_object_column_in_its_own_type(duckdb_cursor, sq assert mapped.fetchall() == source.fetchall() +UNION_AB = "UNION(a INTEGER, b VARCHAR)" +UNION_OVERLAPPING = "UNION(a INTEGER, b BIGINT)" +# The function gets a UNION's member values only, so the tags cannot reach the result through it. +UNION_ROWS = pytest.mark.parametrize( + ("sql_type", "rows", "returned", "tags_cast"), + [ + pytest.param( + UNION_AB, + f"(union_value(a := 1)::{UNION_AB}), (union_value(b := 'x')::{UNION_AB}), (NULL::{UNION_AB})", + "INTEGER", + [("a", 1), ("b", "x"), (None, None)], + id="both_members", + ), + pytest.param( + UNION_AB, + f"(union_value(a := 1)::{UNION_AB}), (union_value(b := NULL::VARCHAR)::{UNION_AB})", + "INTEGER", + [("a", 1), (None, None)], + id="null_member", + ), + pytest.param( + UNION_OVERLAPPING, + f"(union_value(a := 1)::{UNION_OVERLAPPING}), (union_value(b := 7::BIGINT)::{UNION_OVERLAPPING})", + "INTEGER", + [("a", 1), ("a", 7)], + id="overlapping_members", + ), + ], +) + + +@UNION_ROWS +def test_map_identity_of_a_union_column_is_refused(duckdb_cursor, sql_type, rows, returned, tags_cast): + mapped = duckdb_cursor.sql(f"select * from (values {rows}) t(u)").map(lambda df: df) + expected = f"Invalid Input Error: UDF column type mismatch, expected [{sql_type}], got [{returned}]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +@UNION_ROWS +def test_map_union_column_in_the_schema_casts_each_member_value(duckdb_cursor, sql_type, rows, returned, tags_cast): + # The cast picks the member for each value, so a NULL member reads as a NULL UNION and 7 as the INTEGER member. + source = duckdb_cursor.sql(f"select * from (values {rows}) t(u)") + mapped = source.map(lambda df: df, schema={"u": duckdb.sqltype(sql_type)}) + assert mapped.types == source.types + assert mapped.project("union_tag(u), u").fetchall() == tags_cast + + @SCHEMAS def test_map_identity_returns_a_struct_whose_field_is_null_across_a_chunk(duckdb_cursor, explicit_schema): duckdb_cursor.execute("set threads = 1") @@ -217,6 +271,397 @@ def test_map_struct_cell_without_a_field_reads_it_as_null(duckdb_cursor, row): assert mapped.fetchall() == expected +@ODD_ROWS +def test_map_struct_cell_without_a_field_in_a_type_holding_a_map_reads_it_as_null(duckdb_cursor, row): + duckdb_cursor.execute("set threads = 1") + source = duckdb_cursor.sql("select {'x': i::INTEGER, 'm': map(['k'], [i::INTEGER])} as a from range(2048) t(i)") + mapped = source.map(_odd_cell_at(row, {"x": 3})) + expected = [({"x": 3, "m": None} if i == row else {"x": i, "m": {"k": i}},) for i in range(2048)] + assert mapped.types == source.types + assert mapped.fetchall() == expected + + +def _apply_to_a(convert): + return lambda df: df.assign(a=df["a"].apply(convert)) + + +CHANGED_IN_KIND = pytest.mark.parametrize( + ("source", "function", "sql_type", "returned"), + [ + pytest.param( + "{'x': i::BIGINT}", + _apply_to_a(lambda d: {"x": str(d["x"]).zfill(5)}), + "STRUCT(x BIGINT)", + "STRUCT(x VARCHAR)", + id="struct_field_turned_into_text", + ), + pytest.param( + "[i::INTEGER]", + _apply_to_a(lambda values: [str(v).zfill(5) for v in values]), + "INTEGER[]", + "VARCHAR[]", + id="list_elements_turned_into_text", + ), + pytest.param( + """'{"k": ' || i || '}'""", + _apply_to_a(json.loads), + "VARCHAR", + "STRUCT(k INTEGER)", + id="text_parsed_as_json", + ), + pytest.param( + "i", + _apply_to_a(lambda v: Decimal(int(v)) / 3), + "BIGINT", + "DECIMAL(28,28)", + id="integers_divided_as_decimals", + ), + pytest.param( + "map([i::INTEGER], [i::INTEGER])", + _apply_to_a(lambda m: {str(k).zfill(3): v for k, v in m.items()}), + "MAP(INTEGER, INTEGER)", + "MAP(VARCHAR, INTEGER)", + id="map_keys_turned_into_text", + ), + pytest.param( + "{'x': i::TINYINT}", + _apply_to_a(lambda d: {"x": d["x"] + 300}), + "STRUCT(x TINYINT)", + "STRUCT(x INTEGER)", + id="struct_field_outside_its_type", + ), + pytest.param( + "[i::INTEGER]", + _apply_to_a(lambda values: np.asarray(values, dtype=np.float64) / 2), + "INTEGER[]", + "DOUBLE[]", + id="list_given_as_a_float_array", + ), + pytest.param( + "{'x': i::BIGINT}", + _apply_to_a(lambda d: {"x": d["x"], "y": 1}), + "STRUCT(x BIGINT)", + "STRUCT(x INTEGER, y INTEGER)", + id="struct_given_a_key_it_does_not_name", + ), + pytest.param( + "{'t': TIMESTAMP '2020-01-01'}", + _apply_to_a(lambda d: {"t": np.datetime64("2020-01-01T00:00:00", "ns")}), + "STRUCT(t TIMESTAMP)", + "STRUCT(t BIGINT)", + id="nanosecond_datetime64_field", + ), + # A float of more than 15 significant digits, or one the cast into a wide scale rounds in double precision + pytest.param( + "{'x': i::BIGINT}", + _apply_to_a(lambda d: {"x": 9007199254740992.0}), + "STRUCT(x BIGINT)", + "STRUCT(x DOUBLE)", + id="float_of_16_digits", + ), + pytest.param( + "{'x': i::DECIMAL(38,30)}", + _apply_to_a(lambda d: {"x": 0.1}), + "STRUCT(x DECIMAL(38,30))", + "STRUCT(x DOUBLE)", + id="float_the_cast_rounds_at_a_wide_scale", + ), + pytest.param( + "{'x': i::BIGINT}", + _apply_to_a(lambda d: {"x": d["x"], "X": 2}), + "STRUCT(x BIGINT)", + "STRUCT(x INTEGER, X INTEGER)", + id="struct_given_two_keys_for_one_field", + ), + ], +) + + +@CHANGED_IN_KIND +def test_map_column_whose_cells_the_function_changed_in_kind_is_refused( + duckdb_cursor, source, function, sql_type, returned +): + mapped = duckdb_cursor.sql(f"select {source} as a from range(3) t(i)").map(function) + expected = f"Invalid Input Error: UDF column type mismatch, expected [{sql_type}], got [{returned}]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +NEW_YORK = "SET TimeZone = 'America/New_York'" +UTC = datetime.timezone.utc + + +@pytest.mark.parametrize( + ("source", "function", "sql_type", "returned"), + [ + # df() hands the elements over as naive UTC datetimes + pytest.param( + "[TIMESTAMPTZ '2020-01-01 12:00:00+00']", + lambda df: df, + "TIMESTAMP WITH TIME ZONE[]", + "TIMESTAMP[]", + id="list_of_timestamptz_identity", + ), + pytest.param( + "{'t': TIMESTAMP '2020-01-01 12:00:00'}", + _apply_to_a(lambda d: {"t": datetime.datetime(2020, 1, 1, 12, tzinfo=UTC)}), + "STRUCT(t TIMESTAMP)", + "STRUCT(t TIMESTAMP WITH TIME ZONE)", + id="aware_datetime_for_timestamp", + ), + pytest.param( + "{'t': TIMESTAMPTZ '2020-01-01 12:00:00+00'}", + _apply_to_a(lambda d: {"t": datetime.datetime(2020, 1, 1, 7)}), + "STRUCT(t TIMESTAMP WITH TIME ZONE)", + "STRUCT(t TIMESTAMP)", + id="naive_datetime_for_timestamptz", + ), + pytest.param( + "{'t': TIME '01:00:00'}", + _apply_to_a(lambda d: {"t": datetime.time(1, tzinfo=UTC)}), + "STRUCT(t TIME)", + "STRUCT(t TIME WITH TIME ZONE)", + id="aware_time_for_time", + ), + pytest.param( + "{'t': TIMETZ '01:00:00+00'}", + _apply_to_a(lambda d: {"t": datetime.time(1)}), + "STRUCT(t TIME WITH TIME ZONE)", + "STRUCT(t TIME)", + id="naive_time_for_timetz", + ), + ], +) +def test_map_value_with_a_time_zone_and_one_without_are_of_different_kinds( + duckdb_cursor, source, function, sql_type, returned +): + duckdb_cursor.execute(NEW_YORK) + mapped = duckdb_cursor.sql(f"select {source} as a from range(1) t(i)").map(function) + expected = f"Invalid Input Error: UDF column type mismatch, expected [{sql_type}], got [{returned}]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +@pytest.mark.parametrize( + ("source", "function", "sql_type", "returned"), + [ + pytest.param( + "{'t': '2020-01-01'::TIMESTAMP_S}", + _apply_to_a(lambda d: {"t": datetime.datetime(2020, 1, 1, 0, 0, 0, 500000)}), + "STRUCT(t TIMESTAMP_S)", + "STRUCT(t TIMESTAMP)", + id="milliseconds_for_seconds", + ), + pytest.param( + "{'t': '2020-01-01'::TIMESTAMP_MS}", + _apply_to_a(lambda d: {"t": datetime.datetime(2020, 1, 1, 0, 0, 0, 1)}), + "STRUCT(t TIMESTAMP_MS)", + "STRUCT(t TIMESTAMP)", + id="a_microsecond_for_milliseconds", + ), + pytest.param( + "{'t': TIME '01:00:00'}", + _apply_to_a(lambda d: {"t": duckdb.Value("01:00:00.000000001", duckdb.sqltype("TIME_NS"))}), + "STRUCT(t TIME)", + "STRUCT(t TIME_NS)", + id="a_nanosecond_for_microseconds", + ), + ], +) +def test_map_time_value_finer_than_its_column_keeps_only_values_the_unit_holds( + duckdb_cursor, source, function, sql_type, returned +): + mapped = duckdb_cursor.sql(f"select {source} as a from range(1) t(i)").map(function) + expected = f"Invalid Input Error: UDF column type mismatch, expected [{sql_type}], got [{returned}]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +def test_map_time_value_outside_the_range_of_its_columns_unit_is_refused(duckdb_cursor): + # TIMESTAMP_NS covers the years 1677 to 2262 only + mapped = duckdb_cursor.sql("select {'t': '2020-01-01'::TIMESTAMP_NS} as a from range(1) t(i)").map( + _apply_to_a(lambda d: {"t": datetime.datetime(1500, 1, 1)}) + ) + expected = ( + "Invalid Input Error: UDF column type mismatch, expected [STRUCT(t TIMESTAMP_NS)], got [STRUCT(t TIMESTAMP)]" + ) + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +# df() hands a LIST of HUGEINT or DECIMAL over as float64 elements, and a float of more than 15 significant digits +# may already have rounded the value it came from +@pytest.mark.parametrize( + ("values", "sql_type"), + [ + pytest.param("[1152921504606846977::HUGEINT]", "HUGEINT[]", id="hugeint_beyond_2_53"), + pytest.param("[123456789012345678901::HUGEINT]", "HUGEINT[]", id="hugeint_beyond_bigint"), + pytest.param("[1152921504606846977::HUGEINT, NULL]", "HUGEINT[]", id="hugeint_with_a_null"), + pytest.param("[1234567890123456.78::DECIMAL(18,2)]", "DECIMAL(18,2)[]", id="decimal_of_18_digits"), + ], +) +def test_map_identity_of_a_list_whose_floats_may_have_rounded_is_refused(duckdb_cursor, values, sql_type): + mapped = duckdb_cursor.sql(f"select {values} as a").map(lambda df: df) + expected = f"Invalid Input Error: UDF column type mismatch, expected [{sql_type}], got [DOUBLE[]]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +@pytest.mark.parametrize( + "values", + [ + pytest.param("[5::HUGEINT, NULL, (-999999999999999)::HUGEINT]", id="hugeint_of_15_digits"), + pytest.param("[1.25::DECIMAL(9,2), NULL, 0.1::DECIMAL(9,2)]", id="decimal"), + ], +) +def test_map_identity_of_a_list_whose_floats_hold_their_values_round_trips(duckdb_cursor, values): + source = duckdb_cursor.sql(f"select {values} as a") + mapped = source.map(lambda df: df) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + + +def test_map_decimal_column_handed_back_as_objects_of_18_digits_is_refused(duckdb_cursor): + # df() hands a DECIMAL column over as float64, which rounds 1234567890123456.78 + source = duckdb_cursor.sql("select 1234567890123456.78::DECIMAL(18,2) as a") + mapped = source.map(lambda df: df.assign(a=df["a"].astype(object))) + expected = "Invalid Input Error: UDF column type mismatch, expected [DECIMAL(18,2)], got [DOUBLE]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +@pytest.mark.parametrize( + ("source", "cell", "expected"), + [ + pytest.param( + "{'t': '2020-01-01'::TIMESTAMP_S}", + datetime.datetime(2020, 1, 1, 0, 0, 3), + datetime.datetime(2020, 1, 1, 0, 0, 3), + id="whole_seconds_into_seconds", + ), + pytest.param( + "{'t': '2020-01-01'::TIMESTAMP_NS}", + datetime.datetime(2020, 1, 1, 0, 0, 0, 7), + datetime.datetime(2020, 1, 1, 0, 0, 0, 7), + id="microseconds_into_nanoseconds", + ), + pytest.param( + "{'t': TIME '01:00:00'}", + duckdb.Value("01:00:00.000001000", duckdb.sqltype("TIME_NS")), + datetime.time(1, 0, 0, 1), + id="whole_microseconds_into_microseconds", + ), + pytest.param( + "{'t': '01:00:00'::TIME_NS}", + datetime.time(2, 0, 0, 5), + datetime.time(2, 0, 0, 5), + id="time_into_nanoseconds", + ), + ], +) +def test_map_time_value_its_column_unit_holds_reads_unchanged(duckdb_cursor, source, cell, expected): + relation = duckdb_cursor.sql(f"select {source} as a from range(1) t(i)") + mapped = relation.map(_apply_to_a(lambda d: {"t": cell})) + assert mapped.types == relation.types + assert mapped.fetchall() == [({"t": expected},)] + + +@pytest.mark.parametrize( + ("source", "function", "expected"), + [ + pytest.param("[1, 2]", _apply_to_a(lambda values: [1, True]), [1, 1], id="int_and_bool_list"), + pytest.param( + "map(['a'], ['x'])", _apply_to_a(lambda m: {"a": 1, "b": "x"}), {"a": "1", "b": "x"}, id="int_and_str_map" + ), + ], +) +def test_map_container_of_mixed_elements_reads_as_the_conversion_types_it(duckdb_cursor, source, function, expected): + # The conversion types a list or a MAP by the common type of its elements; the base refused these + relation = duckdb_cursor.sql(f"select {source} as a from range(1) t(i)") + mapped = relation.map(function) + assert mapped.types == relation.types + assert mapped.fetchall() == [(expected,)] + + +@pytest.mark.parametrize( + "value", + [ + pytest.param("{'t': TIMESTAMPTZ '2020-01-01 12:00:00+00'}", id="struct_field"), + pytest.param("map(['k'], [TIMESTAMPTZ '2020-01-01 12:00:00+00'])", id="map_value"), + pytest.param("{'t': TIMETZ '01:00:00+02'}", id="timetz_field"), + ], +) +def test_map_identity_of_a_zoned_value_keeps_its_instant_in_another_session_time_zone(duckdb_cursor, value): + duckdb_cursor.execute(NEW_YORK) + source = duckdb_cursor.sql(f"select {value} as a") + mapped = source.map(lambda df: df) + assert mapped.types == source.types + assert mapped.project("a::VARCHAR").fetchall() == source.project("a::VARCHAR").fetchall() + + +@SCHEMAS +def test_map_numpy_scalar_inside_a_struct_raises_the_conversions_error(duckdb_cursor, explicit_schema): + source = duckdb_cursor.sql("select {'x': i::INTEGER} as a from range(1) t(i)") + mapped = source.map(_apply_to_a(lambda d: {"x": np.int64(5)}), schema=_schema_of(source, explicit_schema)) + expected = ( + "Not implemented Error: UDF column 'a': Unable to transform python value of type " + "'' to DuckDB LogicalType" + ) + with pytest.raises(duckdb.NotImplementedException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +def test_map_column_whose_cells_changed_in_kind_casts_into_the_schema_type(duckdb_cursor): + mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map( + _apply_to_a(lambda v: Decimal(int(v)) / 3), schema=BIGINT_A + ) + assert mapped.fetchall() == [(0,), (0,), (1,)] + + +@pytest.mark.parametrize( + ("source", "function", "expected"), + [ + pytest.param("{'x': i::TINYINT}", _apply_to_a(lambda d: {"x": d["x"] + 100}), {"x": 100}, id="int_that_fits"), + pytest.param("[i::DOUBLE]", _apply_to_a(lambda values: [int(v) for v in values]), [0.0], id="int_into_double"), + pytest.param("{'x': i::FLOAT}", _apply_to_a(lambda d: {"x": d["x"] / 2}), {"x": 0.0}, id="float_into_float"), + pytest.param( + "{'x': i::DECIMAL(9,2)}", _apply_to_a(lambda d: {"x": Decimal("1.5")}), {"x": Decimal("1.50")}, id="decimal" + ), + pytest.param( + "{'x': i::BIGINT}", _apply_to_a(lambda d: {"X": d["x"]}), {"x": 0}, id="field_named_in_other_case" + ), + pytest.param( + "map([i::INTEGER], [i::INTEGER])", + _apply_to_a(lambda m: {k + 5: v for k, v in m.items()}), + {5: 0}, + id="map_keys_of_their_kind", + ), + pytest.param( + "[i::INTEGER]", _apply_to_a(lambda values: np.asarray(values, dtype=np.int64) + 3), [3], id="int64_array" + ), + pytest.param( + "[i::DOUBLE]", _apply_to_a(lambda values: np.asarray(values, dtype=np.float32)), [0.0], id="float32_array" + ), + pytest.param("{'x': i::BIGINT}", _apply_to_a(lambda d: {"x": Decimal(5)}), {"x": 5}, id="integral_decimal"), + pytest.param("{'x': i::BIGINT}", _apply_to_a(lambda d: {"x": 2.0}), {"x": 2}, id="integral_float_for_bigint"), + pytest.param( + "{'x': i::DECIMAL(18,2)}", + _apply_to_a(lambda d: {"x": 1234567890123.45}), + {"x": Decimal("1234567890123.45")}, + id="float_of_15_digits", + ), + pytest.param( + "{'x': i::DECIMAL(9,2)}", _apply_to_a(lambda d: {"x": 2.0}), {"x": Decimal("2.00")}, id="integral_float" + ), + ], +) +def test_map_column_whose_cells_keep_their_kind_reads_into_the_input_type(duckdb_cursor, source, function, expected): + relation = duckdb_cursor.sql(f"select {source} as a from range(1) t(i)") + mapped = relation.map(function) + assert mapped.types == relation.types + assert mapped.fetchall() == [(expected,)] + + def test_map_object_column_of_zero_dimensional_arrays_reads_their_values(duckdb_cursor): def arrays(df): cells = np.empty(len(df), dtype=object) @@ -239,17 +684,101 @@ def test_map_declared_column_read_error_names_the_column(duckdb_cursor): @pytest.mark.parametrize( ("cells", "expected"), [ - # Every staged value is an int, so one cast reads them; the str rows between them are written directly. + # A text cell and cells of other kinds in one vector, the NULL among them pytest.param(["a", 5, 6, "b", None, 7, "c"], ["a", "5", "6", "b", None, "7", "c"], id="str_int_and_null"), pytest.param(["a", 5, np.str_("c"), None], ["a", "5", "c", None], id="with_a_str_subclass"), ], ) -def test_map_text_column_mixing_str_with_other_cells_reads_each_as_text(duckdb_cursor, cells, expected): +def test_map_text_column_mixing_str_with_other_cells_reads_each_as_text_in_the_schema_type( + duckdb_cursor, cells, expected +): source = duckdb_cursor.sql(f"select i::VARCHAR as a from range({len(cells)}) t(i)") - mapped = source.map(_text_cells(cells)) + mapped = source.map(_text_cells(cells), schema={"a": duckdb.sqltypes.VARCHAR}) assert mapped.fetchall() == [(value,) for value in expected] +@SCHEMAS +def test_map_text_list_mixing_str_with_an_int_reads_as_text_as_on_the_base(duckdb_cursor, explicit_schema): + # A list converts by the common type of its elements, so the int reads as its text inside a text list + source = duckdb_cursor.sql("select ['x', 'y'] as a from range(1) t(i)") + mapped = source.map( + _apply_to_a(lambda values: np.array(["x", 5], dtype=object)), schema=_schema_of(source, explicit_schema) + ) + assert mapped.types == source.types + assert mapped.fetchall() == [(["x", "5"],)] + + +def test_map_text_column_of_str_subclass_cells_reads_them_without_a_schema(duckdb_cursor): + mapped = duckdb_cursor.sql("select i::VARCHAR as a from range(3) t(i)").map(_text_cells(["a", np.str_("c"), None])) + assert mapped.fetchall() == [("a",), ("c",), (None,)] + + +def test_map_text_column_with_a_cell_of_another_kind_is_refused_without_a_schema(duckdb_cursor): + mapped = duckdb_cursor.sql("select i::VARCHAR as a from range(4) t(i)").map( + _text_cells(["a", np.str_("c"), 5, None]) + ) + expected = "Invalid Input Error: UDF column type mismatch, expected [VARCHAR], got [INTEGER]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +@SCHEMAS +def test_map_blob_column_mixing_bytes_with_other_cells_reads_each(duckdb_cursor, explicit_schema): + # Every Python type a BLOB value reads from, in one vector + cells = [b"a", bytearray(b"b"), None, memoryview(b"c"), b"\x00d"] + source = duckdb_cursor.sql(f"select i::VARCHAR::BLOB as a from range({len(cells)}) t(i)") + mapped = source.map( + lambda df: df.assign(a=pd.Series(cells[: len(df)], dtype=object, index=df.index)), + schema=_schema_of(source, explicit_schema), + ) + assert mapped.fetchall() == [(b"a",), (b"b",), (None,), (b"c",), (b"\x00d",)] + + +@pytest.mark.parametrize( + ("cell", "name"), + [pytest.param(Fraction(1, 3), "fractions.Fraction", id="Fraction"), pytest.param({"b"}, "set", id="set")], +) +def test_map_text_column_with_an_object_without_a_duckdb_type_is_refused_without_a_schema(duckdb_cursor, cell, name): + mapped = duckdb_cursor.sql("select i::VARCHAR as a from range(2) t(i)").map(_text_cells(["a", cell])) + expected = ( + f"Not implemented Error: UDF column 'a': Unable to transform python value of type '' to " + "DuckDB LogicalType" + ) + with pytest.raises(duckdb.NotImplementedException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +def test_map_text_column_with_an_object_without_a_duckdb_type_reads_its_text_in_the_schema_type(duckdb_cursor): + source = duckdb_cursor.sql("select i::VARCHAR as a from range(2) t(i)") + mapped = source.map(_text_cells(["a", Fraction(1, 3)]), schema={"a": duckdb.sqltypes.VARCHAR}) + assert mapped.fetchall() == [("a",), ("1/3",)] + + +@SCHEMAS +@pytest.mark.parametrize( + ("value", "function", "sql_type", "returned"), + [ + pytest.param("'101'::BIT", lambda df: df, "BIT", "BLOB", id="column"), + pytest.param("['101'::BIT]", lambda df: df, "BIT[]", "BLOB[]", id="list_element"), + pytest.param( + "'101'::BIT", + lambda df: df.assign(a=pd.Series([memoryview(bytes(v)) for v in df["a"]], dtype=object, index=df.index)), + "BIT", + "BLOB", + id="memoryview", + ), + ], +) +def test_map_bit_value_reaching_the_function_as_its_storage_bytes_is_refused( + duckdb_cursor, explicit_schema, value, function, sql_type, returned +): + source = duckdb_cursor.sql(f"select {value} as a from range(2) t(i)") + mapped = source.map(function, schema=_schema_of(source, explicit_schema)) + expected = f"Invalid Input Error: UDF column type mismatch, expected [{sql_type}], got [{returned}]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + def test_map_json_column_still_validates_its_text(duckdb_cursor): schema = {"a": duckdb.sqltype("JSON")} good = duckdb_cursor.sql("select 1 as a").map(_text_cells(['{"k": 1}']), schema=schema) @@ -336,39 +865,142 @@ def array_cell(df): @pytest.mark.parametrize( - ("function", "column"), + ("function", "types"), [ pytest.param( lambda df: df.assign(a=pd.Series(df["a"].astype(str).str.zfill(5).tolist(), dtype=object, index=df.index)), - "a", + ("BIGINT, VARCHAR", "VARCHAR, VARCHAR"), id="typed_input_turned_into_text", ), - pytest.param(lambda df: df.rename(columns={"a": "b", "b": "a"}), "a", id="names_swapped"), + # rename keeps the order, so the BIGINT column comes first under the name b + pytest.param( + lambda df: df.rename(columns={"a": "b", "b": "a"}), + ("BIGINT, BIGINT", "BIGINT, VARCHAR"), + id="names_swapped", + ), ], ) -def test_map_typed_input_column_the_function_moved_to_another_type_is_refused(duckdb_cursor, function, column): +def test_map_typed_input_column_the_function_moved_to_another_type_is_refused(duckdb_cursor, function, types): mapped = duckdb_cursor.sql("select i as a, i::VARCHAR as b from range(3) t(i)").map(function) - with pytest.raises(duckdb.InvalidInputException, match=_no_type(column)): + expected = f"Invalid Input Error: UDF column type mismatch, expected [{types[0]}], got [{types[1]}]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): mapped.fetchall() @pytest.mark.parametrize( - ("make_cell", "expected"), + ("sql_type", "offset"), [ - pytest.param(int, [(0,), (1,), (2,)], id="python_ints"), - pytest.param(lambda v: None, [(None,), (None,), (None,)], id="only_none"), + pytest.param("INTEGER", 0, id="INTEGER"), + # Python ints in the INTEGER range for a BIGINT input + pytest.param("BIGINT", 0, id="BIGINT_holding_INTEGER_values"), + pytest.param("BIGINT", 2**40, id="BIGINT_holding_BIGINT_values"), + pytest.param("HUGEINT", 0, id="HUGEINT_holding_INTEGER_values"), ], ) +@pytest.mark.parametrize("only_none", [False, True], ids=["python_ints", "only_none"]) def test_map_typed_input_column_the_function_turned_into_objects_keeps_a_confirmed_type( - duckdb_cursor, make_cell, expected + duckdb_cursor, sql_type, offset, only_none ): def objects(df): - return df.assign(a=pd.Series([make_cell(v) for v in df["a"]], dtype=object, index=df.index)) + return df.assign(a=pd.Series([None if only_none else int(v) for v in df["a"]], dtype=object, index=df.index)) - source = duckdb_cursor.sql("select i::INTEGER as a from range(3) t(i)") + source = duckdb_cursor.sql(f"select (i + {offset})::{sql_type} as a from range(3) t(i)") mapped = source.map(objects) assert mapped.types == source.types - assert mapped.fetchall() == expected + assert mapped.fetchall() == [(None if only_none else i + offset,) for i in range(3)] + + +@ODD_ROWS +def test_map_typed_input_column_turned_into_objects_reads_a_wide_value_at_any_row(duckdb_cursor, row): + duckdb_cursor.execute("set threads = 1") + + def objects(df): + return df.assign(a=pd.Series([2**40 if i == row else int(v) for i, v in enumerate(df["a"])], dtype=object)) + + mapped = duckdb_cursor.sql("select i as a from range(2048) t(i)").map(objects) + assert mapped.types == [duckdb.sqltypes.BIGINT] + assert mapped.fetchall() == [(2**40 if i == row else i,) for i in range(2048)] + + +@pytest.mark.parametrize("sampled", [0, 2**40], ids=["small_sample", "wide_sample"]) +# In a 2048 row chunk the analyzer samples the even rows, so row 6 is inside its sample and row 7 outside it +@pytest.mark.parametrize("row", [6, 7], ids=["sampled_row", "unsampled_row"]) +@pytest.mark.parametrize( + ("cell", "outcome"), [pytest.param(2.0, 2, id="integral_float"), pytest.param(2.5, None, id="fractional_float")] +) +def test_map_typed_input_column_turned_into_objects_reads_a_float_the_same_whatever_the_sample( + duckdb_cursor, sampled, row, cell, outcome +): + duckdb_cursor.execute("set threads = 1") + + def objects(df): + values = [sampled if i == 0 else cell if i == row else int(v) for i, v in enumerate(df["a"])] + return df.assign(a=pd.Series(values, dtype=object, index=df.index)) + + mapped = duckdb_cursor.sql("select i as a from range(2048) t(i)").map(objects) + if outcome is None: + expected = "Invalid Input Error: UDF column type mismatch, expected [BIGINT], got [DOUBLE]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + else: + assert mapped.fetchall() == [(sampled if i == 0 else outcome if i == row else i,) for i in range(2048)] + + +@pytest.mark.parametrize( + ("sql_type", "dtype", "values"), + [ + pytest.param("BIGINT", np.int32, [0, 1, 2], id="int32_chunk_for_bigint"), + pytest.param("DOUBLE", np.float32, [0.5, 1.5, 2.5], id="float32_chunk_for_double"), + ], +) +def test_map_typed_chunk_for_a_guess_widens_into_it(duckdb_cursor, sql_type, dtype, values): + # A Series built from no values is object, so the bind frame's column takes the input's type as a guess + def typed_chunk(df): + cells = np.asarray(df["a"], dtype=dtype) if len(df) else [] + return pd.DataFrame({"a": pd.Series(cells, dtype=None if len(df) else object)}) + + source = duckdb_cursor.sql(f"select (i + {values[0]})::{sql_type} as a from range(3) t(i)") + mapped = source.map(typed_chunk) + assert mapped.types == source.types + assert mapped.fetchall() == [(v,) for v in values] + + +def test_map_double_input_column_turned_into_float_values_reads_them_as_double(duckdb_cursor): + def floats(df): + cells = [duckdb.Value(float(v), duckdb.sqltypes.FLOAT) for v in df["a"]] + return df.assign(a=pd.Series(cells, dtype=object, index=df.index)) + + mapped = duckdb_cursor.sql("select (i + 0.5)::DOUBLE as a from range(3) t(i)").map(floats) + assert mapped.types == [duckdb.sqltypes.DOUBLE] + assert mapped.fetchall() == [(0.5,), (1.5,), (2.5,)] + + +@pytest.mark.parametrize(("offset", "fits"), [(0, True), (40000, False)], ids=["values_it_holds", "values_beyond_it"]) +def test_map_typed_input_column_turned_into_objects_of_a_narrow_type_reads_the_values_it_holds( + duckdb_cursor, offset, fits +): + def objects(df): + return df.assign(a=pd.Series([int(v) + offset for v in df["a"]], dtype=object, index=df.index)) + + mapped = duckdb_cursor.sql("select i::SMALLINT as a from range(3) t(i)").map(objects) + if fits: + assert mapped.types == [duckdb.sqltypes.SMALLINT] + assert mapped.fetchall() == [(0,), (1,), (2,)] + else: + expected = "Invalid Input Error: UDF column type mismatch, expected [SMALLINT], got [INTEGER]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +def test_map_new_object_column_named_otherwise_after_bind_reports_the_name_mismatch(duckdb_cursor): + def late_name(df): + name = "t" if len(df) else "s" + return df.assign(**{name: pd.Series(["x"] * len(df), dtype=object, index=df.index)}) + + mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map(late_name) + expected = "Invalid Input Error: UDF column name mismatch, expected [a, s], got [a, t]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() def test_map_new_object_column_fails_on_the_first_chunk_with_a_value(duckdb_cursor): @@ -395,6 +1027,9 @@ def late_values(df): pytest.param(pd.NA, id="pd.NA"), pytest.param(pd.NaT, id="pandas_NaT"), pytest.param(np.datetime64("NaT", "us"), id="numpy_NaT"), + # pandas reads neither as missing; the conversion reads both as NULL + pytest.param(duckdb.NullValue(), id="duckdb_NullValue"), + pytest.param(np.ma.masked, id="numpy_masked"), ], ) def test_map_new_object_column_of_only_missing_values_stays_null(duckdb_cursor, missing): @@ -404,13 +1039,3 @@ def note(df): mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map(note) assert mapped.types == [duckdb.sqltypes.BIGINT, duckdb.sqltypes.SQLNULL] assert mapped.fetchall() == [(0, None), (1, None), (2, None)] - - -def test_map_new_object_column_of_null_values_is_refused_as_untyped(duckdb_cursor): - # pandas does not read a duckdb.NullValue() as missing, so the column holds values bind could not type - def nulls(df): - return df.assign(s=pd.Series([duckdb.NullValue() for _ in range(len(df))], dtype=object, index=df.index)) - - mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map(nulls) - with pytest.raises(duckdb.InvalidInputException, match=_no_type("s")): - mapped.fetchall() From 2ed45dd15040d77631b28ec5ebda5a6ec1c2e651 Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Fri, 6 Nov 2026 18:11:47 +0000 Subject: [PATCH 3/9] Build the stable ABI, lend no type from a NULL input, shape only container cells toward their key and value types, keep nanoseconds and check narrow integers by range --- .../duckdb_python/python_conversion.hpp | 15 + src/map.cpp | 6 +- src/native/python_conversion.cpp | 143 +++++++- src/numpy/numpy_scan.cpp | 237 +++++++++----- tests/fast/pandas/test_append_df.py | 32 +- tests/fast/test_map_types.py | 304 +++++++++++++++++- 6 files changed, 648 insertions(+), 89 deletions(-) diff --git a/src/include/duckdb_python/python_conversion.hpp b/src/include/duckdb_python/python_conversion.hpp index ac4fa11a..68ffce1f 100644 --- a/src/include/duckdb_python/python_conversion.hpp +++ b/src/include/duckdb_python/python_conversion.hpp @@ -53,6 +53,21 @@ LogicalType SniffPythonIntegerType(nb::handle ele); //! The type of a duckdb.Value, or InvalidInputException when its 'type' attribute is not a DuckDB type LogicalType PythonValueType(nb::handle value); bool DictionaryHasMapFormat(const PyDictionary &dict); +//! A target the declared scan builds from a column's type (a shape), as opposed to a column's own type: a dict for a +//! shaped STRUCT takes its own keys, as without a target; a shaped MAP casts each key and value whose type has no +//! children with the connection's settings; a shaped naive timestamp reads a cell as TransformTimestamp does +LogicalType AsShape(const LogicalType &type); +bool IsShape(const LogicalType &type); +//! Whether a cell is a naive pandas Timestamp or a numpy datetime64 finer than a microsecond, whose nanoseconds no +//! Python datetime holds; they are then in nanoseconds +bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds); +//! A numpy scalar's Python value; a datetime64 or timedelta64 at nanoseconds would read as an integer, so those take +//! the microsecond unit first, as a frame's own datetime columns do +nb::object NumpyScalarItem(nb::handle scalar); +//! A value for a naive timestamp type converted without a target, but keeping the nanoseconds of a pandas Timestamp +//! or a fine datetime64 for a TIMESTAMP_NS, and reading a datetime64 at microseconds for a coarser unit +Value TransformTimestamp(optional_ptr context, nb::handle ele, const LogicalType &target_type, + bool nan_as_null); void TransformPythonObject(optional_ptr context, nb::handle ele, Vector &vector, idx_t result_offset, bool nan_as_null = true); Value TransformPythonValue(optional_ptr context, nb::handle ele, diff --git a/src/map.cpp b/src/map.cpp index b73d286e..b17ed547 100644 --- a/src/map.cpp +++ b/src/map.cpp @@ -45,7 +45,7 @@ struct MapFunctionData : public TableFunctionData { vector in_types, out_types; vector in_names, out_names; vector sources; - //! Per output column, the type a returned object column is read into, INVALID where the analyzer types it + //! Per output column, the type a returned object column is declared: the column's own type, NULL for an untyped one vector declared_types; //! Per input column, the CategoricalDtype an ENUM column is handed over in, built once since it lists every value vector enum_dtypes; @@ -124,7 +124,9 @@ static void OverrideNullType(vector &return_types, const vector context, const PyD dict.ToString(), target_type.ToString()); }; bool struct_target = StructType::IsStruct(target_type); - // A target holding UNKNOWN gives only the shape of the values below it: the dict's own keys name the fields, as - // without a target, and each value converts with the shape of the field of its name - bool shape_only = struct_target && TypeVisitor::Contains(target_type, LogicalTypeId::UNKNOWN); + // A shape gives only the containers of the values below it: the dict's own keys name the fields, as without a + // target, and each value converts with the shape of the field of its name + bool shape_only = struct_target && IsShape(target_type); if (struct_target && !shape_only && dict.len != StructType::GetChildCount(target_type)) { throw conversion_error(); } @@ -213,6 +213,117 @@ Value TransformDictionaryToStruct(optional_ptr context, const PyD return Value::STRUCT(std::move(struct_values)); } +// Short enough to stay in a string's own buffer, since every shaped conversion compares it +static constexpr const char *SHAPE_ALIAS = "duckdb_py_shape"; + +LogicalType AsShape(const LogicalType &type) { + return type.WithAlias(SHAPE_ALIAS); +} + +bool IsShape(const LogicalType &type) { + return type.HasAlias() && type.GetAlias() == SHAPE_ALIAS; +} + +bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds) { + auto &import_cache = *DuckDBPyConnection::ImportCache(); + auto datetime_type = import_cache.datetime.datetime(); + if (Py_TYPE(cell.ptr()) == reinterpret_cast(datetime_type.ptr())) { + return false; + } + if (duckdb::PyUtil::IsInstance(cell, import_cache.numpy.datetime64())) { + auto code = nb::cast(cell.attr("dtype").attr("str")); + if (!StringUtil::EndsWith(code, "ns]") && !StringUtil::EndsWith(code, "ps]") && + !StringUtil::EndsWith(code, "fs]") && !StringUtil::EndsWith(code, "as]")) { + return false; + } + auto value = cell.attr("astype")("datetime64[ns]").attr("item")(); + if (value.is_none()) { + return false; + } + nanoseconds = timestamp_ns_t(nb::cast(value)); + return true; + } + if (!duckdb::PyUtil::IsInstance(cell, datetime_type) || !PandasIsImported(import_cache) || + cell.is(import_cache.pandas.NaT()) || !nb::hasattr(cell, "nanosecond") || + !nb::getattr(cell, "tzinfo").is_none()) { + return false; + } + nanoseconds = timestamp_ns_t(nb::cast(cell.attr("value"))); + return true; +} + +nb::object NumpyScalarItem(nb::handle scalar) { + auto kind = nb::cast(scalar.attr("dtype").attr("kind")); + if (kind == "M") { + return scalar.attr("astype")("datetime64[us]").attr("item")(); + } + if (kind == "m") { + return scalar.attr("astype")("timedelta64[us]").attr("item")(); + } + return scalar.attr("item")(); +} + +Value TransformTimestamp(optional_ptr context, nb::handle ele, const LogicalType &target_type, + bool nan_as_null) { + switch (target_type.id()) { + case LogicalTypeId::TIMESTAMP_NS: { + timestamp_ns_t nanoseconds; + if (HoldsNanoseconds(ele, nanoseconds)) { + return Value::TIMESTAMPNS(nanoseconds); + } + break; + } + case LogicalTypeId::TIMESTAMP: + case LogicalTypeId::TIMESTAMP_MS: + case LogicalTypeId::TIMESTAMP_SEC: + if (duckdb::PyUtil::IsInstance(ele, DuckDBPyConnection::ImportCache()->numpy.datetime64())) { + return TransformPythonValue(context, NumpyScalarItem(ele), LogicalType::UNKNOWN, nan_as_null); + } + break; + default: + break; + } + return TransformPythonValue(context, ele, LogicalType::UNKNOWN, nan_as_null); +} + +//! Whether an object is what a shape is for: a dict for a MAP or a STRUCT, a list, a tuple or a numpy array for a +//! LIST or an ARRAY; any other object converts without a target and meets the engine's cast or the kind check, so a +//! tuple reads as a LIST at every depth, as the appender reads a value of a row +static bool ShapeTakes(nb::handle ele, const LogicalType &shape) { + switch (shape.id()) { + case LogicalTypeId::MAP: + case LogicalTypeId::STRUCT: + return PyDict_Check(ele.ptr()); + case LogicalTypeId::LIST: + case LogicalTypeId::ARRAY: + return PyList_Check(ele.ptr()) || PyTuple_Check(ele.ptr()) || + duckdb::PyUtil::IsInstance(ele, DuckDBPyConnection::ImportCache()->numpy.ndarray()); + default: + return false; + } +} + +//! A key or value of a MAP whose target is a type without children converts without a target and is then cast with +//! the connection's settings, as a value of a column is, so each keeps its own type until the cast instead of meeting +//! its siblings' common type (2**60 + 1 beside 1.0 would read as a DOUBLE, rounded) +static Value TransformMapEntry(optional_ptr context, nb::handle entry, const LogicalType &target_type, + bool nan_as_null) { + if (IsShape(target_type) || target_type.IsNested() || target_type.id() == LogicalTypeId::UNKNOWN || + target_type.id() == LogicalTypeId::UNION) { + return TransformPythonValue(context, entry, target_type, nan_as_null); + } + auto value = TransformTimestamp(context, entry, target_type, nan_as_null); + if (!context) { + return value.DefaultCastAs(target_type); + } + string error; + auto cast = value.TryCastAs(*context, target_type, &error); + if (!cast) { + throw ConversionException(error); + } + return std::move(*cast); +} + Value TransformStructFormatDictionaryToMap(optional_ptr context, const PyDictionary &dict, const LogicalType &target_type, bool nan_as_null) { if (dict.len == 0) { @@ -232,6 +343,7 @@ Value TransformStructFormatDictionaryToMap(optional_ptr context, auto key_target = MapType::KeyType(target_type); auto value_target = MapType::ValueType(target_type); + auto shape = IsShape(target_type); LogicalType key_type = LogicalType::SQLNULL; LogicalType value_type = LogicalType::SQLNULL; @@ -239,8 +351,12 @@ Value TransformStructFormatDictionaryToMap(optional_ptr context, vector elements; for (idx_t i = 0; i < size; i++) { - Value new_key = TransformPythonValue(context, dict.keys.attr("__getitem__")(i), key_target, nan_as_null); - Value new_value = TransformPythonValue(context, dict.values.attr("__getitem__")(i), value_target, nan_as_null); + auto key = dict.keys.attr("__getitem__")(i); + auto value = dict.values.attr("__getitem__")(i); + Value new_key = shape ? TransformMapEntry(context, key, key_target, nan_as_null) + : TransformPythonValue(context, key, key_target, nan_as_null); + Value new_value = shape ? TransformMapEntry(context, value, value_target, nan_as_null) + : TransformPythonValue(context, value, value_target, nan_as_null); key_type = ProxiedForceMaxLogicalType(context, key_type, new_key.type()); value_type = ProxiedForceMaxLogicalType(context, value_type, new_value.type()); @@ -252,10 +368,12 @@ Value TransformStructFormatDictionaryToMap(optional_ptr context, elements.push_back(Value::STRUCT(std::move(struct_values))); } // A target holding UNKNOWN asks for no type there, and a MAP of such a type is no value the engine can hold - if (key_type.id() == LogicalTypeId::SQLNULL && !TypeVisitor::Contains(key_target, LogicalTypeId::UNKNOWN)) { + if (key_type.id() == LogicalTypeId::SQLNULL && !TypeVisitor::Contains(key_target, LogicalTypeId::UNKNOWN) && + !IsShape(key_target)) { key_type = key_target; } - if (value_type.id() == LogicalTypeId::SQLNULL && !TypeVisitor::Contains(value_target, LogicalTypeId::UNKNOWN)) { + if (value_type.id() == LogicalTypeId::SQLNULL && !TypeVisitor::Contains(value_target, LogicalTypeId::UNKNOWN) && + !IsShape(value_target)) { value_type = value_target; } @@ -1222,6 +1340,17 @@ void TransformPythonObject(optional_ptr context, nb::handle ele, Value TransformPythonValue(optional_ptr context, nb::handle ele, const LogicalType &target_type, bool nan_as_null) { + // A shape of no children marks a timestamp, read as TransformTimestamp reads it; a container shape takes only the + // objects it is for, so a scalar inside a container meets the engine's cast as at the top, never a type holding + // UNKNOWN + if (IsShape(target_type)) { + if (!target_type.IsNested()) { + return TransformTimestamp(context, ele, target_type, nan_as_null); + } + if (!ShapeTakes(ele, target_type)) { + return TransformPythonValue(context, ele, LogicalType::UNKNOWN, nan_as_null); + } + } Value result; TransformPythonObjectInternal(context, ele, result, target_type, nan_as_null); return result; diff --git a/src/numpy/numpy_scan.cpp b/src/numpy/numpy_scan.cpp index 9242efad..ece998c8 100644 --- a/src/numpy/numpy_scan.cpp +++ b/src/numpy/numpy_scan.cpp @@ -307,9 +307,7 @@ class DeclaredCasts { //! A declared column converts a cell as the appender converts a value of a row: with the type the value has in //! Python, then the engine's cast into the column's type, which reads the connection's settings; NaN reads NULL at //! every depth, and a value without a DuckDB type of its own reads as its text in a text column -//! A numpy scalar, or an array of no dimensions, reads as its Python value; a datetime64 or timedelta64 at -//! nanoseconds would read as an integer, so those take the microsecond unit first, as the frame's own datetime -//! columns do +//! A numpy scalar, or an array of no dimensions, reads as its Python value (NumpyScalarItem) static bool IsNumpyScalar(nb::handle cell) { auto &import_cache = *DuckDBPyConnection::ImportCache(); if (nb::isinstance(cell, import_cache.numpy.generic())) { @@ -319,17 +317,6 @@ static bool IsNumpyScalar(nb::handle cell) { return cell.type().is(import_cache.numpy.ndarray()) && nb::cast(cell.attr("ndim")) == 0; } -static nb::object NumpyScalarValue(nb::handle scalar) { - auto kind = nb::cast(scalar.attr("dtype").attr("kind")); - if (kind == "M") { - return scalar.attr("astype")("datetime64[us]").attr("item")(); - } - if (kind == "m") { - return scalar.attr("astype")("timedelta64[us]").attr("item")(); - } - return scalar.attr("item")(); -} - //! With a shape (KeyedShape) the cell converts with it as its target, which a type holding a MAP needs: without it a //! dict reads as a STRUCT named by its keys' text, and a key whose text does not cast back, such as bytes, is lost. //! An object without a DuckDB type reads as its text into a stated VARCHAR, and has no kind for an implied one. @@ -338,8 +325,19 @@ static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb if (cell.ptr() == Py_None) { return Value(); } + // A TIMESTAMP_NS column keeps a cell's nanoseconds, and an implied naive timestamp type sees them, so its kind + // check refuses what a coarser unit would drop; a stated coarser type reads the cell at microseconds, as a frame's + // own datetime64 column reads + auto naive_timestamp = + target.id() == LogicalTypeId::TIMESTAMP_NS || + (implied && (target.id() == LogicalTypeId::TIMESTAMP || target.id() == LogicalTypeId::TIMESTAMP_MS || + target.id() == LogicalTypeId::TIMESTAMP_SEC)); + timestamp_ns_t nanoseconds; + if (naive_timestamp && HoldsNanoseconds(cell, nanoseconds)) { + return Value::TIMESTAMPNS(nanoseconds); + } if (IsNumpyScalar(cell)) { - auto value = NumpyScalarValue(cell); + auto value = NumpyScalarItem(cell); // A long double's item() is a numpy scalar again; one step, then the conversion refuses what is left if (!IsNumpyScalar(value)) { return DeclaredValue(context, target, value, shape, implied); @@ -351,24 +349,41 @@ static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb return TransformPythonValue(context, cell, shape ? *shape : LogicalType(LogicalTypeId::UNKNOWN), true); } -//! The type a cell of a type holding a MAP converts with: its containers and no other type, so a dict builds its MAP -//! from its keys and values in the types they have in Python, and every key and value meets the engine's cast with -//! the connection's settings, as a TIMESTAMPTZ in the session time zone -static LogicalType KeyedShape(const LogicalType &type) { +//! The type a cell of a nested type converts with: its containers and its naive timestamps and no other type, so a +//! dict builds its MAP from its keys and values in the types they have in Python and every key and value meets the +//! engine's cast with the connection's settings, as a TIMESTAMPTZ in the session time zone, and a timestamp reads in +//! the unit of its place. With typed_entries a MAP keeps a key or value type that +//! has no children, so its entries are each cast into it rather than into their common type first; a type the caller +//! stated takes that, while an implied type keeps the entries' own types for the kind check. +static LogicalType KeyedShape(const LogicalType &type, bool typed_entries) { + auto entry = [&](const LogicalType &entry_type) { + auto plain = !entry_type.IsNested() && entry_type.id() != LogicalTypeId::UNION; + return typed_entries && plain ? entry_type : KeyedShape(entry_type, typed_entries); + }; switch (type.id()) { case LogicalTypeId::STRUCT: { child_list_t children; for (idx_t i = 0; i < StructType::GetChildCount(type); i++) { - children.emplace_back(StructType::GetChildName(type, i), KeyedShape(StructType::GetChildType(type, i))); + children.emplace_back(StructType::GetChildName(type, i), + KeyedShape(StructType::GetChildType(type, i), typed_entries)); } - return LogicalType::STRUCT(std::move(children)); + return AsShape(LogicalType::STRUCT(std::move(children))); } case LogicalTypeId::LIST: - return LogicalType::LIST(KeyedShape(ListType::GetChildType(type))); + return AsShape(LogicalType::LIST(KeyedShape(ListType::GetChildType(type), typed_entries))); case LogicalTypeId::ARRAY: - return LogicalType::ARRAY(KeyedShape(ArrayType::GetChildType(type)), ArrayType::GetSize(type)); + return AsShape( + LogicalType::ARRAY(KeyedShape(ArrayType::GetChildType(type), typed_entries), ArrayType::GetSize(type))); case LogicalTypeId::MAP: - return LogicalType::MAP(KeyedShape(MapType::KeyType(type)), KeyedShape(MapType::ValueType(type))); + return AsShape(LogicalType::MAP(entry(MapType::KeyType(type)), entry(MapType::ValueType(type)))); + case LogicalTypeId::TIMESTAMP_NS: + return AsShape(type); + case LogicalTypeId::TIMESTAMP: + case LogicalTypeId::TIMESTAMP_MS: + case LogicalTypeId::TIMESTAMP_SEC: + // An implied type sees the nanoseconds a coarser unit would drop, so its kind check refuses them; a stated one + // reads them at microseconds, as at the top + return AsShape(typed_entries ? type : LogicalType::TIMESTAMP_NS); default: return LogicalType::UNKNOWN; } @@ -483,16 +498,7 @@ class ImpliedKinds { if (source == type) { return true; } - for (auto &entry : seen) { - if (entry.source == source) { - return entry.verdict == Verdict::FITS || - (entry.verdict == Verdict::BY_VALUE && FitsByValue(value, type)); - } - } - auto verdict = Judge(source, type); - if (seen.size() < CACHED_SOURCE_TYPES) { - seen.push_back({source, verdict}); - } + auto verdict = JudgeOnce(source, type); return verdict == Verdict::FITS || (verdict == Verdict::BY_VALUE && FitsByValue(value, type)); } @@ -565,8 +571,8 @@ class ImpliedKinds { return Verdict::REFUSED; } claimed[field.GetIndex()] = true; - verdict = Combine(verdict, Judge(StructType::GetChildType(source, i), - StructType::GetChildType(type, field.GetIndex()))); + verdict = Combine(verdict, JudgeOnce(StructType::GetChildType(source, i), + StructType::GetChildType(type, field.GetIndex()))); } return verdict; } @@ -575,22 +581,22 @@ class ImpliedKinds { if (source.id() != LogicalTypeId::MAP) { return Verdict::REFUSED; } - return Combine(Judge(MapType::KeyType(source), MapType::KeyType(type)), - Judge(MapType::ValueType(source), MapType::ValueType(type))); + return Combine(JudgeOnce(MapType::KeyType(source), MapType::KeyType(type)), + JudgeOnce(MapType::ValueType(source), MapType::ValueType(type))); case LogicalTypeId::LIST: if (source.id() != LogicalTypeId::LIST) { return Verdict::REFUSED; } - return Judge(ListType::GetChildType(source), ListType::GetChildType(type)); + return JudgeOnce(ListType::GetChildType(source), ListType::GetChildType(type)); case LogicalTypeId::ARRAY: // A Python list reads as a LIST, which the cast takes into an ARRAY of its length or refuses if (source.id() == LogicalTypeId::LIST) { - return Judge(ListType::GetChildType(source), ArrayType::GetChildType(type)); + return JudgeOnce(ListType::GetChildType(source), ArrayType::GetChildType(type)); } if (source.id() != LogicalTypeId::ARRAY) { return Verdict::REFUSED; } - return Judge(ArrayType::GetChildType(source), ArrayType::GetChildType(type)); + return JudgeOnce(ArrayType::GetChildType(source), ArrayType::GetChildType(type)); default: break; } @@ -640,12 +646,64 @@ class ImpliedKinds { return CastFunctionSet::ImplicitCastCost(context, source, type) >= 0 ? Verdict::FITS : Verdict::REFUSED; } + //! Judge, looked up once per pair of types for the scan of a vector: the elements of a list share a type, and so do + //! the cells of a column, so a pair is judged once rather than per cell or element + Verdict JudgeOnce(const LogicalType &source, const LogicalType &type) { + for (auto &entry : judged) { + if (entry.source == source && entry.type == type) { + return entry.verdict; + } + } + auto verdict = Judge(source, type); + if (judged.size() < CACHED_TYPE_PAIRS) { + judged.push_back({source, type, verdict}); + } + return verdict; + } + + static bool HoldsInteger(int64_t value, LogicalTypeId id) { + switch (id) { + case LogicalTypeId::TINYINT: + return value >= NumericLimits::Minimum() && value <= NumericLimits::Maximum(); + case LogicalTypeId::SMALLINT: + return value >= NumericLimits::Minimum() && value <= NumericLimits::Maximum(); + case LogicalTypeId::INTEGER: + return value >= NumericLimits::Minimum() && value <= NumericLimits::Maximum(); + case LogicalTypeId::UTINYINT: + return value >= 0 && value <= NumericLimits::Maximum(); + case LogicalTypeId::USMALLINT: + return value >= 0 && value <= NumericLimits::Maximum(); + case LogicalTypeId::UINTEGER: + return value >= 0 && value <= NumericLimits::Maximum(); + case LogicalTypeId::UBIGINT: + case LogicalTypeId::UHUGEINT: + return value >= 0; + default: + return true; + } + } + + static bool HoldsInt64(LogicalTypeId id) { + switch (id) { + case LogicalTypeId::TINYINT: + case LogicalTypeId::SMALLINT: + case LogicalTypeId::INTEGER: + case LogicalTypeId::BIGINT: + case LogicalTypeId::UTINYINT: + case LogicalTypeId::USMALLINT: + case LogicalTypeId::UINTEGER: + return true; + default: + return false; + } + } + bool FitsByValue(const Value &value, const LogicalType &type) { if (value.IsNull()) { return true; } auto &source = value.type(); - switch (Judge(source, type)) { + switch (JudgeOnce(source, type)) { case Verdict::FITS: return true; case Verdict::REFUSED: @@ -686,6 +744,10 @@ class ImpliedKinds { return true; } default: { + // A Python int converts to INTEGER or BIGINT, whose range says whether a narrower type holds it + if (HoldsInt64(source.id()) && type.IsIntegral()) { + return HoldsInteger(value.GetValue(), type.id()); + } if ((source.id() == LogicalTypeId::DOUBLE || source.id() == LogicalTypeId::FLOAT) && !FloatHoldsItsText(value, type)) { return false; @@ -759,9 +821,8 @@ class ImpliedKinds { } private: - //! The significant digits of a number's shortest text ("2.0" has one, "1e+20" twenty one): from its first digit - //! that is not zero to its last, or to the units when its integer part ends in zeros - static idx_t SignificantDigits(const string &text) { + //! The digits of the integer part of a number's shortest text: none for "0.1", one for "2.0", 21 for "1e+20" + static idx_t IntegerDigits(const string &text) { string digits; int64_t integer_digits = 0; bool point = false; @@ -776,32 +837,28 @@ class ImpliedKinds { } int64_t exponent = i < text.size() ? std::stoll(text.substr(i + 1)) : 0; auto units = integer_digits + exponent; - auto end = static_cast(digits.size()); - while (end > MaxValue(units, 0) && digits[NumericCast(end - 1)] == '0') { - end--; - } - int64_t start = 0; - while (start < end && digits[NumericCast(start)] == '0') { - start++; - } - if (start == end) { - return 0; + int64_t leading_zeros = 0; + while (leading_zeros < units && leading_zeros < static_cast(digits.size()) && + digits[NumericCast(leading_zeros)] == '0') { + leading_zeros++; } - return NumericCast(MaxValue(end, units) - start); + return NumericCast(MaxValue(units - leading_zeros, 0)); } - //! A double holds every decimal of at most 15 significant digits exactly enough to give it back as its shortest - //! text, so a float whose text has more may be a rounding already: df() hands a HUGEINT or a DECIMAL over as a - //! float64. And the engine's cast of a float into a DECIMAL rounds the value times 10**scale in double precision, - //! which for a wide scale is no longer that value (2.0 reads 2.000000000000000039769249677312 in a DECIMAL(38,30), - //! exactly in a DECIMAL(38,20)), so the float fits only when the cast gives the value of its text. + //! A float stands for the value it was made from only when that value had at most 15 significant digits, the + //! decimals a double holds exactly: df() hands a HUGEINT or a DECIMAL over as a float64, already rounded when + //! the value had more. A float with i integer digits stands for a DECIMAL(p, s) value of at most i + s digits, so + //! it fits only when i + s is at most 15 (s is 0 for an integer type). And the engine's cast of a float into a + //! DECIMAL rounds the value times 10**scale in double precision, so it must also give the value of the float's + //! shortest text. static bool FloatHoldsItsText(const Value &value, const LogicalType &type) { auto number = value.GetValue(); if (!Value::DoubleIsFinite(number)) { return false; } auto text = Value::DOUBLE(number).ToString(); - if (SignificantDigits(text) > MAX_EXACT_FLOAT_DIGITS) { + auto scale = type.id() == LogicalTypeId::DECIMAL ? DecimalType::GetScale(type) : 0; + if (IntegerDigits(text) + scale > MAX_EXACT_FLOAT_DIGITS) { return false; } if (type.id() != LogicalTypeId::DECIMAL) { @@ -812,17 +869,18 @@ class ImpliedKinds { return cast && exact && Value::NotDistinctFrom(*cast, *exact); } - struct Seen { + struct Judged { LogicalType source; + LogicalType type; Verdict verdict; }; //! DBL_DIG: every decimal of this many significant digits survives a round trip through a double static constexpr idx_t MAX_EXACT_FLOAT_DIGITS = 15; //! Bounded like DeclaredCasts, since dicts give a STRUCT type per key set - static constexpr idx_t CACHED_SOURCE_TYPES = 16; + static constexpr idx_t CACHED_TYPE_PAIRS = 32; ClientContext &context; - vector seen; + vector judged; }; // A column without a mask takes the instantiation that tests no mask bits, so it reads like the plain array @@ -867,14 +925,22 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t bool bit_storage_refused) { DeclaredVector staged(count); auto &target = out.GetType(); - const auto keyed = TypeVisitor::Contains(target, LogicalTypeId::MAP); + // A type holding a MAP takes a shape for its keys, and one holding a naive timestamp for the unit a nested value + // reads in; a top level timestamp is read by DeclaredValue itself + const auto keyed = target.IsNested() && (TypeVisitor::Contains(target, LogicalTypeId::MAP) || + TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_NS) || + TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP) || + TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_MS) || + TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_SEC)); const auto holds_bit = TypeVisitor::Contains(target, LogicalTypeId::BIT); - const auto shape = keyed ? KeyedShape(target) : LogicalType(LogicalTypeId::UNKNOWN); - // A str into VARCHAR without an alias, and bytes or a bytearray into BLOB without one, read as TransformPythonValue - // would read them, without building a Value per cell; a collation governs comparison only, an alias such as JSON - // validates the text in its cast + const auto shape = keyed ? KeyedShape(target, !implied_type) : LogicalType(LogicalTypeId::UNKNOWN); + // A str into VARCHAR without an alias, bytes or a bytearray into BLOB without one, and a time without a zone into + // TIME read as TransformPythonValue would read them, without building a Value per cell; a collation governs + // comparison only, an alias such as JSON validates the text in its cast const bool plain_text = target == LogicalType::VARCHAR; const bool plain_blob = target == LogicalType::BLOB; + const bool plain_time = target == LogicalType::TIME; + PyObject *time_type = nullptr; ImpliedKinds kinds(context); ScanObjectCells(context, col, stride, masked, count, offset, out, [&](nb::handle cell, idx_t row) { if (plain_text && PyUnicode_CheckExact(cell.ptr())) { @@ -890,14 +956,39 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t } // df() hands a BLOB cell to the function as a bytearray if (plain_blob && (PyBytes_CheckExact(cell.ptr()) || PyByteArray_CheckExact(cell.ptr()))) { - auto bytes = PyBytes_CheckExact(cell.ptr()); - auto data = bytes ? PyBytes_AS_STRING(cell.ptr()) : PyByteArray_AS_STRING(cell.ptr()); - auto size = bytes ? PyBytes_GET_SIZE(cell.ptr()) : PyByteArray_GET_SIZE(cell.ptr()); + char *data; + Py_ssize_t size; + if (PyBytes_CheckExact(cell.ptr())) { + if (PyBytes_AsStringAndSize(cell.ptr(), &data, &size) < 0) { + throw nb::python_error(); + } + } else { + data = PyByteArray_AsString(cell.ptr()); + size = PyByteArray_Size(cell.ptr()); + if (!data || size < 0) { + throw nb::python_error(); + } + } FlatVector::GetDataMutable(out)[row] = StringVector::AddStringOrBlob(out, data, NumericCast(size)); staged.MarkWritten(row); return; } + if (plain_time) { + if (!time_type) { + time_type = DuckDBPyConnection::ImportCache()->datetime.time().ptr(); + } + if (Py_TYPE(cell.ptr()) == reinterpret_cast(time_type)) { + PyTime time(cell); + // A time with a zone converts as a TIMETZ, which an implied TIME refuses and the cast into a stated one + // reads without its offset + if (time.timezone_obj.is_none()) { + FlatVector::GetDataMutable(out)[row] = time.ToDuckTime(); + staged.MarkWritten(row); + return; + } + } + } auto value = DeclaredValue(context, target, cell, keyed ? &shape : nullptr, implied_type); // A duckdb.Value states its type, as a schema does, so a cast into the column's type is meant if (implied_type && !kinds.Fits(value, target) && diff --git a/tests/fast/pandas/test_append_df.py b/tests/fast/pandas/test_append_df.py index 66007b6b..b5b40fed 100644 --- a/tests/fast/pandas/test_append_df.py +++ b/tests/fast/pandas/test_append_df.py @@ -483,16 +483,18 @@ def test_map_key_meets_the_cast_in_the_session_settings(self, duckdb_cursor, tar @pytest.mark.parametrize( ("target", "key", "expected"), [ - # The engine's cast reads the instant in the session's time zone; the base dropped the offset + # The engine's cast reads the instant in the session's time zone; 2.0.0 dropped the offset pytest.param( "TIMESTAMP", datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc), "{'2019-12-31 19:00:00'=1}", id="aware_datetime_into_timestamp", ), - # The cast from BLOB gives the bytes' text; the base gave the text of the Python bytes object + # The cast from BLOB gives the bytes' text; 2.0.0 gave the text of the Python bytes object pytest.param("VARCHAR", b"x", "{x=1}", id="bytes_into_varchar"), - pytest.param("STRUCT(a INTEGER)", (1,), "{{'a': 1}=1}", id="tuple_into_struct"), + pytest.param("VARCHAR", True, "{true=1}", id="bool_into_varchar"), + pytest.param("VARCHAR", datetime.timedelta(days=1), "{1 day=1}", id="timedelta_into_varchar"), + pytest.param("VARCHAR", (1, 2), "{'[1, 2]'=1}", id="tuple_into_varchar"), ], ) def test_map_key_reads_as_the_cast_from_its_python_type_gives_it(self, duckdb_cursor, target, key, expected): @@ -501,12 +503,36 @@ def test_map_key_reads_as_the_cast_from_its_python_type_gives_it(self, duckdb_cu duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([{key: 1}], dtype=object)})) assert duckdb_cursor.sql("SELECT a::VARCHAR FROM t").fetchall() == [(expected,)] + def test_map_key_given_as_a_tuple_for_a_struct_key_type_reads_as_a_list_and_is_refused(self, duckdb_cursor): + # A tuple is a LIST at every depth, as the appender reads a value of a row; 2.0.0 read this key through its text + duckdb_cursor.execute("CREATE TABLE t (a MAP(STRUCT(a INTEGER), INTEGER))") + frame = pd.DataFrame({"a": pd.Series([{(1,): 1}], dtype=object)}) + expected = "Unimplemented type for cast (INTEGER[] -> STRUCT(a INTEGER))" + with pytest.raises(duckdb.ConversionException, match=re.escape(expected)): + duckdb_cursor.append("t", frame) + assert duckdb_cursor.sql("SELECT 42").fetchall() == [(42,)] + def test_map_keys_a_str_and_its_bytes_are_one_key_and_refused(self, duckdb_cursor): duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, INTEGER))") frame = pd.DataFrame({"a": pd.Series([{"a": 1, b"a": 2}], dtype=object)}) with pytest.raises(duckdb.InvalidInputException, match=re.escape("Map keys must be unique")): duckdb_cursor.append("t", frame) + def test_text_of_a_map_into_a_map_column_reads_as_the_cast_reads_it(self): + con = duckdb.connect() + con.execute("CREATE TABLE t (a MAP(VARCHAR, INTEGER))") + con.append("t", pd.DataFrame({"a": pd.Series(["{a=1}"], dtype=object)})) + assert con.table("t").fetchall() == [({"a": 1},)] + assert con.sql("SELECT 42").fetchall() == [(42,)] + + def test_map_values_of_mixed_numbers_each_keep_their_value(self): + # Converted toward their common type first, 2**60 + 1 beside 1.0 would read as a rounded DOUBLE + con = duckdb.connect() + con.execute("CREATE TABLE t (a MAP(VARCHAR, BIGINT))") + con.append("t", pd.DataFrame({"a": pd.Series([{"a": 2**60 + 1, "b": 1.0}], dtype=object)})) + assert con.table("t").fetchall() == [({"a": 2**60 + 1, "b": 1},)] + assert con.sql("SELECT 42").fetchall() == [(42,)] + def test_map_key_of_none_is_refused_and_the_connection_stays_usable(self, duckdb_cursor): duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, INTEGER))") with pytest.raises(duckdb.Error) as error: diff --git a/tests/fast/test_map_types.py b/tests/fast/test_map_types.py index 59c45892..0a54fd85 100644 --- a/tests/fast/test_map_types.py +++ b/tests/fast/test_map_types.py @@ -346,9 +346,9 @@ def _apply_to_a(convert): ), pytest.param( "{'t': TIMESTAMP '2020-01-01'}", - _apply_to_a(lambda d: {"t": np.datetime64("2020-01-01T00:00:00", "ns")}), + _apply_to_a(lambda d: {"t": np.datetime64("2020-01-01T00:00:00.000000001", "ns")}), "STRUCT(t TIMESTAMP)", - "STRUCT(t BIGINT)", + "STRUCT(t TIMESTAMP_NS)", id="nanosecond_datetime64_field", ), # A float of more than 15 significant digits, or one the cast into a wide scale rounds in double precision @@ -498,6 +498,8 @@ def test_map_time_value_outside_the_range_of_its_columns_unit_is_refused(duckdb_ pytest.param("[123456789012345678901::HUGEINT]", "HUGEINT[]", id="hugeint_beyond_bigint"), pytest.param("[1152921504606846977::HUGEINT, NULL]", "HUGEINT[]", id="hugeint_with_a_null"), pytest.param("[1234567890123456.78::DECIMAL(18,2)]", "DECIMAL(18,2)[]", id="decimal_of_18_digits"), + pytest.param("[1.000000000000000001::DECIMAL(19,18)]", "DECIMAL(19,18)[]", id="decimal_of_scale_18"), + pytest.param("[100000000000000.001::DECIMAL(18,3)]", "DECIMAL(18,3)[]", id="decimal_of_15_integer_digits"), ], ) def test_map_identity_of_a_list_whose_floats_may_have_rounded_is_refused(duckdb_cursor, values, sql_type): @@ -512,6 +514,10 @@ def test_map_identity_of_a_list_whose_floats_may_have_rounded_is_refused(duckdb_ [ pytest.param("[5::HUGEINT, NULL, (-999999999999999)::HUGEINT]", id="hugeint_of_15_digits"), pytest.param("[1.25::DECIMAL(9,2), NULL, 0.1::DECIMAL(9,2)]", id="decimal"), + pytest.param("[1.5::DECIMAL(18,3), NULL]", id="decimal_of_scale_3"), + pytest.param("[(-32768)::SMALLINT, NULL, 32767::SMALLINT]", id="smallint"), + pytest.param("[(-128)::TINYINT, 127::TINYINT]", id="tinyint"), + pytest.param("{'x': 32767::SMALLINT}", id="smallint_field"), ], ) def test_map_identity_of_a_list_whose_floats_hold_their_values_round_trips(duckdb_cursor, values): @@ -521,6 +527,240 @@ def test_map_identity_of_a_list_whose_floats_hold_their_values_round_trips(duckd assert mapped.fetchall() == source.fetchall() +@pytest.mark.parametrize( + ("source", "values", "returned"), + [ + pytest.param("[1::SMALLINT]", [40000], "INTEGER[]", id="smallint_list"), + pytest.param("[1::TINYINT]", [-129], "INTEGER[]", id="tinyint_list"), + ], +) +def test_map_narrow_integer_list_with_an_int_beyond_its_range_is_refused(duckdb_cursor, source, values, returned): + mapped = duckdb_cursor.sql(f"select {source} as a").map(_apply_to_a(lambda v: values)) + sql_type = duckdb_cursor.sql(f"select {source} as a").types[0] + expected = f"Invalid Input Error: UDF column type mismatch, expected [{sql_type}], got [{returned}]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +@pytest.mark.parametrize("cell", [0, "x"], ids=["int", "str"]) +def test_map_scalar_returned_for_a_map_column_is_refused_and_the_database_stays_usable(cell): + con = duckdb.connect() + mapped = con.sql("select map(['k'], [1]) as a").map( + lambda df: df.assign(a=pd.Series([cell] * len(df), dtype=object, index=df.index)) + ) + with pytest.raises(duckdb.InvalidInputException, match=re.escape("UDF column type mismatch")): + mapped.fetchall() + assert con.sql("select 42").fetchall() == [(42,)] + + +def test_map_text_of_a_map_for_a_stated_map_column_reads_as_the_cast_reads_it(): + con = duckdb.connect() + mapped = con.sql("select map(['k'], [1]) as a").map( + lambda df: df.assign(a=pd.Series(["{k=2}"] * len(df), dtype=object, index=df.index)), + schema={"a": duckdb.sqltype("MAP(VARCHAR, INTEGER)")}, + ) + assert mapped.fetchall() == [({"k": 2},)] + assert con.sql("select 42").fetchall() == [(42,)] + + +@pytest.mark.parametrize( + ("value", "nanoseconds"), + [ + pytest.param("'2020-01-01 00:00:00.000000001'::TIMESTAMP_NS", 1, id="pandas_timestamp"), + pytest.param("NULL::TIMESTAMP_NS", None, id="null"), + ], +) +def test_map_timestamp_ns_column_handed_back_as_objects_keeps_its_nanoseconds(duckdb_cursor, value, nanoseconds): + source = duckdb_cursor.sql(f"select {value} as a") + mapped = source.map(lambda df: df.assign(a=df["a"].astype(object))) + assert mapped.types == [duckdb.sqltype("TIMESTAMP_NS")] + assert mapped.project("a::VARCHAR").fetchall() == source.project("a::VARCHAR").fetchall() + + +def test_map_timestamp_ns_column_of_numpy_datetime64_cells_keeps_their_nanoseconds(duckdb_cursor): + cell = np.datetime64("2020-01-01T00:00:00.000000002", "ns") + mapped = duckdb_cursor.sql("select '2020-01-01'::TIMESTAMP_NS as a").map( + lambda df: df.assign(a=pd.Series([cell] * len(df), dtype=object, index=df.index)) + ) + assert mapped.project("a::VARCHAR").fetchall() == [("2020-01-01 00:00:00.000000002",)] + + +def test_map_timestamp_column_given_a_timestamp_with_nanoseconds_is_refused(duckdb_cursor): + # The cast into microseconds would drop the nanosecond + stamp = pd.Timestamp("2020-01-01 00:00:00.000000001") + mapped = duckdb_cursor.sql("select TIMESTAMP '2020-01-01' as a").map( + lambda df: df.assign(a=pd.Series([stamp] * len(df), dtype=object, index=df.index)) + ) + expected = "Invalid Input Error: UDF column type mismatch, expected [TIMESTAMP], got [TIMESTAMP_NS]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +# A scalar inside a container of a type holding a MAP, which meets the engine's cast or the kind check, never a +# conversion target holding no type +NESTED_SCALARS = pytest.mark.parametrize( + ("sql_type", "cell", "cast_into"), + [ + pytest.param("MAP(VARCHAR, INTEGER)[]", [0], "MAP(VARCHAR, INTEGER)", id="int_in_a_list_of_maps"), + pytest.param("MAP(VARCHAR, INTEGER)[]", ["x"], "MAP(VARCHAR, INTEGER)", id="str_in_a_list_of_maps"), + pytest.param("STRUCT(m MAP(VARCHAR, INTEGER))", {"m": 5}, "MAP(VARCHAR, INTEGER)", id="int_for_a_map_field"), + pytest.param("MAP(VARCHAR, INTEGER)[]", np.array([0, 1]), "MAP(VARCHAR, INTEGER)", id="int_array_for_a_list"), + pytest.param("MAP(VARCHAR, INTEGER)[]", (0,), "MAP(VARCHAR, INTEGER)", id="tuple_for_a_list"), + pytest.param("MAP(VARCHAR, STRUCT(x INTEGER))", {"k": "abc"}, "STRUCT(x INTEGER)", id="str_for_a_struct_value"), + ], +) + + +def _returning(cell): + def replace(df): + return df.assign(a=pd.Series([cell] * len(df), dtype=object, index=df.index)) + + return replace + + +@NESTED_SCALARS +@pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append"]) +def test_map_scalar_inside_a_container_of_a_type_holding_a_map_is_refused_and_the_database_stays_usable( + sql_type, cell, cast_into, mode +): + con = duckdb.connect() + con.execute(f"create table t (a {sql_type})") + con.execute("insert into t values (NULL)") + if mode == "append": + run = lambda: con.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) # noqa: E731 + else: + schema = {"a": duckdb.sqltype(sql_type)} if mode == "explicit_schema" else None + run = con.table("t").map(_returning(cell), schema=schema).fetchall + with pytest.raises(duckdb.Error) as error: + run() + message = str(error.value) + assert not isinstance(error.value, duckdb.FatalException), message + if mode == "inferred": + assert message.startswith(f"Invalid Input Error: UDF column type mismatch, expected [{sql_type}], got ["), ( + message + ) + else: + # The message names the column's own types, never a conversion target + assert cast_into in message, message + assert "shape" not in message, message + assert con.sql("select 42").fetchall() == [(42,)] + + +def test_map_cast_error_for_a_nested_scalar_names_the_columns_type(): + con = duckdb.connect() + con.execute("create table t (a MAP(VARCHAR, INTEGER)[])") + expected = "Conversion Error: Unimplemented type for cast (INTEGER -> MAP(VARCHAR, INTEGER))" + with pytest.raises(duckdb.ConversionException, match=re.escape(expected)): + con.append("t", pd.DataFrame({"a": pd.Series([[0]], dtype=object)})) + assert con.sql("select 42").fetchall() == [(42,)] + + +NANOSECOND = pd.Timestamp("2020-01-01 00:00:00.000000001") + + +@pytest.mark.parametrize( + ("value", "cell"), + [ + pytest.param("{'t': '2020-01-01'::TIMESTAMP_NS}", {"t": NANOSECOND}, id="struct_field"), + pytest.param("['2020-01-01'::TIMESTAMP_NS]", [NANOSECOND], id="list_element"), + ], +) +def test_map_nested_timestamp_ns_value_keeps_its_nanoseconds(duckdb_cursor, value, cell): + source = duckdb_cursor.sql(f"select {value} as a") + mapped = source.map(_returning(cell)) + assert mapped.types == source.types + assert "00:00:00.000000001" in mapped.project("a::VARCHAR").fetchall()[0][0] + + +def test_map_nested_timestamp_with_nanoseconds_for_an_implied_timestamp_is_refused(duckdb_cursor): + mapped = duckdb_cursor.sql("select {'t': TIMESTAMP '2020-01-01'} as a").map(_returning({"t": NANOSECOND})) + expected = ( + "Invalid Input Error: UDF column type mismatch, expected [STRUCT(t TIMESTAMP)], got [STRUCT(t TIMESTAMP_NS)]" + ) + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +FIVE_MICROSECONDS = np.datetime64("2020-01-01T00:00:00.000005123", "ns") + + +@pytest.mark.parametrize( + ("sql_type", "cell", "expected"), + [ + pytest.param( + "STRUCT(t TIMESTAMP)", + {"t": FIVE_MICROSECONDS}, + "{'t': '2020-01-01 00:00:00.000005'}", + id="datetime64_in_a_struct", + ), + pytest.param( + "STRUCT(t TIMESTAMP)", + {"t": pd.Timestamp(FIVE_MICROSECONDS)}, + "{'t': '2020-01-01 00:00:00.000005'}", + id="pandas_timestamp_in_a_struct", + ), + pytest.param("TIMESTAMP[]", [FIVE_MICROSECONDS], "['2020-01-01 00:00:00.000005']", id="datetime64_in_a_list"), + pytest.param( + "MAP(VARCHAR, TIMESTAMP)", + {"k": FIVE_MICROSECONDS}, + "{k='2020-01-01 00:00:00.000005'}", + id="datetime64_as_a_map_value", + ), + pytest.param( + "TIMESTAMP_MS[]", + [np.datetime64("2020-01-01T00:00:00.005123", "ns")], + "['2020-01-01 00:00:00.005']", + id="datetime64_in_a_list_of_milliseconds", + ), + ], +) +@pytest.mark.parametrize("mode", ["explicit_schema", "append"]) +def test_map_nested_timestamp_with_nanoseconds_for_a_stated_coarser_unit_reads_at_microseconds_as_at_the_top( + sql_type, cell, expected, mode +): + con = duckdb.connect() + con.execute(f"create table t (a {sql_type})") + if mode == "append": + con.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) + rows = con.sql("select a::VARCHAR from t").fetchall() + else: + con.execute("insert into t values (NULL)") + schema = {"a": duckdb.sqltype(sql_type)} + rows = con.table("t").map(_returning(cell), schema=schema).project("a::VARCHAR").fetchall() + assert rows == [(expected,)] + + +@pytest.mark.parametrize( + "sql_type", + ["STRUCT(a INTEGER, b INTEGER)", "STRUCT(a INTEGER, t TIMESTAMP)", "STRUCT(a INTEGER, m MAP(VARCHAR, INTEGER))"], + ids=["plain_fields", "timestamp_field", "map_field"], +) +@pytest.mark.parametrize("depth", ["top", "struct_field", "list_element"]) +@pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append"]) +def test_map_tuple_reads_as_a_list_at_every_depth_so_a_struct_refuses_it(sql_type, depth, mode): + column_type, cell = { + "top": (sql_type, (3, 4)), + "struct_field": (f"STRUCT(s {sql_type})", {"s": (3, 4)}), + "list_element": (f"{sql_type}[]", [(3, 4)]), + }[depth] + con = duckdb.connect() + con.execute(f"create table t (a {column_type})") + con.execute("insert into t values (NULL)") + if mode == "append": + run = lambda: con.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) # noqa: E731 + else: + schema = {"a": duckdb.sqltype(column_type)} if mode == "explicit_schema" else None + run = con.table("t").map(_returning(cell), schema=schema).fetchall + if mode == "inferred": + expected = f"UDF column type mismatch, expected [{column_type}], got [" + with pytest.raises(duckdb.InvalidInputException, match=re.escape(expected)): + run() + else: + with pytest.raises(duckdb.ConversionException, match=re.escape(f"(INTEGER[] -> {sql_type})")): + run() + assert con.sql("select 42").fetchall() == [(42,)] + + def test_map_decimal_column_handed_back_as_objects_of_18_digits_is_refused(duckdb_cursor): # df() hands a DECIMAL column over as float64, which rounds 1234567890123456.78 source = duckdb_cursor.sql("select 1234567890123456.78::DECIMAL(18,2) as a") @@ -576,7 +816,7 @@ def test_map_time_value_its_column_unit_holds_reads_unchanged(duckdb_cursor, sou ], ) def test_map_container_of_mixed_elements_reads_as_the_conversion_types_it(duckdb_cursor, source, function, expected): - # The conversion types a list or a MAP by the common type of its elements; the base refused these + # The conversion types a list or a MAP by the common type of its elements; 2.0.0 refused these relation = duckdb_cursor.sql(f"select {source} as a from range(1) t(i)") mapped = relation.map(function) assert mapped.types == relation.types @@ -644,6 +884,12 @@ def test_map_column_whose_cells_changed_in_kind_casts_into_the_schema_type(duckd ), pytest.param("{'x': i::BIGINT}", _apply_to_a(lambda d: {"x": Decimal(5)}), {"x": 5}, id="integral_decimal"), pytest.param("{'x': i::BIGINT}", _apply_to_a(lambda d: {"x": 2.0}), {"x": 2}, id="integral_float_for_bigint"), + pytest.param( + "{'t': TIMESTAMP '2020-01-01'}", + _apply_to_a(lambda d: {"t": np.datetime64("2020-01-02T00:00:00", "ns")}), + {"t": datetime.datetime(2020, 1, 2)}, + id="nanosecond_datetime64_of_whole_microseconds", + ), pytest.param( "{'x': i::DECIMAL(18,2)}", _apply_to_a(lambda d: {"x": 1234567890123.45}), @@ -698,7 +944,7 @@ def test_map_text_column_mixing_str_with_other_cells_reads_each_as_text_in_the_s @SCHEMAS -def test_map_text_list_mixing_str_with_an_int_reads_as_text_as_on_the_base(duckdb_cursor, explicit_schema): +def test_map_text_list_mixing_str_with_an_int_reads_as_text_as_2_0_0_reads_it(duckdb_cursor, explicit_schema): # A list converts by the common type of its elements, so the int reads as its text inside a text list source = duckdb_cursor.sql("select ['x', 'y'] as a from range(1) t(i)") mapped = source.map( @@ -779,6 +1025,29 @@ def test_map_bit_value_reaching_the_function_as_its_storage_bytes_is_refused( mapped.fetchall() +@SCHEMAS +def test_map_time_column_of_times_with_and_without_a_zone(duckdb_cursor, explicit_schema): + source = duckdb_cursor.sql( + "select (case when i = 1 then NULL else TIME '01:02:03.000004' end) as a from range(3) t(i)" + ) + schema = _schema_of(source, explicit_schema) + mapped = source.map(lambda df: df, schema=schema) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + zone = datetime.timezone(datetime.timedelta(hours=2)) + zoned = source.map( + lambda df: df.assign(a=pd.Series([datetime.time(1, tzinfo=zone)] * len(df), dtype=object, index=df.index)), + schema=schema, + ) + if explicit_schema: + # The cast of a TIMETZ into TIME keeps the wall time and drops the offset + assert zoned.fetchall() == [(datetime.time(1),)] * 3 + else: + expected = "Invalid Input Error: UDF column type mismatch, expected [TIME], got [TIME WITH TIME ZONE]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + zoned.fetchall() + + def test_map_json_column_still_validates_its_text(duckdb_cursor): schema = {"a": duckdb.sqltype("JSON")} good = duckdb_cursor.sql("select 1 as a").map(_text_cells(['{"k": 1}']), schema=schema) @@ -1003,6 +1272,33 @@ def late_name(df): mapped.fetchall() +def test_map_column_filled_from_a_null_typed_input_names_the_column_and_schema(duckdb_cursor): + def fill(df): + return df.assign(c=pd.Series(["x"] * len(df), dtype=object, index=df.index)) + + mapped = duckdb_cursor.sql("select i as a, NULL as c from range(3) t(i)").map(fill) + with pytest.raises(duckdb.InvalidInputException, match=_no_type("c")): + mapped.fetchall() + + +def test_map_column_a_first_map_added_as_none_and_a_second_fills_names_the_column_and_schema(duckdb_cursor): + def add_none(df): + return df.assign(c=pd.Series([None] * len(df), dtype=object, index=df.index)) + + def fill(df): + return df.assign(c=pd.Series([1] * len(df), dtype=object, index=df.index)) + + mapped = duckdb_cursor.sql("select i as a from range(3) t(i)").map(add_none).map(fill) + with pytest.raises(duckdb.InvalidInputException, match=_no_type("c")): + mapped.fetchall() + + +def test_map_identity_of_a_null_typed_input_column_reads_null(duckdb_cursor): + mapped = duckdb_cursor.sql("select i as a, NULL as c from range(3) t(i)").map(lambda df: df) + assert mapped.types == [duckdb.sqltypes.BIGINT, duckdb.sqltypes.SQLNULL] + assert mapped.fetchall() == [(0, None), (1, None), (2, None)] + + def test_map_new_object_column_fails_on_the_first_chunk_with_a_value(duckdb_cursor): duckdb_cursor.execute("set threads = 1") rows_per_call = [] From 6746e39887f93cbe8b75151cc040b20e8708cb6f Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Fri, 6 Nov 2026 22:23:05 +0000 Subject: [PATCH 4/9] Give the key and value list form of a MAP the entry rules, count a Timestamp as nanoseconds only when it has them, pick the declared cell reader once per vector --- .../duckdb_python/python_conversion.hpp | 6 +- src/native/python_conversion.cpp | 179 ++++++++++++------ src/numpy/numpy_scan.cpp | 133 +++++++++---- tests/fast/pandas/test_append_df.py | 28 +++ tests/fast/test_map_types.py | 134 +++++++++++++ 5 files changed, 382 insertions(+), 98 deletions(-) diff --git a/src/include/duckdb_python/python_conversion.hpp b/src/include/duckdb_python/python_conversion.hpp index 68ffce1f..1337fa92 100644 --- a/src/include/duckdb_python/python_conversion.hpp +++ b/src/include/duckdb_python/python_conversion.hpp @@ -58,8 +58,10 @@ bool DictionaryHasMapFormat(const PyDictionary &dict); //! children with the connection's settings; a shaped naive timestamp reads a cell as TransformTimestamp does LogicalType AsShape(const LogicalType &type); bool IsShape(const LogicalType &type); -//! Whether a cell is a naive pandas Timestamp or a numpy datetime64 finer than a microsecond, whose nanoseconds no -//! Python datetime holds; they are then in nanoseconds +//! The nanosecond a datetime subclass such as a pandas Timestamp gives as an int, or -1 when it gives none +int64_t NanosecondOf(nb::handle cell); +//! Whether a cell is a naive pandas Timestamp with nanoseconds or a numpy datetime64 finer than a microsecond, whose +//! nanoseconds no Python datetime holds; they are then in nanoseconds bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds); //! A numpy scalar's Python value; a datetime64 or timedelta64 at nanoseconds would read as an integer, so those take //! the microsecond unit first, as a frame's own datetime columns do diff --git a/src/native/python_conversion.cpp b/src/native/python_conversion.cpp index a03d3839..74baab01 100644 --- a/src/native/python_conversion.cpp +++ b/src/native/python_conversion.cpp @@ -111,6 +111,17 @@ static bool IsValidMapComponent(const nb::handle &component) { return true; } +//! The length of a MAP's key or value component, or -1 for one that has none although it claims one, as a numpy +//! array of no dimensions does +static int64_t ComponentLength(nb::handle component) { + auto length = PyObject_Length(component.ptr()); + if (length < 0) { + PyErr_Clear(); + return -1; + } + return length; +} + bool DictionaryHasMapFormat(const PyDictionary &dict) { if (dict.len != 2) { return false; @@ -137,12 +148,11 @@ bool DictionaryHasMapFormat(const PyDictionary &dict) { return true; } - // Verify that both the keys and values are of the same length - auto size = nb::len(keys); - if (size != nb::len(values)) { - return false; - } - return true; + // Verify that both the keys and values are of the same length; a component without one is refused by the + // conversion + auto key_size = ComponentLength(keys); + auto value_size = ComponentLength(values); + return key_size < 0 || value_size < 0 || key_size == value_size; } Value TransformDictionaryToStruct(optional_ptr context, const PyDictionary &dict, @@ -224,12 +234,37 @@ bool IsShape(const LogicalType &type) { return type.HasAlias() && type.GetAlias() == SHAPE_ALIAS; } +int64_t NanosecondOf(nb::handle cell) { + // Interned once, since every cell of a timestamp column asks + static PyObject *const name = PyUnicode_InternFromString("nanosecond"); + if (!name) { + PyErr_Clear(); + return -1; + } + auto nanosecond = nb::steal(PyObject_GetAttr(cell.ptr(), name)); + if (!nanosecond.is_valid()) { + PyErr_Clear(); + return -1; + } + // NaT's nanosecond is NaN + return PyLong_Check(nanosecond.ptr()) ? PyLong_AsLongLong(nanosecond.ptr()) : -1; +} + bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds) { auto &import_cache = *DuckDBPyConnection::ImportCache(); auto datetime_type = import_cache.datetime.datetime(); if (Py_TYPE(cell.ptr()) == reinterpret_cast(datetime_type.ptr())) { return false; } + if (duckdb::PyUtil::IsInstance(cell, datetime_type)) { + // A Timestamp without nanoseconds reads as the datetime it is: its value in nanoseconds overflows outside the + // years 1677 to 2262, and building it costs more than the datetime's own conversion + if (NanosecondOf(cell) <= 0 || !nb::getattr(cell, "tzinfo").is_none()) { + return false; + } + nanoseconds = timestamp_ns_t(nb::cast(cell.attr("value"))); + return true; + } if (duckdb::PyUtil::IsInstance(cell, import_cache.numpy.datetime64())) { auto code = nb::cast(cell.attr("dtype").attr("str")); if (!StringUtil::EndsWith(code, "ns]") && !StringUtil::EndsWith(code, "ps]") && @@ -243,13 +278,7 @@ bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds) { nanoseconds = timestamp_ns_t(nb::cast(value)); return true; } - if (!duckdb::PyUtil::IsInstance(cell, datetime_type) || !PandasIsImported(import_cache) || - cell.is(import_cache.pandas.NaT()) || !nb::hasattr(cell, "nanosecond") || - !nb::getattr(cell, "tzinfo").is_none()) { - return false; - } - nanoseconds = timestamp_ns_t(nb::cast(cell.attr("value"))); - return true; + return false; } nb::object NumpyScalarItem(nb::handle scalar) { @@ -286,9 +315,19 @@ Value TransformTimestamp(optional_ptr context, nb::handle ele, co return TransformPythonValue(context, ele, LogicalType::UNKNOWN, nan_as_null); } -//! Whether an object is what a shape is for: a dict for a MAP or a STRUCT, a list, a tuple or a numpy array for a -//! LIST or an ARRAY; any other object converts without a target and meets the engine's cast or the kind check, so a -//! tuple reads as a LIST at every depth, as the appender reads a value of a row +//! Whether an object is a list a shape takes: a list, a tuple or a numpy array of at least one dimension, since an +//! array of none holds a single value +static bool ShapeTakesAsList(nb::handle ele) { + if (PyList_Check(ele.ptr()) || PyTuple_Check(ele.ptr())) { + return true; + } + return duckdb::PyUtil::IsInstance(ele, DuckDBPyConnection::ImportCache()->numpy.ndarray()) && + nb::cast(ele.attr("ndim")) > 0; +} + +//! Whether an object is what a shape is for: a dict for a MAP or a STRUCT, a list for a LIST or an ARRAY; any other +//! object converts without a target and meets the engine's cast or the kind check, so a tuple reads as a LIST at +//! every depth, as the appender reads a value of a row static bool ShapeTakes(nb::handle ele, const LogicalType &shape) { switch (shape.id()) { case LogicalTypeId::MAP: @@ -296,8 +335,7 @@ static bool ShapeTakes(nb::handle ele, const LogicalType &shape) { return PyDict_Check(ele.ptr()); case LogicalTypeId::LIST: case LogicalTypeId::ARRAY: - return PyList_Check(ele.ptr()) || PyTuple_Check(ele.ptr()) || - duckdb::PyUtil::IsInstance(ele, DuckDBPyConnection::ImportCache()->numpy.ndarray()); + return ShapeTakesAsList(ele); default: return false; } @@ -324,23 +362,10 @@ static Value TransformMapEntry(optional_ptr context, nb::handle e return std::move(*cast); } -Value TransformStructFormatDictionaryToMap(optional_ptr context, const PyDictionary &dict, - const LogicalType &target_type, bool nan_as_null) { - if (dict.len == 0) { - return EmptyMapValue(); - } - - if (target_type.id() != LogicalTypeId::MAP) { - throw InvalidInputException("Please provide a valid target type for transform from Python to Value"); - } - - if (nb::none().is(dict.keys) || nb::none().is(dict.values)) { - return Value(LogicalType::MAP(LogicalTypeId::SQLNULL, LogicalTypeId::SQLNULL)); - } - - auto size = nb::len(dict.keys); - D_ASSERT(size == nb::len(dict.values)); - +//! A MAP of the entries of a sequence of keys and one of values, both of the given size; under a shape each key and +//! value converts as TransformMapEntry reads it +static Value MapFromEntries(optional_ptr context, nb::handle keys, nb::handle values, idx_t size, + const LogicalType &target_type, bool nan_as_null) { auto key_target = MapType::KeyType(target_type); auto value_target = MapType::ValueType(target_type); auto shape = IsShape(target_type); @@ -351,8 +376,8 @@ Value TransformStructFormatDictionaryToMap(optional_ptr context, vector elements; for (idx_t i = 0; i < size; i++) { - auto key = dict.keys.attr("__getitem__")(i); - auto value = dict.values.attr("__getitem__")(i); + auto key = keys.attr("__getitem__")(i); + auto value = values.attr("__getitem__")(i); Value new_key = shape ? TransformMapEntry(context, key, key_target, nan_as_null) : TransformPythonValue(context, key, key_target, nan_as_null); Value new_value = shape ? TransformMapEntry(context, value, value_target, nan_as_null) @@ -382,6 +407,25 @@ Value TransformStructFormatDictionaryToMap(optional_ptr context, return Value::MAP(ListType::GetChildType(map_type), std::move(elements)); } +Value TransformStructFormatDictionaryToMap(optional_ptr context, const PyDictionary &dict, + const LogicalType &target_type, bool nan_as_null) { + if (dict.len == 0) { + return EmptyMapValue(); + } + + if (target_type.id() != LogicalTypeId::MAP) { + throw InvalidInputException("Please provide a valid target type for transform from Python to Value"); + } + + if (nb::none().is(dict.keys) || nb::none().is(dict.values)) { + return Value(LogicalType::MAP(LogicalTypeId::SQLNULL, LogicalTypeId::SQLNULL)); + } + + auto size = nb::len(dict.keys); + D_ASSERT(size == nb::len(dict.values)); + return MapFromEntries(context, dict.keys, dict.values, size, target_type, nan_as_null); +} + Value TransformDictionaryToMap(optional_ptr context, const PyDictionary &dict, const LogicalType &target_type, bool nan_as_null) { if (target_type.id() != LogicalTypeId::UNKNOWN && !DictionaryHasMapFormat(dict)) { @@ -397,9 +441,9 @@ Value TransformDictionaryToMap(optional_ptr context, const PyDict return Value(LogicalType::MAP(LogicalTypeId::SQLNULL, LogicalTypeId::SQLNULL)); } - auto key_size = nb::len(keys); - D_ASSERT(key_size == nb::len(values)); - if (key_size == 0) { + auto key_size = ComponentLength(keys); + auto sized = key_size >= 0 && ComponentLength(values) >= 0; + if (sized && key_size == 0) { // dict == { 'key': [], 'value': [] } return EmptyMapValue(); } @@ -408,19 +452,35 @@ Value TransformDictionaryToMap(optional_ptr context, const PyDict LogicalType key_target = LogicalTypeId::UNKNOWN; LogicalType value_target = LogicalTypeId::UNKNOWN; - if (target_type.id() != LogicalTypeId::UNKNOWN) { + if (IsShape(target_type)) { + // A shape reads the two lists entry by entry, as it reads a dict's own entries; a component that is no list + // converts without a target and is refused below + if (sized && ShapeTakesAsList(keys) && ShapeTakesAsList(values)) { + auto items = [](nb::handle list) { + return PyList_Check(list.ptr()) || PyTuple_Check(list.ptr()) ? nb::borrow(list) + : list.attr("tolist")(); + }; + return MapFromEntries(context, items(keys), items(values), key_size, target_type, nan_as_null); + } + } else if (target_type.id() != LogicalTypeId::UNKNOWN) { key_target = LogicalType::LIST(MapType::KeyType(target_type)); value_target = LogicalType::LIST(MapType::ValueType(target_type)); } auto key_list = TransformPythonValue(context, keys, key_target, nan_as_null); auto value_list = TransformPythonValue(context, values, value_target, nan_as_null); + if (key_list.type().id() != LogicalTypeId::LIST || value_list.type().id() != LogicalTypeId::LIST || + ListValue::GetChildren(key_list).size() != ListValue::GetChildren(value_list).size()) { + throw InvalidInputException("A dict of 'key' and 'value' for a MAP takes a list of keys and a list of values, " + "not %s and %s", + key_list.type().ToString(), value_list.type().ToString()); + } LogicalType key_type = LogicalType::SQLNULL; LogicalType value_type = LogicalType::SQLNULL; vector elements; - for (idx_t i = 0; i < key_size; i++) { + for (idx_t i = 0; i < ListValue::GetChildren(key_list).size(); i++) { Value new_key = ListValue::GetChildren(key_list)[i]; Value new_value = ListValue::GetChildren(value_list)[i]; @@ -1338,21 +1398,32 @@ void TransformPythonObject(optional_ptr context, nb::handle ele, TransformPythonObjectInternal(context, ele, vector, result_offset, nan_as_null); } +//! A shape of no children marks a timestamp, read as TransformTimestamp reads it; a container shape takes only the +//! objects it is for, so a scalar inside a container meets the engine's cast as at the top, never a type holding +//! UNKNOWN +static Value TransformShapedValue(optional_ptr context, nb::handle ele, const LogicalType &shape, + bool nan_as_null) { + if (!shape.IsNested()) { + return TransformTimestamp(context, ele, shape, nan_as_null); + } + if (!ShapeTakes(ele, shape)) { + return TransformPythonValue(context, ele, LogicalType::UNKNOWN, nan_as_null); + } + Value result; + TransformPythonObjectInternal(context, ele, result, shape, nan_as_null); + return result; +} + Value TransformPythonValue(optional_ptr context, nb::handle ele, const LogicalType &target_type, bool nan_as_null) { - // A shape of no children marks a timestamp, read as TransformTimestamp reads it; a container shape takes only the - // objects it is for, so a scalar inside a container meets the engine's cast as at the top, never a type holding - // UNKNOWN - if (IsShape(target_type)) { - if (!target_type.IsNested()) { - return TransformTimestamp(context, ele, target_type, nan_as_null); - } - if (!ShapeTakes(ele, target_type)) { - return TransformPythonValue(context, ele, LogicalType::UNKNOWN, nan_as_null); - } - } + // One named result on every path, so it is built in place rather than moved out for every value converted; a + // conversion without a target, the most common, asks for no alias Value result; - TransformPythonObjectInternal(context, ele, result, target_type, nan_as_null); + if (target_type.id() != LogicalTypeId::UNKNOWN && IsShape(target_type)) { + result = TransformShapedValue(context, ele, target_type, nan_as_null); + } else { + TransformPythonObjectInternal(context, ele, result, target_type, nan_as_null); + } return result; } diff --git a/src/numpy/numpy_scan.cpp b/src/numpy/numpy_scan.cpp index ece998c8..2d82fde9 100644 --- a/src/numpy/numpy_scan.cpp +++ b/src/numpy/numpy_scan.cpp @@ -317,11 +317,13 @@ static bool IsNumpyScalar(nb::handle cell) { return cell.type().is(import_cache.numpy.ndarray()) && nb::cast(cell.attr("ndim")) == 0; } -//! With a shape (KeyedShape) the cell converts with it as its target, which a type holding a MAP needs: without it a -//! dict reads as a STRUCT named by its keys' text, and a key whose text does not cast back, such as bytes, is lost. +//! The cell converts with the column type's shape (KeyedShape) as its target, or with UNKNOWN for a type that takes +//! none; a type holding a MAP needs one: without it a dict reads as a STRUCT named by its keys' text, and a key whose +//! text does not cast back, such as bytes, is lost. Both are built once per vector, since building a type looks up +//! its type information each time. //! An object without a DuckDB type reads as its text into a stated VARCHAR, and has no kind for an implied one. -static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb::handle cell, - optional_ptr shape, bool implied) { +static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb::handle cell, const LogicalType &shape, + bool implied) { if (cell.ptr() == Py_None) { return Value(); } @@ -346,7 +348,7 @@ static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb if (!implied && target.id() == LogicalTypeId::VARCHAR && GetPythonObjectType(cell) == PythonObjectType::Other) { return Value(TransformPythonString(nb::str(cell))); } - return TransformPythonValue(context, cell, shape ? *shape : LogicalType(LogicalTypeId::UNKNOWN), true); + return TransformPythonValue(context, cell, shape, true); } //! The type a cell of a nested type converts with: its containers and its naive timestamps and no other type, so a @@ -404,7 +406,7 @@ class DeclaredVector { written_count++; } - void Set(idx_t row, Value value) { + void Set(idx_t row, Value &&value) { if (value.IsNull()) { return; } @@ -420,6 +422,17 @@ class DeclaredVector { if (written_count == count) { return; } + // A vector the scan wrote no row of skips the test for one, which would cost every row of the plain path + if (written_count == 0) { + WriteStaged(context, out, count); + } else { + WriteStaged(context, out, count); + } + } + +private: + template + void WriteStaged(ClientContext &context, Vector &out, idx_t count) { auto &target = out.GetType(); if (uniform && source.id() != LogicalTypeId::SQLNULL && source != target) { idx_t present = 0; @@ -441,7 +454,7 @@ class DeclaredVector { VectorOperations::TryCast(context, staged, casted, present, nullptr); idx_t copied = 0; for (idx_t row = 0; row < count;) { - if (written[row]) { + if (SKIP_WRITTEN && written[row]) { row++; continue; } @@ -462,18 +475,21 @@ class DeclaredVector { } DeclaredCasts casts(context, target); for (idx_t row = 0; row < count; row++) { - if (written[row]) { + if (SKIP_WRITTEN && written[row]) { continue; } - if (values[row].IsNull()) { + auto &value = values[row]; + if (value.IsNull()) { FlatVector::SetNull(out, row, true); + } else if (value.type() == target) { + // Written as it is, without the copy of the Value that a cast of it would return + out.SetValue(row, value); } else { - out.SetValue(row, casts.Cast(values[row])); + out.SetValue(row, casts.Cast(value)); } } } -private: vector values; vector written; idx_t written_count = 0; @@ -932,18 +948,34 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP) || TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_MS) || TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_SEC)); - const auto holds_bit = TypeVisitor::Contains(target, LogicalTypeId::BIT); + const auto bits_refused = bit_storage_refused && TypeVisitor::Contains(target, LogicalTypeId::BIT); const auto shape = keyed ? KeyedShape(target, !implied_type) : LogicalType(LogicalTypeId::UNKNOWN); - // A str into VARCHAR without an alias, bytes or a bytearray into BLOB without one, and a time without a zone into - // TIME read as TransformPythonValue would read them, without building a Value per cell; a collation governs - // comparison only, an alias such as JSON validates the text in its cast - const bool plain_text = target == LogicalType::VARCHAR; - const bool plain_blob = target == LogicalType::BLOB; - const bool plain_time = target == LogicalType::TIME; - PyObject *time_type = nullptr; ImpliedKinds kinds(context); - ScanObjectCells(context, col, stride, masked, count, offset, out, [&](nb::handle cell, idx_t row) { - if (plain_text && PyUnicode_CheckExact(cell.ptr())) { + auto read = [&](nb::handle cell, idx_t row) { + // The row already holds the NULL a None reads as + if (cell.ptr() == Py_None) { + return; + } + auto value = DeclaredValue(context, target, cell, shape, implied_type); + // A duckdb.Value states its type, as a schema does, so a cast into the column's type is meant + if (implied_type && !kinds.Fits(value, target) && + !duckdb::PyUtil::IsInstance(cell, DuckDBPyConnection::ImportCache()->duckdb.Value())) { + throw ImpliedTypeRefused(value.type()); + } + if (bits_refused && ImpliedKinds::HoldsBitStorage(value, target)) { + throw ImpliedTypeRefused(value.type()); + } + staged.Set(row, std::move(value)); + }; + // A str into VARCHAR without an alias, bytes or a bytearray into BLOB without one, and a time or a datetime without + // a zone into TIME or TIMESTAMP read as TransformPythonValue would read them, without building a Value per cell; a + // collation governs comparison only, an alias such as JSON validates the text in its cast. The reader is chosen + // once per vector, since a test per cell for every fast path costs the plain path more than the tests themselves. + if (target == LogicalType::VARCHAR) { + ScanObjectCells(context, col, stride, masked, count, offset, out, [&](nb::handle cell, idx_t row) { + if (!PyUnicode_CheckExact(cell.ptr())) { + return read(cell, row); + } Py_ssize_t size; auto text = PyUnicode_AsUTF8AndSize(cell.ptr(), &size); if (!text) { @@ -952,33 +984,36 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t FlatVector::GetDataMutable(out)[row] = StringVector::AddString(out, text, NumericCast(size)); staged.MarkWritten(row); - return; - } + }); + } else if (target == LogicalType::BLOB) { // df() hands a BLOB cell to the function as a bytearray - if (plain_blob && (PyBytes_CheckExact(cell.ptr()) || PyByteArray_CheckExact(cell.ptr()))) { + ScanObjectCells(context, col, stride, masked, count, offset, out, [&](nb::handle cell, idx_t row) { char *data; Py_ssize_t size; if (PyBytes_CheckExact(cell.ptr())) { if (PyBytes_AsStringAndSize(cell.ptr(), &data, &size) < 0) { throw nb::python_error(); } - } else { + } else if (PyByteArray_CheckExact(cell.ptr())) { data = PyByteArray_AsString(cell.ptr()); size = PyByteArray_Size(cell.ptr()); if (!data || size < 0) { throw nb::python_error(); } + } else { + return read(cell, row); } FlatVector::GetDataMutable(out)[row] = StringVector::AddStringOrBlob(out, data, NumericCast(size)); staged.MarkWritten(row); - return; - } - if (plain_time) { + }); + } else if (target == LogicalType::TIME) { + PyTypeObject *time_type = nullptr; + ScanObjectCells(context, col, stride, masked, count, offset, out, [&](nb::handle cell, idx_t row) { if (!time_type) { - time_type = DuckDBPyConnection::ImportCache()->datetime.time().ptr(); + time_type = reinterpret_cast(DuckDBPyConnection::ImportCache()->datetime.time().ptr()); } - if (Py_TYPE(cell.ptr()) == reinterpret_cast(time_type)) { + if (Py_TYPE(cell.ptr()) == time_type) { PyTime time(cell); // A time with a zone converts as a TIMETZ, which an implied TIME refuses and the cast into a stated one // reads without its offset @@ -988,18 +1023,32 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t return; } } - } - auto value = DeclaredValue(context, target, cell, keyed ? &shape : nullptr, implied_type); - // A duckdb.Value states its type, as a schema does, so a cast into the column's type is meant - if (implied_type && !kinds.Fits(value, target) && - !duckdb::PyUtil::IsInstance(cell, DuckDBPyConnection::ImportCache()->duckdb.Value())) { - throw ImpliedTypeRefused(value.type()); - } - if (bit_storage_refused && holds_bit && ImpliedKinds::HoldsBitStorage(value, target)) { - throw ImpliedTypeRefused(value.type()); - } - staged.Set(row, std::move(value)); - }); + read(cell, row); + }); + } else if (target == LogicalType::TIMESTAMP) { + PyTypeObject *datetime_type = nullptr; + ScanObjectCells(context, col, stride, masked, count, offset, out, [&](nb::handle cell, idx_t row) { + if (!datetime_type) { + datetime_type = + reinterpret_cast(DuckDBPyConnection::ImportCache()->datetime.datetime().ptr()); + } + // A datetime, or a pandas Timestamp without nanoseconds, holds what the conversion reads + auto type = Py_TYPE(cell.ptr()); + if (type == datetime_type || (PyType_IsSubtype(type, datetime_type) && NanosecondOf(cell) == 0)) { + PyDateTime datetime(cell); + // A datetime with a zone converts as a TIMESTAMPTZ, which an implied TIMESTAMP refuses and the cast + // into a stated one reads in the session's time zone + if (datetime.tzone_obj.is_none()) { + FlatVector::GetDataMutable(out)[row] = datetime.ToTimestamp(); + staged.MarkWritten(row); + return; + } + } + read(cell, row); + }); + } else { + ScanObjectCells(context, col, stride, masked, count, offset, out, read); + } staged.Write(context, out, count); VerifyTypeConstraints(out, count); } diff --git a/tests/fast/pandas/test_append_df.py b/tests/fast/pandas/test_append_df.py index b5b40fed..93f81ca7 100644 --- a/tests/fast/pandas/test_append_df.py +++ b/tests/fast/pandas/test_append_df.py @@ -512,6 +512,34 @@ def test_map_key_given_as_a_tuple_for_a_struct_key_type_reads_as_a_list_and_is_r duckdb_cursor.append("t", frame) assert duckdb_cursor.sql("SELECT 42").fetchall() == [(42,)] + @pytest.mark.parametrize( + ("target", "value", "expected"), + [ + pytest.param( + "TIMESTAMPTZ", + datetime.datetime(2024, 5, 6), + "{k='2024-05-06 00:00:00-04'}", + id="naive_datetime_into_timestamptz", + ), + pytest.param("VARCHAR", b"x", "{k=x}", id="bytes_into_varchar"), + pytest.param("DECIMAL(9,2)", 1.5, "{k=1.50}", id="float_into_decimal"), + pytest.param( + "TIMESTAMP_NS", + datetime.datetime(2024, 5, 6), + "{k='2024-05-06 00:00:00'}", + id="datetime_into_timestamp_ns", + ), + ], + ) + def test_map_of_key_and_value_lists_reads_each_entry_as_the_same_dict_does( + self, duckdb_cursor, target, value, expected + ): + duckdb_cursor.execute("SET TimeZone = 'America/New_York'") + duckdb_cursor.execute(f"CREATE TABLE t (a MAP(VARCHAR, {target}))") + frame = pd.DataFrame({"a": pd.Series([{"key": ["k"], "value": [value]}, {"k": value}], dtype=object)}) + duckdb_cursor.append("t", frame) + assert duckdb_cursor.sql("SELECT a::VARCHAR FROM t").fetchall() == [(expected,), (expected,)] + def test_map_keys_a_str_and_its_bytes_are_one_key_and_refused(self, duckdb_cursor): duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, INTEGER))") frame = pd.DataFrame({"a": pd.Series([{"a": 1, b"a": 2}], dtype=object)}) diff --git a/tests/fast/test_map_types.py b/tests/fast/test_map_types.py index 0a54fd85..d541635c 100644 --- a/tests/fast/test_map_types.py +++ b/tests/fast/test_map_types.py @@ -596,6 +596,71 @@ def test_map_timestamp_column_given_a_timestamp_with_nanoseconds_is_refused(duck mapped.fetchall() +def _far_from_1970(value): + return pytest.param(value, id=value[:4]) + + +FAR_YEARS = pytest.mark.parametrize( + "value", [_far_from_1970("1500-01-01"), _far_from_1970("2500-01-01"), _far_from_1970("9999-12-31")] +) + + +@FAR_YEARS +@pytest.mark.parametrize("sql_type", ["TIMESTAMP", "TIMESTAMP_S"]) +def test_map_timestamp_column_handed_back_as_timestamps_outside_the_nanosecond_range_round_trips( + duckdb_cursor, sql_type, value +): + source = duckdb_cursor.sql(f"select '{value}'::{sql_type} as a") + mapped = source.map(lambda df: df.assign(a=df["a"].astype(object))) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + + +@FAR_YEARS +def test_map_struct_field_given_a_timestamp_outside_the_nanosecond_range_reads_it(duckdb_cursor, value): + source = duckdb_cursor.sql(f"select {{'t': '{value}'::TIMESTAMP}} as a") + mapped = source.map(_returning({"t": pd.Timestamp(value)})) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + + +def test_map_timestamp_column_of_2020_and_9999_handed_back_as_timestamps_round_trips(duckdb_cursor): + source = duckdb_cursor.sql( + "select (case when i = 0 then '2020-01-01' else '9999-12-31' end)::TIMESTAMP as a from range(2) t(i)" + ) + mapped = source.map(lambda df: df.assign(a=df["a"].astype(object))) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + + +@SCHEMAS +def test_map_timestamp_ns_column_given_a_timestamp_beyond_its_range_is_refused(duckdb_cursor, explicit_schema): + source = duckdb_cursor.sql("select TIMESTAMP_NS '2020-01-01' as a") + mapped = source.map(_returning(pd.Timestamp("9999-12-31")), schema=_schema_of(source, explicit_schema)) + if explicit_schema: + expected = "UDF column 'a': Could not convert Timestamp to higher precision" + with pytest.raises(duckdb.ConversionException, match=re.escape(expected)): + mapped.fetchall() + else: + expected = "Invalid Input Error: UDF column type mismatch, expected [TIMESTAMP_NS], got [TIMESTAMP]" + with pytest.raises(duckdb.InvalidInputException, match=f"^{re.escape(expected)}$"): + mapped.fetchall() + + +@pytest.mark.parametrize( + "value", + ["infinity", "-infinity", "2262-04-11 23:47:16.854775806"], + ids=["infinity", "minus_infinity", "largest_finite_nanosecond_but_one"], +) +@SCHEMAS +def test_map_timestamp_ns_identity_handed_back_as_timestamps_keeps_infinity(duckdb_cursor, value, explicit_schema): + # df() hands an infinity over as pandas' Timestamp.max or Timestamp.min, which read back as that infinity + source = duckdb_cursor.sql(f"select '{value}'::TIMESTAMP_NS as a") + mapped = source.map(lambda df: df.assign(a=df["a"].astype(object)), schema=_schema_of(source, explicit_schema)) + assert mapped.types == source.types + assert mapped.project("a::VARCHAR").fetchall() == [(value,)] + + # A scalar inside a container of a type holding a MAP, which meets the engine's cast or the kind check, never a # conversion target holding no type NESTED_SCALARS = pytest.mark.parametrize( @@ -607,6 +672,12 @@ def test_map_timestamp_column_given_a_timestamp_with_nanoseconds_is_refused(duck pytest.param("MAP(VARCHAR, INTEGER)[]", np.array([0, 1]), "MAP(VARCHAR, INTEGER)", id="int_array_for_a_list"), pytest.param("MAP(VARCHAR, INTEGER)[]", (0,), "MAP(VARCHAR, INTEGER)", id="tuple_for_a_list"), pytest.param("MAP(VARCHAR, STRUCT(x INTEGER))", {"k": "abc"}, "STRUCT(x INTEGER)", id="str_for_a_struct_value"), + pytest.param( + "MAP(VARCHAR, INTEGER)[][]", + [np.array(5)], + "MAP(VARCHAR, INTEGER)[]", + id="zero_dimensional_array_for_a_list", + ), ], ) @@ -646,6 +717,69 @@ def test_map_scalar_inside_a_container_of_a_type_holding_a_map_is_refused_and_th assert con.sql("select 42").fetchall() == [(42,)] +@pytest.mark.parametrize( + ("sql_type", "cell", "converted"), + [ + pytest.param("MAP(VARCHAR, VARCHAR)", {"key": "ab", "value": "cd"}, "VARCHAR and VARCHAR", id="map"), + pytest.param( + "STRUCT(m MAP(VARCHAR, VARCHAR))", + {"m": {"key": "ab", "value": "cd"}}, + "VARCHAR and VARCHAR", + id="map_field", + ), + pytest.param( + "MAP(VARCHAR, INTEGER)", + {"key": np.array("a"), "value": np.array(1)}, + "VARCHAR and INTEGER", + id="zero_dimensional_arrays", + ), + ], +) +@pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append", "appender", "parameter"]) +def test_map_of_a_key_and_a_value_that_are_no_lists_is_refused_and_the_database_stays_usable( + sql_type, cell, converted, mode +): + con = duckdb.connect() + con.execute(f"create table t (a {sql_type})") + if mode == "append": + run = lambda: con.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) # noqa: E731 + elif mode == "appender": + + def run(): + with con.appender("t") as appender: + appender.append_row(cell) + + elif mode == "parameter": + run = lambda: con.execute("insert into t values (?)", [cell]) # noqa: E731 + else: + con.execute("insert into t values (NULL)") + schema = {"a": duckdb.sqltype(sql_type)} if mode == "explicit_schema" else None + run = con.table("t").map(_returning(cell), schema=schema).fetchall + expected = f"A dict of 'key' and 'value' for a MAP takes a list of keys and a list of values, not {converted}" + with pytest.raises(duckdb.InvalidInputException, match=re.escape(expected)): + run() + assert con.sql("select 42").fetchall() == [(42,)] + + +@pytest.mark.parametrize( + ("cell", "expected"), + [ + pytest.param({"key": ["k"], "value": [5]}, "{k=5}", id="list_of_an_int"), + pytest.param({"key": ("k",), "value": (5,)}, "{k=5}", id="tuple_of_an_int"), + pytest.param({"key": np.array(["k"]), "value": np.array([5])}, "{k=5}", id="numpy_arrays"), + pytest.param({"key": ["a", "b"], "value": [2**60 + 1, 1]}, "{a=1152921504606846977, b=1}", id="wide_ints"), + ], +) +@SCHEMAS +def test_map_of_key_and_value_lists_reads_each_entry_as_a_dicts_entry_reads( + duckdb_cursor, cell, expected, explicit_schema +): + source = duckdb_cursor.sql("select MAP {'k': 1::BIGINT} as a") + mapped = source.map(_returning(cell), schema=_schema_of(source, explicit_schema)) + assert mapped.types == source.types + assert mapped.project("a::VARCHAR").fetchall() == [(expected,)] + + def test_map_cast_error_for_a_nested_scalar_names_the_columns_type(): con = duckdb.connect() con.execute("create table t (a MAP(VARCHAR, INTEGER)[])") From 66a35e7ea22f0b628a9faae077264f42191327ea Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Fri, 6 Nov 2026 23:48:36 +0000 Subject: [PATCH 5/9] Read an aware datetime once in the TIMESTAMP fast path and detect the key and value form of a MAP by its components --- src/native/python_conversion.cpp | 35 ++++++------- src/numpy/numpy_scan.cpp | 19 ++++--- tests/fast/pandas/test_append_df.py | 6 +++ tests/fast/test_map_types.py | 77 +++++++++++++++++++++++++++-- 4 files changed, 104 insertions(+), 33 deletions(-) diff --git a/src/native/python_conversion.cpp b/src/native/python_conversion.cpp index 74baab01..1e5f4ad6 100644 --- a/src/native/python_conversion.cpp +++ b/src/native/python_conversion.cpp @@ -97,18 +97,11 @@ static vector TransformStructKeys(nb::handle keys, idx_t size) { return res; } -static bool IsValidMapComponent(const nb::handle &component) { - // The component is either NULL - if (nb::none().is(component)) { - return true; - } - if (!nb::hasattr(component, "__getitem__")) { - return false; - } - if (!nb::hasattr(component, "__len__")) { - return false; - } - return true; +//! Whether a MAP's key or value component is a column of keys or values: a list, a tuple or a numpy array, which the +//! conversion reads as a LIST (an array of no dimensions it refuses) +static bool IsMapComponentList(nb::handle component) { + return PyList_Check(component.ptr()) || PyTuple_Check(component.ptr()) || + duckdb::PyUtil::IsInstance(component, DuckDBPyConnection::ImportCache()->numpy.ndarray()); } //! The length of a MAP's key or value component, or -1 for one that has none although it claims one, as a numpy @@ -136,15 +129,17 @@ bool DictionaryHasMapFormat(const PyDictionary &dict) { return false; } - if (!IsValidMapComponent(keys)) { - return false; - } - if (!IsValidMapComponent(values)) { + // A list on either side, or None on both, makes the dict a MAP of the two lists: None a NULL MAP, and the + // conversion refuses the other side when it is no list. Without a list it is a plain dict, a MAP of the entries + // 'key' and 'value', so a NULL entry beside a plain value stays an entry + auto key_none = nb::none().is(keys); + auto value_none = nb::none().is(values); + auto key_list = IsMapComponentList(keys); + auto value_list = IsMapComponentList(values); + if (!key_list && !value_list && !(key_none && value_none)) { return false; } - - // If either of the components is NULL, return early - if (nb::none().is(keys) || nb::none().is(values)) { + if (key_none || value_none || !key_list || !value_list) { return true; } @@ -258,7 +253,7 @@ bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds) { } if (duckdb::PyUtil::IsInstance(cell, datetime_type)) { // A Timestamp without nanoseconds reads as the datetime it is: its value in nanoseconds overflows outside the - // years 1677 to 2262, and building it costs more than the datetime's own conversion + // years 1677 to 2262, and a coarser column would judge that value by a cast where the datetime fits outright if (NanosecondOf(cell) <= 0 || !nb::getattr(cell, "tzinfo").is_none()) { return false; } diff --git a/src/numpy/numpy_scan.cpp b/src/numpy/numpy_scan.cpp index 2d82fde9..4e724b13 100644 --- a/src/numpy/numpy_scan.cpp +++ b/src/numpy/numpy_scan.cpp @@ -951,12 +951,7 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t const auto bits_refused = bit_storage_refused && TypeVisitor::Contains(target, LogicalTypeId::BIT); const auto shape = keyed ? KeyedShape(target, !implied_type) : LogicalType(LogicalTypeId::UNKNOWN); ImpliedKinds kinds(context); - auto read = [&](nb::handle cell, idx_t row) { - // The row already holds the NULL a None reads as - if (cell.ptr() == Py_None) { - return; - } - auto value = DeclaredValue(context, target, cell, shape, implied_type); + auto accept = [&](nb::handle cell, idx_t row, Value &&value) { // A duckdb.Value states its type, as a schema does, so a cast into the column's type is meant if (implied_type && !kinds.Fits(value, target) && !duckdb::PyUtil::IsInstance(cell, DuckDBPyConnection::ImportCache()->duckdb.Value())) { @@ -967,6 +962,13 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t } staged.Set(row, std::move(value)); }; + auto read = [&](nb::handle cell, idx_t row) { + // The row already holds the NULL a None reads as + if (cell.ptr() == Py_None) { + return; + } + accept(cell, row, DeclaredValue(context, target, cell, shape, implied_type)); + }; // A str into VARCHAR without an alias, bytes or a bytearray into BLOB without one, and a time or a datetime without // a zone into TIME or TIMESTAMP read as TransformPythonValue would read them, without building a Value per cell; a // collation governs comparison only, an alias such as JSON validates the text in its cast. The reader is chosen @@ -1036,13 +1038,14 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t auto type = Py_TYPE(cell.ptr()); if (type == datetime_type || (PyType_IsSubtype(type, datetime_type) && NanosecondOf(cell) == 0)) { PyDateTime datetime(cell); - // A datetime with a zone converts as a TIMESTAMPTZ, which an implied TIMESTAMP refuses and the cast - // into a stated one reads in the session's time zone if (datetime.tzone_obj.is_none()) { FlatVector::GetDataMutable(out)[row] = datetime.ToTimestamp(); staged.MarkWritten(row); return; } + // A datetime with a zone converts as a TIMESTAMPTZ, as the conversion reads the one built here, which + // an implied TIMESTAMP refuses and the cast into a stated one reads in the session's time zone + return accept(cell, row, datetime.ToDuckValue(LogicalType::UNKNOWN)); } read(cell, row); }); diff --git a/tests/fast/pandas/test_append_df.py b/tests/fast/pandas/test_append_df.py index 93f81ca7..1083016f 100644 --- a/tests/fast/pandas/test_append_df.py +++ b/tests/fast/pandas/test_append_df.py @@ -540,6 +540,12 @@ def test_map_of_key_and_value_lists_reads_each_entry_as_the_same_dict_does( duckdb_cursor.append("t", frame) assert duckdb_cursor.sql("SELECT a::VARCHAR FROM t").fetchall() == [(expected,), (expected,)] + def test_map_of_two_plain_values_named_key_and_value_appends_its_entries(self, duckdb_cursor): + duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, VARCHAR))") + frame = pd.DataFrame({"a": pd.Series([{"key": "ab", "value": "cd"}], dtype=object)}) + duckdb_cursor.append("t", frame) + assert duckdb_cursor.table("t").fetchall() == [({"key": "ab", "value": "cd"},)] + def test_map_keys_a_str_and_its_bytes_are_one_key_and_refused(self, duckdb_cursor): duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, INTEGER))") frame = pd.DataFrame({"a": pd.Series([{"a": 1, b"a": 2}], dtype=object)}) diff --git a/tests/fast/test_map_types.py b/tests/fast/test_map_types.py index d541635c..1b4f9b34 100644 --- a/tests/fast/test_map_types.py +++ b/tests/fast/test_map_types.py @@ -606,7 +606,7 @@ def _far_from_1970(value): @FAR_YEARS -@pytest.mark.parametrize("sql_type", ["TIMESTAMP", "TIMESTAMP_S"]) +@pytest.mark.parametrize("sql_type", ["TIMESTAMP", "TIMESTAMP_MS", "TIMESTAMP_S"]) def test_map_timestamp_column_handed_back_as_timestamps_outside_the_nanosecond_range_round_trips( duckdb_cursor, sql_type, value ): @@ -720,11 +720,11 @@ def test_map_scalar_inside_a_container_of_a_type_holding_a_map_is_refused_and_th @pytest.mark.parametrize( ("sql_type", "cell", "converted"), [ - pytest.param("MAP(VARCHAR, VARCHAR)", {"key": "ab", "value": "cd"}, "VARCHAR and VARCHAR", id="map"), + pytest.param("MAP(VARCHAR, VARCHAR)", {"key": ["a"], "value": "b"}, "VARCHAR[] and VARCHAR", id="map"), pytest.param( "STRUCT(m MAP(VARCHAR, VARCHAR))", - {"m": {"key": "ab", "value": "cd"}}, - "VARCHAR and VARCHAR", + {"m": {"key": ["a"], "value": "b"}}, + "VARCHAR[] and VARCHAR", id="map_field", ), pytest.param( @@ -736,7 +736,7 @@ def test_map_scalar_inside_a_container_of_a_type_holding_a_map_is_refused_and_th ], ) @pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append", "appender", "parameter"]) -def test_map_of_a_key_and_a_value_that_are_no_lists_is_refused_and_the_database_stays_usable( +def test_map_of_a_key_list_and_a_value_that_is_no_list_is_refused_and_the_database_stays_usable( sql_type, cell, converted, mode ): con = duckdb.connect() @@ -761,6 +761,73 @@ def run(): assert con.sql("select 42").fetchall() == [(42,)] +@pytest.mark.parametrize( + "value", + [ + "MAP {'key': 'ab', 'value': 'cd'}", + "MAP {'key': 'ab', 'value': 'cde'}", + "MAP {'key': 1, 'value': 2}", + "MAP {'key': 'x', 'value': NULL}", + "MAP {'key': NULL, 'value': 'y'}", + ], + ids=["text_of_one_length", "text_of_two_lengths", "ints", "null_value", "null_key_entry"], +) +@SCHEMAS +def test_map_identity_of_a_map_keyed_by_key_and_value_reads_its_entries(duckdb_cursor, value, explicit_schema): + # df() hands such a MAP over as the dict of its two entries, which no list makes a dict of key and value lists + source = duckdb_cursor.sql(f"select {value} as a") + mapped = source.map(lambda df: df, schema=_schema_of(source, explicit_schema)) + assert mapped.types == source.types + assert mapped.fetchall() == source.fetchall() + + +@pytest.mark.parametrize( + "value", ["MAP {'key': NULL, 'value': NULL}", "MAP {'key': [1], 'value': NULL}"], ids=["both_null", "list_and_null"] +) +@SCHEMAS +def test_map_identity_of_a_map_keyed_by_key_and_value_of_a_null_beside_a_list_or_a_null_reads_null( + duckdb_cursor, value, explicit_schema +): + # df() hands it over as {"key": None, "value": None} or {"key": [1], "value": None}, the forms of a NULL MAP + source = duckdb_cursor.sql(f"select {value} as a") + mapped = source.map(lambda df: df, schema=_schema_of(source, explicit_schema)) + assert mapped.types == source.types + assert mapped.fetchall() == [(None,)] + + +@pytest.mark.parametrize("mode", ["append", "appender", "parameter"]) +def test_map_of_a_null_key_entry_beside_a_plain_value_reads_its_two_entries(mode): + cell = {"key": None, "value": 5} + con = duckdb.connect() + con.execute("create table t (a MAP(VARCHAR, INTEGER))") + if mode == "append": + con.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) + elif mode == "appender": + with con.appender("t") as appender: + appender.append_row(cell) + else: + con.execute("insert into t values (?)", [cell]) + assert con.table("t").fetchall() == [({"key": None, "value": 5},)] + + +@pytest.mark.parametrize( + "cell", + [{"key": None, "value": None}, {"key": ["a"], "value": None}, {"key": None, "value": ["a"]}], + ids=["both_none", "values_none", "keys_none"], +) +@pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append"]) +def test_map_of_key_and_value_with_a_none_reads_a_null_map(cell, mode): + con = duckdb.connect() + con.execute("create table t (a MAP(VARCHAR, VARCHAR))") + if mode == "append": + con.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) + assert con.table("t").fetchall() == [(None,)] + else: + con.execute("insert into t values (MAP {'k': 'v'})") + schema = {"a": duckdb.sqltype("MAP(VARCHAR, VARCHAR)")} if mode == "explicit_schema" else None + assert con.table("t").map(_returning(cell), schema=schema).fetchall() == [(None,)] + + @pytest.mark.parametrize( ("cell", "expected"), [ From 0da3cbfc8bc660c20e2507d83cce4b48cfc2c044 Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Sat, 7 Nov 2026 00:43:50 +0000 Subject: [PATCH 6/9] Take a dict as MAP keys and values only for two lists, read a declared STRUCT by its field names --- src/native/python_conversion.cpp | 92 ++++++++--------- src/numpy/numpy_scan.cpp | 6 +- tests/fast/pandas/test_append_df.py | 10 ++ tests/fast/test_map_types.py | 151 ++++++++++++++++++++++------ 4 files changed, 174 insertions(+), 85 deletions(-) diff --git a/src/native/python_conversion.cpp b/src/native/python_conversion.cpp index 1e5f4ad6..77f613b6 100644 --- a/src/native/python_conversion.cpp +++ b/src/native/python_conversion.cpp @@ -97,22 +97,28 @@ static vector TransformStructKeys(nb::handle keys, idx_t size) { return res; } -//! Whether a MAP's key or value component is a column of keys or values: a list, a tuple or a numpy array, which the -//! conversion reads as a LIST (an array of no dimensions it refuses) -static bool IsMapComponentList(nb::handle component) { - return PyList_Check(component.ptr()) || PyTuple_Check(component.ptr()) || - duckdb::PyUtil::IsInstance(component, DuckDBPyConnection::ImportCache()->numpy.ndarray()); +//! Whether an object is a list the conversion reads as a LIST: a list, a tuple or a numpy array of at least one +//! dimension, since an array of none holds a single value +static bool IsListLike(nb::handle ele) { + if (PyList_Check(ele.ptr()) || PyTuple_Check(ele.ptr())) { + return true; + } + return duckdb::PyUtil::IsInstance(ele, DuckDBPyConnection::ImportCache()->numpy.ndarray()) && + nb::cast(ele.attr("ndim")) > 0; } -//! The length of a MAP's key or value component, or -1 for one that has none although it claims one, as a numpy -//! array of no dimensions does -static int64_t ComponentLength(nb::handle component) { - auto length = PyObject_Length(component.ptr()); - if (length < 0) { - PyErr_Clear(); - return -1; +//! A dict of exactly 'key' and 'value' with a list on one side and a plain value on the other: a STRUCT of those +//! fields without a target, and refused as a MAP's key and value lists for a MAP target +static bool MixesAListAndAValue(const PyDictionary &dict) { + if (dict.len != 2) { + return false; } - return length; + auto keys = dict[nb::str("key")]; + auto values = dict[nb::str("value")]; + if (!keys || !values || nb::none().is(keys) || nb::none().is(values)) { + return false; + } + return IsListLike(keys) != IsListLike(values); } bool DictionaryHasMapFormat(const PyDictionary &dict) { @@ -129,25 +135,14 @@ bool DictionaryHasMapFormat(const PyDictionary &dict) { return false; } - // A list on either side, or None on both, makes the dict a MAP of the two lists: None a NULL MAP, and the - // conversion refuses the other side when it is no list. Without a list it is a plain dict, a MAP of the entries - // 'key' and 'value', so a NULL entry beside a plain value stays an entry + // Two lists of one length are a MAP's keys and values, and None on both sides, or beside a list, a NULL MAP; any + // other dict is a plain one, so a NULL entry or a list beside a plain value stays an entry or a field auto key_none = nb::none().is(keys); auto value_none = nb::none().is(values); - auto key_list = IsMapComponentList(keys); - auto value_list = IsMapComponentList(values); - if (!key_list && !value_list && !(key_none && value_none)) { - return false; + if (key_none || value_none) { + return (key_none || IsListLike(keys)) && (value_none || IsListLike(values)); } - if (key_none || value_none || !key_list || !value_list) { - return true; - } - - // Verify that both the keys and values are of the same length; a component without one is refused by the - // conversion - auto key_size = ComponentLength(keys); - auto value_size = ComponentLength(values); - return key_size < 0 || value_size < 0 || key_size == value_size; + return IsListLike(keys) && IsListLike(values) && nb::len(keys) == nb::len(values); } Value TransformDictionaryToStruct(optional_ptr context, const PyDictionary &dict, @@ -310,16 +305,6 @@ Value TransformTimestamp(optional_ptr context, nb::handle ele, co return TransformPythonValue(context, ele, LogicalType::UNKNOWN, nan_as_null); } -//! Whether an object is a list a shape takes: a list, a tuple or a numpy array of at least one dimension, since an -//! array of none holds a single value -static bool ShapeTakesAsList(nb::handle ele) { - if (PyList_Check(ele.ptr()) || PyTuple_Check(ele.ptr())) { - return true; - } - return duckdb::PyUtil::IsInstance(ele, DuckDBPyConnection::ImportCache()->numpy.ndarray()) && - nb::cast(ele.attr("ndim")) > 0; -} - //! Whether an object is what a shape is for: a dict for a MAP or a STRUCT, a list for a LIST or an ARRAY; any other //! object converts without a target and meets the engine's cast or the kind check, so a tuple reads as a LIST at //! every depth, as the appender reads a value of a row @@ -330,7 +315,7 @@ static bool ShapeTakes(nb::handle ele, const LogicalType &shape) { return PyDict_Check(ele.ptr()); case LogicalTypeId::LIST: case LogicalTypeId::ARRAY: - return ShapeTakesAsList(ele); + return IsListLike(ele); default: return false; } @@ -424,6 +409,13 @@ Value TransformStructFormatDictionaryToMap(optional_ptr context, Value TransformDictionaryToMap(optional_ptr context, const PyDictionary &dict, const LogicalType &target_type, bool nan_as_null) { if (target_type.id() != LogicalTypeId::UNKNOWN && !DictionaryHasMapFormat(dict)) { + if (MixesAListAndAValue(dict)) { + auto keys = TransformPythonValue(context, dict[nb::str("key")], LogicalType::UNKNOWN, nan_as_null); + auto values = TransformPythonValue(context, dict[nb::str("value")], LogicalType::UNKNOWN, nan_as_null); + throw InvalidInputException("A dict of 'key' and 'value' for a MAP takes a list of keys and a list of " + "values, not %s and %s", + keys.type().ToString(), values.type().ToString()); + } // dict == { 'k1': v1, 'k2': v2, ..., 'kn': vn } return TransformStructFormatDictionaryToMap(context, dict, target_type, nan_as_null); } @@ -436,9 +428,8 @@ Value TransformDictionaryToMap(optional_ptr context, const PyDict return Value(LogicalType::MAP(LogicalTypeId::SQLNULL, LogicalTypeId::SQLNULL)); } - auto key_size = ComponentLength(keys); - auto sized = key_size >= 0 && ComponentLength(values) >= 0; - if (sized && key_size == 0) { + auto key_size = nb::len(keys); + if (key_size == 0) { // dict == { 'key': [], 'value': [] } return EmptyMapValue(); } @@ -448,15 +439,12 @@ Value TransformDictionaryToMap(optional_ptr context, const PyDict LogicalType value_target = LogicalTypeId::UNKNOWN; if (IsShape(target_type)) { - // A shape reads the two lists entry by entry, as it reads a dict's own entries; a component that is no list - // converts without a target and is refused below - if (sized && ShapeTakesAsList(keys) && ShapeTakesAsList(values)) { - auto items = [](nb::handle list) { - return PyList_Check(list.ptr()) || PyTuple_Check(list.ptr()) ? nb::borrow(list) - : list.attr("tolist")(); - }; - return MapFromEntries(context, items(keys), items(values), key_size, target_type, nan_as_null); - } + // A shape reads the two lists entry by entry, as it reads a dict's own entries + auto items = [](nb::handle list) { + return PyList_Check(list.ptr()) || PyTuple_Check(list.ptr()) ? nb::borrow(list) + : list.attr("tolist")(); + }; + return MapFromEntries(context, items(keys), items(values), key_size, target_type, nan_as_null); } else if (target_type.id() != LogicalTypeId::UNKNOWN) { key_target = LogicalType::LIST(MapType::KeyType(target_type)); value_target = LogicalType::LIST(MapType::ValueType(target_type)); diff --git a/src/numpy/numpy_scan.cpp b/src/numpy/numpy_scan.cpp index 4e724b13..e0a28203 100644 --- a/src/numpy/numpy_scan.cpp +++ b/src/numpy/numpy_scan.cpp @@ -941,9 +941,11 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t bool bit_storage_refused) { DeclaredVector staged(count); auto &target = out.GetType(); - // A type holding a MAP takes a shape for its keys, and one holding a naive timestamp for the unit a nested value - // reads in; a top level timestamp is read by DeclaredValue itself + // A type holding a MAP takes a shape for its keys, one holding a STRUCT for a dict to name its fields, not to be + // guessed a MAP's key and value lists, and one holding a naive timestamp for the unit a nested value reads in; a + // top level timestamp is read by DeclaredValue itself const auto keyed = target.IsNested() && (TypeVisitor::Contains(target, LogicalTypeId::MAP) || + TypeVisitor::Contains(target, LogicalTypeId::STRUCT) || TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_NS) || TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP) || TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_MS) || diff --git a/tests/fast/pandas/test_append_df.py b/tests/fast/pandas/test_append_df.py index 1083016f..b0c0819f 100644 --- a/tests/fast/pandas/test_append_df.py +++ b/tests/fast/pandas/test_append_df.py @@ -540,6 +540,16 @@ def test_map_of_key_and_value_lists_reads_each_entry_as_the_same_dict_does( duckdb_cursor.append("t", frame) assert duckdb_cursor.sql("SELECT a::VARCHAR FROM t").fetchall() == [(expected,), (expected,)] + def test_aware_timestamp_without_nanoseconds_reads_in_the_session_time_zone(self, duckdb_cursor): + duckdb_cursor.execute("SET TimeZone = 'America/New_York'") + duckdb_cursor.execute("CREATE TABLE t (a TIMESTAMP)") + stamp = pd.Timestamp("2020-01-01 00:00", tz="UTC") + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([stamp, stamp.to_pydatetime()], dtype=object)})) + assert duckdb_cursor.sql("SELECT a::VARCHAR FROM t").fetchall() == [ + ("2019-12-31 19:00:00",), + ("2019-12-31 19:00:00",), + ] + def test_map_of_two_plain_values_named_key_and_value_appends_its_entries(self, duckdb_cursor): duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, VARCHAR))") frame = pd.DataFrame({"a": pd.Series([{"key": "ab", "value": "cd"}], dtype=object)}) diff --git a/tests/fast/test_map_types.py b/tests/fast/test_map_types.py index 1b4f9b34..2bfcc378 100644 --- a/tests/fast/test_map_types.py +++ b/tests/fast/test_map_types.py @@ -717,48 +717,137 @@ def test_map_scalar_inside_a_container_of_a_type_holding_a_map_is_refused_and_th assert con.sql("select 42").fetchall() == [(42,)] +def _through(mode, con, sql_type, cell): + con.execute(f"create or replace table t (a {sql_type})") + if mode == "append": + con.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) + elif mode == "appender": + with con.appender("t") as appender: + appender.append_row(cell) + elif mode == "parameter": + con.execute("insert into t values (?)", [cell]) + else: + con.execute("insert into t values (NULL)") + schema = {"a": duckdb.sqltype(sql_type)} if mode == "explicit_schema" else None + return con.table("t").map(_returning(cell), schema=schema).fetchall() + return con.table("t").fetchall() + + @pytest.mark.parametrize( - ("sql_type", "cell", "converted"), + ("sql_type", "cell"), [ - pytest.param("MAP(VARCHAR, VARCHAR)", {"key": ["a"], "value": "b"}, "VARCHAR[] and VARCHAR", id="map"), - pytest.param( - "STRUCT(m MAP(VARCHAR, VARCHAR))", - {"m": {"key": ["a"], "value": "b"}}, - "VARCHAR[] and VARCHAR", - id="map_field", - ), - pytest.param( - "MAP(VARCHAR, INTEGER)", - {"key": np.array("a"), "value": np.array(1)}, - "VARCHAR and INTEGER", - id="zero_dimensional_arrays", - ), + pytest.param("MAP(VARCHAR, VARCHAR)", {"key": ["a"], "value": "b"}, id="map"), + pytest.param("STRUCT(m MAP(VARCHAR, VARCHAR))", {"m": {"key": ["a"], "value": "b"}}, id="map_field"), ], ) -@pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append", "appender", "parameter"]) -def test_map_of_a_key_list_and_a_value_that_is_no_list_is_refused_and_the_database_stays_usable( - sql_type, cell, converted, mode +@pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append"]) +def test_map_of_a_key_list_beside_a_plain_value_for_a_map_type_is_refused_and_the_database_stays_usable( + sql_type, cell, mode ): con = duckdb.connect() - con.execute(f"create table t (a {sql_type})") + expected = ( + "A dict of 'key' and 'value' for a MAP takes a list of keys and a list of values, not VARCHAR[] and VARCHAR" + ) + with pytest.raises(duckdb.InvalidInputException, match=re.escape(expected)): + _through(mode, con, sql_type, cell) + assert con.sql("select 42").fetchall() == [(42,)] + + +@pytest.mark.parametrize("mode", ["appender", "parameter"]) +def test_map_of_a_key_list_beside_a_plain_value_without_a_target_converts_as_a_struct(mode): + # The conversion has no MAP to read for, so the dict is the STRUCT of its fields, which the engine casts into the + # MAP of the two entries, the list as its text + con = duckdb.connect() + assert _through(mode, con, "MAP(VARCHAR, VARCHAR)", {"key": ["a"], "value": "b"}) == [ + ({"key": "[a]", "value": "b"},) + ] + + +@pytest.mark.parametrize( + ("cell", "struct_type"), + [ + pytest.param({"key": [], "value": "ab"}, 'STRUCT("key" "NULL"[], "value" VARCHAR)', id="empty_list_and_text"), + pytest.param({"key": [1], "value": 5}, 'STRUCT("key" INTEGER[], "value" INTEGER)', id="list_and_int"), + pytest.param({"key": "x", "value": None}, 'STRUCT("key" VARCHAR, "value" "NULL")', id="text_and_null"), + ], +) +@pytest.mark.parametrize("mode", ["parameter", "dataframe_scan"]) +def test_dict_of_key_and_value_without_two_lists_is_inferred_as_a_struct(cell, struct_type, mode): + con = duckdb.connect() + if mode == "parameter": + rows = con.execute("select ?, typeof(?)", [cell, cell]).fetchall() + else: + frame = pd.DataFrame({"a": pd.Series([cell], dtype=object)}) # noqa: F841 + rows = con.sql("select a, typeof(a) from frame").fetchall() + assert rows == [(cell, struct_type)] + + +KEY_AND_VALUE_STRUCTS = pytest.mark.parametrize( + "value", + [ + "{'key': 'color', 'value': ['red', 'blue']}::STRUCT(key VARCHAR, value VARCHAR[])", + "map_entries(MAP {'color': ['red', 'blue']})", + "{'key': [1], 'value': 5}::STRUCT(key INTEGER[], value INTEGER)", + "{'key': 1, 'value': [5]}::STRUCT(key INTEGER, value INTEGER[])", + ], + ids=["text_and_list", "map_entries", "list_and_int", "int_and_list"], +) + + +@KEY_AND_VALUE_STRUCTS +@pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append", "appender", "parameter", "dataframe_scan"]) +def test_struct_with_fields_key_and_value_one_a_list_converts_as_its_struct(value, mode): + con = duckdb.connect() + source = con.sql(f"select {value} as a") + expected = source.fetchall() + frame = source.df() + if mode == "dataframe_scan": + assert duckdb.connect().sql("select a from frame").fetchall() == expected + return + if mode in ("inferred", "explicit_schema"): + schema = _schema_of(source, mode == "explicit_schema") + assert source.map(lambda df: df, schema=schema).fetchall() == expected + return + con.execute(f"create table t (a {source.types[0]})") if mode == "append": - run = lambda: con.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) # noqa: E731 + con.append("t", frame) elif mode == "appender": + with con.appender("t") as appender: + appender.append_row(frame["a"][0]) + else: + con.execute("insert into t values (?)", [frame["a"][0]]) + assert con.table("t").fetchall() == expected - def run(): - with con.appender("t") as appender: - appender.append_row(cell) - elif mode == "parameter": - run = lambda: con.execute("insert into t values (?)", [cell]) # noqa: E731 +@pytest.mark.parametrize( + "value", + [ + "{'key': NULL, 'value': NULL}::STRUCT(key VARCHAR, value VARCHAR)", + "{'key': [1], 'value': NULL}::STRUCT(key INTEGER[], value INTEGER)", + ], + ids=["both_null", "list_and_null"], +) +@pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append"]) +def test_declared_struct_with_fields_key_and_value_reads_a_null_field_as_its_struct(value, mode): + # A declared STRUCT names its fields, so such a dict is no NULL MAP of key and value lists + con = duckdb.connect() + source = con.sql(f"select {value} as a") + expected = source.fetchall() + if mode == "append": + con.execute(f"create table t (a {source.types[0]})") + con.append("t", source.df()) + assert con.table("t").fetchall() == expected else: - con.execute("insert into t values (NULL)") - schema = {"a": duckdb.sqltype(sql_type)} if mode == "explicit_schema" else None - run = con.table("t").map(_returning(cell), schema=schema).fetchall - expected = f"A dict of 'key' and 'value' for a MAP takes a list of keys and a list of values, not {converted}" - with pytest.raises(duckdb.InvalidInputException, match=re.escape(expected)): - run() - assert con.sql("select 42").fetchall() == [(42,)] + mapped = source.map(lambda df: df, schema=_schema_of(source, mode == "explicit_schema")) + assert mapped.fetchall() == expected + + +@pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append", "appender", "parameter"]) +def test_map_of_key_and_value_arrays_of_no_dimensions_reads_their_two_entries(mode): + # An array of no dimensions holds a single value, so the dict is a plain one + con = duckdb.connect() + cell = {"key": np.array("a"), "value": np.array("b")} + assert _through(mode, con, "MAP(VARCHAR, VARCHAR)", cell) == [({"key": "a", "value": "b"},)] @pytest.mark.parametrize( From fee081e93f072c7de815ca651d0f93c341c9d7aa Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Sat, 7 Nov 2026 04:52:45 +0000 Subject: [PATCH 7/9] Stage a MAP column's keys and values once per vector and index a STRUCT shape's fields once per scan --- .../duckdb_python/python_conversion.hpp | 11 +++ src/native/python_conversion.cpp | 98 ++++++++++++++----- src/numpy/numpy_scan.cpp | 83 +++++++++++++--- tests/fast/pandas/test_append_df.py | 19 ++++ tests/fast/test_map_types.py | 46 +++++++++ 5 files changed, 217 insertions(+), 40 deletions(-) diff --git a/src/include/duckdb_python/python_conversion.hpp b/src/include/duckdb_python/python_conversion.hpp index 3d96301c..dec847c8 100644 --- a/src/include/duckdb_python/python_conversion.hpp +++ b/src/include/duckdb_python/python_conversion.hpp @@ -53,6 +53,17 @@ LogicalType SniffPythonIntegerType(nb::handle ele); //! The type of a duckdb.Value, or InvalidInputException when its 'type' attribute is not a DuckDB type LogicalType PythonValueType(nb::handle value); bool DictionaryHasMapFormat(const PyDictionary &dict); +//! A dict of exactly 'key' and 'value' with a list on one side and a plain value on the other: a STRUCT of those +//! fields without a target, and refused as a MAP's key and value lists for a MAP target +bool MixesAListAndAValue(const PyDictionary &dict); +//! Drops the field indexes the conversion of shaped STRUCT dicts kept on this thread; the declared scan calls it when +//! the scan of a vector ends, so no index outlives the shape it was built for +void ClearStructFieldIndexes(); +//! A key or value of a MAP whose target is a type without children converts without a target and is then cast with +//! the connection's settings, as a value of a column is, so each keeps its own type until the cast instead of meeting +//! its siblings' common type (2**60 + 1 beside 1.0 would read as a DOUBLE, rounded) +Value TransformMapEntry(optional_ptr context, nb::handle entry, const LogicalType &target_type, + bool nan_as_null); //! A target the declared scan builds from a column's type (a shape), as opposed to a column's own type: a dict for a //! shaped STRUCT takes its own keys, as without a target; a shaped MAP casts each key and value whose type has no //! children with the connection's settings; a shaped naive timestamp reads a cell as TransformTimestamp does diff --git a/src/native/python_conversion.cpp b/src/native/python_conversion.cpp index 76d63e9e..8d884a20 100644 --- a/src/native/python_conversion.cpp +++ b/src/native/python_conversion.cpp @@ -191,9 +191,7 @@ static bool IsListLike(nb::handle ele) { nb::cast(ele.attr("ndim")) > 0; } -//! A dict of exactly 'key' and 'value' with a list on one side and a plain value on the other: a STRUCT of those -//! fields without a target, and refused as a MAP's key and value lists for a MAP target -static bool MixesAListAndAValue(const PyDictionary &dict) { +bool MixesAListAndAValue(const PyDictionary &dict) { if (dict.len != 2) { return false; } @@ -229,6 +227,40 @@ bool DictionaryHasMapFormat(const PyDictionary &dict) { return IsListLike(keys) && IsListLike(values) && nb::len(keys) == nb::len(values); } +namespace { +//! The field index of each shaped STRUCT a declared scan meets, kept per thread until the scan of the vector ends, +//! since each takes a pass over the fields to build and every dict of the vector asks the same one; an entry holds +//! its type, so the type information it is found by cannot be freed and given to another type while it is kept, and +//! an index stays where it is while others are added +struct FieldIndex { + LogicalType type; + identifier_map_t fields; +}; + +unordered_map &FieldIndexes() { + thread_local unordered_map indexes; + return indexes; +} + +const identifier_map_t &FieldIndexOf(const LogicalType &struct_type) { + auto &indexes = FieldIndexes(); + auto entry = indexes.find(&struct_type.GetTypeInfo()); + if (entry != indexes.end()) { + return entry->second.fields; + } + auto &index = indexes[&struct_type.GetTypeInfo()]; + index.type = struct_type; + for (idx_t child = 0; child < StructType::GetChildCount(struct_type); child++) { + index.fields.emplace(StructType::GetChildName(struct_type, child), child); + } + return index.fields; +} +} // namespace + +void ClearStructFieldIndexes() { + FieldIndexes().clear(); +} + Value TransformDictionaryToStruct(optional_ptr context, const PyDictionary &dict, const LogicalType &target_type, bool nan_as_null) { auto conversion_error = [&]() { @@ -256,15 +288,30 @@ Value TransformDictionaryToStruct(optional_ptr context, const PyD auto struct_keys = TransformStructKeys(dict.keys, dict.len); if (!struct_target || shape_only) { + const LogicalType unknown(LogicalTypeId::UNKNOWN); + auto field_count = shape_only ? StructType::GetChildCount(target_type) : 0; + // The field at a key's own position is tried first, as df() hands a STRUCT over in field order; a key found + // elsewhere is looked up in an index of the field names, so a wide STRUCT costs no search for every key + optional_ptr> fields; for (idx_t i = 0; i < dict.len; i++) { - LogicalType child_type = LogicalType::UNKNOWN; - for (idx_t child = 0; shape_only && child < StructType::GetChildCount(target_type); child++) { - if (StructType::GetChildName(target_type, child) == struct_keys[i]) { - child_type = StructType::GetChildType(target_type, child); - break; + optional_idx field; + if (i < field_count && StructType::GetChildName(target_type, i) == struct_keys[i]) { + field = i; + } else if (field_count > 0) { + if (!fields) { + fields = &FieldIndexOf(target_type); + } + auto entry = fields->find(struct_keys[i]); + if (entry != fields->end()) { + field = entry->second; } } - auto val = TransformPythonValue(context, dict.values.attr("__getitem__")(i), child_type, nan_as_null); + auto &child_type = field.IsValid() ? StructType::GetChildType(target_type, field.GetIndex()) : unknown; + auto value = nb::steal(PySequence_GetItem(dict.values.ptr(), NumericCast(i))); + if (!value.is_valid()) { + throw nb::python_error(); + } + auto val = TransformPythonValue(context, value, child_type, nan_as_null); struct_values.emplace_back(make_pair(std::move(struct_keys[i]), std::move(val))); } return Value::STRUCT(std::move(struct_values)); @@ -368,11 +415,8 @@ static bool ShapeTakes(nb::handle ele, const LogicalType &shape) { } } -//! A key or value of a MAP whose target is a type without children converts without a target and is then cast with -//! the connection's settings, as a value of a column is, so each keeps its own type until the cast instead of meeting -//! its siblings' common type (2**60 + 1 beside 1.0 would read as a DOUBLE, rounded) -static Value TransformMapEntry(optional_ptr context, nb::handle entry, const LogicalType &target_type, - bool nan_as_null) { +Value TransformMapEntry(optional_ptr context, nb::handle entry, const LogicalType &target_type, + bool nan_as_null) { if (IsShape(target_type) || target_type.IsNested() || target_type.id() == LogicalTypeId::UNKNOWN || target_type.id() == LogicalTypeId::UNION) { return TransformPythonValue(context, entry, target_type, nan_as_null); @@ -400,11 +444,16 @@ static Value MapFromEntries(optional_ptr context, nb::handle keys LogicalType key_type = LogicalType::SQLNULL; LogicalType value_type = LogicalType::SQLNULL; - vector elements; + vector map_keys; + vector map_values; + map_keys.reserve(size); + map_values.reserve(size); for (idx_t i = 0; i < size; i++) { - - auto key = keys.attr("__getitem__")(i); - auto value = values.attr("__getitem__")(i); + auto key = nb::steal(PySequence_GetItem(keys.ptr(), NumericCast(i))); + auto value = nb::steal(PySequence_GetItem(values.ptr(), NumericCast(i))); + if (!key.is_valid() || !value.is_valid()) { + throw nb::python_error(); + } Value new_key = shape ? TransformMapEntry(context, key, key_target, nan_as_null) : TransformPythonValue(context, key, key_target, nan_as_null); Value new_value = shape ? TransformMapEntry(context, value, value_target, nan_as_null) @@ -412,12 +461,8 @@ static Value MapFromEntries(optional_ptr context, nb::handle keys key_type = ProxiedForceMaxLogicalType(context, key_type, new_key.type()); value_type = ProxiedForceMaxLogicalType(context, value_type, new_value.type()); - - child_list_t struct_values; - struct_values.emplace_back(make_pair("key", std::move(new_key))); - struct_values.emplace_back(make_pair("value", std::move(new_value))); - - elements.push_back(Value::STRUCT(std::move(struct_values))); + map_keys.push_back(std::move(new_key)); + map_values.push_back(std::move(new_value)); } // A target holding UNKNOWN asks for no type there, and a MAP of such a type is no value the engine can hold if (key_type.id() == LogicalTypeId::SQLNULL && !TypeVisitor::Contains(key_target, LogicalTypeId::UNKNOWN) && @@ -429,9 +474,8 @@ static Value MapFromEntries(optional_ptr context, nb::handle keys value_type = value_target; } - LogicalType map_type = LogicalType::MAP(key_type, value_type); - - return Value::MAP(ListType::GetChildType(map_type), std::move(elements)); + // Built from the keys and values themselves, as the STRUCT of each entry would only be taken apart again + return Value::MAP(key_type, value_type, std::move(map_keys), std::move(map_values)); } Value TransformStructFormatDictionaryToMap(optional_ptr context, const PyDictionary &dict, diff --git a/src/numpy/numpy_scan.cpp b/src/numpy/numpy_scan.cpp index 9938f88a..d30026bd 100644 --- a/src/numpy/numpy_scan.cpp +++ b/src/numpy/numpy_scan.cpp @@ -207,18 +207,10 @@ void ScanNumpyObject(optional_ptr context, PyObject *object, idx_ TransformPythonObject(context, object, out, offset); } +//! Raises what the engine raises building a MAP Value of such keys, so a MAP written straight into its vector and one +//! built as a Value are refused alike static void VerifyMapConstraints(Vector &vec, idx_t count) { - auto invalid_reason = MapVector::CheckMapValidity(vec, count); - switch (invalid_reason) { - case MapInvalidReason::VALID: - return; - case MapInvalidReason::DUPLICATE_KEY: - throw InvalidInputException("Dict->Map conversion failed because 'key' list contains duplicates"); - case MapInvalidReason::NULL_KEY: - throw InvalidInputException("Dict->Map conversion failed because 'key' list contains None"); - default: - throw InvalidInputException("Option not implemented for MapInvalidReason"); - } + MapVector::EvalMapInvalidReason(MapVector::CheckMapValidity(vec, count)); } //! A MAP's keys must be unique at every depth; a cast of a STRUCT into a MAP, as of a dict whose keys collapse into @@ -343,6 +335,11 @@ static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb return TransformPythonValue(context, value, shape, true); } +//! A key or value type a stated MAP casts its entries into one by one +static bool IsPlainEntryType(const LogicalType &type) { + return !type.IsNested() && type.id() != LogicalTypeId::UNION; +} + //! The type a cell of a nested type converts with: its containers and its naive timestamps and no other type, so a //! dict builds its MAP from its keys and values in the types they have in Python and every key and value meets the //! engine's cast with the connection's settings, as a TIMESTAMPTZ in the session time zone, and a timestamp reads in @@ -351,8 +348,7 @@ static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb //! stated takes that, while an implied type keeps the entries' own types for the kind check. static LogicalType KeyedShape(const LogicalType &type, bool typed_entries) { auto entry = [&](const LogicalType &entry_type) { - auto plain = !entry_type.IsNested() && entry_type.id() != LogicalTypeId::UNION; - return typed_entries && plain ? entry_type : KeyedShape(entry_type, typed_entries); + return typed_entries && IsPlainEntryType(entry_type) ? entry_type : KeyedShape(entry_type, typed_entries); }; switch (type.id()) { case LogicalTypeId::STRUCT: { @@ -928,10 +924,30 @@ static void ScanObjectCells(ClientContext &context, PyObject **col, idx_t stride } } +//! Writes the converted keys or values of a MAP vector's entries from first on, cast into the child's type as a +//! declared column's values are +static void WriteEntries(ClientContext &context, vector &entries, Vector &child, idx_t first) { + if (entries.empty()) { + return; + } + DeclaredVector staged(entries.size()); + for (idx_t i = 0; i < entries.size(); i++) { + staged.Set(i, std::move(entries[i])); + } + Vector written(child.GetType(), entries.size()); + staged.Write(context, written, entries.size()); + VectorOperations::Copy(written, child, entries.size(), 0, first); +} + void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t stride, const uint8_t *masked, idx_t count, idx_t offset, Vector &out, bool implied_type, bool bit_storage_refused) { DeclaredVector staged(count); + struct FieldIndexesOfTheScan { + ~FieldIndexesOfTheScan() { + ClearStructFieldIndexes(); + } + } field_indexes; auto &target = out.GetType(); // A type holding a MAP takes a shape for its keys, one holding a STRUCT for a dict to name its fields, not to be // guessed a MAP's key and value lists, and one holding a naive timestamp for the unit a nested value reads in; a @@ -1043,6 +1059,47 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t } read(cell, row); }); + } else if (target.id() == LogicalTypeId::MAP && !implied_type && !bits_refused && + IsPlainEntryType(MapType::KeyType(target)) && IsPlainEntryType(MapType::ValueType(target))) { + // A dict of its entries goes straight into the MAP's keys and values, each converted as a MAP entry is and all + // of them cast together once the vector is read, since a MAP Value per cell builds and takes apart a STRUCT + // Value per entry; the key and value lists and every other cell take the general reader + auto &key_type = MapType::KeyType(target); + auto &value_type = MapType::ValueType(target); + auto first_entry = ListVector::GetListSize(out); + vector keys; + vector values; + ScanObjectCells(context, col, stride, masked, count, offset, out, [&](nb::handle cell, idx_t row) { + if (!PyDict_Check(cell.ptr())) { + return read(cell, row); + } + auto size = NumericCast(PyDict_Size(cell.ptr())); + if (size == 2) { + PyDictionary dict(nb::borrow(cell)); + if (DictionaryHasMapFormat(dict) || MixesAListAndAValue(dict)) { + return read(cell, row); + } + } + auto start = ListVector::GetListSize(out); + Py_ssize_t position = 0; + PyObject *key; + PyObject *value; + while (PyDict_Next(cell.ptr(), &position, &key, &value)) { + // Converting runs Python, which could drop the dict's references to them + auto key_object = nb::borrow(key); + auto value_object = nb::borrow(value); + keys.push_back(TransformTimestamp(context, key_object, key_type, true)); + values.push_back(TransformTimestamp(context, value_object, value_type, true)); + } + auto &list_entry = FlatVector::GetDataMutable(out)[row]; + list_entry.offset = start; + list_entry.length = first_entry + keys.size() - start; + ListVector::Reserve(out, first_entry + keys.size()); + ListVector::SetListSize(out, first_entry + keys.size()); + staged.MarkWritten(row); + }); + WriteEntries(context, keys, MapVector::GetKeys(out), first_entry); + WriteEntries(context, values, MapVector::GetValues(out), first_entry); } else { ScanObjectCells(context, col, stride, masked, count, offset, out, read); } diff --git a/tests/fast/pandas/test_append_df.py b/tests/fast/pandas/test_append_df.py index c35e9182..61b31802 100644 --- a/tests/fast/pandas/test_append_df.py +++ b/tests/fast/pandas/test_append_df.py @@ -586,6 +586,25 @@ def test_aware_timestamp_without_nanoseconds_reads_in_the_session_time_zone(self ("2019-12-31 19:00:00",), ] + def test_struct_dicts_with_keys_out_of_field_order_read_each_field_by_name(self, duckdb_cursor): + # Twelve STRUCT types, more than the per thread field indexes keep, each dict's keys in reverse field order + inner = [f"s{i} STRUCT(a{i} INTEGER, b{i} INTEGER, c{i} INTEGER)" for i in range(12)] + duckdb_cursor.execute(f"CREATE TABLE t (i INTEGER, a STRUCT({', '.join(inner)}))") + cells = [ + {f"s{i}": {f"c{i}": row + 3, f"b{i}": row + 2, f"a{i}": row + 1} for i in reversed(range(12))} + for row in range(3000) + ] + duckdb_cursor.append("t", pd.DataFrame({"i": range(3000), "a": pd.Series(cells, dtype=object)})) + assert [row[0] for row in duckdb_cursor.sql("SELECT a FROM t ORDER BY i").fetchall()] == cells + + def test_map_column_mixing_dicts_with_other_cells_keeps_every_row_in_place(self, duckdb_cursor): + # Dicts are written into the MAP's entries directly and the other cells after them; 3000 rows span two vectors + cells = [{"a": 1}, "{b=2}", {"key": ["c"], "value": [3]}, None, {"d": 4, "e": 5}, {}] * 500 + expected = [{"a": 1}, {"b": 2}, {"c": 3}, None, {"d": 4, "e": 5}, {}] * 500 + duckdb_cursor.execute("CREATE TABLE t (i INTEGER, a MAP(VARCHAR, INTEGER))") + duckdb_cursor.append("t", pd.DataFrame({"i": range(len(cells)), "a": pd.Series(cells, dtype=object)})) + assert [row[0] for row in duckdb_cursor.sql("SELECT a FROM t ORDER BY i").fetchall()] == expected + def test_map_of_two_plain_values_named_key_and_value_appends_its_entries(self, duckdb_cursor): duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, VARCHAR))") frame = pd.DataFrame({"a": pd.Series([{"key": "ab", "value": "cd"}], dtype=object)}) diff --git a/tests/fast/test_map_types.py b/tests/fast/test_map_types.py index 2bfcc378..683ee0a0 100644 --- a/tests/fast/test_map_types.py +++ b/tests/fast/test_map_types.py @@ -842,6 +842,15 @@ def test_declared_struct_with_fields_key_and_value_reads_a_null_field_as_its_str assert mapped.fetchall() == expected +@pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append"]) +def test_struct_given_an_empty_dict_is_refused(mode): + # A missing field reads NULL, but a dict of no keys is no STRUCT of the column's fields + con = duckdb.connect() + with pytest.raises(duckdb.TypeMismatchException, match="Cannot cast STRUCTs of different size"): + _through(mode, con, "STRUCT(a INTEGER, b INTEGER)", {}) + assert con.sql("select 42").fetchall() == [(42,)] + + @pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append", "appender", "parameter"]) def test_map_of_key_and_value_arrays_of_no_dimensions_reads_their_two_entries(mode): # An array of no dimensions holds a single value, so the dict is a plain one @@ -1020,6 +1029,43 @@ def test_map_nested_timestamp_with_nanoseconds_for_a_stated_coarser_unit_reads_a assert rows == [(expected,)] +@pytest.mark.parametrize( + ("value", "text"), + [ + pytest.param( + np.datetime64("2020-01-01T00:00:00.000005789", "ns"), + "2020-01-01 00:00:00.000006", + id="datetime64_rounds_up", + ), + pytest.param(np.datetime64(-1, "ns"), "1970-01-01 00:00:00", id="datetime64_before_the_epoch"), + pytest.param( + pd.Timestamp("2020-01-01 00:00:00.000005789"), "2020-01-01 00:00:00.000005", id="timestamp_as_its_datetime" + ), + pytest.param(pd.Timestamp(-1, unit="ns"), "1969-12-31 23:59:59.999999", id="timestamp_before_the_epoch"), + ], +) +@pytest.mark.parametrize("place", ["top", "list_element", "struct_field", "map_value"]) +@pytest.mark.parametrize("mode", ["explicit_schema", "append"]) +def test_map_timestamp_finer_than_a_stated_coarser_unit_reads_at_every_depth_as_at_the_top(value, text, place, mode): + # A datetime64 is a TIMESTAMP_NS the engine's cast rounds into microseconds; a Timestamp is the datetime it is + sql_type, cell, expected = { + "top": ("TIMESTAMP", value, text), + "list_element": ("TIMESTAMP[]", [value], f"['{text}']"), + "struct_field": ("STRUCT(t TIMESTAMP)", {"t": value}, f"{{'t': '{text}'}}"), + "map_value": ("MAP(VARCHAR, TIMESTAMP)", {"k": value}, f"{{k='{text}'}}"), + }[place] + con = duckdb.connect() + con.execute(f"create table t (a {sql_type})") + if mode == "append": + con.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) + rows = con.sql("select a::VARCHAR from t").fetchall() + else: + con.execute("insert into t values (NULL)") + schema = {"a": duckdb.sqltype(sql_type)} + rows = con.table("t").map(_returning(cell), schema=schema).project("a::VARCHAR").fetchall() + assert rows == [(expected,)] + + @pytest.mark.parametrize( "sql_type", ["STRUCT(a INTEGER, b INTEGER)", "STRUCT(a INTEGER, t TIMESTAMP)", "STRUCT(a INTEGER, m MAP(VARCHAR, INTEGER))"], From 5ee6ecc552016f4ad62226dd6e9895567cc8b96a Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Sat, 7 Nov 2026 07:10:20 +0000 Subject: [PATCH 8/9] Clear the STRUCT field index from the outermost scan only, take exact dicts in the MAP reader, keep a zoned Timestamp's nanoseconds for TIMESTAMPTZ_NS --- .../duckdb_python/python_conversion.hpp | 24 ++-- src/native/python_conversion.cpp | 24 +++- src/numpy/numpy_scan.cpp | 52 ++++++--- tests/fast/pandas/test_append_df.py | 109 ++++++++++++++++-- tests/fast/test_map_types.py | 78 +++++++++++++ 5 files changed, 244 insertions(+), 43 deletions(-) diff --git a/src/include/duckdb_python/python_conversion.hpp b/src/include/duckdb_python/python_conversion.hpp index dec847c8..33133d42 100644 --- a/src/include/duckdb_python/python_conversion.hpp +++ b/src/include/duckdb_python/python_conversion.hpp @@ -56,9 +56,15 @@ bool DictionaryHasMapFormat(const PyDictionary &dict); //! A dict of exactly 'key' and 'value' with a list on one side and a plain value on the other: a STRUCT of those //! fields without a target, and refused as a MAP's key and value lists for a MAP target bool MixesAListAndAValue(const PyDictionary &dict); -//! Drops the field indexes the conversion of shaped STRUCT dicts kept on this thread; the declared scan calls it when -//! the scan of a vector ends, so no index outlives the shape it was built for -void ClearStructFieldIndexes(); +//! Keeps the field indexes the conversion of shaped STRUCT dicts builds on this thread while a declared scan runs; the +//! outermost scope drops them when it ends, so no index outlives the shapes it was built for +class StructFieldIndexScope { +public: + StructFieldIndexScope(); + ~StructFieldIndexScope(); + StructFieldIndexScope(const StructFieldIndexScope &) = delete; + StructFieldIndexScope &operator=(const StructFieldIndexScope &) = delete; +}; //! A key or value of a MAP whose target is a type without children converts without a target and is then cast with //! the connection's settings, as a value of a column is, so each keeps its own type until the cast instead of meeting //! its siblings' common type (2**60 + 1 beside 1.0 would read as a DOUBLE, rounded) @@ -71,12 +77,12 @@ LogicalType AsShape(const LogicalType &type); bool IsShape(const LogicalType &type); //! The nanosecond a datetime subclass such as a pandas Timestamp gives as an int, or -1 when it gives none int64_t NanosecondOf(nb::handle cell); -//! Whether a cell is a naive pandas Timestamp with nanoseconds, which no Python datetime holds; they are then in -//! nanoseconds -bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds); -//! A value for a naive timestamp type converted without a target, but keeping the nanoseconds of a pandas Timestamp -//! for a TIMESTAMP_NS, and reading a numpy scalar as PythonValueOfNumpyScalar does, so a datetime64 finer than a -//! microsecond is a TIMESTAMP_NS that a coarser unit takes through the engine's cast +//! Whether a cell is a pandas Timestamp with nanoseconds, which no Python datetime holds, naive or, when zoned, zone +//! aware; they are then the nanoseconds since the epoch, in UTC for a zone aware one +bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds, bool zoned); +//! A value for a timestamp type converted without a target, but keeping the nanoseconds of a pandas Timestamp for a +//! TIMESTAMP_NS or a TIMESTAMPTZ_NS, and reading a numpy scalar as PythonValueOfNumpyScalar does, so a datetime64 finer +//! than a microsecond is a TIMESTAMP_NS that a coarser unit takes through the engine's cast Value TransformTimestamp(optional_ptr context, nb::handle ele, const LogicalType &target_type, bool nan_as_null); void TransformPythonObject(optional_ptr context, nb::handle ele, Vector &vector, idx_t result_offset, diff --git a/src/native/python_conversion.cpp b/src/native/python_conversion.cpp index 8d884a20..dab25344 100644 --- a/src/native/python_conversion.cpp +++ b/src/native/python_conversion.cpp @@ -257,8 +257,18 @@ const identifier_map_t &FieldIndexOf(const LogicalType &struct_type) { } } // namespace -void ClearStructFieldIndexes() { - FieldIndexes().clear(); +// Python a conversion runs, as a tzinfo's utcoffset, can start another declared scan on this thread while an outer +// conversion holds an index, so only the outermost scope drops them +static thread_local idx_t struct_field_index_scopes = 0; + +StructFieldIndexScope::StructFieldIndexScope() { + struct_field_index_scopes++; +} + +StructFieldIndexScope::~StructFieldIndexScope() { + if (--struct_field_index_scopes == 0) { + FieldIndexes().clear(); + } } Value TransformDictionaryToStruct(optional_ptr context, const PyDictionary &dict, @@ -371,7 +381,7 @@ int64_t NanosecondOf(nb::handle cell) { return PyLong_Check(nanosecond.ptr()) ? PyLong_AsLongLong(nanosecond.ptr()) : -1; } -bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds) { +bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds, bool zoned) { auto datetime_type = DuckDBPyConnection::ImportCache()->datetime.datetime(); if (Py_TYPE(cell.ptr()) == reinterpret_cast(datetime_type.ptr()) || !duckdb::PyUtil::IsInstance(cell, datetime_type)) { @@ -379,9 +389,10 @@ bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds) { } // A Timestamp without nanoseconds reads as the datetime it is: its value in nanoseconds overflows outside the // years 1677 to 2262, and a coarser column would judge that value by a cast where the datetime fits outright - if (NanosecondOf(cell) <= 0 || !nb::getattr(cell, "tzinfo").is_none()) { + if (NanosecondOf(cell) <= 0 || nb::getattr(cell, "tzinfo").is_none() == zoned) { return false; } + // A zone aware Timestamp's value counts the nanoseconds since the epoch in UTC, as a TIMESTAMPTZ_NS does nanoseconds = timestamp_ns_t(nb::cast(cell.attr("value"))); return true; } @@ -389,9 +400,12 @@ bool HoldsNanoseconds(nb::handle cell, timestamp_ns_t &nanoseconds) { Value TransformTimestamp(optional_ptr context, nb::handle ele, const LogicalType &target_type, bool nan_as_null) { timestamp_ns_t nanoseconds; - if (target_type.id() == LogicalTypeId::TIMESTAMP_NS && HoldsNanoseconds(ele, nanoseconds)) { + if (target_type.id() == LogicalTypeId::TIMESTAMP_NS && HoldsNanoseconds(ele, nanoseconds, false)) { return Value::TIMESTAMPNS(nanoseconds); } + if (target_type.id() == LogicalTypeId::TIMESTAMP_TZ_NS && HoldsNanoseconds(ele, nanoseconds, true)) { + return Value::TIMESTAMPTZNS(timestamp_tz_ns_t(nanoseconds.value)); + } auto plain = PythonValueOfNumpyScalar(ele, nanoseconds); if (!plain.is_valid()) { return Value::TIMESTAMPNS(nanoseconds); diff --git a/src/numpy/numpy_scan.cpp b/src/numpy/numpy_scan.cpp index d30026bd..671269dc 100644 --- a/src/numpy/numpy_scan.cpp +++ b/src/numpy/numpy_scan.cpp @@ -309,17 +309,22 @@ static Value DeclaredValue(ClientContext &context, const LogicalType &target, nb if (cell.ptr() == Py_None) { return Value(); } - // A TIMESTAMP_NS column keeps a pandas Timestamp's nanoseconds, and an implied naive timestamp type sees them, so - // its kind check refuses what a coarser unit would drop; a stated coarser type reads the Timestamp as the datetime - // it also is + // A TIMESTAMP_NS or TIMESTAMPTZ_NS column keeps a pandas Timestamp's nanoseconds, and an implied timestamp type of + // a coarser unit sees them, so its kind check refuses what the unit would drop; a stated coarser type reads the + // Timestamp as the datetime it also is auto naive_timestamp = target.id() == LogicalTypeId::TIMESTAMP_NS || (implied && (target.id() == LogicalTypeId::TIMESTAMP || target.id() == LogicalTypeId::TIMESTAMP_MS || target.id() == LogicalTypeId::TIMESTAMP_SEC)); + auto zoned_timestamp = + target.id() == LogicalTypeId::TIMESTAMP_TZ_NS || (implied && target.id() == LogicalTypeId::TIMESTAMP_TZ); timestamp_ns_t nanoseconds; - if (naive_timestamp && HoldsNanoseconds(cell, nanoseconds)) { + if (naive_timestamp && HoldsNanoseconds(cell, nanoseconds, false)) { return Value::TIMESTAMPNS(nanoseconds); } + if (zoned_timestamp && HoldsNanoseconds(cell, nanoseconds, true)) { + return Value::TIMESTAMPTZNS(timestamp_tz_ns_t(nanoseconds.value)); + } auto plain = PythonValueOfNumpyScalar(cell, nanoseconds); if (!plain.is_valid()) { return Value::TIMESTAMPNS(nanoseconds); @@ -367,6 +372,7 @@ static LogicalType KeyedShape(const LogicalType &type, bool typed_entries) { case LogicalTypeId::MAP: return AsShape(LogicalType::MAP(entry(MapType::KeyType(type)), entry(MapType::ValueType(type)))); case LogicalTypeId::TIMESTAMP_NS: + case LogicalTypeId::TIMESTAMP_TZ_NS: return AsShape(type); case LogicalTypeId::TIMESTAMP: case LogicalTypeId::TIMESTAMP_MS: @@ -374,6 +380,9 @@ static LogicalType KeyedShape(const LogicalType &type, bool typed_entries) { // An implied type sees the nanoseconds a coarser unit would drop, so its kind check refuses them; a stated one // takes them through the engine's cast into its unit, as at the top return AsShape(typed_entries ? type : LogicalType::TIMESTAMP_NS); + case LogicalTypeId::TIMESTAMP_TZ: + // A stated one reads a zone aware Timestamp as the datetime it also is, as at the top + return typed_entries ? LogicalType(LogicalTypeId::UNKNOWN) : AsShape(LogicalType::TIMESTAMP_TZ_NS); default: return LogicalType::UNKNOWN; } @@ -622,7 +631,11 @@ class ImpliedKinds { auto zoned_timestamps = (source.id() == LogicalTypeId::TIMESTAMP_TZ || source.id() == LogicalTypeId::TIMESTAMP_TZ_NS) && (type.id() == LogicalTypeId::TIMESTAMP_TZ || type.id() == LogicalTypeId::TIMESTAMP_TZ_NS); - if (zoned_timestamps || (source.id() == LogicalTypeId::VARCHAR && type.id() == LogicalTypeId::VARCHAR) || + // Another unit of a zoned timestamp, as of a naive one, holds only the values the cast keeps unchanged + if (zoned_timestamps) { + return source.id() == type.id() ? Verdict::FITS : Verdict::BY_VALUE; + } + if ((source.id() == LogicalTypeId::VARCHAR && type.id() == LogicalTypeId::VARCHAR) || (source.id() == LogicalTypeId::DOUBLE && type.id() == LogicalTypeId::FLOAT)) { return Verdict::FITS; } @@ -943,21 +956,20 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t idx_t count, idx_t offset, Vector &out, bool implied_type, bool bit_storage_refused) { DeclaredVector staged(count); - struct FieldIndexesOfTheScan { - ~FieldIndexesOfTheScan() { - ClearStructFieldIndexes(); - } - } field_indexes; + StructFieldIndexScope field_indexes; auto &target = out.GetType(); // A type holding a MAP takes a shape for its keys, one holding a STRUCT for a dict to name its fields, not to be - // guessed a MAP's key and value lists, and one holding a naive timestamp for the unit a nested value reads in; a - // top level timestamp is read by DeclaredValue itself - const auto keyed = target.IsNested() && (TypeVisitor::Contains(target, LogicalTypeId::MAP) || - TypeVisitor::Contains(target, LogicalTypeId::STRUCT) || - TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_NS) || - TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP) || - TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_MS) || - TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_SEC)); + // guessed a MAP's key and value lists, and one holding a timestamp for the unit a nested value reads in; a + // top level timestamp is read by DeclaredValue itself, and a stated TIMESTAMPTZ takes no mark in the shape + const auto keyed = + target.IsNested() && + (TypeVisitor::Contains(target, LogicalTypeId::MAP) || TypeVisitor::Contains(target, LogicalTypeId::STRUCT) || + TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_NS) || + TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_TZ_NS) || + (implied_type && TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_TZ)) || + TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP) || + TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_MS) || + TypeVisitor::Contains(target, LogicalTypeId::TIMESTAMP_SEC)); const auto bits_refused = bit_storage_refused && TypeVisitor::Contains(target, LogicalTypeId::BIT); const auto shape = keyed ? KeyedShape(target, !implied_type) : LogicalType(LogicalTypeId::UNKNOWN); ImpliedKinds kinds(context); @@ -1070,7 +1082,9 @@ void NumpyScan::ScanDeclaredColumn(ClientContext &context, PyObject **col, idx_t vector keys; vector values; ScanObjectCells(context, col, stride, masked, count, offset, out, [&](nb::handle cell, idx_t row) { - if (!PyDict_Check(cell.ptr())) { + // A dict subclass, as an OrderedDict, may iterate in another order than its storage, which is what + // PyDict_Next walks + if (!PyDict_CheckExact(cell.ptr())) { return read(cell, row); } auto size = NumericCast(PyDict_Size(cell.ptr())); diff --git a/tests/fast/pandas/test_append_df.py b/tests/fast/pandas/test_append_df.py index 61b31802..86a4e601 100644 --- a/tests/fast/pandas/test_append_df.py +++ b/tests/fast/pandas/test_append_df.py @@ -1,3 +1,4 @@ +import collections import datetime import enum import math @@ -587,15 +588,104 @@ def test_aware_timestamp_without_nanoseconds_reads_in_the_session_time_zone(self ] def test_struct_dicts_with_keys_out_of_field_order_read_each_field_by_name(self, duckdb_cursor): - # Twelve STRUCT types, more than the per thread field indexes keep, each dict's keys in reverse field order - inner = [f"s{i} STRUCT(a{i} INTEGER, b{i} INTEGER, c{i} INTEGER)" for i in range(12)] - duckdb_cursor.execute(f"CREATE TABLE t (i INTEGER, a STRUCT({', '.join(inner)}))") - cells = [ - {f"s{i}": {f"c{i}": row + 3, f"b{i}": row + 2, f"a{i}": row + 1} for i in reversed(range(12))} - for row in range(3000) - ] + # Each dict's keys in reverse field order, for twelve STRUCT types and one inside a LIST; the TIMESTAMP_NS and + # BLOB keyed MAP fields read their nanoseconds and bytes only when the key finds its field + nanosecond = pd.Timestamp("2020-01-01 00:00:00.000000001") + + def fields(i): + return f"a{i} INTEGER, t{i} TIMESTAMP_NS, m{i} MAP(BLOB, INTEGER)" + + def cell(i, row): + return {f"m{i}": {b"\x00k": row}, f"t{i}": nanosecond, f"a{i}": row} + + inner = [f"s{i} STRUCT({fields(i)})" for i in range(12)] + duckdb_cursor.execute(f"CREATE TABLE t (i INTEGER, a STRUCT({', '.join(inner)}, l STRUCT({fields(12)})[]))") + cells = [{"l": [cell(12, row)], **{f"s{i}": cell(i, row) for i in reversed(range(12))}} for row in range(3000)] duckdb_cursor.append("t", pd.DataFrame({"i": range(3000), "a": pd.Series(cells, dtype=object)})) - assert [row[0] for row in duckdb_cursor.sql("SELECT a FROM t ORDER BY i").fetchall()] == cells + places = [(f"a.s{i}", i) for i in range(12)] + [("a.l[1]", 12)] + columns = ", ".join(f"{place}.a{i}, {place}.t{i}::VARCHAR, {place}.m{i}" for place, i in places) + expected = [(row, "2020-01-01 00:00:00.000000001", {b"\x00k": row}) * len(places) for row in range(3000)] + assert duckdb_cursor.sql(f"SELECT {columns} FROM t ORDER BY i").fetchall() == expected + + def test_struct_append_inside_a_conversion_keeps_the_outer_scans_field_index(self, duckdb_cursor): + # A tzinfo's utcoffset that appends on another connection runs a declared scan inside this one's, between + # this dict's out of order keys + other = duckdb.connect() + other.execute("CREATE TABLE y (c BIGINT)") + + class Appending(datetime.tzinfo): + def utcoffset(self, dt): + other.append("y", pd.DataFrame({"c": pd.Series([1, 2, 3], dtype=object)})) + return datetime.timedelta(0) + + def dst(self, dt): + return datetime.timedelta(0) + + def tzname(self, dt): + return "UTC" + + fields = [f"f{j}" for j in range(30)] + duckdb_cursor.execute(f"CREATE TABLE t (a STRUCT(t TIMESTAMPTZ, {', '.join(f'{f} INTEGER' for f in fields)}))") + moment = datetime.datetime(2020, 1, 1, tzinfo=Appending()) + cells = [{"f29": row, "t": moment, **dict.fromkeys(reversed(fields[:-1]), row)} for row in range(20)] + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series(cells, dtype=object)})) + assert duckdb_cursor.sql("SELECT count(*), sum(a.f0), sum(a.f29) FROM t").fetchall() == [(20, 190, 190)] + assert other.sql("SELECT count(*) > 0 FROM y").fetchall() == [(True,)] + + @pytest.mark.parametrize( + ("target", "cell", "expected"), + [ + pytest.param("TIMESTAMPTZ_NS", "stamp", "2020-01-01 00:00:00.000000001+00", id="top"), + pytest.param( + "STRUCT(t TIMESTAMPTZ_NS)", {"t": "stamp"}, "{'t': '2020-01-01 00:00:00.000000001+00'}", id="field" + ), + pytest.param("TIMESTAMPTZ_NS[]", ["stamp"], "['2020-01-01 00:00:00.000000001+00']", id="list_element"), + pytest.param( + "MAP(VARCHAR, TIMESTAMPTZ_NS)", {"k": "stamp"}, "{k='2020-01-01 00:00:00.000000001+00'}", id="map_value" + ), + ], + ) + def test_zoned_timestamp_with_nanoseconds_keeps_them_in_a_timestamptz_ns( + self, duckdb_cursor, target, cell, expected + ): + stamp = pd.Timestamp("2020-01-01 00:00:00.000000001", tz="UTC") + + def with_stamp(value): + if isinstance(value, dict): + return {key: with_stamp(item) for key, item in value.items()} + if isinstance(value, list): + return [with_stamp(item) for item in value] + return stamp if value == "stamp" else value + + duckdb_cursor.execute("SET TimeZone = 'UTC'") + duckdb_cursor.execute(f"CREATE TABLE t (a {target})") + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([with_stamp(cell)], dtype=object)})) + assert duckdb_cursor.sql("SELECT a::VARCHAR FROM t").fetchall() == [(expected,)] + + @pytest.mark.parametrize( + ("stamp", "text"), + [ + pytest.param(pd.Timestamp.max.tz_localize("UTC"), "infinity", id="max"), + pytest.param(pd.Timestamp.min.tz_localize("UTC"), "-infinity", id="min"), + ], + ) + @pytest.mark.parametrize("place", ["top", "field", "list_element"]) + def test_zoned_timestamp_max_and_min_store_infinity_in_a_timestamptz_ns(self, duckdb_cursor, stamp, text, place): + target, cell, expected = { + "top": ("TIMESTAMPTZ_NS", stamp, text), + "field": ("STRUCT(t TIMESTAMPTZ_NS)", {"t": stamp}, f"{{'t': {text}}}"), + "list_element": ("TIMESTAMPTZ_NS[]", [stamp], f"[{text}]"), + }[place] + duckdb_cursor.execute(f"CREATE TABLE t (a {target})") + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) + assert duckdb_cursor.sql("SELECT a::VARCHAR FROM t").fetchall() == [(expected,)] + + def test_map_of_an_ordered_dict_keeps_its_iteration_order(self, duckdb_cursor): + ordered = collections.OrderedDict([("a", 1), ("b", 2), ("c", 3)]) + ordered.move_to_end("a") + duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, INTEGER))") + duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([ordered], dtype=object)})) + assert duckdb_cursor.sql("SELECT map_keys(a) FROM t").fetchall() == [(["b", "c", "a"],)] def test_map_column_mixing_dicts_with_other_cells_keeps_every_row_in_place(self, duckdb_cursor): # Dicts are written into the MAP's entries directly and the other cells after them; 3000 rows span two vectors @@ -634,9 +724,8 @@ def test_map_values_of_mixed_numbers_each_keep_their_value(self): def test_map_key_of_none_is_refused_and_the_connection_stays_usable(self, duckdb_cursor): duckdb_cursor.execute("CREATE TABLE t (a MAP(VARCHAR, INTEGER))") - with pytest.raises(duckdb.Error) as error: + with pytest.raises(duckdb.InvalidInputException, match=re.escape("Map keys can not be NULL.")): duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series([{None: 1}], dtype=object)})) - assert not isinstance(error.value, duckdb.FatalException), error.value assert duckdb_cursor.sql("SELECT 42").fetchall() == [(42,)] def test_text_with_a_lone_surrogate_names_it(self, duckdb_cursor): diff --git a/tests/fast/test_map_types.py b/tests/fast/test_map_types.py index 683ee0a0..043ff41f 100644 --- a/tests/fast/test_map_types.py +++ b/tests/fast/test_map_types.py @@ -1,5 +1,6 @@ """map() fixes one type per output column at bind and reads each chunk's returned column into it.""" +import collections import datetime import json import re @@ -842,6 +843,14 @@ def test_declared_struct_with_fields_key_and_value_reads_a_null_field_as_its_str assert mapped.fetchall() == expected +def test_map_of_an_ordered_dict_for_a_stated_map_keeps_its_iteration_order(duckdb_cursor): + ordered = collections.OrderedDict([("a", 1), ("b", 2), ("c", 3)]) + ordered.move_to_end("a") + source = duckdb_cursor.sql("select MAP {'k': 1} as a") + mapped = source.map(_returning(ordered), schema={"a": duckdb.sqltype("MAP(VARCHAR, INTEGER)")}) + assert mapped.project("map_keys(a)").fetchall() == [(["b", "c", "a"],)] + + @pytest.mark.parametrize("mode", ["inferred", "explicit_schema", "append"]) def test_struct_given_an_empty_dict_is_refused(mode): # A missing field reads NULL, but a dict of no keys is no STRUCT of the column's fields @@ -971,6 +980,75 @@ def test_map_nested_timestamp_ns_value_keeps_its_nanoseconds(duckdb_cursor, valu assert "00:00:00.000000001" in mapped.project("a::VARCHAR").fetchall()[0][0] +AWARE_NANOSECOND = pd.Timestamp("2020-01-01 00:00:00.000000001", tz="UTC") + + +@pytest.mark.parametrize( + ("value", "cell"), + [ + pytest.param("'2020-01-01 00:00:00.000000001+00'::TIMESTAMPTZ_NS", AWARE_NANOSECOND, id="top"), + pytest.param("{'t': '2020-01-01 00:00:00.000000001+00'::TIMESTAMPTZ_NS}", {"t": AWARE_NANOSECOND}, id="field"), + pytest.param("['2020-01-01 00:00:00.000000001+00'::TIMESTAMPTZ_NS]", [AWARE_NANOSECOND], id="list_element"), + ], +) +@SCHEMAS +def test_map_zoned_timestamp_ns_value_keeps_its_nanoseconds(duckdb_cursor, value, cell, explicit_schema): + duckdb_cursor.execute("SET TimeZone = 'UTC'") + source = duckdb_cursor.sql(f"select {value} as a") + mapped = source.map(_returning(cell), schema=_schema_of(source, explicit_schema)) + assert mapped.types == source.types + assert mapped.project("a::VARCHAR").fetchall() == source.project("a::VARCHAR").fetchall() + + +@pytest.mark.parametrize( + ("value", "cell", "returned"), + [ + pytest.param("'2020-01-01 00:00:00+00'::TIMESTAMPTZ", AWARE_NANOSECOND, "TIMESTAMPTZ_NS", id="top"), + pytest.param( + "{'t': '2020-01-01 00:00:00+00'::TIMESTAMPTZ}", + {"t": AWARE_NANOSECOND}, + "STRUCT(t TIMESTAMPTZ_NS)", + id="field", + ), + pytest.param( + "['2020-01-01 00:00:00+00'::TIMESTAMPTZ]", [AWARE_NANOSECOND], "TIMESTAMPTZ_NS[]", id="list_element" + ), + ], +) +def test_map_zoned_timestamp_with_nanoseconds_for_an_implied_timestamptz_is_refused( + duckdb_cursor, value, cell, returned +): + duckdb_cursor.execute("SET TimeZone = 'UTC'") + source = duckdb_cursor.sql(f"select {value} as a") + expected = f"UDF column type mismatch, expected [{source.types[0]}], got [{returned}]" + with pytest.raises(duckdb.InvalidInputException, match=re.escape(expected)): + source.map(_returning(cell)).fetchall() + + +@pytest.mark.parametrize("place", ["top", "field", "list_element", "map_value"]) +@pytest.mark.parametrize("mode", ["explicit_schema", "append"]) +def test_map_zoned_timestamp_with_nanoseconds_for_a_stated_timestamptz_reads_as_its_datetime(place, mode): + stamp = pd.Timestamp("2020-01-01 00:00:00.000005789", tz="UTC") + text = "2020-01-01 00:00:00.000005+00" + sql_type, cell, expected = { + "top": ("TIMESTAMPTZ", stamp, text), + "field": ("STRUCT(t TIMESTAMPTZ)", {"t": stamp}, f"{{'t': '{text}'}}"), + "list_element": ("TIMESTAMPTZ[]", [stamp], f"['{text}']"), + "map_value": ("MAP(VARCHAR, TIMESTAMPTZ)", {"k": stamp}, f"{{k='{text}'}}"), + }[place] + con = duckdb.connect() + con.execute("SET TimeZone = 'UTC'") + con.execute(f"create table t (a {sql_type})") + if mode == "append": + con.append("t", pd.DataFrame({"a": pd.Series([cell], dtype=object)})) + rows = con.sql("select a::VARCHAR from t").fetchall() + else: + con.execute("insert into t values (NULL)") + schema = {"a": duckdb.sqltype(sql_type)} + rows = con.table("t").map(_returning(cell), schema=schema).project("a::VARCHAR").fetchall() + assert rows == [(expected,)] + + def test_map_nested_timestamp_with_nanoseconds_for_an_implied_timestamp_is_refused(duckdb_cursor): mapped = duckdb_cursor.sql("select {'t': TIMESTAMP '2020-01-01'} as a").map(_returning({"t": NANOSECOND})) expected = ( From e4bb2a7bb4737eb0aba43b3f60934b481ad4d4d9 Mon Sep 17 00:00:00 2001 From: Evert Lammerts Date: Sat, 7 Nov 2026 08:06:56 +0000 Subject: [PATCH 9/9] Nest the reentrancy test's scans on one thread and pin the zoned nanoseconds under a New York zone --- tests/fast/pandas/test_append_df.py | 26 +++++++++++++++++++------- tests/fast/test_map_types.py | 16 +++++++++++----- 2 files changed, 30 insertions(+), 12 deletions(-) diff --git a/tests/fast/pandas/test_append_df.py b/tests/fast/pandas/test_append_df.py index 86a4e601..0ebe0b2b 100644 --- a/tests/fast/pandas/test_append_df.py +++ b/tests/fast/pandas/test_append_df.py @@ -609,8 +609,8 @@ def cell(i, row): def test_struct_append_inside_a_conversion_keeps_the_outer_scans_field_index(self, duckdb_cursor): # A tzinfo's utcoffset that appends on another connection runs a declared scan inside this one's, between - # this dict's out of order keys - other = duckdb.connect() + # this dict's out of order keys; with one thread that scan runs on this thread, which holds the index + other = duckdb.connect(config={"threads": 1}) other.execute("CREATE TABLE y (c BIGINT)") class Appending(datetime.tzinfo): @@ -625,12 +625,22 @@ def tzname(self, dt): return "UTC" fields = [f"f{j}" for j in range(30)] - duckdb_cursor.execute(f"CREATE TABLE t (a STRUCT(t TIMESTAMPTZ, {', '.join(f'{f} INTEGER' for f in fields)}))") + integers = ", ".join(f"{f} INTEGER" for f in fields) + duckdb_cursor.execute( + f"CREATE TABLE t (a STRUCT(t TIMESTAMPTZ, {integers}, n TIMESTAMP_NS, m MAP(BLOB, INTEGER)))" + ) moment = datetime.datetime(2020, 1, 1, tzinfo=Appending()) - cells = [{"f29": row, "t": moment, **dict.fromkeys(reversed(fields[:-1]), row)} for row in range(20)] + stamp = pd.Timestamp("2020-01-01 00:00:00.000000001") + # The fields after the callback read right only through the index: n keeps its nanosecond, m its bytes keys + cells = [ + {"f29": row, "t": moment, "n": stamp, "m": {b"\x00k": row}, **dict.fromkeys(reversed(fields[:-1]), row)} + for row in range(20) + ] duckdb_cursor.append("t", pd.DataFrame({"a": pd.Series(cells, dtype=object)})) - assert duckdb_cursor.sql("SELECT count(*), sum(a.f0), sum(a.f29) FROM t").fetchall() == [(20, 190, 190)] + rows = duckdb_cursor.sql("SELECT a.f0, a.f29, a.n::VARCHAR, a.m FROM t ORDER BY a.f0").fetchall() + assert rows == [(row, row, "2020-01-01 00:00:00.000000001", {b"\x00k": row}) for row in range(20)] assert other.sql("SELECT count(*) > 0 FROM y").fetchall() == [(True,)] + other.close() @pytest.mark.parametrize( ("target", "cell", "expected"), @@ -645,10 +655,12 @@ def tzname(self, dt): ), ], ) + @pytest.mark.parametrize("zone", ["UTC", "America/New_York"], ids=["UTC", "America_New_York"]) def test_zoned_timestamp_with_nanoseconds_keeps_them_in_a_timestamptz_ns( - self, duckdb_cursor, target, cell, expected + self, duckdb_cursor, target, cell, expected, zone ): - stamp = pd.Timestamp("2020-01-01 00:00:00.000000001", tz="UTC") + # The same instant in either zone, so a reading of the wall time instead would differ + stamp = pd.Timestamp("2020-01-01 00:00:00.000000001", tz="UTC").tz_convert(zoneinfo.ZoneInfo(zone)) def with_stamp(value): if isinstance(value, dict): diff --git a/tests/fast/test_map_types.py b/tests/fast/test_map_types.py index 043ff41f..9b71b9bc 100644 --- a/tests/fast/test_map_types.py +++ b/tests/fast/test_map_types.py @@ -4,6 +4,7 @@ import datetime import json import re +import zoneinfo from decimal import Decimal from fractions import Fraction @@ -986,16 +987,21 @@ def test_map_nested_timestamp_ns_value_keeps_its_nanoseconds(duckdb_cursor, valu @pytest.mark.parametrize( ("value", "cell"), [ - pytest.param("'2020-01-01 00:00:00.000000001+00'::TIMESTAMPTZ_NS", AWARE_NANOSECOND, id="top"), - pytest.param("{'t': '2020-01-01 00:00:00.000000001+00'::TIMESTAMPTZ_NS}", {"t": AWARE_NANOSECOND}, id="field"), - pytest.param("['2020-01-01 00:00:00.000000001+00'::TIMESTAMPTZ_NS]", [AWARE_NANOSECOND], id="list_element"), + pytest.param("'2020-01-01 00:00:00.000000001+00'::TIMESTAMPTZ_NS", lambda stamp: stamp, id="top"), + pytest.param( + "{'t': '2020-01-01 00:00:00.000000001+00'::TIMESTAMPTZ_NS}", lambda stamp: {"t": stamp}, id="field" + ), + pytest.param("['2020-01-01 00:00:00.000000001+00'::TIMESTAMPTZ_NS]", lambda stamp: [stamp], id="list_element"), ], ) +@pytest.mark.parametrize("zone", ["UTC", "America/New_York"], ids=["UTC", "America_New_York"]) @SCHEMAS -def test_map_zoned_timestamp_ns_value_keeps_its_nanoseconds(duckdb_cursor, value, cell, explicit_schema): +def test_map_zoned_timestamp_ns_value_keeps_its_nanoseconds(duckdb_cursor, value, cell, zone, explicit_schema): + # The same instant in either zone, so a reading of the wall time instead would differ + stamp = AWARE_NANOSECOND.tz_convert(zoneinfo.ZoneInfo(zone)) duckdb_cursor.execute("SET TimeZone = 'UTC'") source = duckdb_cursor.sql(f"select {value} as a") - mapped = source.map(_returning(cell), schema=_schema_of(source, explicit_schema)) + mapped = source.map(_returning(cell(stamp)), schema=_schema_of(source, explicit_schema)) assert mapped.types == source.types assert mapped.project("a::VARCHAR").fetchall() == source.project("a::VARCHAR").fetchall()