import copy from abc import ABC, abstractmethod from dataclasses import dataclass from decimal import Decimal, InvalidOperation from typing import Any, Dict, List, Optional from .checkpointable_prng import CheckpointablePRNG @dataclass class ShuffleBlock(ABC): """Sampling block base class""" block_id: int # Unique ID for each block, used for random selection @abstractmethod def get_remaining_capacity(self) -> int: """Get remaining capacity""" @abstractmethod def sample(self, random_instance) -> Optional[Any]: """Sample a value""" @abstractmethod def reset(self): """Reset state""" @dataclass class TopNBlock(ShuffleBlock): """TopN sampling block - memory efficient using id/count/value format""" topn_items: List[Dict[str, Any]] # [{'id': 0, 'value': val, 'count': cnt}] _original_items: List[Dict[str, Any]] = None # Save original items for reset def __post_init__(self): # Strictly validate TopN data integrity for i, item in enumerate(self.topn_items): if not isinstance(item, dict): raise TypeError(f"TopN item {i} must be dict, got {type(item)}") if "count" not in item: raise KeyError(f"TopN item {i} missing required field 'count'") if "value" not in item: raise KeyError(f"TopN item {i} missing required field 'value'") if not isinstance(item["count"], int) or item["count"] <= 0: raise ValueError( f"TopN item {i} count must be positive integer, got {item['count']}" ) # Save original count for reset item["original_count"] = item["count"] self._original_items = copy.deepcopy(self.topn_items) def get_remaining_capacity(self) -> int: return sum(item["count"] for item in self.topn_items) def sample(self, random_instance) -> Optional[Any]: if not self.topn_items: return None # Directly select from available items randomly index = int(random_instance.uniform(0, len(self.topn_items))) chosen = self.topn_items[index] value = chosen["value"] chosen["count"] -= 1 # Remove from list if this item is used up if chosen["count"] == 0: self.topn_items.remove(chosen) return value def reset(self): # Restore to original state (including removed items) self.topn_items = copy.deepcopy(self._original_items) @dataclass class BucketBlock(ShuffleBlock): """Bucket sampling block - layered shuffle, memory efficient""" source_bucket_id: int lower_bound: float upper_bound: float count: int repeats: int = 0 column_name: Optional[str] = None data_type: Optional[str] = None raw_lower_bound: Optional[str] = None raw_upper_bound: Optional[str] = None is_integer_bucket: bool = False integer_lower_bound: Optional[int] = None integer_upper_bound: Optional[int] = None # Layered shuffle state bucket_indices: Optional[List[int]] = None remaining_count: Optional[int] = None def __post_init__(self): # Strictly validate Bucket data integrity if not isinstance(self.source_bucket_id, int): raise TypeError( f"source_bucket_id must be int, got {type(self.source_bucket_id)}" ) if not isinstance(self.count, int) or self.count <= 0: raise ValueError(f"count must be positive integer, got {self.count}") if not isinstance(self.repeats, int) or self.repeats < 0: raise ValueError( f"repeats must be non-negative integer, got {self.repeats}" ) if self.repeats > self.count: raise ValueError( f"repeats ({self.repeats}) cannot exceed count ({self.count})" ) if not isinstance(self.lower_bound, (int, float)): raise TypeError( f"lower_bound must be numeric, got {type(self.lower_bound)}" ) if not isinstance(self.upper_bound, (int, float)): raise TypeError( f"upper_bound ({self.upper_bound}) cannot exceed upper_bound ({self.upper_bound})" ) if self.lower_bound > self.upper_bound: raise ValueError( f"lower_bound ({self.lower_bound}) cannot exceed upper_bound ({self.upper_bound})" ) # Initialize integer metadata if needed if self.is_integer_bucket: if self.raw_lower_bound is None or self.raw_upper_bound is None: raise ValueError( "Integer bucket requires raw lower/upper bounds for precise sampling" ) try: lower_dec = Decimal(self.raw_lower_bound) upper_dec = Decimal(self.raw_upper_bound) except InvalidOperation as exc: raise ValueError( f"Failed to parse integer bounds for bucket {self.source_bucket_id}: " f"lower={self.raw_lower_bound}, upper={self.raw_upper_bound}" ) from exc self.integer_lower_bound = int(lower_dec) self.integer_upper_bound = int(upper_dec) if self.integer_upper_bound < self.integer_lower_bound: raise ValueError( f"Integer bucket bounds invalid: lower={self.integer_lower_bound} " f"upper={self.integer_upper_bound}" ) # Initialize state if isinstance(self.repeats, int) and self.repeats > 0: self.remaining_count = self.count // self.repeats else: self.remaining_count = self.count def get_remaining_capacity(self) -> int: if self.remaining_count is None: raise RuntimeError( "BucketBlock not properly initialized - remaining_count is None" ) return self.remaining_count def sample(self, random_instance) -> Optional[tuple]: if self.remaining_count <= 0: return None # Create and shuffle indices once if not initialized if self.bucket_indices is None: if self.repeats > 1: self.bucket_indices = list(range(self.count // self.repeats)) else: self.bucket_indices = list(range(self.count)) random_instance.shuffle(self.bucket_indices) if not self.bucket_indices: return None # Take one index index = self.bucket_indices.pop() self.remaining_count -= 1 # Compute linear interpolation value value = self._compute_linear_interpolation(index) return (value, index) # Return value and index tuple def _compute_linear_interpolation(self, index: int) -> float: """Compute bucket linear interpolation value""" if ( self.is_integer_bucket and self.integer_lower_bound is not None and self.repeats <= 1 ): value = self.integer_lower_bound + index if ( self.integer_upper_bound is not None and value > self.integer_upper_bound ): raise ValueError( f"Integer bucket {self.column_name} column in table {self.table_name} index overflow: value={value} exceeds upper=" f"{self.integer_upper_bound} (bucket_id={self.source_bucket_id})" ) return value if self.repeats > 1: # Fix: Use left-closed, right-open interval [lower_bound, upper_bound) # Remove the +1 to avoid including upper_bound range_size = self.upper_bound - self.lower_bound offset = range_size * self.repeats / self.count interpolated_value = self.lower_bound + offset * index elif self.repeats == 1: # Fix: Ensure each index maps to unique integer by using integer step range_size = self.upper_bound - self.lower_bound + 1 # Use integer division to ensure unique mapping interpolated_value = self.lower_bound + range_size / self.count * index else: # Fix: Use left-closed, right-open interval [lower_bound, upper_bound) range_size = self.upper_bound - self.lower_bound + 1 interpolated_value = self.lower_bound + range_size / self.count * index return interpolated_value def reset(self): self.remaining_count = self.count // self.repeats self.bucket_indices = None @dataclass class ShuffleResult: """Universal result type for samplers""" value: Any # Sampled output value: single column value or combination tuple source_type: ( str # BlockShuffler: "topn"/"bucket", CompositeBlockShuffler: "composite" ) debug_info: Optional[tuple] = ( None # Debug info: bucket is (bucket_id, index), others defined as needed ) shuffler_results: Optional[List["ShuffleResult"]] = ( None # For composite results: individual column ShuffleResults ) class BlockShuffler: """ No-repeat mathematical sampler based on TopN/Bucket Features: - Guarantees absolute no-repeat (compared to InverseTransformSampler's probabilistic no-repeat) - Supports mixed TopNItem and HistogramBucket input - TopN values are shuffled directly, Bucket generates (bucket_id, index) coordinates - Pure mathematical layer, output coordinates for business layer conversion """ def __init__( self, buckets: Optional[List] = None, # HistogramBucket list top_n: Optional[List] = None, # TopNItem list seed: Optional[int] = 42, max_samples: Optional[int] = None, # Maximum number of samples ): """ Initialize memory-efficient BlockShuffler Args: buckets: HistogramBucket object list, using layered shuffle top_n: TopNItem object list, using id/count/value structure seed: Random seed max_samples: Maximum number of samples """ # Validate input parameters if buckets is not None and not isinstance(buckets, list): raise TypeError(f"buckets must be list or None, got {type(buckets)}") if top_n is not None and not isinstance(top_n, list): raise TypeError(f"top_n must be list or None, got {type(top_n)}") if seed is not None and not isinstance(seed, int): raise TypeError(f"seed must be int or None, got {type(seed)}") if buckets is None and top_n is None: raise ValueError("At least one of buckets or top_n must be provided") # Save original parameters for reset self._original_seed = seed self._random = CheckpointablePRNG(seed) # Build sampling blocks - pass actual lists, no fallback to empty lists self.blocks = self._build_blocks( buckets if buckets is not None else [], top_n if top_n is not None else [] ) # Validate build result if not self.blocks: raise ValueError( "No valid blocks were created from provided buckets and top_n items" ) # Calculate total capacity self.total_capacity = sum( block.get_remaining_capacity() for block in self.blocks ) if self.total_capacity <= 0: raise ValueError("Total capacity must be positive") self.remaining_capacity = self.total_capacity # Layered shuffle: outer shuffle block_id order self.block_id_sequence = [block.block_id for block in self.blocks] self._random.shuffle(self.block_id_sequence) def _build_blocks(self, buckets, top_n_items) -> List[ShuffleBlock]: """Build memory-efficient sampling blocks - strictly validate input data, assign block_id""" blocks = [] block_id_counter = 0 # Process TopN items - strictly validate input if top_n_items: topn_data = [] for i, item in enumerate(top_n_items): # Strictly validate TopN item structure if not hasattr(item, "count"): raise AttributeError( f"TopN item {i} missing required attribute 'count'" ) if not hasattr(item, "value"): raise AttributeError( f"TopN item {i} missing required attribute 'value'" ) count = 1 # item.count # tricky: topn shouldn't be duplicated value = item.value if not isinstance(count, int) or count <= 0: raise ValueError( f"TopN item {i} count must be positive integer, got {count}" ) topn_data.append({"id": i, "value": value, "count": count}) if topn_data: blocks.append( TopNBlock(block_id=block_id_counter, topn_items=topn_data) ) block_id_counter += 1 # Process Bucket items - strictly validate input for i, bucket in enumerate(buckets): # Strictly validate Bucket structure required_attrs = ["count", "bucket_id", "lower_bound", "upper_bound"] for attr in required_attrs: if not hasattr(bucket, attr): raise AttributeError( f"Bucket {i} missing required attribute '{attr}'" ) if bucket.count <= 0: raise ValueError( f"Bucket {i} count must be positive, got {bucket.count}" ) # Get repeats, allow missing (default to 0) repeats = getattr(bucket, "repeats", 0) data_type = getattr(bucket, "data_type", None) raw_lower = getattr(bucket, "raw_lower_bound", None) raw_upper = getattr(bucket, "raw_upper_bound", None) is_integer_bucket = False if data_type: normalized_type = str(data_type).upper().replace(" UNSIGNED", "") if normalized_type in { "TINYINT", "SMALLINT", "MEDIUMINT", "INT", "INTEGER", "BIGINT", }: is_integer_bucket = True blocks.append( BucketBlock( block_id=block_id_counter, source_bucket_id=bucket.bucket_id, lower_bound=bucket.lower_bound, upper_bound=bucket.upper_bound, count=bucket.count, repeats=repeats, column_name=getattr(bucket, "column_name", None), data_type=data_type, raw_lower_bound=raw_lower, raw_upper_bound=raw_upper, is_integer_bucket=is_integer_bucket, ) ) block_id_counter += 1 return blocks def sample(self) -> Optional[ShuffleResult]: """ Layered shuffle sampling strategy: outer shuffle block_id, inner shuffle index Returns: ShuffleResult: sampling result, returns None if all values are consumed """ if self.remaining_capacity <= 0: return None # Find the next available block chosen_block = self._get_next_available_block() if chosen_block is None: return None # Sample from the selected block result = self._sample_from_block(chosen_block) if result is not None: self.remaining_capacity -= 1 return result def _get_next_available_block(self) -> Optional["ShuffleBlock"]: """Find the next available block in the shuffled block_id order, pop the block if it's used up""" # Check if the current block has remaining capacity if hasattr(self, "_current_active_block") and self._current_active_block: if self._current_active_block.get_remaining_capacity() > 0: return self._current_active_block # If the current block is used up or there is no current block, find the next one while self.block_id_sequence: block_id = self.block_id_sequence[-1] # Always take the last one # Find the corresponding block block = next((b for b in self.blocks if b.block_id == block_id), None) if block and block.get_remaining_capacity() > 0: self._current_active_block = block return block else: # This block is used up, remove it from the sequence self.block_id_sequence.pop() # O(1) # All blocks are used up return None def _sample_from_block(self, block: ShuffleBlock) -> Optional[ShuffleResult]: """Sample from the specified block""" if isinstance(block, TopNBlock): # TopN block: take one by one until used up value = block.sample(self._random) if value is not None: return ShuffleResult( value=value, source_type="topn", debug_info=( "block_id", block.block_id, "remaining", block.get_remaining_capacity(), ), ) elif isinstance(block, BucketBlock): result = block.sample(self._random) if result is not None: value, index = result return ShuffleResult( value=value, source_type="bucket", debug_info=( "block_id", block.block_id, "bucket_id", block.source_bucket_id, "index", index, ), ) return None def sample_batch(self, n: int) -> List[ShuffleResult]: """Batch sampling""" results = [] for _ in range(n): result = self.sample() if result is None: break results.append(result) return results def get_remaining_capacity(self) -> int: """Get remaining sampling capacity""" return self.remaining_capacity def reset(self): """Reset the memory-efficient sampler - ensure complete reproducibility""" # print(f"DEBUG: BlockShuffler RESET called! capacity was {self.remaining_capacity}/{self.total_capacity}") # Reset the random number generator to the initial state self._random = CheckpointablePRNG(self._original_seed) # Reset all blocks for block in self.blocks: block.reset() # Reset capacity self.remaining_capacity = self.total_capacity # Reinitialize the layered shuffle state self.block_id_sequence = [block.block_id for block in self.blocks] self._random.shuffle(self.block_id_sequence) # Clear the current active block state if hasattr(self, "_current_active_block"): self._current_active_block = None @dataclass class _LazyColumnState: iterator: "_ColumnBucketIterator" pending_reset: bool = True current_result: Optional[ShuffleResult] = None class _ColumnBucketIterator: """Lazy iterator that shuffles buckets and consumes values bucket-by-bucket.""" def __init__(self, shuffler: BlockShuffler): self._blocks_template = [] self.total_capacity = 0 for block in shuffler.blocks: if isinstance(block, TopNBlock): has_original = hasattr(block, "_original_items") original_items = getattr(block, "_original_items", None) if has_original and original_items is not None: items_source = copy.deepcopy(original_items) else: items_source = copy.deepcopy(block.topn_items) block_capacity = sum(item.get("count", 0) for item in items_source) if block_capacity <= 0: continue self._blocks_template.append( { "type": "topn", "block_id": block.block_id, "items": items_source, } ) self.total_capacity += block_capacity elif isinstance(block, BucketBlock): repeats = max(1, block.repeats) block_capacity = block.count // repeats if block_capacity <= 0: continue self._blocks_template.append( { "type": "bucket", "block_id": block.block_id, "source_bucket_id": block.source_bucket_id, "lower_bound": block.lower_bound, "upper_bound": block.upper_bound, "count": block.count, "repeats": repeats, } ) self.total_capacity += block_capacity self._rng: Optional[CheckpointablePRNG] = None self._blocks_state: List[Dict[str, Any]] = [] self._block_stack: List[int] = [] self._current_block_idx: Optional[int] = None def reset(self, seed: Optional[int]) -> None: seed_value = seed if seed is not None else 42 self._rng = CheckpointablePRNG(seed_value) self._blocks_state = [] for tpl in self._blocks_template: if tpl["type"] == "topn": items_state = [ {"value": item["value"], "count": item["count"]} for item in tpl["items"] ] self._blocks_state.append( { "type": "topn", "block_id": tpl["block_id"], "items": items_state, "order": None, } ) else: self._blocks_state.append( { "type": "bucket", "block_id": tpl["block_id"], "source_bucket_id": tpl["source_bucket_id"], "lower_bound": tpl["lower_bound"], "upper_bound": tpl["upper_bound"], "count": tpl["count"], "repeats": tpl["repeats"], "indices": None, } ) self._block_stack = list(range(len(self._blocks_state))) if self._block_stack: self._rng.shuffle(self._block_stack) self._current_block_idx = None def next(self) -> Optional[ShuffleResult]: if not self._blocks_state: return None while True: if self._current_block_idx is None: if not self._block_stack: return None self._current_block_idx = self._block_stack.pop() state = self._blocks_state[self._current_block_idx] if state["type"] == "bucket": if state["indices"] is None: indices = self._build_bucket_indices(state) self._rng.shuffle(indices) state["indices"] = indices if not state["indices"]: state["indices"] = None self._current_block_idx = None continue index = state["indices"].pop() value = self._compute_bucket_value( state["lower_bound"], state["upper_bound"], state["count"], state["repeats"], index, ) if not state["indices"]: state["indices"] = None self._current_block_idx = None debug_info = ( "block_id", state["block_id"], "bucket_id", state["source_bucket_id"], "index", index, ) return ShuffleResult( value=value, source_type="bucket", debug_info=debug_info, ) else: # TopN if state["order"] is None: order = [] for item_idx, item in enumerate(state["items"]): count = int(item.get("count", 0)) if count <= 0: continue order.extend([item_idx] * count) self._rng.shuffle(order) state["order"] = order if not state["order"]: state["order"] = None self._current_block_idx = None continue item_idx = state["order"].pop() item = state["items"][item_idx] value = item["value"] item["count"] = max(0, int(item.get("count", 0)) - 1) if not state["order"]: state["order"] = None self._current_block_idx = None debug_info = ("block_id", state["block_id"], "topn_index", item_idx) return ShuffleResult( value=value, source_type="topn", debug_info=debug_info ) @staticmethod def _build_bucket_indices(state: Dict[str, Any]) -> List[int]: repeats = max(1, int(state["repeats"])) count = int(state["count"]) if repeats > 1: return list(range(count // repeats)) return list(range(count)) @staticmethod def _compute_bucket_value( lower_bound: float, upper_bound: float, count: int, repeats: int, index: int, ) -> float: repeats = max(1, repeats) if repeats > 1: range_size = upper_bound - lower_bound offset = range_size * repeats / max(1, count) return lower_bound + offset * index range_size = upper_bound - lower_bound + 1 return lower_bound + range_size / max(1, count) * index class CompositeBlockShuffler: """Lazy composite BlockShuffler that consumes bucket values on demand.""" def __init__(self, shufflers: List[BlockShuffler], seed: Optional[int] = 42): self._original_seed = seed if seed is not None else 42 self._random = CheckpointablePRNG(self._original_seed) self.shufflers = shufflers self._column_states: List[_LazyColumnState] = [] self.total_capacity = 1 for shuffler in shufflers: iterator = _ColumnBucketIterator(shuffler) if iterator.total_capacity <= 0: self.total_capacity = 0 break self.total_capacity *= iterator.total_capacity self._column_states.append(_LazyColumnState(iterator=iterator)) if not self._column_states or self.total_capacity == 0: self.total_capacity = 0 self.remaining_capacity = 0 self._exhausted = True self._current_results: List[Optional[ShuffleResult]] = [] return self.remaining_capacity = self.total_capacity self._exhausted = False self._current_results: List[Optional[ShuffleResult]] = [None] * len( self._column_states ) self._initialize_first_combination() def _initialize_first_combination(self) -> None: for idx in range(len(self._column_states)): if not self._advance_column(idx): self._exhausted = True self.remaining_capacity = 0 return def _generate_seed(self) -> int: return int(self._random.uniform(0, 2**63 - 1)) def _advance_column(self, col_idx: int) -> bool: state = self._column_states[col_idx] if state.pending_reset: state.iterator.reset(self._generate_seed()) state.pending_reset = False result = state.iterator.next() if result is None: state.pending_reset = True state.current_result = None return False state.current_result = result self._current_results[col_idx] = result return True def _advance_to_next_combination(self) -> None: if self._exhausted: return for col_idx in range(len(self._column_states) - 1, -1, -1): if self._advance_column(col_idx): for reset_idx in range(col_idx + 1, len(self._column_states)): self._column_states[reset_idx].pending_reset = True if not self._advance_column(reset_idx): self._exhausted = True self.remaining_capacity = 0 return return else: if col_idx == 0: self._exhausted = True self.remaining_capacity = 0 return def sample(self) -> Optional[ShuffleResult]: if self._exhausted or self.remaining_capacity <= 0: return None while not self._exhausted and self.remaining_capacity > 0: raw_combination = tuple( result.value if result is not None else None for result in self._current_results ) shuffler_results = list(self._current_results) final_result = ShuffleResult( value=raw_combination, source_type="composite", debug_info=None, shuffler_results=shuffler_results, ) self.remaining_capacity -= 1 self._advance_to_next_combination() return final_result return None def sample_batch(self, n: int) -> List[ShuffleResult]: results = [] for _ in range(n): res = self.sample() if res is None: break results.append(res) return results def get_remaining_capacity(self) -> int: return self.remaining_capacity def has_capacity(self) -> bool: return self.remaining_capacity > 0 def reset(self) -> None: self._random = CheckpointablePRNG(self._original_seed) for state in self._column_states: state.pending_reset = True state.current_result = None self.remaining_capacity = self.total_capacity self._exhausted = False self._current_results = [None] * len(self._column_states) self._initialize_first_combination() class DeprecatedCompositeBlockShuffler: """ Multi-column composite BlockShuffler - Cartesian product generator Uses "layered shuffle" mode to generate complete Cartesian product: 1. First n columns are shuffled to get the initial combination 2. The last column is fully shuffled 3. The second-to-last column is shuffled once (change the value) 4. The last column is shuffled again 5. Continue this process until all combinations are generated Ensures complete coverage while maintaining column internal randomness """ def __init__(self, shufflers: List[BlockShuffler], seed: Optional[int] = 42): """ Initialize the multi-column composite BlockShuffler Args: shufflers: List of BlockShuffler, each responsible for sampling one column seed: Random seed """ self._random = CheckpointablePRNG(seed) self.shufflers = shufflers # Get all possible values for each column (shuffle order) self.column_values = [] for shuffler in shufflers: column_vals = [] shuffler.reset() while shuffler.get_remaining_capacity() > 0: result = shuffler.sample() if result is None: break column_vals.append(result.value) self.column_values.append(column_vals) shuffler.reset() # Calculate the true Cartesian product capacity self.column_capacities = [len(vals) for vals in self.column_values] self.total_capacity = 1 for cap in self.column_capacities: self.total_capacity *= cap # Initialize layered shuffle state self._init_layered_shuffle() self.remaining_capacity = self.total_capacity def _init_layered_shuffle(self): """Initialize layered shuffle state""" if not self.column_values: return # Current index position for each column (pointing to the shuffled value) self.current_indices = [0] * len(self.column_values) # Index sequence after shuffling for each column self.shuffled_indices = [] for i, column_vals in enumerate(self.column_values): indices = list(range(len(column_vals))) self._random.shuffle(indices) self.shuffled_indices.append(indices) def sample(self) -> Optional[ShuffleResult]: """ Sample a combination value - layered shuffle mode Algorithm: 1. Build the combination from the current indices 2. Move the index of the last column forward 3. If the last column is exhausted, move the index of the second-to-last column forward, and shuffle the last column 4. Continue this process until all combinations are generated """ if self.remaining_capacity <= 0: return None # Build the current combination combination = tuple( self.column_values[col_idx][ self.shuffled_indices[col_idx][self.current_indices[col_idx]] ] for col_idx in range(len(self.column_values)) ) # Move to the next combination (layered shuffle carry) self._advance_to_next_combination() self.remaining_capacity -= 1 return ShuffleResult( value=combination, source_type="composite", debug_info=None ) def _advance_to_next_combination(self): """Move to the next combination - layered shuffle carry logic""" if not self.current_indices: return # Start carrying from the last column for col_idx in range(len(self.current_indices) - 1, -1, -1): self.current_indices[col_idx] += 1 # If the current column still has values, the carry ends if self.current_indices[col_idx] < self.column_capacities[col_idx]: # If it's not the last column, shuffle all columns after it if col_idx < len(self.current_indices) - 1: for reset_col_idx in range(col_idx + 1, len(self.current_indices)): # Shuffle this column indices = list(range(self.column_capacities[reset_col_idx])) self._random.shuffle(indices) self.shuffled_indices[reset_col_idx] = indices self.current_indices[reset_col_idx] = 0 break else: # If the current column is exhausted, reset to 0 and continue carrying to the previous column self.current_indices[col_idx] = 0 def sample_batch(self, n: int) -> List[ShuffleResult]: """Batch sample combinations""" results = [] for _ in range(n): result = self.sample() if result is None: break results.append(result) return results def get_remaining_capacity(self) -> int: """Get remaining sampling capacity""" return self.remaining_capacity def reset(self): """Reset all sub-samplers""" for shuffler in self.shufflers: shuffler.reset() # Reinitialize the layered shuffle state self._init_layered_shuffle() self.remaining_capacity = self.total_capacity class DeprecatedCompositeBlockShufflerV2: """ Multi-column composite BlockShuffler - Memory-efficient true Cartesian product generator Core idea: 1. For each column, pre-generate all possible values using single-column sampling (respects TopN counts) 2. Use layered shuffle approach similar to DeprecatedCompositeBlockShuffler 3. Memory: O(total values across all columns) instead of O(Cartesian product size) Performance considerations: - Memory usage scales linearly with total column capacity (sum of all column values) - Initialization time scales with total column capacity (must sample all values once) - For very large columns (>millions of values), consider using streaming or batch approaches - reset() operation re-extracts all column values, which can be expensive for large datasets """ def __init__(self, shufflers: List[BlockShuffler], seed: Optional[int] = 42): """ Initialize the composite shuffler Args: shufflers: List of BlockShuffler, each responsible for sampling one column seed: Random seed """ # Save original seed for reset reproducibility self._original_seed = seed self._random = CheckpointablePRNG(seed) self.shufflers = shufflers # self._init_layered_shuffle() # Extract all possible values from each column and cache them self._extract_column_values_and_initialize() # Remove verbose creation logging # def _check_any_topn_blocks(self) -> bool: # """Check if any shufflers contain TopN blocks""" # for shuffler in self.shufflers: # for block in shuffler.blocks: # if isinstance(block, TopNBlock): # return True # return False def _extract_column_values_and_initialize(self): """Extract all values from each column and initialize the shuffler state""" # Use standard strategy (original logic) self._init_standard_strategy() def _init_standard_strategy(self): """Original strategy: extract all values respecting counts""" # Cache original column values and results for reset() optimization if not hasattr(self, "_original_column_values"): self._original_column_values = [] self._original_column_results = [] # Extract all possible values from each column (respecting TopN counts and randomness) for shuffler in self.shufflers: column_vals = [] column_results = [] seen_values = set() # Track seen values to avoid duplicates shuffler.reset() # Ensure clean state while shuffler.get_remaining_capacity() > 0: result = shuffler.sample() if result is None: break # Convert to business-like value for deduplication business_value = ( int(result.value) if isinstance(result.value, (int, float)) else result.value ) # Only add if not seen before if business_value not in seen_values: seen_values.add(business_value) column_vals.append(result.value) column_results.append(result) # If duplicate, skip this result but continue sampling self._original_column_values.append(column_vals) self._original_column_results.append(column_results) shuffler.reset() # Reset for potential reuse # Use cached values and results (for reset() efficiency) self.column_values = [list(vals) for vals in self._original_column_values] self.column_results = [ list(results) for results in self._original_column_results ] # Calculate the true Cartesian product capacity self.column_capacities = [len(vals) for vals in self.column_values] self.total_capacity = 1 for cap in self.column_capacities: self.total_capacity *= cap if self.total_capacity == 0: raise ValueError( "Cannot create composite shuffler: one or more columns have zero capacity" ) # Initialize layered shuffle state self._init_layered_shuffle() self.remaining_capacity = self.total_capacity def _init_layered_shuffle(self): """Initialize layered shuffle state""" if not self.column_values: return # Current index position for each column (pointing to the shuffled value) self.current_indices = [0] * len(self.column_values) # Index sequence after shuffling for each column self.shuffled_indices = [] for i, column_vals in enumerate(self.column_values): indices = list(range(len(column_vals))) self._random.shuffle(indices) self.shuffled_indices.append(indices) def sample(self) -> Optional[ShuffleResult]: """ Sample a combination value using appropriate strategy """ # Use standard layered shuffle strategy return self._sample_standard() def _sample_standard(self) -> Optional[ShuffleResult]: """ Standard sampling strategy - layered shuffle mode Algorithm (same as DeprecatedCompositeBlockShuffler): 1. Build the combination from the current indices 2. Move the index of the last column forward 3. If the last column is exhausted, move the index of the second-to-last column forward, and shuffle the last column 4. Continue this process until all combinations are generated """ if self.remaining_capacity <= 0: # Exhaustion logging reduced to prevent spam return None # Build the current combination and collect corresponding ShuffleResults combination = [] shuffler_results = [] # Track indices before advance for col_idx in range(len(self.column_values)): value_index = self.shuffled_indices[col_idx][self.current_indices[col_idx]] value = self.column_values[col_idx][value_index] result = self.column_results[col_idx][value_index] combination.append(value) shuffler_results.append(result) # Move to the next combination (layered shuffle carry) self._advance_to_next_combination() self.remaining_capacity -= 1 # Log potential duplicates by tracking index usage final_combination = tuple(combination) # Remove duplicate detection logging since main issue is fixed # Remove progress logging to reduce noise return ShuffleResult( value=final_combination, source_type="composite", debug_info=None, shuffler_results=shuffler_results, ) def _advance_to_next_combination(self): """Move to the next combination - layered shuffle carry logic""" if not self.current_indices: return # Start carrying from the last column for col_idx in range(len(self.current_indices) - 1, -1, -1): self.current_indices[col_idx] += 1 # If the current column still has values, the carry ends if self.current_indices[col_idx] < self.column_capacities[col_idx]: # If it's not the last column, shuffle all columns after it if col_idx < len(self.current_indices) - 1: for reset_col_idx in range(col_idx + 1, len(self.current_indices)): # Shuffle this column indices = list(range(self.column_capacities[reset_col_idx])) self._random.shuffle(indices) self.shuffled_indices[reset_col_idx] = indices self.current_indices[reset_col_idx] = 0 break else: # If the current column is exhausted, reset to 0 and continue carrying to the previous column self.current_indices[col_idx] = 0 def get_remaining_capacity(self) -> int: """Get remaining sampling capacity""" return self.remaining_capacity def reset(self): """Reset the composite shuffler to initial state with reproducible randomness""" # Reset all column shufflers to ensure they're in clean state for shuffler in self.shufflers: shuffler.reset() # Reset the random number generator to initial state for reproducibility self._random = CheckpointablePRNG(self._original_seed) # Reset standard strategy self._reset_standard_strategy() def _reset_topn_unique_strategy(self): """Reset TopN unique strategy state""" # Re-initialize TopN strategy from scratch self._init_topn_unique_strategy() def _reset_standard_strategy(self): """Reset standard strategy state""" # Use cached column values and results (avoid expensive re-extraction) if hasattr(self, "_original_column_values"): # Fast path: reuse cached values and results self.column_values = [list(vals) for vals in self._original_column_values] self.column_results = [ list(results) for results in self._original_column_results ] self.column_capacities = [len(vals) for vals in self.column_values] self.total_capacity = 1 for cap in self.column_capacities: self.total_capacity *= cap else: # Fallback: re-extract (shouldn't happen in normal usage) self._extract_column_values_and_initialize() return # Reinitialize the layered shuffle state with new random seed self._init_layered_shuffle() self.remaining_capacity = self.total_capacity if __name__ == "__main__": print("=== BlockShuffler Test Cases ===") # Import necessary types from .inverse_transform_sampler import HistogramBucket, TopNItem # Create simple test buckets # Column 1: 0-3 (4 values) col1_buckets = [HistogramBucket(lower_bound=0, upper_bound=3, count=4, bucket_id=1)] # Column 2: 0-2 (3 values) col2_buckets = [HistogramBucket(lower_bound=0, upper_bound=2, count=3, bucket_id=2)] # Column 3: 0-1 (2 values) col3_buckets = [HistogramBucket(lower_bound=0, upper_bound=1, count=2, bucket_id=3)] # Create single column samplers seed = 42 col1_shuffler = BlockShuffler(buckets=col1_buckets, seed=seed) col2_shuffler = BlockShuffler(buckets=col2_buckets, seed=seed + 1) col3_shuffler = BlockShuffler(buckets=col3_buckets, seed=seed + 2) print("Single column sampler test:") print(f"Column 1 capacity: {col1_shuffler.get_remaining_capacity()}") print(f"Column 2 capacity: {col2_shuffler.get_remaining_capacity()}") print(f"Column 3 capacity: {col3_shuffler.get_remaining_capacity()}") # Test composite sampler print("\n=== Composite primary key sampler test ===") composite_shuffler = CompositeBlockShuffler( [col1_shuffler, col2_shuffler, col3_shuffler], seed=seed ) print("Generating all combinations (seed=42):") results_seed42 = [] for i in range(24): # 4*3*2=24 combinations result = composite_shuffler.sample() if result is None: break results_seed42.append(result.value) print(f"Combination {i + 1}: {result.value}") print(f"Total generated: {len(results_seed42)} combinations") # Verify reproducibility print("\n=== Reproducibility test ===") # Recreate samplers with same seed col1_shuffler_repeat = BlockShuffler(buckets=col1_buckets, seed=seed) col2_shuffler_repeat = BlockShuffler(buckets=col2_buckets, seed=seed + 1) col3_shuffler_repeat = BlockShuffler(buckets=col3_buckets, seed=seed + 2) composite_shuffler_repeat = CompositeBlockShuffler( [col1_shuffler_repeat, col2_shuffler_repeat, col3_shuffler_repeat], seed=seed ) print("Repeat generation (same seed=42):") results_repeat = [] for i in range(24): result = composite_shuffler_repeat.sample() if result is None: break results_repeat.append(result.value) # Verify result consistency print(f"Original result count: {len(results_seed42)}") print(f"Repeat result count: {len(results_repeat)}") print(f"Results consistent: {results_seed42 == results_repeat}") if results_seed42 != results_repeat: print("Inconsistent positions:") for i, (r1, r2) in enumerate(zip(results_seed42, results_repeat)): if r1 != r2: print(f" Position {i}: {r1} != {r2}") # Verify no duplicates print(f"Duplicates present: {len(set(results_seed42)) != len(results_seed42)}") print(f"Unique combinations: {len(set(results_seed42))}") print("\n=== Different seed test ===") col1_shuffler_diff = BlockShuffler(buckets=col1_buckets, seed=999) col2_shuffler_diff = BlockShuffler(buckets=col2_buckets, seed=1000) col3_shuffler_diff = BlockShuffler(buckets=col3_buckets, seed=1001) composite_shuffler_diff = CompositeBlockShuffler( [col1_shuffler_diff, col2_shuffler_diff, col3_shuffler_diff], seed=999 ) print("Different seed generated first 10 combinations:") for i in range(10): result = composite_shuffler_diff.sample() if result is None: break print(f"Combination {i + 1}: {result.value}") print("\n=== TopN + Bucket mixed test ===") # Create mixed sampler with TopN and Bucket topn_items = [ TopNItem(value=100, count=2), # High-frequency value 100 appears 2 times TopNItem(value=200, count=1), # High-frequency value 200 appears 1 time ] mixed_buckets = [ HistogramBucket( lower_bound=0, upper_bound=10, count=3, bucket_id=10 ) # Normal interval ] mixed_shuffler = BlockShuffler(buckets=mixed_buckets, top_n=topn_items, seed=seed) print("Mixed sampler test (TopN + Bucket):") print( f"Mixed column capacity: {mixed_shuffler.get_remaining_capacity()}" ) # Should be 2+1+3=6 print("Mixed column all values:") mixed_values = [] for i in range(10): # Try up to 10 times result = mixed_shuffler.sample() if result is None: break mixed_values.append((result.value, result.source_type)) print(f" {result.value} (source: {result.source_type})") print(f"Mixed column actual generated: {len(mixed_values)} values") print("\n=== Large scale performance test ===") import sys import time def get_memory_usage(): """Simple memory usage estimation""" return sys.getsizeof([]) / 1024 / 1024 # Baseline MB # Test performance at different scales test_cases = [ {"name": "Small scale", "counts": [100, 50, 20], "expected": 100 * 50 * 20}, { "name": "Medium scale", "counts": [1000, 500, 100], "expected": 1000 * 500 * 100, }, { "name": "Large scale", "counts": [5000, 1000, 200], "expected": 5000 * 1000 * 200, }, ] for case in test_cases: print(f"\n--- {case['name']} test ---") counts = case["counts"] # Create large capacity buckets large_buckets = [] for i, count in enumerate(counts): buckets = [ HistogramBucket( lower_bound=0, upper_bound=count - 1, count=count, bucket_id=i + 100 ) ] large_buckets.append(buckets) # Test single column sampler creation time start_time = time.time() large_shufflers = [] for i, buckets in enumerate(large_buckets): shuffler = BlockShuffler(buckets=buckets, seed=seed + i + 100) large_shufflers.append(shuffler) single_creation_time = time.time() - start_time # Test composite sampler creation time start_time = time.time() composite = CompositeBlockShuffler(large_shufflers, seed=seed + 1000) composite_creation_time = time.time() - start_time print( f"Column capacities: {[s.get_remaining_capacity() for s in large_shufflers]}" ) print(f"Theoretical combinations: {case['expected']:,}") print(f"Actual total capacity: {composite.get_remaining_capacity():,}") print(f"Single column creation time: {single_creation_time:.4f} seconds") print(f"Composite creation time: {composite_creation_time:.4f} seconds") # Test sampling performance (sample 1000 samples) sample_count = min(1000, case["expected"]) start_time = time.time() sampled_results = [] for i in range(sample_count): result = composite.sample() if result is None: break sampled_results.append(result.value) sampling_time = time.time() - start_time print( f"Sampled {len(sampled_results)} combinations, time: {sampling_time:.4f} seconds" ) print( f"Average per combination: {sampling_time / len(sampled_results) * 1000:.2f}ms" ) # Verify no duplicates unique_count = len(set(sampled_results)) print( f"No duplicate verification: {unique_count}/{len(sampled_results)} = {unique_count == len(sampled_results)}" ) # Show first few samples print(f"First 5 samples: {sampled_results[:5]}") print("\n=== Memory optimization comparison ===") # Compare: storing all combinations directly vs current approach large_case = test_cases[-1] # Large scale case total_combinations = large_case["expected"] # Assume each combination takes 3 * 8 bytes (3 columns * 8 byte floats) naive_memory = total_combinations * 3 * 8 / 1024 / 1024 / 1024 # GB # Current approach: column value sequences current_memory = sum(large_case["counts"]) * 8 / 1024 / 1024 # MB print(f"Storing all combinations directly: {naive_memory:.2f}GB") print(f"Current layered approach: {current_memory:.2f}MB") print(f"Memory savings: {naive_memory * 1024 / current_memory:.1f}x") print("\nTest completed!")