Spaces:
Running on Zero
Running on Zero
File size: 3,128 Bytes
ac8ed28 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | import os
import sqlite3
from typing import Tuple, List, Optional, Union, Any
import chardet
import pandas as pd
# Prevent pandas from using pyarrow's string backend.
# The pyarrow DLL from user site-packages is incompatible with this conda env
# and causes an access violation (segfault) in Jupyter's multi-threaded kernel.
pd.options.mode.string_storage = "python"
try:
from prompt_builder import normalize_text
except ImportError:
from src.prompt_builder import normalize_text
def detect_encoding(file_input: Union[str, Any]) -> Optional[str]:
"""Detect the encoding of a CSV file or file-like object using chardet."""
try:
if isinstance(file_input, str):
with open(file_input, 'rb') as f:
raw_data = f.read(100000)
else:
# It's a file-like object (e.g., Streamlit UploadedFile)
pos = file_input.tell()
raw_data = file_input.read(100000)
file_input.seek(pos)
return chardet.detect(raw_data)['encoding']
except Exception:
return None
def load_csv_to_sqlite(file_input: Union[str, Any]) -> Tuple[sqlite3.Connection, str, List[str], List[str]]:
"""
Loads a CSV file into an in-memory SQLite database.
Normalizes column names and maps pandas dtypes to WikiSQL types (real/text).
Args:
file_input (str or file-like): Path to the CSV file, or a file-like object
Returns:
tuple: (conn, table_name, columns, types)
"""
# 1. Try to load the CSV with detected encoding, fallback to others
encodings_to_try = []
detected = detect_encoding(file_input)
if detected:
encodings_to_try.append(detected)
encodings_to_try.extend(['utf-8', 'latin-1', 'cp1252'])
df = None
for enc in encodings_to_try:
try:
if not isinstance(file_input, str):
file_input.seek(0)
df = pd.read_csv(file_input, encoding=enc)
break
except (UnicodeDecodeError, FileNotFoundError, pd.errors.ParserError, pd.errors.EmptyDataError):
continue
if df is None:
file_desc = file_input if isinstance(file_input, str) else getattr(file_input, "name", "input buffer")
raise ValueError(f"Could not read CSV file '{file_desc}'. Tried encodings: {encodings_to_try}")
# 2. Normalize column names (essential for prompt building consistency)
normalized_cols = [normalize_text(col) for col in df.columns]
df.columns = normalized_cols
# 3. Infer WikiSQL types (real vs text)
wiki_types = []
for col in df.columns:
if pd.api.types.is_numeric_dtype(df[col]):
wiki_types.append("real")
else:
wiki_types.append("text")
# 4. Load into SQLite in-memory DB
conn = sqlite3.connect(":memory:", check_same_thread=False)
table_name = "data"
# Write to SQLite
# We use df.to_sql, which automatically maps pandas types to SQLite types safely
df.to_sql(table_name, conn, index=False, if_exists="replace")
return conn, table_name, normalized_cols, wiki_types
|