{

FunctionArrayAUC(const char * name_, bool is_pr_) : function_name(name_), is_pr(is_pr_) {}

static FunctionPtr create(const char * name, bool is_pr)

{

return std::make_shared<FunctionArrayAUC>(name, is_pr);

}

String getName() const override { return function_name; }

bool isVariadic() const override { return true; }

size_t getNumberOfArguments() const override { return 0; }

bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo &) const override { return true; }

DataTypePtr getReturnTypeImpl(const ColumnsWithTypeAndName & arguments) const override

{

size_t argument_count = arguments.size();

if ((is_pr && (argument_count < 2 || argument_count > 3))

|| (!is_pr && (argument_count < 2 || argument_count > 4)))

throw Exception(

ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH,

"Number of arguments for function {} doesn't match: passed {}, should be {}",

getName(),

argument_count,

is_pr ? "2 or 3" : "2, 3 or 4");

/// Validate 'arr_scores' argument

const DataTypeArray * array_scores_type = checkAndGetDataType<DataTypeArray>(arguments[0].type.get());

if (!array_scores_type)

throw Exception(

ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "First argument (arr_scores) for function {} must be of type Array", getName());

const auto & nested_array_scores_type = array_scores_type->getNestedType();

if (!isNativeNumber(nested_array_scores_type))

throw Exception(

ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,

"First argument (arr_scores) for function {} is of type {}, should be a number",

getName(),

nested_array_scores_type->getName());

/// Validate 'arr_labels' argument

const DataTypeArray * array_labels_type = checkAndGetDataType<DataTypeArray>(arguments[1].type.get());

if (!array_labels_type)

throw Exception(

ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Second argument (arr_labels) for function {} must be of type Array", getName());

const auto & nested_arary_labels_type = array_labels_type->getNestedType();

if (!isNativeNumber(nested_arary_labels_type) && !isEnum(nested_arary_labels_type))

throw Exception(

ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,

"Second argument (arr_labels) for function {} is of type {}, should be a number or Enum",

getName(),

nested_arary_labels_type->getName());

/// Validate 'scale' argument

if (!is_pr && argument_count >= 3 && !isConstBoolColumn(arguments[2]))

throw Exception(

ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "Third argument (scale) for function {} must be of type const Bool", getName());

/// Validate 'arr_partial_offsets' argument

if (argument_count >= array_partial_offsets_arg_index + 1)

{

const DataTypeArray * array_offsets_type = checkAndGetDataType<DataTypeArray>(arguments[array_partial_offsets_arg_index].type.get());

if (!array_offsets_type)

throw Exception(

ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,

"{} argument (arr_partial_offsets) for function {} must be of type Array",

array_partial_offsets_arg_index == 2 ? "Third" : "Fourth",

getName());

const auto & nested_array_offsets_type = array_offsets_type->getNestedType();

/// Last argument (arr_partial_offsets) must be an array of integers

if (!isInteger(nested_array_offsets_type))

throw Exception(

ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT,

"{} argument (arr_partial_offsets) for function {} is of type {}, should be integer",

array_partial_offsets_arg_index == 2 ? "Third" : "Fourth",

getName(),

nested_array_offsets_type->getName());

}

return std::make_shared<DataTypeFloat64>();

}

DataTypePtr getReturnTypeForDefaultImplementationForDynamic() const override { return std::make_shared<DataTypeFloat64>(); }

ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override

{

size_t argument_count = arguments.size();

ColumnPtr col1 = arguments[0].column->convertToFullColumnIfConst();

ColumnPtr col2 = arguments[1].column->convertToFullColumnIfConst();

const ColumnArray * col_array1 = checkAndGetColumn<ColumnArray>(col1.get());

if (!col_array1)

throw Exception(

ErrorCodes::ILLEGAL_COLUMN,

"Illegal column {} of first argument of function {}",

arguments[0].column->getName(),

getName());

const ColumnArray * col_array2 = checkAndGetColumn<ColumnArray>(col2.get());

if (!col_array2)

throw Exception(

ErrorCodes::ILLEGAL_COLUMN,

"Illegal column {} of second argument of function {}",

arguments[1].column->getName(),

getName());

if (!col_array1->hasEqualOffsets(*col_array2))

throw Exception(ErrorCodes::BAD_ARGUMENTS, "First two arguments for function {} must be Arrays of equal sizes", getName());

/// Handle the 'scale' argument (if passed and the function is arrayROCAUC, otherwise default to true)

bool scale = true;

if (!is_pr && argument_count >= 3 && input_rows_count > 0)

scale = arguments[2].column->getBool(0);

/// Handle the 'arr_partial_offsets' argument (if passed)

ColumnPtr col_offsets = nullptr;

const ColumnArray * col_array_offsets = nullptr;

if (argument_count == array_partial_offsets_arg_index + 1)

{

col_offsets = arguments[array_partial_offsets_arg_index].column->convertToFullColumnIfConst();

col_array_offsets = checkAndGetColumn<ColumnArray>(col_offsets.get());

if (!col_array_offsets)

throw Exception(

ErrorCodes::ILLEGAL_COLUMN,

"Illegal column {} of {} argument of function {}",

arguments[array_partial_offsets_arg_index].column->getName(),

array_partial_offsets_arg_index == 2 ? "third" : "fourth",

getName());

/// The partial offsets argument must be a column containing 3-elements (PR AUC) or 4-elements (ROC AUC) arrays on each row

const auto & offsets = col_array_offsets->getOffsets();

if ((col_array_offsets->getData().size() != array_partial_offsets_size * input_rows_count) || (offsets.size() != input_rows_count))

throw Exception(

ErrorCodes::BAD_ARGUMENTS,

"{} argument (arr_partial_offsets) for function {} must contain Arrays of size {}",

array_partial_offsets_arg_index == 2 ? "Third" : "Fourth",

getName(),

array_partial_offsets_size);

for (size_t i = 0; i < input_rows_count; ++i)

{

auto current = offsets[i];

auto previous = i == 0 ? 0 : offsets[i - 1];

if (current - previous != array_partial_offsets_size)

throw Exception(

ErrorCodes::BAD_ARGUMENTS,

"{} argument (arr_partial_offsets) for function {} must contain Arrays of size {}, not {}",

array_partial_offsets_arg_index == 2 ? "Third" : "Fourth",

getName(),

array_partial_offsets_size,

current - previous);

}

for (size_t i = 0; i < col_array_offsets->getData().size(); ++i)

{

if (col_array_offsets->getData().getInt(i) < 0)

throw Exception(

ErrorCodes::BAD_ARGUMENTS,

"{} argument (arr_partial_offsets) for function {} must not contain negative values",

array_partial_offsets_arg_index == 2 ? "Third" : "Fourth",

getName());

}

}

auto col_res = ColumnVector<Float64>::create();

vector(

col_array1->getData(),

col_array2->getData(),

col_array1->getOffsets(),

col_res->getData(),

input_rows_count,

scale,

col_array_offsets);

return col_res;

}

const char * function_name;

const bool is_pr;

const size_t array_partial_offsets_arg_index = is_pr ? 2 : 3;

const size_t array_partial_offsets_size = is_pr ? 3 : 4;

static bool isConstBoolColumn(ColumnWithTypeAndName argument)

{

if (!isBool(argument.type))

return false;

if (argument.column.get() == nullptr)

return false;

if (!isColumnConst(*argument.column))

return false;

return true;

}

Float64 increase_unscaled_area(size_t prev_fp, size_t prev_tp, size_t curr_fp, size_t curr_tp) const

{

if (is_pr)

/// PR curve plots Precision x Recall

///

/// Precision = TP / (TP + FP)

/// Recall = TP / (TP + FN)

///

/// The AUC is calculated using the Right Riemann Sum.

///

/// Instead of calculating

/// area += (Recall_n - Recall_{n-1}) * Precision_n,

/// we simplify it to

/// area += (TP_n - TP_{n-1}) * Precision_n

/// and compute the "unscaled" area.

///

/// The unscaled area represents the AUC of the Precision x TP curve.

/// Later we can divide it by (TP + FN) to obtain the correct AUC.

///

/// This can be done because (TP + FN) is constant and equal to total positive labels.

return static_cast<Float64>(curr_tp) / static_cast<Float64>(curr_tp + curr_fp) * static_cast<Float64>(curr_tp - prev_tp);

else

/// ROC curve plots TPR x FPR

///

/// TPR = TP / (TP + FN)

/// FPR = FP / (FP + TN)

///

/// The AUC is calculated using the Trapezoidal Rule.

///

/// Instead of calculating

/// area += (FPR_n - FPR_{n-1}) * (TPR_n + TPR_{n-1}) / 2,

/// we simplify it to

/// area += (FP_n - FP_{n-1}) * (TP_n + TP_{n-1}) / 2,

/// and compute the "unscaled" area.

///

/// The unscaled area represents the AUC of the TP x FP curve.

/// Later we can divide it by (TP + FN) and (FP + TN) to obtain the correct AUC.

///

/// This can be done because both (TP + FN) and (FP + TN) are constant and

/// equal to total positive labels and total negative labels, respectively.

return static_cast<Float64>(curr_fp - prev_fp) * static_cast<Float64>(curr_tp + prev_tp) / 2.0;

}

Float64 scale_back_area(Float64 area, size_t total_positive_labels, size_t total_negative_labels) const

{

if (is_pr)

/// To simplify the calculations, previously we calculated the AUC for the Precision x TP curve.

/// This scales back to Precision x Recall by dividing the area by (TP + FN).

return area / static_cast<Float64>(total_positive_labels);

else

/// To simplify the calculations, previously we calculated the AUC for the TP x FP curve.

/// This scales back to TPR x FPR by dividing the area by (TP + FN) and (FP + TN).

return area / static_cast<Float64>(total_positive_labels) / static_cast<Float64>(total_negative_labels);

}

Float64 apply(

const IColumn & scores,

const IColumn & labels,

ColumnArray::Offset current_offset,

ColumnArray::Offset next_offset,

bool scale,

size_t higher_partitions_tp = 0,

size_t higher_partitions_fp = 0,

size_t total_positives = 0,

size_t total_negatives = 0) const

{

struct ScoreLabel

{

Float64 score;

bool label;

};

size_t size = next_offset - current_offset;

if (size == 0)

return is_pr ? 0.0 : std::numeric_limits<Float64>::quiet_NaN();

PODArrayWithStackMemory<ScoreLabel, 1024> sorted_labels(size);

for (size_t i = 0; i < size; ++i)

{

sorted_labels[i].label = labels.getFloat64(current_offset + i) > 0;

sorted_labels[i].score = scores.getFloat64(current_offset + i);

}

/// A NaN score has no place in the order below and belongs to no threshold, so the area is undefined.

for (size_t i = 0; i < size; ++i)

if (isNaN(sorted_labels[i].score))

return std::numeric_limits<Float64>::quiet_NaN();

/// Sorting scores in descending order to traverse the ROC / Precision-Recall curve from left to right

std::sort(sorted_labels.begin(), sorted_labels.end(), [](const auto & lhs, const auto & rhs) { return lhs.score > rhs.score; });

Float64 area = 0.0;

Float64 threshold = sorted_labels[0].score;

size_t prev_fp = higher_partitions_fp;

size_t prev_tp = higher_partitions_tp;

size_t curr_fp = higher_partitions_fp; /// False positives predictions (label <= 0 and score > threshold)

size_t curr_tp = higher_partitions_tp; /// True positives predictions (label > 0 and score > threshold)

/// Traversing the sorted labels, changing the threshold and incrementing the area accordingly

for (size_t i = 0; i < size; ++i)

{

/// Only increment the area when the threshold (score) changes

if (sorted_labels[i].score != threshold)

{

area += increase_unscaled_area(prev_fp, prev_tp, curr_fp, curr_tp);

prev_fp = curr_fp;

prev_tp = curr_tp;

threshold = sorted_labels[i].score;

}

if (sorted_labels[i].label)

curr_tp += 1;

else

curr_fp += 1;

}

area += increase_unscaled_area(prev_fp, prev_tp, curr_fp, curr_tp);

/// Unless scale is false, we scale the area back to the [0..1] range

if (scale)

{

/// Degenerate cases where we would divide by zero when scaling back the area

if (!is_pr && ((total_positives == 0 && curr_tp == 0) || (total_negatives == 0 && curr_fp == 0)))

/// If no positive or negative labels, TPR or FPR is undefined and we return NaN

return std::numeric_limits<Float64>::quiet_NaN();

if (is_pr && (total_positives == 0 && curr_tp == 0))

/// Precision did not change, PR curve degenerates into a single point (0, 1) and we return 0.0

return 0.0;

/// If we're calculating the partial AUC, the user needs to pass the total_positives or total_negatives labels

/// Otherwise, we assume that the whole curve was traversed above, which means that the threshold is at minimum and all

/// labels were predicted as true. This means that curr_tp = total positive labels and curr_fp = total negative labels

return scale_back_area(

area, total_positives != 0 ? total_positives : curr_tp, total_negatives != 0 ? total_negatives : curr_fp);

}

return area;

}

void vector(

const IColumn & scores,

const IColumn & labels,

const ColumnArray::Offsets & offsets,

PaddedPODArray<Float64> & result,

size_t input_rows_count,

bool scale,

const ColumnArray * partial_auc_offsets) const

{

result.resize(input_rows_count);

ColumnArray::Offset current_offset = 0;

for (size_t i = 0; i < input_rows_count; ++i)

{

auto next_offset = offsets[i];

if (partial_auc_offsets)

result[i] = apply(

scores,

labels,

current_offset,

next_offset,

scale,

partial_auc_offsets->getData().getUInt(array_partial_offsets_size * i),

partial_auc_offsets->getData().getUInt(array_partial_offsets_size * i + 1),

partial_auc_offsets->getData().getUInt(array_partial_offsets_size * i + 2),

array_partial_offsets_size == 4 ? partial_auc_offsets->getData().getUInt(array_partial_offsets_size * i + 3) : 0);

else

result[i] = apply(scores, labels, current_offset, next_offset, scale);

current_offset = next_offset;

}

}

{

/// ROC AUC

FunctionDocumentation::Description description_roc = R"(

)";

FunctionDocumentation::Syntax syntax_roc = "arrayROCAUC(scores, labels[, scale[, partial_offsets]])";

FunctionDocumentation::Arguments arguments_roc = {

{"scores", "Scores prediction model gives.", {"Array((U)Int*)", "Array(Float*)"}},

{"labels", "Labels of samples, usually 1 for positive sample and 0 for negative sample.", {"Array((U)Int*)", "Enum"}},

{"scale", "Optional. Decides whether to return the normalized area. If false, returns the area under the TP (true positives) x FP (false positives) curve instead. Default value: true.", {"Bool"}},

{"partial_offsets", R"(

)"}

};

FunctionDocumentation::ReturnedValue returned_value_roc = {"Returns area under the receiver operating characteristic (ROC) curve. Returns `NaN` if a score is `NaN`, and, when `scale` is true, also when the ROC AUC is undefined because there are no positive or no negative labels.", {"Float64"}};

FunctionDocumentation::Examples examples_roc = {{"Usage example", "SELECT arrayROCAUC([0.1, 0.4, 0.35, 0.8], [0, 0, 1, 1]);", "0.75"}};

FunctionDocumentation::IntroducedIn introduced_in_roc = {20, 4};

FunctionDocumentation::Category category_roc = FunctionDocumentation::Category::Array;

FunctionDocumentation documentation_roc = {description_roc, syntax_roc, arguments_roc, {}, returned_value_roc, examples_roc, introduced_in_roc, category_roc};

factory.registerFunction("arrayROCAUC", [](ContextPtr){ return FunctionArrayAUC::create("arrayROCAUC", false); }, documentation_roc);

factory.registerAlias("arrayAUC", "arrayROCAUC"); /// Backward compatibility, also ROC AUC is often shorted to just AUC

/// PR AUC

FunctionDocumentation::Description description_pr = R"(

)";

FunctionDocumentation::Syntax syntax_pr = "arrayAUCPR(scores, labels[, partial_offsets])";

FunctionDocumentation::Arguments arguments_pr = {

{"cores", "Scores prediction model gives.", {"Array((U)Int*)", "Array(Float*)"}},

{"labels", "Labels of samples, usually 1 for positive sample and 0 for negative sample.", {"Array((U)Int*)", "Array(Enum)"}},

{"partial_offsets", R"(

)"}

};

FunctionDocumentation::ReturnedValue returned_value_pr = {"Returns area under the precision-recall (PR) curve, or `NaN` if any score is `NaN`.", {"Float64"}};

FunctionDocumentation::Examples examples_pr = {{"Usage example", "SELECT arrayAUCPR([0.1, 0.4, 0.35, 0.8], [0, 0, 1, 1]);", R"(

)"}};

FunctionDocumentation::IntroducedIn introduced_in_pr = {20, 4};

FunctionDocumentation::Category category_pr = FunctionDocumentation::Category::Array;

FunctionDocumentation documentation_pr = {description_pr, syntax_pr, arguments_pr, {}, returned_value_pr, examples_pr, introduced_in_pr, category_pr};

factory.registerFunction("arrayAUCPR", [](ContextPtr){ return FunctionArrayAUC::create("arrayAUCPR", true); }, documentation_pr);

factory.registerAlias("arrayPRAUC", "arrayAUCPR");