|
8 | 8 | from dataclasses import dataclass |
9 | 9 | import math |
10 | 10 | from numbers import Real |
| 11 | +from operator import itemgetter |
11 | 12 | from typing import Any |
12 | 13 | from urllib.parse import urlparse |
13 | 14 |
|
| 15 | +import numpy as np |
| 16 | + |
14 | 17 | from ..core import UnsupportedOperationError |
15 | 18 |
|
16 | 19 |
|
@@ -211,7 +214,11 @@ def preflight( |
211 | 214 | f"{profile.name} permits {profile.max_row_bytes}.", |
212 | 215 | declared_row_bytes=declared_row_bytes, maximum=profile.max_row_bytes, |
213 | 216 | ) |
214 | | - for row_ordinal, row in enumerate(rows or (), start=1): |
| 217 | + rows = list(rows or ()) |
| 218 | + if not rows or _rows_within_limits(profile, variables, rows): |
| 219 | + return |
| 220 | + # Slow path only reports the first failing row in row-major order. |
| 221 | + for row_ordinal, row in enumerate(rows, start=1): |
215 | 222 | row_bytes = 8 |
216 | 223 | for variable in variables: |
217 | 224 | if variable.get("storage_kind") == "numeric": |
@@ -312,6 +319,50 @@ def preflight( |
312 | 319 | ) |
313 | 320 |
|
314 | 321 |
|
| 322 | +def _rows_within_limits( |
| 323 | + profile: SqlProfile, variables: list[Mapping[str, Any]], |
| 324 | + rows: list[Mapping[str, Any]], |
| 325 | +) -> bool: |
| 326 | + """Column-wise proof that the row-by-row checks in preflight cannot fail. |
| 327 | +
|
| 328 | + Per-column maxima bound every row's row and statement size from above, so |
| 329 | + True means no row exceeds a limit; False defers to the exact slow path. |
| 330 | + """ |
| 331 | + row_bytes, statement_bytes = 8, 32 |
| 332 | + for variable in variables: |
| 333 | + name = str(variable["physical_name"]) |
| 334 | + try: |
| 335 | + column = list(map(itemgetter(name), rows)) |
| 336 | + except KeyError: |
| 337 | + return False |
| 338 | + if variable.get("storage_kind") == "numeric": |
| 339 | + if not set(map(type, column)) <= _PLAIN_NUMERIC_TYPES: |
| 340 | + return False |
| 341 | + values = np.array(column, dtype=float) # None becomes NaN |
| 342 | + if np.isinf(values).any() or np.isnan(values).sum() != column.count(None): |
| 343 | + return False |
| 344 | + row_bytes += 8 |
| 345 | + statement_bytes += 24 |
| 346 | + continue |
| 347 | + if set(map(type, column)) != {str}: |
| 348 | + return False |
| 349 | + longest = max(map(len, map(str.encode, column))) |
| 350 | + if longest > profile.max_text_value_bytes: |
| 351 | + return False |
| 352 | + row_bytes += 20 if profile.name in MYSQL_WIRE_PROFILES else longest |
| 353 | + statement_bytes += longest + 8 |
| 354 | + return ( |
| 355 | + (profile.max_row_bytes is None or row_bytes <= profile.max_row_bytes) |
| 356 | + and ( |
| 357 | + profile.max_statement_bytes is None |
| 358 | + or statement_bytes <= profile.max_statement_bytes |
| 359 | + ) |
| 360 | + ) |
| 361 | + |
| 362 | + |
| 363 | +_PLAIN_NUMERIC_TYPES = {float, type(None)} |
| 364 | + |
| 365 | + |
315 | 366 | def _row_storage_bytes(profile: SqlProfile, variable: Mapping[str, Any]) -> int: |
316 | 367 | if variable.get("storage_kind") == "numeric": |
317 | 368 | return 8 |
|
0 commit comments