be/src/exprs/function/function_variant_element.cpp
Line | Count | Source |
1 | | // Licensed to the Apache Software Foundation (ASF) under one |
2 | | // or more contributor license agreements. See the NOTICE file |
3 | | // distributed with this work for additional information |
4 | | // regarding copyright ownership. The ASF licenses this file |
5 | | // to you under the Apache License, Version 2.0 (the |
6 | | // "License"); you may not use this file except in compliance |
7 | | // with the License. You may obtain a copy of the License at |
8 | | // |
9 | | // http://www.apache.org/licenses/LICENSE-2.0 |
10 | | // |
11 | | // Unless required by applicable law or agreed to in writing, |
12 | | // software distributed under the License is distributed on an |
13 | | // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
14 | | // KIND, either express or implied. See the License for the |
15 | | // specific language governing permissions and limitations |
16 | | // under the License. |
17 | | |
18 | | #include <glog/logging.h> |
19 | | #include <stddef.h> |
20 | | |
21 | | #include <memory> |
22 | | #include <ostream> |
23 | | #include <string> |
24 | | #include <string_view> |
25 | | #include <utility> |
26 | | #include <vector> |
27 | | |
28 | | #include "common/status.h" |
29 | | #include "core/assert_cast.h" |
30 | | #include "core/block/block.h" |
31 | | #include "core/column/column.h" |
32 | | #include "core/column/column_nullable.h" |
33 | | #include "core/column/column_string.h" |
34 | | #include "core/column/column_variant.h" |
35 | | #include "core/column/subcolumn_tree.h" |
36 | | #include "core/data_type/data_type.h" |
37 | | #include "core/data_type/data_type_nothing.h" |
38 | | #include "core/data_type/data_type_nullable.h" |
39 | | #include "core/data_type/data_type_string.h" |
40 | | #include "core/data_type/data_type_variant.h" |
41 | | #include "core/string_ref.h" |
42 | | #include "exprs/function/function.h" |
43 | | #include "exprs/function/function_helpers.h" |
44 | | #include "exprs/function/simple_function_factory.h" |
45 | | #include "exprs/json_functions.h" |
46 | | #include "simdjson.h" |
47 | | #include "util/defer_op.h" |
48 | | #include "util/json/path_in_data.h" |
49 | | |
50 | | namespace doris { |
51 | | |
52 | | class FunctionVariantElement : public IFunction { |
53 | | public: |
54 | | static constexpr auto name = "element_at"; |
55 | 2 | static FunctionPtr create() { return std::make_shared<FunctionVariantElement>(); } |
56 | | |
57 | | // Get function name. |
58 | 1 | String get_name() const override { return name; } |
59 | | |
60 | 0 | bool use_default_implementation_for_nulls() const override { return false; } |
61 | | |
62 | 0 | size_t get_number_of_arguments() const override { return 2; } |
63 | | |
64 | 0 | ColumnNumbers get_arguments_that_are_always_constant() const override { return {1}; } |
65 | | |
66 | 1 | DataTypes get_variadic_argument_types_impl() const override { |
67 | 1 | return {std::make_shared<DataTypeVariant>(), std::make_shared<DataTypeString>()}; |
68 | 1 | } |
69 | | |
70 | 0 | DataTypePtr get_return_type_impl(const DataTypes& arguments) const override { |
71 | 0 | DCHECK_EQ(arguments[0]->get_primitive_type(), TYPE_VARIANT) |
72 | 0 | << "First argument for function: " << name |
73 | 0 | << " should be DataTypeVariant but it has type " << arguments[0]->get_name() << "."; |
74 | 0 | DCHECK(is_string_type(arguments[1]->get_primitive_type())) |
75 | 0 | << "Second argument for function: " << name << " should be String but it has type " |
76 | 0 | << arguments[1]->get_name() << "."; |
77 | 0 | auto arg_variant = remove_nullable(arguments[0]); |
78 | 0 | const auto& data_type_object = assert_cast<const DataTypeVariant&>(*arg_variant); |
79 | 0 | return make_nullable( |
80 | 0 | std::make_shared<DataTypeVariant>(data_type_object.variant_max_subcolumns_count(), |
81 | 0 | data_type_object.enable_doc_mode())); |
82 | 0 | } |
83 | | |
84 | | // wrap variant column with nullable |
85 | | // 1. if variant is null root(empty or nothing as root), then nullable map is all null |
86 | | // 2. if variant is scalar variant, then use the root's nullable map |
87 | | // 3. if variant is hierarchical variant, then create a nullable map with all none null |
88 | 0 | ColumnPtr wrap_variant_nullable(ColumnPtr col) const { |
89 | 0 | const auto& var = assert_cast<const ColumnVariant&>(*col); |
90 | 0 | if (var.is_null_root()) { |
91 | 0 | return make_nullable(col, true); |
92 | 0 | } |
93 | 0 | if (var.is_scalar_variant() && var.get_root()->is_nullable()) { |
94 | 0 | const auto* nullable = assert_cast<const ColumnNullable*>(var.get_root().get()); |
95 | 0 | return ColumnNullable::create( |
96 | 0 | col, nullable->get_null_map_column_ptr()->clone_resized(col->size())); |
97 | 0 | } |
98 | 0 | return make_nullable(col); |
99 | 0 | } |
100 | | |
101 | | Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments, |
102 | 0 | uint32_t result, size_t input_rows_count) const override { |
103 | 0 | const auto* variant_col = check_and_get_column<ColumnVariant>( |
104 | 0 | remove_nullable(block.get_by_position(arguments[0]).column).get()); |
105 | 0 | if (!variant_col) { |
106 | 0 | return Status::RuntimeError( |
107 | 0 | fmt::format("unsupported types for function {}({}, {})", get_name(), |
108 | 0 | block.get_by_position(arguments[0]).type->get_name(), |
109 | 0 | block.get_by_position(arguments[1]).type->get_name())); |
110 | 0 | } |
111 | 0 | if (block.empty()) { |
112 | 0 | block.replace_by_position(result, block.get_by_position(result).type->create_column()); |
113 | 0 | return Status::OK(); |
114 | 0 | } |
115 | | |
116 | 0 | auto index_column = block.get_by_position(arguments[1]).column; |
117 | 0 | ColumnPtr result_column; |
118 | 0 | RETURN_IF_ERROR(get_element_column(*variant_col, index_column, &result_column)); |
119 | 0 | if (block.get_by_position(result).type->is_nullable()) { |
120 | 0 | result_column = wrap_variant_nullable(result_column); |
121 | 0 | } |
122 | 0 | block.replace_by_position(result, result_column); |
123 | 0 | return Status::OK(); |
124 | 0 | } |
125 | | |
126 | | private: |
127 | | // Return sub-path by specified prefix. |
128 | | // For example, for prefix a.b: |
129 | | // a.b.c.d -> c.d, a.b.c -> c |
130 | | static std::optional<std::string_view> get_sub_path(const std::string_view& path, |
131 | 3 | const std::string_view& prefix) { |
132 | 3 | if (path.size() <= prefix.size() || path[prefix.size()] != '.') { |
133 | 1 | return std::nullopt; |
134 | 1 | } |
135 | 2 | return path.substr(prefix.size() + 1); |
136 | 3 | } |
137 | | |
138 | | // Extract and populate sparse column data with given path prefix |
139 | | // Copies data from source sparse column, extracting only the sub-paths that match the prefix |
140 | | static void _extract_sparse_column_from_source(ColumnVariant* src_ptr, const PathInData& path, |
141 | 1 | ColumnVariant::MutablePtr& target_ptr) { |
142 | 1 | ColumnVariant::Subcolumn root {0, true, true}; |
143 | | // no root, no sparse column |
144 | 1 | const auto& sparse_data_map = assert_cast<const ColumnMap&>(*src_ptr->get_sparse_column()); |
145 | 1 | const auto& src_sparse_data_offsets = sparse_data_map.get_offsets(); |
146 | 1 | const auto& src_sparse_data_paths = |
147 | 1 | assert_cast<const ColumnString&>(sparse_data_map.get_keys()); |
148 | 1 | const auto& src_sparse_data_values = |
149 | 1 | assert_cast<const ColumnString&>(sparse_data_map.get_values()); |
150 | 1 | auto& sparse_data_offsets = |
151 | 1 | assert_cast<ColumnMap&>(*target_ptr->get_sparse_column()->assume_mutable()) |
152 | 1 | .get_offsets(); |
153 | 1 | auto [sparse_data_paths, sparse_data_values] = |
154 | 1 | target_ptr->get_sparse_data_paths_and_values(); |
155 | 1 | StringRef prefix_ref(path.get_path()); |
156 | 1 | std::string_view path_prefix(prefix_ref.data, prefix_ref.size); |
157 | 2 | for (size_t i = 0; i != src_sparse_data_offsets.size(); ++i) { |
158 | 1 | size_t start = src_sparse_data_offsets[ssize_t(i) - 1]; |
159 | 1 | size_t end = src_sparse_data_offsets[ssize_t(i)]; |
160 | 1 | size_t lower_bound_index = ColumnVariant::find_path_lower_bound_in_sparse_data( |
161 | 1 | prefix_ref, src_sparse_data_paths, start, end); |
162 | 4 | for (; lower_bound_index != end; ++lower_bound_index) { |
163 | 3 | auto path_ref = src_sparse_data_paths.get_data_at(lower_bound_index); |
164 | 3 | std::string_view nested_path(path_ref.data, path_ref.size); |
165 | 3 | if (!nested_path.starts_with(path_prefix)) { |
166 | 0 | break; |
167 | 0 | } |
168 | | // Don't include path that is equal to the prefix. |
169 | 3 | if (nested_path.size() != path_prefix.size()) { |
170 | 3 | auto sub_path_optional = get_sub_path(nested_path, path_prefix); |
171 | 3 | if (!sub_path_optional.has_value()) { |
172 | 1 | continue; |
173 | 1 | } |
174 | 2 | std::string_view sub_path = *sub_path_optional; |
175 | 2 | sparse_data_paths->insert_data(sub_path.data(), sub_path.size()); |
176 | 2 | sparse_data_values->insert_from(src_sparse_data_values, lower_bound_index); |
177 | 2 | } else { |
178 | | // insert into root column, example: access v['b'] and b is in sparse column |
179 | | // data example: |
180 | | // {"b" : 123} |
181 | | // {"b" : {"c" : 456}} |
182 | | // b maybe in sparse column, and b.c is in subolumn, put `b` into root column to distinguish |
183 | | // from "" which is empty path and root |
184 | 0 | root.deserialize_from_binary_column(&src_sparse_data_values, lower_bound_index); |
185 | 0 | } |
186 | 3 | } |
187 | 1 | if (root.size() == sparse_data_offsets.size()) { |
188 | 1 | root.insert_default(); |
189 | 1 | } |
190 | 1 | sparse_data_offsets.push_back(sparse_data_paths->size()); |
191 | 1 | } |
192 | 1 | target_ptr->get_subcolumns().create_root(root); |
193 | 1 | target_ptr->get_doc_value_column()->assume_mutable()->resize(src_ptr->size()); |
194 | 1 | target_ptr->set_num_rows(src_ptr->size()); |
195 | 1 | } |
196 | | |
197 | | // Extract and populate sparse column data from doc_value column with given path prefix |
198 | | // Copies data from source doc_value column, extracting only the sub-paths that match the prefix |
199 | | static void _extract_doc_value_column_from_source(ColumnVariant* src_ptr, |
200 | | const PathInData& path, |
201 | 0 | ColumnVariant::MutablePtr& target_ptr) { |
202 | 0 | ColumnVariant::Subcolumn root {0, true, true}; |
203 | 0 | const auto& doc_value_data_map = |
204 | 0 | assert_cast<const ColumnMap&>(*src_ptr->get_doc_value_column()); |
205 | 0 | const auto& src_doc_value_data_offsets = doc_value_data_map.get_offsets(); |
206 | 0 | const auto& src_doc_value_data_paths = |
207 | 0 | assert_cast<const ColumnString&>(doc_value_data_map.get_keys()); |
208 | 0 | const auto& src_doc_value_data_values = |
209 | 0 | assert_cast<const ColumnString&>(doc_value_data_map.get_values()); |
210 | | // Write extracted data into target's doc_value column (not sparse) to preserve |
211 | | // doc mode invariant: doc_mode columns must not have sparse data. |
212 | 0 | auto& doc_value_offsets = |
213 | 0 | assert_cast<ColumnMap&>(*target_ptr->get_doc_value_column()->assume_mutable()) |
214 | 0 | .get_offsets(); |
215 | 0 | auto [doc_value_paths, doc_value_values] = |
216 | 0 | target_ptr->get_doc_value_data_paths_and_values(); |
217 | 0 | StringRef prefix_ref(path.get_path()); |
218 | 0 | std::string_view path_prefix(prefix_ref.data, prefix_ref.size); |
219 | 0 | for (size_t i = 0; i != src_doc_value_data_offsets.size(); ++i) { |
220 | 0 | size_t start = src_doc_value_data_offsets[ssize_t(i) - 1]; |
221 | 0 | size_t end = src_doc_value_data_offsets[ssize_t(i)]; |
222 | 0 | size_t lower_bound_index = ColumnVariant::find_path_lower_bound_in_sparse_data( |
223 | 0 | prefix_ref, src_doc_value_data_paths, start, end); |
224 | 0 | for (; lower_bound_index != end; ++lower_bound_index) { |
225 | 0 | auto path_ref = src_doc_value_data_paths.get_data_at(lower_bound_index); |
226 | 0 | std::string_view nested_path(path_ref.data, path_ref.size); |
227 | 0 | if (!nested_path.starts_with(path_prefix)) { |
228 | 0 | break; |
229 | 0 | } |
230 | 0 | if (nested_path.size() != path_prefix.size()) { |
231 | 0 | auto sub_path_optional = get_sub_path(nested_path, path_prefix); |
232 | 0 | if (!sub_path_optional.has_value()) { |
233 | 0 | continue; |
234 | 0 | } |
235 | 0 | std::string_view sub_path = *sub_path_optional; |
236 | 0 | doc_value_paths->insert_data(sub_path.data(), sub_path.size()); |
237 | 0 | doc_value_values->insert_from(src_doc_value_data_values, lower_bound_index); |
238 | 0 | } else { |
239 | 0 | root.deserialize_from_binary_column(&src_doc_value_data_values, |
240 | 0 | lower_bound_index); |
241 | 0 | } |
242 | 0 | } |
243 | 0 | if (root.size() == doc_value_offsets.size()) { |
244 | 0 | root.insert_default(); |
245 | 0 | } |
246 | 0 | doc_value_offsets.push_back(doc_value_paths->size()); |
247 | 0 | } |
248 | 0 | target_ptr->get_subcolumns().create_root(root); |
249 | 0 | target_ptr->get_sparse_column()->assume_mutable()->resize(src_ptr->size()); |
250 | 0 | target_ptr->set_num_rows(src_ptr->size()); |
251 | 0 | } |
252 | | |
253 | | static Status get_element_column(const ColumnVariant& src, const ColumnPtr& index_column, |
254 | 1 | ColumnPtr* result) { |
255 | 1 | std::string field_name = index_column->get_data_at(0).to_string(); |
256 | 1 | if (src.empty()) { |
257 | 0 | *result = ColumnVariant::create(src.max_subcolumns_count(), src.enable_doc_mode()); |
258 | | // src subcolumns empty but src row count may not be 0 |
259 | 0 | (*result)->assume_mutable()->insert_many_defaults(src.size()); |
260 | | // ColumnVariant should be finalized before parsing, finalize maybe modify original column structure |
261 | 0 | (*result)->assume_mutable()->finalize(); |
262 | 0 | return Status::OK(); |
263 | 0 | } |
264 | 1 | if (src.is_scalar_variant() && is_string_type(src.get_root_type()->get_primitive_type())) { |
265 | | // use parser to extract from root |
266 | 0 | auto type = std::make_shared<DataTypeString>(); |
267 | 0 | MutableColumnPtr result_column = type->create_column(); |
268 | 0 | const ColumnString& docs = |
269 | 0 | *check_and_get_column<ColumnString>(remove_nullable(src.get_root()).get()); |
270 | 0 | simdjson::ondemand::parser parser; |
271 | 0 | std::vector<JsonPath> parsed_paths; |
272 | 0 | if (field_name.empty() || field_name[0] != '$') { |
273 | 0 | field_name = "$." + field_name; |
274 | 0 | } |
275 | 0 | JsonFunctions::parse_json_paths(field_name, &parsed_paths); |
276 | 0 | ColumnString* col_str = assert_cast<ColumnString*>(result_column.get()); |
277 | 0 | for (size_t i = 0; i < docs.size(); ++i) { |
278 | 0 | if (!extract_from_document(parser, docs.get_data_at(i), parsed_paths, col_str)) { |
279 | 0 | VLOG_DEBUG << "failed to parse " << docs.get_data_at(i) << ", field " |
280 | 0 | << field_name; |
281 | 0 | result_column->insert_default(); |
282 | 0 | } |
283 | 0 | } |
284 | 0 | *result = ColumnVariant::create(src.max_subcolumns_count(), src.enable_doc_mode(), type, |
285 | 0 | std::move(result_column)); |
286 | 0 | (*result)->assume_mutable()->finalize(); |
287 | 0 | return Status::OK(); |
288 | 1 | } else { |
289 | 1 | auto mutable_src = src.clone_finalized(); |
290 | 1 | auto* mutable_ptr = assert_cast<ColumnVariant*>(mutable_src.get()); |
291 | 1 | PathInData path(field_name); |
292 | 1 | ColumnVariant::Subcolumns subcolumns = mutable_ptr->get_subcolumns(); |
293 | 1 | const auto* node = subcolumns.find_exact(path); |
294 | 1 | MutableColumnPtr result_col = |
295 | 1 | ColumnVariant::create(src.max_subcolumns_count(), src.enable_doc_mode()); |
296 | 1 | ColumnVariant::Subcolumns new_subcolumns; |
297 | | |
298 | 1 | if (node != nullptr) { |
299 | 0 | std::vector<decltype(node)> nodes; |
300 | 0 | PathsInData paths; |
301 | 0 | ColumnVariant::Subcolumns::get_leaves_of_node(node, nodes, paths); |
302 | 0 | for (const auto* n : nodes) { |
303 | 0 | PathInData new_path = n->path.copy_pop_front(); |
304 | 0 | VLOG_DEBUG << "add node " << new_path.get_path() |
305 | 0 | << ", data size: " << n->data.size() |
306 | 0 | << ", finalized size: " << n->data.get_finalized_column().size() |
307 | 0 | << ", common type: " << n->data.get_least_common_type()->get_name(); |
308 | | // if new_path is empty, indicate it's the root column, but adding a root will return false when calling add |
309 | 0 | if (!new_subcolumns.add(new_path, n->data)) { |
310 | 0 | VLOG_DEBUG << "failed to add node " << new_path.get_path(); |
311 | 0 | } |
312 | 0 | } |
313 | | |
314 | | // handle the root node |
315 | 0 | if (new_subcolumns.empty() && !nodes.empty()) { |
316 | 0 | CHECK_EQ(nodes.size(), 1); |
317 | 0 | new_subcolumns.create_root(ColumnVariant::Subcolumn { |
318 | 0 | nodes[0]->data.get_finalized_column_ptr()->assume_mutable(), |
319 | 0 | nodes[0]->data.get_least_common_type(), true, true}); |
320 | 0 | auto container = |
321 | 0 | ColumnVariant::create(src.max_subcolumns_count(), src.enable_doc_mode(), |
322 | 0 | std::move(new_subcolumns)); |
323 | 0 | result_col->insert_range_from(*container, 0, container->size()); |
324 | 0 | } else { |
325 | 0 | auto container = |
326 | 0 | ColumnVariant::create(src.max_subcolumns_count(), src.enable_doc_mode(), |
327 | 0 | std::move(new_subcolumns)); |
328 | 0 | container->clear_sparse_column(); |
329 | 0 | _extract_sparse_column_from_source(mutable_ptr, path, container); |
330 | 0 | result_col->insert_range_from(*container, 0, container->size()); |
331 | 0 | } |
332 | 1 | } else { |
333 | 1 | auto container = |
334 | 1 | ColumnVariant::create(src.max_subcolumns_count(), src.enable_doc_mode(), |
335 | 1 | std::move(new_subcolumns)); |
336 | 1 | const auto& sparse_offsets = mutable_ptr->serialized_sparse_column_offsets(); |
337 | 1 | if (sparse_offsets.back() == sparse_offsets[-1]) { |
338 | 0 | _extract_doc_value_column_from_source(mutable_ptr, path, container); |
339 | 1 | } else { |
340 | 1 | _extract_sparse_column_from_source(mutable_ptr, path, container); |
341 | 1 | } |
342 | 1 | result_col->insert_range_from(*container, 0, container->size()); |
343 | 1 | } |
344 | 1 | *result = result_col->get_ptr(); |
345 | | // ColumnVariant should be finalized before parsing, finalize maybe modify original column structure |
346 | 1 | (*result)->assume_mutable()->finalize(); |
347 | 1 | VLOG_DEBUG << "dump new object " |
348 | 0 | << static_cast<const ColumnVariant*>(result_col.get())->debug_string() |
349 | 0 | << ", path " << path.get_path(); |
350 | 1 | return Status::OK(); |
351 | 1 | } |
352 | 1 | } |
353 | | |
354 | | static Status extract_from_document(simdjson::ondemand::parser& parser, const StringRef& doc, |
355 | 0 | const std::vector<JsonPath>& paths, ColumnString* column) { |
356 | 0 | try { |
357 | 0 | simdjson::padded_string json_str {doc.data, doc.size}; |
358 | 0 | simdjson::ondemand::document document = parser.iterate(json_str); |
359 | 0 | simdjson::ondemand::object object = document.get_object(); |
360 | 0 | simdjson::ondemand::value value; |
361 | 0 | RETURN_IF_ERROR(JsonFunctions::extract_from_object(object, paths, &value)); |
362 | 0 | _write_data_to_column(value, column); |
363 | 0 | } catch (simdjson::simdjson_error& e) { |
364 | 0 | VLOG_DEBUG << "simdjson parse exception: " << e.what(); |
365 | 0 | return Status::DataQualityError("simdjson parse exception {}", e.what()); |
366 | 0 | } |
367 | 0 | return Status::OK(); |
368 | 0 | } |
369 | | |
370 | 0 | static void _write_data_to_column(simdjson::ondemand::value& value, ColumnString* column) { |
371 | 0 | switch (value.type()) { |
372 | 0 | case simdjson::ondemand::json_type::null: { |
373 | 0 | column->insert_default(); |
374 | 0 | break; |
375 | 0 | } |
376 | 0 | case simdjson::ondemand::json_type::boolean: { |
377 | 0 | if (value.get_bool()) { |
378 | 0 | column->insert_data("1", 1); |
379 | 0 | } else { |
380 | 0 | column->insert_data("0", 1); |
381 | 0 | } |
382 | 0 | break; |
383 | 0 | } |
384 | 0 | default: { |
385 | 0 | auto value_str = simdjson::to_json_string(value).value(); |
386 | 0 | column->insert_data(value_str.data(), value_str.length()); |
387 | 0 | } |
388 | 0 | } |
389 | 0 | } |
390 | | }; |
391 | | |
392 | 1 | void register_function_variant_element(SimpleFunctionFactory& factory) { |
393 | 1 | factory.register_function<FunctionVariantElement>(); |
394 | 1 | } |
395 | | |
396 | | } // namespace doris |