1  import pandas as pd
   2  
   3  class Sampler:
   4      def __init__(self, sampling_config: dict):
   5          self.sampling_config = sampling_config
   6  
   7      def sample(self, df: pd.DataFrame) -> pd.DataFrame:
   8          strategy = self.sampling_config.get("type")
   9  
  10          if strategy == "random":
  11              return self.random_sample(df)
  12          elif strategy == "stratified":
  13              return self.stratified_sample(df)
  14          elif strategy == "stratified-proportional":
  15              return self.stratified_proportional_sample(df)
  16          else:
  17              raise ValueError(f"Unsupported sampling strategy: {strategy}")
  18  
  19      def random_sample(self, df: pd.DataFrame) -> pd.DataFrame:
  20          sample_size = self.sampling_config.get("sample_size")
  21          if sample_size is None:
  22              raise ValueError("`sample_size` must be specified for random sampling.")
  23          return df.sample(n=min(sample_size, len(df)), random_state=42)
  24  
  25      def stratified_sample(self, df: pd.DataFrame) -> pd.DataFrame:
  26          by = self.sampling_config.get("by")
  27          ratio = self.sampling_config.get("ratio")
  28  
  29          if not by:
  30              raise ValueError("`by` must be specified for stratified sampling.")
  31  
  32          grouped = df.groupby(by)
  33          sampled_groups = []
  34  
  35          # Auto-compute equal ratio if ratio not passed
  36          if ratio is None:
  37              default_ratio = self.sampling_config.get("default_per_group", 10)
  38              print(f"[INFO] No ratio provided. Using default ratio of {default_ratio} per group.")
  39              for group_key, group_df in grouped:
  40                  n = min(default_ratio, len(group_df))
  41                  if n > 0:
  42                      sampled_groups.append(group_df.sample(n=n, random_state=42))
  43          else:
  44              for group_key, group_df in grouped:
  45                  lookup_key = ":".join(map(str, group_key)) if isinstance(group_key, tuple) else str(group_key)
  46                  group_n = ratio.get(lookup_key, 0) if isinstance(ratio, dict) else ratio
  47                  n = min(group_n, len(group_df))
  48                  if n > 0:
  49                      sampled_groups.append(group_df.sample(n=n, random_state=42))
  50                      print(f"[STRATIFIED] {lookup_key} → requested: {group_n}, sampled: {n}")
  51  
  52          return pd.concat(sampled_groups, ignore_index=True) if sampled_groups else pd.DataFrame(columns=df.columns)
  53  
  54      def stratified_proportional_sample(self, df: pd.DataFrame) -> pd.DataFrame:
  55          by = self.sampling_config.get("by")
  56          proportion = self.sampling_config.get("proportion")
  57  
  58          if not by:
  59              raise ValueError("`by` must be specified for stratified-proportional sampling.")
  60  
  61          grouped = df.groupby(by)
  62          sampled_groups = []
  63  
  64          # Auto-equal proportions if not passed
  65          if proportion is None:
  66              equal_p = 1 / len(grouped)
  67              print(f"[INFO] No proportion provided. Using equal {round(equal_p * 100)}% per group.")
  68              for group_key, group_df in grouped:
  69                  n = int(len(group_df) * equal_p)
  70                  if n > 0:
  71                      sampled_groups.append(group_df.sample(n=n, random_state=42))
  72          else:
  73              for group_key, group_df in grouped:
  74                  lookup_key = ":".join(map(str, group_key)) if isinstance(group_key, tuple) else str(group_key)
  75                  p = proportion.get(lookup_key, 0) if isinstance(proportion, dict) else proportion
  76                  n = int(len(group_df) * p)
  77                  if n > 0:
  78                      sampled_groups.append(group_df.sample(n=n, random_state=42))
  79                      print(f"[STRAT-PROP] {lookup_key} → {p*100}% → sampled: {n}")
  80  
  81          return pd.concat(sampled_groups, ignore_index=True) if sampled_groups else pd.DataFrame(columns=df.columns)
  82