Skip to content

Commit

Permalink
Format JSON array append
Browse files Browse the repository at this point in the history
  • Loading branch information
@PhictionalOne
PhictionalOne committed Feb 22, 2024
1 parent 8bc9be4 commit df72b16
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 58 deletions.
31 changes: 19 additions & 12 deletions extension/json/include/json_executors.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,25 +125,32 @@ struct JSONExecutors {

//! Two-argument JSON manipulation function
template <class T>
static void BinaryMutExecute(DataChunk &args, ExpressionState &state, Vector &result,
std::function<yyjson_mut_val*(yyjson_mut_val*, yyjson_mut_doc*, T, yyjson_alc*, Vector&)> fun) {
static void BinaryMutExecute(
DataChunk &args, ExpressionState &state, Vector &result,
std::function<yyjson_mut_val *(yyjson_mut_val *, yyjson_mut_doc *, T, yyjson_alc *, Vector &)> fun) {
auto &lstate = JSONFunctionLocalState::ResetAndGet(state);
auto alc = lstate.json_allocator.GetYYAlc();

auto &inputs_left = args.data[0];
auto &inputs_right = args.data[1];
auto &inputs_left = args.data[0];
auto &inputs_right = args.data[1];

BinaryExecutor::Execute<string_t, T, string_t>(
inputs_left, inputs_right, result, args.size(),
[&](string_t left, T right) {
auto ldoc = JSONCommon::ReadDocument(left, JSONCommon::READ_FLAG, alc);
auto mut_ldoc = yyjson_doc_mut_copy(ldoc, alc);
inputs_left, inputs_right, result, args.size(), [&](string_t left, T right) {
// Convert String to immutable yyjson document
auto ldoc = JSONCommon::ReadDocument(left, JSONCommon::READ_FLAG, alc);

auto new_val = fun(mut_ldoc->root, mut_ldoc, right, alc, result);
ldoc = yyjson_mut_val_imut_copy(new_val, alc);
// Convert immutable document into mutable document
auto mut_ldoc = yyjson_doc_mut_copy(ldoc, alc);

return JSONCommon::WriteVal<yyjson_val>(ldoc->root, alc);
});
// Compute mutable value
auto new_val = fun(mut_ldoc->root, mut_ldoc, right, alc, result);

// Convert mutable value back to immutable document
ldoc = yyjson_mut_val_imut_copy(new_val, alc);

// Convert immutable yyjson document back into string
return JSONCommon::WriteVal<yyjson_val>(ldoc->root, alc);
});
}


Expand Down
2 changes: 1 addition & 1 deletion extension/json/json_functions.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -160,7 +160,7 @@ vector<ScalarFunctionSet> JSONFunctions::GetScalarFunctions() {
functions.push_back(GetSerializeSqlFunction());
functions.push_back(GetDeserializeSqlFunction());

functions.push_back(GetArrayAppendFunction());
functions.push_back(GetArrayAppendFunction());
return functions;
}

Expand Down
89 changes: 44 additions & 45 deletions extension/json/json_functions/json_array_append.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,31 +2,30 @@

namespace duckdb {


static void ArrayAppendFunction(DataChunk &args, ExpressionState &state, Vector &result) {
auto left_type = args.data[0].GetType();
auto left_type = args.data[0].GetType();
D_ASSERT(left_type == LogicalType::VARCHAR || left_type == JSONCommon::JSONType());

auto right_type = args.data[1].GetType();
auto right_type = args.data[1].GetType();

// String or JSON value
if (right_type == LogicalType::VARCHAR || right_type == JSONCommon::JSONType()) {
JSONExecutors::BinaryMutExecute<string_t>(
args, state, result,
[](yyjson_mut_val *arr, yyjson_mut_doc *doc, string_t element, yyjson_alc *alc, Vector &result) {
if (!yyjson_mut_is_arr(arr)) {
throw InvalidInputException("JSON input not an JSON Array");
}
if (!yyjson_mut_is_arr(arr)) {
throw InvalidInputException("JSON input not an JSON Array");
}

auto edoc = JSONCommon::ReadDocument(element, JSONCommon::READ_FLAG, alc);
auto mut_edoc = yyjson_doc_mut_copy(edoc, alc);
auto edoc = JSONCommon::ReadDocument(element, JSONCommon::READ_FLAG, alc);
auto mut_edoc = yyjson_doc_mut_copy(edoc, alc);

yyjson_mut_arr_append(arr, mut_edoc->root);
yyjson_mut_arr_append(arr, mut_edoc->root);

return arr;
});
// Boolean
} else if (right_type == LogicalType::BOOLEAN) {
return arr;
});
// Boolean
} else if (right_type == LogicalType::BOOLEAN) {
JSONExecutors::BinaryMutExecute<bool>(
args, state, result,
[&](yyjson_mut_val *arr, yyjson_mut_doc *doc, bool element, yyjson_alc *alc, Vector &result) {
Expand All @@ -37,25 +36,25 @@ static void ArrayAppendFunction(DataChunk &args, ExpressionState &state, Vector
yyjson_mut_arr_add_bool(doc, arr, element);
return arr;
});
// Integer value
} else if (right_type == LogicalType::BIGINT || right_type == LogicalType::UTINYINT ||
right_type == LogicalType::USMALLINT || right_type == LogicalType::UINTEGER ||
right_type == LogicalType::UBIGINT) {
JSONExecutors::BinaryMutExecute<int64_t>(
args, state, result,
[&](yyjson_mut_val *arr, yyjson_mut_doc *doc, int64_t el, yyjson_alc *alc, Vector &result) {
if (!yyjson_mut_is_arr(arr)) {
throw InvalidInputException("JSON input not a JSON Array");
}

int64_t element = static_cast<int64_t>(el);

yyjson_mut_arr_add_int(doc, arr, element);
return arr;
});
// Floating value
} else if (right_type == LogicalType::FLOAT || right_type == LogicalType::DOUBLE) {
JSONExecutors::BinaryMutExecute<double>(
// Integer value
} else if (right_type == LogicalType::BIGINT || right_type == LogicalType::UTINYINT ||
right_type == LogicalType::USMALLINT || right_type == LogicalType::UINTEGER ||
right_type == LogicalType::UBIGINT) {
JSONExecutors::BinaryMutExecute<int64_t>(
args, state, result,
[&](yyjson_mut_val *arr, yyjson_mut_doc *doc, int64_t el, yyjson_alc *alc, Vector &result) {
if (!yyjson_mut_is_arr(arr)) {
throw InvalidInputException("JSON input not a JSON Array");
}

int64_t element = static_cast<int64_t>(el);

yyjson_mut_arr_add_int(doc, arr, element);
return arr;
});
// Floating value
} else if (right_type == LogicalType::FLOAT || right_type == LogicalType::DOUBLE) {
JSONExecutors::BinaryMutExecute<double>(
args, state, result,
[&](yyjson_mut_val *arr, yyjson_mut_doc *doc, double el, yyjson_alc *alc, Vector &result) {
if (!yyjson_mut_is_arr(arr)) {
Expand All @@ -67,7 +66,7 @@ static void ArrayAppendFunction(DataChunk &args, ExpressionState &state, Vector
yyjson_mut_arr_add_real(doc, arr, element);
return arr;
});
}
}
}

static void GetArrayAppendFunctionInternal(ScalarFunctionSet &set, const LogicalType &lhs, const LogicalType &rhs) {
Expand All @@ -76,25 +75,25 @@ static void GetArrayAppendFunctionInternal(ScalarFunctionSet &set, const Logical
}

ScalarFunctionSet JSONFunctions::GetArrayAppendFunction() {
ScalarFunctionSet set("json_array_append");
ScalarFunctionSet set("json_array_append");

// Use different executor for these
// Allows booleans directly
GetArrayAppendFunctionInternal(set, LogicalType::VARCHAR, LogicalType::BOOLEAN);
GetArrayAppendFunctionInternal(set, JSONCommon::JSONType(), LogicalType::BOOLEAN);
GetArrayAppendFunctionInternal(set, LogicalType::VARCHAR, LogicalType::BOOLEAN);
GetArrayAppendFunctionInternal(set, JSONCommon::JSONType(), LogicalType::BOOLEAN);

// Allows for Integer types
// TINYINT, SMALLINT, INTEGER, UTINYINT, USMALLINT, UINTEGER are captured by BIGINT
// relies on consistant casting strategy upfront
GetArrayAppendFunctionInternal(set, LogicalType::VARCHAR, LogicalType::BIGINT);
GetArrayAppendFunctionInternal(set, JSONCommon::JSONType(), LogicalType::BIGINT);
GetArrayAppendFunctionInternal(set, LogicalType::VARCHAR, LogicalType::UBIGINT);
GetArrayAppendFunctionInternal(set, JSONCommon::JSONType(), LogicalType::UBIGINT);
GetArrayAppendFunctionInternal(set, LogicalType::VARCHAR, LogicalType::BIGINT);
GetArrayAppendFunctionInternal(set, JSONCommon::JSONType(), LogicalType::BIGINT);
GetArrayAppendFunctionInternal(set, LogicalType::VARCHAR, LogicalType::UBIGINT);
GetArrayAppendFunctionInternal(set, JSONCommon::JSONType(), LogicalType::UBIGINT);

// Allows for floating types
// FLOAT is covered by automatic upfront casting to double
GetArrayAppendFunctionInternal(set, LogicalType::VARCHAR, LogicalType::DOUBLE);
GetArrayAppendFunctionInternal(set, JSONCommon::JSONType(), LogicalType::DOUBLE);
// Allows for floating types
// FLOAT is covered by automatic upfront casting to double
GetArrayAppendFunctionInternal(set, LogicalType::VARCHAR, LogicalType::DOUBLE);
GetArrayAppendFunctionInternal(set, JSONCommon::JSONType(), LogicalType::DOUBLE);

// Allows for json and string values
GetArrayAppendFunctionInternal(set, LogicalType::VARCHAR, LogicalType::VARCHAR);
Expand All @@ -105,4 +104,4 @@ ScalarFunctionSet JSONFunctions::GetArrayAppendFunction() {
return set;
}

}
} // namespace duckdb

0 comments on commit df72b16

Please sign in to comment.