| # Licensed to the Apache Software Foundation (ASF) under one | |
| # or more contributor license agreements. See the NOTICE file | |
| # distributed with this work for additional information | |
| # regarding copyright ownership. The ASF licenses this file | |
| # to you under the Apache License, Version 2.0 (the | |
| # "License"); you may not use this file except in compliance | |
| # with the License. You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, | |
| # software distributed under the License is distributed on an | |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
| # KIND, either express or implied. See the License for the | |
| # specific language governing permissions and limitations | |
| # under the License. | |
| """Java TypeUtils conversions for private overwrite and truncate specs.""" | |
| import re | |
| import struct | |
| from calendar import monthrange | |
| from datetime import date, datetime, timedelta | |
| _UNREPRESENTABLE_PARTITION_VALUE = object() | |
| _JAVA_WHITESPACE = ''.join(chr(value) for value in range(33)) | |
| def _partition_spec_value(value, data_type): | |
| """TypeUtils.castFromString for the scalar types accepted as partitions. | |
| Partition specs are not SQL casts: integer text must be integral and | |
| DECIMAL uses HALF_UP with precision overflow represented as NULL. | |
| Already typed Python values retain the existing public API contract. | |
| """ | |
| if not isinstance(value, str): | |
| return value | |
| import pyarrow as pa | |
| from pypaimon.execution.python.casting.decimal_cast import cast_decimal | |
| from pypaimon.utils.numeric_utils import java_floating_string | |
| from pypaimon.utils.temporal_utils import _local_to_utc_seconds, _signed | |
| from pypaimon.schema.data_types import PyarrowFieldParser | |
| from pypaimon.table.row.generic_row import _parse_type_precision_scale | |
| from unicodedata import decimal | |
| target = PyarrowFieldParser.from_paimon_type(data_type) | |
| if pa.types.is_integer(target) or pa.types.is_decimal(target): | |
| # Integer.valueOf and BigDecimal scan UTF-16 chars. BMP decimal | |
| # digits are accepted; a supplementary-plane digit is two surrogates. | |
| value = ''.join(str(decimal(char)) if ord(char) <= 0xffff and char.isdecimal() else char for char in value) | |
| if pa.types.is_string(target) or pa.types.is_binary(target): | |
| result = value.encode('utf-8') if pa.types.is_binary(target) else value | |
| match = re.fullmatch(r'(CHAR|VARCHAR|VARBINARY)\((\d+)\)', data_type.type.upper()) | |
| if match is not None and len(result) > int(match.group(2)): | |
| raise ValueError('Partition value exceeds the declared type length') | |
| return result | |
| if pa.types.is_integer(target): | |
| if re.fullmatch(r'[+-]?[0-9]+', value) is None: | |
| raise ValueError('Partition integer text must be integral') | |
| result = int(value) | |
| if not -(1 << (target.bit_width - 1)) <= result < (1 << (target.bit_width - 1)): | |
| raise ValueError('Partition integer value is out of range') | |
| return result | |
| if pa.types.is_boolean(target): | |
| lowered = value.lower() | |
| if lowered in ('t', 'true', 'y', 'yes', '1'): | |
| return True | |
| if lowered in ('f', 'false', 'n', 'no', '0'): | |
| return False | |
| raise ValueError('Invalid BOOLEAN partition value') | |
| if pa.types.is_floating(target): | |
| text = value.strip(_JAVA_WHITESPACE) | |
| if re.fullmatch(r'[+-]?(?:NaN|Infinity|(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+)' | |
| r'(?:[eE][+-]?[0-9]+)?[fFdD]?|0[xX](?:[0-9a-fA-F]+(?:\.[0-9a-fA-F]*)?' | |
| r'|\.[0-9a-fA-F]+)' | |
| r'[pP][+-]?[0-9]+[fFdD]?)', text) is None: | |
| raise ValueError('Invalid floating partition value') | |
| text = text[:-1] if text[-1:] in ('f', 'F', 'd', 'D') else text | |
| try: | |
| result = float.fromhex(text) if 'p' in text.lower() else float(text) | |
| except OverflowError: | |
| # Double.parseDouble saturates overflowing hexadecimal text to | |
| # infinity, while Python's float.fromhex raises instead. | |
| result = float('-inf') if text.startswith('-') else float('inf') | |
| if target == pa.float32(): | |
| try: | |
| single = struct.unpack('>f', struct.pack('>f', result))[0] | |
| except OverflowError: | |
| single = float('-inf') if result < 0 else float('inf') | |
| if result != single and java_floating_string(result) != java_floating_string(single, True): | |
| raise ValueError('Partition value cannot be cast to FLOAT due to precision loss') | |
| return single | |
| return result | |
| if pa.types.is_decimal(target): | |
| return cast_decimal(pa.array([value]), target)[0].as_py() | |
| if pa.types.is_date32(target): | |
| if re.fullmatch(r'-?[0-9]+', value): | |
| days = int(value) | |
| if not -(1 << 31) <= days < (1 << 31): | |
| raise ValueError('DATE partition value is out of range') | |
| try: | |
| return date(1970, 1, 1) + timedelta(days=days) | |
| except OverflowError: | |
| # Pure Python cannot write partition dates outside datetime's | |
| # range. A valid Java spec outside that range cannot match one. | |
| return _UNREPRESENTABLE_PARTITION_VALUE | |
| text = value.split(' ', 1)[0] if value.find(' ') > 0 else value | |
| parts = text.split('-') | |
| parts = [part.strip(_JAVA_WHITESPACE) for part in parts] | |
| if not 1 <= len(parts) <= 3 or any(re.fullmatch(r'[0-9]+', part) is None for part in parts): | |
| raise ValueError('Invalid DATE partition value') | |
| return date(*([int(part) for part in parts] + [1] * (3 - len(parts)))) | |
| if pa.types.is_time(target): | |
| if re.fullmatch(r'-?[0-9]+', value): | |
| millis = int(value) | |
| else: | |
| # Java parseTime permits partial clocks and adds a signed offset. | |
| match = re.fullmatch(r'[\x00-\x20]*([0-9]+)[\x00-\x20]*(?::[\x00-\x20]*([0-9]+)[\x00-\x20]*' | |
| r'(?::[\x00-\x20]*([0-9]+)(?:\.([^+\-]*))?[\x00-\x20]*)?)?' | |
| r'(?:([+-])[\x00-\x20]*([0-9]+)[\x00-\x20]*' | |
| r'(?::[\x00-\x20]*([0-9]+)[\x00-\x20]*)?)?', value) | |
| if match is None: | |
| raise ValueError('Invalid TIME partition value') | |
| hours, minutes, seconds, fraction, sign, zone_hour, zone_minute = match.groups() | |
| if zone_minute is not None and seconds is None: | |
| raise ValueError('Invalid TIME partition value') | |
| fraction = (fraction or '').strip(_JAVA_WHITESPACE) | |
| # DateTimeUtils.parseFraction scans UTF-16 chars. A supplementary | |
| # character occupies two positions, including for rounding. | |
| encoded = fraction[:4].encode('utf-16-le', 'surrogatepass') | |
| units = [encoded[index] | encoded[index + 1] << 8 for index in range(0, len(encoded), 2)] | |
| millis = sum((unit - ord('0') if ord('0') <= unit <= ord('9') else 0) * factor | |
| for unit, factor in zip(units, (100, 10, 1))) | |
| if len(units) > 3 and units[3] >= ord('5'): | |
| millis += 1 | |
| components = [int(item or 0) for item in (hours, minutes, seconds, zone_hour, zone_minute)] | |
| if any(item >= 1 << 31 for item in components): | |
| raise ValueError('TIME partition component is out of range') | |
| millis += ((int(hours) * 60 + int(minutes or 0)) * 60 + int(seconds or 0)) * 1000 | |
| if sign is not None: | |
| millis += (1 if sign == '+' else -1) * (int(zone_hour) * 60 + int(zone_minute or 0)) * 60000 | |
| millis = _signed(millis, 32) | |
| if not -(1 << 31) <= millis < (1 << 31): | |
| raise ValueError('TIME partition value is out of range') | |
| if not 0 <= millis < 86400000: | |
| return _UNREPRESENTABLE_PARTITION_VALUE | |
| return (datetime(1970, 1, 1) + timedelta(milliseconds=millis)).time() | |
| if pa.types.is_timestamp(target): | |
| precision = _parse_type_precision_scale(data_type)[0] | |
| if target.tz is None and re.fullmatch(r'-?[0-9]+', value): | |
| if precision not in (0, 3, 6, 9): | |
| raise ValueError('Numeric TIMESTAMP partition values require precision 0, 3, 6 or 9') | |
| ticks = int(value) | |
| if not -(1 << 63) <= ticks < (1 << 63): | |
| raise ValueError('TIMESTAMP partition value is out of range') | |
| # Java first converts numeric seconds to signed-long millis. | |
| # The multiplication wraps before it constructs the timestamp. | |
| storage_type = target | |
| if precision == 0: | |
| ticks = _signed(ticks * 1000, 64) | |
| # A wrapped product need not retain whole-second alignment. | |
| # Timestamp.fromEpochMillis keeps that millisecond remainder. | |
| storage_type = pa.timestamp('ms') | |
| try: | |
| return pa.scalar(ticks, type=storage_type).as_py() | |
| except OverflowError: | |
| return _UNREPRESENTABLE_PARTITION_VALUE | |
| match = re.fullmatch(r'(\+?[0-9]+)-([0-9]*)-([0-9]*)' | |
| r'(?: ([0-9]*):([0-9]*):([0-9]*)(?:\.([0-9]{0,9}))?)?' | |
| r'(?:T([0-9]*):([0-9]*)(?::([0-9]*)(?:\.([0-9]{0,9}))?)?)?', value) | |
| if match is None: | |
| raise ValueError('Invalid TIMESTAMP partition value') | |
| year_text, month_text, day_text, hour, minute, second, fraction, t_hour, t_minute, t_second, t_fraction = ( | |
| match.groups()) | |
| year_digits = year_text.lstrip('+') | |
| if ((not year_text.startswith('+') and len(year_digits) != 4) | |
| or year_text.startswith('+') and not 5 <= len(year_digits) <= 19 | |
| or not 1 <= int(year_digits) <= 999999999 | |
| or not month_text and day_text): | |
| raise ValueError('Invalid TIMESTAMP partition value') | |
| def component(text): | |
| if not text: | |
| return None | |
| # DateTimeFormatter's adjacent optional [MM][M] patterns can | |
| # repeat a component, provided both parsed values agree. | |
| if len(text) > 21 or len(text) > 2 and int(text[:2]) != int(text[2:]): | |
| raise ValueError('Conflicting TIMESTAMP partition components') | |
| return int(text[:2]) | |
| clock = [component(text) for text in (hour, minute, second)] + [ | |
| None if fraction is None else int(fraction.ljust(9, '0'))] | |
| other = [component(text) for text in (t_hour, t_minute, t_second)] + [ | |
| None if t_fraction is None else int(t_fraction.ljust(9, '0'))] | |
| for index, (left, right) in enumerate(zip(clock, other)): | |
| if left is not None and right is not None and left != right: | |
| raise ValueError('Conflicting TIMESTAMP partition components') | |
| if right is not None: | |
| clock[index] = right | |
| if clock[0] == 24 and (clock[1] is None and clock[2] is not None | |
| or clock[2] is None and clock[3] is not None): | |
| raise ValueError('Invalid TIMESTAMP partition value') | |
| hour, minute, second, nanos = [item or 0 for item in clock] | |
| year, month, day = int(year_digits), component(month_text), component(day_text) | |
| month = 1 if month is None else month | |
| day = 1 if day is None else day | |
| if (not 1 <= month <= 12 or not 1 <= day <= 31 or not 0 <= hour <= 24 | |
| or hour == 24 and (minute or second or nanos) or not 0 <= minute <= 59 or not 0 <= second <= 59): | |
| raise ValueError('Invalid TIMESTAMP partition value') | |
| if year > 9999: | |
| return _UNREPRESENTABLE_PARTITION_VALUE | |
| # DateTimeFormatter's default SMART resolver clamps day-of-month and | |
| # accepts exactly midnight as 24:00:00 on the following date. | |
| whole = datetime(year, month, min(day, monthrange(year, month)[1]), hour % 24, minute, second) | |
| if hour == 24 and month_text and day_text: | |
| whole += timedelta(days=1) | |
| nanos = nanos // 10 ** (9 - precision) * 10 ** (9 - precision) | |
| epoch = whole - datetime(1970, 1, 1) | |
| seconds = epoch.days * 86400 + epoch.seconds | |
| if target.tz is not None: | |
| seconds = _local_to_utc_seconds(seconds) | |
| ticks_per_second = {'s': 1, 'ms': 1000, 'us': 1000000, 'ns': 1000000000}[target.unit] | |
| return pa.scalar(seconds * ticks_per_second + nanos // (1000000000 // ticks_per_second), type=target).as_py() | |
| raise ValueError('Unsupported partition spec type %s' % data_type) |