(prompt_db, db_id, db_path, limit_value=3, normalization=True)
| 289 | |
| 290 | |
| 291 | def extract_create_table_prompt_column_example(prompt_db, db_id, db_path, limit_value=3, normalization=True): |
| 292 | table_query = "SELECT * FROM sqlite_master WHERE type='table';" |
| 293 | tables = sqlite3.connect(db_path).cursor().execute(table_query).fetchall() |
| 294 | prompt = "" |
| 295 | table_names = [] |
| 296 | for table in tables: |
| 297 | table_name = table[1] |
| 298 | if table_name == 'sqlite_sequence': |
| 299 | continue |
| 300 | table_names.append(table_name) |
| 301 | if normalization: |
| 302 | table_name = table_name.lower() |
| 303 | create_table_statement = table[-1] |
| 304 | |
| 305 | table_info_query = f"PRAGMA table_info({table_name});" |
| 306 | headers = [x[1] for x in sqlite3.connect(db_path).cursor().execute(table_info_query).fetchall()] |
| 307 | if normalization: |
| 308 | create_table_statement = normalize_create_table(table_name, create_table_statement) |
| 309 | create_table_statement = create_table_statement.lower() |
| 310 | headers = [x.lower() for x in headers] |
| 311 | prompt += create_table_statement + ";\n" |
| 312 | if limit_value > 0: |
| 313 | prompt_columns = [] |
| 314 | for col_name in headers: |
| 315 | if col_name.lower() == "index": |
| 316 | top_k_rows = list(range(limit_value)) |
| 317 | top_k_rows = ' '.join([str(x) for x in top_k_rows]) |
| 318 | else: |
| 319 | top_k_row_query = f"SELECT distinct \"{col_name}\" FROM {table_name} LIMIT {limit_value};" |
| 320 | top_k_rows = sqlite3.connect(db_path).cursor().execute(top_k_row_query).fetchall() |
| 321 | top_k_rows = [x[0].strip() if isinstance(x[0], str) else x[0] |
| 322 | for x in top_k_rows] # remove \n and space prefix and suffix in cell value |
| 323 | top_k_rows = [x if x is not None else "" for x in top_k_rows] |
| 324 | top_k_rows = ', '.join([str(x) if is_number(x) else '"' + str(x) + '"' for x in top_k_rows][:limit_value]) |
| 325 | |
| 326 | prompt_columns.append(f"{col_name}: {top_k_rows};") |
| 327 | |
| 328 | prompt += "/*\n" |
| 329 | prompt += f"Columns in {table_name} and {limit_value} distinct examples in each column:\n" |
| 330 | prompt += "\n".join(prompt_columns) |
| 331 | prompt += "\n*/\n" |
| 332 | prompt += "\n" |
| 333 | |
| 334 | return prompt, table_names |
| 335 | |
| 336 | |
| 337 | def generate_db_prompt_spider(root_dir, dataset, db_id, prompt_db="CreateTableSelect", limit_value=3, normalization=True): |
no test coverage detected