From 2985d6aae54bf855199704613c4a4999ef945f34 Mon Sep 17 00:00:00 2001 From: Ray Tien Date: Mon, 10 Aug 2026 11:57:45 +0800 Subject: [PATCH 01/10] fix(wren): preserve wide MySQL decimal precision --- core/wren/src/wren/connector/mysql.py | 25 +++++++++---------- .../tests/connectors/test_mysql_connector.py | 16 ++++++++++++ core/wren/tests/unit/test_mysql_helpers.py | 13 +++------- 3 files changed, 32 insertions(+), 22 deletions(-) diff --git a/core/wren/src/wren/connector/mysql.py b/core/wren/src/wren/connector/mysql.py index f6ccdc2838..19a286d8e7 100644 --- a/core/wren/src/wren/connector/mysql.py +++ b/core/wren/src/wren/connector/mysql.py @@ -133,12 +133,10 @@ def create_connector(data_source: DataSource, connection_info) -> MySqlConnector # Arrow conversion helpers # --------------------------------------------------------------------------- -# MySQL ``DECIMAL(M, D)`` allows ``M`` up to 65 and ``D`` up to 30, while -# PyArrow's ``decimal128`` only supports precision up to 38. We clamp the -# precision derived from ``cursor.description`` to ``38`` and the scale to -# ``min(precision, 30)`` so PyArrow can still represent the value. A future -# change could switch to ``decimal256`` when MySQL exceeds 38 digits. +# MySQL ``DECIMAL(M, D)`` allows ``M`` up to 65 and ``D`` up to 30. Use +# ``decimal128`` through precision 38 and ``decimal256`` above that boundary. _ARROW_DECIMAL128_MAX_PRECISION = 38 +_MYSQL_DECIMAL_MAX_PRECISION = 65 _MYSQL_DECIMAL_MAX_SCALE = 30 # Fallback used when ``cursor.description`` does not carry precision/scale # (e.g. for the legacy ``FIELD_TYPE.DECIMAL`` code or non-MySQLdb cursors). @@ -268,7 +266,7 @@ def _arrow_decimal_from_mysql_field( scale: int | None, is_unsigned: bool = False, ) -> pa.DataType: - """Derive a ``pa.decimal128`` type from a MySQLdb ``cursor.description`` entry. + """Derive an Arrow decimal type from a MySQLdb cursor description entry. MySQLdb populates ``description[4]`` (PEP 249 ``precision``) with the ``MYSQL_FIELD.length`` — i.e. the *display length*, which includes one @@ -278,11 +276,10 @@ def _arrow_decimal_from_mysql_field( M = length - (1 if unsigned else 0) - (1 if D > 0 else 0) - MySQL allows precision up to 65 and scale up to 30, but Arrow - ``decimal128`` caps precision at 38. We clamp precision to 38 and clamp - scale to ``min(scale, precision, 30)`` so any value MySQL accepts (within - the 38-digit Arrow ceiling) round-trips correctly. The previous - hard-coded ``decimal128(38, 9)`` would silently lose digits when ``D > 9``. + MySQL allows precision up to 65 and scale up to 30. Arrow ``decimal128`` + covers precision up to 38, while ``decimal256`` covers every wider MySQL + decimal. The previous hard-coded ``decimal128(38, 9)`` would silently lose + digits when ``D > 9``. """ if display_length is None or display_length <= 0: precision = _MYSQL_DECIMAL_FALLBACK_PRECISION @@ -295,9 +292,11 @@ def _arrow_decimal_from_mysql_field( precision = _MYSQL_DECIMAL_FALLBACK_PRECISION if scale is None or scale < 0: scale = _MYSQL_DECIMAL_FALLBACK_SCALE - precision = min(int(precision), _ARROW_DECIMAL128_MAX_PRECISION) + precision = min(int(precision), _MYSQL_DECIMAL_MAX_PRECISION) scale = min(int(scale), _MYSQL_DECIMAL_MAX_SCALE, precision) - return pa.decimal128(precision, scale) + if precision <= _ARROW_DECIMAL128_MAX_PRECISION: + return pa.decimal128(precision, scale) + return pa.decimal256(precision, scale) def _mysql_field_arrow_type( diff --git a/core/wren/tests/connectors/test_mysql_connector.py b/core/wren/tests/connectors/test_mysql_connector.py index 7decbdd0e3..22ae1e9995 100644 --- a/core/wren/tests/connectors/test_mysql_connector.py +++ b/core/wren/tests/connectors/test_mysql_connector.py @@ -118,6 +118,22 @@ def test_decimal_large_scale(connector: MySqlConnector) -> None: assert tbl.column("a").to_pylist()[0] == Decimal("12345.123456789012345") +def test_decimal_above_decimal128_uses_decimal256(connector: MySqlConnector) -> None: + value = Decimal( + "12345678901234567890123456789012345.123456789012345678901234567890" + ) + + _exec(connector, "DROP TABLE IF EXISTS t_dec_wide") + _exec(connector, "CREATE TABLE t_dec_wide (a DECIMAL(65, 30))") + with closing(connector.connection.cursor()) as cursor: + cursor.execute("INSERT INTO t_dec_wide VALUES (%s)", (value,)) + + tbl = connector.query("SELECT a FROM t_dec_wide") + + assert tbl.schema.field("a").type == pa.decimal256(65, 30) + assert tbl.column("a").to_pylist() == [value] + + def test_float_and_double(connector: MySqlConnector) -> None: _exec(connector, "DROP TABLE IF EXISTS t_real") _exec(connector, "CREATE TABLE t_real (a FLOAT, b DOUBLE)") diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index 06120c14bc..c543852388 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -180,9 +180,7 @@ def test_decimal_type_passthrough() -> None: # DECIMAL(12, 4) signed → MySQLdb description length = 12 + 1 (sign) + 1 # (decimal point) = 14. t = _arrow_decimal_from_mysql_field(14, 4, is_unsigned=False) - assert pa.types.is_decimal(t) - assert t.precision == 12 - assert t.scale == 4 + assert t == pa.decimal128(12, 4) def test_decimal_type_unsigned_recovers_precision() -> None: @@ -204,16 +202,13 @@ def test_decimal_type_high_scale() -> None: precision >= 30.""" # DECIMAL(38, 30) signed → length = 38 + 1 + 1 = 40. t = _arrow_decimal_from_mysql_field(40, 30, is_unsigned=False) - assert t.precision == 38 - assert t.scale == 30 + assert t == pa.decimal128(38, 30) -def test_decimal_type_clamps_above_arrow_max_precision() -> None: - """MySQL precision tops at 65; Arrow decimal128 tops at 38. We clamp.""" +def test_decimal_type_above_decimal128_uses_decimal256() -> None: # DECIMAL(65, 30) signed → length = 65 + 1 + 1 = 67. t = _arrow_decimal_from_mysql_field(67, 30, is_unsigned=False) - assert t.precision == 38 - assert t.scale == 30 + assert t == pa.decimal256(65, 30) def test_decimal_type_none_uses_fallback() -> None: From 3c768f57a565eaa1eec7210e32ac8c82b6ce20e0 Mon Sep 17 00:00:00 2001 From: Ray Tien Date: Mon, 10 Aug 2026 12:10:05 +0800 Subject: [PATCH 02/10] fix(wren): support Doris Decimal256 range --- .github/workflows/wren-ci.yml | 4 +++- core/wren/src/wren/connector/mysql.py | 22 ++++++++++++---------- core/wren/tests/unit/test_mysql_helpers.py | 15 +++++++++++++++ 3 files changed, 30 insertions(+), 11 deletions(-) diff --git a/.github/workflows/wren-ci.yml b/.github/workflows/wren-ci.yml index ea64515bfa..b1645de737 100644 --- a/.github/workflows/wren-ci.yml +++ b/.github/workflows/wren-ci.yml @@ -89,7 +89,9 @@ jobs: marker: postgres - datasource: mysql extra: mysql - test_file: tests/connectors/test_mysql.py + test_file: >- + tests/connectors/test_mysql.py + tests/connectors/test_mysql_connector.py marker: mysql defaults: run: diff --git a/core/wren/src/wren/connector/mysql.py b/core/wren/src/wren/connector/mysql.py index 19a286d8e7..78cb5205c8 100644 --- a/core/wren/src/wren/connector/mysql.py +++ b/core/wren/src/wren/connector/mysql.py @@ -133,11 +133,12 @@ def create_connector(data_source: DataSource, connection_info) -> MySqlConnector # Arrow conversion helpers # --------------------------------------------------------------------------- -# MySQL ``DECIMAL(M, D)`` allows ``M`` up to 65 and ``D`` up to 30. Use -# ``decimal128`` through precision 38 and ``decimal256`` above that boundary. +# MySQL ``DECIMAL(M, D)`` allows ``M`` up to 65 and ``D`` up to 30. Doris uses +# the same protocol and supports precision and scale up to 76 when Decimal256 is +# enabled. Use Arrow's full Decimal256 range so the shared conversion preserves +# both data sources. _ARROW_DECIMAL128_MAX_PRECISION = 38 -_MYSQL_DECIMAL_MAX_PRECISION = 65 -_MYSQL_DECIMAL_MAX_SCALE = 30 +_ARROW_DECIMAL256_MAX_PRECISION = 76 # Fallback used when ``cursor.description`` does not carry precision/scale # (e.g. for the legacy ``FIELD_TYPE.DECIMAL`` code or non-MySQLdb cursors). _MYSQL_DECIMAL_FALLBACK_PRECISION = 38 @@ -276,10 +277,11 @@ def _arrow_decimal_from_mysql_field( M = length - (1 if unsigned else 0) - (1 if D > 0 else 0) - MySQL allows precision up to 65 and scale up to 30. Arrow ``decimal128`` - covers precision up to 38, while ``decimal256`` covers every wider MySQL - decimal. The previous hard-coded ``decimal128(38, 9)`` would silently lose - digits when ``D > 9``. + MySQL allows precision up to 65 and scale up to 30. Doris uses the same + protocol and supports precision and scale up to 76 when Decimal256 is + enabled. Arrow ``decimal128`` covers precision up to 38, while + ``decimal256`` covers both data sources' wider decimals. The previous + hard-coded ``decimal128(38, 9)`` would silently lose digits when ``D > 9``. """ if display_length is None or display_length <= 0: precision = _MYSQL_DECIMAL_FALLBACK_PRECISION @@ -292,8 +294,8 @@ def _arrow_decimal_from_mysql_field( precision = _MYSQL_DECIMAL_FALLBACK_PRECISION if scale is None or scale < 0: scale = _MYSQL_DECIMAL_FALLBACK_SCALE - precision = min(int(precision), _MYSQL_DECIMAL_MAX_PRECISION) - scale = min(int(scale), _MYSQL_DECIMAL_MAX_SCALE, precision) + precision = min(int(precision), _ARROW_DECIMAL256_MAX_PRECISION) + scale = min(int(scale), precision) if precision <= _ARROW_DECIMAL128_MAX_PRECISION: return pa.decimal128(precision, scale) return pa.decimal256(precision, scale) diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index c543852388..de18545fcb 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -205,12 +205,27 @@ def test_decimal_type_high_scale() -> None: assert t == pa.decimal128(38, 30) +def test_decimal_type_precision_39_starts_decimal256() -> None: + # DECIMAL(39, 30) signed → length = 39 + 1 + 1 = 41. + t = _arrow_decimal_from_mysql_field(41, 30, is_unsigned=False) + assert t == pa.decimal256(39, 30) + + def test_decimal_type_above_decimal128_uses_decimal256() -> None: # DECIMAL(65, 30) signed → length = 65 + 1 + 1 = 67. t = _arrow_decimal_from_mysql_field(67, 30, is_unsigned=False) assert t == pa.decimal256(65, 30) +def test_decimal_type_preserves_doris_decimal256_maximum() -> None: + # Doris DECIMAL(76, 76) signed → length = 76 + 1 + 1 = 78. + t = _arrow_decimal_from_mysql_field(78, 76, is_unsigned=False) + value_decimal = "0." + "9" * 76 + + assert t == pa.decimal256(76, 76) + assert str(_build_mysql_column([value_decimal], t)[0].as_py()) == value_decimal + + def test_decimal_type_none_uses_fallback() -> None: t = _arrow_decimal_from_mysql_field(None, None) assert t.precision == 38 From afc4fa889d23d182f137c80f68d8d017879d6251 Mon Sep 17 00:00:00 2001 From: Ray Tien Date: Mon, 10 Aug 2026 12:51:37 +0800 Subject: [PATCH 03/10] fix(wren): preserve value-wide MySQL decimals --- core/wren/src/wren/connector/mysql.py | 93 ++++++++++++++++- .../tests/connectors/test_mysql_connector.py | 50 ++++++++++ core/wren/tests/unit/test_mysql_helpers.py | 99 +++++++++++++++++++ 3 files changed, 237 insertions(+), 5 deletions(-) diff --git a/core/wren/src/wren/connector/mysql.py b/core/wren/src/wren/connector/mysql.py index 78cb5205c8..8afde76337 100644 --- a/core/wren/src/wren/connector/mysql.py +++ b/core/wren/src/wren/connector/mysql.py @@ -275,7 +275,7 @@ def _arrow_decimal_from_mysql_field( when the column is signed. The declared ``DECIMAL(M, D)`` precision ``M`` is recovered as:: - M = length - (1 if unsigned else 0) - (1 if D > 0 else 0) + M = length - (0 if unsigned else 1) - (1 if D > 0 else 0) MySQL allows precision up to 65 and scale up to 30. Doris uses the same protocol and supports precision and scale up to 76 when Decimal256 is @@ -283,6 +283,20 @@ def _arrow_decimal_from_mysql_field( ``decimal256`` covers both data sources' wider decimals. The previous hard-coded ``decimal128(38, 9)`` would silently lose digits when ``D > 9``. """ + precision, scale = _mysql_decimal_precision_scale( + display_length, scale, is_unsigned + ) + precision = min(precision, _ARROW_DECIMAL256_MAX_PRECISION) + scale = min(scale, precision) + return _arrow_decimal_type(precision, scale) + + +def _mysql_decimal_precision_scale( + display_length: int | None, + scale: int | None, + is_unsigned: bool, +) -> tuple[int, int]: + """Return raw MySQL decimal precision and scale without Arrow clamping.""" if display_length is None or display_length <= 0: precision = _MYSQL_DECIMAL_FALLBACK_PRECISION else: @@ -294,13 +308,71 @@ def _arrow_decimal_from_mysql_field( precision = _MYSQL_DECIMAL_FALLBACK_PRECISION if scale is None or scale < 0: scale = _MYSQL_DECIMAL_FALLBACK_SCALE - precision = min(int(precision), _ARROW_DECIMAL256_MAX_PRECISION) + precision = int(precision) scale = min(int(scale), precision) + return precision, scale + + +def _arrow_decimal_type(precision: int, scale: int) -> pa.DataType: if precision <= _ARROW_DECIMAL128_MAX_PRECISION: return pa.decimal128(precision, scale) return pa.decimal256(precision, scale) +def _mysql_decimal_value_shape(value) -> tuple[int, int] | None: + """Return integer digits and scale needed to preserve one decimal value.""" + value = value if isinstance(value, PyDecimal) else PyDecimal(str(value)) + if not value.is_finite(): + return None + _, digits, exponent = value.as_tuple() + if exponent >= 0: + return len(digits) + exponent, 0 + value_scale = -exponent + return max(len(digits) - value_scale, 0), value_scale + + +def _mysql_decimal_type_for_values( + display_length: int | None, + scale: int | None, + is_unsigned: bool, + values: list, +) -> pa.DataType: + """Choose an exact Arrow type from decimal metadata and concrete values.""" + precision, scale = _mysql_decimal_precision_scale( + display_length, scale, is_unsigned + ) + shapes = [ + _mysql_decimal_value_shape(value) for value in values if value is not None + ] + if not shapes: + if precision > _ARROW_DECIMAL256_MAX_PRECISION: + return pa.string() + return _arrow_decimal_type(precision, scale) + if any(shape is None for shape in shapes): + return pa.string() + + integer_digits = max(shape[0] for shape in shapes if shape is not None) + value_scale = max(shape[1] for shape in shapes if shape is not None) + if integer_digits + value_scale > _ARROW_DECIMAL256_MAX_PRECISION: + return pa.string() + + if precision > _ARROW_DECIMAL256_MAX_PRECISION: + # Preserve as much metadata scale as possible while reserving enough + # integer capacity for every fetched value. + target_scale = max( + value_scale, + min(scale, _ARROW_DECIMAL256_MAX_PRECISION - integer_digits), + ) + return pa.decimal256(_ARROW_DECIMAL256_MAX_PRECISION, target_scale) + + target_scale = max(scale, value_scale) + target_integer_digits = max(precision - scale, integer_digits) + target_precision = target_integer_digits + target_scale + if target_precision > _ARROW_DECIMAL256_MAX_PRECISION: + return pa.string() + return _arrow_decimal_type(target_precision, target_scale) + + def _mysql_field_arrow_type( type_code: int, flags: int = 0, @@ -367,9 +439,20 @@ def _build_mysql_arrow_table(cursor) -> pa.Table: # a hard-coded ``decimal128(38, 9)``. precision = col[4] if len(col) > 4 else None scale = col[5] if len(col) > 5 else None - arrow_type = _mysql_field_arrow_type( - col[1], flag_list[i] or 0, precision=precision, scale=scale - ) + flags = flag_list[i] or 0 + if col[1] in _mysql_decimal_codes(): + from MySQLdb.constants import FLAG # noqa: PLC0415 + + arrow_type = _mysql_decimal_type_for_values( + precision, + scale, + is_unsigned=bool(flags & FLAG.UNSIGNED), + values=[row[i] for row in rows], + ) + else: + arrow_type = _mysql_field_arrow_type( + col[1], flags, precision=precision, scale=scale + ) fields.append(pa.field(col[0], arrow_type, nullable=True)) schema = pa.schema(fields) diff --git a/core/wren/tests/connectors/test_mysql_connector.py b/core/wren/tests/connectors/test_mysql_connector.py index 22ae1e9995..566a0e54fd 100644 --- a/core/wren/tests/connectors/test_mysql_connector.py +++ b/core/wren/tests/connectors/test_mysql_connector.py @@ -134,6 +134,56 @@ def test_decimal_above_decimal128_uses_decimal256(connector: MySqlConnector) -> assert tbl.column("a").to_pylist() == [value] +def test_decimal_addition_widens_for_concrete_66_digit_value( + connector: MySqlConnector, +) -> None: + left = Decimal("9" * 65) + expected = left + 1 + + tbl = connector.query( + f"SELECT CAST('{left}' AS DECIMAL(65, 0)) + CAST('1' AS DECIMAL(1, 0)) AS total" + ) + + assert tbl.schema.field("total").type == pa.decimal256(66, 0) + assert tbl.column("total").to_pylist() == [expected] + + +@pytest.mark.parametrize("left_digits,right_digits", [(65, 12), (40, 40)]) +def test_decimal_multiplication_above_arrow_limit_uses_exact_string( + connector: MySqlConnector, + left_digits: int, + right_digits: int, +) -> None: + left = Decimal("9" * left_digits) + right = Decimal("9" * right_digits) + expected = str(int(left) * int(right)) + + tbl = connector.query( + f"SELECT CAST('{left}' AS DECIMAL({left_digits}, 0)) " + f"* CAST('{right}' AS DECIMAL({right_digits}, 0)) AS product" + ) + + assert len(expected) in {77, 80} + assert tbl.schema.field("product").type == pa.string() + assert tbl.column("product").to_pylist() == [expected] + + +def test_decimal_wide_sum_metadata_with_fitting_value_stays_numeric( + connector: MySqlConnector, +) -> None: + value = Decimal("9" * 65) + + tbl = connector.query( + "SELECT SUM(value) AS total FROM " + f"(SELECT CAST('{value}' AS DECIMAL(65, 0)) AS value) AS values_" + ) + + # MySQL reports precision 88 for SUM(DECIMAL(65, 0)). The concrete value + # fits Decimal256, so retain a numeric schema at Arrow's precision limit. + assert tbl.schema.field("total").type == pa.decimal256(76, 0) + assert tbl.column("total").to_pylist() == [value] + + def test_float_and_double(connector: MySqlConnector) -> None: _exec(connector, "DROP TABLE IF EXISTS t_real") _exec(connector, "CREATE TABLE t_real (a FLOAT, b DOUBLE)") diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index de18545fcb..7ccfa07c0d 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -14,6 +14,7 @@ from wren.connector.mysql import ( _apply_limit, _arrow_decimal_from_mysql_field, + _build_mysql_arrow_table, _build_mysql_column, _build_mysql_connect_kwargs, _mysql_blob_codes, @@ -41,6 +42,33 @@ def __init__(self, url: str, kwargs: dict[str, str] | None = None) -> None: self.kwargs = kwargs +class _FakeDecimalCursor: + def __init__( + self, + rows: list[tuple], + display_length: int, + scale: int = 0, + ) -> None: + from MySQLdb.constants import FIELD_TYPE # noqa: PLC0415 + + self.description = ( + ( + "value", + FIELD_TYPE.NEWDECIMAL, + None, + None, + display_length, + scale, + True, + ), + ) + self.description_flags = (0,) + self._rows = rows + + def fetchall(self) -> list[tuple]: + return self._rows + + # ── coerce_limit (shared base helper; mysql private removed) ───────────── @@ -239,6 +267,77 @@ def test_decimal_type_scale_not_greater_than_precision() -> None: assert t.scale <= t.precision +# ── value-aware DECIMAL conversion ──────────────────────────────── + + +def test_decimal_column_widens_for_concrete_integer_digits() -> None: + from decimal import Decimal # noqa: PLC0415 + + value = Decimal("1" + "0" * 65) + cursor = _FakeDecimalCursor([(value,)], display_length=66) + + table = _build_mysql_arrow_table(cursor) + + assert table.schema.field("value").type == pa.decimal256(66, 0) + assert table.column("value").to_pylist() == [value] + + +def test_decimal_column_widens_scale_without_losing_integer_capacity() -> None: + from decimal import Decimal # noqa: PLC0415 + + value = Decimal("12345678.1234") + # Signed DECIMAL(10, 2): 10 digits + sign + decimal point. + cursor = _FakeDecimalCursor([(value,)], display_length=12, scale=2) + + table = _build_mysql_arrow_table(cursor) + + assert table.schema.field("value").type == pa.decimal128(12, 4) + assert table.column("value").to_pylist() == [value] + + +def test_decimal_column_above_arrow_limit_uses_exact_strings() -> None: + from decimal import Decimal # noqa: PLC0415 + + values = [Decimal("9" * 77), None, Decimal("9" * 80)] + cursor = _FakeDecimalCursor([(value,) for value in values], display_length=66) + + table = _build_mysql_arrow_table(cursor) + + assert table.schema.field("value").type == pa.string() + assert table.column("value").to_pylist() == [ + "9" * 77, + None, + "9" * 80, + ] + + +def test_decimal_metadata_above_arrow_limit_with_fitting_value_stays_numeric() -> None: + from decimal import Decimal # noqa: PLC0415 + + value = Decimal("9" * 65) + # Signed precision 89 is not representable in Arrow Decimal256. + cursor = _FakeDecimalCursor([(value,)], display_length=90) + + table = _build_mysql_arrow_table(cursor) + + assert table.schema.field("value").type == pa.decimal256(76, 0) + assert table.column("value").to_pylist() == [value] + + +@pytest.mark.parametrize("rows", [[], [(None,)]], ids=["empty", "all_null"]) +def test_decimal_metadata_above_arrow_limit_without_values_uses_string( + rows: list[tuple], +) -> None: + # Signed precision 89 is not representable in Arrow Decimal256, and no + # concrete value proves that a narrower numeric schema would be safe. + cursor = _FakeDecimalCursor(rows, display_length=90) + + table = _build_mysql_arrow_table(cursor) + + assert table.schema.field("value").type == pa.string() + assert table.column("value").to_pylist() == [row[0] for row in rows] + + # ── TIME → duration round-trip ──────────────────────────────────────────── From 12dc2bd6993905e1d1d4130489e006a7cf9d9ef0 Mon Sep 17 00:00:00 2001 From: Ray Tien Date: Mon, 10 Aug 2026 13:47:19 +0800 Subject: [PATCH 04/10] fix(wren): keep decimal unit tests driver-free --- core/wren/tests/unit/test_mysql_helpers.py | 75 +++++++--------------- 1 file changed, 22 insertions(+), 53 deletions(-) diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index 7ccfa07c0d..d57295d840 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -14,11 +14,11 @@ from wren.connector.mysql import ( _apply_limit, _arrow_decimal_from_mysql_field, - _build_mysql_arrow_table, _build_mysql_column, _build_mysql_connect_kwargs, _mysql_blob_codes, _mysql_decimal_codes, + _mysql_decimal_type_for_values, _mysql_field_type_map, _mysql_string_codes, _mysql_unsigned_variant_map, @@ -42,33 +42,6 @@ def __init__(self, url: str, kwargs: dict[str, str] | None = None) -> None: self.kwargs = kwargs -class _FakeDecimalCursor: - def __init__( - self, - rows: list[tuple], - display_length: int, - scale: int = 0, - ) -> None: - from MySQLdb.constants import FIELD_TYPE # noqa: PLC0415 - - self.description = ( - ( - "value", - FIELD_TYPE.NEWDECIMAL, - None, - None, - display_length, - scale, - True, - ), - ) - self.description_flags = (0,) - self._rows = rows - - def fetchall(self) -> list[tuple]: - return self._rows - - # ── coerce_limit (shared base helper; mysql private removed) ───────────── @@ -274,12 +247,11 @@ def test_decimal_column_widens_for_concrete_integer_digits() -> None: from decimal import Decimal # noqa: PLC0415 value = Decimal("1" + "0" * 65) - cursor = _FakeDecimalCursor([(value,)], display_length=66) - - table = _build_mysql_arrow_table(cursor) + arrow_type = _mysql_decimal_type_for_values(66, 0, False, [value]) + column = _build_mysql_column([value], arrow_type) - assert table.schema.field("value").type == pa.decimal256(66, 0) - assert table.column("value").to_pylist() == [value] + assert arrow_type == pa.decimal256(66, 0) + assert column.to_pylist() == [value] def test_decimal_column_widens_scale_without_losing_integer_capacity() -> None: @@ -287,24 +259,22 @@ def test_decimal_column_widens_scale_without_losing_integer_capacity() -> None: value = Decimal("12345678.1234") # Signed DECIMAL(10, 2): 10 digits + sign + decimal point. - cursor = _FakeDecimalCursor([(value,)], display_length=12, scale=2) + arrow_type = _mysql_decimal_type_for_values(12, 2, False, [value]) + column = _build_mysql_column([value], arrow_type) - table = _build_mysql_arrow_table(cursor) - - assert table.schema.field("value").type == pa.decimal128(12, 4) - assert table.column("value").to_pylist() == [value] + assert arrow_type == pa.decimal128(12, 4) + assert column.to_pylist() == [value] def test_decimal_column_above_arrow_limit_uses_exact_strings() -> None: from decimal import Decimal # noqa: PLC0415 values = [Decimal("9" * 77), None, Decimal("9" * 80)] - cursor = _FakeDecimalCursor([(value,) for value in values], display_length=66) - - table = _build_mysql_arrow_table(cursor) + arrow_type = _mysql_decimal_type_for_values(66, 0, False, values) + column = _build_mysql_column(values, arrow_type) - assert table.schema.field("value").type == pa.string() - assert table.column("value").to_pylist() == [ + assert arrow_type == pa.string() + assert column.to_pylist() == [ "9" * 77, None, "9" * 80, @@ -316,12 +286,11 @@ def test_decimal_metadata_above_arrow_limit_with_fitting_value_stays_numeric() - value = Decimal("9" * 65) # Signed precision 89 is not representable in Arrow Decimal256. - cursor = _FakeDecimalCursor([(value,)], display_length=90) + arrow_type = _mysql_decimal_type_for_values(90, 0, False, [value]) + column = _build_mysql_column([value], arrow_type) - table = _build_mysql_arrow_table(cursor) - - assert table.schema.field("value").type == pa.decimal256(76, 0) - assert table.column("value").to_pylist() == [value] + assert arrow_type == pa.decimal256(76, 0) + assert column.to_pylist() == [value] @pytest.mark.parametrize("rows", [[], [(None,)]], ids=["empty", "all_null"]) @@ -330,12 +299,12 @@ def test_decimal_metadata_above_arrow_limit_without_values_uses_string( ) -> None: # Signed precision 89 is not representable in Arrow Decimal256, and no # concrete value proves that a narrower numeric schema would be safe. - cursor = _FakeDecimalCursor(rows, display_length=90) - - table = _build_mysql_arrow_table(cursor) + values = [row[0] for row in rows] + arrow_type = _mysql_decimal_type_for_values(90, 0, False, values) + column = _build_mysql_column(values, arrow_type) - assert table.schema.field("value").type == pa.string() - assert table.column("value").to_pylist() == [row[0] for row in rows] + assert arrow_type == pa.string() + assert column.to_pylist() == values # ── TIME → duration round-trip ──────────────────────────────────────────── From b1df52050da940a8f3f1008dd38e229b476d7c18 Mon Sep 17 00:00:00 2001 From: Ray Tien Date: Mon, 10 Aug 2026 22:04:30 +0800 Subject: [PATCH 05/10] fix(wren): rebalance decimal scale for Arrow --- core/wren/src/wren/connector/mysql.py | 5 ++++- core/wren/tests/unit/test_mysql_helpers.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/core/wren/src/wren/connector/mysql.py b/core/wren/src/wren/connector/mysql.py index 8afde76337..00c6a033f4 100644 --- a/core/wren/src/wren/connector/mysql.py +++ b/core/wren/src/wren/connector/mysql.py @@ -369,7 +369,10 @@ def _mysql_decimal_type_for_values( target_integer_digits = max(precision - scale, integer_digits) target_precision = target_integer_digits + target_scale if target_precision > _ARROW_DECIMAL256_MAX_PRECISION: - return pa.string() + target_scale = _ARROW_DECIMAL256_MAX_PRECISION - target_integer_digits + if target_scale < value_scale: + return pa.string() + return pa.decimal256(_ARROW_DECIMAL256_MAX_PRECISION, target_scale) return _arrow_decimal_type(target_precision, target_scale) diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index d57295d840..f73cf4fc64 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -266,6 +266,19 @@ def test_decimal_column_widens_scale_without_losing_integer_capacity() -> None: assert column.to_pylist() == [value] +def test_decimal_column_rebalances_metadata_scale_for_observed_integer_digits() -> None: + from decimal import Decimal # noqa: PLC0415 + + value = Decimal("12") + # Signed DECIMAL(76, 75) reserves one integer digit. The observed value + # needs two, so retain the maximum scale that still fits Decimal256. + arrow_type = _mysql_decimal_type_for_values(78, 75, False, [value]) + column = _build_mysql_column([value], arrow_type) + + assert arrow_type == pa.decimal256(76, 74) + assert column.to_pylist() == [value] + + def test_decimal_column_above_arrow_limit_uses_exact_strings() -> None: from decimal import Decimal # noqa: PLC0415 From c62e574f7c140677ee5cad469748059afef2c013 Mon Sep 17 00:00:00 2001 From: Ray Tien Date: Mon, 10 Aug 2026 22:10:33 +0800 Subject: [PATCH 06/10] fix(wren): rebalance decimal integer capacity --- core/wren/src/wren/connector/mysql.py | 7 ++++--- core/wren/tests/unit/test_mysql_helpers.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/core/wren/src/wren/connector/mysql.py b/core/wren/src/wren/connector/mysql.py index 00c6a033f4..16b82a7f0f 100644 --- a/core/wren/src/wren/connector/mysql.py +++ b/core/wren/src/wren/connector/mysql.py @@ -369,9 +369,10 @@ def _mysql_decimal_type_for_values( target_integer_digits = max(precision - scale, integer_digits) target_precision = target_integer_digits + target_scale if target_precision > _ARROW_DECIMAL256_MAX_PRECISION: - target_scale = _ARROW_DECIMAL256_MAX_PRECISION - target_integer_digits - if target_scale < value_scale: - return pa.string() + target_scale = max( + value_scale, + min(scale, _ARROW_DECIMAL256_MAX_PRECISION - integer_digits), + ) return pa.decimal256(_ARROW_DECIMAL256_MAX_PRECISION, target_scale) return _arrow_decimal_type(target_precision, target_scale) diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index f73cf4fc64..9b1e032e4d 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -279,6 +279,19 @@ def test_decimal_column_rebalances_metadata_scale_for_observed_integer_digits() assert column.to_pylist() == [value] +def test_decimal_column_rebalances_integer_capacity_for_observed_scale() -> None: + from decimal import Decimal # noqa: PLC0415 + + value = Decimal("0.12") + # Signed DECIMAL(76, 1) reserves 75 integer digits. The observed value + # needs scale 2, so release unused integer capacity to preserve it exactly. + arrow_type = _mysql_decimal_type_for_values(78, 1, False, [value]) + column = _build_mysql_column([value], arrow_type) + + assert arrow_type == pa.decimal256(76, 2) + assert column.to_pylist() == [value] + + def test_decimal_column_above_arrow_limit_uses_exact_strings() -> None: from decimal import Decimal # noqa: PLC0415 From c7bdd43f9e4ed0ddce3bb63cff00e6cf3433e913 Mon Sep 17 00:00:00 2001 From: Ray Tien Date: Mon, 17 Aug 2026 14:24:04 +0800 Subject: [PATCH 07/10] fix(wren): address MySQL decimal review feedback --- core/wren/src/wren/connector/mysql.py | 38 +++++++++++++--------- core/wren/tests/unit/test_mysql_helpers.py | 32 +++++++++++++++--- 2 files changed, 50 insertions(+), 20 deletions(-) diff --git a/core/wren/src/wren/connector/mysql.py b/core/wren/src/wren/connector/mysql.py index 16b82a7f0f..68c5068294 100644 --- a/core/wren/src/wren/connector/mysql.py +++ b/core/wren/src/wren/connector/mysql.py @@ -13,6 +13,7 @@ import json from contextlib import closing from decimal import Decimal as PyDecimal +from decimal import InvalidOperation from functools import cache import pyarrow as pa @@ -321,7 +322,10 @@ def _arrow_decimal_type(precision: int, scale: int) -> pa.DataType: def _mysql_decimal_value_shape(value) -> tuple[int, int] | None: """Return integer digits and scale needed to preserve one decimal value.""" - value = value if isinstance(value, PyDecimal) else PyDecimal(str(value)) + try: + value = value if isinstance(value, PyDecimal) else PyDecimal(str(value)) + except (InvalidOperation, TypeError, ValueError): + return None if not value.is_finite(): return None _, digits, exponent = value.as_tuple() @@ -345,8 +349,8 @@ def _mysql_decimal_type_for_values( _mysql_decimal_value_shape(value) for value in values if value is not None ] if not shapes: - if precision > _ARROW_DECIMAL256_MAX_PRECISION: - return pa.string() + precision = min(precision, _ARROW_DECIMAL256_MAX_PRECISION) + scale = min(scale, precision) return _arrow_decimal_type(precision, scale) if any(shape is None for shape in shapes): return pa.string() @@ -382,6 +386,7 @@ def _mysql_field_arrow_type( flags: int = 0, precision: int | None = None, scale: int | None = None, + values: list | None = None, ) -> pa.DataType: from MySQLdb.constants import FLAG # noqa: PLC0415 @@ -401,6 +406,13 @@ def _mysql_field_arrow_type( # — the previous hard-coded ``decimal128(38, 9)`` would lose digits when # ``D > 9``. if type_code in decimal_codes: + if values is not None: + return _mysql_decimal_type_for_values( + precision, + scale, + is_unsigned=bool(flags & FLAG.UNSIGNED), + values=values, + ) return _arrow_decimal_from_mysql_field( precision, scale, is_unsigned=bool(flags & FLAG.UNSIGNED) ) @@ -444,19 +456,13 @@ def _build_mysql_arrow_table(cursor) -> pa.Table: precision = col[4] if len(col) > 4 else None scale = col[5] if len(col) > 5 else None flags = flag_list[i] or 0 - if col[1] in _mysql_decimal_codes(): - from MySQLdb.constants import FLAG # noqa: PLC0415 - - arrow_type = _mysql_decimal_type_for_values( - precision, - scale, - is_unsigned=bool(flags & FLAG.UNSIGNED), - values=[row[i] for row in rows], - ) - else: - arrow_type = _mysql_field_arrow_type( - col[1], flags, precision=precision, scale=scale - ) + arrow_type = _mysql_field_arrow_type( + col[1], + flags, + precision=precision, + scale=scale, + values=[row[i] for row in rows], + ) fields.append(pa.field(col[0], arrow_type, nullable=True)) schema = pa.schema(fields) diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index 9b1e032e4d..73d7ccb694 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -19,6 +19,7 @@ _mysql_blob_codes, _mysql_decimal_codes, _mysql_decimal_type_for_values, + _mysql_field_arrow_type, _mysql_field_type_map, _mysql_string_codes, _mysql_unsigned_variant_map, @@ -240,6 +241,22 @@ def test_decimal_type_scale_not_greater_than_precision() -> None: assert t.scale <= t.precision +def test_decimal_field_type_uses_concrete_values() -> None: + from decimal import Decimal # noqa: PLC0415 + + type_code = next(iter(_mysql_decimal_codes())) + value = Decimal("1" + "0" * 65) + + arrow_type = _mysql_field_arrow_type( + type_code, + precision=66, + scale=0, + values=[value], + ) + + assert arrow_type == pa.decimal256(66, 0) + + # ── value-aware DECIMAL conversion ──────────────────────────────── @@ -307,6 +324,15 @@ def test_decimal_column_above_arrow_limit_uses_exact_strings() -> None: ] +def test_decimal_unparseable_value_uses_exact_strings() -> None: + values = [b"1.5"] + arrow_type = _mysql_decimal_type_for_values(4, 1, False, values) + column = _build_mysql_column(values, arrow_type) + + assert arrow_type == pa.string() + assert column.to_pylist() == ["1.5"] + + def test_decimal_metadata_above_arrow_limit_with_fitting_value_stays_numeric() -> None: from decimal import Decimal # noqa: PLC0415 @@ -320,16 +346,14 @@ def test_decimal_metadata_above_arrow_limit_with_fitting_value_stays_numeric() - @pytest.mark.parametrize("rows", [[], [(None,)]], ids=["empty", "all_null"]) -def test_decimal_metadata_above_arrow_limit_without_values_uses_string( +def test_decimal_metadata_above_arrow_limit_without_values_stays_numeric( rows: list[tuple], ) -> None: - # Signed precision 89 is not representable in Arrow Decimal256, and no - # concrete value proves that a narrower numeric schema would be safe. values = [row[0] for row in rows] arrow_type = _mysql_decimal_type_for_values(90, 0, False, values) column = _build_mysql_column(values, arrow_type) - assert arrow_type == pa.string() + assert arrow_type == pa.decimal256(76, 0) assert column.to_pylist() == values From 3da3dad581c645f732981ab9771f4556fa73aabd Mon Sep 17 00:00:00 2001 From: Ray Tien Date: Mon, 17 Aug 2026 16:05:27 +0800 Subject: [PATCH 08/10] test(wren): keep MySQL unit tests driver-free --- core/wren/tests/unit/test_mysql_helpers.py | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index 73d7ccb694..0f83f0cfa2 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -19,7 +19,6 @@ _mysql_blob_codes, _mysql_decimal_codes, _mysql_decimal_type_for_values, - _mysql_field_arrow_type, _mysql_field_type_map, _mysql_string_codes, _mysql_unsigned_variant_map, @@ -241,22 +240,6 @@ def test_decimal_type_scale_not_greater_than_precision() -> None: assert t.scale <= t.precision -def test_decimal_field_type_uses_concrete_values() -> None: - from decimal import Decimal # noqa: PLC0415 - - type_code = next(iter(_mysql_decimal_codes())) - value = Decimal("1" + "0" * 65) - - arrow_type = _mysql_field_arrow_type( - type_code, - precision=66, - scale=0, - values=[value], - ) - - assert arrow_type == pa.decimal256(66, 0) - - # ── value-aware DECIMAL conversion ──────────────────────────────── From df252833a225eacf95dc80d560e6d529188c5fe2 Mon Sep 17 00:00:00 2001 From: Ray Tien Date: Tue, 25 Aug 2026 12:12:18 +0800 Subject: [PATCH 09/10] fix(wren): refine MySQL decimal fallback --- core/wren/src/wren/connector/mysql.py | 92 +++++++++---- .../tests/connectors/test_mysql_connector.py | 8 +- core/wren/tests/unit/test_mysql_helpers.py | 128 ++++++++++++++++++ 3 files changed, 197 insertions(+), 31 deletions(-) diff --git a/core/wren/src/wren/connector/mysql.py b/core/wren/src/wren/connector/mysql.py index 68c5068294..375c195a72 100644 --- a/core/wren/src/wren/connector/mysql.py +++ b/core/wren/src/wren/connector/mysql.py @@ -341,7 +341,12 @@ def _mysql_decimal_type_for_values( is_unsigned: bool, values: list, ) -> pa.DataType: - """Choose an exact Arrow type from decimal metadata and concrete values.""" + """Choose an exact Arrow type for one result set's concrete values. + + The returned type may differ between result sets when expression metadata + understates the fetched values. Values beyond Arrow Decimal256, invalid + values, and non-finite values use exact strings instead of losing data. + """ precision, scale = _mysql_decimal_precision_scale( display_length, scale, is_unsigned ) @@ -360,24 +365,16 @@ def _mysql_decimal_type_for_values( if integer_digits + value_scale > _ARROW_DECIMAL256_MAX_PRECISION: return pa.string() - if precision > _ARROW_DECIMAL256_MAX_PRECISION: - # Preserve as much metadata scale as possible while reserving enough - # integer capacity for every fetched value. - target_scale = max( - value_scale, - min(scale, _ARROW_DECIMAL256_MAX_PRECISION - integer_digits), - ) - return pa.decimal256(_ARROW_DECIMAL256_MAX_PRECISION, target_scale) - - target_scale = max(scale, value_scale) - target_integer_digits = max(precision - scale, integer_digits) - target_precision = target_integer_digits + target_scale - if target_precision > _ARROW_DECIMAL256_MAX_PRECISION: - target_scale = max( - value_scale, - min(scale, _ARROW_DECIMAL256_MAX_PRECISION - integer_digits), - ) - return pa.decimal256(_ARROW_DECIMAL256_MAX_PRECISION, target_scale) + # Preserve as much metadata scale as possible, but release unused declared + # integer capacity before escalating from Decimal128 to Decimal256. + target_scale = max( + value_scale, + min(scale, _ARROW_DECIMAL256_MAX_PRECISION - integer_digits), + ) + target_precision = max( + min(precision, _ARROW_DECIMAL256_MAX_PRECISION), + integer_digits + target_scale, + ) return _arrow_decimal_type(target_precision, target_scale) @@ -430,8 +427,47 @@ def _mysql_field_arrow_type( return base +def _build_mysql_result_column( + name: str, + type_code: int, + flags: int, + precision: int | None, + scale: int | None, + values: list, + arrow_type: pa.DataType, +) -> tuple[pa.Array, pa.DataType]: + """Build one result column, scanning decimal values only after failure.""" + if not values: + return pa.array([], type=arrow_type), arrow_type + if not pa.types.is_decimal(arrow_type): + return _build_mysql_column(values, arrow_type), arrow_type + + try: + return _build_mysql_column(values, arrow_type), arrow_type + except (InvalidOperation, pa.ArrowInvalid): + arrow_type = _mysql_field_arrow_type( + type_code, + flags, + precision=precision, + scale=scale, + values=values, + ) + if pa.types.is_string(arrow_type): + logger.warning( + "MySQL decimal column {!r} cannot be safely converted to an " + "Arrow decimal; returning exact strings for this result set", + name, + ) + return _build_mysql_column(values, arrow_type), arrow_type + + def _build_mysql_arrow_table(cursor) -> pa.Table: - """Convert a MySQLdb cursor result to a PyArrow table.""" + """Convert a MySQLdb cursor result to a PyArrow table. + + Decimal columns optimistically use their metadata type. If concrete values + do not fit, their type is derived from that result set and may fall back to + exact strings; that fallback emits a warning naming the affected column. + """ if cursor.description is None: return pa.table({}) @@ -446,7 +482,9 @@ def _build_mysql_arrow_table(cursor) -> pa.Table: flag_list = (flag_list + [0] * len(cursor.description))[: len(cursor.description)] rows = cursor.fetchall() + columns = [[row[i] for row in rows] for i in range(len(cursor.description))] fields = [] + arrays = [] for i, col in enumerate(cursor.description): # PEP 249 ``description`` tuple: # (name, type_code, display_size, internal_size, precision, scale, null_ok) @@ -461,18 +499,14 @@ def _build_mysql_arrow_table(cursor) -> pa.Table: flags, precision=precision, scale=scale, - values=[row[i] for row in rows], + ) + array, arrow_type = _build_mysql_result_column( + col[0], col[1], flags, precision, scale, columns[i], arrow_type ) fields.append(pa.field(col[0], arrow_type, nullable=True)) + arrays.append(array) schema = pa.schema(fields) - if not rows: - arrays = [pa.array([], type=field.type) for field in schema] - else: - arrays = [ - _build_mysql_column([row[i] for row in rows], schema.field(i).type) - for i in range(len(fields)) - ] # ``pa.table(dict(...), schema=...)`` silently drops a column when two # fields share the same name (the dict collapses the duplicate). Use # ``pa.Table.from_arrays`` so a query like @@ -509,7 +543,7 @@ def _build_mysql_column(values: list, arrow_type: pa.DataType) -> pa.Array: else (v if isinstance(v, PyDecimal) else PyDecimal(str(v))) for v in values ] - return pa.array(processed, type=arrow_type, from_pandas=True) + return pa.array(processed, type=arrow_type, from_pandas=False) if pa.types.is_timestamp(arrow_type): return pa.array(values, type=arrow_type, from_pandas=True) diff --git a/core/wren/tests/connectors/test_mysql_connector.py b/core/wren/tests/connectors/test_mysql_connector.py index 566a0e54fd..b17aed0275 100644 --- a/core/wren/tests/connectors/test_mysql_connector.py +++ b/core/wren/tests/connectors/test_mysql_connector.py @@ -148,11 +148,15 @@ def test_decimal_addition_widens_for_concrete_66_digit_value( assert tbl.column("total").to_pylist() == [expected] -@pytest.mark.parametrize("left_digits,right_digits", [(65, 12), (40, 40)]) +@pytest.mark.parametrize( + ("left_digits", "right_digits", "expected_digits"), + [(65, 12, 77), (40, 40, 80)], +) def test_decimal_multiplication_above_arrow_limit_uses_exact_string( connector: MySqlConnector, left_digits: int, right_digits: int, + expected_digits: int, ) -> None: left = Decimal("9" * left_digits) right = Decimal("9" * right_digits) @@ -163,7 +167,7 @@ def test_decimal_multiplication_above_arrow_limit_uses_exact_string( f"* CAST('{right}' AS DECIMAL({right_digits}, 0)) AS product" ) - assert len(expected) in {77, 80} + assert len(expected) == expected_digits assert tbl.schema.field("product").type == pa.string() assert tbl.column("product").to_pylist() == [expected] diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index 0f83f0cfa2..2e52860083 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -10,6 +10,7 @@ import pyarrow as pa import pytest +import wren.connector.mysql as mysql_connector from wren.connector.base import coerce_limit from wren.connector.mysql import ( _apply_limit, @@ -42,6 +43,15 @@ def __init__(self, url: str, kwargs: dict[str, str] | None = None) -> None: self.kwargs = kwargs +class _FakeCursor: + def __init__(self, description: list[tuple], rows: list[tuple]) -> None: + self.description = description + self._rows = rows + + def fetchall(self) -> list[tuple]: + return self._rows + + # ── coerce_limit (shared base helper; mysql private removed) ───────────── @@ -266,6 +276,19 @@ def test_decimal_column_widens_scale_without_losing_integer_capacity() -> None: assert column.to_pylist() == [value] +def test_decimal_column_reuses_decimal128_when_observed_value_fits() -> None: + from decimal import Decimal # noqa: PLC0415 + + value = Decimal("1.234567") + # Signed DECIMAL(38, 4) has spare integer capacity that can be reassigned + # to the observed scale without escalating to Decimal256. + arrow_type = _mysql_decimal_type_for_values(40, 4, False, [value]) + column = _build_mysql_column([value], arrow_type) + + assert arrow_type == pa.decimal128(38, 6) + assert column.to_pylist() == [value] + + def test_decimal_column_rebalances_metadata_scale_for_observed_integer_digits() -> None: from decimal import Decimal # noqa: PLC0415 @@ -307,6 +330,26 @@ def test_decimal_column_above_arrow_limit_uses_exact_strings() -> None: ] +@pytest.mark.parametrize( + ("digits", "expected_type"), + [(76, pa.decimal256(76, 0)), (77, pa.string())], +) +def test_negative_decimal_values_follow_arrow_precision_limit( + digits: int, + expected_type: pa.DataType, +) -> None: + from decimal import Decimal # noqa: PLC0415 + + value = Decimal("-" + "9" * digits) + arrow_type = _mysql_decimal_type_for_values(66, 0, False, [value]) + column = _build_mysql_column([value], arrow_type) + + assert arrow_type == expected_type + assert column.to_pylist() == ( + [str(value)] if pa.types.is_string(arrow_type) else [value] + ) + + def test_decimal_unparseable_value_uses_exact_strings() -> None: values = [b"1.5"] arrow_type = _mysql_decimal_type_for_values(4, 1, False, values) @@ -340,6 +383,91 @@ def test_decimal_metadata_above_arrow_limit_without_values_stays_numeric( assert column.to_pylist() == values +def test_decimal_table_uses_metadata_type_before_scanning_values( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from decimal import Decimal # noqa: PLC0415 + + values_seen = [] + + def get_arrow_type( + type_code: int, + flags: int = 0, + precision: int | None = None, + scale: int | None = None, + values: list | None = None, + ) -> pa.DataType: + values_seen.append(values) + return pa.decimal128(3, 2) + + monkeypatch.setattr(mysql_connector, "_mysql_field_arrow_type", get_arrow_type) + cursor = _FakeCursor( + [("amount", 0, None, None, 5, 2, True)], + [(Decimal("1.23"),)], + ) + + table = mysql_connector._build_mysql_arrow_table(cursor) + + assert values_seen == [None] + assert table.schema.field("amount").type == pa.decimal128(3, 2) + assert table.column("amount").to_pylist() == [Decimal("1.23")] + + +@pytest.mark.parametrize( + ("value", "display_length", "scale", "expected"), + [ + pytest.param(None, 66, 0, "9" * 77, id="above-arrow-limit"), + pytest.param(b"1.5", 4, 1, "1.5", id="unparseable"), + pytest.param("NaN", 4, 1, "NaN", id="non-finite"), + ], +) +def test_decimal_table_warns_when_falling_back_to_exact_strings( + monkeypatch: pytest.MonkeyPatch, + value, + display_length: int, + scale: int, + expected: str, +) -> None: + from decimal import Decimal # noqa: PLC0415 + + value = Decimal(expected) if value is None else value + warnings = [] + + class _WarningRecorder: + def warning(self, message: str, *args) -> None: + warnings.append(message.format(*args)) + + def get_arrow_type( + type_code: int, + flags: int = 0, + precision: int | None = None, + scale: int | None = None, + values: list | None = None, + ) -> pa.DataType: + if values is not None: + return _mysql_decimal_type_for_values( + precision, + scale, + is_unsigned=False, + values=values, + ) + return mysql_connector._arrow_decimal_from_mysql_field(precision, scale) + + monkeypatch.setattr(mysql_connector, "logger", _WarningRecorder()) + monkeypatch.setattr(mysql_connector, "_mysql_field_arrow_type", get_arrow_type) + cursor = _FakeCursor( + [("product", 0, None, None, display_length, scale, True)], + [(value,)], + ) + + table = mysql_connector._build_mysql_arrow_table(cursor) + + assert table.schema.field("product").type == pa.string() + assert table.column("product").to_pylist() == [expected] + assert len(warnings) == 1 + assert "product" in warnings[0] + + # ── TIME → duration round-trip ──────────────────────────────────────────── From 60fe1777d7498a7aff6434f0b107409692061411 Mon Sep 17 00:00:00 2001 From: Ray Tien Date: Tue, 25 Aug 2026 12:18:35 +0800 Subject: [PATCH 10/10] fix(wren): retain Decimal128 under scale pressure --- core/wren/src/wren/connector/mysql.py | 15 +++++++++++++-- core/wren/tests/unit/test_mysql_helpers.py | 13 +++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/core/wren/src/wren/connector/mysql.py b/core/wren/src/wren/connector/mysql.py index 375c195a72..e595b1a187 100644 --- a/core/wren/src/wren/connector/mysql.py +++ b/core/wren/src/wren/connector/mysql.py @@ -365,8 +365,19 @@ def _mysql_decimal_type_for_values( if integer_digits + value_scale > _ARROW_DECIMAL256_MAX_PRECISION: return pa.string() - # Preserve as much metadata scale as possible, but release unused declared - # integer capacity before escalating from Decimal128 to Decimal256. + # Release unused metadata scale before crossing the Decimal128 boundary. + if ( + precision <= _ARROW_DECIMAL128_MAX_PRECISION + and integer_digits + value_scale <= _ARROW_DECIMAL128_MAX_PRECISION + ): + target_scale = max( + value_scale, + min(scale, _ARROW_DECIMAL128_MAX_PRECISION - integer_digits), + ) + target_precision = max(precision, integer_digits + target_scale) + return pa.decimal128(target_precision, target_scale) + + # Apply the same rebalancing at the Decimal256 boundary. target_scale = max( value_scale, min(scale, _ARROW_DECIMAL256_MAX_PRECISION - integer_digits), diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index 2e52860083..7d30941f21 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -289,6 +289,19 @@ def test_decimal_column_reuses_decimal128_when_observed_value_fits() -> None: assert column.to_pylist() == [value] +def test_decimal_column_rebalances_scale_to_stay_decimal128() -> None: + from decimal import Decimal # noqa: PLC0415 + + value = Decimal("12") + # Signed DECIMAL(38, 37) reserves one integer digit. The observed value + # needs two, so release one unused scale digit without using Decimal256. + arrow_type = _mysql_decimal_type_for_values(40, 37, False, [value]) + column = _build_mysql_column([value], arrow_type) + + assert arrow_type == pa.decimal128(38, 36) + assert column.to_pylist() == [value] + + def test_decimal_column_rebalances_metadata_scale_for_observed_integer_digits() -> None: from decimal import Decimal # noqa: PLC0415