Skip to content

[improve](function)Refactor distance function return types to float #55184

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

Open
wants to merge 4 commits into
base: master
Choose a base branch
from
+86 −131
Open

[improve](function)Refactor distance function return types to float #55184

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
9e61b51
Change return type of l2_distance_approximate and inner_product_appro…
uchenily Aug 21, 2025
ba4abda
Change return type of all distance functions
uchenily Aug 22, 2025
62484c4
check style
uchenily Aug 25, 2025
cb6209c
remove unnecessary argument
uchenily Aug 25, 2025
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
33 changes: 10 additions & 23 deletions be/src/olap/rowset/segment_v2/ann_index/ann_topn_runtime.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,10 +54,7 @@ Status AnnTopNRuntime::prepare(RuntimeState* state, const RowDescriptor& row_des
|----------------
| |
| |
CastToArray ArrayLiteral
|
|
SlotRef
SlotRef ArrayLiteral
*/
std::shared_ptr<vectorized::VirtualSlotRef> vir_slot_ref =
std::dynamic_pointer_cast<vectorized::VirtualSlotRef>(_order_by_expr_ctx->root());
Expand All @@ -78,19 +75,11 @@ Status AnnTopNRuntime::prepare(RuntimeState* state, const RowDescriptor& row_des
vir_col_expr->debug_string());
}

std::shared_ptr<vectorized::VCastExpr> cast_to_array_expr =
std::dynamic_pointer_cast<vectorized::VCastExpr>(distance_fn_call->children()[0]);

if (cast_to_array_expr == nullptr) {
return Status::InternalError("Ann topn expr expect cast_to_array_expr, got\n{}",
distance_fn_call->children()[0]->debug_string());
}

std::shared_ptr<vectorized::VSlotRef> slot_ref =
std::dynamic_pointer_cast<vectorized::VSlotRef>(cast_to_array_expr->children()[0]);
std::dynamic_pointer_cast<vectorized::VSlotRef>(distance_fn_call->children()[0]);
Copy link
Contributor

@zhiqiang-hhhh zhiqiang-hhhh Aug 25, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why cast_to_array is removed?

if (slot_ref == nullptr) {
return Status::InternalError("Ann topn expr expect SlotRef, got\n{}",
cast_to_array_expr->children()[0]->debug_string());
distance_fn_call->children()[0]->debug_string());
}

// slot_ref->column_id() is acutually the columnd idx in block.
Expand Down Expand Up @@ -139,11 +128,11 @@ Status AnnTopNRuntime::evaluate_vector_ann_search(segment_v2::IndexIterator* ann
assert_cast<const vectorized::ColumnArray*>(const_column->get_data_column_ptr().get());
const vectorized::ColumnNullable* column_nullable =
assert_cast<const vectorized::ColumnNullable*>(column_array->get_data_ptr().get());
const vectorized::ColumnFloat64* cf64 = assert_cast<const vectorized::ColumnFloat64*>(
const vectorized::ColumnFloat32* cf32 = assert_cast<const vectorized::ColumnFloat32*>(
column_nullable->get_nested_column_ptr().get());

const double* query_value = cf64->get_data().data();
const size_t query_value_size = cf64->get_data().size();
const float* query_value = cf32->get_data().data();
const size_t query_value_size = cf32->get_data().size();

std::unique_ptr<float[]> query_value_f32 = std::make_unique<float[]>(query_value_size);
for (size_t i = 0; i < query_value_size; ++i) {
Expand All @@ -167,15 +156,13 @@ Status AnnTopNRuntime::evaluate_vector_ann_search(segment_v2::IndexIterator* ann
DCHECK(ann_query_params.row_ids != nullptr);

size_t num_results = ann_query_params.distance->size();
auto result_column_double = vectorized::ColumnFloat64::create(num_results);
auto result_null_map = vectorized::ColumnUInt8::create(num_results, 0);
auto result_column_float = vectorized::ColumnFloat32::create(num_results);

for (size_t i = 0; i < num_results; ++i) {
result_column_double->get_data()[i] = (*ann_query_params.distance)[i];
result_column_float->get_data()[i] = (*ann_query_params.distance)[i];
}

result_column = vectorized::ColumnNullable::create(std::move(result_column_double),
std::move(result_null_map));
result_column = std::move(result_column_float);
row_ids = std::move(ann_query_params.row_ids);
ann_index_stats = *ann_query_params.stats;
return Status::OK();
Expand All @@ -188,4 +175,4 @@ std::string AnnTopNRuntime::debug_string() const {
_limit, _src_column_idx, _dest_column_idx, _asc, _user_params.to_string(),
segment_v2::metric_to_string(_metric_type), _order_by_expr_ctx->root()->debug_string());
}
} // namespace doris::segment_v2
} // namespace doris::segment_v2
2 changes: 1 addition & 1 deletion be/src/vec/exprs/vdirect_in_predicate.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,4 +148,4 @@ class VDirectInPredicate final : public VExpr {
};

#include "common/compile_check_end.h"
} // namespace doris::vectorized
} // namespace doris::vectorized
22 changes: 9 additions & 13 deletions be/src/vec/exprs/vectorized_fn_call.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -369,7 +369,7 @@ void VectorizedFnCall::prepare_ann_range_search(
auto left_child = get_child(0);
auto right_child = get_child(1);

// Return type of L2Distance is always double.
// Return type of L2Distance is always float.
auto right_literal = std::dynamic_pointer_cast<VLiteral>(right_child);
if (right_literal == nullptr) {
suitable_for_ann_index = false;
Expand Down Expand Up @@ -457,10 +457,10 @@ void VectorizedFnCall::prepare_ann_range_search(
range_search_runtime.query_value = std::make_unique<float[]>(dim);

const ColumnNullable* cn = assert_cast<const ColumnNullable*>(array_col->get_data_ptr().get());
const ColumnFloat64* cf64 =
assert_cast<const ColumnFloat64*>(cn->get_nested_column_ptr().get());
const ColumnFloat32* cf32 =
assert_cast<const ColumnFloat32*>(cn->get_nested_column_ptr().get());
for (size_t i = 0; i < dim; ++i) {
range_search_runtime.query_value[i] = static_cast<Float32>(cf64->get_data()[i]);
range_search_runtime.query_value[i] = cf32->get_data()[i];
}
range_search_runtime.is_ann_range_search = true;
range_search_runtime.user_params = user_params;
Expand Down Expand Up @@ -552,17 +552,13 @@ Status VectorizedFnCall::evaluate_ann_range_search(
DCHECK(virtual_column_iterator != nullptr);
// Now convert distance to column
size_t size = result.roaring->cardinality();
auto distance_col = ColumnFloat64::create(size);
auto null_map = ColumnUInt8::create(size, 0);
// TODO: Return type of L2DistanceApproximate/InnerProductApproximate should be changed to float.
const float* src = reinterpret_cast<const float*>(result.distance.get());
double* dst = distance_col->get_data().data();
auto distance_col = ColumnFloat32::create(size);
const float* src = result.distance.get();
float* dst = distance_col->get_data().data();
for (size_t i = 0; i < size; ++i) {
dst[i] = static_cast<double>(src[i]);
dst[i] = src[i];
}
auto nullable_distance_col =
ColumnNullable::create(std::move(distance_col), std::move(null_map));
virtual_column_iterator->prepare_materialization(std::move(nullable_distance_col),
virtual_column_iterator->prepare_materialization(std::move(distance_col),
std::move(result.row_ids));
} else {
DCHECK(this->op() != TExprOpcode::LE && this->op() != TExprOpcode::LT)
Expand Down
1 change: 1 addition & 0 deletions be/src/vec/exprs/vtopn_pred.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ class VTopNPred : public VExpr {
node.__set_is_nullable(target_ctx->root()->is_nullable());
expr = vectorized::VTopNPred::create_shared(node, source_node_id, target_ctx);

DCHECK(target_ctx->root() != nullptr);
expr->add_child(target_ctx->root());

return Status::OK();
Expand Down
12 changes: 6 additions & 6 deletions be/src/vec/functions/array/function_array_distance.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,12 @@
namespace doris::vectorized {

void register_function_array_distance(SimpleFunctionFactory& factory) {
factory.register_function<FunctionArrayDistance<L1Distance> >();
factory.register_function<FunctionArrayDistance<L2Distance> >();
factory.register_function<FunctionArrayDistance<CosineDistance> >();
factory.register_function<FunctionArrayDistance<InnerProduct> >();
factory.register_function<FunctionArrayDistance<L2DistanceApproximate> >();
factory.register_function<FunctionArrayDistance<InnerProductApproximate> >();
factory.register_function<FunctionArrayDistance<L1Distance>>();
factory.register_function<FunctionArrayDistance<L2Distance>>();
factory.register_function<FunctionArrayDistance<CosineDistance>>();
factory.register_function<FunctionArrayDistance<InnerProduct>>();
factory.register_function<FunctionArrayDistance<L2DistanceApproximate>>();
factory.register_function<FunctionArrayDistance<InnerProductApproximate>>();
}

} // namespace doris::vectorized
87 changes: 29 additions & 58 deletions be/src/vec/functions/array/function_array_distance.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@

#pragma once

#include <gen_cpp/Types_types.h>

#include "vec/columns/column.h"
#include "vec/columns/column_array.h"
#include "vec/common/assert_cast.h"
Expand All @@ -35,73 +37,66 @@ class L1Distance {
public:
static constexpr auto name = "l1_distance";
struct State {
double sum = 0;
float sum = 0;
};
static void accumulate(State& state, double x, double y) { state.sum += fabs(x - y); }
static double finalize(const State& state) { return state.sum; }
static void accumulate(State& state, float x, float y) { state.sum += fabs(x - y); }
static float finalize(const State& state) { return state.sum; }
};

class L2Distance {
public:
static constexpr auto name = "l2_distance";
struct State {
double sum = 0;
float sum = 0;
};
static void accumulate(State& state, double x, double y) { state.sum += (x - y) * (x - y); }
static double finalize(const State& state) { return sqrt(state.sum); }
static void accumulate(State& state, float x, float y) { state.sum += (x - y) * (x - y); }
static float finalize(const State& state) { return sqrt(state.sum); }
};

class InnerProduct {
public:
static constexpr auto name = "inner_product";
struct State {
double sum = 0;
float sum = 0;
};
static void accumulate(State& state, double x, double y) { state.sum += x * y; }
static double finalize(const State& state) { return state.sum; }
static void accumulate(State& state, float x, float y) { state.sum += x * y; }
static float finalize(const State& state) { return state.sum; }
};

class CosineDistance {
public:
static constexpr auto name = "cosine_distance";
struct State {
double dot_prod = 0;
double squared_x = 0;
double squared_y = 0;
float dot_prod = 0;
float squared_x = 0;
float squared_y = 0;
};
static void accumulate(State& state, double x, double y) {
static void accumulate(State& state, float x, float y) {
state.dot_prod += x * y;
state.squared_x += x * x;
state.squared_y += y * y;
}
static double finalize(const State& state) {
static float finalize(const State& state) {
return 1 - state.dot_prod / sqrt(state.squared_x * state.squared_y);
}
};

class L2DistanceApproximate {
class L2DistanceApproximate : public L2Distance {
public:
static constexpr auto name = "l2_distance_approximate";
struct State {
double sum = 0;
};
static void accumulate(State& state, double x, double y) { state.sum += (x - y) * (x - y); }
static double finalize(const State& state) { return sqrt(state.sum); }
};

class InnerProductApproximate {
class InnerProductApproximate : public InnerProduct {
public:
static constexpr auto name = "inner_product_approximate";
struct State {
double sum = 0;
};
static void accumulate(State& state, double x, double y) { state.sum += x * y; }
static double finalize(const State& state) { return state.sum; }
};

template <typename DistanceImpl>
class FunctionArrayDistance : public IFunction {
public:
using DataType = PrimitiveTypeTraits<TYPE_FLOAT>::DataType;
using ColumnType = PrimitiveTypeTraits<TYPE_FLOAT>::ColumnType;

static constexpr auto name = DistanceImpl::name;
String get_name() const override { return name; }
static FunctionPtr create() { return std::make_shared<FunctionArrayDistance<DistanceImpl>>(); }
Expand All @@ -110,7 +105,7 @@ class FunctionArrayDistance : public IFunction {
bool use_default_implementation_for_nulls() const override { return false; }

DataTypePtr get_return_type_impl(const DataTypes& arguments) const override {
return make_nullable(std::make_shared<DataTypeFloat64>());
return std::make_shared<DataType>();
}

Status execute_impl(FunctionContext* context, Block& block, const ColumnNumbers& arguments,
Expand Down Expand Up @@ -140,27 +135,14 @@ class FunctionArrayDistance : public IFunction {
}

// prepare return data
auto dst = ColumnFloat64::create(input_rows_count);
auto dst = ColumnType::create(input_rows_count);
auto& dst_data = dst->get_data();
auto dst_null_column = ColumnUInt8::create(input_rows_count, 0);
auto& dst_null_data = dst_null_column->get_data();

const auto& offsets1 = *arr1.offsets_ptr;
const auto& offsets2 = *arr2.offsets_ptr;
const auto& nested_col1 = assert_cast<const ColumnFloat64*>(arr1.nested_col.get());
const auto& nested_col2 = assert_cast<const ColumnFloat64*>(arr2.nested_col.get());
const auto& nested_col1 = assert_cast<const ColumnType*>(arr1.nested_col.get());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

what if input column is nullable

const auto& nested_col2 = assert_cast<const ColumnType*>(arr2.nested_col.get());
for (ssize_t row = 0; row < offsets1.size(); ++row) {
if (arr1.array_nullmap_data && arr1.array_nullmap_data[row]) {
dst_null_data[row] = true;
continue;
}
if (arr2.array_nullmap_data && arr2.array_nullmap_data[row]) {
dst_null_data[row] = true;
continue;
}

dst_null_data[row] = false;

// Calculate actual array sizes for current row.
// For nullable arrays, we cannot compare absolute offset values directly because:
// 1. When a row is null, its offset might equal the previous offset (no elements added)
Expand All @@ -179,25 +161,14 @@ class FunctionArrayDistance : public IFunction {
for (ssize_t pos = offsets1[row - 1]; pos < offsets1[row]; ++pos) {
// Calculate corresponding position in the second array
ssize_t pos2 = offsets2[row - 1] + (pos - offsets1[row - 1]);
if (arr1.nested_nullmap_data && arr1.nested_nullmap_data[pos]) {
dst_null_data[row] = true;
break;
}
if (arr2.nested_nullmap_data && arr2.nested_nullmap_data[pos2]) {
dst_null_data[row] = true;
break;
}
DistanceImpl::accumulate(st, nested_col1->get_element(pos),
nested_col2->get_element(pos2));
}
if (!dst_null_data[row]) {
dst_data[row] = DistanceImpl::finalize(st);
dst_null_data[row] = std::isnan(dst_data[row]);
}

dst_data[row] = DistanceImpl::finalize(st);
}

block.replace_by_position(
result, ColumnNullable::create(std::move(dst), std::move(dst_null_column)));
block.replace_by_position(result, std::move(dst));
return Status::OK();
}

Expand All @@ -209,7 +180,7 @@ class FunctionArrayDistance : public IFunction {
}
auto nested_type =
remove_nullable(assert_cast<const DataTypeArray&>(*array_type).get_nested_type());
return nested_type->get_primitive_type() == TYPE_DOUBLE;
return nested_type->get_primitive_type() == TYPE_FLOAT;
}
};

Expand Down
10 changes: 5 additions & 5 deletions ...main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/CosineDistance.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,12 +19,12 @@

import org.apache.doris.catalog.FunctionSignature;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable;
import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable;
import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.types.ArrayType;
import org.apache.doris.nereids.types.DoubleType;
import org.apache.doris.nereids.types.FloatType;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
Expand All @@ -35,11 +35,11 @@
* cosine_distance function
*/
public class CosineDistance extends ScalarFunction implements ExplicitlyCastableSignature,
BinaryExpression, AlwaysNullable {
BinaryExpression, PropagateNullable {

public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
FunctionSignature.ret(DoubleType.INSTANCE)
.args(ArrayType.of(DoubleType.INSTANCE), ArrayType.of(DoubleType.INSTANCE))
FunctionSignature.ret(FloatType.INSTANCE)
.args(ArrayType.of(FloatType.INSTANCE), ArrayType.of(FloatType.INSTANCE))
);

/**
Expand Down
10 changes: 5 additions & 5 deletions ...c/main/java/org/apache/doris/nereids/trees/expressions/functions/scalar/InnerProduct.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,12 +19,12 @@

import org.apache.doris.catalog.FunctionSignature;
import org.apache.doris.nereids.trees.expressions.Expression;
import org.apache.doris.nereids.trees.expressions.functions.AlwaysNullable;
import org.apache.doris.nereids.trees.expressions.functions.ExplicitlyCastableSignature;
import org.apache.doris.nereids.trees.expressions.functions.PropagateNullable;
import org.apache.doris.nereids.trees.expressions.shape.BinaryExpression;
import org.apache.doris.nereids.trees.expressions.visitor.ExpressionVisitor;
import org.apache.doris.nereids.types.ArrayType;
import org.apache.doris.nereids.types.DoubleType;
import org.apache.doris.nereids.types.FloatType;

import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableList;
Expand All @@ -35,11 +35,11 @@
* inner_product function
*/
public class InnerProduct extends ScalarFunction implements ExplicitlyCastableSignature,
BinaryExpression, AlwaysNullable {
BinaryExpression, PropagateNullable {

public static final List<FunctionSignature> SIGNATURES = ImmutableList.of(
FunctionSignature.ret(DoubleType.INSTANCE)
.args(ArrayType.of(DoubleType.INSTANCE), ArrayType.of(DoubleType.INSTANCE))
FunctionSignature.ret(FloatType.INSTANCE)
.args(ArrayType.of(FloatType.INSTANCE), ArrayType.of(FloatType.INSTANCE))
);

/**
Expand Down
Loading