DayuanX commented on code in PR #54036:
URL: https://github.com/apache/doris/pull/54036#discussion_r2325726509
##########
be/src/vec/functions/function_map.cpp:
##########
@@ -440,13 +518,400 @@ class FunctionStrToMap : public IFunction {
}
};
+class FunctionMapContainsEntry : public IFunction {
+public:
+ static constexpr auto name = "map_contains_entry";
+ static FunctionPtr create() { return
std::make_shared<FunctionMapContainsEntry>(); }
+
+ String get_name() const override { return name; }
+ size_t get_number_of_arguments() const override { return 3; }
+ bool use_default_implementation_for_nulls() const override { return false;
}
+
+ DataTypePtr get_return_type_impl(const DataTypes& arguments) const
override {
+ DataTypePtr datatype = arguments[0];
+ if (datatype->is_nullable()) {
+ datatype = assert_cast<const
DataTypeNullable*>(datatype.get())->get_nested_type();
+ }
+ DCHECK(datatype->get_primitive_type() == TYPE_MAP)
+ << "first argument for function: " << name << " should be
DataTypeMap";
+
+ if (arguments[0]->is_nullable()) {
+ return make_nullable(std::make_shared<DataTypeBool>());
+ } else {
+ return std::make_shared<DataTypeBool>();
+ }
+ }
+
+ Status execute_impl(FunctionContext* context, Block& block, const
ColumnNumbers& arguments,
+ uint32_t result, size_t input_rows_count) const
override {
+ return _execute_type_check_and_dispatch(block, arguments, result);
+ }
+
+private:
+ // assume result_matches is initialized to all 1s
+ template <typename ColumnType>
+ void _execute_column_comparison(const IColumn& map_entry_column, const
UInt8* map_entry_nullmap,
+ const IColumn& search_column, const UInt8*
search_nullmap,
+ const ColumnArray::Offsets64& map_offsets,
+ const UInt8* map_row_nullmap, bool
search_is_const,
+ ColumnUInt8& result_matches) const {
+ auto& result_data = result_matches.get_data();
+ for (size_t row = 0; row < map_offsets.size(); ++row) {
+ if (map_row_nullmap && map_row_nullmap[row]) {
+ continue;
+ }
+ size_t map_start = row == 0 ? 0 : map_offsets[row - 1];
+ size_t map_end = map_offsets[row];
+ // const column always uses index 0
+ size_t search_idx = search_is_const ? 0 : row;
+ for (size_t i = map_start; i < map_end; ++i) {
+ result_data[i] &=
+ compare_values<ColumnType>(map_entry_column, i,
map_entry_nullmap,
+ search_column, search_idx,
search_nullmap)
+ ? 1
+ : 0;
+ }
+ }
+ }
+
+ // dispatch column comparison by type, map_entry_column is the column of
map's key or value, search_column is the column of search key or value
+ void _dispatch_column_comparison(PrimitiveType type, const IColumn&
map_entry_column,
+ const UInt8* map_entry_nullmap, const
IColumn& search_column,
+ const UInt8* search_nullmap,
+ const ColumnArray::Offsets64& map_offsets,
+ const UInt8* map_row_nullmap, bool
search_is_const,
+ ColumnUInt8& result_matches) const {
+ switch (type) {
+ case TYPE_BOOLEAN:
+ _execute_column_comparison<ColumnUInt8>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_TINYINT:
+ _execute_column_comparison<ColumnInt8>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_SMALLINT:
+ _execute_column_comparison<ColumnInt16>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_INT:
+ _execute_column_comparison<ColumnInt32>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_BIGINT:
+ _execute_column_comparison<ColumnInt64>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_LARGEINT:
+ _execute_column_comparison<ColumnInt128>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_FLOAT:
+ _execute_column_comparison<ColumnFloat32>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_DOUBLE:
+ _execute_column_comparison<ColumnFloat64>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_DECIMAL32:
+ _execute_column_comparison<ColumnDecimal32>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_DECIMAL64:
+ _execute_column_comparison<ColumnDecimal64>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_DECIMAL128I:
+ _execute_column_comparison<ColumnDecimal128V3>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_DECIMALV2:
+ _execute_column_comparison<ColumnDecimal128V2>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_DECIMAL256:
+ _execute_column_comparison<ColumnDecimal256>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_STRING:
+ case TYPE_CHAR:
+ case TYPE_VARCHAR:
+ _execute_column_comparison<ColumnString>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_DATE:
+ _execute_column_comparison<ColumnDate>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_DATETIME:
+ _execute_column_comparison<ColumnDateTime>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_DATEV2:
+ _execute_column_comparison<ColumnDateV2>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_DATETIMEV2:
+ _execute_column_comparison<ColumnDateTimeV2>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ case TYPE_TIMEV2:
+ _execute_column_comparison<ColumnTimeV2>(
+ map_entry_column, map_entry_nullmap, search_column,
search_nullmap, map_offsets,
+ map_row_nullmap, search_is_const, result_matches);
+ break;
+ default:
Review Comment:
DCHECK(false) added.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
