duckdb 0.7.2-dev12.0 → 0.7.2-dev1244.0

package/src/duckdb/extension/json/include/json_common.hpp

@@ -100,6 +100,7 @@ public:
 	static constexpr auto READ_FLAG = YYJSON_READ_ALLOW_INF_AND_NAN | YYJSON_READ_ALLOW_TRAILING_COMMAS;
 	static constexpr auto STOP_READ_FLAG = READ_FLAG | YYJSON_READ_STOP_WHEN_DONE | YYJSON_READ_INSITU;
 	static constexpr auto WRITE_FLAG = YYJSON_WRITE_ALLOW_INF_AND_NAN;
+	static constexpr auto WRITE_PRETTY_FLAG = YYJSON_WRITE_ALLOW_INF_AND_NAN | YYJSON_WRITE_PRETTY;
 
 public:
 	//! Constant JSON type strings

package/src/duckdb/extension/json/include/json_functions.hpp

@@ -54,6 +54,7 @@ public:
 
 struct JSONFunctionLocalState : public FunctionLocalState {
 public:
+	explicit JSONFunctionLocalState(Allocator &allocator);
 	explicit JSONFunctionLocalState(ClientContext &context);
 	static unique_ptr<FunctionLocalState> Init(ExpressionState &state, const BoundFunctionExpression &expr,
 	                                           FunctionData *bind_data);
@@ -94,6 +95,7 @@ private:
 	static CreateScalarFunctionInfo GetKeysFunction();
 	static CreateScalarFunctionInfo GetTypeFunction();
 	static CreateScalarFunctionInfo GetValidFunction();
+	static CreateScalarFunctionInfo GetSerializeSqlFunction();
 
 	template <class FUNCTION_INFO>
 	static void AddAliases(const vector<string> &names, FUNCTION_INFO fun, vector<FUNCTION_INFO> &functions) {

package/src/duckdb/extension/json/include/json_serializer.hpp

@@ -0,0 +1,77 @@
+#pragma once
+
+#include "json_common.hpp"
+#include "duckdb/common/serializer/format_serializer.hpp"
+
+namespace duckdb {
+
+struct JsonSerializer : FormatSerializer {
+private:
+	yyjson_mut_doc *doc;
+	yyjson_mut_val *current_tag;
+	vector<yyjson_mut_val *> stack;
+
+	// Skip writing property if null
+	bool skip_if_null = false;
+	// Skip writing property if empty string, empty list or empty map.
+	bool skip_if_empty = false;
+
+	// Get the current json value
+	inline yyjson_mut_val *Current() {
+		return stack.back();
+	};
+
+	// Either adds a value to the current object with the current tag, or appends it to the current array
+	void PushValue(yyjson_mut_val *val);
+
+public:
+	explicit JsonSerializer(yyjson_mut_doc *doc, bool skip_if_null, bool skip_if_empty)
+	    : doc(doc), stack({yyjson_mut_obj(doc)}), skip_if_null(skip_if_null), skip_if_empty(skip_if_empty) {
+		serialize_enum_as_string = true;
+	}
+
+	yyjson_mut_val *GetRootObject() {
+		D_ASSERT(stack.size() == 1); // or we forgot to pop somewhere
+		return stack.front();
+	};
+
+	void SetTag(const char *tag) final;
+
+	//===--------------------------------------------------------------------===//
+	// Nested Types Hooks
+	//===--------------------------------------------------------------------===//
+	void OnOptionalBegin(bool present) final;
+	void OnListBegin(idx_t count) final;
+	void OnListEnd(idx_t count) final;
+	void OnMapBegin(idx_t count) final;
+	void OnMapEntryBegin() final;
+	void OnMapEntryEnd() final;
+	void OnMapKeyBegin() final;
+	void OnMapValueBegin() final;
+	void OnMapEnd(idx_t count) final;
+	void OnObjectBegin() final;
+	void OnObjectEnd() final;
+
+	//===--------------------------------------------------------------------===//
+	// Primitive Types
+	//===--------------------------------------------------------------------===//
+	void WriteNull() final;
+	void WriteValue(uint8_t value) final;
+	void WriteValue(int8_t value) final;
+	void WriteValue(uint16_t value) final;
+	void WriteValue(int16_t value) final;
+	void WriteValue(uint32_t value) final;
+	void WriteValue(int32_t value) final;
+	void WriteValue(uint64_t value) final;
+	void WriteValue(int64_t value) final;
+	void WriteValue(hugeint_t value) final;
+	void WriteValue(float value) final;
+	void WriteValue(double value) final;
+	void WriteValue(interval_t value) final;
+	void WriteValue(const string &value) final;
+	void WriteValue(const string_t value) final;
+	void WriteValue(const char *value) final;
+	void WriteValue(bool value) final;
+};
+
+} // namespace duckdb

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

@@ -0,0 +1,147 @@
+#include "json_functions.hpp"
+#include "json_serializer.hpp"
+#include "duckdb/parser/parser.hpp"
+#include "duckdb/execution/expression_executor.hpp"
+
+namespace duckdb {
+
+struct JsonSerializeBindData : public FunctionData {
+	bool skip_if_null = false;
+	bool skip_if_empty = false;
+	bool format = false;
+
+	JsonSerializeBindData(bool skip_if_null_p, bool skip_if_empty_p, bool format_p)
+	    : skip_if_null(skip_if_null_p), skip_if_empty(skip_if_empty_p), format(format_p) {
+	}
+
+public:
+	unique_ptr<FunctionData> Copy() const override {
+		return make_unique<JsonSerializeBindData>(skip_if_null, skip_if_empty, format);
+	}
+	bool Equals(const FunctionData &other_p) const override {
+		return true;
+	}
+};
+
+static unique_ptr<FunctionData> JsonSerializeBind(ClientContext &context, ScalarFunction &bound_function,
+                                                  vector<unique_ptr<Expression>> &arguments) {
+	if (arguments.empty()) {
+		throw BinderException("json_serialize_sql takes at least one argument");
+	}
+
+	if (arguments[0]->return_type != LogicalType::VARCHAR) {
+		throw InvalidTypeException("json_serialize_sql first argument must be a VARCHAR");
+	}
+
+	// Optional arguments
+
+	bool skip_if_null = false;
+	bool skip_if_empty = false;
+	bool format = false;
+
+	for (idx_t i = 1; i < arguments.size(); i++) {
+		auto &arg = arguments[i];
+		if (arg->HasParameter()) {
+			throw ParameterNotResolvedException();
+		}
+		if (!arg->IsFoldable()) {
+			throw InvalidInputException("arguments to json_serialize_sql must be constant");
+		}
+		if (arg->alias == "skip_null") {
+			if (arg->return_type.id() != LogicalTypeId::BOOLEAN) {
+				throw InvalidTypeException("skip_null argument must be a boolean");
+			}
+			skip_if_null = BooleanValue::Get(ExpressionExecutor::EvaluateScalar(context, *arg));
+		} else if (arg->alias == "skip_empty") {
+			if (arg->return_type.id() != LogicalTypeId::BOOLEAN) {
+				throw InvalidTypeException("skip_empty argument must be a boolean");
+			}
+			skip_if_empty = BooleanValue::Get(ExpressionExecutor::EvaluateScalar(context, *arg));
+		} else if (arg->alias == "format") {
+			if (arg->return_type.id() != LogicalTypeId::BOOLEAN) {
+				throw InvalidTypeException("indent argument must be a boolean");
+			}
+			format = BooleanValue::Get(ExpressionExecutor::EvaluateScalar(context, *arg));
+		} else {
+			throw BinderException(StringUtil::Format("Unknown argument to json_serialize_sql: %s", arg->alias.c_str()));
+		}
+	}
+	return make_unique<JsonSerializeBindData>(skip_if_null, skip_if_empty, format);
+}
+
+static void JsonSerializeFunction(DataChunk &args, ExpressionState &state, Vector &result) {
+	auto &local_state = JSONFunctionLocalState::ResetAndGet(state);
+	auto alc = local_state.json_allocator.GetYYJSONAllocator();
+	auto &inputs = args.data[0];
+
+	auto &func_expr = (BoundFunctionExpression &)state.expr;
+	const auto &info = (JsonSerializeBindData &)*func_expr.bind_info;
+
+	UnaryExecutor::Execute<string_t, string_t>(inputs, result, args.size(), [&](string_t input) {
+		auto doc = JSONCommon::CreateDocument(alc);
+		auto result_obj = yyjson_mut_obj(doc);
+		yyjson_mut_doc_set_root(doc, result_obj);
+
+		try {
+			auto parser = Parser();
+			parser.ParseQuery(input.GetString());
+
+			auto statements_arr = yyjson_mut_arr(doc);
+
+			for (auto &statement : parser.statements) {
+				if (statement->type != StatementType::SELECT_STATEMENT) {
+					throw NotImplementedException("Only SELECT statements can be serialized to json!");
+				}
+				auto &select = (SelectStatement &)*statement;
+				auto serializer = JsonSerializer(doc, info.skip_if_null, info.skip_if_empty);
+				select.FormatSerialize(serializer);
+				auto json = serializer.GetRootObject();
+
+				yyjson_mut_arr_append(statements_arr, json);
+			}
+
+			yyjson_mut_obj_add_false(doc, result_obj, "error");
+			yyjson_mut_obj_add_val(doc, result_obj, "statements", statements_arr);
+			idx_t len;
+			auto data = yyjson_mut_val_write_opts(result_obj,
+			                                      info.format ? JSONCommon::WRITE_PRETTY_FLAG : JSONCommon::WRITE_FLAG,
+			                                      alc, (size_t *)&len, nullptr);
+			return StringVector::AddString(result, data, len);
+
+		} catch (Exception &exception) {
+			yyjson_mut_obj_add_true(doc, result_obj, "error");
+			yyjson_mut_obj_add_strcpy(doc, result_obj, "error_type",
+			                          StringUtil::Lower(exception.ExceptionTypeToString(exception.type)).c_str());
+			yyjson_mut_obj_add_strcpy(doc, result_obj, "error_message", exception.RawMessage().c_str());
+
+			idx_t len;
+			auto data = yyjson_mut_val_write_opts(result_obj,
+			                                      info.format ? JSONCommon::WRITE_PRETTY_FLAG : JSONCommon::WRITE_FLAG,
+			                                      alc, (size_t *)&len, nullptr);
+			return StringVector::AddString(result, data, len);
+		}
+	});
+}
+
+CreateScalarFunctionInfo JSONFunctions::GetSerializeSqlFunction() {
+	ScalarFunctionSet set("json_serialize_sql");
+	set.AddFunction(ScalarFunction({LogicalType::VARCHAR}, JSONCommon::JSONType(), JsonSerializeFunction,
+	                               JsonSerializeBind, nullptr, nullptr, JSONFunctionLocalState::Init));
+
+	set.AddFunction(ScalarFunction({LogicalType::VARCHAR, LogicalType::BOOLEAN}, JSONCommon::JSONType(),
+	                               JsonSerializeFunction, JsonSerializeBind, nullptr, nullptr,
+	                               JSONFunctionLocalState::Init));
+
+	set.AddFunction(ScalarFunction({LogicalType::VARCHAR, LogicalType::BOOLEAN, LogicalType::BOOLEAN},
+	                               JSONCommon::JSONType(), JsonSerializeFunction, JsonSerializeBind, nullptr, nullptr,
+	                               JSONFunctionLocalState::Init));
+
+	set.AddFunction(
+	    ScalarFunction({LogicalType::VARCHAR, LogicalType::BOOLEAN, LogicalType::BOOLEAN, LogicalType::BOOLEAN},
+	                   JSONCommon::JSONType(), JsonSerializeFunction, JsonSerializeBind, nullptr, nullptr,
+	                   JSONFunctionLocalState::Init));
+
+	return CreateScalarFunctionInfo(set);
+}
+
+} // namespace duckdb

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

@@ -20,6 +20,12 @@ void JSONScan::AutoDetect(ClientContext &context, JSONScanData &bind_data, vecto
 	idx_t remaining = bind_data.sample_size;
 	while (remaining != 0) {
 		allocator.Reset();
+
+		if (gstate.file_index >= 10) {
+			// We really shouldn't open more than 10 files when sampling
+			break;
+		}
+
 		auto read_count = lstate.ReadNext(gstate);
 		if (lstate.scan_count > 1) {
 			more_than_one = true;
@@ -46,11 +52,6 @@ void JSONScan::AutoDetect(ClientContext &context, JSONScanData &bind_data, vecto
 		node.InitializeCandidateTypes(bind_data.max_depth);
 		node.RefineCandidateTypes(values, next, string_vector, allocator, bind_data.date_format_map);
 		remaining -= next;
-
-		if (gstate.file_index == 10) {
-			// We really shouldn't open more than 10 files when sampling
-			break;
-		}
 	}
 	bind_data.type = original_scan_type;
 

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

@@ -103,7 +103,10 @@ unique_ptr<FunctionData> JSONReadManyFunctionData::Bind(ClientContext &context,
 	return make_unique<JSONReadManyFunctionData>(std::move(paths), std::move(lens));
 }
 
-JSONFunctionLocalState::JSONFunctionLocalState(ClientContext &context) : json_allocator(BufferAllocator::Get(context)) {
+JSONFunctionLocalState::JSONFunctionLocalState(Allocator &allocator) : json_allocator(allocator) {
+}
+JSONFunctionLocalState::JSONFunctionLocalState(ClientContext &context)
+    : JSONFunctionLocalState(BufferAllocator::Get(context)) {
 }
 
 unique_ptr<FunctionLocalState> JSONFunctionLocalState::Init(ExpressionState &state, const BoundFunctionExpression &expr,
@@ -143,6 +146,7 @@ vector<CreateScalarFunctionInfo> JSONFunctions::GetScalarFunctions() {
 	functions.push_back(GetKeysFunction());
 	functions.push_back(GetTypeFunction());
 	functions.push_back(GetValidFunction());
+	functions.push_back(GetSerializeSqlFunction());
 
 	return functions;
 }
@@ -184,8 +188,12 @@ unique_ptr<TableRef> JSONFunctions::ReadJSONReplacement(ClientContext &context,
 	return std::move(table_function);
 }
 
-static unique_ptr<FunctionLocalState> InitJSONCastLocalState(ClientContext &context) {
-	return make_unique<JSONFunctionLocalState>(context);
+static unique_ptr<FunctionLocalState> InitJSONCastLocalState(CastLocalStateParameters &parameters) {
+	if (parameters.context) {
+		return make_unique<JSONFunctionLocalState>(*parameters.context);
+	} else {
+		return make_unique<JSONFunctionLocalState>(Allocator::DefaultAllocator());
+	}
 }
 
 static bool CastVarcharToJSON(Vector &source, Vector &result, idx_t count, CastParameters &parameters) {
@@ -208,9 +216,9 @@ static bool CastVarcharToJSON(Vector &source, Vector &result, idx_t count, CastP
 			    mask.SetInvalid(idx);
 			    success = false;
 		    }
-
 		    return input;
 	    });
+	result.Reinterpret(source);
 	return success;
 }
 

package/src/duckdb/extension/json/json_scan.cpp

@@ -221,7 +221,7 @@ unique_ptr<GlobalTableFunctionState> JSONGlobalTableFunctionState::Init(ClientCo
 		}
 		bind_data.names = std::move(names);
 	}
-	return result;
+	return std::move(result);
 }
 
 idx_t JSONGlobalTableFunctionState::MaxThreads() const {
@@ -262,7 +262,7 @@ unique_ptr<LocalTableFunctionState> JSONLocalTableFunctionState::Init(ExecutionC
 	result->state.transform_options = gstate.state.bind_data.transform_options;
 	result->state.transform_options.date_format_map = &result->state.date_format_map;
 
-	return result;
+	return std::move(result);
 }
 
 idx_t JSONLocalTableFunctionState::GetBatchIndex() const {

package/src/duckdb/extension/json/json_serializer.cpp

@@ -0,0 +1,217 @@
+#include "json_serializer.hpp"
+
+namespace duckdb {
+
+void JsonSerializer::PushValue(yyjson_mut_val *val) {
+	auto current = Current();
+	// Array case, just append the value
+	if (yyjson_mut_is_arr(current)) {
+		yyjson_mut_arr_append(current, val);
+	}
+	// Object case, use the currently set tag.
+	else if (yyjson_mut_is_obj(current)) {
+		yyjson_mut_obj_add(current, current_tag, val);
+	}
+	// Else throw
+	else {
+		throw InternalException("Cannot add value to non-array/object json value");
+	}
+}
+
+void JsonSerializer::SetTag(const char *tag) {
+	current_tag = yyjson_mut_strcpy(doc, tag);
+}
+
+//===--------------------------------------------------------------------===//
+// Nested types
+//===--------------------------------------------------------------------===//
+void JsonSerializer::OnOptionalBegin(bool present) {
+	if (!present && !skip_if_null) {
+		WriteNull();
+	}
+}
+
+void JsonSerializer::OnListBegin(idx_t count) {
+	auto new_value = yyjson_mut_arr(doc);
+	// We always push a value to the stack, we just don't add it as a child to the current value
+	// if skipping empty. Even though it is "unnecessary" to create an empty value just to discard it,
+	// this allows the rest of the code to keep on like normal.
+	if (!(count == 0 && skip_if_empty)) {
+		PushValue(new_value);
+	}
+	stack.push_back(new_value);
+}
+
+void JsonSerializer::OnListEnd(idx_t count) {
+	stack.pop_back();
+}
+
+// Serialize maps as arrays of objects with "key" and "value" properties.
+void JsonSerializer::OnMapBegin(idx_t count) {
+	auto new_value = yyjson_mut_arr(doc);
+	if (!(count == 0 && skip_if_empty)) {
+		PushValue(new_value);
+	}
+	stack.push_back(new_value);
+}
+
+void JsonSerializer::OnMapEntryBegin() {
+	auto new_value = yyjson_mut_obj(doc);
+	stack.push_back(new_value);
+}
+
+void JsonSerializer::OnMapKeyBegin() {
+	SetTag("key");
+}
+
+void JsonSerializer::OnMapValueBegin() {
+	SetTag("value");
+}
+
+void JsonSerializer::OnMapEntryEnd() {
+	stack.pop_back();
+}
+
+void JsonSerializer::OnMapEnd(idx_t count) {
+	stack.pop_back();
+}
+
+void JsonSerializer::OnObjectBegin() {
+	auto new_value = yyjson_mut_obj(doc);
+	PushValue(new_value);
+	stack.push_back(new_value);
+}
+
+void JsonSerializer::OnObjectEnd() {
+	auto obj = Current();
+	auto count = yyjson_mut_obj_size(obj);
+
+	stack.pop_back();
+
+	if (count == 0 && skip_if_empty && !stack.empty()) {
+		// remove obj from parent since it was empty
+		auto parent = Current();
+		if (yyjson_mut_is_arr(parent)) {
+			size_t idx;
+			size_t max;
+			yyjson_mut_val *item;
+			size_t found;
+			yyjson_mut_arr_foreach(parent, idx, max, item) {
+				if (item == obj) {
+					found = idx;
+				}
+			}
+			yyjson_mut_arr_remove(parent, found);
+		} else if (yyjson_mut_is_obj(parent)) {
+			size_t idx;
+			size_t max;
+			yyjson_mut_val *item;
+			yyjson_mut_val *key;
+			const char *found;
+			yyjson_mut_obj_foreach(parent, idx, max, key, item) {
+				if (item == obj) {
+					found = yyjson_mut_get_str(key);
+				}
+			}
+			yyjson_mut_obj_remove_key(parent, found);
+		}
+	}
+}
+
+//===--------------------------------------------------------------------===//
+// Primitive types
+//===--------------------------------------------------------------------===//
+void JsonSerializer::WriteNull() {
+	auto val = yyjson_mut_null(doc);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(uint8_t value) {
+	auto val = yyjson_mut_uint(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(int8_t value) {
+	auto val = yyjson_mut_sint(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(uint16_t value) {
+	auto val = yyjson_mut_uint(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(int16_t value) {
+	auto val = yyjson_mut_sint(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(uint32_t value) {
+	auto val = yyjson_mut_uint(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(int32_t value) {
+	auto val = yyjson_mut_int(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(uint64_t value) {
+	auto val = yyjson_mut_uint(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(int64_t value) {
+	auto val = yyjson_mut_sint(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(hugeint_t value) {
+	throw NotImplementedException("Cannot serialize hugeint_t to json yet!");
+}
+
+void JsonSerializer::WriteValue(float value) {
+	auto val = yyjson_mut_real(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(double value) {
+	auto val = yyjson_mut_real(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(interval_t value) {
+	throw NotImplementedException("Cannot serialize interval_t to json yet!");
+}
+
+void JsonSerializer::WriteValue(const string &value) {
+	if (skip_if_empty && value.empty()) {
+		return;
+	}
+	auto val = yyjson_mut_strcpy(doc, value.c_str());
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(const string_t value) {
+	if (skip_if_empty && value.GetSize() == 0) {
+		return;
+	}
+	auto str = value.GetString();
+	auto val = yyjson_mut_strcpy(doc, str.c_str());
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(const char *value) {
+	if (skip_if_empty && (value == nullptr || value[0] == '\0')) {
+		return;
+	}
+	auto val = yyjson_mut_strcpy(doc, value);
+	PushValue(val);
+}
+
+void JsonSerializer::WriteValue(bool value) {
+	auto val = yyjson_mut_bool(doc, value);
+	PushValue(val);
+}
+
+} // namespace duckdb