MetonaSqlark/src/engine/aria/index/sstable_builder.ts

/**
 * AriaEngine SSTable Builder — 构建有序字符串表
 * @module engine/aria/index/sstable_builder
 *
 * 将排序后的 key-value 数据写入 SSTable 格式。
 *
 * SSTable 文件布局:
 * ┌──────────────────────────────────────────────┐
 * │  Data Block 0                               │
 * │  Data Block 1                               │
 * │  ...                                        │
 * │  Index Block (block offset → key range)     │
 * │  Bloom Filter                               │
 * │  Footer (32 bytes)                          │
 * │    - index_offset  (u32)                    │
 * │    - index_size    (u32)                    │
 * │    - bloom_offset  (u32)                    │
 * │    - bloom_size    (u32)                    │
 * │    - bloom_hash_count (u32)                 │
 * │    - entry_count   (u32)                    │
 * │    - magic_number  (u32, 0x53535442 ="SSTB")│
 * │    - checksum      (u32)                    │
 * └──────────────────────────────────────────────┘
 */

import { BloomFilter } from './bloom';
import type { IndexEntry, DataBlock } from '../types';

const SSTABLE_MAGIC = 0x53535442; // "SSTB"
const SSTABLE_FOOTER_SIZE = 32;

// ---------------------------------------------------------------------------
//  SSTableBuilder
// ---------------------------------------------------------------------------

export class SSTableBuilder {
  private entries: [string, Record<string, unknown>][] = [];
  private currentBlock: [string, Record<string, unknown>][] = [];
  private currentBlockStartKey = '';
  private blockSizeLimit: number;

  constructor(blockSizeLimit: number = 4096) {
    this.blockSizeLimit = blockSizeLimit;
  }

  /** 添加一个 key-value 条目（必须按键排序添加） */
  add(key: string, value: Record<string, unknown>): void {
    if (this.currentBlock.length === 0) {
      this.currentBlockStartKey = key;
    }

    this.currentBlock.push([key, value]);
    this.entries.push([key, value]);

    // 如果当前 Block 达到大小限制，切割
    const estimated = this.estimateBlockSize();
    if (estimated >= this.blockSizeLimit && this.currentBlock.length > 1) {
      // 当前 block 结束（不在这里切割，在 build 时统一处理）
    }
  }

  /**
   * 构建 SSTable 文件的二进制数据。
   * 返回 { data: Uint8Array, indexEntries: IndexEntry[], bloomFilter: BloomFilter }
   */
  build(): { sstableData: Uint8Array; indexEntries: IndexEntry[] } {
    const blocks = this.splitIntoBlocks();
    const bloomFilter = new BloomFilter(this.entries.length);

    // 预计算总大小
    let totalSize = 0;
    const blockOffsets: number[] = [];

    for (const block of blocks) {
      blockOffsets.push(totalSize);
      const blockSize = this.computeBlockSize(block);
      totalSize += blockSize;
    }

    // 索引块
    const indexEntries: IndexEntry[] = [];
    for (let i = 0; i < blocks.length; i++) {
      const block = blocks[i];
      const lastKey = block[block.length - 1][0];
      const blockSize = this.computeBlockSize(block);
      indexEntries.push({
        key: lastKey,
        blockOffset: blockOffsets[i],
        blockSize,
      });
    }

    const indexBlockSize = this.estimateIndexBlockSize(indexEntries);

    // 写入到 buffer
    const finalSize = totalSize + indexBlockSize + SSTABLE_FOOTER_SIZE;
    const buf = new ArrayBuffer(finalSize);
    const view = new DataView(buf);

    let offset = 0;

    // ---- Data Blocks ----
    for (const block of blocks) {
      offset = this.writeDataBlock(view, offset, block, bloomFilter);
    }

    // ---- Index Block ----
    const indexOffset = offset;
    offset = this.writeIndexBlock(view, offset, indexEntries);

    // ---- Footer ----
    const footerOffset = offset;
    view.setUint32(footerOffset, indexOffset, false);      // index_offset
    view.setUint32(footerOffset + 4, indexBlockSize, false); // index_size
    view.setUint32(footerOffset + 8, 0, false);  // bloom_offset (embedded in footer)
    view.setUint32(footerOffset + 12, 0, false); // bloom_size
    view.setUint32(footerOffset + 16, bloomFilter.getHashCount(), false);
    view.setUint32(footerOffset + 20, this.entries.length, false);
    view.setUint32(footerOffset + 24, SSTABLE_MAGIC, false);
    view.setUint32(footerOffset + 28, 0, false); // checksum (simplified: 0)

    return {
      sstableData: new Uint8Array(buf),
      indexEntries,
    };
  }

  /** 获取条目数 */
  getEntryCount(): number {
    return this.entries.length;
  }

  // -----------------------------------------------------------------------
  //  内部
  // -----------------------------------------------------------------------

  private splitIntoBlocks(): [string, Record<string, unknown>][][] {
    const blocks: [string, Record<string, unknown>][][] = [];
    let current: [string, Record<string, unknown>][] = [];

    for (const entry of this.entries) {
      current.push(entry);
      if (this.estimateBlockSizeFromEntries(current) >= this.blockSizeLimit && current.length > 1) {
        blocks.push(current.slice(0, -1));
        current = [entry];
      }
    }
    if (current.length > 0) blocks.push(current);

    return blocks;
  }

  private estimateBlockSize(): number {
    return this.estimateBlockSizeFromEntries(this.currentBlock);
  }

  private estimateBlockSizeFromEntries(entries: [string, unknown][]): number {
    let size = 0;
    for (const [key, value] of entries) {
      size += 4 + key.length + JSON.stringify(value).length;
    }
    return size;
  }

  private computeBlockSize(block: [string, unknown][]): number {
    // entryCount (u32) + 每对: keyLen(u16) + key + valueLen(u16) + value json
    let size = 4;
    for (const [key, value] of block) {
      const json = JSON.stringify(value);
      size += 2 + key.length + 2 + json.length;
    }
    return size;
  }

  private writeDataBlock(
    view: DataView,
    offset: number,
    block: [string, Record<string, unknown>][],
    bloomFilter: BloomFilter,
  ): number {
    const start = offset;

    // entry count
    view.setUint32(offset, block.length, false);
    offset += 4;

    for (const [key, value] of block) {
      const encoder = new TextEncoder();
      const keyBytes = encoder.encode(key);
      const valueBytes = encoder.encode(JSON.stringify(value));

      // key length
      view.setUint16(offset, keyBytes.length, false);
      offset += 2;
      // key
      new Uint8Array(view.buffer).set(keyBytes, offset);
      offset += keyBytes.length;
      // value length
      view.setUint16(offset, valueBytes.length, false);
      offset += 2;
      // value
      new Uint8Array(view.buffer).set(valueBytes, offset);
      offset += valueBytes.length;

      // 插入 bloom filter
      bloomFilter.insert(key);
    }

    return offset;
  }

  private estimateIndexBlockSize(entries: IndexEntry[]): number {
    // entryCount(u32) + each: keyLen(u16)+key+blockOffset(u32)+blockSize(u32)
    let size = 4;
    for (const entry of entries) {
      size += 2 + entry.key.length + 8;
    }
    return size;
  }

  private writeIndexBlock(view: DataView, offset: number, entries: IndexEntry[]): number {
    view.setUint32(offset, entries.length, false);
    offset += 4;

    for (const entry of entries) {
      const encoder = new TextEncoder();
      const keyBytes = encoder.encode(entry.key);
      view.setUint16(offset, keyBytes.length, false);
      offset += 2;
      new Uint8Array(view.buffer).set(keyBytes, offset);
      offset += keyBytes.length;
      view.setUint32(offset, entry.blockOffset, false);
      offset += 4;
      view.setUint32(offset, entry.blockSize, false);
      offset += 4;
    }

    return offset;
  }
}