(self,
data_file: str,
db_file: Optional[str] = None,
db_password: str = 'password',
max_round: int = 20,
env_driver: str = 'docker',
env_options: Optional[dict] = None,
**configs)
| 37 | class DBBenchTask(Task): |
| 38 | |
| 39 | def __init__(self, |
| 40 | data_file: str, |
| 41 | db_file: Optional[str] = None, |
| 42 | db_password: str = 'password', |
| 43 | max_round: int = 20, |
| 44 | env_driver: str = 'docker', |
| 45 | env_options: Optional[dict] = None, |
| 46 | **configs): |
| 47 | super().__init__(**configs) |
| 48 | self.full_async = True |
| 49 | self.logger = logging.getLogger(__name__) |
| 50 | |
| 51 | self.max_round = max_round |
| 52 | self.data_file = data_file |
| 53 | self.db_root_dir = db_file |
| 54 | |
| 55 | self.dataset = [] |
| 56 | # Load dataset |
| 57 | with open(self.data_file) as f: |
| 58 | raw_data = f.read() |
| 59 | if self.data_file.endswith("json"): |
| 60 | data = json.loads(raw_data) |
| 61 | else: # Assuming jsonl |
| 62 | data = [json.loads(line) for line in raw_data.strip().split('\n')] |
| 63 | |
| 64 | for entry in data: |
| 65 | ans_key = "answer_md5" if entry["type"][0] in ("INSERT", "DELETE", "UPDATE") else "label" |
| 66 | ans = entry.pop(ans_key, None) # Use pop with default |
| 67 | inp = entry |
| 68 | self.dataset.append((inp, ans)) |
| 69 | |
| 70 | self.env_delegation = DBBenchEnvironmentDelegation(db_password) |
| 71 | self.env_controller = create_controller(env_driver, self.env_delegation, **env_options) |
| 72 | self.env_controller_background_task = None |
| 73 | |
| 74 | self.logger.info(f"DBBench initialized with {len(self.dataset)} samples. Root dir: {self.db_root_dir}") |
| 75 | |
| 76 | def get_indices(self) -> List[SampleIndex]: |
| 77 | return list(range(len(self.dataset))) |
nothing calls this directly
no test coverage detected