duckdb 0.8.2-dev33.0 → 0.8.2-dev3334.0

package/src/duckdb/src/function/scalar_function.cpp

@@ -13,7 +13,8 @@ ScalarFunction::ScalarFunction(string name, vector<LogicalType> arguments, Logic
     : BaseScalarFunction(std::move(name), std::move(arguments), std::move(return_type), side_effects,
                          std::move(varargs), null_handling),
       function(std::move(function)), bind(bind), init_local_state(init_local_state), dependency(dependency),
-      statistics(statistics), serialize(nullptr), deserialize(nullptr) {
+      statistics(statistics), serialize(nullptr), deserialize(nullptr), format_serialize(nullptr),
+      format_deserialize(nullptr) {
 }
 
 ScalarFunction::ScalarFunction(vector<LogicalType> arguments, LogicalType return_type, scalar_function_t function,
@@ -26,9 +27,10 @@ ScalarFunction::ScalarFunction(vector<LogicalType> arguments, LogicalType return
 }
 
 bool ScalarFunction::operator==(const ScalarFunction &rhs) const {
-	return CompareScalarFunctionT(rhs.function) && bind == rhs.bind && dependency == rhs.dependency &&
-	       statistics == rhs.statistics;
+	return name == rhs.name && arguments == rhs.arguments && return_type == rhs.return_type && varargs == rhs.varargs &&
+	       bind == rhs.bind && dependency == rhs.dependency && statistics == rhs.statistics;
 }
+
 bool ScalarFunction::operator!=(const ScalarFunction &rhs) const {
 	return !(*this == rhs);
 }
@@ -56,23 +58,6 @@ bool ScalarFunction::Equal(const ScalarFunction &rhs) const {
 	return true; // they are equal
 }
 
-bool ScalarFunction::CompareScalarFunctionT(const scalar_function_t &other) const {
-	typedef void(scalar_function_ptr_t)(DataChunk &, ExpressionState &, Vector &);
-
-	auto func_ptr = (scalar_function_ptr_t **)function.template target<scalar_function_ptr_t *>(); // NOLINT
-	auto other_ptr = (scalar_function_ptr_t **)other.template target<scalar_function_ptr_t *>();   // NOLINT
-
-	// Case the functions were created from lambdas the target will return a nullptr
-	if (!func_ptr && !other_ptr) {
-		return true;
-	}
-	if (func_ptr == nullptr || other_ptr == nullptr) {
-		// scalar_function_t (std::functions) from lambdas cannot be compared
-		return false;
-	}
-	return CastPointerToValue(*func_ptr) == CastPointerToValue(*other_ptr);
-}
-
 void ScalarFunction::NopFunction(DataChunk &input, ExpressionState &state, Vector &result) {
 	D_ASSERT(input.ColumnCount() >= 1);
 	result.Reference(input.data[0]);

package/src/duckdb/src/function/table/arrow/arrow_duck_schema.cpp

@@ -0,0 +1,57 @@
+#include "duckdb/function/table/arrow/arrow_duck_schema.hpp"
+#include "duckdb/common/arrow/arrow.hpp"
+#include "duckdb/common/exception.hpp"
+
+namespace duckdb {
+
+void ArrowTableType::AddColumn(idx_t index, unique_ptr<ArrowType> type) {
+	D_ASSERT(arrow_convert_data.find(index) == arrow_convert_data.end());
+	arrow_convert_data.emplace(std::make_pair(index, std::move(type)));
+}
+
+const arrow_column_map_t &ArrowTableType::GetColumns() const {
+	return arrow_convert_data;
+}
+
+void ArrowType::AddChild(unique_ptr<ArrowType> child) {
+	children.emplace_back(std::move(child));
+}
+
+void ArrowType::AssignChildren(vector<unique_ptr<ArrowType>> children) {
+	D_ASSERT(this->children.empty());
+	this->children = std::move(children);
+}
+
+void ArrowType::SetDictionary(unique_ptr<ArrowType> dictionary) {
+	D_ASSERT(!this->dictionary_type);
+	dictionary_type = std::move(dictionary);
+}
+
+const ArrowType &ArrowType::GetDictionary() const {
+	D_ASSERT(dictionary_type);
+	return *dictionary_type;
+}
+
+const LogicalType &ArrowType::GetDuckType() const {
+	return type;
+}
+
+ArrowVariableSizeType ArrowType::GetSizeType() const {
+	return size_type;
+}
+
+ArrowDateTimeType ArrowType::GetDateTimeType() const {
+	return date_time_precision;
+}
+
+const ArrowType &ArrowType::operator[](idx_t index) const {
+	D_ASSERT(index < children.size());
+	return *children[index];
+}
+
+idx_t ArrowType::FixedSize() const {
+	D_ASSERT(size_type == ArrowVariableSizeType::FIXED_SIZE);
+	return fixed_size;
+}
+
+} // namespace duckdb

package/src/duckdb/src/function/table/arrow.cpp

@@ -9,41 +9,37 @@
 #include "duckdb/function/table/arrow.hpp"
 #include "duckdb/function/table_function.hpp"
 #include "duckdb/parser/parsed_data/create_table_function_info.hpp"
+#include "duckdb/function/table/arrow/arrow_duck_schema.hpp"
 #include "utf8proc_wrapper.hpp"
 
 namespace duckdb {
 
-LogicalType ArrowTableFunction::GetArrowLogicalType(
-    ArrowSchema &schema, std::unordered_map<idx_t, unique_ptr<ArrowConvertData>> &arrow_convert_data, idx_t col_idx) {
+unique_ptr<ArrowType> ArrowTableFunction::GetArrowLogicalType(ArrowSchema &schema) {
 	auto format = string(schema.format);
-	if (arrow_convert_data.find(col_idx) == arrow_convert_data.end()) {
-		arrow_convert_data[col_idx] = make_uniq<ArrowConvertData>();
-	}
-	auto &convert_data = *arrow_convert_data[col_idx];
 	if (format == "n") {
-		return LogicalType::SQLNULL;
+		return make_uniq<ArrowType>(LogicalType::SQLNULL);
 	} else if (format == "b") {
-		return LogicalType::BOOLEAN;
+		return make_uniq<ArrowType>(LogicalType::BOOLEAN);
 	} else if (format == "c") {
-		return LogicalType::TINYINT;
+		return make_uniq<ArrowType>(LogicalType::TINYINT);
 	} else if (format == "s") {
-		return LogicalType::SMALLINT;
+		return make_uniq<ArrowType>(LogicalType::SMALLINT);
 	} else if (format == "i") {
-		return LogicalType::INTEGER;
+		return make_uniq<ArrowType>(LogicalType::INTEGER);
 	} else if (format == "l") {
-		return LogicalType::BIGINT;
+		return make_uniq<ArrowType>(LogicalType::BIGINT);
 	} else if (format == "C") {
-		return LogicalType::UTINYINT;
+		return make_uniq<ArrowType>(LogicalType::UTINYINT);
 	} else if (format == "S") {
-		return LogicalType::USMALLINT;
+		return make_uniq<ArrowType>(LogicalType::USMALLINT);
 	} else if (format == "I") {
-		return LogicalType::UINTEGER;
+		return make_uniq<ArrowType>(LogicalType::UINTEGER);
 	} else if (format == "L") {
-		return LogicalType::UBIGINT;
+		return make_uniq<ArrowType>(LogicalType::UBIGINT);
 	} else if (format == "f") {
-		return LogicalType::FLOAT;
+		return make_uniq<ArrowType>(LogicalType::FLOAT);
 	} else if (format == "g") {
-		return LogicalType::DOUBLE;
+		return make_uniq<ArrowType>(LogicalType::DOUBLE);
 	} else if (format[0] == 'd') { //! this can be either decimal128 or decimal 256 (e.g., d:38,0)
 		std::string parameters = format.substr(format.find(':'));
 		uint8_t width = std::stoi(parameters.substr(1, parameters.find(',')));
@@ -51,116 +47,138 @@ LogicalType ArrowTableFunction::GetArrowLogicalType(
 		if (width > 38) {
 			throw NotImplementedException("Unsupported Internal Arrow Type for Decimal %s", format);
 		}
-		return LogicalType::DECIMAL(width, scale);
+		return make_uniq<ArrowType>(LogicalType::DECIMAL(width, scale));
 	} else if (format == "u") {
-		convert_data.variable_sz_type.emplace_back(ArrowVariableSizeType::NORMAL, 0);
-		return LogicalType::VARCHAR;
+		return make_uniq<ArrowType>(LogicalType::VARCHAR, ArrowVariableSizeType::NORMAL);
 	} else if (format == "U") {
-		convert_data.variable_sz_type.emplace_back(ArrowVariableSizeType::SUPER_SIZE, 0);
-		return LogicalType::VARCHAR;
+		return make_uniq<ArrowType>(LogicalType::VARCHAR, ArrowVariableSizeType::SUPER_SIZE);
 	} else if (format == "tsn:") {
-		return LogicalTypeId::TIMESTAMP_NS;
+		return make_uniq<ArrowType>(LogicalTypeId::TIMESTAMP_NS);
 	} else if (format == "tsu:") {
-		return LogicalTypeId::TIMESTAMP;
+		return make_uniq<ArrowType>(LogicalTypeId::TIMESTAMP);
 	} else if (format == "tsm:") {
-		return LogicalTypeId::TIMESTAMP_MS;
+		return make_uniq<ArrowType>(LogicalTypeId::TIMESTAMP_MS);
 	} else if (format == "tss:") {
-		return LogicalTypeId::TIMESTAMP_SEC;
+		return make_uniq<ArrowType>(LogicalTypeId::TIMESTAMP_SEC);
 	} else if (format == "tdD") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::DAYS);
-		return LogicalType::DATE;
+		return make_uniq<ArrowType>(LogicalType::DATE, ArrowDateTimeType::DAYS);
 	} else if (format == "tdm") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::MILLISECONDS);
-		return LogicalType::DATE;
+		return make_uniq<ArrowType>(LogicalType::DATE, ArrowDateTimeType::MILLISECONDS);
 	} else if (format == "tts") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::SECONDS);
-		return LogicalType::TIME;
+		return make_uniq<ArrowType>(LogicalType::TIME, ArrowDateTimeType::SECONDS);
 	} else if (format == "ttm") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::MILLISECONDS);
-		return LogicalType::TIME;
+		return make_uniq<ArrowType>(LogicalType::TIME, ArrowDateTimeType::MILLISECONDS);
 	} else if (format == "ttu") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::MICROSECONDS);
-		return LogicalType::TIME;
+		return make_uniq<ArrowType>(LogicalType::TIME, ArrowDateTimeType::MICROSECONDS);
 	} else if (format == "ttn") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::NANOSECONDS);
-		return LogicalType::TIME;
+		return make_uniq<ArrowType>(LogicalType::TIME, ArrowDateTimeType::NANOSECONDS);
 	} else if (format == "tDs") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::SECONDS);
-		return LogicalType::INTERVAL;
+		return make_uniq<ArrowType>(LogicalType::INTERVAL, ArrowDateTimeType::SECONDS);
 	} else if (format == "tDm") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::MILLISECONDS);
-		return LogicalType::INTERVAL;
+		return make_uniq<ArrowType>(LogicalType::INTERVAL, ArrowDateTimeType::MILLISECONDS);
 	} else if (format == "tDu") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::MICROSECONDS);
-		return LogicalType::INTERVAL;
+		return make_uniq<ArrowType>(LogicalType::INTERVAL, ArrowDateTimeType::MICROSECONDS);
 	} else if (format == "tDn") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::NANOSECONDS);
-		return LogicalType::INTERVAL;
+		return make_uniq<ArrowType>(LogicalType::INTERVAL, ArrowDateTimeType::NANOSECONDS);
 	} else if (format == "tiD") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::DAYS);
-		return LogicalType::INTERVAL;
+		return make_uniq<ArrowType>(LogicalType::INTERVAL, ArrowDateTimeType::DAYS);
 	} else if (format == "tiM") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::MONTHS);
-		return LogicalType::INTERVAL;
+		return make_uniq<ArrowType>(LogicalType::INTERVAL, ArrowDateTimeType::MONTHS);
 	} else if (format == "tin") {
-		convert_data.date_time_precision.emplace_back(ArrowDateTimeType::MONTH_DAY_NANO);
-		return LogicalType::INTERVAL;
+		return make_uniq<ArrowType>(LogicalType::INTERVAL, ArrowDateTimeType::MONTH_DAY_NANO);
 	} else if (format == "+l") {
-		convert_data.variable_sz_type.emplace_back(ArrowVariableSizeType::NORMAL, 0);
-		auto child_type = GetArrowLogicalType(*schema.children[0], arrow_convert_data, col_idx);
-		return LogicalType::LIST(child_type);
+		auto child_type = GetArrowLogicalType(*schema.children[0]);
+		auto list_type =
+		    make_uniq<ArrowType>(LogicalType::LIST(child_type->GetDuckType()), ArrowVariableSizeType::NORMAL);
+		list_type->AddChild(std::move(child_type));
+		return list_type;
 	} else if (format == "+L") {
-		convert_data.variable_sz_type.emplace_back(ArrowVariableSizeType::SUPER_SIZE, 0);
-		auto child_type = GetArrowLogicalType(*schema.children[0], arrow_convert_data, col_idx);
-		return LogicalType::LIST(child_type);
+		auto child_type = GetArrowLogicalType(*schema.children[0]);
+		auto list_type =
+		    make_uniq<ArrowType>(LogicalType::LIST(child_type->GetDuckType()), ArrowVariableSizeType::SUPER_SIZE);
+		list_type->AddChild(std::move(child_type));
+		return list_type;
 	} else if (format[0] == '+' && format[1] == 'w') {
 		std::string parameters = format.substr(format.find(':') + 1);
 		idx_t fixed_size = std::stoi(parameters);
-		convert_data.variable_sz_type.emplace_back(ArrowVariableSizeType::FIXED_SIZE, fixed_size);
-		auto child_type = GetArrowLogicalType(*schema.children[0], arrow_convert_data, col_idx);
-		return LogicalType::LIST(child_type);
+		auto child_type = GetArrowLogicalType(*schema.children[0]);
+		auto list_type = make_uniq<ArrowType>(LogicalType::LIST(child_type->GetDuckType()), fixed_size);
+		list_type->AddChild(std::move(child_type));
+		return list_type;
 	} else if (format == "+s") {
 		child_list_t<LogicalType> child_types;
+		vector<unique_ptr<ArrowType>> children;
 		for (idx_t type_idx = 0; type_idx < (idx_t)schema.n_children; type_idx++) {
-			auto child_type = GetArrowLogicalType(*schema.children[type_idx], arrow_convert_data, col_idx);
-			child_types.push_back({schema.children[type_idx]->name, child_type});
+			children.emplace_back(GetArrowLogicalType(*schema.children[type_idx]));
+			child_types.emplace_back(schema.children[type_idx]->name, children.back()->GetDuckType());
+		}
+		auto struct_type = make_uniq<ArrowType>(LogicalType::STRUCT(std::move(child_types)));
+		struct_type->AssignChildren(std::move(children));
+		return struct_type;
+	} else if (format[0] == '+' && format[1] == 'u') {
+		if (format[2] != 's') {
+			throw NotImplementedException("Unsupported Internal Arrow Type: \"%c\" Union", format[2]);
 		}
-		return LogicalType::STRUCT(child_types);
+		D_ASSERT(format[3] == ':');
 
-	} else if (format == "+m") {
-		convert_data.variable_sz_type.emplace_back(ArrowVariableSizeType::NORMAL, 0);
+		std::string prefix = "+us:";
+		// TODO: what are these type ids actually for?
+		auto type_ids = StringUtil::Split(format.substr(prefix.size()), ',');
 
+		child_list_t<LogicalType> members;
+		vector<unique_ptr<ArrowType>> children;
+		for (idx_t type_idx = 0; type_idx < (idx_t)schema.n_children; type_idx++) {
+			auto type = schema.children[type_idx];
+
+			children.emplace_back(GetArrowLogicalType(*type));
+			members.emplace_back(type->name, children.back()->GetDuckType());
+		}
+
+		auto union_type = make_uniq<ArrowType>(LogicalType::UNION(members));
+		union_type->AssignChildren(std::move(children));
+		return union_type;
+	} else if (format == "+m") {
 		auto &arrow_struct_type = *schema.children[0];
 		D_ASSERT(arrow_struct_type.n_children == 2);
-		auto key_type = GetArrowLogicalType(*arrow_struct_type.children[0], arrow_convert_data, col_idx);
-		auto value_type = GetArrowLogicalType(*arrow_struct_type.children[1], arrow_convert_data, col_idx);
-		return LogicalType::MAP(key_type, value_type);
+		auto key_type = GetArrowLogicalType(*arrow_struct_type.children[0]);
+		auto value_type = GetArrowLogicalType(*arrow_struct_type.children[1]);
+		auto map_type = make_uniq<ArrowType>(LogicalType::MAP(key_type->GetDuckType(), value_type->GetDuckType()),
+		                                     ArrowVariableSizeType::NORMAL);
+		child_list_t<LogicalType> key_value;
+		key_value.emplace_back(std::make_pair("key", key_type->GetDuckType()));
+		key_value.emplace_back(std::make_pair("value", value_type->GetDuckType()));
+
+		auto inner_struct =
+		    make_uniq<ArrowType>(LogicalType::STRUCT(std::move(key_value)), ArrowVariableSizeType::NORMAL);
+		vector<unique_ptr<ArrowType>> children;
+		children.reserve(2);
+		children.push_back(std::move(key_type));
+		children.push_back(std::move(value_type));
+		inner_struct->AssignChildren(std::move(children));
+		map_type->AddChild(std::move(inner_struct));
+		return map_type;
 	} else if (format == "z") {
-		convert_data.variable_sz_type.emplace_back(ArrowVariableSizeType::NORMAL, 0);
-		return LogicalType::BLOB;
+		return make_uniq<ArrowType>(LogicalType::BLOB, ArrowVariableSizeType::NORMAL);
 	} else if (format == "Z") {
-		convert_data.variable_sz_type.emplace_back(ArrowVariableSizeType::SUPER_SIZE, 0);
-		return LogicalType::BLOB;
+		return make_uniq<ArrowType>(LogicalType::BLOB, ArrowVariableSizeType::SUPER_SIZE);
 	} else if (format[0] == 'w') {
 		std::string parameters = format.substr(format.find(':') + 1);
 		idx_t fixed_size = std::stoi(parameters);
-		convert_data.variable_sz_type.emplace_back(ArrowVariableSizeType::FIXED_SIZE, fixed_size);
-		return LogicalType::BLOB;
+		return make_uniq<ArrowType>(LogicalType::BLOB, fixed_size);
 	} else if (format[0] == 't' && format[1] == 's') {
 		// Timestamp with Timezone
+		// TODO right now we just get the UTC value. We probably want to support this properly in the future
 		if (format[2] == 'n') {
-			convert_data.date_time_precision.emplace_back(ArrowDateTimeType::NANOSECONDS);
+			return make_uniq<ArrowType>(LogicalType::TIMESTAMP_TZ, ArrowDateTimeType::NANOSECONDS);
 		} else if (format[2] == 'u') {
-			convert_data.date_time_precision.emplace_back(ArrowDateTimeType::MICROSECONDS);
+			return make_uniq<ArrowType>(LogicalType::TIMESTAMP_TZ, ArrowDateTimeType::MICROSECONDS);
 		} else if (format[2] == 'm') {
-			convert_data.date_time_precision.emplace_back(ArrowDateTimeType::MILLISECONDS);
+			return make_uniq<ArrowType>(LogicalType::TIMESTAMP_TZ, ArrowDateTimeType::MILLISECONDS);
 		} else if (format[2] == 's') {
-			convert_data.date_time_precision.emplace_back(ArrowDateTimeType::SECONDS);
+			return make_uniq<ArrowType>(LogicalType::TIMESTAMP_TZ, ArrowDateTimeType::SECONDS);
 		} else {
 			throw NotImplementedException(" Timestamptz precision of not accepted");
 		}
-		// TODO right now we just get the UTC value. We probably want to support this properly in the future
-		return LogicalType::TIMESTAMP_TZ;
 	} else {
 		throw NotImplementedException("Unsupported Internal Arrow Type %s", format);
 	}
@@ -205,13 +223,17 @@ unique_ptr<FunctionData> ArrowTableFunction::ArrowScanBind(ClientContext &contex
 		if (!schema.release) {
 			throw InvalidInputException("arrow_scan: released schema passed");
 		}
+		auto arrow_type = GetArrowLogicalType(schema);
 		if (schema.dictionary) {
-			auto logical_type = GetArrowLogicalType(schema, res->arrow_convert_data, col_idx);
-			res->arrow_convert_data[col_idx] = make_uniq<ArrowConvertData>(std::move(logical_type));
-			return_types.emplace_back(GetArrowLogicalType(*schema.dictionary, res->arrow_convert_data, col_idx));
+			auto logical_type = arrow_type->GetDuckType();
+			auto dictionary = GetArrowLogicalType(*schema.dictionary);
+			return_types.emplace_back(dictionary->GetDuckType());
+			// The dictionary might have different attributes (size type, datetime precision, etc..)
+			arrow_type->SetDictionary(std::move(dictionary));
 		} else {
-			return_types.emplace_back(GetArrowLogicalType(schema, res->arrow_convert_data, col_idx));
+			return_types.emplace_back(arrow_type->GetDuckType());
 		}
+		res->arrow_table.AddColumn(col_idx, std::move(arrow_type));
 		auto format = string(schema.format);
 		auto name = string(schema.name);
 		if (name.empty()) {
@@ -329,11 +351,11 @@ void ArrowTableFunction::ArrowScanFunction(ClientContext &context, TableFunction
 	if (global_state.CanRemoveFilterColumns()) {
 		state.all_columns.Reset();
 		state.all_columns.SetCardinality(output_size);
-		ArrowToDuckDB(state, data.arrow_convert_data, state.all_columns, data.lines_read - output_size);
+		ArrowToDuckDB(state, data.arrow_table.GetColumns(), state.all_columns, data.lines_read - output_size);
 		output.ReferenceColumns(state.all_columns, global_state.projection_ids);
 	} else {
 		output.SetCardinality(output_size);
-		ArrowToDuckDB(state, data.arrow_convert_data, output, data.lines_read - output_size);
+		ArrowToDuckDB(state, data.arrow_table.GetColumns(), output, data.lines_read - output_size);
 	}
 
 	output.Verify();