Skip to content

fix(ci): enlarge mem pool only around merge_insert index build #8191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Xuanwo merged 2 commits into from
Aug 4, 2026
Merged

fix(ci): enlarge mem pool only around merge_insert index build #8191

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
2f83943
fix: raise sort spill reservation ceiling for default memory pool
Xuanwo Aug 3, 2026
89b92e2
fix(ci): enlarge mem pool only around merge_insert index build
Xuanwo Aug 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 28 additions & 3 deletions python/python/ci_benchmarks/datagen/merge_insert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,9 @@

from __future__ import annotations

import os
from contextlib import contextmanager

import lance
import numpy as np
import pyarrow as pa
Expand All @@ -53,6 +56,13 @@
# Tag marking the pristine version that benchmarks restore to.
BASE_TAG = "merge_insert_base"

# BTREE training sorts the full column under DataFusion's FairSpillPool. The
# default 150 MiB pool is tight for 10M 32-byte string keys (uuid columns):
# ExternalSorterMerge has peaked near the pool edge and aborted datagen with
# ResourcesExhausted. Index build is one-shot setup, not a measured path, so
# temporarily enlarge the pool only around create_scalar_index.
_INDEX_BUILD_MEM_POOL_SIZE = str(1024 * 1024 * 1024)

NARROW_NUM_ROWS = 10_000_000
NARROW_ROWS_PER_FRAGMENT = 1_000_000

Expand Down Expand Up @@ -272,6 +282,20 @@ def _already_generated(uri: str, expected_rows: int) -> bool:
return ds.checkout_version(base_version).count_rows() == expected_rows


@contextmanager
def _enlarged_mem_pool_for_index_build():
"""Raise LANCE_MEM_POOL_SIZE only for the duration of index training."""
previous = os.environ.get("LANCE_MEM_POOL_SIZE")
os.environ["LANCE_MEM_POOL_SIZE"] = _INDEX_BUILD_MEM_POOL_SIZE
try:
yield
finally:
if previous is None:
os.environ.pop("LANCE_MEM_POOL_SIZE", None)
else:
os.environ["LANCE_MEM_POOL_SIZE"] = previous


def _gen(
name: str,
data,
Expand All @@ -294,9 +318,10 @@ def _gen(
max_rows_per_file=rows_per_fragment,
max_rows_per_group=min(rows_per_fragment, 100_000),
)
for column in indexed_columns:
LOGGER.info("Building BTREE index on %s.%s", name, column)
ds.create_scalar_index(column, "BTREE")
with _enlarged_mem_pool_for_index_build():
for column in indexed_columns:
LOGGER.info("Building BTREE index on %s.%s", name, column)
ds.create_scalar_index(column, "BTREE")
_tag_base(ds)
return ds

Expand Down
Loading