duckdb 0.6.2-dev971.0 → 0.7.1-dev2.0

package/src/duckdb/extension/json/json_functions/json_create.cpp

@@ -1,13 +1,13 @@
+#include "duckdb/planner/expression/bound_parameter_expression.hpp"
 #include "json_common.hpp"
 #include "json_functions.hpp"
-#include "duckdb/planner/expression/bound_parameter_expression.hpp"
 
 namespace duckdb {
 
 struct JSONCreateFunctionData : public FunctionData {
 public:
-	JSONCreateFunctionData(unordered_map<string, unique_ptr<Vector>> const_struct_names)
-	    : const_struct_names(move(const_struct_names)) {
+	explicit JSONCreateFunctionData(unordered_map<string, unique_ptr<Vector>> const_struct_names)
+	    : const_struct_names(std::move(const_struct_names)) {
 	}
 	unique_ptr<FunctionData> Copy() const override {
 		// Have to do this because we can't implicitly copy Vector
@@ -16,7 +16,7 @@ public:
 			// The vectors are const vectors of the key value
 			map_copy[kv.first] = make_unique<Vector>(Value(kv.first));
 		}
-		return make_unique<JSONCreateFunctionData>(move(map_copy));
+		return make_unique<JSONCreateFunctionData>(std::move(map_copy));
 	}
 	bool Equals(const FunctionData &other_p) const override {
 		return true;
@@ -28,10 +28,13 @@ public:
 };
 
 static LogicalType GetJSONType(unordered_map<string, unique_ptr<Vector>> &const_struct_names, const LogicalType &type) {
+	if (JSONCommon::LogicalTypeIsJSON(type)) {
+		return type;
+	}
+
 	switch (type.id()) {
 	// These types can go directly into JSON
 	case LogicalTypeId::SQLNULL:
-	case LogicalTypeId::JSON:
 	case LogicalTypeId::BOOLEAN:
 	case LogicalTypeId::BIGINT:
 	case LogicalTypeId::UBIGINT:
@@ -100,7 +103,7 @@ static unique_ptr<FunctionData> JSONCreateBindParams(ScalarFunction &bound_funct
 			bound_function.arguments.push_back(GetJSONType(const_struct_names, type));
 		}
 	}
-	return make_unique<JSONCreateFunctionData>(move(const_struct_names));
+	return make_unique<JSONCreateFunctionData>(std::move(const_struct_names));
 }
 
 static unique_ptr<FunctionData> JSONObjectBind(ClientContext &context, ScalarFunction &bound_function,
@@ -185,7 +188,7 @@ inline yyjson_mut_val *CreateJSONValue(yyjson_mut_doc *doc, const string_t &valu
 }
 
 inline yyjson_mut_val *CreateJSONValueFromJSON(yyjson_mut_doc *doc, const string_t &value) {
-	auto value_doc = JSONCommon::ReadDocument(value);
+	auto value_doc = JSONCommon::ReadDocument(value, JSONCommon::READ_FLAG, &doc->alc);
 	auto result = yyjson_val_mut_copy(doc, value_doc->root);
 	return result;
 }
@@ -228,12 +231,12 @@ static void TemplatedCreateValues(yyjson_mut_doc *doc, yyjson_mut_val *vals[], V
 	value_v.ToUnifiedFormat(count, value_data);
 	auto values = (T *)value_data.data;
 
-	const auto value_type = value_v.GetType().id();
+	const auto type_is_json = JSONCommon::LogicalTypeIsJSON(value_v.GetType());
 	for (idx_t i = 0; i < count; i++) {
 		idx_t val_idx = value_data.sel->get_index(i);
 		if (!value_data.validity.RowIsValid(val_idx)) {
 			vals[i] = yyjson_mut_null(doc);
-		} else if (value_type == LogicalTypeId::JSON) {
+		} else if (type_is_json) {
 			vals[i] = CreateJSONValueFromJSON(doc, (string_t &)values[val_idx]);
 		} else {
 			vals[i] = CreateJSONValue<T>(doc, values[val_idx]);
@@ -249,8 +252,8 @@ static void CreateValuesStruct(const JSONCreateFunctionData &info, yyjson_mut_do
 		vals[i] = yyjson_mut_obj(doc);
 	}
 	// Initialize re-usable array for the nested values
-	auto nested_vals_ptr = unique_ptr<yyjson_mut_val *[]>(new yyjson_mut_val *[count]);
-	auto nested_vals = nested_vals_ptr.get();
+	auto nested_vals = (yyjson_mut_val **)doc->alc.malloc(doc->alc.ctx, sizeof(yyjson_mut_val *) * count);
+
 	// Add the key/value pairs to the objects
 	auto &entries = StructVector::GetEntries(value_v);
 	for (idx_t entry_i = 0; entry_i < entries.size(); entry_i++) {
@@ -274,14 +277,12 @@ static void CreateValuesMap(const JSONCreateFunctionData &info, yyjson_mut_doc *
 	// Create nested keys
 	auto &map_key_v = MapVector::GetKeys(value_v);
 	auto map_key_count = ListVector::GetListSize(value_v);
-	auto nested_keys_ptr = unique_ptr<yyjson_mut_val *[]>(new yyjson_mut_val *[map_key_count]);
-	auto nested_keys = nested_keys_ptr.get();
+	auto nested_keys = (yyjson_mut_val **)doc->alc.malloc(doc->alc.ctx, sizeof(yyjson_mut_val *) * map_key_count);
 	TemplatedCreateValues<string_t>(doc, nested_keys, map_key_v, map_key_count);
 	// Create nested values
 	auto &map_val_v = MapVector::GetValues(value_v);
 	auto map_val_count = ListVector::GetListSize(value_v);
-	auto nested_vals_ptr = unique_ptr<yyjson_mut_val *[]>(new yyjson_mut_val *[map_val_count]);
-	auto nested_vals = nested_vals_ptr.get();
+	auto nested_vals = (yyjson_mut_val **)doc->alc.malloc(doc->alc.ctx, sizeof(yyjson_mut_val *) * map_val_count);
 	CreateValues(info, doc, nested_vals, map_val_v, map_val_count);
 	// Add the key/value pairs to the objects
 	UnifiedVectorFormat map_data;
@@ -313,8 +314,7 @@ static void CreateValuesUnion(const JSONCreateFunctionData &info, yyjson_mut_doc
 	}
 
 	// Initialize re-usable array for the nested values
-	auto nested_vals_ptr = unique_ptr<yyjson_mut_val *[]>(new yyjson_mut_val *[count]);
-	auto nested_vals = nested_vals_ptr.get();
+	auto nested_vals = (yyjson_mut_val **)doc->alc.malloc(doc->alc.ctx, sizeof(yyjson_mut_val *) * count);
 
 	auto &tag_v = UnionVector::GetTags(value_v);
 	UnifiedVectorFormat tag_data;
@@ -360,8 +360,7 @@ static void CreateValuesList(const JSONCreateFunctionData &info, yyjson_mut_doc
 	// Initialize array for the nested values
 	auto &child_v = ListVector::GetEntry(value_v);
 	auto child_count = ListVector::GetListSize(value_v);
-	auto nested_vals_ptr = unique_ptr<yyjson_mut_val *[]>(new yyjson_mut_val *[child_count]);
-	auto nested_vals = nested_vals_ptr.get();
+	auto nested_vals = (yyjson_mut_val **)doc->alc.malloc(doc->alc.ctx, sizeof(yyjson_mut_val *) * child_count);
 	// Fill nested_vals with list values
 	CreateValues(info, doc, nested_vals, child_v, child_count);
 	// Now we add the values to the appropriate JSON arrays
@@ -401,7 +400,6 @@ static void CreateValues(const JSONCreateFunctionData &info, yyjson_mut_doc *doc
 		TemplatedCreateValues<double>(doc, vals, value_v, count);
 		break;
 	case LogicalTypeId::VARCHAR:
-	case LogicalTypeId::JSON:
 		TemplatedCreateValues<string_t>(doc, vals, value_v, count);
 		break;
 	case LogicalTypeId::STRUCT:
@@ -424,12 +422,15 @@ static void CreateValues(const JSONCreateFunctionData &info, yyjson_mut_doc *doc
 static void ObjectFunction(DataChunk &args, ExpressionState &state, Vector &result) {
 	auto &func_expr = (BoundFunctionExpression &)state.expr;
 	const auto &info = (JSONCreateFunctionData &)*func_expr.bind_info;
+	auto &lstate = JSONFunctionLocalState::ResetAndGet(state);
+	auto alc = lstate.json_allocator.GetYYJSONAllocator();
+
 	// Initialize objects
 	const idx_t count = args.size();
-	auto doc = JSONCommon::CreateDocument();
+	auto doc = JSONCommon::CreateDocument(alc);
 	yyjson_mut_val *objs[STANDARD_VECTOR_SIZE];
 	for (idx_t i = 0; i < count; i++) {
-		objs[i] = yyjson_mut_obj(*doc);
+		objs[i] = yyjson_mut_obj(doc);
 	}
 	// Initialize a re-usable value array
 	yyjson_mut_val *vals[STANDARD_VECTOR_SIZE];
@@ -437,13 +438,12 @@ static void ObjectFunction(DataChunk &args, ExpressionState &state, Vector &resu
 	for (idx_t pair_idx = 0; pair_idx < args.data.size() / 2; pair_idx++) {
 		Vector &key_v = args.data[pair_idx * 2];
 		Vector &value_v = args.data[pair_idx * 2 + 1];
-		CreateKeyValuePairs(info, *doc, objs, vals, key_v, value_v, count);
+		CreateKeyValuePairs(info, doc, objs, vals, key_v, value_v, count);
 	}
 	// Write JSON objects to string
 	auto objects = FlatVector::GetData<string_t>(result);
 	for (idx_t i = 0; i < count; i++) {
-		yyjson_mut_doc_set_root(*doc, objs[i]);
-		objects[i] = JSONCommon::WriteDoc(*doc, result);
+		objects[i] = JSONCommon::WriteVal<yyjson_mut_val>(objs[i], alc);
 	}
 
 	if (args.AllConstant()) {
@@ -454,18 +454,21 @@ static void ObjectFunction(DataChunk &args, ExpressionState &state, Vector &resu
 static void ArrayFunction(DataChunk &args, ExpressionState &state, Vector &result) {
 	auto &func_expr = (BoundFunctionExpression &)state.expr;
 	const auto &info = (JSONCreateFunctionData &)*func_expr.bind_info;
+	auto &lstate = JSONFunctionLocalState::ResetAndGet(state);
+	auto alc = lstate.json_allocator.GetYYJSONAllocator();
+
 	// Initialize arrays
 	const idx_t count = args.size();
-	auto doc = JSONCommon::CreateDocument();
+	auto doc = JSONCommon::CreateDocument(alc);
 	yyjson_mut_val *arrs[STANDARD_VECTOR_SIZE];
 	for (idx_t i = 0; i < count; i++) {
-		arrs[i] = yyjson_mut_arr(*doc);
+		arrs[i] = yyjson_mut_arr(doc);
 	}
 	// Initialize a re-usable value array
 	yyjson_mut_val *vals[STANDARD_VECTOR_SIZE];
 	// Loop through args
 	for (auto &v : args.data) {
-		CreateValues(info, *doc, vals, v, count);
+		CreateValues(info, doc, vals, v, count);
 		for (idx_t i = 0; i < count; i++) {
 			yyjson_mut_arr_append(arrs[i], vals[i]);
 		}
@@ -473,8 +476,7 @@ static void ArrayFunction(DataChunk &args, ExpressionState &state, Vector &resul
 	// Write JSON arrays to string
 	auto objects = FlatVector::GetData<string_t>(result);
 	for (idx_t i = 0; i < count; i++) {
-		yyjson_mut_doc_set_root(*doc, arrs[i]);
-		objects[i] = JSONCommon::WriteDoc(*doc, result);
+		objects[i] = JSONCommon::WriteVal<yyjson_mut_val>(arrs[i], alc);
 	}
 
 	if (args.AllConstant()) {
@@ -485,11 +487,14 @@ static void ArrayFunction(DataChunk &args, ExpressionState &state, Vector &resul
 static void ToJSONFunction(DataChunk &args, ExpressionState &state, Vector &result) {
 	auto &func_expr = (BoundFunctionExpression &)state.expr;
 	const auto &info = (JSONCreateFunctionData &)*func_expr.bind_info;
+	auto &lstate = JSONFunctionLocalState::ResetAndGet(state);
+	auto alc = lstate.json_allocator.GetYYJSONAllocator();
+
 	// Initialize array for values
 	const idx_t count = args.size();
-	auto doc = JSONCommon::CreateDocument();
+	auto doc = JSONCommon::CreateDocument(alc);
 	yyjson_mut_val *vals[STANDARD_VECTOR_SIZE];
-	CreateValues(info, *doc, vals, args.data[0], count);
+	CreateValues(info, doc, vals, args.data[0], count);
 	// Write JSON values to string
 	auto objects = FlatVector::GetData<string_t>(result);
 	auto &result_validity = FlatVector::Validity(result);
@@ -498,8 +503,7 @@ static void ToJSONFunction(DataChunk &args, ExpressionState &state, Vector &resu
 	for (idx_t i = 0; i < count; i++) {
 		idx_t idx = input_data.sel->get_index(i);
 		if (input_data.validity.RowIsValid(idx)) {
-			yyjson_mut_doc_set_root(*doc, vals[i]);
-			objects[i] = JSONCommon::WriteDoc(*doc, result);
+			objects[i] = JSONCommon::WriteVal<yyjson_mut_val>(vals[i], alc);
 		} else {
 			result_validity.SetInvalid(i);
 		}
@@ -511,33 +515,38 @@ static void ToJSONFunction(DataChunk &args, ExpressionState &state, Vector &resu
 }
 
 CreateScalarFunctionInfo JSONFunctions::GetObjectFunction() {
-	auto fun = ScalarFunction("json_object", {}, LogicalType::JSON, ObjectFunction, JSONObjectBind);
+	auto fun = ScalarFunction("json_object", {}, JSONCommon::JSONType(), ObjectFunction, JSONObjectBind, nullptr,
+	                          nullptr, JSONFunctionLocalState::Init);
 	fun.varargs = LogicalType::ANY;
 	fun.null_handling = FunctionNullHandling::SPECIAL_HANDLING;
 	return CreateScalarFunctionInfo(fun);
 }
 
 CreateScalarFunctionInfo JSONFunctions::GetArrayFunction() {
-	auto fun = ScalarFunction("json_array", {}, LogicalType::JSON, ArrayFunction, JSONArrayBind);
+	auto fun = ScalarFunction("json_array", {}, JSONCommon::JSONType(), ArrayFunction, JSONArrayBind, nullptr, nullptr,
+	                          JSONFunctionLocalState::Init);
 	fun.varargs = LogicalType::ANY;
 	fun.null_handling = FunctionNullHandling::SPECIAL_HANDLING;
 	return CreateScalarFunctionInfo(fun);
 }
 
 CreateScalarFunctionInfo JSONFunctions::GetToJSONFunction() {
-	auto fun = ScalarFunction("to_json", {}, LogicalType::JSON, ToJSONFunction, ToJSONBind);
+	auto fun = ScalarFunction("to_json", {}, JSONCommon::JSONType(), ToJSONFunction, ToJSONBind, nullptr, nullptr,
+	                          JSONFunctionLocalState::Init);
 	fun.varargs = LogicalType::ANY;
 	return CreateScalarFunctionInfo(fun);
 }
 
 CreateScalarFunctionInfo JSONFunctions::GetArrayToJSONFunction() {
-	auto fun = ScalarFunction("array_to_json", {}, LogicalType::JSON, ToJSONFunction, ArrayToJSONBind);
+	auto fun = ScalarFunction("array_to_json", {}, JSONCommon::JSONType(), ToJSONFunction, ArrayToJSONBind, nullptr,
+	                          nullptr, JSONFunctionLocalState::Init);
 	fun.varargs = LogicalType::ANY;
 	return CreateScalarFunctionInfo(fun);
 }
 
 CreateScalarFunctionInfo JSONFunctions::GetRowToJSONFunction() {
-	auto fun = ScalarFunction("row_to_json", {}, LogicalType::JSON, ToJSONFunction, RowToJSONBind);
+	auto fun = ScalarFunction("row_to_json", {}, JSONCommon::JSONType(), ToJSONFunction, RowToJSONBind, nullptr,
+	                          nullptr, JSONFunctionLocalState::Init);
 	fun.varargs = LogicalType::ANY;
 	return CreateScalarFunctionInfo(fun);
 }

package/src/duckdb/extension/json/json_functions/json_extract.cpp

@@ -1,53 +1,62 @@
-#include "json_common.hpp"
-#include "json_functions.hpp"
+#include "json_executors.hpp"
 
 namespace duckdb {
 
-static inline string_t ExtractFromVal(yyjson_val *val, Vector &result) {
-	return JSONCommon::WriteVal(val, result);
+static inline string_t ExtractFromVal(yyjson_val *val, yyjson_alc *alc, Vector &result) {
+	return JSONCommon::WriteVal<yyjson_val>(val, alc);
 }
 
-static inline string_t ExtractStringFromVal(yyjson_val *val, Vector &result) {
-	return yyjson_is_str(val) ? StringVector::AddString(result, unsafe_yyjson_get_str(val), unsafe_yyjson_get_len(val))
-	                          : JSONCommon::WriteVal(val, result);
+static inline string_t ExtractStringFromVal(yyjson_val *val, yyjson_alc *alc, Vector &result) {
+	return yyjson_is_str(val) ? string_t(unsafe_yyjson_get_str(val), unsafe_yyjson_get_len(val))
+	                          : JSONCommon::WriteVal<yyjson_val>(val, alc);
 }
 
 static void ExtractFunction(DataChunk &args, ExpressionState &state, Vector &result) {
-	JSONCommon::BinaryExecute<string_t>(args, state, result, ExtractFromVal);
+	JSONExecutors::BinaryExecute<string_t>(args, state, result, ExtractFromVal);
 }
 
 static void ExtractManyFunction(DataChunk &args, ExpressionState &state, Vector &result) {
-	JSONCommon::ExecuteMany<string_t>(args, state, result, ExtractFromVal);
+	JSONExecutors::ExecuteMany<string_t>(args, state, result, ExtractFromVal);
 }
 
 static void ExtractStringFunction(DataChunk &args, ExpressionState &state, Vector &result) {
-	JSONCommon::BinaryExecute<string_t>(args, state, result, ExtractStringFromVal);
+	JSONExecutors::BinaryExecute<string_t>(args, state, result, ExtractStringFromVal);
 }
 
 static void ExtractStringManyFunction(DataChunk &args, ExpressionState &state, Vector &result) {
-	JSONCommon::ExecuteMany<string_t>(args, state, result, ExtractStringFromVal);
+	JSONExecutors::ExecuteMany<string_t>(args, state, result, ExtractStringFromVal);
+}
+
+static void GetExtractFunctionsInternal(ScalarFunctionSet &set, const LogicalType &input_type) {
+	set.AddFunction(ScalarFunction({input_type, LogicalType::VARCHAR}, JSONCommon::JSONType(), ExtractFunction,
+	                               JSONReadFunctionData::Bind, nullptr, nullptr, JSONFunctionLocalState::Init));
+	set.AddFunction(ScalarFunction({input_type, LogicalType::LIST(LogicalType::VARCHAR)},
+	                               LogicalType::LIST(JSONCommon::JSONType()), ExtractManyFunction,
+	                               JSONReadManyFunctionData::Bind, nullptr, nullptr, JSONFunctionLocalState::Init));
 }
 
 CreateScalarFunctionInfo JSONFunctions::GetExtractFunction() {
 	// Generic extract function
 	ScalarFunctionSet set("json_extract");
-	set.AddFunction(ScalarFunction({LogicalType::JSON, LogicalType::VARCHAR}, LogicalType::JSON, ExtractFunction,
-	                               JSONReadFunctionData::Bind));
-	set.AddFunction(ScalarFunction({LogicalType::JSON, LogicalType::LIST(LogicalType::VARCHAR)},
-	                               LogicalType::LIST(LogicalType::JSON), ExtractManyFunction,
-	                               JSONReadManyFunctionData::Bind));
+	GetExtractFunctionsInternal(set, LogicalType::VARCHAR);
+	GetExtractFunctionsInternal(set, JSONCommon::JSONType());
 
 	return CreateScalarFunctionInfo(set);
 }
 
+static void GetExtractStringFunctionsInternal(ScalarFunctionSet &set, const LogicalType &input_type) {
+	set.AddFunction(ScalarFunction({input_type, LogicalType::VARCHAR}, LogicalType::VARCHAR, ExtractStringFunction,
+	                               JSONReadFunctionData::Bind, nullptr, nullptr, JSONFunctionLocalState::Init));
+	set.AddFunction(ScalarFunction({input_type, LogicalType::LIST(LogicalType::VARCHAR)},
+	                               LogicalType::LIST(LogicalType::VARCHAR), ExtractStringManyFunction,
+	                               JSONReadManyFunctionData::Bind, nullptr, nullptr, JSONFunctionLocalState::Init));
+}
+
 CreateScalarFunctionInfo JSONFunctions::GetExtractStringFunction() {
 	// String extract function
 	ScalarFunctionSet set("json_extract_string");
-	set.AddFunction(ScalarFunction({LogicalType::JSON, LogicalType::VARCHAR}, LogicalType::VARCHAR,
-	                               ExtractStringFunction, JSONReadFunctionData::Bind));
-	set.AddFunction(ScalarFunction({LogicalType::JSON, LogicalType::LIST(LogicalType::VARCHAR)},
-	                               LogicalType::LIST(LogicalType::VARCHAR), ExtractStringManyFunction,
-	                               JSONReadManyFunctionData::Bind));
+	GetExtractStringFunctionsInternal(set, LogicalType::VARCHAR);
+	GetExtractStringFunctionsInternal(set, JSONCommon::JSONType());
 
 	return CreateScalarFunctionInfo(set);
 }

package/src/duckdb/extension/json/json_functions/json_keys.cpp

@@ -0,0 +1,60 @@
+#include "json_executors.hpp"
+
+namespace duckdb {
+
+static inline list_entry_t GetJSONKeys(yyjson_val *val, yyjson_alc *alc, Vector &result) {
+	auto num_keys = yyjson_obj_size(val);
+	auto current_size = ListVector::GetListSize(result);
+	auto new_size = current_size + num_keys;
+
+	// Grow list if needed
+	if (ListVector::GetListCapacity(result) < new_size) {
+		ListVector::Reserve(result, new_size);
+	}
+
+	// Write the strings to the child vector
+	auto keys = FlatVector::GetData<string_t>(ListVector::GetEntry(result));
+	size_t idx, max;
+	yyjson_val *key, *child_val;
+	yyjson_obj_foreach(val, idx, max, key, child_val) {
+		keys[current_size + idx] = string_t(unsafe_yyjson_get_str(key), unsafe_yyjson_get_len(key));
+	}
+
+	// Update size
+	ListVector::SetListSize(result, current_size + num_keys);
+
+	return {current_size, num_keys};
+}
+
+static void UnaryJSONKeysFunction(DataChunk &args, ExpressionState &state, Vector &result) {
+	JSONExecutors::UnaryExecute<list_entry_t>(args, state, result, GetJSONKeys);
+}
+
+static void BinaryJSONKeysFunction(DataChunk &args, ExpressionState &state, Vector &result) {
+	JSONExecutors::BinaryExecute<list_entry_t>(args, state, result, GetJSONKeys);
+}
+
+static void ManyJSONKeysFunction(DataChunk &args, ExpressionState &state, Vector &result) {
+	JSONExecutors::ExecuteMany<list_entry_t>(args, state, result, GetJSONKeys);
+}
+
+static void GetJSONKeysFunctionsInternal(ScalarFunctionSet &set, const LogicalType &input_type) {
+	set.AddFunction(ScalarFunction({input_type}, LogicalType::LIST(LogicalType::VARCHAR), UnaryJSONKeysFunction,
+	                               nullptr, nullptr, nullptr, JSONFunctionLocalState::Init));
+	set.AddFunction(ScalarFunction({input_type, LogicalType::VARCHAR}, LogicalType::LIST(LogicalType::VARCHAR),
+	                               BinaryJSONKeysFunction, JSONReadFunctionData::Bind, nullptr, nullptr,
+	                               JSONFunctionLocalState::Init));
+	set.AddFunction(ScalarFunction({input_type, LogicalType::LIST(LogicalType::VARCHAR)},
+	                               LogicalType::LIST(LogicalType::LIST(LogicalType::VARCHAR)), ManyJSONKeysFunction,
+	                               JSONReadManyFunctionData::Bind, nullptr, nullptr, JSONFunctionLocalState::Init));
+}
+
+CreateScalarFunctionInfo JSONFunctions::GetKeysFunction() {
+	ScalarFunctionSet set("json_keys");
+	GetJSONKeysFunctionsInternal(set, LogicalType::VARCHAR);
+	GetJSONKeysFunctionsInternal(set, JSONCommon::JSONType());
+
+	return CreateScalarFunctionInfo(std::move(set));
+}
+
+} // namespace duckdb

package/src/duckdb/extension/json/json_functions/json_merge_patch.cpp

@@ -3,6 +3,24 @@
 
 namespace duckdb {
 
+static unique_ptr<FunctionData> JSONMergePatchBind(ClientContext &context, ScalarFunction &bound_function,
+                                                   vector<unique_ptr<Expression>> &arguments) {
+	if (arguments.size() < 2) {
+		throw InvalidInputException("json_merge_patch requires at least two parameters");
+	}
+	bound_function.arguments.reserve(arguments.size());
+	for (auto &arg : arguments) {
+		const auto &arg_type = arg->return_type;
+		if (arg_type == LogicalTypeId::SQLNULL || arg_type == LogicalType::VARCHAR ||
+		    JSONCommon::LogicalTypeIsJSON(arg_type)) {
+			bound_function.arguments.push_back(arg_type);
+		} else {
+			throw InvalidInputException("Arguments to json_merge_patch must be of type VARCHAR or JSON");
+		}
+	}
+	return nullptr;
+}
+
 static inline yyjson_mut_val *MergePatch(yyjson_mut_doc *doc, yyjson_mut_val *orig, yyjson_mut_val *patch) {
 	if ((yyjson_mut_get_tag(orig) != (YYJSON_TYPE_OBJ | YYJSON_SUBTYPE_NONE)) ||
 	    (yyjson_mut_get_tag(patch) != (YYJSON_TYPE_OBJ | YYJSON_SUBTYPE_NONE))) {
@@ -26,24 +44,28 @@ static inline void ReadObjects(yyjson_mut_doc *doc, Vector &input, yyjson_mut_va
 		if (!input_data.validity.RowIsValid(idx)) {
 			objs[i] = nullptr;
 		} else {
-			objs[i] = yyjson_val_mut_copy(doc, JSONCommon::ReadDocument(inputs[idx])->root);
+			objs[i] =
+			    yyjson_val_mut_copy(doc, JSONCommon::ReadDocument(inputs[idx], JSONCommon::READ_FLAG, &doc->alc)->root);
 		}
 	}
 }
 
 //! Follows MySQL behaviour
 static void MergePatchFunction(DataChunk &args, ExpressionState &state, Vector &result) {
-	auto doc = JSONCommon::CreateDocument();
+	auto &lstate = JSONFunctionLocalState::ResetAndGet(state);
+	auto alc = lstate.json_allocator.GetYYJSONAllocator();
+
+	auto doc = JSONCommon::CreateDocument(lstate.json_allocator.GetYYJSONAllocator());
 	const auto count = args.size();
 
 	// Read the first json arg
 	yyjson_mut_val *origs[STANDARD_VECTOR_SIZE];
-	ReadObjects(*doc, args.data[0], origs, count);
+	ReadObjects(doc, args.data[0], origs, count);
 
 	// Read the next json args one by one and merge them into the first json arg
 	yyjson_mut_val *patches[STANDARD_VECTOR_SIZE];
 	for (idx_t arg_idx = 1; arg_idx < args.data.size(); arg_idx++) {
-		ReadObjects(*doc, args.data[arg_idx], patches, count);
+		ReadObjects(doc, args.data[arg_idx], patches, count);
 		for (idx_t i = 0; i < count; i++) {
 			if (patches[i] == nullptr) {
 				// Next json arg is NULL, obj becomes NULL
@@ -53,7 +75,7 @@ static void MergePatchFunction(DataChunk &args, ExpressionState &state, Vector &
 				origs[i] = patches[i];
 			} else {
 				// Neither is NULL, merge them
-				origs[i] = MergePatch(*doc, origs[i], patches[i]);
+				origs[i] = MergePatch(doc, origs[i], patches[i]);
 			}
 		}
 	}
@@ -65,7 +87,7 @@ static void MergePatchFunction(DataChunk &args, ExpressionState &state, Vector &
 		if (origs[i] == nullptr) {
 			result_validity.SetInvalid(i);
 		} else {
-			result_data[i] = JSONCommon::WriteVal(origs[i], result);
+			result_data[i] = JSONCommon::WriteVal<yyjson_mut_val>(origs[i], alc);
 		}
 	}
 
@@ -75,12 +97,12 @@ static void MergePatchFunction(DataChunk &args, ExpressionState &state, Vector &
 }
 
 CreateScalarFunctionInfo JSONFunctions::GetMergePatchFunction() {
-	// Needs at least two json inputs, but supports merging vararg json inputs
-	ScalarFunction fun("json_merge_patch", {LogicalType::JSON, LogicalType::JSON}, LogicalType::JSON,
-	                   MergePatchFunction);
-	fun.varargs = LogicalType::JSON;
+	ScalarFunction fun("json_merge_patch", {}, JSONCommon::JSONType(), MergePatchFunction, JSONMergePatchBind, nullptr,
+	                   nullptr, JSONFunctionLocalState::Init);
+	fun.varargs = LogicalType::ANY;
 	fun.null_handling = FunctionNullHandling::SPECIAL_HANDLING;
-	return CreateScalarFunctionInfo(move(fun));
+
+	return CreateScalarFunctionInfo(fun);
 }
 
 } // namespace duckdb