from datetime import datetime from typing import List, Optional from utils.tidb_time_decoder import scalar_to_tidb_time 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 TiDBTimeBucket, TiDBTimeTopNItem class DatetimeSampler: """ Datetime sampler using InverseTransformSampler Handles TiDB datetime column statistics with deltaT format This sampler expects TiDB-style input where datetime values are converted to deltaT (time differences in seconds from a baseline datetime). It supports string datetime parsing and returns formatted datetime strings. """ def __init__( self, buckets: List[TiDBTimeBucket], # TiDB time bucket objects from schema_manager top_n: Optional[ List[TiDBTimeTopNItem] ] = None, # TiDB time TopN objects from schema_manager null_count: int = 0, baseline_timestamp: Optional[int] = None, datetime_type: str = "DATETIME", # 'DATETIME', 'TIMESTAMP', or 'DATE' is_cumulative_count: bool = False, # False for absolute counts constraint_type: str = "normal", # "normal" or "unique" seed: Optional[int] = 42, column_name: str = None, # Column name for logging ): """ Initialize DatetimeSampler from standardized TiDB objects Args: buckets: List of TiDB time bucket objects (standardized by schema_manager) top_n: List of TiDB time TopN objects (standardized by schema_manager) null_count: Number of null values baseline_timestamp: Baseline timestamp (Unix timestamp in seconds) If None, uses TiDB-appropriate baseline based on datetime_type datetime_type: Type of datetime column ('DATETIME', 'TIMESTAMP', 'DATE') is_cumulative_count: False for absolute counts constraint_type: "normal" for InverseTransformSampler, "unique" for BlockShuffler seed: Random seed """ self.constraint_type = constraint_type self.null_count = null_count self.datetime_type = datetime_type.upper() self.column_name = column_name or "unknown_column" self._seed = seed self._null_rng = None if null_count > 0: self._null_rng = CheckpointablePRNG(seed=seed) # Set TiDB-appropriate baseline based on datetime type # Use safer baselines to avoid timestamp overflow issues with large duration scalars if baseline_timestamp is None: if self.datetime_type == "TIMESTAMP": # TiDB MinTimestamp: 1970-01-01 00:00:01 UTC self.baseline_timestamp = int(datetime(1970, 1, 1, 0, 0, 1).timestamp()) else: # Use 1900 as safer baseline instead of year 1 to avoid overflow # This matches the updated safe_baselines in _parse_datetime_value self.baseline_timestamp = int(datetime(1900, 1, 1, 0, 0, 0).timestamp()) else: self.baseline_timestamp = baseline_timestamp # Store TiDB buckets for datetime-specific processing self.tidb_buckets = buckets or [] self.tidb_topn = top_n or [] # Convert TiDB objects to mathematical objects math_buckets = [bucket.to_histogram_bucket() for bucket in self.tidb_buckets] math_topn = [item.to_topn_item() for item in self.tidb_topn] # Select 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 datetime column '{self.column_name}' has no statistical data, using pure new value generation") self.sampler = None # No BlockShuffler needed else: try: self.sampler = BlockShuffler( buckets=math_buckets, top_n=math_topn, seed=seed, ) except ValueError as e: print( f"ERROR: Failed to create BlockShuffler for unique datetime column '{self.column_name}': {e}" ) print(f" math_buckets: {len(math_buckets)} items") print(f" math_topn: {len(math_topn)} items") print(f" tidb_buckets: {len(self.tidb_buckets)} items") print(f" tidb_topn: {len(self.tidb_topn)} items") raise # Initialize new value generation strategy for unique columns self._init_new_value_generation() 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( "DatetimeSampler in 'unique' constraint mode uses BlockShuffler, " "not InverseTransformSampler. Use sampler.sampler instead." ) return self.sampler def _init_new_value_generation(self): """Initialize new value generation strategy for unique columns when BlockShuffler is exhausted""" # Analyze time range from TiDB buckets and TopN all_timestamps = [] # Collect timestamps from TopN for topn_item in self.tidb_topn: try: # Use TiDB time decoder to extract timestamp from utils.tidb_time_decoder import tidb_encoded_to_scalar scalar_value = tidb_encoded_to_scalar(topn_item.encoded) if scalar_value is not None: all_timestamps.append(scalar_value) except Exception: # Skip problematic entries continue # Collect timestamps from buckets for bucket in self.tidb_buckets: try: from utils.tidb_time_decoder import tidb_encoded_to_scalar # Extract scalar from lower and upper bounds lower_scalar = tidb_encoded_to_scalar(bucket.lower_bound) upper_scalar = tidb_encoded_to_scalar(bucket.upper_bound) if lower_scalar is not None: all_timestamps.append(lower_scalar) if upper_scalar is not None: all_timestamps.append(upper_scalar) except Exception: # Skip problematic entries continue if all_timestamps: self._min_timestamp = min(all_timestamps) self._max_timestamp = max(all_timestamps) self._next_new_timestamp = self._max_timestamp + self._get_increment() else: # No statistical data, start from current time now = datetime.now().timestamp() self._min_timestamp = now self._max_timestamp = now self._next_new_timestamp = now # Flag to log exhaustion message only once self._block_exhausted_logged = False def _get_increment(self) -> float: """Get increment value based on datetime type""" if self.datetime_type == "DATE": return 86400 # 1 day in seconds else: # DATETIME or TIMESTAMP return 1 # 1 second def _generate_new_unique_value(self) -> str: """Generate new unique datetime value using incremental strategy""" # Convert scalar to datetime using TiDB time decoder dt = scalar_to_tidb_time(self._next_new_timestamp, self.datetime_type) if dt is None: # Fallback: use direct timestamp conversion dt = datetime.fromtimestamp(self._next_new_timestamp) # Increment for next generation self._next_new_timestamp += self._get_increment() # Format based on datetime type if self.datetime_type == "DATE": return dt.strftime("%Y-%m-%d") else: return dt.strftime("%Y-%m-%d %H:%M:%S") 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 sample_result = self.sampler.sample() if sample_result is None or sample_result.value is None: return None try: # Convert scalar value back to datetime using TiDB time decoder dt = scalar_to_tidb_time(sample_result.value, self.datetime_type) if dt is None: raise ValueError( f"Failed to convert scalar {sample_result.value} to datetime" ) # Format based on datetime type if self.datetime_type == "DATE": return dt.strftime("%Y-%m-%d") else: return dt.strftime("%Y-%m-%d %H:%M:%S") except Exception as e: raise ValueError( f"Invalid datetime value {sample_result.value}: {e}" ) from e elif self.constraint_type == "unique": if phase == SamplingPhase.PREPARE: # PREPARE phase: BlockShuffler + new value generation if self.sampler is None: # No statistical data, use pure new value generation return self._generate_new_unique_value() sample_result = self.sampler.sample() if sample_result is None: # BlockShuffler exhausted, log once and switch to new value generation if not self._block_exhausted_logged: # remaining_capacity = getattr( # self.sampler, "remaining_capacity", "unknown" # ) # print( # f"INFO: BlockShuffler exhausted for unique datetime column '{self.column_name}', switching to new value generation. Remaining capacity: {remaining_capacity}" # ) self._block_exhausted_logged = True return self._generate_new_unique_value() # Handle null_count for BlockShuffler (business layer) 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 if sample_result.value is None: return None try: # Convert scalar value back to datetime using TiDB time decoder dt = scalar_to_tidb_time(sample_result.value, self.datetime_type) if dt is None: raise ValueError( f"Failed to convert scalar {sample_result.value} to datetime" ) # Format based on datetime type if self.datetime_type == "DATE": return dt.strftime("%Y-%m-%d") else: return dt.strftime("%Y-%m-%d %H:%M:%S") except Exception as e: raise ValueError( f"Invalid datetime value {sample_result.value}: {e}" ) from e elif phase == SamplingPhase.REPLAY: # REPLAY phase: continue from where PREPARE left off (no duplicates due to state persistence) return self._generate_new_unique_value() else: raise ValueError(f"Unknown constraint_type: {self.constraint_type}") def sample(self) -> Optional[str]: """ Sample a datetime value and return as formatted string (backward compatibility, defaults to PREPARE phase) Returns: str: Formatted datetime string like "2024-03-19 10:22:42", or None for null values """ return self.sample_with_phase(SamplingPhase.PREPARE) 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 (formatted datetime string) """ if shuffler_result is None or shuffler_result.value is None: return None try: # Convert scalar value back to datetime using TiDB time decoder dt = scalar_to_tidb_time(shuffler_result.value, self.datetime_type) if dt is None: raise ValueError( f"Failed to convert scalar {shuffler_result.value} to datetime" ) # Format based on datetime type if self.datetime_type == "DATE": return dt.strftime("%Y-%m-%d") else: return dt.strftime("%Y-%m-%d %H:%M:%S") except Exception as e: raise ValueError( f"Invalid datetime value {shuffler_result.value}: {e}" ) from e def sample_batch(self, n: int) -> List[Optional[str]]: """ Sample multiple datetime values Args: n: Number of samples to generate Returns: List of formatted datetime strings """ return [self.sample() for _ in range(n)] def set_baseline(self, baseline_datetime: datetime): """ Set the baseline datetime for deltaT calculations Args: baseline_datetime: The baseline datetime to use """ self.baseline_timestamp = int(baseline_datetime.timestamp()) def get_baseline_datetime(self) -> datetime: """ Get the current baseline datetime Returns: datetime: Current baseline datetime """ return datetime.fromtimestamp(self.baseline_timestamp)