Bitpacking TaxID index

[fawaz-dabbaghieh]

What's New

Bitpacking, and how it works

I have implemented bitpacking for the TaxID index for Metabuli. The idea is that instead of storing 32 bit raw integers for each metamer, we can bitpack them because the number of unique IDs does not need the complete range of the 32bits. The current implementation can pack based on several cases (8, 10, 12, 16 or 21) bits, so depending on what the maximum tax ID, it chooses which level of packing to do.

This implementation basically packs them always into uint64_t, for example, if the packing rate is 10, then we put 5 integers in a 64 bit chunk and waste 4 bit, and so on.

This implementation is backward compatible, so if we give Metabuli an old database, it works as intended, it looks at db.parameters and if there is no mention about bitpacking, then it assumes it's the original 32 bit integers and continues with a plain buffer. If the db.parameters mentions the bitpacking rate, then that is used and a special reader that can return the individual integers from the 64 bit with masking and shifting. For example, for a 12 bit packing, we do something like this:

constexpr uint8_t ID_BITS = 12;
constexpr uint64_t MASK = (uint64_t{1} << ID_BITS) - 1;  // 0xFFF

uint64_t word = packedWord; // read 64 bits

uint32_t id = (word >> (lane * ID_BITS)) & MASK;  // where ID_BITS is the bitpacking rate and lane is the slot number

With this lane concept, we can easily seek a certain element with its index number with

size_t wordOffset = logicalOffset / idsPerWord;  // which uint64_t word
size_t lane = logicalOffset % idsPerWord;  // which ID slot inside that word

Compact old Databases subcommand

Also added the subcommand compact-info which basically takes an old databases and compacts the info index, so we don't need to rebuild the database all over again, we can just compact in place when possible

Testing

The regression test seems to run. Further, I tested on GTDB release RS80 (it's an old one) which contains around 3000 genomes or so, I built the database with the original Metabuli and the bitpacking one, the resulting info index was around 34% smaller in size, with negligible increase in database building time (around 5% more). For classification, I simulated a million single-end short reads from these genomes and classified them, both version ran in about 16 seconds, therefore, there was no noticible decrease in time after bitpacking the IDs.

[fawaz-dabbaghieh]

It does not pass the regression tests now because the regression test is checking byte-for-byte the info index which is now smaller. I added a --pack-info flag that can be set to 0 to deactivate the bitpacking. If this flag is used, then the regression test passes. We probably need to update the regression test to make this work. Probably have regression tests for both original and bitpacked version

[fawaz-dabbaghieh]

I updated the run_build.sh regression test to now test both unpacked and packed versions of Metabuli. This is an updated bash script for the regression test.

#!/bin/sh -ex

BASE="${BASE:-$(cd "$(dirname "$0")/.." && pwd)}"
DATADIR="${DATADIR:-${BASE}/data}"

IN_INPUT="${DATADIR}/files_for_db/inclusion.txt"
EX_INPUT="${DATADIR}/files_for_db/exclusion.txt"
ACC2TAXID="${DATADIR}/files_for_db/test.accession2taxid"
TAXONOMY="${DATADIR}/files_for_db/taxonomy"
REF="${DATADIR}/reference/"

cd "${DATADIR}/files_for_db/"

build_db() {
    "${METABULI}" build "$1" "$2" "${ACC2TAXID}" --threads 1 --max-ram 8 --taxonomy-path "${TAXONOMY}" --pack-info "$3"
}

build_custom_db() {
    "${METABULI}" build "$1" "$2" "${ACC2TAXID}" --threads 1 --max-ram 8 --taxonomy-path "${TAXONOMY}" \
        --space-mask 11101110111 --custom-metamer "${DATADIR}/files_for_db/reduced_15_pattern.txt" --pack-info "$3"
}

build_db "${RESULTS}/inclusionDB" "${IN_INPUT}" 0

build_db "${RESULTS}/exclusionDB" "${EX_INPUT}" 0

build_custom_db "${RESULTS}/inclusionDB-custom" "${IN_INPUT}" 0

build_custom_db "${RESULTS}/exclusionDB-custom" "${EX_INPUT}" 0

build_db "${RESULTS}/inclusionDB-packed" "${IN_INPUT}" 1

build_db "${RESULTS}/exclusionDB-packed" "${EX_INPUT}" 1

build_custom_db "${RESULTS}/inclusionDB-custom-packed" "${IN_INPUT}" 1

build_custom_db "${RESULTS}/exclusionDB-custom-packed" "${EX_INPUT}" 1

cd "${BASE}"

## Evaluate

## 1. Compare built DBs with expected DBs.
CMP_NUM=32
ID_CNT=0

count_pass() {
    ID_CNT=$((ID_CNT+1))
}

compare_file() {
    if diff -q "$1" "$2"; then
        count_pass
    fi
}

info_count() {
    BYTES=$(wc -c < "$1")
    echo $((BYTES / 4))
}

metadata_value() {
    awk -F '\t' -v key="$2" '$1 == key { print $2; found=1; exit } END { if (!found) exit 1 }' "$1/db.parameters"
}

check_raw_metadata() {
    COUNT=$(info_count "$2/info")
    FORMAT=$(metadata_value "$1" "Info_format") || FORMAT=""
    BITS=$(metadata_value "$1" "Info_id_bits") || BITS=""
    ID_COUNT=$(metadata_value "$1" "Info_id_count") || ID_COUNT=""
    if [ "${FORMAT}" = "uint32" ] && [ "${BITS}" = "32" ] && [ "${ID_COUNT}" = "${COUNT}" ]; then
        count_pass
    fi
}

check_packed_metadata() {
    COUNT=$(info_count "$2/info")
    FORMAT=$(metadata_value "$1" "Info_format") || FORMAT=""
    BITS=$(metadata_value "$1" "Info_id_bits") || BITS=""
    ID_COUNT=$(metadata_value "$1" "Info_id_count") || ID_COUNT=""
    case "${BITS}" in
        ''|*[!0-9]*)
            return
            ;;
    esac
    if [ "${FORMAT}" = "packed_uint" ] && [ "${BITS}" -lt 32 ] && [ "${ID_COUNT}" = "${COUNT}" ]; then
        count_pass
    fi
}

compare_decoded_info() {
    COUNT=$(info_count "$2/info")
    ACTUAL_INFO="${RESULTS}/$(basename "$1").info.txt"
    EXPECTED_INFO="${RESULTS}/$(basename "$1").expected-info.txt"

    "${METABULI}" printInfo "$1/info" --info-end "${COUNT}" | awk '/^[0-9][0-9]*$/ { print }' > "${ACTUAL_INFO}"
    "${METABULI}" printInfo "$2/info" --info-end "${COUNT}" | awk '/^[0-9][0-9]*$/ { print }' > "${EXPECTED_INFO}"
    if diff -q "${ACTUAL_INFO}" "${EXPECTED_INFO}"; then
        count_pass
    fi
}

write_report() {
    awk -v actual="$ID_CNT" -v target="$CMP_NUM" \
        'BEGIN { print (actual == target) ? "GOOD" : "BAD"; print "Expected: ", target; print "Actual: ", actual; }' \
        > "${RESULTS}.report"
}

## Legacy uint32 builds: keep the existing byte-for-byte reference checks.
compare_file "${RESULTS}/inclusionDB/diffIdx" "${REF}/in/diffIdx"
compare_file "${RESULTS}/inclusionDB/info" "${REF}/in/info"
compare_file "${RESULTS}/inclusionDB/split" "${REF}/in/split"
check_raw_metadata "${RESULTS}/inclusionDB" "${REF}/in"

compare_file "${RESULTS}/inclusionDB-custom/diffIdx" "${REF}/in-custom/diffIdx"
compare_file "${RESULTS}/inclusionDB-custom/info" "${REF}/in-custom/info"
compare_file "${RESULTS}/inclusionDB-custom/split" "${REF}/in-custom/split"
check_raw_metadata "${RESULTS}/inclusionDB-custom" "${REF}/in-custom"

compare_file "${RESULTS}/exclusionDB/diffIdx" "${REF}/ex/diffIdx"
compare_file "${RESULTS}/exclusionDB/info" "${REF}/ex/info"
compare_file "${RESULTS}/exclusionDB/split" "${REF}/ex/split"
check_raw_metadata "${RESULTS}/exclusionDB" "${REF}/ex"

compare_file "${RESULTS}/exclusionDB-custom/diffIdx" "${REF}/ex-custom/diffIdx"
compare_file "${RESULTS}/exclusionDB-custom/info" "${REF}/ex-custom/info"
compare_file "${RESULTS}/exclusionDB-custom/split" "${REF}/ex-custom/split"
check_raw_metadata "${RESULTS}/exclusionDB-custom" "${REF}/ex-custom"

## Packed builds: diffIdx/split must stay identical; info is compared after decoding.
compare_file "${RESULTS}/inclusionDB-packed/diffIdx" "${REF}/in/diffIdx"
compare_file "${RESULTS}/inclusionDB-packed/split" "${REF}/in/split"
compare_decoded_info "${RESULTS}/inclusionDB-packed" "${REF}/in"
check_packed_metadata "${RESULTS}/inclusionDB-packed" "${REF}/in"

compare_file "${RESULTS}/inclusionDB-custom-packed/diffIdx" "${REF}/in-custom/diffIdx"
compare_file "${RESULTS}/inclusionDB-custom-packed/split" "${REF}/in-custom/split"
compare_decoded_info "${RESULTS}/inclusionDB-custom-packed" "${REF}/in-custom"
check_packed_metadata "${RESULTS}/inclusionDB-custom-packed" "${REF}/in-custom"

compare_file "${RESULTS}/exclusionDB-packed/diffIdx" "${REF}/ex/diffIdx"
compare_file "${RESULTS}/exclusionDB-packed/split" "${REF}/ex/split"
compare_decoded_info "${RESULTS}/exclusionDB-packed" "${REF}/ex"
check_packed_metadata "${RESULTS}/exclusionDB-packed" "${REF}/ex"

compare_file "${RESULTS}/exclusionDB-custom-packed/diffIdx" "${REF}/ex-custom/diffIdx"
compare_file "${RESULTS}/exclusionDB-custom-packed/split" "${REF}/ex-custom/split"
compare_decoded_info "${RESULTS}/exclusionDB-custom-packed" "${REF}/ex-custom"
check_packed_metadata "${RESULTS}/exclusionDB-custom-packed" "${REF}/ex-custom"

## If ID_CNT is equal to CMP_NUM, the test is passed.
if [ "${ID_CNT}" -eq "${CMP_NUM}" ]; then
    echo "Test passed."
    write_report
else
    echo "Test failed."
    write_report
    exit 1
fi