4646_DEFAULT_DENSE_METRIC = "euclidean"
4747_DEFAULT_SPARSE_METRIC = "ip"
4848_MAX_ANISOTROPY_DIM = 50000
49- _MAX_ANISOTROPY_BYTES = 256 * 1024 * 1024
49+ # One byte budget shared by every sampled-memory guard: the dense float64
50+ # sample, the dense expansion of a sparse sample, and the retained sparse
51+ # sample rows below. Sparse samples must obey the same budget as dense ones
52+ # even though their rows carry a variable number of entries.
53+ _SAMPLE_MEMORY_BUDGET_BYTES = 256 * 1024 * 1024
54+ # A retained sparse sample keeps one (uint32, float32) copy per row and later
55+ # expands into one Python dict entry per finite value; estimate both together
56+ # from nnz with interpreter object overhead so nnz-heavy rows trip the budget
57+ # before any copy or dict expansion happens instead of exhausting memory.
58+ _SPARSE_SAMPLE_ENTRY_BYTES = 32
59+ _SPARSE_SAMPLE_ROW_OVERHEAD_BYTES = 128
5060_MAX_SPARSE_STATS_DIMENSION = 10_000_000
5161_PCA_SVD_MAX_DIM = 4096
5262# Exact norm quantiles/outliers and exact duplicate groups need per-row state.
@@ -279,10 +289,35 @@ def _dense_sample_size(count: int, dim: int, requested: int) -> int:
279289 dataset_limited = min (int (requested ), int (count ))
280290 if dim <= 0 :
281291 return int (dataset_limited )
282- max_rows = _MAX_ANISOTROPY_BYTES // (int (dim ) * 8 )
292+ max_rows = _SAMPLE_MEMORY_BUDGET_BYTES // (int (dim ) * 8 )
283293 return int (min (dataset_limited , max_rows ))
284294
285295
296+ def _sparse_sample_row_bytes (ids : np .ndarray ) -> int :
297+ """Estimate the bytes one sampled sparse row keeps after retention and
298+ after the Python dict expansion performed by :func:`_pairwise_sparse`."""
299+ return (
300+ _SPARSE_SAMPLE_ROW_OVERHEAD_BYTES + int (ids .size ) * _SPARSE_SAMPLE_ENTRY_BYTES
301+ )
302+
303+
304+ def _sparse_sample_budget_reason (
305+ next_row_bytes : int , retained_bytes : int , retained_rows : int
306+ ) -> str :
307+ """Return why retention of the sampled sparse rows stopped."""
308+ return (
309+ "sparse sample limited to %d retained rows: continuing would exceed the "
310+ "%d-byte sample memory budget (estimated %d total bytes so far plus the "
311+ "pending row); the sampled sparse pairwise and distribution metrics are "
312+ "unavailable, while streaming and full-split metrics are still reported"
313+ % (
314+ retained_rows ,
315+ _SAMPLE_MEMORY_BUDGET_BYTES ,
316+ int (retained_bytes ) + int (next_row_bytes ),
317+ )
318+ )
319+
320+
286321def _full_split_state_reason (
287322 metric_name : str , count : int , bytes_per_vector : int
288323) -> Optional [str ]:
@@ -531,6 +566,11 @@ def _profile_sparse(
531566 vectors_with_invalid = 0
532567 all_zero_vectors = 0
533568 sample_records : List [Tuple [int , np .ndarray , np .ndarray ]] = []
569+ # Sample retention is bounded by estimated bytes, not by the row count:
570+ # valid sparse files can carry so many entries per vector that the full
571+ # row count would exhaust memory before any report is produced.
572+ retained_sample_bytes = 0
573+ sample_budget_reason : Optional [str ] = None
534574
535575 for index , raw_ids , raw_values in dataset .iter_sparse_records (
536576 max_chunk_bytes = config .max_chunk_bytes
@@ -582,10 +622,21 @@ def _profile_sparse(
582622 "<f4" , copy = False
583623 ).tobytes ()
584624 hashes [index ] = np .frombuffer (_digest (payload ), dtype = np .uint8 )
585- if index in sample_set :
586- sample_records .append ((index , ids .copy (), values .copy ()))
625+ if sample_budget_reason is None and index in sample_set :
626+ row_bytes = _sparse_sample_row_bytes (ids )
627+ if retained_sample_bytes + row_bytes > _SAMPLE_MEMORY_BUDGET_BYTES :
628+ sample_budget_reason = _sparse_sample_budget_reason (
629+ row_bytes , retained_sample_bytes , len (sample_records )
630+ )
631+ warnings .append (sample_budget_reason )
632+ else :
633+ retained_sample_bytes += row_bytes
634+ sample_records .append ((index , ids .copy (), values .copy ()))
587635
588- if len (sample_records ) != len (sample_set ):
636+ if (
637+ sample_budget_reason is None
638+ and len (sample_records ) != len (sample_set )
639+ ):
589640 raise ValueError ("sparse sampling failed: not all sampled records were found" )
590641 sample_records .sort (key = lambda item : item [0 ])
591642 pair_records = [
@@ -621,20 +672,27 @@ def _profile_sparse(
621672 reason = dimension_stats_reason ,
622673 dim = int (dim ),
623674 )
624- if any (
625- not bool (np .isfinite (values ).all ()) for _ , _ , values in sample_records
626- ):
627- warnings .append (_INVALID_SAMPLE_WARNING )
628- pair_matrix = _pairwise_sparse (pair_records , metric )
629- pair_values = _upper_triangle (pair_matrix )
630- duplicate_pair_values = pair_values
631- duplicate_metric = metric
632- if metric == "ip" :
633- duplicate_pair_values = _upper_triangle (
634- _pairwise_sparse (pair_records , "euclidean" )
635- )
636- duplicate_metric = "euclidean"
637- sample_matrix = _sparse_sample_to_dense (sample_records , dim )
675+ if sample_budget_reason is None :
676+ if any (
677+ not bool (np .isfinite (values ).all ()) for _ , _ , values in sample_records
678+ ):
679+ warnings .append (_INVALID_SAMPLE_WARNING )
680+ pair_matrix = _pairwise_sparse (pair_records , metric )
681+ pair_values = _upper_triangle (pair_matrix )
682+ duplicate_pair_values = pair_values
683+ duplicate_metric = metric
684+ if metric == "ip" :
685+ duplicate_pair_values = _upper_triangle (
686+ _pairwise_sparse (pair_records , "euclidean" )
687+ )
688+ duplicate_metric = "euclidean"
689+ sample_matrix = _sparse_sample_to_dense (sample_records , dim )
690+ else :
691+ pair_matrix = np .zeros ((0 , 0 ), dtype = np .float64 )
692+ pair_values = np .empty (0 , dtype = np .float64 )
693+ duplicate_pair_values = pair_values
694+ duplicate_metric = metric
695+ sample_matrix = None
638696
639697 metrics : Dict [str , Any ] = {
640698 "invalid_values" : _invalid_metric (
@@ -662,6 +720,7 @@ def _profile_sparse(
662720 duplicate_metric ,
663721 int (pair_matrix .shape [0 ]),
664722 duplicates_reason ,
723+ near_unavailable_reason = sample_budget_reason ,
665724 ),
666725 }
667726 metrics .update (
@@ -674,6 +733,7 @@ def _profile_sparse(
674733 rng = rng ,
675734 warnings = warnings ,
676735 sparse_pairwise = True ,
736+ unavailable_reason = sample_budget_reason ,
677737 )
678738 )
679739 return metrics
@@ -1058,6 +1118,7 @@ def _duplicates_metric(
10581118 near_metric : str ,
10591119 near_sample_size : int ,
10601120 unavailable_reason : Optional [str ] = None ,
1121+ near_unavailable_reason : Optional [str ] = None ,
10611122) -> Dict [str , Any ]:
10621123 if hashes is None :
10631124 exact = {
@@ -1107,7 +1168,7 @@ def _duplicates_metric(
11071168 "indicator covers the bounded pairwise sample"
11081169 )
11091170 near = _near_duplicate_summary (
1110- pair_values , epsilon , near_metric , near_sample_size
1171+ pair_values , epsilon , near_metric , near_sample_size , near_unavailable_reason
11111172 )
11121173 return metric_entry (
11131174 DEFINITIONS ["duplicates" ],
@@ -1120,12 +1181,16 @@ def _duplicates_metric(
11201181
11211182
11221183def _near_duplicate_summary (
1123- pair_values : np .ndarray , epsilon : float , metric : str , sample_size : int
1184+ pair_values : np .ndarray ,
1185+ epsilon : float ,
1186+ metric : str ,
1187+ sample_size : int ,
1188+ unavailable_reason : Optional [str ] = None ,
11241189) -> Dict [str , Any ]:
11251190 if pair_values .size == 0 :
11261191 return {
11271192 "available" : False ,
1128- "reason" : "fewer than two sampled vectors" ,
1193+ "reason" : unavailable_reason or "fewer than two sampled vectors" ,
11291194 "scope" : "sample" ,
11301195 "sample_size" : int (sample_size ),
11311196 "pair_count" : 0 ,
@@ -1170,9 +1235,26 @@ def _sampled_metrics(
11701235 rng : np .random .Generator ,
11711236 warnings : List [str ],
11721237 sparse_pairwise : bool = False ,
1238+ unavailable_reason : Optional [str ] = None ,
11731239) -> Dict [str , Any ]:
11741240 sample_size = int (pair_matrix .shape [0 ])
11751241 pair_count = int (pair_values .size )
1242+ if unavailable_reason is not None :
1243+ return {
1244+ name : metric_entry (
1245+ DEFINITIONS [name ],
1246+ "sample" ,
1247+ sample_size ,
1248+ available = False ,
1249+ reason = unavailable_reason ,
1250+ )
1251+ for name in (
1252+ "pairwise_distances" ,
1253+ "nearest_neighbor_distances" ,
1254+ "anisotropy" ,
1255+ "clustering" ,
1256+ )
1257+ }
11761258 pairwise = metric_entry (
11771259 DEFINITIONS ["pairwise_distances" ],
11781260 "sample" ,
@@ -1482,7 +1564,7 @@ def _sparse_sample_to_dense(
14821564 if dim <= 0 or not records :
14831565 return None
14841566 rows = len (records )
1485- if rows * dim * 8 > _MAX_ANISOTROPY_BYTES or dim > _MAX_ANISOTROPY_DIM :
1567+ if rows * dim * 8 > _SAMPLE_MEMORY_BUDGET_BYTES or dim > _MAX_ANISOTROPY_DIM :
14861568 return None
14871569 matrix = np .zeros ((rows , dim ), dtype = np .float64 )
14881570 for position , (_ , ids , values ) in enumerate (records ):
0 commit comments