simple join count optimization

7 files changed, 202 insertions, 0 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index 960f15eac27..844c81acf52 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -77,6 +77,7 @@ vespa_define_module(
     src/tests/instruction/mixed_simple_join_function
     src/tests/instruction/pow_as_map_optimizer
     src/tests/instruction/remove_trivial_dimension_optimizer
+    src/tests/instruction/simple_join_count
     src/tests/instruction/sparse_112_dot_product
     src/tests/instruction/sparse_dot_product_function
     src/tests/instruction/sparse_full_overlap_join_function
diff --git a/eval/src/tests/instruction/simple_join_count/CMakeLists.txt b/eval/src/tests/instruction/simple_join_count/CMakeLists.txt
new file mode 100644
index 00000000000..d1be148dc59
--- /dev/null
+++ b/eval/src/tests/instruction/simple_join_count/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(eval_simple_join_count_test_app TEST
+    SOURCES
+    simple_join_count_test.cpp
+    DEPENDS
+    vespaeval
+    GTest::GTest
+)
+vespa_add_test(NAME eval_simple_join_count_test_app COMMAND eval_simple_join_count_test_app)
diff --git a/eval/src/tests/instruction/simple_join_count/simple_join_count_test.cpp b/eval/src/tests/instruction/simple_join_count/simple_join_count_test.cpp
new file mode 100644
index 00000000000..a040a876609
--- /dev/null
+++ b/eval/src/tests/instruction/simple_join_count/simple_join_count_test.cpp
@@ -0,0 +1,70 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/eval/eval/fast_value.h>
+#include <vespa/eval/eval/simple_value.h>
+#include <vespa/eval/instruction/simple_join_count.h>
+#include <vespa/eval/eval/test/eval_fixture.h>
+#include <vespa/eval/eval/test/gen_spec.h>
+#include <vespa/vespalib/gtest/gtest.h>
+
+using namespace vespalib::eval;
+using namespace vespalib::eval::test;
+
+//-----------------------------------------------------------------------------
+
+struct FunInfo {
+    using LookFor = SimpleJoinCount;
+    uint64_t expected_dense_factor;
+    FunInfo(uint64_t expected_dense_factor_in)
+      : expected_dense_factor(expected_dense_factor_in) {}
+    void verify(const LookFor &fun) const {
+        EXPECT_TRUE(fun.result_is_mutable());
+        EXPECT_EQ(fun.dense_factor(), expected_dense_factor);
+    }
+};
+
+void verify_optimized_cell_types(const vespalib::string &expr, size_t expected_dense_factor = 1) {
+    CellTypeSpace types(CellTypeUtils::list_types(), 2);
+    EvalFixture::verify<FunInfo>(expr, {FunInfo(expected_dense_factor)}, types);
+}
+
+void verify_optimized(const vespalib::string &expr, size_t expected_dense_factor = 1) {
+    CellTypeSpace just_float({CellType::FLOAT}, 2);
+    EvalFixture::verify<FunInfo>(expr, {FunInfo(expected_dense_factor)}, just_float);
+}
+
+void verify_not_optimized(const vespalib::string &expr) {
+    CellTypeSpace just_float({CellType::FLOAT}, 2);
+    EvalFixture::verify<FunInfo>(expr, {}, just_float);
+}
+
+//-----------------------------------------------------------------------------
+
+TEST(SimpleJoinCount, expression_can_be_optimized) {
+    verify_optimized_cell_types("reduce(x5_2*x5_1,count)");
+    verify_optimized_cell_types("reduce(x5_2y3z4*x5_1z4a1,count)", 12);
+}
+
+TEST(SimpleJoinCount, join_operation_does_not_matter) {
+    verify_optimized("reduce(x5_2+x5_1,count)");
+    verify_optimized("reduce(x5_2-x5_1,count)");
+    verify_optimized("reduce(x5_2/x5_1,count)");
+}
+
+TEST(SimpleJoinCount, parameters_must_have_full_mapped_singledim_overlap) {
+    verify_not_optimized("reduce(x5_2y5_2*x5_1y5_2,count)");
+    verify_not_optimized("reduce(x5_2*y5_2,count)");
+    verify_not_optimized("reduce(x5_2y5_2*x5_1z5_2,count)");
+    verify_not_optimized("reduce(x5_2*y5,count)");
+    verify_not_optimized("reduce(x5*y5,count)");
+}
+
+TEST(SimpleJoinCount, similar_expressions_are_not_optimized) {
+    verify_not_optimized("reduce(x5_2y3z4*x5_1z4a1,count,x)");
+    verify_not_optimized("reduce(x5_2y3z4*x5_1z4a1,count,x,y,z)");
+    verify_not_optimized("reduce(x5_2y3*x5_1,sum)");
+}
+
+//-----------------------------------------------------------------------------
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/eval/src/vespa/eval/eval/optimize_tensor_function.cpp b/eval/src/vespa/eval/eval/optimize_tensor_function.cpp
index 45a4c79d2ed..63f0f00a13c 100644
--- a/eval/src/vespa/eval/eval/optimize_tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/optimize_tensor_function.cpp
@@ -34,6 +34,8 @@
 #include <vespa/eval/instruction/dense_tensor_peek_function.h>
 #include <vespa/eval/instruction/dense_hamming_distance.h>
 #include <vespa/eval/instruction/l2_distance.h>
+#include <vespa/eval/instruction/simple_join_count.h>
+
 
 #include <vespa/log/log.h>
 LOG_SETUP(".eval.eval.optimize_tensor_function");
@@ -82,6 +84,7 @@ const TensorFunction &optimize_for_factory(const ValueBuilderFactory &, const Te
                           child.set(DenseMultiMatMulFunction::optimize(child.get(), stash));
                           child.set(MixedInnerProductFunction::optimize(child.get(), stash));
                           child.set(DenseHammingDistance::optimize(child.get(), stash));
+                          child.set(SimpleJoinCount::optimize(child.get(), stash));
                       });
     run_optimize_pass(root, [&stash](const Child &child)
                       {
diff --git a/eval/src/vespa/eval/instruction/CMakeLists.txt b/eval/src/vespa/eval/instruction/CMakeLists.txt
index c4fd0443cbb..2146e3ee8ab 100644
--- a/eval/src/vespa/eval/instruction/CMakeLists.txt
+++ b/eval/src/vespa/eval/instruction/CMakeLists.txt
@@ -37,6 +37,7 @@ vespa_add_library(eval_instruction OBJECT
     pow_as_map_optimizer.cpp
     remove_trivial_dimension_optimizer.cpp
     replace_type_function.cpp
+    simple_join_count.cpp
     sparse_112_dot_product.cpp
     sparse_dot_product_function.cpp
     sparse_full_overlap_join_function.cpp
diff --git a/eval/src/vespa/eval/instruction/simple_join_count.cpp b/eval/src/vespa/eval/instruction/simple_join_count.cpp
new file mode 100644
index 00000000000..1a4750e276b
--- /dev/null
+++ b/eval/src/vespa/eval/instruction/simple_join_count.cpp
@@ -0,0 +1,92 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "simple_join_count.h"
+#include "generic_join.h"
+#include <vespa/eval/eval/fast_value.hpp>
+
+namespace vespalib::eval {
+
+using namespace tensor_function;
+using namespace operation;
+using namespace instruction;
+
+namespace {
+
+size_t my_intersect_count_fallback(const Value::Index &lhs_idx, const Value::Index &rhs_idx) {
+    size_t result = 0.0;
+    SparseJoinPlan plan(1);
+    SparseJoinState sparse(plan, lhs_idx, rhs_idx);
+    auto outer = sparse.first_index.create_view({});
+    auto inner = sparse.second_index.create_view(sparse.second_view_dims);
+    outer->lookup({});
+    while (outer->next_result(sparse.first_address, sparse.first_subspace)) {
+        inner->lookup(sparse.address_overlap);
+        if (inner->next_result(sparse.second_only_address, sparse.second_subspace)) {
+            ++result;
+        }
+    }
+    return result;
+}
+
+size_t my_fast_intersect_count(const FastAddrMap *small_map, const FastAddrMap *big_map) {
+    size_t result = 0;
+    if (big_map->size() < small_map->size()) {
+        std::swap(small_map, big_map);
+    }
+    const auto &labels = small_map->labels();
+    for (size_t i = 0; i < labels.size(); ++i) {
+        if (big_map->lookup_singledim(labels[i]) != FastAddrMap::npos()) {
+            ++result;
+        }
+    }
+    return result;
+}
+
+void my_simple_join_count_op(InterpretedFunction::State &state, uint64_t dense_factor) {
+    const auto &lhs_idx = state.peek(1).index();
+    const auto &rhs_idx = state.peek(0).index();
+    double result = dense_factor * (__builtin_expect(are_fast(lhs_idx, rhs_idx), true)
+                                    ? my_fast_intersect_count(&as_fast(lhs_idx).map, &as_fast(rhs_idx).map)
+                                    : my_intersect_count_fallback(lhs_idx, rhs_idx));
+    state.pop_pop_push(state.stash.create<DoubleValue>(result));
+}
+
+bool check_types(const ValueType &res, const ValueType &lhs, const ValueType &rhs) {
+    return ((res.is_double()) &&
+            (lhs.count_mapped_dimensions() == 1) &&
+            (lhs.mapped_dimensions() == rhs.mapped_dimensions()));
+}
+
+} // namespace <unnamed>
+
+SimpleJoinCount::SimpleJoinCount(const TensorFunction &lhs_in,
+                                 const TensorFunction &rhs_in,
+                                 uint64_t dense_factor_in)
+  : tensor_function::Op2(ValueType::double_type(), lhs_in, rhs_in),
+    _dense_factor(dense_factor_in)
+{
+}
+
+InterpretedFunction::Instruction
+SimpleJoinCount::compile_self(const ValueBuilderFactory &, Stash &) const
+{
+    return InterpretedFunction::Instruction(my_simple_join_count_op, _dense_factor);
+}
+
+const TensorFunction &
+SimpleJoinCount::optimize(const TensorFunction &expr, Stash &stash)
+{
+    auto reduce = as<Reduce>(expr);
+    if (reduce && (reduce->aggr() == Aggr::COUNT)) {
+        if (auto join = as<Join>(reduce->child())) {
+            const TensorFunction &lhs = join->lhs();
+            const TensorFunction &rhs = join->rhs();
+            if (check_types(expr.result_type(), lhs.result_type(), rhs.result_type())) {
+                return stash.create<SimpleJoinCount>(lhs, rhs, join->result_type().dense_subspace_size());
+            }
+        }
+    }
+    return expr;
+}
+
+} // namespace
diff --git a/eval/src/vespa/eval/instruction/simple_join_count.h b/eval/src/vespa/eval/instruction/simple_join_count.h
new file mode 100644
index 00000000000..a566d2a7e68
--- /dev/null
+++ b/eval/src/vespa/eval/instruction/simple_join_count.h
@@ -0,0 +1,26 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/eval/eval/tensor_function.h>
+
+namespace vespalib::eval {
+
+/**
+ * Tensor function that will count the number of cells in the result
+ * of a join between two tensors with full mapped overlap consisting
+ * of a single dimension.
+ **/
+class SimpleJoinCount : public tensor_function::Op2
+{
+private:
+    uint64_t _dense_factor;
+public:
+    SimpleJoinCount(const TensorFunction &lhs_in, const TensorFunction &rhs_in, size_t dense_factor_in);
+    InterpretedFunction::Instruction compile_self(const ValueBuilderFactory &factory, Stash &stash) const override;
+    bool result_is_mutable() const override { return true; }
+    uint64_t dense_factor() const { return _dense_factor; }
+    static const TensorFunction &optimize(const TensorFunction &expr, Stash &stash);
+};
+
+} // namespace