Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feature] Support array functions for JSON type #13155

Merged
merged 8 commits into from
Nov 10, 2022
Merged

[Feature] Support array functions for JSON type #13155

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
0c31d7d
improve performance of array_agg function
murphyatwork Nov 7, 2022
ae35334
fix
murphyatwork Nov 7, 2022
3a5b690
note
murphyatwork Nov 7, 2022
ae8cf58
support several array functions for array type
murphyatwork Nov 9, 2022
5466c4e
support array_agg for json
murphyatwork Nov 9, 2022
e317884
cast parameter for array_slice
murphyatwork Nov 9, 2022
3af8db5
Merge branch 'main' into murphy_agg_json
murphyatwork Nov 9, 2022
5e1d275
format
murphyatwork Nov 9, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
support several array functions for array type
  • Loading branch information
@murphyatwork
murphyatwork committed Nov 9, 2022
commit ae8cf58850d83caa53f10597cab839b90c52624f
5 changes: 5 additions & 0 deletions be/src/column/object_column.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,11 @@ class ObjectColumn : public ColumnFactory<Column, ObjectColumn<T>> {
_buffer.clear();
}

void reset_cache() {
_cache_ok = false;
_cache.clear();
}

Buffer<T>& get_pool() { return _pool; }

const Buffer<T>& get_pool() const { return _pool; }
Expand Down
3 changes: 2 additions & 1 deletion be/src/exprs/agg/factory/aggregate_resolver_avg.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ struct AvgDispatcher {
struct ArrayAggDispatcher {
template <PrimitiveType pt>
void operator()(AggregateFuncResolver* resolver) {
if constexpr (pt_is_aggregate<pt> || pt_is_string<pt>) {
if constexpr (pt_is_aggregate<pt> || pt_is_string<pt> || pt_is_json<pt>) {
auto func = std::make_shared<ArrayAggAggregateFunction<pt>>();
using AggState = ArrayAggAggregateState<pt>;
resolver->add_aggregate_mapping<pt, TYPE_ARRAY, AggState, AggregateFunctionPtr, false>("array_agg", false,
Expand All @@ -38,6 +38,7 @@ void AggregateFuncResolver::register_avg() {
type_dispatch_all(type, AvgDispatcher(), this);
type_dispatch_all(type, ArrayAggDispatcher(), this);
}
type_dispatch_all(TYPE_JSON, ArrayAggDispatcher(), this);
add_decimal_mapping<TYPE_DECIMAL32, TYPE_DECIMAL128, true>("decimal_avg");
add_decimal_mapping<TYPE_DECIMAL64, TYPE_DECIMAL128, true>("decimal_avg");
add_decimal_mapping<TYPE_DECIMAL128, TYPE_DECIMAL128, true>("decimal_avg");
Expand Down
1 change: 1 addition & 0 deletions be/src/exprs/agg/factory/aggregate_resolver_minmaxany.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ void AggregateFuncResolver::register_minmaxany() {
for (auto type : aggregate_types()) {
type_dispatch_all(type, MinMaxAnyDispatcher(), this);
}
type_dispatch_all(TYPE_JSON, MinMaxAnyDispatcher(), this);
}

} // namespace starrocks::vectorized
4 changes: 4 additions & 0 deletions be/src/exprs/vectorized/array_functions.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,13 +63,15 @@ class ArrayFunctions {
return ArraySlice<PT>::process(context, columns); \
}
APPLY_COMMONE_TYPES_FOR_ARRAY(DEFINE_ARRAY_SLICE_FN)
DEFINE_ARRAY_SLICE_FN(json, PrimitiveType::TYPE_JSON)
#undef DEFINE_ARRAY_SLICE_FN

#define DEFINE_ARRAY_CONCAT_FN(NAME, PT) \
static ColumnPtr array_concat_##NAME(FunctionContext* context, const Columns& columns) { \
return ArrayConcat<PT>::process(context, columns); \
}
APPLY_COMMONE_TYPES_FOR_ARRAY(DEFINE_ARRAY_CONCAT_FN)
DEFINE_ARRAY_CONCAT_FN(json, PrimitiveType::TYPE_JSON)
#undef DEFINE_ARRAY_CONCAT_FN

#define DEFINE_ARRAY_OVERLAP_FN(NAME, PT) \
Expand All @@ -91,13 +93,15 @@ class ArrayFunctions {
return ArraySort<PT>::process(context, columns); \
}
APPLY_COMMONE_TYPES_FOR_ARRAY(DEFINE_ARRAY_SORT_FN)
DEFINE_ARRAY_SORT_FN(json, PrimitiveType::TYPE_JSON)
#undef DEFINE_ARRAY_SORT_FN

#define DEFINE_ARRAY_REVERSE_FN(NAME, PT) \
static ColumnPtr array_reverse_##NAME(FunctionContext* context, const Columns& columns) { \
return ArrayReverse<PT>::process(context, columns); \
}
APPLY_COMMONE_TYPES_FOR_ARRAY(DEFINE_ARRAY_REVERSE_FN)
DEFINE_ARRAY_REVERSE_FN(json, PrimitiveType::TYPE_JSON)
#undef DEFINE_ARRAY_REVERSE_FN

#define DEFINE_ARRAY_JOIN_FN(NAME) \
Expand Down
25 changes: 23 additions & 2 deletions be/src/exprs/vectorized/array_functions.tpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,10 @@
#include "column/column_builder.h"
#include "column/column_hash.h"
#include "column/column_viewer.h"
#include "column/json_column.h"
#include "column/type_traits.h"
#include "exprs/vectorized/function_helper.h"
#include "runtime/primitive_type.h"
#include "udf/udf.h"
#include "util/orlp/pdqsort.h"
#include "util/phmap/phmap.h"
Expand Down Expand Up @@ -760,6 +763,13 @@ private:
pdqsort(false, sort_index->begin() + offset, sort_index->begin() + offset + count, less_fn);
}

// For JSON type
static void _sort_column(std::vector<uint32_t>* sort_index, const JsonColumn& src_column, size_t offset,
size_t count) {
auto less_fn = [&](uint32_t l, uint32_t r) -> bool { return src_column.compare_at(l, r, src_column, -1) < 0; };
pdqsort(false, sort_index->begin() + offset, sort_index->begin() + offset + count, less_fn);
}

static void _sort_item(std::vector<uint32_t>* sort_index, const Column& src_column,
const UInt32Column& offset_column, size_t index) {
const auto& offsets = offset_column.get_data();
Expand All @@ -770,7 +780,7 @@ private:
return;
}

_sort_column(sort_index, src_column, start, count);
_sort_column(sort_index, down_cast<const RunTimeColumnType<PT>&>(src_column), start, count);
}

static void _sort_nullable_item(std::vector<uint32_t>* sort_index, const Column& src_data_column,
Expand All @@ -790,7 +800,7 @@ private:
std::partition(sort_index->begin() + start, sort_index->begin() + start + count, null_first_fn);
size_t data_offset = begin_of_not_null - sort_index->begin();
size_t null_count = data_offset - start;
_sort_column(sort_index, src_data_column, start + null_count, count - null_count);
_sort_column(sort_index, down_cast<const RunTimeColumnType<PT>&>(src_data_column), start + null_count, count - null_count);
}

static void _sort_array_column(Column* dest_array_column, std::vector<uint32_t>* sort_index,
Expand Down Expand Up @@ -900,11 +910,22 @@ private:
}
}

static void _reverse_json_column(Column* column, const Buffer<uint32_t>& array_offsets, size_t chunk_size) {
auto json_column = down_cast<JsonColumn*>(column);
auto& pool = json_column->get_pool();
for (size_t i = 0; i < chunk_size; i++) {
std::reverse(pool.begin() + array_offsets[i], pool.begin() + array_offsets[i + 1]);
}
json_column->reset_cache();
}

static void _reverse_data_column(Column* column, const Buffer<uint32_t>& offsets, size_t chunk_size) {
if constexpr (pt_is_fixedlength<PT>) {
_reverse_fixed_column(column, offsets, chunk_size);
} else if constexpr (pt_is_string<PT>) {
_reverse_binary_column(column, offsets, chunk_size);
} else if constexpr (pt_is_json<PT>) {
_reverse_json_column(column, offsets, chunk_size);
} else {
assert(false);
}
Expand Down
5 changes: 5 additions & 0 deletions gensrc/script/vectorized/vectorized_functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -663,6 +663,7 @@
[150119, 'array_sort', 'ARRAY_DECIMALV2', ['ARRAY_DECIMALV2'], 'ArrayFunctions::array_sort_decimalv2'],
[150120, 'array_sort', 'ARRAY_DATETIME', ['ARRAY_DATETIME'], 'ArrayFunctions::array_sort_datetime'],
[150121, 'array_sort', 'ARRAY_DATE', ['ARRAY_DATE'], 'ArrayFunctions::array_sort_date'],
[150122, 'array_sort', 'ARRAY_JSON', ['ARRAY_JSON'], 'ArrayFunctions::array_sort_json'],

[150130, 'reverse', 'ARRAY_BOOLEAN', ['ARRAY_BOOLEAN'], 'ArrayFunctions::array_reverse_boolean'],
[150131, 'reverse', 'ARRAY_TINYINT', ['ARRAY_TINYINT'], 'ArrayFunctions::array_reverse_tinyint'],
Expand All @@ -676,6 +677,7 @@
[150139, 'reverse', 'ARRAY_DECIMALV2', ['ARRAY_DECIMALV2'], 'ArrayFunctions::array_reverse_decimalv2'],
[150140, 'reverse', 'ARRAY_DATETIME', ['ARRAY_DATETIME'], 'ArrayFunctions::array_reverse_datetime'],
[150141, 'reverse', 'ARRAY_DATE', ['ARRAY_DATE'], 'ArrayFunctions::array_reverse_date'],
[150142, 'reverse', 'ARRAY_JSON', ['ARRAY_JSON'], 'ArrayFunctions::array_reverse_json'],

[150150, 'array_join', 'VARCHAR', ['ARRAY_VARCHAR', 'VARCHAR'], 'ArrayFunctions::array_join_varchar'],
[150151, 'array_join', 'VARCHAR', ['ARRAY_VARCHAR', 'VARCHAR', 'VARCHAR'], 'ArrayFunctions::array_join_varchar'],
Expand All @@ -702,6 +704,7 @@
[150179, 'array_slice', 'ARRAY_DOUBLE', ['ARRAY_DOUBLE', 'BIGINT', 'BIGINT'], 'ArrayFunctions::array_slice_double'],
[150180, 'array_slice', 'ARRAY_DECIMALV2', ['ARRAY_DECIMALV2', 'BIGINT', 'BIGINT'], 'ArrayFunctions::array_slice_decimalv2'],
[150181, 'array_slice', 'ARRAY_VARCHAR', ['ARRAY_VARCHAR', 'BIGINT', 'BIGINT'], 'ArrayFunctions::array_slice_varchar'],
[150182, 'array_slice', 'ARRAY_JSON', ['ARRAY_JSON', 'BIGINT', 'BIGINT'], 'ArrayFunctions::array_slice_json'],

[150190, 'array_concat', 'ARRAY_DATE', ['ARRAY_DATE', "..."], 'ArrayFunctions::array_concat_date'],
[150191, 'array_concat', 'ARRAY_DATETIME', ['ARRAY_DATETIME', "..."], 'ArrayFunctions::array_concat_datetime'],
Expand All @@ -715,6 +718,7 @@
[150199, 'array_concat', 'ARRAY_DOUBLE', ['ARRAY_DOUBLE', "..."], 'ArrayFunctions::array_concat_double'],
[150200, 'array_concat', 'ARRAY_DECIMALV2', ['ARRAY_DECIMALV2', "..."], 'ArrayFunctions::array_concat_decimalv2'],
[150201, 'array_concat', 'ARRAY_VARCHAR', ['ARRAY_VARCHAR', "..."], 'ArrayFunctions::array_concat_varchar'],
[150202, 'array_concat', 'ARRAY_JSON', ['ARRAY_JSON', "..."], 'ArrayFunctions::array_concat_json'],

[150210, 'arrays_overlap', 'BOOLEAN', ['ARRAY_DATE', 'ARRAY_DATE'], 'ArrayFunctions::array_overlap_date'],
[150211, 'arrays_overlap', 'BOOLEAN', ['ARRAY_DATETIME', 'ARRAY_DATETIME'], 'ArrayFunctions::array_overlap_datetime'],
Expand Down Expand Up @@ -754,6 +758,7 @@
[150259, 'array_slice', 'ARRAY_DOUBLE', ['ARRAY_DOUBLE', 'BIGINT'], 'ArrayFunctions::array_slice_double'],
[150260, 'array_slice', 'ARRAY_DECIMALV2', ['ARRAY_DECIMALV2', 'BIGINT'], 'ArrayFunctions::array_slice_decimalv2'],
[150261, 'array_slice', 'ARRAY_VARCHAR', ['ARRAY_VARCHAR', 'BIGINT'], 'ArrayFunctions::array_slice_varchar'],
[150262, 'array_slice', 'ARRAY_JSON', ['ARRAY_JSON', 'BIGINT'], 'ArrayFunctions::array_slice_json'],

[150270, 'array_cum_sum', 'ARRAY_BIGINT', ['ARRAY_BIGINT'], 'ArrayFunctions::array_cum_sum_bigint'],
[150271, 'array_cum_sum', 'ARRAY_DOUBLE', ['ARRAY_DOUBLE'], 'ArrayFunctions::array_cum_sum_double'],
Expand Down