Memory limited FTS training

from tempfile import TemporaryDirectory

import pyarrow as pa

import lance
from lance._datagen import rand_batches
import memtest

def measure_peak_memory(
    data_size: int,
    index_type: str,
) -> int:
    if index_type == "btree":
        schema = pa.schema([pa.field("col", pa.string())])
        data = rand_batches(schema, num_batches=data_size // (1024 * 1024), batch_size_bytes=1024 * 1024)
    elif index_type == "bitmap":
        schema = pa.schema([pa.field("col", pa.string(), metadata={b"lance-datagen:cardinality": b"1000"})])
        data = rand_batches(schema, num_batches=data_size // (1024 * 1024), batch_size_bytes=1024 * 1024)
    elif index_type == "inverted":
        schema = pa.schema([pa.field("col", pa.string(), metadata={"lance-datagen:content-type": "sentence"})])
        data = rand_batches(schema, num_batches=data_size // (1024 * 1024), batch_size_bytes=1024 * 1024)
    else:
        raise ValueError(f"Unsupported index type: {index_type}")
    with TemporaryDirectory() as tmpdir:
        ds = lance.write_dataset(data, tmpdir)
        with memtest.track() as get_stats:
            if index_type == "btree":
                ds.create_scalar_index("col", "btree", replace=True)
            elif index_type == "bitmap":
                ds.create_scalar_index("col", "bitmap", replace=True)
            elif index_type == "inverted":
                ds.create_scalar_index("col", "INVERTED", with_position=True, replace=True)
            stats = get_stats()
    return stats["peak_bytes"]

for size_mb in [1, 10, 100, 1024, 5 * 1024, 10 * 1024]:
    size_bytes = size_mb * 1024 * 1024
    for index in ["inverted"]:
        peak_mem = measure_peak_memory(size_bytes, index)
        print(f"Data Size: {size_mb} MB, Index: {index}, Peak Memory: {peak_mem / (1024 * 1024):.2f} MB")