// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements. See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership. The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License. You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied. See the License for the

// specific language governing permissions and limitations

// under the License.

#include <glog/logging.h>

#include <algorithm>

#include <boost/iterator/iterator_facade.hpp>

#include <cstddef>

#include <memory>

#include <utility>

#include "common/status.h"

#include "core/assert_cast.h"

#include "core/block/block.h"

#include "core/block/column_numbers.h"

#include "core/block/column_with_type_and_name.h"

#include "core/column/column.h"

#include "core/column/column_array.h"

#include "core/column/column_nullable.h"

#include "core/column/column_vector.h"

#include "core/data_type/data_type.h"

#include "core/data_type/data_type_array.h"

#include "core/data_type/data_type_date_or_datetime_v2.h"

#include "core/data_type/data_type_date_time.h"

#include "core/data_type/data_type_nullable.h"

#include "core/data_type/define_primitive_type.h"

#include "core/pod_array_fwd.h"

#include "core/types.h"

#include "core/value/vdatetime_value.h"

#include "exprs/aggregate/aggregate_function.h"

#include "exprs/function/function.h"

#include "exprs/function/function_date_or_datetime_computation.h"

#include "exprs/function/simple_function_factory.h"

namespace doris {

class FunctionContext;

} // namespace doris

namespace doris {

template <typename Impl>

class FunctionArrayRange : public IFunction {

public:

static constexpr auto name = Impl::name;

static FunctionPtr create() { return std::make_shared<FunctionArrayRange>(); }

/// Get function name.

String get_name() const override { return name; }

bool is_variadic() const override { return true; }

size_t get_number_of_arguments() const override {

return get_variadic_argument_types_impl().size();

}

DataTypes get_variadic_argument_types_impl() const override {

return Impl::get_variadic_argument_types();

}

DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {

auto nested_type = make_nullable(Impl::get_data_type());

auto res = std::make_shared<DataTypeArray>(nested_type);

return make_nullable(res);

}

Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,

uint32_t result, size_t input_rows_count) const override {

return Impl::execute_impl(context, block, arguments, result, input_rows_count);

}

};

template <PrimitiveType SourceDataPType, typename TimeUnitOrVoid = void>

struct RangeImplUtil {

using SourceDataType = typename PrimitiveTypeTraits<SourceDataPType>::CppType;

using DataType = typename PrimitiveTypeTraits<SourceDataPType>::DataType;

static DataTypePtr get_data_type() { return std::make_shared<DataType>(); }

static constexpr const char* get_function_name() {

if constexpr (std::is_same_v<SourceDataType, DateV2Value<DateTimeV2ValueType>> &&

!std::is_same_v<TimeUnitOrVoid, void>) {

if constexpr (std::is_same_v<TimeUnitOrVoid,

std::integral_constant<TimeUnit, TimeUnit::YEAR>>) {

return "array_range_year_unit";

} else if constexpr (std::is_same_v<

TimeUnitOrVoid,

std::integral_constant<TimeUnit, TimeUnit::QUARTER>>) {

return "array_range_quarter_unit";

} else if constexpr (std::is_same_v<

TimeUnitOrVoid,

std::integral_constant<TimeUnit, TimeUnit::MONTH>>) {

return "array_range_month_unit";

} else if constexpr (std::is_same_v<TimeUnitOrVoid,

std::integral_constant<TimeUnit, TimeUnit::WEEK>>) {

return "array_range_week_unit";

} else if constexpr (std::is_same_v<TimeUnitOrVoid,

std::integral_constant<TimeUnit, TimeUnit::DAY>>) {

return "array_range_day_unit";

} else if constexpr (std::is_same_v<TimeUnitOrVoid,

std::integral_constant<TimeUnit, TimeUnit::HOUR>>) {

return "array_range_hour_unit";

} else if constexpr (std::is_same_v<

TimeUnitOrVoid,

std::integral_constant<TimeUnit, TimeUnit::MINUTE>>) {

return "array_range_minute_unit";

} else if constexpr (std::is_same_v<

TimeUnitOrVoid,

std::integral_constant<TimeUnit, TimeUnit::SECOND>>) {

return "array_range_second_unit";

}

} else {

return "array_range";

}

}

static constexpr auto name = get_function_name();

static Status range_execute(Block& block, const ColumnNumbers& arguments, uint32_t result,

size_t input_rows_count) {

DCHECK_EQ(arguments.size(), 3);

auto return_nested_type = make_nullable(std::make_shared<DataType>());

auto dest_array_column_ptr = ColumnArray::create(return_nested_type->create_column(),

ColumnArray::ColumnOffsets::create());

IColumn* dest_nested_column = &dest_array_column_ptr->get_data();

auto* dest_nested_nullable_col = assert_cast<ColumnNullable*>(dest_nested_column);

dest_nested_column = dest_nested_nullable_col->get_nested_column_ptr().get();

auto& dest_nested_null_map = dest_nested_nullable_col->get_null_map_column().get_data();

auto args_null_map = ColumnUInt8::create(input_rows_count, 0);

ColumnPtr argument_columns[3];

for (int i = 0; i < 3; ++i) {

argument_columns[i] =

block.get_by_position(arguments[i]).column->convert_to_full_column_if_const();

}

auto start_column =

assert_cast<const ColumnVector<SourceDataPType>*>(argument_columns[0].get());

auto end_column =

assert_cast<const ColumnVector<SourceDataPType>*>(argument_columns[1].get());

const auto* step_column = assert_cast<const ColumnInt32*>(argument_columns[2].get());

DCHECK(dest_nested_column != nullptr);

auto& dest_offsets = dest_array_column_ptr->get_offsets();

auto nested_column = reinterpret_cast<ColumnVector<SourceDataPType>*>(dest_nested_column);

dest_offsets.reserve(input_rows_count);

dest_nested_column->reserve(input_rows_count);

dest_nested_null_map.reserve(input_rows_count);

RETURN_IF_ERROR(vector(start_column->get_data(), end_column->get_data(),

step_column->get_data(), args_null_map->get_data(),

nested_column->get_data(), dest_nested_null_map, dest_offsets));

block.get_by_position(result).column =

ColumnNullable::create(std::move(dest_array_column_ptr), std::move(args_null_map));

return Status::OK();

}

private:

static Status vector(const PaddedPODArray<SourceDataType>& start,

const PaddedPODArray<SourceDataType>& end,

const PaddedPODArray<Int32>& step, NullMap& args_null_map,

PaddedPODArray<SourceDataType>& nested_column,

PaddedPODArray<UInt8>& dest_nested_null_map,

ColumnArray::Offsets64& dest_offsets) {

size_t rows = start.size();

for (auto row = 0; row < rows; ++row) {

auto idx = start[row];

auto end_row = end[row];

auto step_row = step[row];

auto args_null_map_row = args_null_map[row];

if constexpr (std::is_same_v<SourceDataType, Int32>) {

if (args_null_map_row || idx < 0 || end_row < 0 || step_row <= 0) {

args_null_map[row] = 1;

dest_offsets.push_back(dest_offsets.back());

continue;

} else {

if (idx < end_row && step_row > 0 &&

((static_cast<__int128_t>(end_row) - static_cast<__int128_t>(idx) - 1) /

static_cast<__int128_t>(step_row) +

1) > max_array_size_as_field) {

return Status::InvalidArgument("Array size exceeds the limit {}",

max_array_size_as_field);

}

size_t offset = dest_offsets.back();

while (idx < end[row]) {

nested_column.push_back(idx);

dest_nested_null_map.push_back(0);

offset++;

idx = idx + step_row;

}

dest_offsets.push_back(offset);

}

} else {

const auto& idx_0 = reinterpret_cast<const DateV2Value<DateTimeV2ValueType>&>(idx);

const auto& end_row_cast =

reinterpret_cast<const DateV2Value<DateTimeV2ValueType>&>(end_row);

bool is_null = !idx_0.is_valid_date();

bool is_end_row_invalid = !end_row_cast.is_valid_date();

if (args_null_map_row || step_row <= 0 || is_null || is_end_row_invalid) {

args_null_map[row] = 1;

dest_offsets.push_back(dest_offsets.back());

continue;

} else {

size_t offset = dest_offsets.back();

using UNIT = std::conditional_t<std::is_same_v<TimeUnitOrVoid, void>,

std::integral_constant<TimeUnit, TimeUnit::DAY>,

TimeUnitOrVoid>;

int move = 0;

while (doris::datetime_diff<UNIT::value, DateTimeV2ValueType,

DateTimeV2ValueType>(idx, end_row) > 0) {

if (move > max_array_size_as_field) {

return Status::InvalidArgument("Array size exceeds the limit {}",

max_array_size_as_field);

}

nested_column.push_back(idx);

dest_nested_null_map.push_back(0);

offset++;

move++;

idx = doris::date_time_add<UNIT::value, TYPE_DATETIMEV2, Int32>(idx,

step_row);

}

dest_offsets.push_back(offset);

}

}

}

return Status::OK();

}

};

template <PrimitiveType SourceDataPType, typename TimeUnitOrVoid = void>

struct RangeOneImpl : public RangeImplUtil<SourceDataPType, TimeUnitOrVoid> {

static DataTypes get_variadic_argument_types() {

return {std::make_shared<typename RangeImplUtil<SourceDataPType>::DataType>()};

}

static Status execute_impl(FunctionContext* context, Block& block,

const ColumnNumbers& arguments, uint32_t result,

size_t input_rows_count) {

using ColumnType =

std::conditional_t<SourceDataPType == TYPE_INT, ColumnInt32, ColumnDateTimeV2>;

auto start_column = ColumnType::create(input_rows_count, 0);

auto step_column = ColumnInt32::create(input_rows_count, 1);

block.insert({std::move(start_column),

std::make_shared<typename RangeImplUtil<SourceDataPType>::DataType>(),

"start_column"});

block.insert({std::move(step_column), std::make_shared<DataTypeInt32>(), "step_column"});

ColumnNumbers temp_arguments = {block.columns() - 2, arguments[0], block.columns() - 1};

return (RangeImplUtil<SourceDataPType, TimeUnitOrVoid>::range_execute)(

block, temp_arguments, result, input_rows_count);

}

};

template <PrimitiveType SourceDataPType, typename TimeUnitOrVoid = void>

struct RangeTwoImpl : public RangeImplUtil<SourceDataPType, TimeUnitOrVoid> {

static DataTypes get_variadic_argument_types() {

return {std::make_shared<typename RangeImplUtil<SourceDataPType>::DataType>(),

std::make_shared<typename RangeImplUtil<SourceDataPType>::DataType>()};

}

static Status execute_impl(FunctionContext* context, Block& block,

const ColumnNumbers& arguments, uint32_t result,

size_t input_rows_count) {

auto step_column = ColumnInt32::create(input_rows_count, 1);

block.insert({std::move(step_column), std::make_shared<DataTypeInt32>(), "step_column"});

ColumnNumbers temp_arguments = {arguments[0], arguments[1], block.columns() - 1};

return (RangeImplUtil<SourceDataPType, TimeUnitOrVoid>::range_execute)(

block, temp_arguments, result, input_rows_count);

}

};

template <PrimitiveType SourceDataPType, typename TimeUnitOrVoid = void>

struct RangeThreeImpl : public RangeImplUtil<SourceDataPType, TimeUnitOrVoid> {

static DataTypes get_variadic_argument_types() {

return {std::make_shared<typename RangeImplUtil<SourceDataPType>::DataType>(),

std::make_shared<typename RangeImplUtil<SourceDataPType>::DataType>(),

std::make_shared<DataTypeInt32>()};

}

static Status execute_impl(FunctionContext* context, Block& block,

const ColumnNumbers& arguments, uint32_t result,

size_t input_rows_count) {

return (RangeImplUtil<SourceDataPType, TimeUnitOrVoid>::range_execute)(

block, arguments, result, input_rows_count);

}

};

void register_function_array_range(SimpleFunctionFactory& factory) {

/// One argument, just for Int32

factory.register_function<FunctionArrayRange<RangeOneImpl<TYPE_INT>>>();

/// Two arguments, for Int32 and DateTimeV2 without Interval

factory.register_function<FunctionArrayRange<RangeTwoImpl<TYPE_INT>>>();

factory.register_function<FunctionArrayRange<RangeTwoImpl<TYPE_DATETIMEV2>>>();

/// Three arguments, for Int32 and DateTimeV2 with YEAR to SECOND Interval

factory.register_function<FunctionArrayRange<RangeThreeImpl<TYPE_INT>>>();

factory.register_function<FunctionArrayRange<

RangeThreeImpl<TYPE_DATETIMEV2, std::integral_constant<TimeUnit, TimeUnit::YEAR>>>>();

factory.register_function<FunctionArrayRange<RangeThreeImpl<

TYPE_DATETIMEV2, std::integral_constant<TimeUnit, TimeUnit::QUARTER>>>>();

factory.register_function<FunctionArrayRange<

RangeThreeImpl<TYPE_DATETIMEV2, std::integral_constant<TimeUnit, TimeUnit::MONTH>>>>();

factory.register_function<FunctionArrayRange<

RangeThreeImpl<TYPE_DATETIMEV2, std::integral_constant<TimeUnit, TimeUnit::WEEK>>>>();

factory.register_function<FunctionArrayRange<

RangeThreeImpl<TYPE_DATETIMEV2, std::integral_constant<TimeUnit, TimeUnit::DAY>>>>();

factory.register_function<FunctionArrayRange<

RangeThreeImpl<TYPE_DATETIMEV2, std::integral_constant<TimeUnit, TimeUnit::HOUR>>>>();

factory.register_function<FunctionArrayRange<

RangeThreeImpl<TYPE_DATETIMEV2, std::integral_constant<TimeUnit, TimeUnit::MINUTE>>>>();

factory.register_function<FunctionArrayRange<

RangeThreeImpl<TYPE_DATETIMEV2, std::integral_constant<TimeUnit, TimeUnit::SECOND>>>>();

// alias

factory.register_alias("array_range", "sequence");

}

} // namespace doris