60 lines
2.4 KiB
C++
60 lines
2.4 KiB
C++
#include "json_executors.hpp"
|
|
|
|
namespace duckdb {
|
|
|
|
static inline list_entry_t GetJSONKeys(yyjson_val *val, yyjson_alc *, Vector &result, ValidityMask &, idx_t) {
|
|
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));
|
|
}
|
|
|
|
ScalarFunctionSet JSONFunctions::GetKeysFunction() {
|
|
ScalarFunctionSet set("json_keys");
|
|
GetJSONKeysFunctionsInternal(set, LogicalType::VARCHAR);
|
|
GetJSONKeysFunctionsInternal(set, LogicalType::JSON());
|
|
return set;
|
|
}
|
|
|
|
} // namespace duckdb
|