import string from decimal import Decimal, getcontext from typing import List, Optional from utils.scalar_to_string import scalar_to_string from utils.tidb_topn_decoder import decode_tidb_topn_value_safe, is_tidb_encoded_value from .block_shuffler import BlockShuffler from .checkpointable_prng import CheckpointablePRNG from .inverse_transform_sampler import InverseTransformSampler from .sampling_phase import SamplingPhase from .tidb_types import TiDBStringBucket, TiDBStringTopNItem # Set high precision for Decimal operations getcontext().prec = 50 class TextSampler: """ Text sampler that combines InverseTransformSampler with scalar_to_string for realistic text generation from TiDB string column statistics This is the production version that expects pre-computed TiDB statistics. Now includes TiDB-style common prefix handling for improved string sampling accuracy. Uses TiDBStringBucket internally for business logic, but converts to pure HistogramBucket for the mathematical InverseTransformSampler. """ def __init__( self, buckets: List[ TiDBStringBucket ], # TiDB string bucket objects from schema_manager top_n: Optional[ List[TiDBStringTopNItem] ] = None, # TiDB string TopN objects from schema_manager null_count: int = 0, charset: str = "utf8mb4", topn_string_mappings: dict = None, is_cumulative_count: bool = False, # False for absolute counts, TiDB is False too in stats table character_maximum_length: Optional[ int ] = None, # Maximum length for generated strings constraint_type: str = "normal", # "normal" or "unique" seed: Optional[int] = 42, column_name: str = None, # Column name for logging ): """ Initialize TextSampler from standardized TiDB string objects Args: buckets: List of TiDB string bucket objects (standardized by schema_manager) top_n: List of TiDB string TopN objects (standardized by schema_manager) null_count: Number of null values charset: Character set for scalar_to_string topn_string_mappings: Dict mapping TopN scalar values to original strings character_maximum_length: Maximum length for generated strings constraint_type: "normal" for InverseTransformSampler, "unique" for BlockShuffler seed: Random seed """ self.constraint_type = constraint_type self.null_count = null_count self.charset = charset self.topn_string_mappings = topn_string_mappings or {} self.character_maximum_length = character_maximum_length self.column_name = column_name or "unknown_column" self._seed = seed self._null_rng = CheckpointablePRNG(seed=seed) if null_count > 0 else None # Store original data for business logic self.tidb_buckets = buckets or [] self.tidb_topn = top_n or [] # Build complete scalar->string mapping and bucket boundaries self._build_text_mappings() # Create mathematical objects for sampler math_buckets, math_topn = self._create_math_objects() # Choose underlying sampler based on constraint type if constraint_type == "unique": # Primary key/unique key uses BlockShuffler (no duplicate sampling) # Handle case where column has no statistical data if not math_buckets and not math_topn: # print(f"INFO: Unique text column '{self.column_name}' has no statistical data, using pure new value generation") self.sampler = None # No BlockShuffler available else: self.sampler = BlockShuffler( buckets=math_buckets, top_n=math_topn, seed=seed, ) # Initialize new value generation strategy for unique columns self._init_new_value_generation(seed) else: # Normal column uses InverseTransformSampler (probability sampling) self.sampler = InverseTransformSampler( buckets=math_buckets, top_n=math_topn, null_count=null_count, is_cumulative_count=is_cumulative_count, seed=seed, ) @property def its(self): """Backward compatibility: Access underlying InverseTransformSampler""" if self.constraint_type == "unique": # BlockShuffler mode, return None or raise error for backward compatibility raise AttributeError( "TextSampler in 'unique' constraint mode uses BlockShuffler, " "not InverseTransformSampler. Use sampler.sampler instead." ) return self.sampler def _init_new_value_generation(self, seed): """Initialize new value generation strategy for unique columns when BlockShuffler is exhausted""" # Analyze existing string length patterns self._analyze_length_patterns() # Determine available characters based on charset if self.charset == "utf8mb4": self._available_chars = string.ascii_letters + string.digits + "_-" else: # Other charset handling self._available_chars = string.ascii_letters + string.digits # Initialize deterministic random generator using CheckpointablePRNG self._text_prng = CheckpointablePRNG(seed=seed) self._next_counter = 1 # Flag to log exhaustion message only once self._block_exhausted_logged = False def _analyze_length_patterns(self): """Analyze existing string length patterns""" length_distribution = [] # Analyze TopN string lengths for topn in self.tidb_topn: original_str = self._extract_original_string(topn.encoded) if original_str: length_distribution.append(len(original_str)) # Analyze bucket string lengths for bucket in self.tidb_buckets: if hasattr(bucket, "lower_bound") and hasattr(bucket, "upper_bound"): lower_str = self._extract_original_string(bucket.lower_bound) upper_str = self._extract_original_string(bucket.upper_bound) if lower_str: length_distribution.append(len(lower_str)) if upper_str: length_distribution.append(len(upper_str)) # Calculate average length self._avg_length = ( int(sum(length_distribution) / len(length_distribution)) if length_distribution else 8 ) def _generate_new_unique_value(self) -> str: """Generate deterministic new text value with fixed random number consumption""" # Create dedicated PRNG instance based on current counter for deterministic generation # Each generation consumes fixed amount of random numbers, ensuring determinism generation_seed = ( self._text_prng.seed if self._text_prng.seed is not None else 42 ) + self._next_counter local_prng = CheckpointablePRNG(seed=generation_seed) # Determine length - consumes 1 random number length_variation = int(local_prng.uniform(-2, 3)) # [-2, 2] integer range target_length = max(1, self._avg_length + length_variation) # Consider counter suffix length counter_suffix = str(self._next_counter) main_content_length = max(0, target_length - len(counter_suffix)) # Apply character length limit if self.character_maximum_length: if len(counter_suffix) >= self.character_maximum_length: # If counter already exceeds limit, return truncated counter result = counter_suffix[: self.character_maximum_length] self._next_counter += 1 return result else: main_content_length = min( main_content_length, self.character_maximum_length - len(counter_suffix), ) # Generate main content - each character consumes 1 random number chars = [] char_count = len(self._available_chars) for i in range(main_content_length): # Use uniform to generate [0, 1) then map to character index rand_val = local_prng.uniform(0, 1) char_index = int(rand_val * char_count) chars.append(self._available_chars[char_index]) # Combine result: main content + counter suffix result = "".join(chars) + counter_suffix # Increment counter self._next_counter += 1 # Final length check if ( self.character_maximum_length and len(result) > self.character_maximum_length ): result = result[: self.character_maximum_length] return result def has_statistical_capacity(self) -> bool: """Check if statistical information still has capacity""" if self.constraint_type == "unique": if self.sampler is None: return False # No statistical data, no capacity return hasattr(self.sampler, "has_capacity") and self.sampler.has_capacity() return True # Normal columns always have capacity def sample_with_phase(self, phase: SamplingPhase) -> Optional[str]: """Sample with phase awareness""" if self.constraint_type == "normal": # Normal columns use InverseTransformSampler in both phases return self._sample_normal() elif self.constraint_type == "unique": if phase == SamplingPhase.PREPARE: if self.sampler is None: # No statistical data, use pure new value generation return self._generate_unique_fallback_value() sample_result = self.sampler.sample() if sample_result is None: if not self._block_exhausted_logged: remaining_capacity = getattr( self.sampler, "remaining_capacity", "unknown" ) print( f"INFO: BlockShuffler exhausted for unique text column '{self.column_name}', switching to new value generation. Remaining capacity: {remaining_capacity}" ) self._block_exhausted_logged = True return self._generate_unique_fallback_value() if self.null_count > 0: if self._null_rng is None: self._null_rng = CheckpointablePRNG(seed=self._seed) total_capacity = self.sampler.total_capacity + self.null_count null_probability = self.null_count / total_capacity if self._null_rng.random() < null_probability: return None return self._sample_from_blockshuffler_result(sample_result) elif phase == SamplingPhase.REPLAY: return self._generate_unique_fallback_value() else: raise ValueError(f"Unknown constraint_type: {self.constraint_type}") def _sample_normal(self) -> Optional[str]: """Sample from normal columns using existing logic""" # Get sample result from underlying sampler sample_result = self.sampler.sample() if sample_result is None or sample_result.value is None: return None scalar_value = sample_result.value if sample_result.source_type == "topn": candidate = self._find_nearest_topn_string(scalar_value) elif sample_result.source_type == "bucket": candidate = self._generate_bucket_string(scalar_value, sample_result) else: raise ValueError( f"Unknown sample source type: {sample_result.source_type}. " f"Expected 'topn' or 'bucket', got: {sample_result.source_type}" ) return self._enforce_unique_block_value(candidate, sample_result) def _sample_from_blockshuffler_result(self, sample_result) -> Optional[str]: """Reuse BlockShuffler result while enforcing uniqueness.""" if sample_result is None or sample_result.value is None: return None scalar_value = sample_result.value if sample_result.source_type == "topn": candidate = self._find_nearest_topn_string(scalar_value) elif sample_result.source_type == "bucket": candidate = self._generate_bucket_string(scalar_value, sample_result) else: raise ValueError( f"Unknown sample source type: {sample_result.source_type}. " f"Expected 'topn' or 'bucket', got: {sample_result.source_type}" ) return self._enforce_unique_block_value(candidate, sample_result) def _build_text_mappings(self): """Build complete scalar->string mappings and bucket boundaries for business logic""" # Initialize mappings - use Decimal keys for high precision self.topn_string_mappings = {} # Decimal -> string mappings self.bucket_boundaries = {} # bucket_id -> boundary info for O(1) lookup # Process original TopN items for topn_item in self.tidb_topn: # Extract original string from TiDBStringTopNItem if hasattr(topn_item, "encoded"): original_string = self._extract_original_string(topn_item.encoded) if original_string: # Get scalar value for this TopN and convert to Decimal scalar_value = topn_item.parse_encoded_value(topn_item.encoded) decimal_key = Decimal( str(scalar_value) ) # Convert float to high-precision Decimal self.topn_string_mappings[decimal_key] = original_string # Process buckets - extract equal-bound as TopN, keep normal buckets for bucket in self.tidb_buckets: if hasattr(bucket, "parse_bound_value"): # Compare original bounds directly (more reliable than scalar comparison) bounds_equal = bucket.lower_bound == bucket.upper_bound if bounds_equal: # Extract as TopN - get original string and scalar original_string = self._extract_original_string(bucket.lower_bound) if original_string: # For equal bounds, we should use the same logic as normal buckets # to ensure consistency with common prefix handling from utils.bytes_scalar_converter import convert_bytes_to_scalar # TiDBStringBucket bounds should be bytes if isinstance(bucket.lower_bound, bytes): scalar_value = convert_bytes_to_scalar(bucket.lower_bound) else: # Fallback: convert string to bytes first scalar_value = convert_bytes_to_scalar( bucket.lower_bound.encode("utf-8") ) decimal_key = Decimal( str(scalar_value) ) # Convert to high-precision Decimal self.topn_string_mappings[decimal_key] = original_string else: # Normal bucket - calculate common prefix first, then scalar values (TiDB style) # TiDBStringBucket bounds should already be bytes if isinstance(bucket.lower_bound, bytes): lower_bytes = bucket.lower_bound else: # Fallback: convert string to bytes first (should not happen in normal use) lower_bytes = bucket.lower_bound.encode("utf-8") if isinstance(bucket.upper_bound, bytes): upper_bytes = bucket.upper_bound else: # Fallback: convert string to bytes first (should not happen in normal use) upper_bytes = bucket.upper_bound.encode("utf-8") # Calculate common prefix for this bucket using optimized function from utils.bytes_scalar_converter import common_prefix_length common_prefix_len = common_prefix_length(lower_bytes, upper_bytes) common_prefix_bytes = ( lower_bytes[:common_prefix_len] if common_prefix_len > 0 else b"" ) # Calculate scalar values AFTER removing common prefix (TiDB style) from utils.bytes_scalar_converter import ( convert_bytes_to_scalar_with_prefix, ) lower_scalar = convert_bytes_to_scalar_with_prefix( lower_bytes, common_prefix_len ) upper_scalar = convert_bytes_to_scalar_with_prefix( upper_bytes, common_prefix_len ) bucket_id = getattr(bucket, "bucket_id", None) if bucket_id is None: raise ValueError( f"bucket_id is required for optimization but not found in bucket: " f"{bucket.lower_bound} - {bucket.upper_bound}" ) boundary_data = { "lower_scalar": lower_scalar, "upper_scalar": upper_scalar, "lower_bound": bucket.lower_bound, "upper_bound": bucket.upper_bound, "count": bucket.count, "repeats": bucket.repeats, "common_prefix_bytes": common_prefix_bytes, "bucket_id": bucket_id, "column_name": getattr(bucket, "column_name", None), } self.bucket_boundaries[bucket_id] = boundary_data def _extract_original_string(self, value): """Extract original string from various formats""" if isinstance(value, str): return value elif isinstance(value, bytes): # First try TiDB decoding for encoded values if is_tidb_encoded_value(value): decoded = decode_tidb_topn_value_safe(value) if isinstance(decoded, str): return decoded elif isinstance(decoded, bytes): return decoded.decode("utf-8", errors="ignore") try: # Then try UTF-8 decoding return value.decode("utf-8") except UnicodeDecodeError: # Fallback for problematic bytes return value.decode("utf-8", errors="ignore") else: return str(value) def _create_math_objects(self): """Create mathematical objects for ITS""" from .inverse_transform_sampler import HistogramBucket, TopNItem # Create buckets from bucket_boundaries (already contains bucket_id and column_name) math_buckets = [] for boundary in self.bucket_boundaries.values(): math_buckets.append( HistogramBucket( lower_bound=boundary["lower_scalar"], upper_bound=boundary["upper_scalar"], count=boundary["count"], repeats=boundary["repeats"], bucket_id=boundary["bucket_id"], column_name=boundary["column_name"], ) ) # Create TopN from all scalar mappings (original + extracted) math_topn = [] # First add original TopN items for topn_item in self.tidb_topn: math_topn.append(topn_item.to_topn_item()) # Then add extracted TopN items from equal-bound buckets for bucket in self.tidb_buckets: if hasattr(bucket, "parse_bound_value"): # Compare original bounds directly (more reliable than scalar comparison) bounds_equal = bucket.lower_bound == bucket.upper_bound if bounds_equal: scalar_value = bucket.parse_bound_value(bucket.lower_bound) math_topn.append(TopNItem(value=scalar_value, count=bucket.count)) return math_buckets, math_topn def sample(self) -> Optional[str]: """ Sample a text value (backward compatibility, defaults to PREPARE phase) Returns: str: Generated text, or None for null values """ return self.sample_with_phase(SamplingPhase.PREPARE) def _find_nearest_topn_string(self, target_scalar: float) -> str: """Find the TopN string with scalar value closest to target_scalar""" if not self.topn_string_mappings: raise ValueError( "No TopN string mappings available for scalar lookup. " "This indicates a data initialization problem." ) # Convert target to Decimal for high-precision comparison target_decimal = Decimal(str(target_scalar)) # Find the Decimal key with minimum absolute difference closest_decimal = min( self.topn_string_mappings.keys(), key=lambda x: abs(x - target_decimal) ) result = self.topn_string_mappings[closest_decimal] return self._truncate_string(result) def _generate_bucket_string(self, scalar_value: float, sample_result=None) -> str: """Generate string from bucket using scalar_to_string with boundary info""" # Use bucket_id from sample_result if available for optimization target_boundary = None if ( sample_result and hasattr(sample_result, "bucket_id") and sample_result.bucket_id is not None ): target_boundary = self.bucket_boundaries.get(sample_result.bucket_id) if target_boundary is None: # Fallback to original search method if not self.bucket_boundaries: raise ValueError( "No bucket boundaries available for string generation. " "This indicates a data initialization problem." ) # Find the bucket that should contain this scalar value for boundary in self.bucket_boundaries.values(): if boundary["lower_scalar"] <= scalar_value <= boundary["upper_scalar"]: target_boundary = boundary break # If not found in any bucket, use the bucket with closest range if target_boundary is None: target_boundary = min( self.bucket_boundaries.values(), key=lambda b: min( abs(scalar_value - b["lower_scalar"]), abs(scalar_value - b["upper_scalar"]), ), ) # Use scalar_to_string directly result, _ = scalar_to_string( scalar_value, target_boundary["lower_scalar"], target_boundary["upper_scalar"], charset=self.charset, ) if result is not None: # Add back the common prefix (TiDB-style) if target_boundary["common_prefix_bytes"]: try: prefix_str = target_boundary["common_prefix_bytes"].decode( "utf-8", errors="ignore" ) final_result = prefix_str + result return self._truncate_string(final_result) except UnicodeDecodeError: raise ValueError( f"Invalid common prefix bytes: {target_boundary['common_prefix_bytes']}" ) else: return self._truncate_string(result) # Ultimate fallback raise ValueError( f"Failed to generate string for scalar {scalar_value} " f"within bounds {target_boundary['lower_scalar']} - {target_boundary['upper_scalar']}. " "This indicates a data integrity issue." ) def sample_batch(self, n: int) -> List[Optional[str]]: """ Sample multiple text values Args: n: Number of samples to generate Returns: List of generated text values """ return [self.sample() for _ in range(n)] def convert_scalar_to_business_value(self, shuffler_result) -> str: """Convert BlockShuffler result to business value Args: shuffler_result: The detailed result object from BlockShuffler Returns: str: The converted business value (string) """ # Reuse existing _sample_from_blockshuffler_result logic return self._sample_from_blockshuffler_result(shuffler_result) def _enforce_unique_block_value( self, candidate: Optional[str], sample_result ) -> Optional[str]: """Ensure BlockShuffler outputs remain unique for unique text columns.""" if self.constraint_type != "unique" or candidate is None: return candidate if not self._is_valid_block_candidate(candidate): return self._generate_unique_fallback_value() return candidate def _generate_unique_fallback_value(self) -> str: """Generate a clean unique value when BlockShuffler output is unusable.""" max_attempts = 100 for _ in range(max_attempts): new_value = self._generate_new_unique_value() if not self._is_valid_block_candidate(new_value): continue return new_value raise RuntimeError( f"Failed to generate unique fallback value for column {self.column_name}" ) @staticmethod def _is_valid_block_candidate(candidate: str) -> bool: """Check whether the generated candidate contains only printable ASCII characters.""" if candidate is None: return False disallowed_codepoints = {0x07FF, 0x7F} for ch in candidate: code = ord(ch) # Reject C0/C1 control characters (except common whitespace) if code < 32 and ch not in {"\t", "\n", "\r"}: return False if 0x7F <= code <= 0x9F: return False # Reject Unicode replacement character and surrogate halves if code == 0xFFFD or 0xD800 <= code <= 0xDFFF: return False # Reject Unicode noncharacters (U+FDD0-U+FDEF and plane endings) if 0xFDD0 <= code <= 0xFDEF: return False if code & 0xFFFF in {0xFFFE, 0xFFFF}: return False if code in disallowed_codepoints: return False return True def _truncate_string(self, text: str) -> str: """Truncate string to character_maximum_length if specified Args: text: Input string to truncate Returns: str: Truncated string if character_maximum_length is set, otherwise original string """ if self.character_maximum_length is None or text is None: return text # Ensure character_maximum_length is an integer try: max_length = int(self.character_maximum_length) except (ValueError, TypeError) as e: # This should be validated at construction time, not masked here raise ValueError( f"Invalid character_maximum_length value: {self.character_maximum_length}. " f"Must be a valid integer. Error: {e}" ) from e if len(text) <= max_length: return text # Truncate to maximum length return text[:max_length]