kll_sketch_merge.sqlx

config { hasOutput: true }
/*
 * Copyright 2024 Google LLC
 *
 * Licensed 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.
 */

CREATE OR REPLACE AGGREGATE FUNCTION ${self()}(sketch BYTES, k INT64 NOT AGGREGATE)
RETURNS BYTES
LANGUAGE js
OPTIONS (
  library=["${JS_BUCKET}/kll_sketch.mjs"],
  description='''Aggregates KLL sketches, k arg, performs a union op and returns a merged KLL sketch.
For more details: https://datasketches.apache.org/docs/KLL/KLLSketch.html'''
) AS '''
import ModuleFactory from "${JS_BUCKET}/kll_sketch.mjs";
var Module = await ModuleFactory();

// Helper definitions

// shared buffer for serialization and deserialization
var BUFFER = {
  ptr: 0,
  size: 0,
};
// This function returns expected size of sketch when serialized
function maxSize(sketch) {
  // https://github.com/apache/datasketches-cpp/issues/4
  // https://datasketches.apache.org/docs/KLL/KLLAccuracyAndSize.html
  return Module._kll_sketch_serialized_size_bytes(sketch);
}

function requireBuffer(size) {
  if (BUFFER.size < size) {
    releaseBuffer();
    BUFFER.ptr = Module._malloc(size);
    BUFFER.size = size;
  }
  return BUFFER;
}

function releaseBuffer() {
  if (BUFFER.ptr) {
    Module._free(BUFFER.ptr);
  }
  BUFFER.ptr = 0;
  BUFFER.size = 0;
}

function destroyState(state) {
  if (state.sketch) {
    Module._kll_sketch_destroy(state.sketch);
    state.sketch = 0;
  }
  state.serialized = null;
}

function updateSketch(sketch, bytes) {
  var buffer = requireBuffer(bytes.length);
  Module.HEAPU8.subarray(buffer.ptr, buffer.ptr + buffer.size).set(bytes);
  var deserialized_sketch = Module._kll_sketch_deserialize(buffer.ptr, bytes.length);
  Module._merge_sketch(sketch, deserialized_sketch);
  Module._kll_sketch_destroy(deserialized_sketch);
}

// UDAF interface

export function initialState(k) {
  k = Module._clamp_k(k);
  return {
    sketch: Module._kll_sketch_initialize(k),
    k: k,
    serialized: null
  };
}

// Upstream: initialState, aggregate, merge, deserialize
export function aggregate(state, bytes) {
  if (!state.sketch) {
    state.sketch = Module._kll_sketch_initialize(state.k);
  }
  updateSketch(state.sketch, bytes);
}

// Upstream: initialState, aggregate, merge, deserialize
export function serialize(state) {
  try {
    var buffer = 0;
    var len = 0;
    if ( state.sketch && state.serialized) {
      updateSketch(state.sketch, state.serialized);
      state.serialized = null;
      buffer = requireBuffer(maxSize(state.sketch));
      len = Module._kll_sketch_serialize(
          state.sketch, buffer.ptr, buffer.size);
    } else if (state.sketch) {
      buffer = requireBuffer(maxSize(state.sketch));
      len = Module._kll_sketch_serialize(
          state.sketch, buffer.ptr, buffer.size);
    } else if (state.serialized) {
      return {
        k: state.k,
        bytes: state.serialized
      }
    } else {
      throw new Error(
          "Unexpected state in serialization " + JSON.stringify(state));
    }

    return {
      k: state.k,
      bytes: Module.HEAPU8.slice(buffer.ptr, buffer.ptr + len)
    }
  } finally {
    // clean up kll sketch
    Module._kll_sketch_destroy(state.sketch);
    state.sketch = 0;
    state.serialized = null;
  }
}

// Keeping state same as agrgegation state
export function deserialize(serialized) {
  return {
    sketch: 0,
    k: serialized.k,
    serialized: serialized.bytes
  };
}

// Assuming Merge can only be called after deserialize
export function merge(state, other_state) {
  if (!state.sketch) {
    state.sketch = Module._kll_sketch_initialize(state.k);
  }

  if (other_state.sketch ) {
    throw new Error("Did not expect sketch in other state");
  }

  if (state.serialized) {
    // consume it
    updateSketch(state.sketch, state.serialized);
    state.serialized = null;
  }

  if (other_state.serialized) {
    updateSketch(state.sketch, other_state.serialized);
    other_state.serialized = null;
  } else {
    throw new Error("Expected serialized sketch on other_state");
  }
}


export function finalize(state) {
  var result = serialize(state);
  return result.bytes;
}


''';