Coverage for src/dataknobs_data/pandas/type_mapper.py: 0%

1"""Type mapping between DataKnobs Field types and Pandas dtypes.""" 

2 

3import json 

4from dataclasses import dataclass 

5from datetime import datetime, date, time 

6from typing import Any, Dict, Optional, Type, Union 

7 

8import numpy as np 

9import pandas as pd 

10from pandas.api.types import is_datetime64_any_dtype, is_numeric_dtype, is_bool_dtype 

11 

12from dataknobs_data.fields import FieldType 

13 

14 

15@dataclass 

16class PandasTypeMapping: 

17 """Mapping configuration for type conversion.""" 

18 field_type: FieldType 

19 pandas_dtype: Union[str, Type, np.dtype] 

20 nullable: bool = True 

21 converter: Optional[callable] = None 

22 reverse_converter: Optional[callable] = None 

23 

24 

25class TypeMapper: 

26 """Handles type mapping between DataKnobs Field types and Pandas dtypes.""" 

27 

28 def __init__(self): 

29 """Initialize type mapper with default mappings.""" 

30 self._init_mappings() 

31 

32 def _init_mappings(self): 

33 """Initialize type mappings.""" 

34 self._field_to_pandas: Dict[FieldType, PandasTypeMapping] = { 

35 FieldType.STRING: PandasTypeMapping( 

36 field_type=FieldType.STRING, 

37 pandas_dtype="string", # pd.StringDtype() 

38 nullable=True 

39 ), 

40 FieldType.INTEGER: PandasTypeMapping( 

41 field_type=FieldType.INTEGER, 

42 pandas_dtype="Int64", # pd.Int64Dtype() 

43 nullable=True 

44 ), 

45 FieldType.FLOAT: PandasTypeMapping( 

46 field_type=FieldType.FLOAT, 

47 pandas_dtype="Float64", # pd.Float64Dtype() 

48 nullable=True 

49 ), 

50 FieldType.BOOLEAN: PandasTypeMapping( 

51 field_type=FieldType.BOOLEAN, 

52 pandas_dtype="boolean", # pd.BooleanDtype() 

53 nullable=True 

54 ), 

55 FieldType.DATETIME: PandasTypeMapping( 

56 field_type=FieldType.DATETIME, 

57 pandas_dtype="datetime64[ns]", 

58 nullable=True, 

59 converter=self._to_datetime, 

60 reverse_converter=self._from_datetime 

61 ), 

62 FieldType.JSON: PandasTypeMapping( 

63 field_type=FieldType.JSON, 

64 pandas_dtype="object", 

65 nullable=True, 

66 converter=self._to_json_object, 

67 reverse_converter=self._from_json_object 

68 ), 

69 FieldType.BINARY: PandasTypeMapping( 

70 field_type=FieldType.BINARY, 

71 pandas_dtype="object", 

72 nullable=True 

73 ), 

74 FieldType.TEXT: PandasTypeMapping( 

75 field_type=FieldType.TEXT, 

76 pandas_dtype="string", 

77 nullable=True 

78 ), 

79 } 

80 

81 # Reverse mapping from pandas to field types 

82 self._pandas_to_field: Dict[str, FieldType] = { 

83 "string": FieldType.STRING, 

84 "int64": FieldType.INTEGER, 

85 "float64": FieldType.FLOAT, 

86 "boolean": FieldType.BOOLEAN, 

87 "datetime64[ns]": FieldType.DATETIME, 

88 "object": FieldType.STRING, # Default object to STRING, not JSON 

89 } 

90 

91 def field_type_to_pandas(self, field_type: FieldType) -> Union[str, Type, np.dtype]: 

92 """Convert FieldType to pandas dtype. 

93  

94 Args: 

95 field_type: DataKnobs FieldType 

96  

97 Returns: 

98 Corresponding pandas dtype 

99 """ 

100 mapping = self._field_to_pandas.get(field_type) 

101 if mapping: 

102 return mapping.pandas_dtype 

103 return "object" # Default fallback 

104 

105 def pandas_to_field_type(self, dtype: Union[str, np.dtype, Type]) -> FieldType: 

106 """Infer FieldType from pandas dtype. 

107  

108 Args: 

109 dtype: Pandas dtype 

110  

111 Returns: 

112 Corresponding FieldType 

113 """ 

114 dtype_str = str(dtype).lower() 

115 

116 # Direct mapping 

117 if dtype_str in self._pandas_to_field: 

118 return self._pandas_to_field[dtype_str] 

119 

120 # Infer from dtype categories 

121 if "int" in dtype_str: 

122 return FieldType.INTEGER 

123 elif "float" in dtype_str: 

124 return FieldType.FLOAT 

125 elif "bool" in dtype_str: 

126 return FieldType.BOOLEAN 

127 elif "datetime" in dtype_str or "timestamp" in dtype_str: 

128 return FieldType.DATETIME 

129 elif dtype_str == "string": 

130 return FieldType.STRING 

131 elif dtype_str == "object": 

132 return FieldType.STRING 

133 elif "bytes" in dtype_str: 

134 return FieldType.BINARY 

135 

136 return FieldType.STRING # Default fallback 

137 

138 def convert_value_to_pandas(self, value: Any, field_type: FieldType) -> Any: 

139 """Convert a field value to pandas-compatible format. 

140  

141 Args: 

142 value: Value to convert 

143 field_type: Source field type 

144  

145 Returns: 

146 Pandas-compatible value 

147 """ 

148 if value is None: 

149 return pd.NA 

150 

151 mapping = self._field_to_pandas.get(field_type) 

152 if mapping and mapping.converter: 

153 return mapping.converter(value) 

154 

155 return value 

156 

157 def convert_value_from_pandas(self, value: Any, field_type: FieldType) -> Any: 

158 """Convert a pandas value to field-compatible format. 

159  

160 Args: 

161 value: Pandas value 

162 field_type: Target field type 

163  

164 Returns: 

165 Field-compatible value 

166 """ 

167 # Handle pandas NA/NaN/None 

168 # Use try-except to handle arrays and other special cases 

169 try: 

170 if pd.isna(value): 

171 return None 

172 except (TypeError, ValueError): 

173 # pd.isna can fail on arrays/lists 

174 pass 

175 

176 mapping = self._field_to_pandas.get(field_type) 

177 if mapping and mapping.reverse_converter: 

178 return mapping.reverse_converter(value) 

179 

180 # Handle numpy types 

181 if isinstance(value, (np.integer, np.floating, np.bool_)): 

182 return value.item() 

183 

184 return value 

185 

186 def infer_field_type_from_value(self, value: Any) -> FieldType: 

187 """Infer FieldType from a Python value. 

188  

189 Args: 

190 value: Value to analyze 

191  

192 Returns: 

193 Inferred FieldType 

194 """ 

195 if value is None: 

196 return FieldType.STRING # Default for null 

197 

198 # Check for pandas NA separately to avoid array ambiguity 

199 try: 

200 if pd.isna(value): 

201 return FieldType.STRING 

202 except (TypeError, ValueError): 

203 # pd.isna can fail on some types like lists 

204 pass 

205 

206 if isinstance(value, bool) or isinstance(value, np.bool_): 

207 return FieldType.BOOLEAN 

208 elif isinstance(value, (int, np.integer)): 

209 return FieldType.INTEGER 

210 elif isinstance(value, (float, np.floating)): 

211 return FieldType.FLOAT 

212 elif isinstance(value, (datetime, pd.Timestamp)): 

213 return FieldType.DATETIME 

214 elif isinstance(value, bytes): 

215 return FieldType.BINARY 

216 elif isinstance(value, (dict, list)): 

217 return FieldType.JSON 

218 elif isinstance(value, str): 

219 if len(value) > 1000: 

220 return FieldType.TEXT 

221 return FieldType.STRING 

222 

223 return FieldType.JSON # Complex objects as JSON 

224 

225 def cast_series(self, series: pd.Series, field_type: FieldType) -> pd.Series: 

226 """Cast a pandas Series to the appropriate dtype for a FieldType. 

227  

228 Args: 

229 series: Series to cast 

230 field_type: Target field type 

231  

232 Returns: 

233 Casted Series 

234 """ 

235 target_dtype = self.field_type_to_pandas(field_type) 

236 

237 try: 

238 # Special handling for datetime 

239 if field_type == FieldType.DATETIME: 

240 return pd.to_datetime(series, errors='coerce') 

241 

242 # Special handling for JSON 

243 if field_type == FieldType.JSON: 

244 return series.apply(self._ensure_json_serializable) 

245 

246 # Standard casting 

247 return series.astype(target_dtype) 

248 except (TypeError, ValueError): 

249 # If casting fails, return as object dtype 

250 return series.astype("object") 

251 

252 @staticmethod 

253 def _to_datetime(value: Any) -> pd.Timestamp: 

254 """Convert value to pandas Timestamp.""" 

255 if isinstance(value, str): 

256 return pd.Timestamp(value) 

257 elif isinstance(value, datetime): 

258 return pd.Timestamp(value) 

259 elif isinstance(value, (int, float)): 

260 # Assume Unix timestamp 

261 return pd.Timestamp(value, unit='s') 

262 return value 

263 

264 @staticmethod 

265 def _from_datetime(value: Any) -> datetime: 

266 """Convert pandas Timestamp to datetime.""" 

267 if isinstance(value, pd.Timestamp): 

268 return value.to_pydatetime() 

269 elif isinstance(value, str): 

270 return pd.Timestamp(value).to_pydatetime() 

271 return value 

272 

273 @staticmethod 

274 def _to_json_object(value: Any) -> Any: 

275 """Ensure value is JSON-serializable object.""" 

276 if isinstance(value, str): 

277 try: 

278 return json.loads(value) 

279 except (json.JSONDecodeError, TypeError): 

280 return value 

281 return value 

282 

283 @staticmethod 

284 def _from_json_object(value: Any) -> Any: 

285 """Convert object to JSON-compatible format.""" 

286 if isinstance(value, (dict, list)): 

287 return value 

288 elif isinstance(value, str): 

289 try: 

290 return json.loads(value) 

291 except (json.JSONDecodeError, TypeError): 

292 return value 

293 return value 

294 

295 @staticmethod 

296 def _ensure_json_serializable(value: Any) -> Any: 

297 """Ensure value is JSON-serializable.""" 

298 if pd.isna(value): 

299 return None 

300 if isinstance(value, (dict, list)): 

301 return value 

302 if isinstance(value, str): 

303 try: 

304 return json.loads(value) 

305 except (json.JSONDecodeError, TypeError): 

306 return value 

307 # Convert other types to string representation 

308 return str(value) 

309 

310 def infer_field_type(self, series: pd.Series) -> str: 

311 """Infer field type from a pandas Series. 

312  

313 Args: 

314 series: Series to analyze 

315  

316 Returns: 

317 Field type string 

318 """ 

319 # Remove nulls for analysis 

320 non_null = series.dropna() 

321 

322 if len(non_null) == 0: 

323 return "string" # Default for empty 

324 

325 # Check dtypes 

326 if is_bool_dtype(non_null): 

327 return "boolean" 

328 elif is_datetime64_any_dtype(non_null): 

329 return "datetime" 

330 elif is_numeric_dtype(non_null): 

331 # Check if all values are integers 

332 if non_null.apply(lambda x: isinstance(x, (int, np.integer)) or (isinstance(x, float) and x.is_integer())).all(): 

333 return "integer" 

334 else: 

335 return "number" 

336 else: 

337 # Check values for special types 

338 sample = non_null.iloc[0] if len(non_null) > 0 else None 

339 if sample is not None: 

340 if isinstance(sample, (datetime, pd.Timestamp)): 

341 return "datetime" 

342 elif isinstance(sample, (pd.Timestamp, datetime)): 

343 return "datetime" 

344 elif isinstance(sample, pd._libs.tslibs.timestamps.Timestamp): 

345 return "datetime" 

346 elif isinstance(sample, (pd._libs.tslibs.nattype.NaTType)): 

347 return "datetime" 

348 elif hasattr(sample, '__class__') and 'date' in sample.__class__.__name__.lower(): 

349 return "date" 

350 elif hasattr(sample, '__class__') and 'time' in sample.__class__.__name__.lower(): 

351 return "time" 

352 

353 return "string" # Default 

354 

355 def get_pandas_dtype(self, field_type: str) -> str: 

356 """Get pandas dtype for a field type string. 

357  

358 Args: 

359 field_type: Field type string 

360  

361 Returns: 

362 Pandas dtype string 

363 """ 

364 dtype_map = { 

365 "string": "object", 

366 "integer": "int64", 

367 "number": "float64", 

368 "float": "float64", 

369 "boolean": "bool", 

370 "datetime": "datetime64[ns]", 

371 "date": "object", 

372 "time": "object", 

373 "json": "object", 

374 "binary": "object", 

375 "text": "object", 

376 } 

377 return dtype_map.get(field_type.lower(), "object") 

378 

379 def convert_value(self, value: Any, target_type: str) -> Any: 

380 """Convert a value to target type. 

381  

382 Args: 

383 value: Value to convert 

384 target_type: Target type string 

385  

386 Returns: 

387 Converted value 

388 """ 

389 if value is None or pd.isna(value): 

390 return None 

391 

392 target_type = target_type.lower() 

393 

394 if target_type == "integer": 

395 if isinstance(value, str): 

396 return int(float(value)) 

397 return int(value) 

398 elif target_type == "number" or target_type == "float": 

399 return float(value) 

400 elif target_type == "string": 

401 return str(value) 

402 elif target_type == "boolean": 

403 if isinstance(value, str): 

404 return value.lower() in ('true', '1', 'yes') 

405 return bool(value) 

406 elif target_type == "datetime": 

407 if isinstance(value, str): 

408 return pd.Timestamp(value) 

409 return value 

410 elif target_type == "date": 

411 if isinstance(value, str): 

412 return pd.Timestamp(value).date() 

413 elif hasattr(value, 'date'): 

414 return value.date() 

415 return value 

416 elif target_type == "time": 

417 if isinstance(value, str): 

418 return pd.Timestamp(value).time() 

419 elif hasattr(value, 'time'): 

420 return value.time() 

421 return value 

422 

423 return value 

424 

425 def cast_dataframe_dtypes(self, df: pd.DataFrame, dtype_map: Dict[str, str]) -> pd.DataFrame: 

426 """Cast DataFrame columns to specified dtypes. 

427  

428 Args: 

429 df: DataFrame to cast 

430 dtype_map: Dictionary of column: dtype 

431  

432 Returns: 

433 DataFrame with casted dtypes 

434 """ 

435 result_df = df.copy() 

436 

437 for col, dtype in dtype_map.items(): 

438 if col in result_df.columns: 

439 try: 

440 if dtype == "string": 

441 # Use string dtype for nullable strings 

442 result_df[col] = result_df[col].astype("string") 

443 else: 

444 result_df[col] = result_df[col].astype(dtype) 

445 except (TypeError, ValueError): 

446 # If casting fails, leave as is 

447 pass 

448 

449 return result_df 

450 

451 def normalize_timezone(self, series: pd.Series, target_tz: str) -> pd.Series: 

452 """Normalize timezone for datetime series. 

453  

454 Args: 

455 series: Datetime series 

456 target_tz: Target timezone 

457  

458 Returns: 

459 Series with normalized timezone 

460 """ 

461 if not is_datetime64_any_dtype(series): 

462 # Try to convert to datetime first 

463 series = pd.to_datetime(series, errors='coerce') 

464 

465 # If series is timezone-naive, localize it 

466 if series.dt.tz is None: 

467 return series.dt.tz_localize(target_tz) 

468 else: 

469 # If timezone-aware, convert to target timezone 

470 return series.dt.tz_convert(target_tz) 

471 

472 def get_optimal_dtype(self, series: pd.Series) -> str: 

473 """Determine optimal dtype for a Series based on its values. 

474  

475 Args: 

476 series: Series to analyze 

477  

478 Returns: 

479 Optimal dtype string 

480 """ 

481 # Remove nulls for analysis 

482 non_null = series.dropna() 

483 

484 if len(non_null) == 0: 

485 return "string" # Default for empty 

486 

487 # Try to infer the best dtype 

488 try: 

489 # Check for boolean 

490 if non_null.apply(lambda x: isinstance(x, bool)).all(): 

491 return "bool" 

492 

493 # Check for integer 

494 if non_null.apply(lambda x: isinstance(x, (int, np.integer)) or (isinstance(x, float) and x.is_integer())).all(): 

495 # Determine the smallest int type that can hold the values 

496 min_val = non_null.min() 

497 max_val = non_null.max() 

498 

499 if min_val >= -128 and max_val <= 127: 

500 return "int8" 

501 elif min_val >= -32768 and max_val <= 32767: 

502 return "int16" 

503 elif min_val >= -2147483648 and max_val <= 2147483647: 

504 return "int32" 

505 else: 

506 return "int64" 

507 

508 # Check for float 

509 if non_null.apply(lambda x: isinstance(x, (int, float, np.number))).all(): 

510 # For floats, prefer float32 for small ranges 

511 max_val = non_null.abs().max() 

512 if max_val <= 3.4e38: 

513 return "float32" 

514 else: 

515 return "float64" 

516 

517 # Check for datetime 

518 try: 

519 pd.to_datetime(non_null) 

520 return "datetime64[ns]" 

521 except (ValueError, TypeError): 

522 pass 

523 

524 # Default to object for strings and mixed types 

525 return "object" 

526 except Exception: 

527 return "object"