Line data    Source code

       1              : #include "backend/engine/semantic/functions/kll_funcs.h"
       2              : 
       3              : #include <optional>
       4              : #include <vector>
       5              : 
       6              : #include "backend/engine/semantic/value.h"
       7              : #include "googlesql/public/types/array_type.h"
       8              : #include "gtest/gtest.h"
       9              : 
      10              : namespace bigquery_emulator {
      11              : namespace backend {
      12              : namespace engine {
      13              : namespace semantic {
      14              : namespace functions {
      15              : namespace {
      16              : 
      17            5 : const ::googlesql::ArrayType* Int64ArrayType() {
      18            5 :   return ::googlesql::types::Int64ArrayType();
      19            5 : }
      20              : 
      21              : absl::StatusOr<Value> InitInt64(const std::vector<int64_t>& values,
      22            7 :                                 int64_t precision = 1000) {
      23            7 :   std::vector<Value> inputs;
      24            7 :   inputs.reserve(values.size());
      25           25 :   for (int64_t v : values) {
      26           25 :     inputs.push_back(Value::Int64(v));
      27           25 :   }
      28            7 :   return KllQuantilesInitInt64Values(inputs, precision);
      29            7 : }
      30              : 
      31              : }  // namespace
      32              : 
      33            1 : TEST(KllFuncsTest, InitInt64ProducesNonEmptySketch) {
      34            1 :   auto sketch = InitInt64({1, 2, 3, 4, 5});
      35            2 :   ASSERT_TRUE(sketch.ok()) << sketch.status();
      36            1 :   EXPECT_FALSE(sketch->is_null());
      37            1 :   EXPECT_EQ(sketch->type_kind(), ::googlesql::TYPE_BYTES);
      38            1 :   EXPECT_GT(sketch->bytes_value().size(), 4u);
      39            1 : }
      40              : 
      41            1 : TEST(KllFuncsTest, ExtractInt64HalvesOnSmallSet) {
      42            1 :   auto sketch = InitInt64({1, 2, 3, 4, 5});
      43            2 :   ASSERT_TRUE(sketch.ok()) << sketch.status();
      44              : 
      45            1 :   auto halves = KllQuantilesExtractInt64Scalar({*sketch, Value::Int64(2)},
      46            1 :                                                Int64ArrayType());
      47            2 :   ASSERT_TRUE(halves.ok()) << halves.status();
      48            1 :   ASSERT_TRUE(halves->type()->IsArray());
      49            1 :   ASSERT_EQ(halves->elements().size(), 3u);
      50            1 :   EXPECT_EQ(halves->elements()[0].int64_value(), 1);
      51            1 :   EXPECT_EQ(halves->elements()[1].int64_value(), 3);
      52            1 :   EXPECT_EQ(halves->elements()[2].int64_value(), 5);
      53            1 : }
      54              : 
      55            1 : TEST(KllFuncsTest, ExtractPointInt64P80) {
      56            1 :   auto sketch = InitInt64({1, 2, 3, 4, 5});
      57            2 :   ASSERT_TRUE(sketch.ok()) << sketch.status();
      58              : 
      59            1 :   auto p80 = KllQuantilesExtractPointInt64Scalar({*sketch, Value::Double(0.8)});
      60            2 :   ASSERT_TRUE(p80.ok()) << p80.status();
      61            1 :   EXPECT_EQ(p80->int64_value(), 4);
      62            1 : }
      63              : 
      64            1 : TEST(KllFuncsTest, MergePartialRoundTrip) {
      65            1 :   auto sketch_a = InitInt64({1, 2, 3});
      66            1 :   auto sketch_b = InitInt64({4, 5});
      67            2 :   ASSERT_TRUE(sketch_a.ok()) << sketch_a.status();
      68            2 :   ASSERT_TRUE(sketch_b.ok()) << sketch_b.status();
      69              : 
      70            1 :   auto merged = KllQuantilesMergePartialAggregate({{*sketch_a, *sketch_b}});
      71            2 :   ASSERT_TRUE(merged.ok()) << merged.status();
      72            1 :   EXPECT_FALSE(merged->is_null());
      73              : 
      74            1 :   auto halves = KllQuantilesExtractInt64Scalar({*merged, Value::Int64(2)},
      75            1 :                                                Int64ArrayType());
      76            2 :   ASSERT_TRUE(halves.ok()) << halves.status();
      77            1 :   EXPECT_EQ(halves->elements()[0].int64_value(), 1);
      78            1 :   EXPECT_EQ(halves->elements()[2].int64_value(), 5);
      79            1 : }
      80              : 
      81            1 : TEST(KllFuncsTest, MergeInt64AggregateReturnsHalves) {
      82            1 :   auto sketch_a = InitInt64({1, 2, 3});
      83            1 :   auto sketch_b = InitInt64({4, 5});
      84            2 :   ASSERT_TRUE(sketch_a.ok()) << sketch_a.status();
      85            2 :   ASSERT_TRUE(sketch_b.ok()) << sketch_b.status();
      86              : 
      87            1 :   auto merged = KllQuantilesMergeInt64Aggregate(
      88            1 :       {{*sketch_a, *sketch_b}, {Value::Int64(2)}}, Int64ArrayType());
      89            2 :   ASSERT_TRUE(merged.ok()) << merged.status();
      90            1 :   ASSERT_TRUE(merged->type()->IsArray());
      91            1 :   EXPECT_EQ(merged->elements().size(), 3u);
      92            1 :   EXPECT_EQ(merged->elements()[0].int64_value(), 1);
      93            1 :   EXPECT_EQ(merged->elements()[2].int64_value(), 5);
      94            1 : }
      95              : 
      96            1 : TEST(KllFuncsTest, InitWithWeightSkewsMedian) {
      97            1 :   std::vector<Value> inputs = {Value::Int64(1),
      98            1 :                                Value::Int64(2),
      99            1 :                                Value::Int64(3),
     100            1 :                                Value::Int64(4),
     101            1 :                                Value::Int64(5)};
     102            1 :   std::vector<Value> weights = {Value::Int64(1),
     103            1 :                                 Value::Int64(3),
     104            1 :                                 Value::Int64(1),
     105            1 :                                 Value::Int64(1),
     106            1 :                                 Value::Int64(1)};
     107            1 :   auto sketch = KllQuantilesInitInt64Values(inputs, 1000, &weights);
     108            2 :   ASSERT_TRUE(sketch.ok()) << sketch.status();
     109              : 
     110            1 :   auto halves = KllQuantilesExtractInt64Scalar({*sketch, Value::Int64(2)},
     111            1 :                                                Int64ArrayType());
     112            2 :   ASSERT_TRUE(halves.ok()) << halves.status();
     113            1 :   EXPECT_EQ(halves->elements()[0].int64_value(), 1);
     114            1 :   EXPECT_EQ(halves->elements()[1].int64_value(), 2);
     115            1 :   EXPECT_EQ(halves->elements()[2].int64_value(), 5);
     116            1 : }
     117              : 
     118            1 : TEST(KllFuncsTest, InvalidSketchBytesRejected) {
     119            1 :   Value invalid_sketch = Value::Bytes("not-a-sketch");
     120            1 :   Value num_quantiles = Value::Int64(2);
     121            1 :   auto got = KllQuantilesExtractInt64Scalar({invalid_sketch, num_quantiles},
     122            1 :                                             Int64ArrayType());
     123            1 :   EXPECT_FALSE(got.ok());
     124            1 : }
     125              : 
     126              : }  // namespace functions
     127              : }  // namespace semantic
     128              : }  // namespace engine
     129              : }  // namespace backend
     130              : }  // namespace bigquery_emulator