import collections import random from dataclasses import InitVar, dataclass, field from typing import Any, Callable, Dict, Optional import pandas as pd from samplers.sampling_phase import SamplingPhase from samplers.state_manager import ( BlockShufflerState, BusinessSamplerState, ColumnState, StateManager, ) from .graph import Graph @dataclass class Column: name: str table_schema: str table_name: str ordinal_position: int data_type: str character_maximum_length: int stats_buckets_df: InitVar[pd.DataFrame] stats_top_n_df: InitVar[pd.DataFrame] graph: InitVar[Graph] sampler_factory: InitVar[Callable] # Delay initialization of sampler related fields _sampler: Any = field(init=False, default=None) _sampler_factory: Callable = field(init=False, default=None) _sampler_args: dict = field(init=False, default_factory=dict) def __post_init__(self, stats_buckets_df, stats_top_n_df, graph, sampler_factory): self.cache = collections.deque([], maxlen=4096) # Generate default seed based on column identity column_key = f"{self.table_schema}.{self.table_name}.{self.name}" default_seed = self._generate_column_seed(column_key, graph) # Save creation parameters, do not create sampler immediately self._sampler_factory = sampler_factory self._sampler_args = { "data_type": self.data_type, "stats_buckets_df": stats_buckets_df, "stats_top_n_df": stats_top_n_df, "character_maximum_length": self.character_maximum_length, "seed": default_seed, # Each column has its own deterministic seed } # Reference context for correlated columns (set by Table) self._reference_stats_buckets = None self._reference_stats_top_n = None self._reference_seed = None def _generate_column_seed(self, column_key: str, graph) -> int: """Generate deterministic seed from column identity using Graph's method""" column_hash = graph.get_column_identity(column_key) return int(column_hash[:8], 16) % (2**31) @property def sampler(self): """Access sampler, raise error if not initialized""" if self._sampler is None: raise RuntimeError( f"Sampler for column '{self.name}' not initialized. " f"Call initialize_sampler() first." ) return self._sampler def set_reference_context(self, ref_stats_buckets, ref_stats_top_n, ref_seed): """Set reference context from primary column in correlation cluster This method stores reference statistics and seed that will be applied when initialize_sampler() is called. Allows future extension of preprocessing strategies (e.g., merging own stats with reference stats). Args: ref_stats_buckets: Reference column's stats_buckets_df ref_stats_top_n: Reference column's stats_top_n_df ref_seed: Reference column's random seed """ self._reference_stats_buckets = ref_stats_buckets self._reference_stats_top_n = ref_stats_top_n self._reference_seed = ref_seed def apply_reference_context(self): """Apply reference context to sampler arguments This method merges reference statistics into _sampler_args. Currently implements direct override strategy, but can be extended to support more sophisticated preprocessing (e.g., stats merging, weighted combination, etc.). """ if self._reference_stats_buckets is not None: self._sampler_args["stats_buckets_df"] = self._reference_stats_buckets if self._reference_stats_top_n is not None: self._sampler_args["stats_top_n_df"] = self._reference_stats_top_n if self._reference_seed is not None: self._sampler_args["seed"] = self._reference_seed def initialize_sampler(self, constraint_type: str = "normal"): """Initialize sampler after Table creation, set constraint type and create sampler""" if self._sampler is not None: # Already initialized, return directly return # Add constraint type parameter self._sampler_args["constraint_type"] = constraint_type # Add column name for logging self._sampler_args["column_name"] = self.name # Apply reference context if set (must be called before sampler creation) self.apply_reference_context() # Create sampler self._sampler = self._sampler_factory(**self._sampler_args) def sample_for_insert(self, phase: SamplingPhase = SamplingPhase.PREPARE): """Sample for insert operation with phase awareness""" if hasattr(self.sampler, "sample_with_phase"): sample_value = self.sampler.sample_with_phase(phase) else: # Backward compatibility: if sampler doesn't support phase awareness, use original method sample_value = self.sampler.sample() self.cache.append(sample_value) return self.cache[-1] def sample_for_update(self): self.cache.append(self.sampler.sample()) return self.cache[-1] def sample_for_select(self): if len(self.cache) == 0: return self.sampler.sample() return random.choice(self.cache) # seed def sample_for_delete(self): if len(self.cache) == 0: return self.sampler.sample() return self.cache.pop() def save_state(self, state_manager: StateManager, table_name: str) -> bool: """Save column state using StateManager""" if self._sampler is None: print(f"⚠️ Column {self.name} sampler not initialized, skipping state save") return False try: # Extract sampler state sampler_state = self._extract_sampler_state() # Create column state column_state = ColumnState( sampler_state=sampler_state, cache_contents=list(self.cache), cache_maxlen=self.cache.maxlen, data_type=self.data_type, character_maximum_length=self.character_maximum_length, is_initialized=True, ) return state_manager.save_column_state(table_name, self.name, column_state) except Exception as e: print(f"❌ Failed to save state for column {self.name}: {e}") return False def load_state(self, state_manager: StateManager, table_name: str) -> bool: """Load column state using StateManager""" try: column_state = state_manager.load_column_state(table_name, self.name) if column_state is None: print(f"📄 No saved state found for column {self.name}") return False # Restore cache self.cache = collections.deque( column_state.cache_contents, maxlen=column_state.cache_maxlen ) # Restore sampler state if column_state.is_initialized: if self._sampler is None: # Initialize sampler first if not already done self.initialize_sampler(column_state.sampler_state.constraint_type) self._restore_sampler_state(column_state.sampler_state) return True except Exception as e: print(f"❌ Failed to load state for column {self.name}: {e}") return False def _extract_sampler_state(self) -> BusinessSamplerState: """Extract state from business sampler""" # Get BlockShuffler state if exists block_shuffler_state = None if hasattr(self._sampler, "sampler") and self._sampler.sampler is not None: block_shuffler_state = self._extract_block_shuffler_state() # Extract business layer state constraint_type = getattr(self._sampler, "constraint_type", "normal") column_name = getattr(self._sampler, "column_name", self.name) # Extract unique constraint state block_exhausted_logged = getattr( self._sampler, "_block_exhausted_logged", False ) generated_values = getattr(self._sampler, "_generated_values", set()) new_value_counter = getattr(self._sampler, "_next_new_value", 0) # Extract null handling state null_count = getattr(self._sampler, "null_count", 0) total_capacity = ( getattr(self._sampler.sampler, "total_capacity", 1) if hasattr(self._sampler, "sampler") else 1 ) null_probability = ( null_count / (total_capacity + null_count) if total_capacity > 0 else 0 ) return BusinessSamplerState( block_shuffler_state=block_shuffler_state, constraint_type=constraint_type, column_name=column_name, block_exhausted_logged=block_exhausted_logged, generated_values=generated_values, new_value_counter=new_value_counter, null_count=null_count, null_probability=null_probability, ) def _extract_block_shuffler_state(self) -> Optional[BlockShufflerState]: """Extract state from BlockShuffler""" if not hasattr(self._sampler, "sampler") or self._sampler.sampler is None: return None block_shuffler = self._sampler.sampler # Get PRNG position prng_position = getattr(block_shuffler._random, "position", 0) # Get capacity state remaining_capacity = getattr(block_shuffler, "remaining_capacity", 0) is_exhausted = remaining_capacity <= 0 # Get block states - this is complex, need to extract from each block block_states = {} for block in getattr(block_shuffler, "blocks", []): block_states[block.block_id] = self._extract_block_state(block) # Get sequence state block_id_sequence = getattr(block_shuffler, "block_id_sequence", []) current_block_index = getattr(block_shuffler, "_current_block_index", 0) return BlockShufflerState( prng_position=prng_position, remaining_capacity=remaining_capacity, is_exhausted=is_exhausted, block_states=block_states, block_id_sequence=block_id_sequence.copy() if block_id_sequence else [], current_block_index=current_block_index, ) def _extract_block_state(self, block) -> Dict[str, Any]: """Extract state from a specific block""" from samplers.block_shuffler import BucketBlock, TopNBlock if isinstance(block, TopNBlock): return { "type": "TopNBlock", "topn_items": block.topn_items.copy() if block.topn_items else [], "original_items": getattr(block, "_original_items", []), } elif isinstance(block, BucketBlock): return { "type": "BucketBlock", "remaining_count": getattr(block, "remaining_count", block.count), "bucket_indices": getattr(block, "bucket_indices", None), } else: return {"type": "unknown"} def _restore_sampler_state(self, sampler_state: BusinessSamplerState): """Restore business sampler state""" # Restore business layer state if hasattr(self._sampler, "constraint_type"): self._sampler.constraint_type = sampler_state.constraint_type if hasattr(self._sampler, "column_name"): self._sampler.column_name = sampler_state.column_name if hasattr(self._sampler, "_block_exhausted_logged"): self._sampler._block_exhausted_logged = sampler_state.block_exhausted_logged if hasattr(self._sampler, "_generated_values"): self._sampler._generated_values = sampler_state.generated_values.copy() if hasattr(self._sampler, "_next_new_value"): self._sampler._next_new_value = sampler_state.new_value_counter if hasattr(self._sampler, "null_count"): self._sampler.null_count = sampler_state.null_count # Restore BlockShuffler state if exists if sampler_state.block_shuffler_state and hasattr(self._sampler, "sampler"): self._restore_block_shuffler_state(sampler_state.block_shuffler_state) def _restore_block_shuffler_state(self, block_state: BlockShufflerState): """Restore BlockShuffler state""" if not hasattr(self._sampler, "sampler") or self._sampler.sampler is None: return block_shuffler = self._sampler.sampler # Restore PRNG position if hasattr(block_shuffler, "_random") and hasattr( block_shuffler._random, "restore" ): block_shuffler._random.restore(block_state.prng_position) # Restore capacity state if hasattr(block_shuffler, "remaining_capacity"): block_shuffler.remaining_capacity = block_state.remaining_capacity # Restore sequence state if hasattr(block_shuffler, "block_id_sequence"): block_shuffler.block_id_sequence = block_state.block_id_sequence.copy() if hasattr(block_shuffler, "_current_block_index"): block_shuffler._current_block_index = block_state.current_block_index # Restore individual block states for block in getattr(block_shuffler, "blocks", []): if block.block_id in block_state.block_states: self._restore_block_state( block, block_state.block_states[block.block_id] ) def _restore_block_state(self, block, block_state: Dict[str, Any]): """Restore state for a specific block""" from samplers.block_shuffler import BucketBlock, TopNBlock if isinstance(block, TopNBlock) and block_state["type"] == "TopNBlock": block.topn_items = block_state["topn_items"].copy() if "_original_items" in block_state: block._original_items = block_state["original_items"].copy() elif isinstance(block, BucketBlock) and block_state["type"] == "BucketBlock": block.remaining_count = block_state["remaining_count"] block.bucket_indices = block_state["bucket_indices"]