From 77ccfd249571765e0251f07a69361dbcefe6fc2c Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 03:02:02 -0700 Subject: [PATCH 01/23] [python] Avoid full decoding for VARIANT path updates --- docs/docs/pypaimon/python-api.mdx | 21 + paimon-python/pypaimon/data/__init__.py | 17 +- paimon-python/pypaimon/data/variant_path.py | 673 ++++++++++++++++++ .../pypaimon/tests/variant_path_test.py | 237 ++++++ 4 files changed, 947 insertions(+), 1 deletion(-) create mode 100644 paimon-python/pypaimon/data/variant_path.py create mode 100644 paimon-python/pypaimon/tests/variant_path_test.py diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index 1921a9764974..152e9939a19b 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -1126,6 +1126,27 @@ Supported Paimon type strings for shredded sub-fields: `BOOLEAN`, `INT`, `BIGINT +### VARIANT Path Updates + +Use path helpers when only a few scalar fields change. They avoid decoding +unrelated fields into Python objects: + +```python +import operator + +from pypaimon.data import variant_transform + +updated_payload = variant_transform(payload, { + '$.velocity.y': operator.neg, + '$.velocity.z': operator.neg, +}) +``` + +`variant_get` extracts one path as an Arrow array. `variant_get_many` and +`variant_set_many` process several paths in one pass. `variant_set` replaces +one path. Set operations require existing paths and scalar replacement values; +they keep the original VARIANT metadata. + **`GenericVariant` API:** diff --git a/paimon-python/pypaimon/data/__init__.py b/paimon-python/pypaimon/data/__init__.py index 97f36f3d5278..d50cf1f77230 100644 --- a/paimon-python/pypaimon/data/__init__.py +++ b/paimon-python/pypaimon/data/__init__.py @@ -17,5 +17,20 @@ from pypaimon.data.timestamp import Timestamp from pypaimon.data.decimal import Decimal +from pypaimon.data.variant_path import ( + variant_get, + variant_get_many, + variant_set, + variant_set_many, + variant_transform, +) -__all__ = ['Timestamp', 'Decimal'] +__all__ = [ + 'Timestamp', + 'Decimal', + 'variant_get', + 'variant_get_many', + 'variant_set', + 'variant_set_many', + 'variant_transform', +] diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py new file mode 100644 index 000000000000..6616f4bd2412 --- /dev/null +++ b/paimon-python/pypaimon/data/variant_path.py @@ -0,0 +1,673 @@ +# 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. + +"""Arrow helpers for accessing VARIANT paths without decoding whole values.""" + +import functools +import re +import struct +from typing import Dict, List, Mapping, Optional, Sequence, Tuple + +import pyarrow as pa + +from pypaimon.data._variant_binary import ( + _ARRAY, + _OBJECT, + _U8_MAX, + _U32_SIZE, + _array_header, + _get_int_size, + _object_header, + _read_unsigned, +) +from pypaimon.data.generic_variant import ( + GenericVariant, + _Type, + _value_size, + _variant_get_type, +) + + +_INDEX_PATTERN = re.compile(r"\[(\d+)]") +_KEY_PATTERN = re.compile(r"\.([^.\[]+)|\['([^']+)']|\[\"([^\"]+)\"]") + +_Path = Tuple[Tuple[str, object], ...] +_ObjectLayout = Tuple[int, int, int, int, int, int] +_ArrayLayout = Tuple[int, int, int, int] +_DOUBLE_HEADER = 7 << 2 + + +@functools.lru_cache(maxsize=256) +def _parse_path(path: str) -> _Path: + if not isinstance(path, str) or not path.startswith('$'): + raise ValueError(f"Invalid VARIANT path: {path}") + + pos = 1 + segments = [] + while pos < len(path): + match = _INDEX_PATTERN.match(path, pos) + if match is not None: + segments.append(('index', int(match.group(1)))) + pos = match.end() + continue + + match = _KEY_PATTERN.match(path, pos) + if match is not None: + key = next(value for value in match.groups() if value is not None) + segments.append(('key', key)) + pos = match.end() + continue + raise ValueError(f"Invalid VARIANT path: {path}") + return tuple(segments) + + +@functools.lru_cache(maxsize=256) +def _metadata_key_ids(metadata: bytes) -> Dict[str, int]: + if not metadata: + raise ValueError("MALFORMED_VARIANT: empty metadata") + offset_size = ((metadata[0] >> 6) & 0x3) + 1 + size = _read_unsigned(metadata, 1, offset_size) + offset_start = 1 + offset_size + string_start = offset_start + (size + 1) * offset_size + result = {} + for key_id in range(size): + start = _read_unsigned( + metadata, offset_start + key_id * offset_size, offset_size) + end = _read_unsigned( + metadata, offset_start + (key_id + 1) * offset_size, offset_size) + key = metadata[string_start + start:string_start + end].decode('utf-8') + result[key] = key_id + return result + + +def _object_layout(value: bytes, pos: int) -> _ObjectLayout: + header = value[pos] + if (header & 0x3) != _OBJECT: + raise ValueError("VARIANT path expects an object") + type_info = (header >> 2) & 0x3F + size_bytes = _U32_SIZE if ((type_info >> 4) & 0x1) else 1 + size = _read_unsigned(value, pos + 1, size_bytes) + id_size = ((type_info >> 2) & 0x3) + 1 + offset_size = (type_info & 0x3) + 1 + id_start = pos + 1 + size_bytes + offset_start = id_start + size * id_size + data_start = offset_start + (size + 1) * offset_size + return size, id_size, offset_size, id_start, offset_start, data_start + + +def _array_layout(value: bytes, pos: int) -> _ArrayLayout: + header = value[pos] + if (header & 0x3) != _ARRAY: + raise ValueError("VARIANT path expects an array") + type_info = (header >> 2) & 0x3F + size_bytes = _U32_SIZE if ((type_info >> 2) & 0x1) else 1 + size = _read_unsigned(value, pos + 1, size_bytes) + offset_size = (type_info & 0x3) + 1 + offset_start = pos + 1 + size_bytes + data_start = offset_start + (size + 1) * offset_size + return size, offset_size, offset_start, data_start + + +@functools.lru_cache(maxsize=2048) +def _field_slot(id_table: bytes, id_size: int, key_id: int) -> Optional[int]: + for slot in range(len(id_table) // id_size): + if _read_unsigned(id_table, slot * id_size, id_size) == key_id: + return slot + return None + + +def _object_field_position( + value: bytes, + pos: int, + key_id: int, +) -> Optional[int]: + size, id_size, offset_size, id_start, offset_start, data_start = ( + _object_layout(value, pos)) + id_table = value[id_start:id_start + size * id_size] + slot = _field_slot(id_table, id_size, key_id) + if slot is None: + return None + offset = _read_unsigned( + value, offset_start + slot * offset_size, offset_size) + return data_start + offset + + +def _array_element_position( + value: bytes, + pos: int, + index: int, +) -> Optional[int]: + size, offset_size, offset_start, data_start = _array_layout(value, pos) + if index >= size: + return None + offset = _read_unsigned( + value, offset_start + index * offset_size, offset_size) + return data_start + offset + + +@functools.lru_cache(maxsize=256) +def _compile_paths(paths: Tuple[_Path, ...]): + nodes = [(None, None, None)] + node_by_prefix = {(): 0} + results = [] + for path in paths: + for length in range(1, len(path) + 1): + prefix = path[:length] + if prefix not in node_by_prefix: + parent = node_by_prefix[prefix[:-1]] + kind, segment = prefix[-1] + node_by_prefix[prefix] = len(nodes) + nodes.append((parent, kind, segment)) + results.append(node_by_prefix[path]) + return tuple(nodes), tuple(results) + + +class _PositionPlan: + + def __init__(self, metadata, checks, positions): + self.metadata = metadata + self.checks = checks + self.positions = positions + + def matches(self, value: bytes, metadata: bytes) -> bool: + if metadata != self.metadata: + return False + for pos, expected in self.checks: + if value[pos:pos + len(expected)] != expected: + return False + return True + + +def _position_plan( + value: bytes, + metadata: bytes, + paths: Sequence[_Path], +) -> _PositionPlan: + key_ids = _metadata_key_ids(metadata) + nodes, result_nodes = _compile_paths(tuple(paths)) + positions = [0] + checks = [] + checked = set() + for parent_node, kind, segment in nodes[1:]: + parent = positions[parent_node] + if parent is None: + positions.append(None) + continue + if parent not in checked: + basic_type = value[parent] & 0x3 + if basic_type == _OBJECT: + data_start = _object_layout(value, parent)[-1] + elif basic_type == _ARRAY: + data_start = _array_layout(value, parent)[-1] + else: + data_start = parent + 1 + checks.append((parent, value[parent:data_start])) + checked.add(parent) + if kind == 'key': + key_id = key_ids.get(segment) + if key_id is None or (value[parent] & 0x3) != _OBJECT: + positions.append(None) + else: + positions.append(_object_field_position( + value, parent, key_id)) + elif (value[parent] & 0x3) != _ARRAY: + positions.append(None) + else: + positions.append(_array_element_position( + value, parent, segment)) + return _PositionPlan( + metadata, + tuple(checks), + tuple(positions[node] for node in result_nodes), + ) + + +def _cached_path_positions( + value: bytes, + metadata: bytes, + paths: Sequence[_Path], + plans: List[_PositionPlan], +) -> Sequence[Optional[int]]: + for index, plan in enumerate(plans): + if plan.matches(value, metadata): + if index: + plans.insert(0, plans.pop(index)) + return plan.positions + plan = _position_plan(value, metadata, paths) + plans.insert(0, plan) + del plans[8:] + return plan.positions + + +def _append_unsigned(buf: bytearray, value: int, size: int) -> None: + buf.extend(value.to_bytes(size, 'little')) + + +def _build_object(ids: Sequence[int], children: Sequence[bytes]) -> bytes: + size = len(ids) + data_size = sum(len(child) for child in children) + size_bytes = _U32_SIZE if size > _U8_MAX else 1 + id_size = _get_int_size(max(ids)) if ids else 1 + offset_size = _get_int_size(data_size) if data_size else 1 + + buf = bytearray([_object_header(size > _U8_MAX, id_size, offset_size)]) + _append_unsigned(buf, size, size_bytes) + for key_id in ids: + _append_unsigned(buf, key_id, id_size) + offset = 0 + for child in children: + _append_unsigned(buf, offset, offset_size) + offset += len(child) + _append_unsigned(buf, offset, offset_size) + for child in children: + buf.extend(child) + return bytes(buf) + + +def _build_array(children: Sequence[bytes]) -> bytes: + size = len(children) + data_size = sum(len(child) for child in children) + size_bytes = _U32_SIZE if size > _U8_MAX else 1 + offset_size = _get_int_size(data_size) if data_size else 1 + + buf = bytearray([_array_header(size > _U8_MAX, offset_size)]) + _append_unsigned(buf, size, size_bytes) + offset = 0 + for child in children: + _append_unsigned(buf, offset, offset_size) + offset += len(child) + _append_unsigned(buf, offset, offset_size) + for child in children: + buf.extend(child) + return bytes(buf) + + +def _replace_path( + value: bytes, + metadata: bytes, + pos: int, + path: _Path, + replacement: bytes, +) -> bytes: + if not path: + return replacement + + kind, segment = path[0] + if kind == 'key': + if (value[pos] & 0x3) != _OBJECT: + raise ValueError("VARIANT path expects an object") + key_id = _metadata_key_ids(metadata).get(segment) + if key_id is None: + raise ValueError(f"VARIANT path does not exist: {segment}") + size, id_size, offset_size, id_start, offset_start, data_start = ( + _object_layout(value, pos)) + ids = [ + _read_unsigned(value, id_start + i * id_size, id_size) + for i in range(size) + ] + try: + slot = ids.index(key_id) + except ValueError: + raise ValueError(f"VARIANT path does not exist: {segment}") + children = [] + for i in range(size): + child_pos = data_start + _read_unsigned( + value, offset_start + i * offset_size, offset_size) + child = value[child_pos:child_pos + _value_size(value, child_pos)] + if i == slot: + child = _replace_path( + value, metadata, child_pos, path[1:], replacement) + children.append(child) + return _build_object(ids, children) + + if (value[pos] & 0x3) != _ARRAY: + raise ValueError("VARIANT path expects an array") + size, offset_size, offset_start, data_start = _array_layout(value, pos) + if segment >= size: + raise ValueError(f"VARIANT array index does not exist: {segment}") + children = [] + for i in range(size): + child_pos = data_start + _read_unsigned( + value, offset_start + i * offset_size, offset_size) + child = value[child_pos:child_pos + _value_size(value, child_pos)] + if i == segment: + child = _replace_path( + value, metadata, child_pos, path[1:], replacement) + children.append(child) + return _build_array(children) + + +def _variant_chunks(column): + if isinstance(column, pa.ChunkedArray): + chunks = column.chunks + chunked = True + data_type = column.type + elif isinstance(column, pa.Array): + chunks = [column] + chunked = False + data_type = column.type + else: + raise TypeError( + "VARIANT input must be a PyArrow Array or ChunkedArray") + if not pa.types.is_struct(data_type): + raise TypeError( + "VARIANT input must use the canonical Arrow struct type") + names = [field.name for field in data_type] + if names != ['value', 'metadata']: + raise TypeError("VARIANT input must contain value and metadata fields") + return chunks, chunked, data_type + + +def _build_variant_array(values, metadatas, nulls, data_type): + return pa.StructArray.from_arrays( + [pa.array(values, type=data_type[0].type), + pa.array(metadatas, type=data_type[1].type)], + fields=list(data_type), + mask=pa.array(nulls, type=pa.bool_()), + ) + + +def variant_get_many(column, paths: Mapping[str, object]): + """Extract several VARIANT paths in one pass. + + ``paths`` maps each path to its requested PyArrow type. Use ``None`` to + keep a result in the canonical VARIANT Arrow representation. + """ + parsed_paths = [ + (path, _parse_path(path), target_type) + for path, target_type in paths.items() + ] + chunks, chunked, data_type = _variant_chunks(column) + result_chunks = {path: [] for path in paths} + position_plans = [] + for chunk in chunks: + input_values = chunk.field(0).to_pylist() + input_metadatas = chunk.field(1).to_pylist() + valid = chunk.is_valid().to_pylist() + rows = { + path: ([], [], []) if target_type is None else [] + for path, _, target_type in parsed_paths + } + for row in range(len(chunk)): + if not valid[row]: + for path, _, target_type in parsed_paths: + if target_type is None: + values, metadatas, nulls = rows[path] + values.append(b'') + metadatas.append(b'') + nulls.append(True) + else: + rows[path].append(None) + continue + value = input_values[row] + metadata = input_metadatas[row] + positions = _cached_path_positions( + value, metadata, + [parsed for _, parsed, _ in parsed_paths], + position_plans) + for (path, _, target_type), pos in zip(parsed_paths, positions): + if pos is None: + if target_type is None: + values, metadatas, nulls = rows[path] + values.append(b'') + metadatas.append(b'') + nulls.append(True) + else: + rows[path].append(None) + elif target_type is None: + values, metadatas, nulls = rows[path] + values.append(value[pos:pos + _value_size(value, pos)]) + metadatas.append(metadata) + nulls.append(False) + else: + rows[path].append(_decode_scalar(value, metadata, pos)) + + for path, _, target_type in parsed_paths: + if target_type is None: + values, metadatas, nulls = rows[path] + result = _build_variant_array( + values, metadatas, nulls, data_type) + else: + result = pa.array(rows[path], type=target_type) + result_chunks[path].append(result) + + results = {} + for path, _, target_type in parsed_paths: + if not chunked: + results[path] = result_chunks[path][0] + else: + result_type = data_type if target_type is None else target_type + results[path] = pa.chunked_array( + result_chunks[path], type=result_type) + return results + + +def variant_get(column, path: str, target_type=None): + """Extract one VARIANT path without decoding unrelated fields. + + If ``target_type`` is omitted, the result remains a VARIANT Arrow struct. + Otherwise values are converted to the requested PyArrow type. + Missing paths and SQL NULL inputs produce NULL outputs. + """ + return variant_get_many(column, {path: target_type})[path] + + +class _Replacement: + + def __init__(self, value, length: int): + if isinstance(value, (pa.Array, pa.ChunkedArray)): + if len(value) != length: + raise ValueError( + "VARIANT replacement length must match the input column") + self._array = value + self._value = None + elif isinstance(value, pa.Scalar): + self._array = None + self._value = value.as_py() + else: + self._array = None + self._value = value + + def values(self, offset: int, length: int): + if self._array is None: + return [self._value] * length + return self._array.slice(offset, length).to_pylist() + + +def _encode_replacement(value) -> bytes: + if isinstance(value, (dict, list, tuple)): + raise TypeError( + "variant_set only accepts scalar replacements; nested values " + "may require new metadata keys") + if isinstance(value, float): + return bytes([_DOUBLE_HEADER]) + struct.pack(' bool: + limit = min(len(first), len(second)) + return first[:limit] == second[:limit] + + +def _validate_distinct_paths(paths: Sequence[_Path]) -> None: + for i, first in enumerate(paths): + for second in paths[i + 1:]: + if _paths_overlap(first, second): + raise ValueError("VARIANT replacement paths must not overlap") + + +def _replace_encoded_values( + value: bytes, + metadata: bytes, + paths: Sequence[_Path], + path_names: Sequence[str], + replacements: Sequence[bytes], + position_plans: List[_PositionPlan], + found_positions=None, +) -> bytes: + if found_positions is None: + found_positions = _cached_path_positions( + value, metadata, paths, position_plans) + positions = [] + fixed_size = True + for path, replacement, pos in zip( + path_names, replacements, found_positions): + if pos is None: + raise ValueError(f"VARIANT path does not exist: {path}") + size = _value_size(value, pos) + positions.append((pos, size, replacement)) + fixed_size = fixed_size and size == len(replacement) + + if fixed_size: + result = bytearray(value) + for pos, size, replacement in positions: + result[pos:pos + size] = replacement + return bytes(result) + + for path, replacement in zip(paths, replacements): + value = _replace_path(value, metadata, 0, path, replacement) + return value + + +def variant_set_many(column, replacements: Mapping[str, object]): + """Replace existing scalar VARIANT paths without a full object round trip. + + ``replacements`` maps paths to PyArrow arrays or scalar values. Array + replacements must have the same row count as ``column``. All paths are + updated in one pass, so common fixed-width updates copy each VARIANT value + only once. Paths must already exist; this function never changes metadata. + """ + chunks, chunked, data_type = _variant_chunks(column) + total_length = len(column) + parsed_replacements = [ + (path, _parse_path(path), _Replacement(value, total_length)) + for path, value in replacements.items() + ] + paths = [parsed for _, parsed, _ in parsed_replacements] + path_names = [path for path, _, _ in parsed_replacements] + _validate_distinct_paths(paths) + + result_chunks = [] + position_plans = [] + global_row = 0 + for chunk in chunks: + input_values = chunk.field(0).to_pylist() + input_metadatas = chunk.field(1).to_pylist() + valid = chunk.is_valid().to_pylist() + replacement_values = [ + replacement.values(global_row, len(chunk)) + for _, _, replacement in parsed_replacements + ] + values = [] + metadatas = [] + nulls = [] + for local_row in range(len(chunk)): + if not valid[local_row]: + values.append(b'') + metadatas.append(b'') + nulls.append(True) + global_row += 1 + continue + value = input_values[local_row] + metadata = input_metadatas[local_row] + row_replacements = [ + _encode_replacement(replacement_values[i][local_row]) + for i in range(len(parsed_replacements)) + ] + value = _replace_encoded_values( + value, metadata, paths, path_names, row_replacements, + position_plans) + values.append(value) + metadatas.append(metadata) + nulls.append(False) + global_row += 1 + result_chunks.append( + _build_variant_array(values, metadatas, nulls, data_type)) + + if not chunked: + return result_chunks[0] + return pa.chunked_array(result_chunks, type=data_type) + + +def variant_set(column, path: str, values): + """Replace one existing scalar VARIANT path for every input row.""" + return variant_set_many(column, {path: values}) + + +def variant_transform(column, transforms: Mapping[str, object]): + """Transform existing scalar VARIANT paths in one pass. + + Each callable receives only the scalar at its path. Unrelated fields are + neither decoded nor rebuilt. This is preferable to ``variant_get`` plus + ``variant_set_many`` when the operation can be expressed as a Python + scalar function and no intermediate Arrow arrays are needed. + """ + chunks, chunked, data_type = _variant_chunks(column) + parsed_transforms = [ + (path, _parse_path(path), transform) + for path, transform in transforms.items() + ] + for path, _, transform in parsed_transforms: + if not callable(transform): + raise TypeError(f"VARIANT transform for {path} must be callable") + paths = [parsed for _, parsed, _ in parsed_transforms] + path_names = [path for path, _, _ in parsed_transforms] + _validate_distinct_paths(paths) + + result_chunks = [] + position_plans = [] + for chunk in chunks: + input_values = chunk.field(0).to_pylist() + input_metadatas = chunk.field(1).to_pylist() + valid = chunk.is_valid().to_pylist() + values = [] + metadatas = [] + nulls = [] + for row in range(len(chunk)): + if not valid[row]: + values.append(b'') + metadatas.append(b'') + nulls.append(True) + continue + value = input_values[row] + metadata = input_metadatas[row] + positions = _cached_path_positions( + value, metadata, paths, position_plans) + replacements = [] + for (path, _, transform), pos in zip( + parsed_transforms, positions): + if pos is None: + raise ValueError(f"VARIANT path does not exist: {path}") + current = _decode_scalar(value, metadata, pos) + replacements.append(_encode_replacement(transform(current))) + value = _replace_encoded_values( + value, metadata, paths, path_names, replacements, + position_plans, positions) + values.append(value) + metadatas.append(metadata) + nulls.append(False) + result_chunks.append( + _build_variant_array(values, metadatas, nulls, data_type)) + + if not chunked: + return result_chunks[0] + return pa.chunked_array(result_chunks, type=data_type) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py new file mode 100644 index 000000000000..8988361c5637 --- /dev/null +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -0,0 +1,237 @@ +# 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. + +import unittest +import operator +from unittest.mock import patch + +import pyarrow as pa + +from pypaimon.data.generic_variant import GenericVariant +from pypaimon.data.variant_path import ( + variant_get, + variant_get_many, + variant_set, + variant_set_many, + variant_transform, +) + + +def _variants(values): + return GenericVariant.to_arrow_array([ + GenericVariant.from_python(value) if value is not None else None + for value in values + ]) + + +def _decode(column): + return [ + None if value is None + else GenericVariant.from_arrow_struct(value).to_python() + for value in column.to_pylist() + ] + + +class TestVariantGet(unittest.TestCase): + + def test_get_nested_scalars(self): + column = _variants([ + {'imu': {'velocity': {'y': 1.5}}}, + {'imu': {'velocity': {'y': -2.0}}}, + {'imu': {}}, + None, + ]) + + result = variant_get( + column, '$.imu.velocity.y', pa.float64()) + + self.assertEqual(result.to_pylist(), [1.5, -2.0, None, None]) + + def test_get_quoted_key_and_array_index(self): + column = _variants([{'a.b': [{'value': 3}, {'value': 7}]}]) + + result = variant_get( + column, '$["a.b"][1].value', pa.int64()) + + self.assertEqual(result.to_pylist(), [7]) + + def test_get_variant_subvalue(self): + column = _variants([{'nested': {'value': 3}}, None]) + + result = variant_get(column, '$.nested') + + self.assertEqual(_decode(result), [{'value': 3}, None]) + + def test_get_preserves_chunks(self): + first = _variants([{'value': 1}]) + second = _variants([{'value': 2}, {'value': 3}]) + column = pa.chunked_array([first, second]) + + result = variant_get(column, '$.value', pa.int64()) + + self.assertIsInstance(result, pa.ChunkedArray) + self.assertEqual(result.num_chunks, 2) + self.assertEqual(result.to_pylist(), [1, 2, 3]) + + def test_get_multiple_paths(self): + column = _variants([{ + 'angular': {'y': 1.0, 'z': 2.0}, + 'linear': {'y': 3.0, 'z': 4.0}, + }]) + + result = variant_get_many(column, { + '$.angular.y': pa.float64(), + '$.angular.z': pa.float64(), + '$.linear.y': pa.float64(), + '$.linear.z': pa.float64(), + }) + + self.assertEqual(result['$.angular.y'].to_pylist(), [1.0]) + self.assertEqual(result['$.linear.z'].to_pylist(), [4.0]) + + +class TestVariantSet(unittest.TestCase): + + def test_set_multiple_double_paths(self): + column = _variants([ + { + 'angular_velocity': {'x': 1.0, 'y': 2.0, 'z': 3.0}, + 'linear_acceleration': {'x': 4.0, 'y': 5.0, 'z': 6.0}, + 'unchanged': 'keep', + }, + { + 'angular_velocity': {'x': 7.0, 'y': 8.0, 'z': 9.0}, + 'linear_acceleration': {'x': 10.0, 'y': 11.0, 'z': 12.0}, + 'unchanged': 'also keep', + }, + ]) + metadata_before = column.field('metadata').to_pylist() + + result = variant_set_many(column, { + '$.angular_velocity.y': pa.array([-2.0, -8.0]), + '$.angular_velocity.z': pa.array([-3.0, -9.0]), + '$.linear_acceleration.y': pa.array([-5.0, -11.0]), + '$.linear_acceleration.z': pa.array([-6.0, -12.0]), + }) + + decoded = _decode(result) + self.assertEqual(decoded[0]['angular_velocity'], + {'x': 1.0, 'y': -2.0, 'z': -3.0}) + self.assertEqual(decoded[1]['linear_acceleration'], + {'x': 10.0, 'y': -11.0, 'z': -12.0}) + self.assertEqual(decoded[0]['unchanged'], 'keep') + self.assertEqual( + result.field('metadata').to_pylist(), metadata_before) + + def test_transform_multiple_double_paths(self): + column = _variants([{ + 'angular': {'y': 1.0, 'z': -2.0}, + 'linear': {'y': 3.0, 'z': -4.0}, + 'other': 'keep', + }]) + + result = variant_transform(column, { + '$.angular.y': operator.neg, + '$.angular.z': operator.neg, + '$.linear.y': operator.neg, + '$.linear.z': operator.neg, + }) + + self.assertEqual(_decode(result), [{ + 'angular': {'y': -1.0, 'z': 2.0}, + 'linear': {'y': -3.0, 'z': 4.0}, + 'other': 'keep', + }]) + + def test_set_finds_path_after_variable_length_value(self): + column = _variants([ + {'prefix': 'x', 'nested': {'value': 1.0}}, + {'prefix': 'x' * 200, 'nested': {'value': 2.0}}, + ]) + + result = variant_set( + column, '$.nested.value', pa.array([-1.0, -2.0])) + + decoded = _decode(result) + self.assertEqual(decoded[0]['nested']['value'], -1.0) + self.assertEqual(decoded[1]['nested']['value'], -2.0) + self.assertEqual(decoded[1]['prefix'], 'x' * 200) + + def test_set_rebuilds_offsets_for_different_size(self): + column = _variants([ + {'before': 'a', 'target': 'x', 'after': {'value': 7}}, + ]) + + result = variant_set(column, '$.target', 'a much longer value') + + self.assertEqual(_decode(result), [{ + 'before': 'a', + 'target': 'a much longer value', + 'after': {'value': 7}, + }]) + + def test_set_array_element(self): + column = _variants([{'values': [1, 2, 3]}]) + + result = variant_set(column, '$.values[1]', 100000) + + self.assertEqual(_decode(result), [{'values': [1, 100000, 3]}]) + + def test_set_preserves_sql_null_and_chunks(self): + first = _variants([{'value': 1}, None]) + second = _variants([{'value': 3}]) + column = pa.chunked_array([first, second]) + + result = variant_set( + column, '$.value', pa.array([10, 20, 30])) + + self.assertIsInstance(result, pa.ChunkedArray) + self.assertEqual(result.num_chunks, 2) + self.assertEqual(_decode(result), [{'value': 10}, None, {'value': 30}]) + + def test_set_does_not_decode_whole_variant(self): + column = _variants([{'nested': {'value': 1.0}, 'other': [1, 2, 3]}]) + + with patch.object( + GenericVariant, + 'to_python', + side_effect=AssertionError("full decode is not allowed")): + result = variant_set(column, '$.nested.value', -1.0) + + self.assertEqual(_decode(result), [ + {'nested': {'value': -1.0}, 'other': [1, 2, 3]}, + ]) + + def test_set_rejects_missing_and_overlapping_paths(self): + column = _variants([{'nested': {'value': 1}}]) + + with self.assertRaisesRegex(ValueError, "path does not exist"): + variant_set(column, '$.missing', 2) + with self.assertRaisesRegex(ValueError, "must not overlap"): + variant_set_many(column, { + '$.nested': 2, + '$.nested.value': 3, + }) + + def test_set_rejects_nested_replacement(self): + column = _variants([{'nested': {'value': 1}}]) + + with self.assertRaisesRegex(TypeError, "scalar replacements"): + variant_set(column, '$.nested', {'value': 2}) + + +if __name__ == '__main__': + unittest.main() From 7095ff1620497a58a5a32687b024e1ae9dab28aa Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 03:18:55 -0700 Subject: [PATCH 02/23] [python] Limit VARIANT path transforms to DOUBLE --- docs/docs/pypaimon/python-api.mdx | 10 +- paimon-python/pypaimon/data/__init__.py | 12 +- paimon-python/pypaimon/data/variant_path.py | 471 ++---------------- .../pypaimon/tests/variant_path_test.py | 229 ++------- 4 files changed, 113 insertions(+), 609 deletions(-) diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index 152e9939a19b..699fe4a41f44 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -1126,10 +1126,9 @@ Supported Paimon type strings for shredded sub-fields: `BOOLEAN`, `INT`, `BIGINT -### VARIANT Path Updates +### VARIANT DOUBLE Updates -Use path helpers when only a few scalar fields change. They avoid decoding -unrelated fields into Python objects: +Transform existing DOUBLE paths without decoding unrelated fields: ```python import operator @@ -1142,10 +1141,7 @@ updated_payload = variant_transform(payload, { }) ``` -`variant_get` extracts one path as an Arrow array. `variant_get_many` and -`variant_set_many` process several paths in one pass. `variant_set` replaces -one path. Set operations require existing paths and scalar replacement values; -they keep the original VARIANT metadata. +The paths must exist and contain DOUBLE values. The metadata is unchanged. **`GenericVariant` API:** diff --git a/paimon-python/pypaimon/data/__init__.py b/paimon-python/pypaimon/data/__init__.py index d50cf1f77230..f30d8ab71d9b 100644 --- a/paimon-python/pypaimon/data/__init__.py +++ b/paimon-python/pypaimon/data/__init__.py @@ -17,20 +17,10 @@ from pypaimon.data.timestamp import Timestamp from pypaimon.data.decimal import Decimal -from pypaimon.data.variant_path import ( - variant_get, - variant_get_many, - variant_set, - variant_set_many, - variant_transform, -) +from pypaimon.data.variant_path import variant_transform __all__ = [ 'Timestamp', 'Decimal', - 'variant_get', - 'variant_get_many', - 'variant_set', - 'variant_set_many', 'variant_transform', ] diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 6616f4bd2412..9b9a9a34901c 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Arrow helpers for accessing VARIANT paths without decoding whole values.""" +"""Transform existing DOUBLE paths in Arrow VARIANT columns.""" import functools import re @@ -26,28 +26,17 @@ from pypaimon.data._variant_binary import ( _ARRAY, _OBJECT, - _U8_MAX, _U32_SIZE, - _array_header, - _get_int_size, - _object_header, _read_unsigned, ) -from pypaimon.data.generic_variant import ( - GenericVariant, - _Type, - _value_size, - _variant_get_type, -) +from pypaimon.data.generic_variant import _Type, _variant_get_type _INDEX_PATTERN = re.compile(r"\[(\d+)]") -_KEY_PATTERN = re.compile(r"\.([^.\[]+)|\['([^']+)']|\[\"([^\"]+)\"]") - +_KEY_PATTERN = re.compile(r"\.([^\.\[]+)|\['([^']+)']|\[\"([^\"]+)\"]") _Path = Tuple[Tuple[str, object], ...] _ObjectLayout = Tuple[int, int, int, int, int, int] _ArrayLayout = Tuple[int, int, int, int] -_DOUBLE_HEADER = 7 << 2 @functools.lru_cache(maxsize=256) @@ -61,16 +50,14 @@ def _parse_path(path: str) -> _Path: match = _INDEX_PATTERN.match(path, pos) if match is not None: segments.append(('index', int(match.group(1)))) - pos = match.end() - continue - - match = _KEY_PATTERN.match(path, pos) - if match is not None: - key = next(value for value in match.groups() if value is not None) + else: + match = _KEY_PATTERN.match(path, pos) + if match is None: + raise ValueError(f"Invalid VARIANT path: {path}") + key = next(value for value in match.groups() + if value is not None) segments.append(('key', key)) - pos = match.end() - continue - raise ValueError(f"Invalid VARIANT path: {path}") + pos = match.end() return tuple(segments) @@ -87,9 +74,11 @@ def _metadata_key_ids(metadata: bytes) -> Dict[str, int]: start = _read_unsigned( metadata, offset_start + key_id * offset_size, offset_size) end = _read_unsigned( - metadata, offset_start + (key_id + 1) * offset_size, offset_size) - key = metadata[string_start + start:string_start + end].decode('utf-8') - result[key] = key_id + metadata, offset_start + (key_id + 1) * offset_size, + offset_size) + result[ + metadata[string_start + start:string_start + end].decode('utf-8') + ] = key_id return result @@ -130,10 +119,7 @@ def _field_slot(id_table: bytes, id_size: int, key_id: int) -> Optional[int]: def _object_field_position( - value: bytes, - pos: int, - key_id: int, -) -> Optional[int]: + value: bytes, pos: int, key_id: int) -> Optional[int]: size, id_size, offset_size, id_start, offset_start, data_start = ( _object_layout(value, pos)) id_table = value[id_start:id_start + size * id_size] @@ -146,10 +132,7 @@ def _object_field_position( def _array_element_position( - value: bytes, - pos: int, - index: int, -) -> Optional[int]: + value: bytes, pos: int, index: int) -> Optional[int]: size, offset_size, offset_start, data_start = _array_layout(value, pos) if index >= size: return None @@ -167,10 +150,8 @@ def _compile_paths(paths: Tuple[_Path, ...]): for length in range(1, len(path) + 1): prefix = path[:length] if prefix not in node_by_prefix: - parent = node_by_prefix[prefix[:-1]] - kind, segment = prefix[-1] node_by_prefix[prefix] = len(nodes) - nodes.append((parent, kind, segment)) + nodes.append((node_by_prefix[prefix[:-1]],) + prefix[-1]) results.append(node_by_prefix[path]) return tuple(nodes), tuple(results) @@ -183,12 +164,10 @@ def __init__(self, metadata, checks, positions): self.positions = positions def matches(self, value: bytes, metadata: bytes) -> bool: - if metadata != self.metadata: - return False - for pos, expected in self.checks: - if value[pos:pos + len(expected)] != expected: - return False - return True + return metadata == self.metadata and all( + value[pos:pos + len(expected)] == expected + for pos, expected in self.checks + ) def _position_plan( @@ -235,7 +214,7 @@ def _position_plan( ) -def _cached_path_positions( +def _path_positions( value: bytes, metadata: bytes, paths: Sequence[_Path], @@ -252,126 +231,23 @@ def _cached_path_positions( return plan.positions -def _append_unsigned(buf: bytearray, value: int, size: int) -> None: - buf.extend(value.to_bytes(size, 'little')) - - -def _build_object(ids: Sequence[int], children: Sequence[bytes]) -> bytes: - size = len(ids) - data_size = sum(len(child) for child in children) - size_bytes = _U32_SIZE if size > _U8_MAX else 1 - id_size = _get_int_size(max(ids)) if ids else 1 - offset_size = _get_int_size(data_size) if data_size else 1 - - buf = bytearray([_object_header(size > _U8_MAX, id_size, offset_size)]) - _append_unsigned(buf, size, size_bytes) - for key_id in ids: - _append_unsigned(buf, key_id, id_size) - offset = 0 - for child in children: - _append_unsigned(buf, offset, offset_size) - offset += len(child) - _append_unsigned(buf, offset, offset_size) - for child in children: - buf.extend(child) - return bytes(buf) - - -def _build_array(children: Sequence[bytes]) -> bytes: - size = len(children) - data_size = sum(len(child) for child in children) - size_bytes = _U32_SIZE if size > _U8_MAX else 1 - offset_size = _get_int_size(data_size) if data_size else 1 - - buf = bytearray([_array_header(size > _U8_MAX, offset_size)]) - _append_unsigned(buf, size, size_bytes) - offset = 0 - for child in children: - _append_unsigned(buf, offset, offset_size) - offset += len(child) - _append_unsigned(buf, offset, offset_size) - for child in children: - buf.extend(child) - return bytes(buf) - - -def _replace_path( - value: bytes, - metadata: bytes, - pos: int, - path: _Path, - replacement: bytes, -) -> bytes: - if not path: - return replacement - - kind, segment = path[0] - if kind == 'key': - if (value[pos] & 0x3) != _OBJECT: - raise ValueError("VARIANT path expects an object") - key_id = _metadata_key_ids(metadata).get(segment) - if key_id is None: - raise ValueError(f"VARIANT path does not exist: {segment}") - size, id_size, offset_size, id_start, offset_start, data_start = ( - _object_layout(value, pos)) - ids = [ - _read_unsigned(value, id_start + i * id_size, id_size) - for i in range(size) - ] - try: - slot = ids.index(key_id) - except ValueError: - raise ValueError(f"VARIANT path does not exist: {segment}") - children = [] - for i in range(size): - child_pos = data_start + _read_unsigned( - value, offset_start + i * offset_size, offset_size) - child = value[child_pos:child_pos + _value_size(value, child_pos)] - if i == slot: - child = _replace_path( - value, metadata, child_pos, path[1:], replacement) - children.append(child) - return _build_object(ids, children) - - if (value[pos] & 0x3) != _ARRAY: - raise ValueError("VARIANT path expects an array") - size, offset_size, offset_start, data_start = _array_layout(value, pos) - if segment >= size: - raise ValueError(f"VARIANT array index does not exist: {segment}") - children = [] - for i in range(size): - child_pos = data_start + _read_unsigned( - value, offset_start + i * offset_size, offset_size) - child = value[child_pos:child_pos + _value_size(value, child_pos)] - if i == segment: - child = _replace_path( - value, metadata, child_pos, path[1:], replacement) - children.append(child) - return _build_array(children) - - def _variant_chunks(column): if isinstance(column, pa.ChunkedArray): - chunks = column.chunks - chunked = True - data_type = column.type + chunks, chunked, data_type = column.chunks, True, column.type elif isinstance(column, pa.Array): - chunks = [column] - chunked = False - data_type = column.type + chunks, chunked, data_type = [column], False, column.type else: raise TypeError( "VARIANT input must be a PyArrow Array or ChunkedArray") - if not pa.types.is_struct(data_type): + if (not pa.types.is_struct(data_type) + or [field.name for field in data_type] + != ['value', 'metadata']): raise TypeError( - "VARIANT input must use the canonical Arrow struct type") - names = [field.name for field in data_type] - if names != ['value', 'metadata']: - raise TypeError("VARIANT input must contain value and metadata fields") + "VARIANT input must contain value and metadata fields") return chunks, chunked, data_type -def _build_variant_array(values, metadatas, nulls, data_type): +def _variant_array(values, metadatas, nulls, data_type): return pa.StructArray.from_arrays( [pa.array(values, type=data_type[0].type), pa.array(metadatas, type=data_type[1].type)], @@ -380,293 +256,56 @@ def _build_variant_array(values, metadatas, nulls, data_type): ) -def variant_get_many(column, paths: Mapping[str, object]): - """Extract several VARIANT paths in one pass. - - ``paths`` maps each path to its requested PyArrow type. Use ``None`` to - keep a result in the canonical VARIANT Arrow representation. - """ - parsed_paths = [ - (path, _parse_path(path), target_type) - for path, target_type in paths.items() - ] - chunks, chunked, data_type = _variant_chunks(column) - result_chunks = {path: [] for path in paths} - position_plans = [] - for chunk in chunks: - input_values = chunk.field(0).to_pylist() - input_metadatas = chunk.field(1).to_pylist() - valid = chunk.is_valid().to_pylist() - rows = { - path: ([], [], []) if target_type is None else [] - for path, _, target_type in parsed_paths - } - for row in range(len(chunk)): - if not valid[row]: - for path, _, target_type in parsed_paths: - if target_type is None: - values, metadatas, nulls = rows[path] - values.append(b'') - metadatas.append(b'') - nulls.append(True) - else: - rows[path].append(None) - continue - value = input_values[row] - metadata = input_metadatas[row] - positions = _cached_path_positions( - value, metadata, - [parsed for _, parsed, _ in parsed_paths], - position_plans) - for (path, _, target_type), pos in zip(parsed_paths, positions): - if pos is None: - if target_type is None: - values, metadatas, nulls = rows[path] - values.append(b'') - metadatas.append(b'') - nulls.append(True) - else: - rows[path].append(None) - elif target_type is None: - values, metadatas, nulls = rows[path] - values.append(value[pos:pos + _value_size(value, pos)]) - metadatas.append(metadata) - nulls.append(False) - else: - rows[path].append(_decode_scalar(value, metadata, pos)) - - for path, _, target_type in parsed_paths: - if target_type is None: - values, metadatas, nulls = rows[path] - result = _build_variant_array( - values, metadatas, nulls, data_type) - else: - result = pa.array(rows[path], type=target_type) - result_chunks[path].append(result) - - results = {} - for path, _, target_type in parsed_paths: - if not chunked: - results[path] = result_chunks[path][0] - else: - result_type = data_type if target_type is None else target_type - results[path] = pa.chunked_array( - result_chunks[path], type=result_type) - return results - - -def variant_get(column, path: str, target_type=None): - """Extract one VARIANT path without decoding unrelated fields. - - If ``target_type`` is omitted, the result remains a VARIANT Arrow struct. - Otherwise values are converted to the requested PyArrow type. - Missing paths and SQL NULL inputs produce NULL outputs. - """ - return variant_get_many(column, {path: target_type})[path] - - -class _Replacement: - - def __init__(self, value, length: int): - if isinstance(value, (pa.Array, pa.ChunkedArray)): - if len(value) != length: - raise ValueError( - "VARIANT replacement length must match the input column") - self._array = value - self._value = None - elif isinstance(value, pa.Scalar): - self._array = None - self._value = value.as_py() - else: - self._array = None - self._value = value - - def values(self, offset: int, length: int): - if self._array is None: - return [self._value] * length - return self._array.slice(offset, length).to_pylist() - - -def _encode_replacement(value) -> bytes: - if isinstance(value, (dict, list, tuple)): - raise TypeError( - "variant_set only accepts scalar replacements; nested values " - "may require new metadata keys") - if isinstance(value, float): - return bytes([_DOUBLE_HEADER]) + struct.pack(' bool: - limit = min(len(first), len(second)) - return first[:limit] == second[:limit] - - -def _validate_distinct_paths(paths: Sequence[_Path]) -> None: - for i, first in enumerate(paths): - for second in paths[i + 1:]: - if _paths_overlap(first, second): - raise ValueError("VARIANT replacement paths must not overlap") - - -def _replace_encoded_values( - value: bytes, - metadata: bytes, - paths: Sequence[_Path], - path_names: Sequence[str], - replacements: Sequence[bytes], - position_plans: List[_PositionPlan], - found_positions=None, -) -> bytes: - if found_positions is None: - found_positions = _cached_path_positions( - value, metadata, paths, position_plans) - positions = [] - fixed_size = True - for path, replacement, pos in zip( - path_names, replacements, found_positions): - if pos is None: - raise ValueError(f"VARIANT path does not exist: {path}") - size = _value_size(value, pos) - positions.append((pos, size, replacement)) - fixed_size = fixed_size and size == len(replacement) - - if fixed_size: - result = bytearray(value) - for pos, size, replacement in positions: - result[pos:pos + size] = replacement - return bytes(result) - - for path, replacement in zip(paths, replacements): - value = _replace_path(value, metadata, 0, path, replacement) - return value - - -def variant_set_many(column, replacements: Mapping[str, object]): - """Replace existing scalar VARIANT paths without a full object round trip. - - ``replacements`` maps paths to PyArrow arrays or scalar values. Array - replacements must have the same row count as ``column``. All paths are - updated in one pass, so common fixed-width updates copy each VARIANT value - only once. Paths must already exist; this function never changes metadata. - """ - chunks, chunked, data_type = _variant_chunks(column) - total_length = len(column) - parsed_replacements = [ - (path, _parse_path(path), _Replacement(value, total_length)) - for path, value in replacements.items() - ] - paths = [parsed for _, parsed, _ in parsed_replacements] - path_names = [path for path, _, _ in parsed_replacements] - _validate_distinct_paths(paths) - - result_chunks = [] - position_plans = [] - global_row = 0 - for chunk in chunks: - input_values = chunk.field(0).to_pylist() - input_metadatas = chunk.field(1).to_pylist() - valid = chunk.is_valid().to_pylist() - replacement_values = [ - replacement.values(global_row, len(chunk)) - for _, _, replacement in parsed_replacements - ] - values = [] - metadatas = [] - nulls = [] - for local_row in range(len(chunk)): - if not valid[local_row]: - values.append(b'') - metadatas.append(b'') - nulls.append(True) - global_row += 1 - continue - value = input_values[local_row] - metadata = input_metadatas[local_row] - row_replacements = [ - _encode_replacement(replacement_values[i][local_row]) - for i in range(len(parsed_replacements)) - ] - value = _replace_encoded_values( - value, metadata, paths, path_names, row_replacements, - position_plans) - values.append(value) - metadatas.append(metadata) - nulls.append(False) - global_row += 1 - result_chunks.append( - _build_variant_array(values, metadatas, nulls, data_type)) - - if not chunked: - return result_chunks[0] - return pa.chunked_array(result_chunks, type=data_type) - - -def variant_set(column, path: str, values): - """Replace one existing scalar VARIANT path for every input row.""" - return variant_set_many(column, {path: values}) - - def variant_transform(column, transforms: Mapping[str, object]): - """Transform existing scalar VARIANT paths in one pass. - - Each callable receives only the scalar at its path. Unrelated fields are - neither decoded nor rebuilt. This is preferable to ``variant_get`` plus - ``variant_set_many`` when the operation can be expressed as a Python - scalar function and no intermediate Arrow arrays are needed. - """ - chunks, chunked, data_type = _variant_chunks(column) - parsed_transforms = [ + """Transform existing DOUBLE paths without decoding the whole VARIANT.""" + parsed = [ (path, _parse_path(path), transform) for path, transform in transforms.items() ] - for path, _, transform in parsed_transforms: + for path, _, transform in parsed: if not callable(transform): raise TypeError(f"VARIANT transform for {path} must be callable") - paths = [parsed for _, parsed, _ in parsed_transforms] - path_names = [path for path, _, _ in parsed_transforms] - _validate_distinct_paths(paths) + if not parsed: + return column + chunks, chunked, data_type = _variant_chunks(column) + paths = [path for _, path, _ in parsed] + plans = [] result_chunks = [] - position_plans = [] for chunk in chunks: input_values = chunk.field(0).to_pylist() input_metadatas = chunk.field(1).to_pylist() valid = chunk.is_valid().to_pylist() - values = [] - metadatas = [] - nulls = [] + values, metadatas, nulls = [], [], [] for row in range(len(chunk)): if not valid[row]: values.append(b'') metadatas.append(b'') nulls.append(True) continue + value = input_values[row] metadata = input_metadatas[row] - positions = _cached_path_positions( - value, metadata, paths, position_plans) - replacements = [] - for (path, _, transform), pos in zip( - parsed_transforms, positions): + result = bytearray(value) + positions = _path_positions(value, metadata, paths, plans) + for (path, _, transform), pos in zip(parsed, positions): if pos is None: raise ValueError(f"VARIANT path does not exist: {path}") - current = _decode_scalar(value, metadata, pos) - replacements.append(_encode_replacement(transform(current))) - value = _replace_encoded_values( - value, metadata, paths, path_names, replacements, - position_plans, positions) - values.append(value) + if _variant_get_type(value, pos) != _Type.DOUBLE: + raise TypeError(f"VARIANT path is not DOUBLE: {path}") + current = struct.unpack_from(' Date: Mon, 10 Aug 2026 03:22:21 -0700 Subject: [PATCH 03/23] [python] Tighten DOUBLE path transformations --- paimon-python/pypaimon/data/variant_path.py | 21 ++++++++++--------- .../pypaimon/tests/variant_path_test.py | 6 +++++- 2 files changed, 16 insertions(+), 11 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 9b9a9a34901c..486c477aa08d 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -247,10 +247,9 @@ def _variant_chunks(column): return chunks, chunked, data_type -def _variant_array(values, metadatas, nulls, data_type): +def _variant_array(values, metadata, nulls, data_type): return pa.StructArray.from_arrays( - [pa.array(values, type=data_type[0].type), - pa.array(metadatas, type=data_type[1].type)], + [pa.array(values, type=data_type[0].type), metadata], fields=list(data_type), mask=pa.array(nulls, type=pa.bool_()), ) @@ -274,13 +273,13 @@ def variant_transform(column, transforms: Mapping[str, object]): result_chunks = [] for chunk in chunks: input_values = chunk.field(0).to_pylist() - input_metadatas = chunk.field(1).to_pylist() + metadata_array = chunk.field(1) + input_metadatas = metadata_array.to_pylist() valid = chunk.is_valid().to_pylist() - values, metadatas, nulls = [], [], [] + values, nulls = [], [] for row in range(len(chunk)): if not valid[row]: values.append(b'') - metadatas.append(b'') nulls.append(True) continue @@ -294,18 +293,20 @@ def variant_transform(column, transforms: Mapping[str, object]): if _variant_get_type(value, pos) != _Type.DOUBLE: raise TypeError(f"VARIANT path is not DOUBLE: {path}") current = struct.unpack_from(' Date: Mon, 10 Aug 2026 03:30:01 -0700 Subject: [PATCH 04/23] [python] Reject duplicate VARIANT transform paths --- paimon-python/pypaimon/data/variant_path.py | 2 ++ paimon-python/pypaimon/tests/variant_path_test.py | 2 ++ 2 files changed, 4 insertions(+) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 486c477aa08d..762c136f2534 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -261,6 +261,8 @@ def variant_transform(column, transforms: Mapping[str, object]): (path, _parse_path(path), transform) for path, transform in transforms.items() ] + if len({path for _, path, _ in parsed}) != len(parsed): + raise ValueError("VARIANT transform paths must be unique") for path, _, transform in parsed: if not callable(transform): raise TypeError(f"VARIANT transform for {path} must be callable") diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index b9f57a7c2ae0..1e66df9b42f8 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -104,6 +104,8 @@ def test_rejects_invalid_transform(self): column = _variants([{'number': 1.0, 'text': 'value'}]) cases = [ ({'number': operator.neg}, ValueError, "Invalid VARIANT path"), + ({'$.number': operator.neg, "$['number']": operator.neg}, + ValueError, "paths must be unique"), ({'$.missing': operator.neg}, ValueError, "path does not exist"), ({'$.text': operator.neg}, TypeError, "path is not DOUBLE"), ({'$.number': 'negate'}, TypeError, "must be callable"), From c232052786af7da5998dc124080de0cd1dd1010b Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 05:28:21 -0700 Subject: [PATCH 05/23] [python] Support FLOAT VARIANT path transforms --- docs/docs/pypaimon/python-api.mdx | 7 ++-- paimon-python/pypaimon/data/variant_path.py | 26 ++++++++---- .../pypaimon/tests/variant_path_test.py | 42 ++++++++++++++++++- 3 files changed, 63 insertions(+), 12 deletions(-) diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index 699fe4a41f44..250e58035a03 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -1126,9 +1126,9 @@ Supported Paimon type strings for shredded sub-fields: `BOOLEAN`, `INT`, `BIGINT -### VARIANT DOUBLE Updates +### VARIANT Numeric Updates -Transform existing DOUBLE paths without decoding unrelated fields: +Transform existing FLOAT or DOUBLE paths without decoding unrelated fields: ```python import operator @@ -1141,7 +1141,8 @@ updated_payload = variant_transform(payload, { }) ``` -The paths must exist and contain DOUBLE values. The metadata is unchanged. +The paths must exist and contain FLOAT or DOUBLE values. Each transform must +return a `float`; the original numeric type and metadata are unchanged. **`GenericVariant` API:** diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 762c136f2534..713c32a78537 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Transform existing DOUBLE paths in Arrow VARIANT columns.""" +"""Transform existing FLOAT and DOUBLE paths in Arrow VARIANT columns.""" import functools import re @@ -256,7 +256,7 @@ def _variant_array(values, metadata, nulls, data_type): def variant_transform(column, transforms: Mapping[str, object]): - """Transform existing DOUBLE paths without decoding the whole VARIANT.""" + """Transform existing FLOAT and DOUBLE paths without full decoding.""" parsed = [ (path, _parse_path(path), transform) for path, transform in transforms.items() @@ -292,18 +292,28 @@ def variant_transform(column, transforms: Mapping[str, object]): for (path, _, transform), pos in zip(parsed, positions): if pos is None: raise ValueError(f"VARIANT path does not exist: {path}") - if _variant_get_type(value, pos) != _Type.DOUBLE: - raise TypeError(f"VARIANT path is not DOUBLE: {path}") - current = struct.unpack_from(' Date: Mon, 10 Aug 2026 06:26:22 -0700 Subject: [PATCH 06/23] [python] Add composable VARIANT get and replace APIs --- docs/docs/pypaimon/python-api.mdx | 23 +- paimon-python/pypaimon/data/__init__.py | 5 +- paimon-python/pypaimon/data/variant_path.py | 399 +++++++++++++++--- .../pypaimon/tests/variant_path_test.py | 271 ++++++++---- 4 files changed, 555 insertions(+), 143 deletions(-) diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index 250e58035a03..2e782adc9646 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -1126,23 +1126,26 @@ Supported Paimon type strings for shredded sub-fields: `BOOLEAN`, `INT`, `BIGINT -### VARIANT Numeric Updates +### VARIANT Path Updates -Transform existing FLOAT or DOUBLE paths without decoding unrelated fields: +Read a path as an Arrow array, use Arrow compute, and replace its existing +values without decoding unrelated fields: ```python -import operator +import pyarrow as pa +import pyarrow.compute as pc -from pypaimon.data import variant_transform +from pypaimon.data import variant_get, variant_replace -updated_payload = variant_transform(payload, { - '$.velocity.y': operator.neg, - '$.velocity.z': operator.neg, -}) +current = variant_get(payload, '$.velocity.y', pa.float64()) +updated_payload = variant_replace( + payload, '$.velocity.y', pc.negate(current)) ``` -The paths must exist and contain FLOAT or DOUBLE values. Each transform must -return a `float`; the original numeric type and metadata are unchanged. +`variant_get` returns NULL for a missing path. `variant_replace` accepts an +Arrow Scalar, Array, or ChunkedArray. A missing path is unchanged unless +`strict=True` is specified. Pass `{path: type}` to `variant_get` and +`{path: values}` to `variant_replace` to process multiple paths in one pass. **`GenericVariant` API:** diff --git a/paimon-python/pypaimon/data/__init__.py b/paimon-python/pypaimon/data/__init__.py index f30d8ab71d9b..88fc282dbd24 100644 --- a/paimon-python/pypaimon/data/__init__.py +++ b/paimon-python/pypaimon/data/__init__.py @@ -17,10 +17,11 @@ from pypaimon.data.timestamp import Timestamp from pypaimon.data.decimal import Decimal -from pypaimon.data.variant_path import variant_transform +from pypaimon.data.variant_path import variant_get, variant_replace __all__ = [ 'Timestamp', 'Decimal', - 'variant_transform', + 'variant_get', + 'variant_replace', ] diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 713c32a78537..596772c3e598 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Transform existing FLOAT and DOUBLE paths in Arrow VARIANT columns.""" +"""Read and replace paths in Arrow VARIANT columns.""" import functools import re @@ -27,9 +27,22 @@ _ARRAY, _OBJECT, _U32_SIZE, + _primitive_header, _read_unsigned, ) -from pypaimon.data.generic_variant import _Type, _variant_get_type +from pypaimon.data.generic_variant import ( + GenericVariant, + _DOUBLE, + _FLOAT, + _Type, + _value_size, + _variant_get_type, +) +from pypaimon.data.variant_shredding import ( + _build_array_value, + _build_object_value, + _encode_scalar_to_value_bytes, +) _INDEX_PATTERN = re.compile(r"\[(\d+)]") @@ -122,7 +135,7 @@ def _object_field_position( value: bytes, pos: int, key_id: int) -> Optional[int]: size, id_size, offset_size, id_start, offset_start, data_start = ( _object_layout(value, pos)) - id_table = value[id_start:id_start + size * id_size] + id_table = bytes(value[id_start:id_start + size * id_size]) slot = _field_slot(id_table, id_size, key_id) if slot is None: return None @@ -231,6 +244,61 @@ def _path_positions( return plan.positions +def _replace_path( + value: bytes, + metadata: bytes, + pos: int, + path: _Path, + replacement: bytes, +) -> bytes: + if not path: + return replacement + + kind, segment = path[0] + if kind == 'key': + if (value[pos] & 0x3) != _OBJECT: + raise ValueError("VARIANT path expects an object") + key_id = _metadata_key_ids(metadata).get(segment) + if key_id is None: + raise ValueError(f"VARIANT path does not exist: {segment}") + size, id_size, offset_size, id_start, offset_start, data_start = ( + _object_layout(value, pos)) + ids = [ + _read_unsigned(value, id_start + i * id_size, id_size) + for i in range(size) + ] + try: + slot = ids.index(key_id) + except ValueError: + raise ValueError(f"VARIANT path does not exist: {segment}") + children = [] + for i in range(size): + child_pos = data_start + _read_unsigned( + value, offset_start + i * offset_size, offset_size) + child = value[child_pos:child_pos + _value_size(value, child_pos)] + if i == slot: + child = _replace_path( + value, metadata, child_pos, path[1:], replacement) + children.append(child) + return _build_object_value(list(zip(ids, children))) + + if (value[pos] & 0x3) != _ARRAY: + raise ValueError("VARIANT path expects an array") + size, offset_size, offset_start, data_start = _array_layout(value, pos) + if segment >= size: + raise ValueError(f"VARIANT array index does not exist: {segment}") + children = [] + for i in range(size): + child_pos = data_start + _read_unsigned( + value, offset_start + i * offset_size, offset_size) + child = value[child_pos:child_pos + _value_size(value, child_pos)] + if i == segment: + child = _replace_path( + value, metadata, child_pos, path[1:], replacement) + children.append(child) + return _build_array_value(children) + + def _variant_chunks(column): if isinstance(column, pa.ChunkedArray): chunks, chunked, data_type = column.chunks, True, column.type @@ -244,81 +312,300 @@ def _variant_chunks(column): != ['value', 'metadata']): raise TypeError( "VARIANT input must contain value and metadata fields") + if not (pa.types.is_binary(data_type[0].type) + or pa.types.is_large_binary(data_type[0].type)): + raise TypeError("VARIANT value field must be binary") return chunks, chunked, data_type -def _variant_array(values, metadata, nulls, data_type): +class _BinaryValues: + + def __init__(self, array: pa.Array): + self.array = array + if pa.types.is_binary(array.type): + self.width, self.value_format = 4, ' Tuple[int, int]: + index = self.array.offset + row + return ( + struct.unpack_from( + self.value_format, self.offsets, index * self.width)[0], + struct.unpack_from( + self.value_format, self.offsets, + (index + 1) * self.width)[0], + ) + + def view(self, row: int) -> memoryview: + start, end = self.bounds(row) + return self.data[start:end] + + +def _decode_scalar(value, metadata: bytes, pos: int): + value_type = _variant_get_type(value, pos) + if value_type == _Type.DOUBLE: + return struct.unpack_from(' pa.StructArray: + values = chunk.field(0) + data_buffer = values.buffers()[2] + data = bytearray(data_buffer) if data_buffer is not None else bytearray() + for absolute_pos, replacement in patches: + data[absolute_pos:absolute_pos + len(replacement)] = replacement + + buffers = list(values.buffers()) + buffers[2] = pa.py_buffer(data) + patched_values = pa.Array.from_buffers( + values.type, + len(values), + buffers, + null_count=values.null_count, + offset=values.offset, + ) + metadata = chunk.field(1) + if chunk.offset == 0: + return pa.Array.from_buffers( + chunk.type, + len(chunk), + [chunk.buffers()[0]], + children=[patched_values, metadata], + null_count=chunk.null_count, + ) return pa.StructArray.from_arrays( - [pa.array(values, type=data_type[0].type), metadata], - fields=list(data_type), - mask=pa.array(nulls, type=pa.bool_()), + [patched_values, metadata], + fields=list(chunk.type), + mask=chunk.is_null(), ) -def variant_transform(column, transforms: Mapping[str, object]): - """Transform existing FLOAT and DOUBLE paths without full decoding.""" +def _rebuilt_chunk( + chunk: pa.StructArray, + values: Sequence[bytes], +) -> pa.StructArray: + return pa.StructArray.from_arrays( + [pa.array(values, type=chunk.type[0].type), chunk.field(1)], + fields=list(chunk.type), + mask=chunk.is_null(), + ) + + +class _Replacement: + + def __init__(self, value, length: int): + if isinstance(value, pa.Scalar): + self._value = value + self._array = None + self.type = value.type + elif isinstance(value, (pa.Array, pa.ChunkedArray)): + if len(value) != length: + raise ValueError( + "VARIANT replacement length must match the input column") + self._value = None + self._array = value + self.type = value.type + else: + raise TypeError( + "VARIANT replacement must be an Arrow Scalar or Array") + if not _supported_replacement_type(self.type): + raise TypeError( + f"Unsupported VARIANT replacement type: {self.type}") + if pa.types.is_float64(self.type): + self._value_format = ' bytes: + if value is not None and self._value_format is not None: + return struct.pack( + self._value_format, self._type_header, value) + return _encode_scalar_to_value_bytes(value, self.type) + + +def _supported_replacement_type(data_type: pa.DataType) -> bool: + return ( + pa.types.is_null(data_type) + or pa.types.is_boolean(data_type) + or pa.types.is_signed_integer(data_type) + or pa.types.is_float32(data_type) + or pa.types.is_float64(data_type) + or pa.types.is_string(data_type) + or pa.types.is_large_string(data_type) + or pa.types.is_binary(data_type) + or pa.types.is_large_binary(data_type) + or pa.types.is_date32(data_type) + or pa.types.is_timestamp(data_type) + or pa.types.is_decimal128(data_type) + ) + + +def _variant_get(column, paths: Mapping[str, pa.DataType]): + parsed = [] + for path, target_type in paths.items(): + if not isinstance(target_type, pa.DataType): + raise TypeError("VARIANT target_type must be a PyArrow data type") + parsed.append((path, _parse_path(path), target_type)) + parsed_paths = [parsed_path for _, parsed_path, _ in parsed] + chunks, chunked, _ = _variant_chunks(column) + plans = [] + result_chunks = {path: [] for path in paths} + for chunk in chunks: + values = _BinaryValues(chunk.field(0)) + metadata = chunk.field(1).to_pylist() + valid = chunk.is_valid().to_pylist() + results = {path: [] for path in paths} + for row in range(len(chunk)): + if not valid[row]: + for path in paths: + results[path].append(None) + continue + value = values.view(row) + positions = _path_positions( + value, + metadata[row], + parsed_paths, + plans, + ) + for (path, _, _), pos in zip(parsed, positions): + results[path].append( + None if pos is None + else _decode_scalar(value, metadata[row], pos) + ) + for path, _, target_type in parsed: + result_chunks[path].append( + pa.array(results[path], type=target_type)) + if not chunked: + return {path: chunks[0] for path, chunks in result_chunks.items()} + return { + path: pa.chunked_array(chunks, type=paths[path]) + for path, chunks in result_chunks.items() + } + + +def variant_get(column, path, target_type=None): + """Read one or more VARIANT paths into Arrow arrays.""" + if isinstance(path, Mapping): + if target_type is not None: + raise TypeError( + "VARIANT target_type must be omitted for path mappings") + return _variant_get(column, path) + if target_type is None: + raise TypeError("VARIANT target_type must be a PyArrow data type") + return _variant_get(column, {path: target_type})[path] + + +def _paths_overlap(first: _Path, second: _Path) -> bool: + limit = min(len(first), len(second)) + return first[:limit] == second[:limit] + + +def _validate_distinct_paths(parsed) -> None: + for index, (_, first, _) in enumerate(parsed): + for _, second, _ in parsed[index + 1:]: + if _paths_overlap(first, second): + raise ValueError( + "VARIANT replacement paths must not overlap") + + +def variant_replace( + column, + path, + replacement=None, + strict: bool = False, +): + """Replace one or more existing VARIANT paths with Arrow values.""" + if not isinstance(strict, bool): + raise TypeError("VARIANT strict must be a boolean") + if isinstance(path, Mapping): + if replacement is not None: + raise TypeError( + "VARIANT replacement must be omitted for path mappings") + replacements = path + else: + replacements = {path: replacement} parsed = [ - (path, _parse_path(path), transform) - for path, transform in transforms.items() + (path, _parse_path(path), _Replacement(value, len(column))) + for path, value in replacements.items() ] - if len({path for _, path, _ in parsed}) != len(parsed): - raise ValueError("VARIANT transform paths must be unique") - for path, _, transform in parsed: - if not callable(transform): - raise TypeError(f"VARIANT transform for {path} must be callable") + _validate_distinct_paths(parsed) if not parsed: return column + parsed_paths = [parsed_path for _, parsed_path, _ in parsed] chunks, chunked, data_type = _variant_chunks(column) - paths = [path for _, path, _ in parsed] plans = [] result_chunks = [] + global_row = 0 for chunk in chunks: - input_values = chunk.field(0).to_pylist() - metadata_array = chunk.field(1) - input_metadatas = metadata_array.to_pylist() + values = _BinaryValues(chunk.field(0)) + metadata = chunk.field(1).to_pylist() valid = chunk.is_valid().to_pylist() - values, nulls = [], [] + chunk_replacements = { + path: provider.values(global_row, len(chunk)) + for path, _, provider in parsed + } + row_replacements = [[] for _ in range(len(chunk))] + patches = [] + rebuild = False for row in range(len(chunk)): if not valid[row]: - values.append(b'') - nulls.append(True) continue - - value = input_values[row] - metadata = input_metadatas[row] - result = bytearray(value) - positions = _path_positions(value, metadata, paths, plans) - for (path, _, transform), pos in zip(parsed, positions): + value = values.view(row) + row_start, _ = values.bounds(row) + positions = _path_positions( + value, + metadata[row], + parsed_paths, + plans, + ) + for (path, parsed_path, provider), pos in zip( + parsed, positions): if pos is None: - raise ValueError(f"VARIANT path does not exist: {path}") - value_type = _variant_get_type(value, pos) - if value_type == _Type.DOUBLE: - value_format, type_name = ' Date: Mon, 10 Aug 2026 06:53:16 -0700 Subject: [PATCH 07/23] [python] Bound memory for VARIANT replacements --- paimon-python/pypaimon/data/variant_path.py | 153 ++++++++++++++---- .../pypaimon/tests/variant_path_test.py | 35 ++++ 2 files changed, 157 insertions(+), 31 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 596772c3e598..f205081f312a 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -343,10 +343,62 @@ def bounds(self, row: int) -> Tuple[int, int]: (index + 1) * self.width)[0], ) + def used_bounds(self) -> Tuple[int, int]: + first = self.array.offset + last = first + len(self.array) + return ( + struct.unpack_from( + self.value_format, self.offsets, first * self.width)[0], + struct.unpack_from( + self.value_format, self.offsets, last * self.width)[0], + ) + def view(self, row: int) -> memoryview: start, end = self.bounds(row) return self.data[start:end] + def row(self, row: int) -> Tuple[int, memoryview]: + start, end = self.bounds(row) + return start, self.data[start:end] + + def copy_used_data(self) -> Tuple[bytearray, int]: + start, end = self.used_bounds() + return bytearray(self.data[start:end]), start + + def array_from_data(self, data: bytearray, start: int) -> pa.Array: + buffers = list(self.array.buffers()) + offset = self.array.offset + if offset == 0 and start == 0: + buffers[2] = pa.py_buffer(data) + else: + offsets = bytearray((len(self.array) + 1) * self.width) + for index in range(len(self.array) + 1): + value = struct.unpack_from( + self.value_format, + self.offsets, + (self.array.offset + index) * self.width, + )[0] + struct.pack_into( + self.value_format, + offsets, + index * self.width, + value - start, + ) + buffers = [ + (None if self.array.null_count == 0 + else self.array.is_valid().buffers()[1]), + pa.py_buffer(offsets), + pa.py_buffer(data), + ] + offset = 0 + return pa.Array.from_buffers( + self.array.type, + len(self.array), + buffers, + null_count=self.array.null_count, + offset=offset, + ) + def _decode_scalar(value, metadata: bytes, pos: int): value_type = _variant_get_type(value, pos) @@ -357,24 +409,15 @@ def _decode_scalar(value, metadata: bytes, pos: int): return GenericVariant(bytes(value), metadata, pos).to_python() -def _patched_chunk(chunk: pa.StructArray, patches) -> pa.StructArray: - values = chunk.field(0) - data_buffer = values.buffers()[2] - data = bytearray(data_buffer) if data_buffer is not None else bytearray() - for absolute_pos, replacement in patches: - data[absolute_pos:absolute_pos + len(replacement)] = replacement - - buffers = list(values.buffers()) - buffers[2] = pa.py_buffer(data) - patched_values = pa.Array.from_buffers( - values.type, - len(values), - buffers, - null_count=values.null_count, - offset=values.offset, - ) +def _patched_chunk( + chunk: pa.StructArray, + values: _BinaryValues, + data: bytearray, + start: int, +) -> pa.StructArray: + patched_values = values.array_from_data(data, start) metadata = chunk.field(1) - if chunk.offset == 0: + if chunk.offset == 0 and patched_values.offset == 0: return pa.Array.from_buffers( chunk.type, len(chunk), @@ -429,12 +472,26 @@ def __init__(self, value, length: int): else: self._value_format = None self._type_header = None + self._fixed_size = ( + struct.calcsize(self._value_format) + if self._value_format is not None else None + ) def values(self, offset: int, length: int): if self._array is None: - return [self._value.as_py()] * length + return self._value.as_py() return self._array.slice(offset, length).to_pylist() + def value_at(self, values, row: int): + return values if self._array is None else values[row] + + def fixed_size(self, value) -> Optional[int]: + return self._fixed_size if value is not None else None + + def patch(self, data: bytearray, pos: int, value) -> None: + struct.pack_into( + self._value_format, data, pos, self._type_header, value) + def encode(self, value) -> bytes: if value is not None and self._value_format is not None: return struct.pack( @@ -564,14 +621,13 @@ def variant_replace( path: provider.values(global_row, len(chunk)) for path, _, provider in parsed } - row_replacements = [[] for _ in range(len(chunk))] - patches = [] + patched_data = None + data_start = 0 rebuild = False for row in range(len(chunk)): if not valid[row]: continue - value = values.view(row) - row_start, _ = values.bounds(row) + row_start, value = values.row(row) positions = _path_positions( value, metadata[row], @@ -585,24 +641,59 @@ def variant_replace( raise ValueError( f"VARIANT path does not exist: {path}") continue - encoded = provider.encode(chunk_replacements[path][row]) - row_replacements[row].append((parsed_path, encoded)) - if len(encoded) == _value_size(value, pos): - patches.append((row_start + pos, encoded)) - else: + replacement_value = provider.value_at( + chunk_replacements[path], row) + new_size = provider.fixed_size(replacement_value) + encoded = ( + None if new_size is not None + else provider.encode(replacement_value) + ) + if new_size is None: + new_size = len(encoded) + if new_size != _value_size(value, pos): rebuild = True + break + if patched_data is None: + patched_data, data_start = values.copy_used_data() + patch_pos = row_start - data_start + pos + if encoded is None: + provider.patch( + patched_data, patch_pos, replacement_value) + else: + patched_data[patch_pos:patch_pos + new_size] = encoded + if rebuild: + break if rebuild: rebuilt_values = [] - for row, replacements_for_row in enumerate(row_replacements): + for row in range(len(chunk)): value = bytes(values.view(row)) - for parsed_path, encoded in replacements_for_row: + if not valid[row]: + rebuilt_values.append(value) + continue + positions = _path_positions( + value, + metadata[row], + parsed_paths, + plans, + ) + for (path, parsed_path, provider), pos in zip( + parsed, positions): + if pos is None: + if strict: + raise ValueError( + f"VARIANT path does not exist: {path}") + continue + replacement_value = provider.value_at( + chunk_replacements[path], row) + encoded = provider.encode(replacement_value) value = _replace_path( value, metadata[row], 0, parsed_path, encoded) rebuilt_values.append(value) result_chunks.append(_rebuilt_chunk(chunk, rebuilt_values)) - elif patches: - result_chunks.append(_patched_chunk(chunk, patches)) + elif patched_data is not None: + result_chunks.append(_patched_chunk( + chunk, values, patched_data, data_start)) else: result_chunks.append(chunk) global_row += len(chunk) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index f6ca49a65437..cd01d0b9fdd9 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -136,6 +136,41 @@ def test_equal_length_uses_copy_on_write(self): result.buffers()[0].address, ) + def test_sliced_input_copies_only_visible_values(self): + base = _variants([ + {'number': float(i), 'padding': 'x' * 1000} + for i in range(100) + ]) + + for binary_type in (pa.binary(), pa.large_binary()): + with self.subTest(binary_type=binary_type): + values = base.field('value').cast(binary_type) + converted = pa.StructArray.from_arrays( + [values, base.field('metadata')], + names=['value', 'metadata'], + ) + column = converted.slice(50, 3) + + result = variant_replace( + column, + '$.number', + pa.scalar(-1.0, type=pa.float64()), + ) + + expected_size = sum( + len(value) + for value in column.field('value').to_pylist() + ) + self.assertEqual( + result.field('value').buffers()[2].size, + expected_size, + ) + self.assertEqual(result.field('value').offset, 0) + self.assertEqual( + [row['number'] for row in _decode(result)], + [-1.0, -1.0, -1.0], + ) + def test_get_compute_replace_pipeline(self): column = pa.chunked_array([ _variants([{'y': 1.0, 'z': -2.0}, None]), From 0bfc056c4c6528baa1f02b2beda7c645111b3186 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 07:09:30 -0700 Subject: [PATCH 08/23] [python] Vectorize numeric VARIANT path updates --- paimon-python/pypaimon/data/variant_path.py | 266 ++++++++++++++++++ .../pypaimon/tests/variant_path_test.py | 66 +++++ 2 files changed, 332 insertions(+) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index f205081f312a..1fa6eb321f1d 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -21,6 +21,7 @@ import struct from typing import Dict, List, Mapping, Optional, Sequence, Tuple +import numpy as np import pyarrow as pa from pypaimon.data._variant_binary import ( @@ -343,6 +344,14 @@ def bounds(self, row: int) -> Tuple[int, int]: (index + 1) * self.width)[0], ) + def numpy_offsets(self): + return np.frombuffer( + self.offsets, + dtype=np.dtype(self.value_format), + count=len(self.array) + 1, + offset=self.array.offset * self.width, + ).astype(np.int64, copy=False) + def used_bounds(self) -> Tuple[int, int]: first = self.array.offset last = first + len(self.array) @@ -400,6 +409,177 @@ def array_from_data(self, data: bytearray, start: int) -> pa.Array: ) +def _take_unsigned(data, positions, widths): + result = np.empty(len(positions), dtype=np.int64) + for width in range(1, 5): + selected = widths == width + if not np.any(selected): + continue + selected_positions = positions[selected] + if (np.any(selected_positions < 0) + or np.any(selected_positions + width > len(data))): + raise ValueError("Invalid VARIANT offset") + indices = selected_positions[:, None] + np.arange(width) + values = data[indices].astype(np.int64, copy=False) + result[selected] = np.sum( + values << (np.arange(width, dtype=np.int64) * 8), axis=1) + return result + + +def _all_binary_values_equal(values: _BinaryValues, expected: bytes) -> bool: + offsets = values.numpy_offsets() + lengths = offsets[1:] - offsets[:-1] + if np.any(lengths != len(expected)): + return False + if not len(offsets) > 1 or not expected: + return True + + data = np.frombuffer(values.data, dtype=np.uint8) + expected_array = np.frombuffer(expected, dtype=np.uint8) + rows_per_batch = max(1, (1024 * 1024) // len(expected)) + for row in range(0, len(lengths), rows_per_batch): + end = min(row + rows_per_batch, len(lengths)) + starts = offsets[row:end] + indices = starts[:, None] + np.arange(len(expected)) + if not np.all(data[indices] == expected_array): + return False + return True + + +def _vectorized_path_positions( + values: _BinaryValues, + metadata: pa.Array, + valid_count: int, + paths: Sequence[_Path], +): + if (valid_count != len(values.array) + or values.array.null_count + or metadata.null_count + or not len(values.array)): + return None + + metadata_values = _BinaryValues(metadata) + first_metadata = bytes(metadata_values.view(0)) + if not _all_binary_values_equal(metadata_values, first_metadata): + return None + + key_ids = _metadata_key_ids(first_metadata) + nodes, result_nodes = _compile_paths(tuple(paths)) + row_offsets = values.numpy_offsets() + row_starts = row_offsets[:-1] + row_ends = row_offsets[1:] + data = np.frombuffer(values.data, dtype=np.uint8) + first_value = values.view(0) + positions = [np.zeros(len(values.array), dtype=np.int64)] + + try: + for parent_node, kind, segment in nodes[1:]: + parent = positions[parent_node] + absolute_parent = row_starts + parent + if (np.any(absolute_parent < row_starts) + or np.any(absolute_parent >= row_ends)): + return None + headers = data[absolute_parent] + type_info = (headers >> 2).astype(np.int64, copy=False) + + if kind == 'key': + if np.any((headers & 0x3) != _OBJECT): + return None + key_id = key_ids.get(segment) + if key_id is None: + return None + first_layout = _object_layout( + first_value, int(parent[0])) + size, id_size, _, id_start, _, _ = first_layout + id_table = bytes( + first_value[id_start:id_start + size * id_size]) + slot = _field_slot(id_table, id_size, key_id) + if slot is None: + return None + + size_widths = np.where( + ((type_info >> 4) & 0x1) != 0, _U32_SIZE, 1) + sizes = _take_unsigned( + data, absolute_parent + 1, size_widths) + id_widths = ((type_info >> 2) & 0x3) + 1 + offset_widths = (type_info & 0x3) + 1 + if np.any(sizes <= slot): + return None + id_starts = absolute_parent + 1 + size_widths + ids = _take_unsigned( + data, id_starts + slot * id_widths, id_widths) + if np.any(ids != key_id): + return None + offset_starts = id_starts + sizes * id_widths + offsets = _take_unsigned( + data, + offset_starts + slot * offset_widths, + offset_widths, + ) + data_starts = offset_starts + (sizes + 1) * offset_widths + child = data_starts + offsets - row_starts + else: + if np.any((headers & 0x3) != _ARRAY): + return None + size_widths = np.where( + ((type_info >> 2) & 0x1) != 0, _U32_SIZE, 1) + sizes = _take_unsigned( + data, absolute_parent + 1, size_widths) + if np.any(sizes <= segment): + return None + offset_widths = (type_info & 0x3) + 1 + offset_starts = absolute_parent + 1 + size_widths + offsets = _take_unsigned( + data, + offset_starts + segment * offset_widths, + offset_widths, + ) + data_starts = offset_starts + (sizes + 1) * offset_widths + child = data_starts + offsets - row_starts + + if np.any(child < 0) or np.any(row_starts + child >= row_ends): + return None + positions.append(child) + except (IndexError, ValueError): + return None + + return ( + row_starts, + row_ends, + data, + tuple(positions[node] for node in result_nodes), + ) + + +def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): + planned = _vectorized_path_positions( + values, chunk.field(1), len(chunk) - chunk.null_count, + parsed_paths) + if planned is None: + return None + row_starts, row_ends, data, positions = planned + results = [] + for pos, target_type in zip(positions, target_types): + if not (pa.types.is_float32(target_type) + or pa.types.is_float64(target_type)): + return None + absolute = row_starts + pos + headers = data[absolute] + if np.all(headers == _primitive_header(_DOUBLE)): + value_size, data_type = 8, np.dtype(' row_ends): + return None + indices = absolute[:, None] + 1 + np.arange(value_size) + raw = np.ascontiguousarray(data[indices]) + result = raw.view(data_type).reshape(-1) + results.append(pa.array(result, type=target_type)) + return results + + def _decode_scalar(value, metadata: bytes, pos: int): value_type = _variant_get_type(value, pos) if value_type == _Type.DOUBLE: @@ -492,6 +672,26 @@ def patch(self, data: bytearray, pos: int, value) -> None: struct.pack_into( self._value_format, data, pos, self._type_header, value) + def numpy_values(self, offset: int, length: int): + if self._value_format is None: + return None + data_type = ( + np.dtype(' bytes: if value is not None and self._value_format is not None: return struct.pack( @@ -499,6 +699,51 @@ def encode(self, value) -> bytes: return _encode_scalar_to_value_bytes(value, self.type) +def _vectorized_replace_chunk( + chunk, + values, + parsed, + parsed_paths, + global_row, +): + planned = _vectorized_path_positions( + values, + chunk.field(1), + len(chunk) - chunk.null_count, + parsed_paths, + ) + if planned is None: + return None + row_starts, row_ends, input_data, positions = planned + + replacements = [] + for (_, _, provider), pos in zip(parsed, positions): + replacement = provider.numpy_values(global_row, len(chunk)) + if replacement is None: + return None + absolute = row_starts + pos + expected_header = provider._type_header + if not np.all(input_data[absolute] == expected_header): + return None + if np.any(absolute + provider._fixed_size > row_ends): + return None + replacements.append((pos, provider, replacement)) + + data, data_start = values.copy_used_data() + output_data = np.frombuffer(data, dtype=np.uint8) + relative_starts = row_starts - data_start + for pos, provider, replacement in replacements: + absolute = relative_starts + pos + output_data[absolute] = provider._type_header + value_size = provider._fixed_size - 1 + replacement_bytes = np.ascontiguousarray( + replacement).view(np.uint8).reshape(len(chunk), value_size) + indices = absolute[:, None] + 1 + np.arange(value_size) + output_data[indices] = replacement_bytes + + return _patched_chunk(chunk, values, data, data_start) + + def _supported_replacement_type(data_type: pa.DataType) -> bool: return ( pa.types.is_null(data_type) @@ -528,6 +773,16 @@ def _variant_get(column, paths: Mapping[str, pa.DataType]): result_chunks = {path: [] for path in paths} for chunk in chunks: values = _BinaryValues(chunk.field(0)) + vectorized = _vectorized_get_chunk( + chunk, + values, + parsed_paths, + [target_type for _, _, target_type in parsed], + ) + if vectorized is not None: + for (path, _, _), result in zip(parsed, vectorized): + result_chunks[path].append(result) + continue metadata = chunk.field(1).to_pylist() valid = chunk.is_valid().to_pylist() results = {path: [] for path in paths} @@ -615,6 +870,17 @@ def variant_replace( global_row = 0 for chunk in chunks: values = _BinaryValues(chunk.field(0)) + vectorized = _vectorized_replace_chunk( + chunk, + values, + parsed, + parsed_paths, + global_row, + ) + if vectorized is not None: + result_chunks.append(vectorized) + global_row += len(chunk) + continue metadata = chunk.field(1).to_pylist() valid = chunk.is_valid().to_pylist() chunk_replacements = { diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index cd01d0b9fdd9..03cce7cf52bf 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -193,6 +193,72 @@ def test_get_compute_replace_pipeline(self): {'y': 3.0, 'z': -4.0}, ]) + def test_vectorized_paths_support_varying_offsets(self): + column = _variants([ + { + 'prefix': 'x' * index, + 'items': ['y' * (64 - index), { + 'value': float(index), + 'other': float(-index), + }], + } + for index in range(1, 64) + ]) + paths = { + '$.items[1].value': pa.float64(), + '$.items[1].other': pa.float64(), + } + + with patch( + 'pypaimon.data.variant_path._path_positions', + side_effect=AssertionError("slow path is not allowed")): + current = variant_get(column, paths) + result = variant_replace(column, { + path: pc.negate(values) + for path, values in current.items() + }) + + decoded = _decode(result) + self.assertEqual( + [row['items'][1]['value'] for row in decoded], + [float(-index) for index in range(1, 64)], + ) + self.assertEqual( + [row['items'][1]['other'] for row in decoded], + [float(index) for index in range(1, 64)], + ) + + def test_vectorized_paths_support_wide_containers(self): + rows = [] + for row in range(3): + value = {'field_%03d' % i: i for i in range(300)} + value['field_299'] = float(row) + value['items'] = list(range(299)) + [float(row + 10)] + rows.append(value) + column = _variants(rows) + paths = { + '$.field_299': pa.float64(), + '$.items[299]': pa.float64(), + } + + with patch( + 'pypaimon.data.variant_path._path_positions', + side_effect=AssertionError("slow path is not allowed")): + current = variant_get(column, paths) + result = variant_replace(column, { + path: pc.negate(values) + for path, values in current.items() + }) + + self.assertEqual(current['$.field_299'].to_pylist(), [0.0, 1.0, 2.0]) + self.assertEqual( + current['$.items[299]'].to_pylist(), [10.0, 11.0, 12.0]) + decoded = _decode(result) + self.assertEqual( + [row['field_299'] for row in decoded], [0.0, -1.0, -2.0]) + self.assertEqual( + [row['items'][299] for row in decoded], [-10.0, -11.0, -12.0]) + def test_scalar_and_same_length_string_replacement(self): column = _variants([{'text': 'aa'}, {'text': 'bb'}]) From 5a5f8ac64026429a0d7f45d51e55c852f31051e8 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 07:32:20 -0700 Subject: [PATCH 09/23] [python] Harden VARIANT path access and casts --- .../pypaimon/data/generic_variant.py | 31 +- paimon-python/pypaimon/data/variant_path.py | 555 +++++++++++++----- .../pypaimon/data/variant_shredding.py | 30 +- .../pypaimon/tests/variant_path_test.py | 102 ++++ 4 files changed, 544 insertions(+), 174 deletions(-) diff --git a/paimon-python/pypaimon/data/generic_variant.py b/paimon-python/pypaimon/data/generic_variant.py index 9d0e0a0b0c06..a7edcd048848 100644 --- a/paimon-python/pypaimon/data/generic_variant.py +++ b/paimon-python/pypaimon/data/generic_variant.py @@ -148,6 +148,12 @@ def _write_le(buf, pos, value, n): buf[pos:pos + n] = value.to_bytes(n, 'little') +def _decimal_from_unscaled(unscaled, scale): + sign = 1 if unscaled < 0 else 0 + digits = tuple(int(digit) for digit in str(abs(unscaled))) or (0,) + return _decimal.Decimal((sign, digits, -scale)) + + def _short_str_header(size): return (size << 2) | _SHORT_STR @@ -306,18 +312,25 @@ def append_float(self, f): self._pos += 4 def append_decimal(self, d): - d = d.normalize() sign, digits, exponent = d.as_tuple() - if exponent > 0: - raise ValueError( - f'append_decimal requires a non-positive exponent (got {d!r}); ' - 'use append_double() for Decimal values with positive exponents' - ) unscaled = int(''.join(str(x) for x in digits)) if sign: unscaled = -unscaled - scale = -exponent if exponent < 0 else 0 - precision = len(digits) + if exponent > 0: + unscaled *= 10 ** exponent + scale = 0 + else: + scale = -exponent + self.append_decimal_unscaled( + unscaled, max(1, len(str(abs(unscaled)))), scale) + + def append_decimal_unscaled(self, unscaled, precision, scale): + if not 0 <= scale <= _MAX_DECIMAL16_PRECISION: + raise ValueError(f'Unsupported VARIANT decimal scale: {scale}') + if (not 0 < precision <= _MAX_DECIMAL16_PRECISION + or not -(1 << 127) <= unscaled < (1 << 127)): + raise ValueError( + f'Unsupported VARIANT decimal precision: {precision}') if scale <= _MAX_DECIMAL4_PRECISION and precision <= _MAX_DECIMAL4_PRECISION: self._write_byte(_primitive_header(_DECIMAL4)) @@ -668,7 +681,7 @@ def _to_python_impl(self, value, metadata, pos): else: raw = bytes(value[pos + 2:pos + 18]) unscaled = int.from_bytes(raw, 'little', signed=True) - return _decimal.Decimal(unscaled) / (_decimal.Decimal(10) ** scale) + return _decimal_from_unscaled(unscaled, scale) if vtype == _Type.STRING: if basic_type == _SHORT_STR: return value[pos + 1:pos + 1 + type_info].decode('utf-8') diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 1fa6eb321f1d..0e9b6349687a 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -16,10 +16,15 @@ """Read and replace paths in Arrow VARIANT columns.""" +import base64 +import datetime +import decimal import functools +import json +import math import re import struct -from typing import Dict, List, Mapping, Optional, Sequence, Tuple +from typing import Dict, Mapping, Optional, Sequence, Tuple import numpy as np import pyarrow as pa @@ -27,17 +32,18 @@ from pypaimon.data._variant_binary import ( _ARRAY, _OBJECT, + _SHORT_STR, _U32_SIZE, _primitive_header, _read_unsigned, ) from pypaimon.data.generic_variant import ( + _BINARY, GenericVariant, _DOUBLE, _FLOAT, - _Type, - _value_size, - _variant_get_type, + _LONG_STR, + _PRIMITIVE_FIXED_SIZES, ) from pypaimon.data.variant_shredding import ( _build_array_value, @@ -49,8 +55,6 @@ _INDEX_PATTERN = re.compile(r"\[(\d+)]") _KEY_PATTERN = re.compile(r"\.([^\.\[]+)|\['([^']+)']|\[\"([^\"]+)\"]") _Path = Tuple[Tuple[str, object], ...] -_ObjectLayout = Tuple[int, int, int, int, int, int] -_ArrayLayout = Tuple[int, int, int, int] @functools.lru_cache(maxsize=256) @@ -96,32 +100,93 @@ def _metadata_key_ids(metadata: bytes) -> Dict[str, int]: return result -def _object_layout(value: bytes, pos: int) -> _ObjectLayout: - header = value[pos] - if (header & 0x3) != _OBJECT: - raise ValueError("VARIANT path expects an object") - type_info = (header >> 2) & 0x3F - size_bytes = _U32_SIZE if ((type_info >> 4) & 0x1) else 1 - size = _read_unsigned(value, pos + 1, size_bytes) - id_size = ((type_info >> 2) & 0x3) + 1 - offset_size = (type_info & 0x3) + 1 - id_start = pos + 1 + size_bytes - offset_start = id_start + size * id_size - data_start = offset_start + (size + 1) * offset_size - return size, id_size, offset_size, id_start, offset_start, data_start +def _malformed(message): + raise ValueError(f"MALFORMED_VARIANT: {message}") + + +def _require_range(pos, size, limit): + if pos < 0 or size < 0 or pos + size > limit: + _malformed("value is truncated") + + +def _checked_object_layout(value, pos, limit): + _require_range(pos, 2, limit) + type_info = (value[pos] >> 2) & 0x3F + size_width = _U32_SIZE if ((type_info >> 4) & 0x1) else 1 + _require_range(pos + 1, size_width, limit) + size = _read_unsigned(value, pos + 1, size_width) + id_width = ((type_info >> 2) & 0x3) + 1 + offset_width = (type_info & 0x3) + 1 + id_start = pos + 1 + size_width + offset_start = id_start + size * id_width + data_start = offset_start + (size + 1) * offset_width + _require_range(pos, data_start - pos, limit) + offsets = [] + for index in range(size + 1): + offset = _read_unsigned( + value, offset_start + index * offset_width, offset_width) + offsets.append(offset) + sentinel = offsets[-1] + if ((size and (min(offsets[:-1]) != 0 + or len(set(offsets[:-1])) != size)) + or any(offset >= sentinel for offset in offsets[:-1])): + _malformed("invalid object offsets") + _require_range(data_start, sentinel, limit) + return ( + size, id_width, id_start, data_start, offsets, + data_start + offsets[-1], + ) -def _array_layout(value: bytes, pos: int) -> _ArrayLayout: +def _checked_array_layout(value, pos, limit): + _require_range(pos, 2, limit) + type_info = (value[pos] >> 2) & 0x3F + size_width = _U32_SIZE if ((type_info >> 2) & 0x1) else 1 + _require_range(pos + 1, size_width, limit) + size = _read_unsigned(value, pos + 1, size_width) + offset_width = (type_info & 0x3) + 1 + offset_start = pos + 1 + size_width + data_start = offset_start + (size + 1) * offset_width + _require_range(pos, data_start - pos, limit) + offsets = [] + previous = 0 + for index in range(size + 1): + offset = _read_unsigned( + value, offset_start + index * offset_width, offset_width) + if (index == 0 and offset != 0) or offset < previous: + _malformed("invalid array offsets") + offsets.append(offset) + previous = offset + _require_range(data_start, offsets[-1], limit) + return size, data_start, offsets, data_start + offsets[-1] + + +def _checked_value_size(value, pos, limit=None): + limit = len(value) if limit is None else limit + _require_range(pos, 1, limit) header = value[pos] - if (header & 0x3) != _ARRAY: - raise ValueError("VARIANT path expects an array") + basic_type = header & 0x3 type_info = (header >> 2) & 0x3F - size_bytes = _U32_SIZE if ((type_info >> 2) & 0x1) else 1 - size = _read_unsigned(value, pos + 1, size_bytes) - offset_size = (type_info & 0x3) + 1 - offset_start = pos + 1 + size_bytes - data_start = offset_start + (size + 1) * offset_size - return size, offset_size, offset_start, data_start + if basic_type == _OBJECT: + end = _checked_object_layout(value, pos, limit)[-1] + elif basic_type == _ARRAY: + end = _checked_array_layout(value, pos, limit)[-1] + elif basic_type == _SHORT_STR: + end = pos + 1 + type_info + else: + fixed_size = _PRIMITIVE_FIXED_SIZES.get(type_info) + if fixed_size is not None: + end = pos + fixed_size + elif type_info in (_BINARY, _LONG_STR): + _require_range(pos + 1, _U32_SIZE, limit) + end = ( + pos + 1 + _U32_SIZE + + _read_unsigned(value, pos + 1, _U32_SIZE) + ) + else: + _malformed(f"unknown primitive type {type_info}") + _require_range(pos, end - pos, limit) + return end - pos @functools.lru_cache(maxsize=2048) @@ -132,29 +197,6 @@ def _field_slot(id_table: bytes, id_size: int, key_id: int) -> Optional[int]: return None -def _object_field_position( - value: bytes, pos: int, key_id: int) -> Optional[int]: - size, id_size, offset_size, id_start, offset_start, data_start = ( - _object_layout(value, pos)) - id_table = bytes(value[id_start:id_start + size * id_size]) - slot = _field_slot(id_table, id_size, key_id) - if slot is None: - return None - offset = _read_unsigned( - value, offset_start + slot * offset_size, offset_size) - return data_start + offset - - -def _array_element_position( - value: bytes, pos: int, index: int) -> Optional[int]: - size, offset_size, offset_start, data_start = _array_layout(value, pos) - if index >= size: - return None - offset = _read_unsigned( - value, offset_start + index * offset_size, offset_size) - return data_start + offset - - @functools.lru_cache(maxsize=256) def _compile_paths(paths: Tuple[_Path, ...]): nodes = [(None, None, None)] @@ -170,88 +212,74 @@ def _compile_paths(paths: Tuple[_Path, ...]): return tuple(nodes), tuple(results) -class _PositionPlan: - - def __init__(self, metadata, checks, positions): - self.metadata = metadata - self.checks = checks - self.positions = positions - - def matches(self, value: bytes, metadata: bytes) -> bool: - return metadata == self.metadata and all( - value[pos:pos + len(expected)] == expected - for pos, expected in self.checks - ) - - -def _position_plan( +def _path_positions( value: bytes, metadata: bytes, paths: Sequence[_Path], -) -> _PositionPlan: + plans, +) -> Sequence[Optional[int]]: + del plans + root_size = _checked_value_size(value, 0) + if root_size != len(value): + _malformed("trailing bytes after root value") key_ids = _metadata_key_ids(metadata) nodes, result_nodes = _compile_paths(tuple(paths)) - positions = [0] - checks = [] - checked = set() + bounds = [(0, len(value))] for parent_node, kind, segment in nodes[1:]: - parent = positions[parent_node] + parent = bounds[parent_node] if parent is None: - positions.append(None) + bounds.append(None) continue - if parent not in checked: - basic_type = value[parent] & 0x3 - if basic_type == _OBJECT: - data_start = _object_layout(value, parent)[-1] - elif basic_type == _ARRAY: - data_start = _array_layout(value, parent)[-1] - else: - data_start = parent + 1 - checks.append((parent, value[parent:data_start])) - checked.add(parent) + parent_pos, parent_end = parent + basic_type = value[parent_pos] & 0x3 if kind == 'key': key_id = key_ids.get(segment) - if key_id is None or (value[parent] & 0x3) != _OBJECT: - positions.append(None) - else: - positions.append(_object_field_position( - value, parent, key_id)) - elif (value[parent] & 0x3) != _ARRAY: - positions.append(None) + if key_id is None or basic_type != _OBJECT: + bounds.append(None) + continue + size, id_width, id_start, data_start, offsets, _ = ( + _checked_object_layout(value, parent_pos, parent_end)) + id_table = bytes(value[id_start:id_start + size * id_width]) + slot = _field_slot(id_table, id_width, key_id) + if slot is None: + bounds.append(None) + continue + child_start = data_start + offsets[slot] + child_size = _checked_value_size( + value, child_start, data_start + offsets[-1]) + child_end = child_start + child_size else: - positions.append(_array_element_position( - value, parent, segment)) - return _PositionPlan( - metadata, - tuple(checks), - tuple(positions[node] for node in result_nodes), + if basic_type != _ARRAY: + bounds.append(None) + continue + size, data_start, offsets, _ = _checked_array_layout( + value, parent_pos, parent_end) + if segment >= size: + bounds.append(None) + continue + slot = segment + child_start = data_start + offsets[slot] + child_end = data_start + offsets[slot + 1] + if _checked_value_size(value, child_start, child_end) != ( + child_end - child_start): + _malformed("child size does not match container offsets") + bounds.append((child_start, child_end)) + return tuple( + None if bounds[node] is None else bounds[node][0] + for node in result_nodes ) -def _path_positions( - value: bytes, - metadata: bytes, - paths: Sequence[_Path], - plans: List[_PositionPlan], -) -> Sequence[Optional[int]]: - for index, plan in enumerate(plans): - if plan.matches(value, metadata): - if index: - plans.insert(0, plans.pop(index)) - return plan.positions - plan = _position_plan(value, metadata, paths) - plans.insert(0, plan) - del plans[8:] - return plan.positions - - def _replace_path( value: bytes, metadata: bytes, pos: int, path: _Path, replacement: bytes, + limit=None, ) -> bytes: + limit = len(value) if limit is None else limit + value_end = pos + _checked_value_size(value, pos, limit) if not path: return replacement @@ -262,8 +290,8 @@ def _replace_path( key_id = _metadata_key_ids(metadata).get(segment) if key_id is None: raise ValueError(f"VARIANT path does not exist: {segment}") - size, id_size, offset_size, id_start, offset_start, data_start = ( - _object_layout(value, pos)) + size, id_size, id_start, data_start, offsets, container_end = ( + _checked_object_layout(value, pos, value_end)) ids = [ _read_unsigned(value, id_start + i * id_size, id_size) for i in range(size) @@ -274,28 +302,32 @@ def _replace_path( raise ValueError(f"VARIANT path does not exist: {segment}") children = [] for i in range(size): - child_pos = data_start + _read_unsigned( - value, offset_start + i * offset_size, offset_size) - child = value[child_pos:child_pos + _value_size(value, child_pos)] + child_pos = data_start + offsets[i] + child_end = child_pos + _checked_value_size( + value, child_pos, container_end) + child = value[child_pos:child_end] if i == slot: child = _replace_path( - value, metadata, child_pos, path[1:], replacement) + value, metadata, child_pos, path[1:], replacement, + child_end) children.append(child) return _build_object_value(list(zip(ids, children))) if (value[pos] & 0x3) != _ARRAY: raise ValueError("VARIANT path expects an array") - size, offset_size, offset_start, data_start = _array_layout(value, pos) + size, data_start, offsets, _ = _checked_array_layout( + value, pos, value_end) if segment >= size: raise ValueError(f"VARIANT array index does not exist: {segment}") children = [] for i in range(size): - child_pos = data_start + _read_unsigned( - value, offset_start + i * offset_size, offset_size) - child = value[child_pos:child_pos + _value_size(value, child_pos)] + child_pos = data_start + offsets[i] + child_end = data_start + offsets[i + 1] + child = value[child_pos:child_end] if i == segment: child = _replace_path( - value, metadata, child_pos, path[1:], replacement) + value, metadata, child_pos, path[1:], replacement, + child_end) children.append(child) return _build_array_value(children) @@ -410,6 +442,19 @@ def array_from_data(self, data: bytearray, start: int) -> pa.Array: def _take_unsigned(data, positions, widths): + if len(positions) and np.all(widths == widths[0]): + width = int(widths[0]) + if (width < 1 or width > 4 + or np.any(positions < 0) + or np.any(positions + width > len(data))): + raise ValueError("Invalid VARIANT offset") + if width == 1: + return data[positions].astype(np.int64, copy=False) + indices = positions[:, None] + np.arange(width) + values = data[indices].astype(np.int64, copy=False) + return np.sum( + values << (np.arange(width, dtype=np.int64) * 8), axis=1) + result = np.empty(len(positions), dtype=np.int64) for width in range(1, 5): selected = widths == width @@ -471,13 +516,15 @@ def _vectorized_path_positions( data = np.frombuffer(values.data, dtype=np.uint8) first_value = values.view(0) positions = [np.zeros(len(values.array), dtype=np.int64)] + limits = [row_ends - row_starts] try: for parent_node, kind, segment in nodes[1:]: parent = positions[parent_node] + parent_ends = row_starts + limits[parent_node] absolute_parent = row_starts + parent if (np.any(absolute_parent < row_starts) - or np.any(absolute_parent >= row_ends)): + or np.any(absolute_parent >= parent_ends)): return None headers = data[absolute_parent] type_info = (headers >> 2).astype(np.int64, copy=False) @@ -488,66 +535,98 @@ def _vectorized_path_positions( key_id = key_ids.get(segment) if key_id is None: return None - first_layout = _object_layout( - first_value, int(parent[0])) - size, id_size, _, id_start, _, _ = first_layout + first_layout = _checked_object_layout( + first_value, int(parent[0]), int(limits[parent_node][0])) + size, id_size, id_start, _, first_offsets, _ = first_layout id_table = bytes( first_value[id_start:id_start + size * id_size]) slot = _field_slot(id_table, id_size, key_id) if slot is None: return None + successor_slot = min( + ( + index for index in range(size + 1) + if first_offsets[index] > first_offsets[slot] + ), + key=lambda index: first_offsets[index], + ) size_widths = np.where( ((type_info >> 4) & 0x1) != 0, _U32_SIZE, 1) + if np.any(absolute_parent + 1 + size_widths > parent_ends): + return None sizes = _take_unsigned( data, absolute_parent + 1, size_widths) id_widths = ((type_info >> 2) & 0x3) + 1 offset_widths = (type_info & 0x3) + 1 - if np.any(sizes <= slot): + if np.any(sizes != size): return None id_starts = absolute_parent + 1 + size_widths + offset_starts = id_starts + sizes * id_widths + data_starts = offset_starts + (sizes + 1) * offset_widths + if np.any(data_starts > parent_ends): + return None ids = _take_unsigned( data, id_starts + slot * id_widths, id_widths) if np.any(ids != key_id): return None - offset_starts = id_starts + sizes * id_widths - offsets = _take_unsigned( - data, - offset_starts + slot * offset_widths, - offset_widths, - ) - data_starts = offset_starts + (sizes + 1) * offset_widths - child = data_starts + offsets - row_starts else: if np.any((headers & 0x3) != _ARRAY): return None + size = _checked_array_layout( + first_value, int(parent[0]), + int(limits[parent_node][0]))[0] + if segment >= size: + return None size_widths = np.where( ((type_info >> 2) & 0x1) != 0, _U32_SIZE, 1) + if np.any(absolute_parent + 1 + size_widths > parent_ends): + return None sizes = _take_unsigned( data, absolute_parent + 1, size_widths) - if np.any(sizes <= segment): + if np.any(sizes != size): return None + slot = segment + successor_slot = slot + 1 offset_widths = (type_info & 0x3) + 1 offset_starts = absolute_parent + 1 + size_widths - offsets = _take_unsigned( - data, - offset_starts + segment * offset_widths, - offset_widths, - ) data_starts = offset_starts + (sizes + 1) * offset_widths - child = data_starts + offsets - row_starts + if np.any(data_starts > parent_ends): + return None + offsets = _take_unsigned( + data, + offset_starts + slot * offset_widths, + offset_widths, + ) + next_offsets = _take_unsigned( + data, + offset_starts + successor_slot * offset_widths, + offset_widths, + ) + final_offsets = _take_unsigned( + data, + offset_starts + sizes * offset_widths, + offset_widths, + ) - if np.any(child < 0) or np.any(row_starts + child >= row_ends): + child = data_starts + offsets - row_starts + child_ends = data_starts + next_offsets - row_starts + if (np.any(offsets >= next_offsets) + or np.any(next_offsets > final_offsets) + or np.any(data_starts + final_offsets != parent_ends) + or np.any(child < 0) + or np.any(child >= child_ends)): return None positions.append(child) + limits.append(child_ends) except (IndexError, ValueError): return None return ( row_starts, - row_ends, data, tuple(positions[node] for node in result_nodes), + tuple(limits[node] for node in result_nodes), ) @@ -557,9 +636,9 @@ def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): parsed_paths) if planned is None: return None - row_starts, row_ends, data, positions = planned + row_starts, data, positions, limits = planned results = [] - for pos, target_type in zip(positions, target_types): + for pos, limit, target_type in zip(positions, limits, target_types): if not (pa.types.is_float32(target_type) or pa.types.is_float64(target_type)): return None @@ -571,7 +650,7 @@ def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): value_size, data_type = 4, np.dtype(' row_ends): + if np.any(absolute + 1 + value_size != row_starts + limit): return None indices = absolute[:, None] + 1 + np.arange(value_size) raw = np.ascontiguousarray(data[indices]) @@ -580,13 +659,165 @@ def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): return results -def _decode_scalar(value, metadata: bytes, pos: int): - value_type = _variant_get_type(value, pos) - if value_type == _Type.DOUBLE: - return struct.unpack_from(' target_type.precision: + raise ValueError("decimal precision overflow") + return result + + +def _cast_integer(value, target_type): + number = int(value) + bits = target_type.bit_width + return ((number + (1 << (bits - 1))) % (1 << bits)) - (1 << (bits - 1)) + + +def _cast_python(value, target_type): + if value is None or pa.types.is_null(target_type): + return None + try: + if pa.types.is_struct(target_type): + if isinstance(value, str): + value = json.loads(value) + if not isinstance(value, dict): + raise TypeError + return { + field.name: ( + None if field.name not in value + else _cast_python(value[field.name], field.type) + ) + for field in target_type + } + if (pa.types.is_list(target_type) + or pa.types.is_large_list(target_type) + or pa.types.is_fixed_size_list(target_type)): + if isinstance(value, str): + value = json.loads(value) + if not isinstance(value, list): + raise TypeError + return [ + _cast_python(child, target_type.value_type) + for child in value + ] + if pa.types.is_map(target_type): + if isinstance(value, str): + value = json.loads(value) + if not isinstance(value, dict) or not ( + pa.types.is_string(target_type.key_type) + or pa.types.is_large_string(target_type.key_type)): + raise TypeError + return [ + (key, _cast_python(child, target_type.item_type)) + for key, child in value.items() + ] + if pa.types.is_string(target_type) or pa.types.is_large_string( + target_type): + if isinstance(value, (dict, list)): + return _json_text(value) + if isinstance(value, bool): + return str(value).lower() + if isinstance(value, decimal.Decimal): + return _decimal_text(value) + if isinstance(value, (datetime.date, datetime.datetime)): + return value.isoformat() + return str(value) + if pa.types.is_boolean(target_type): + if isinstance(value, str): + lowered = value.strip().lower() + if lowered not in ('true', 'false'): + raise ValueError + return lowered == 'true' + if isinstance(value, (bool, int, decimal.Decimal)): + return value != 0 + raise TypeError + if pa.types.is_signed_integer(target_type): + if not isinstance(value, (bool, int, float, decimal.Decimal, str)): + raise TypeError + return _cast_integer(value, target_type) + if pa.types.is_floating(target_type): + if not isinstance(value, (bool, int, float, decimal.Decimal, str)): + raise TypeError + return float(value) + if pa.types.is_decimal(target_type): + if not isinstance(value, (bool, int, float, decimal.Decimal, str)): + raise TypeError + return _cast_decimal(value, target_type) + if pa.types.is_binary(target_type) or pa.types.is_large_binary( + target_type): + if not isinstance(value, str): + raise TypeError + return value.encode('utf-8') + if pa.types.is_date32(target_type): + if isinstance(value, datetime.datetime): + return value.date() + if isinstance(value, datetime.date): + return value + if isinstance(value, str): + return datetime.date.fromisoformat(value) + raise TypeError + if pa.types.is_timestamp(target_type): + if isinstance(value, datetime.datetime): + return value + if isinstance(value, datetime.date): + return datetime.datetime.combine(value, datetime.time()) + if isinstance(value, str): + return datetime.datetime.fromisoformat(value) + if isinstance(value, (int, float)): + return datetime.datetime.fromtimestamp( + value, tz=datetime.timezone.utc) + raise TypeError + except (ArithmeticError, TypeError, ValueError): + pass + raise ValueError(f"Invalid cast {value!r} to {target_type}") + + +def _decode_scalar(value, metadata: bytes, pos: int, target_type): + size = _checked_value_size(value, pos) + selected = bytes(value[pos:pos + size]) + decoded = GenericVariant(selected, metadata).to_python() + return _cast_python(decoded, target_type) def _patched_chunk( @@ -714,10 +945,10 @@ def _vectorized_replace_chunk( ) if planned is None: return None - row_starts, row_ends, input_data, positions = planned + row_starts, input_data, positions, limits = planned replacements = [] - for (_, _, provider), pos in zip(parsed, positions): + for (_, _, provider), pos, limit in zip(parsed, positions, limits): replacement = provider.numpy_values(global_row, len(chunk)) if replacement is None: return None @@ -725,7 +956,8 @@ def _vectorized_replace_chunk( expected_header = provider._type_header if not np.all(input_data[absolute] == expected_header): return None - if np.any(absolute + provider._fixed_size > row_ends): + if np.any( + absolute + provider._fixed_size != row_starts + limit): return None replacements.append((pos, provider, replacement)) @@ -798,10 +1030,11 @@ def _variant_get(column, paths: Mapping[str, pa.DataType]): parsed_paths, plans, ) - for (path, _, _), pos in zip(parsed, positions): + for (path, _, target_type), pos in zip(parsed, positions): results[path].append( None if pos is None - else _decode_scalar(value, metadata[row], pos) + else _decode_scalar( + value, metadata[row], pos, target_type) ) for path, _, target_type in parsed: result_chunks[path].append( @@ -916,7 +1149,7 @@ def variant_replace( ) if new_size is None: new_size = len(encoded) - if new_size != _value_size(value, pos): + if new_size != _checked_value_size(value, pos): rebuild = True break if patched_data is None: diff --git a/paimon-python/pypaimon/data/variant_shredding.py b/paimon-python/pypaimon/data/variant_shredding.py index 72a4508dd32a..a8e1b1c81fe0 100644 --- a/paimon-python/pypaimon/data/variant_shredding.py +++ b/paimon-python/pypaimon/data/variant_shredding.py @@ -261,10 +261,32 @@ def _append_scalar(builder, value, arrow_type: pa.DataType) -> None: else: builder.append_timestamp_ntz(int(value)) elif pa.types.is_decimal(arrow_type): - if isinstance(value, _decimal.Decimal): - builder.append_decimal(value) - else: - builder.append_decimal(_decimal.Decimal(str(value))) + decimal = ( + value if isinstance(value, _decimal.Decimal) + else _decimal.Decimal(str(value)) + ) + sign, digits, exponent = decimal.as_tuple() + unscaled = int(''.join(str(digit) for digit in digits)) + if sign: + unscaled = -unscaled + shift = exponent + arrow_type.scale + if shift < 0 and unscaled % (10 ** -shift): + raise ValueError( + f'{decimal} does not have Arrow scale {arrow_type.scale}') + unscaled = ( + unscaled * (10 ** shift) if shift >= 0 + else unscaled // (10 ** -shift) + ) + scale = arrow_type.scale + precision = max(1, len(str(abs(unscaled)))) + if precision > arrow_type.precision: + raise ValueError( + f'{decimal} exceeds Arrow precision {arrow_type.precision}') + if scale < 0: + unscaled *= 10 ** -scale + scale = 0 + precision = max(1, len(str(abs(unscaled)))) + builder.append_decimal_unscaled(unscaled, precision, scale) else: # Fallback: encode as string builder.append_string(str(value)) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index 03cce7cf52bf..6c251c8d435e 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -15,6 +15,7 @@ # limitations under the License. import unittest +from decimal import Decimal from unittest.mock import patch import pyarrow as pa @@ -88,6 +89,69 @@ def test_get_path_mapping_in_one_pass(self): self.assertEqual(result['$.velocity.y'].to_pylist(), [1.0, 3.0]) self.assertEqual(result['$.velocity.z'].to_pylist(), [-2.0, -4.0]) + def test_get_matches_java_cast_semantics(self): + column = _variants([{ + 'long': 123, + 'object': {'age': 2}, + 'array': [1, '2'], + }]) + + self.assertEqual( + variant_get(column, '$.long', pa.string()).to_pylist(), ['123']) + self.assertEqual( + variant_get(column, '$.object', pa.string()).to_pylist(), + ['{"age":2}'], + ) + self.assertEqual( + variant_get(column, '$.array', pa.string()).to_pylist(), + ['[1,"2"]'], + ) + self.assertEqual( + variant_get( + column, + '$.object', + pa.struct([('age', pa.int32()), ('name', pa.string())]), + ).to_pylist(), + [{'age': 2, 'name': None}], + ) + self.assertEqual( + variant_get( + column, '$.array', pa.list_(pa.int32())).to_pylist(), + [[1, 2]], + ) + self.assertEqual( + variant_get( + column, '$.object', + pa.map_(pa.string(), pa.int32())).to_pylist(), + [[('age', 2)]], + ) + with self.assertRaisesRegex(ValueError, "Invalid cast"): + variant_get(column, '$.object', pa.int32()) + + def test_get_decimal_is_exact(self): + expected = Decimal('12345678901234567890123456789012345678') + column = _variants([{'value': expected}]) + + result = variant_get( + column, '$.value', pa.decimal128(38, 0)) + + self.assertEqual(result.to_pylist(), [expected]) + + def test_get_copies_only_selected_subtree(self): + column = _variants([{'small': 'x', 'large': b'x' * (2 * 1024 * 1024)}]) + decoded_sizes = [] + original = GenericVariant.to_python + + def decode(selected): + decoded_sizes.append(len(selected.value())) + return original(selected) + + with patch.object(GenericVariant, 'to_python', decode): + result = variant_get(column, '$.small', pa.string()) + + self.assertEqual(result.to_pylist(), ['x']) + self.assertEqual(decoded_sizes, [2]) + def test_get_rejects_invalid_arguments(self): column = _variants([{'value': 1}]) with self.assertRaisesRegex(ValueError, "Invalid VARIANT path"): @@ -98,6 +162,26 @@ def test_get_rejects_invalid_arguments(self): class TestVariantReplace(unittest.TestCase): + def test_truncated_value_does_not_cross_row_boundary(self): + first = GenericVariant.from_python({'value': 1.0}) + second = GenericVariant.from_python({'value': 2.0}) + column = pa.StructArray.from_arrays( + [ + pa.array([first.value()[:-8], second.value()]), + pa.array([first.metadata(), second.metadata()]), + ], + names=['value', 'metadata'], + ) + + with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"): + variant_replace( + column, '$.value', + pa.array([3.0, 4.0], type=pa.float64())) + self.assertEqual( + GenericVariant.from_arrow_struct(column[1].as_py()).to_python(), + {'value': 2.0}, + ) + def test_equal_length_uses_copy_on_write(self): column = _variants([ {'number': 1.0, 'text': 'keep'}, @@ -271,6 +355,24 @@ def test_scalar_and_same_length_string_replacement(self): result.field('value').buffers()[1].address, ) + def test_decimal_replacement_preserves_arrow_value(self): + for data_type, value, expected_exponent in ( + (pa.decimal128(10, 2), Decimal('100.00'), -2), + (pa.decimal128(10, -2), Decimal('1E+2'), 0)): + with self.subTest(data_type=data_type): + column = _variants([{'value': 0}]) + + result = variant_replace( + column, + '$.value', + pa.array([value], type=data_type), + ) + + decoded = _decode(result)[0]['value'] + self.assertEqual(decoded, Decimal('100.00')) + self.assertEqual( + decoded.as_tuple().exponent, expected_exponent) + def test_different_length_rebuilds_offsets(self): column = _variants([ { From bd14294230a401b11818004cbd3891f0d531cc0c Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 08:45:54 -0700 Subject: [PATCH 10/23] [python] Harden VARIANT path semantics and sparse layouts --- paimon-python/pypaimon/data/variant_path.py | 392 ++++++++++++++---- .../pypaimon/data/variant_shredding.py | 12 +- .../pypaimon/tests/variant_path_test.py | 131 +++++- 3 files changed, 442 insertions(+), 93 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 0e9b6349687a..4ecae8477387 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -34,15 +34,23 @@ _OBJECT, _SHORT_STR, _U32_SIZE, + _VERSION, + _VERSION_MASK, _primitive_header, _read_unsigned, ) from pypaimon.data.generic_variant import ( _BINARY, + _DECIMAL4, + _DECIMAL8, + _DECIMAL16, GenericVariant, _DOUBLE, _FLOAT, _LONG_STR, + _MAX_DECIMAL4_PRECISION, + _MAX_DECIMAL8_PRECISION, + _MAX_DECIMAL16_PRECISION, _PRIMITIVE_FIXED_SIZES, ) from pypaimon.data.variant_shredding import ( @@ -55,6 +63,7 @@ _INDEX_PATTERN = re.compile(r"\[(\d+)]") _KEY_PATTERN = re.compile(r"\.([^\.\[]+)|\['([^']+)']|\[\"([^\"]+)\"]") _Path = Tuple[Tuple[str, object], ...] +_SLOW_PATH_ROWS = 64 @functools.lru_cache(maxsize=256) @@ -79,27 +88,51 @@ def _parse_path(path: str) -> _Path: return tuple(segments) -@functools.lru_cache(maxsize=256) def _metadata_key_ids(metadata: bytes) -> Dict[str, int]: if not metadata: - raise ValueError("MALFORMED_VARIANT: empty metadata") + _malformed("empty metadata") + if (metadata[0] & _VERSION_MASK) != _VERSION: + _malformed("invalid metadata version") offset_size = ((metadata[0] >> 6) & 0x3) + 1 + _require_range(1, offset_size, len(metadata)) size = _read_unsigned(metadata, 1, offset_size) offset_start = 1 + offset_size string_start = offset_start + (size + 1) * offset_size + _require_range(offset_start, (size + 1) * offset_size, len(metadata)) + string_size = len(metadata) - string_start result = {} + previous = 0 for key_id in range(size): start = _read_unsigned( metadata, offset_start + key_id * offset_size, offset_size) end = _read_unsigned( metadata, offset_start + (key_id + 1) * offset_size, offset_size) - result[ - metadata[string_start + start:string_start + end].decode('utf-8') - ] = key_id + if start != previous or end < start or end > string_size: + _malformed("invalid metadata offsets") + try: + key = metadata[string_start + start:string_start + end].decode( + 'utf-8') + except UnicodeDecodeError: + _malformed("invalid metadata string") + if key in result: + _malformed("duplicate metadata key") + result[key] = key_id + previous = end + sentinel = _read_unsigned( + metadata, offset_start + size * offset_size, offset_size) + if sentinel != string_size or sentinel != previous: + _malformed("invalid metadata offsets") return result +def _validate_metadata_version(metadata): + if not metadata: + _malformed("empty metadata") + if (metadata[0] & _VERSION_MASK) != _VERSION: + _malformed("invalid metadata version") + + def _malformed(message): raise ValueError(f"MALFORMED_VARIANT: {message}") @@ -177,6 +210,19 @@ def _checked_value_size(value, pos, limit=None): fixed_size = _PRIMITIVE_FIXED_SIZES.get(type_info) if fixed_size is not None: end = pos + fixed_size + _require_range(pos, fixed_size, limit) + decimal_limit = { + _DECIMAL4: _MAX_DECIMAL4_PRECISION, + _DECIMAL8: _MAX_DECIMAL8_PRECISION, + _DECIMAL16: _MAX_DECIMAL16_PRECISION, + }.get(type_info) + if decimal_limit is not None: + scale = value[pos + 1] + unscaled = int.from_bytes( + value[pos + 2:end], 'little', signed=True) + precision = len(str(abs(unscaled))) if unscaled else 1 + if scale > decimal_limit or precision > decimal_limit: + _malformed("invalid decimal precision or scale") elif type_info in (_BINARY, _LONG_STR): _require_range(pos + 1, _U32_SIZE, limit) end = ( @@ -189,7 +235,6 @@ def _checked_value_size(value, pos, limit=None): return end - pos -@functools.lru_cache(maxsize=2048) def _field_slot(id_table: bytes, id_size: int, key_id: int) -> Optional[int]: for slot in range(len(id_table) // id_size): if _read_unsigned(id_table, slot * id_size, id_size) == key_id: @@ -222,8 +267,12 @@ def _path_positions( root_size = _checked_value_size(value, 0) if root_size != len(value): _malformed("trailing bytes after root value") - key_ids = _metadata_key_ids(metadata) nodes, result_nodes = _compile_paths(tuple(paths)) + _validate_metadata_version(metadata) + key_ids = ( + _metadata_key_ids(metadata) + if any(kind == 'key' for _, kind, _ in nodes[1:]) else {} + ) bounds = [(0, len(value))] for parent_node, kind, segment in nodes[1:]: parent = bounds[parent_node] @@ -471,12 +520,17 @@ def _take_unsigned(data, positions, widths): return result -def _all_binary_values_equal(values: _BinaryValues, expected: bytes) -> bool: +def _all_binary_values_equal( + values: _BinaryValues, expected: bytes, rows=None) -> bool: offsets = values.numpy_offsets() lengths = offsets[1:] - offsets[:-1] + starts = offsets[:-1] + if rows is not None and len(rows) != len(lengths): + lengths = lengths[rows] + starts = starts[rows] if np.any(lengths != len(expected)): return False - if not len(offsets) > 1 or not expected: + if not len(lengths) or not expected: return True data = np.frombuffer(values.data, dtype=np.uint8) @@ -484,38 +538,60 @@ def _all_binary_values_equal(values: _BinaryValues, expected: bytes) -> bool: rows_per_batch = max(1, (1024 * 1024) // len(expected)) for row in range(0, len(lengths), rows_per_batch): end = min(row + rows_per_batch, len(lengths)) - starts = offsets[row:end] - indices = starts[:, None] + np.arange(len(expected)) + indices = starts[row:end, None] + np.arange(len(expected)) if not np.all(data[indices] == expected_array): return False return True +def _valid_row_indices(chunk, values, metadata): + if (chunk.null_count == 0 + and values.array.null_count == 0 + and metadata.null_count == 0): + return np.arange(len(chunk), dtype=np.int64) + valid = np.asarray( + chunk.is_valid().to_numpy(zero_copy_only=False), dtype=bool) + value_valid = np.asarray( + values.array.is_valid().to_numpy(zero_copy_only=False), dtype=bool) + metadata_valid = np.asarray( + metadata.is_valid().to_numpy(zero_copy_only=False), dtype=bool) + if np.any(valid & (~value_valid | ~metadata_valid)): + _malformed("valid VARIANT row has a null child") + return np.flatnonzero(valid) + + def _vectorized_path_positions( values: _BinaryValues, metadata: pa.Array, - valid_count: int, + valid_rows, paths: Sequence[_Path], ): - if (valid_count != len(values.array) - or values.array.null_count - or metadata.null_count - or not len(values.array)): + if not len(valid_rows): return None metadata_values = _BinaryValues(metadata) - first_metadata = bytes(metadata_values.view(0)) - if not _all_binary_values_equal(metadata_values, first_metadata): + first_row = int(valid_rows[0]) + first_metadata = bytes(metadata_values.view(first_row)) + if not _all_binary_values_equal( + metadata_values, first_metadata, valid_rows): return None - key_ids = _metadata_key_ids(first_metadata) nodes, result_nodes = _compile_paths(tuple(paths)) + _validate_metadata_version(first_metadata) + key_ids = ( + _metadata_key_ids(first_metadata) + if any(kind == 'key' for _, kind, _ in nodes[1:]) else {} + ) row_offsets = values.numpy_offsets() - row_starts = row_offsets[:-1] - row_ends = row_offsets[1:] + if len(valid_rows) == len(values.array): + row_starts = row_offsets[:-1] + row_ends = row_offsets[1:] + else: + row_starts = row_offsets[:-1][valid_rows] + row_ends = row_offsets[1:][valid_rows] data = np.frombuffer(values.data, dtype=np.uint8) - first_value = values.view(0) - positions = [np.zeros(len(values.array), dtype=np.int64)] + first_value = values.view(first_row) + positions = [np.zeros(len(valid_rows), dtype=np.int64)] limits = [row_ends - row_starts] try: @@ -623,6 +699,7 @@ def _vectorized_path_positions( return None return ( + valid_rows, row_starts, data, tuple(positions[node] for node in result_nodes), @@ -630,33 +707,101 @@ def _vectorized_path_positions( ) -def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): +def _partition_path_plans(values, metadata, valid_rows, parsed_paths): planned = _vectorized_path_positions( - values, chunk.field(1), len(chunk) - chunk.null_count, - parsed_paths) - if planned is None: + values, metadata, valid_rows, parsed_paths) + if planned is not None: + return [planned], [] + if len(valid_rows) <= _SLOW_PATH_ROWS: + return [], list(valid_rows) + middle = len(valid_rows) // 2 + left_plans, left_rows = _partition_path_plans( + values, metadata, valid_rows[:middle], parsed_paths) + right_plans, right_rows = _partition_path_plans( + values, metadata, valid_rows[middle:], parsed_paths) + return left_plans + right_plans, left_rows + right_rows + + +def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): + if not all(pa.types.is_float32(target_type) + or pa.types.is_float64(target_type) + for target_type in target_types): return None - row_starts, data, positions, limits = planned - results = [] - for pos, limit, target_type in zip(positions, limits, target_types): - if not (pa.types.is_float32(target_type) - or pa.types.is_float64(target_type)): - return None - absolute = row_starts + pos - headers = data[absolute] - if np.all(headers == _primitive_header(_DOUBLE)): - value_size, data_type = 8, np.dtype(' 2 or any( + character < '0' or character > '9' + for part in parts for character in part): + raise ValueError + integral = parts[0] + number = int(integral) if integral else 0 + if negative: + number = -number + bits = target_type.bit_width + minimum = -(1 << (bits - 1)) + maximum = (1 << (bits - 1)) - 1 + if number < minimum or number > maximum: + raise ValueError + return number + + def _cast_python(value, target_type): if value is None or pa.types.is_null(target_type): return None @@ -762,19 +932,26 @@ def _cast_python(value, target_type): return _decimal_text(value) if isinstance(value, (datetime.date, datetime.datetime)): return value.isoformat() + if isinstance(value, bytes): + return value.decode('utf-8') return str(value) if pa.types.is_boolean(target_type): if isinstance(value, str): - lowered = value.strip().lower() - if lowered not in ('true', 'false'): + lowered = value.lower() + if lowered in ('t', 'true', 'y', 'yes', '1'): + return True + if lowered in ('f', 'false', 'n', 'no', '0'): + return False + else: raise ValueError - return lowered == 'true' if isinstance(value, (bool, int, decimal.Decimal)): return value != 0 raise TypeError if pa.types.is_signed_integer(target_type): if not isinstance(value, (bool, int, float, decimal.Decimal, str)): raise TypeError + if isinstance(value, str): + return _cast_string_integer(value, target_type) return _cast_integer(value, target_type) if pa.types.is_floating(target_type): if not isinstance(value, (bool, int, float, decimal.Decimal, str)): @@ -786,9 +963,11 @@ def _cast_python(value, target_type): return _cast_decimal(value, target_type) if pa.types.is_binary(target_type) or pa.types.is_large_binary( target_type): - if not isinstance(value, str): - raise TypeError - return value.encode('utf-8') + if isinstance(value, bytes): + return value + if isinstance(value, str): + return value.encode('utf-8') + raise TypeError if pa.types.is_date32(target_type): if isinstance(value, datetime.datetime): return value.date() @@ -896,6 +1075,11 @@ def values(self, offset: int, length: int): def value_at(self, values, row: int): return values if self._array is None else values[row] + def scalar_at(self, row: int): + if self._array is None: + return self._value.as_py() + return self._array[row].as_py() + def fixed_size(self, value) -> Optional[int]: return self._fixed_size if value is not None else None @@ -903,7 +1087,7 @@ def patch(self, data: bytearray, pos: int, value) -> None: struct.pack_into( self._value_format, data, pos, self._type_header, value) - def numpy_values(self, offset: int, length: int): + def numpy_values(self, offset: int, length: int, rows=None): if self._value_format is None: return None data_type = ( @@ -913,13 +1097,16 @@ def numpy_values(self, offset: int, length: int): if self._array is None: if not self._value.is_valid: return None - return np.full(length, self._value.as_py(), dtype=data_type) + size = length if rows is None else len(rows) + return np.full(size, self._value.as_py(), dtype=data_type) values = self._array.slice(offset, length) - if values.null_count: - return None if isinstance(values, pa.ChunkedArray): values = values.combine_chunks() + if rows is not None: + values = values.take(pa.array(rows, type=pa.int64())) + if values.null_count: + return None return np.asarray( values.to_numpy(zero_copy_only=False), dtype=data_type) @@ -936,42 +1123,68 @@ def _vectorized_replace_chunk( parsed, parsed_paths, global_row, + strict, ): - planned = _vectorized_path_positions( - values, - chunk.field(1), - len(chunk) - chunk.null_count, - parsed_paths, - ) - if planned is None: + if not all(provider._fixed_size is not None + for _, _, provider in parsed): return None - row_starts, input_data, positions, limits = planned - - replacements = [] - for (_, _, provider), pos, limit in zip(parsed, positions, limits): - replacement = provider.numpy_values(global_row, len(chunk)) - if replacement is None: - return None - absolute = row_starts + pos - expected_header = provider._type_header - if not np.all(input_data[absolute] == expected_header): - return None - if np.any( - absolute + provider._fixed_size != row_starts + limit): - return None - replacements.append((pos, provider, replacement)) - + valid_rows = _valid_row_indices(chunk, values, chunk.field(1)) + if not len(valid_rows): + return chunk + plans, slow_rows = _partition_path_plans( + values, chunk.field(1), valid_rows, parsed_paths) data, data_start = values.copy_used_data() output_data = np.frombuffer(data, dtype=np.uint8) - relative_starts = row_starts - data_start - for pos, provider, replacement in replacements: - absolute = relative_starts + pos - output_data[absolute] = provider._type_header - value_size = provider._fixed_size - 1 - replacement_bytes = np.ascontiguousarray( - replacement).view(np.uint8).reshape(len(chunk), value_size) - indices = absolute[:, None] + 1 + np.arange(value_size) - output_data[indices] = replacement_bytes + for planned in plans: + rows, row_starts, input_data, positions, limits = planned + replacements = [] + for (_, _, provider), pos, limit in zip( + parsed, positions, limits): + replacement = provider.numpy_values( + global_row, + len(chunk), + None if len(rows) == len(chunk) else rows, + ) + absolute = row_starts + pos + if (replacement is None + or not np.all( + input_data[absolute] == provider._type_header) + or np.any( + absolute + provider._fixed_size + != row_starts + limit)): + slow_rows.extend(rows) + break + replacements.append((pos, provider, replacement)) + else: + relative_starts = row_starts - data_start + for pos, provider, replacement in replacements: + absolute = relative_starts + pos + output_data[absolute] = provider._type_header + value_size = provider._fixed_size - 1 + replacement_bytes = np.ascontiguousarray( + replacement).view(np.uint8).reshape( + len(rows), value_size) + indices = absolute[:, None] + 1 + np.arange(value_size) + output_data[indices] = replacement_bytes + + metadata = _BinaryValues(chunk.field(1)) + for row in set(int(row) for row in slow_rows): + value = values.view(row) + row_metadata = bytes(metadata.view(row)) + positions = _path_positions(value, row_metadata, parsed_paths, None) + for (path, _, provider), pos in zip(parsed, positions): + if pos is None: + if strict: + raise ValueError(f"VARIANT path does not exist: {path}") + continue + replacement_value = provider.scalar_at(global_row + row) + if replacement_value is None: + return None + if (provider.fixed_size(replacement_value) + != _checked_value_size(value, pos)): + return None + patch_pos = values.bounds(row)[0] - data_start + pos + provider.patch(data, patch_pos, replacement_value) return _patched_chunk(chunk, values, data, data_start) @@ -1109,6 +1322,7 @@ def variant_replace( parsed, parsed_paths, global_row, + strict, ) if vectorized is not None: result_chunks.append(vectorized) diff --git a/paimon-python/pypaimon/data/variant_shredding.py b/paimon-python/pypaimon/data/variant_shredding.py index a8e1b1c81fe0..d04eb6974704 100644 --- a/paimon-python/pypaimon/data/variant_shredding.py +++ b/paimon-python/pypaimon/data/variant_shredding.py @@ -252,11 +252,19 @@ def _append_scalar(builder, value, arrow_type: pa.DataType) -> None: if isinstance(value, datetime.datetime): if value.tzinfo is not None: epoch = datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) - micros = int((value - epoch).total_seconds() * 1_000_000) + delta = value - epoch + micros = ( + (delta.days * 86400 + delta.seconds) * 1_000_000 + + delta.microseconds + ) builder.append_timestamp(micros) else: epoch = datetime.datetime(1970, 1, 1) - micros = int((value - epoch).total_seconds() * 1_000_000) + delta = value - epoch + micros = ( + (delta.days * 86400 + delta.seconds) * 1_000_000 + + delta.microseconds + ) builder.append_timestamp_ntz(micros) else: builder.append_timestamp_ntz(int(value)) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index 6c251c8d435e..3fe30991e568 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -14,6 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import datetime import unittest from decimal import Decimal from unittest.mock import patch @@ -21,8 +22,13 @@ import pyarrow as pa import pyarrow.compute as pc -from pypaimon.data.generic_variant import GenericVariant -from pypaimon.data.variant_path import variant_get, variant_replace +from pypaimon.data._variant_binary import _primitive_header +from pypaimon.data.generic_variant import _DECIMAL4, GenericVariant +from pypaimon.data.variant_path import ( + _path_positions, + variant_get, + variant_replace, +) from pypaimon.data.variant_shredding import _encode_scalar_to_value_bytes @@ -128,6 +134,32 @@ def test_get_matches_java_cast_semantics(self): with self.assertRaisesRegex(ValueError, "Invalid cast"): variant_get(column, '$.object', pa.int32()) + def test_get_matches_java_string_and_binary_casts(self): + column = _variants([{ + 'decimal': '1.9', + 'overflow': '2147483648', + 'truthy': 'yes', + 'falsey': '0', + 'binary': b'abc', + }]) + + self.assertEqual( + variant_get(column, '$.decimal', pa.int32()).to_pylist(), [1]) + self.assertEqual( + variant_get(column, '$.truthy', pa.bool_()).to_pylist(), [True]) + self.assertEqual( + variant_get(column, '$.falsey', pa.bool_()).to_pylist(), [False]) + self.assertEqual( + variant_get(column, '$.binary', pa.binary()).to_pylist(), + [b'abc'], + ) + self.assertEqual( + variant_get(column, '$.binary', pa.string()).to_pylist(), + ['abc'], + ) + with self.assertRaisesRegex(ValueError, "Invalid cast"): + variant_get(column, '$.overflow', pa.int32()) + def test_get_decimal_is_exact(self): expected = Decimal('12345678901234567890123456789012345678') column = _variants([{'value': expected}]) @@ -137,6 +169,28 @@ def test_get_decimal_is_exact(self): self.assertEqual(result.to_pylist(), [expected]) + def test_get_rejects_malformed_metadata_and_decimal(self): + valid = GenericVariant.from_python({'value': 1.0}) + bad_metadata = bytes([2]) + valid.metadata()[1:] + column = pa.StructArray.from_arrays([ + pa.array([valid.value()]), + pa.array([bad_metadata]), + ], names=['value', 'metadata']) + with self.assertRaisesRegex(ValueError, "metadata version"): + variant_get(column, '$.value', pa.float64()) + + for scale, unscaled in ((10, 1), (0, 2147483647)): + with self.subTest(scale=scale, unscaled=unscaled): + value = ( + bytes([_primitive_header(_DECIMAL4), scale]) + + unscaled.to_bytes(4, 'little', signed=True) + ) + malformed = GenericVariant.to_arrow_array([ + GenericVariant(value, b'\x01\x00')]) + with self.assertRaisesRegex( + ValueError, "decimal precision or scale"): + variant_get(malformed, '$', pa.decimal128(38, 0)) + def test_get_copies_only_selected_subtree(self): column = _variants([{'small': 'x', 'large': b'x' * (2 * 1024 * 1024)}]) decoded_sizes = [] @@ -162,6 +216,79 @@ def test_get_rejects_invalid_arguments(self): class TestVariantReplace(unittest.TestCase): + def test_timestamp_replacement_is_exact(self): + for arrow_type, value in ( + (pa.timestamp('us'), + datetime.datetime(9999, 12, 31, 23, 59, 59, 999999)), + (pa.timestamp('us', tz='UTC'), + datetime.datetime( + 2500, 1, 1, 0, 0, 0, 1, + tzinfo=datetime.timezone.utc))): + with self.subTest(arrow_type=arrow_type): + column = _variants([{'value': 0}]) + result = variant_replace( + column, '$.value', pa.scalar(value, type=arrow_type)) + self.assertEqual(_decode(result), [{'value': value}]) + + def test_nullable_fast_path_does_not_devectorize_chunk(self): + size = 50000 + column = _variants( + [None] + [{'value': float(index)} for index in range(1, size)]) + replacement = pa.scalar(-1.0) + + with patch( + 'pypaimon.data.variant_path._path_positions', + wraps=_path_positions, + ) as slow_path: + current = variant_get(column, '$.value', pa.float64()) + result = variant_replace(column, '$.value', replacement) + + self.assertEqual(current[0].as_py(), None) + self.assertEqual(current[-1].as_py(), float(size - 1)) + self.assertIsNone(result[0].as_py()) + self.assertEqual(_decode(result.slice(size - 1, 1)), + [{'value': -1.0}]) + slow_path.assert_not_called() + + def test_sparse_layout_anomalies_keep_slow_path_bounded(self): + size = 4096 + uniform = _variants([ + {'value': float(index)} for index in range(1, size) + ]) + cases = ( + (_variants([{'extra': 1, 'value': 0.0}]), False), + (_variants([{'other': 0.0}]), True), + ) + for first, missing in cases: + with self.subTest(missing=missing): + column = pa.concat_arrays([first, uniform]) + with patch( + 'pypaimon.data.variant_path._path_positions', + wraps=_path_positions, + ) as slow_path: + current = variant_get( + column, '$.value', pa.float64()) + self.assertLessEqual( + slow_path.call_count, 64) + self.assertEqual(current[0].as_py(), None if missing else 0.0) + self.assertEqual(current[-1].as_py(), float(size - 1)) + + with patch( + 'pypaimon.data.variant_path._path_positions', + wraps=_path_positions, + ) as slow_path: + result = variant_replace( + column, '$.value', pa.scalar(-1.0)) + self.assertLessEqual( + slow_path.call_count, 64) + expected_first = {'other': 0.0} if missing else { + 'extra': 1, 'value': -1.0, + } + self.assertEqual(_decode(result.slice(0, 1)), + [expected_first]) + self.assertEqual(_decode(result.slice(size - 1, 1)), + [{'value': -1.0}]) + def test_truncated_value_does_not_cross_row_boundary(self): first = GenericVariant.from_python({'value': 1.0}) second = GenericVariant.from_python({'value': 2.0}) From 94e277f860d7f682535d9b5f9da417ad96009f58 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 09:12:22 -0700 Subject: [PATCH 11/23] [python] Keep sparse VARIANT values on the fast path --- paimon-python/pypaimon/data/variant_path.py | 220 ++++++++++++------ .../pypaimon/tests/variant_path_test.py | 90 +++++++ 2 files changed, 242 insertions(+), 68 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 4ecae8477387..68017f9a588a 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -397,6 +397,9 @@ def _variant_chunks(column): if not (pa.types.is_binary(data_type[0].type) or pa.types.is_large_binary(data_type[0].type)): raise TypeError("VARIANT value field must be binary") + if not (pa.types.is_binary(data_type[1].type) + or pa.types.is_large_binary(data_type[1].type)): + raise TypeError("VARIANT metadata field must be binary") return chunks, chunked, data_type @@ -736,6 +739,7 @@ def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): if len(valid_rows) == len(chunk) and len(plans) == 1 and not slow_rows: _, row_starts, data, positions, limits = plans[0] results = [] + uniform = True for pos, limit, target_type in zip( positions, limits, target_types): absolute = row_starts + pos @@ -745,14 +749,17 @@ def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): elif np.all(headers == _primitive_header(_FLOAT)): value_size, data_type = 4, np.dtype('= maximum: + number = maximum + else: + number = int(value) + else: + number = int(value) return ((number + (1 << (bits - 1))) % (1 << bits)) - (1 << (bits - 1)) @@ -944,7 +972,7 @@ def _cast_python(value, target_type): return False else: raise ValueError - if isinstance(value, (bool, int, decimal.Decimal)): + if isinstance(value, (bool, int)): return value != 0 raise TypeError if pa.types.is_signed_integer(target_type): @@ -983,9 +1011,10 @@ def _cast_python(value, target_type): return datetime.datetime.combine(value, datetime.time()) if isinstance(value, str): return datetime.datetime.fromisoformat(value) - if isinstance(value, (int, float)): - return datetime.datetime.fromtimestamp( + if isinstance(value, int) and not isinstance(value, bool): + result = datetime.datetime.fromtimestamp( value, tz=datetime.timezone.utc) + return result if target_type.tz else result.replace(tzinfo=None) raise TypeError except (ArithmeticError, TypeError, ValueError): pass @@ -1033,6 +1062,47 @@ def _rebuilt_chunk( ) +def _sparse_rebuilt_chunk( + chunk, values, data, data_start, rebuilt_rows): + old_offsets = values.numpy_offsets() + lengths = old_offsets[1:] - old_offsets[:-1] + for row, rebuilt in rebuilt_rows.items(): + lengths[row] = len(rebuilt) + offsets = np.empty(len(chunk) + 1, dtype=np.dtype(values.value_format)) + offsets[0] = 0 + np.cumsum(lengths, out=offsets[1:]) + output = bytearray(int(offsets[-1])) + source_start = int(old_offsets[0]) - data_start + target_start = 0 + for row, rebuilt in sorted(rebuilt_rows.items()): + row_start = int(old_offsets[row]) - data_start + row_end = int(old_offsets[row + 1]) - data_start + prefix_size = row_start - source_start + output[target_start:target_start + prefix_size] = data[ + source_start:row_start] + target_start += prefix_size + output[target_start:target_start + len(rebuilt)] = rebuilt + target_start += len(rebuilt) + source_start = row_end + source_end = int(old_offsets[-1]) - data_start + output[target_start:] = data[source_start:source_end] + validity = ( + None if values.array.null_count == 0 + else values.array.is_valid().buffers()[1] + ) + rebuilt_values = pa.Array.from_buffers( + values.array.type, + len(chunk), + [validity, pa.py_buffer(offsets), pa.py_buffer(output)], + null_count=values.array.null_count, + ) + return pa.StructArray.from_arrays( + [rebuilt_values, chunk.field(1)], + fields=list(chunk.type), + mask=chunk.is_null(), + ) + + class _Replacement: def __init__(self, value, length: int): @@ -1095,20 +1165,26 @@ def numpy_values(self, offset: int, length: int, rows=None): else np.dtype(' bytes: if value is not None and self._value_format is not None: @@ -1133,59 +1209,67 @@ def _vectorized_replace_chunk( return chunk plans, slow_rows = _partition_path_plans( values, chunk.field(1), valid_rows, parsed_paths) + slow_rows = set(int(row) for row in slow_rows) data, data_start = values.copy_used_data() output_data = np.frombuffer(data, dtype=np.uint8) for planned in plans: - rows, row_starts, input_data, positions, limits = planned + rows, row_starts, _, positions, limits = planned replacements = [] + compatible = np.ones(len(rows), dtype=bool) for (_, _, provider), pos, limit in zip( parsed, positions, limits): - replacement = provider.numpy_values( + replacement, replacement_valid = provider.numpy_values( global_row, len(chunk), None if len(rows) == len(chunk) else rows, ) absolute = row_starts + pos - if (replacement is None - or not np.all( - input_data[absolute] == provider._type_header) - or np.any( - absolute + provider._fixed_size - != row_starts + limit)): - slow_rows.extend(rows) - break + compatible &= ( + replacement_valid + & (absolute + provider._fixed_size == row_starts + limit) + ) replacements.append((pos, provider, replacement)) - else: - relative_starts = row_starts - data_start - for pos, provider, replacement in replacements: - absolute = relative_starts + pos - output_data[absolute] = provider._type_header - value_size = provider._fixed_size - 1 - replacement_bytes = np.ascontiguousarray( - replacement).view(np.uint8).reshape( - len(rows), value_size) - indices = absolute[:, None] + 1 + np.arange(value_size) - output_data[indices] = replacement_bytes + slow_rows.update(int(row) for row in rows[~compatible]) + if not np.any(compatible): + continue + relative_starts = row_starts - data_start + compatible_rows = rows[compatible] + for pos, provider, replacement in replacements: + absolute = (relative_starts + pos)[compatible] + output_data[absolute] = provider._type_header + value_size = provider._fixed_size - 1 + replacement_bytes = np.ascontiguousarray( + replacement[compatible]).view(np.uint8).reshape( + len(compatible_rows), value_size) + indices = absolute[:, None] + 1 + np.arange(value_size) + output_data[indices] = replacement_bytes metadata = _BinaryValues(chunk.field(1)) - for row in set(int(row) for row in slow_rows): - value = values.view(row) + rebuilt_rows = {} + for row in slow_rows: + original = bytes(values.view(row)) + value = original row_metadata = bytes(metadata.view(row)) positions = _path_positions(value, row_metadata, parsed_paths, None) - for (path, _, provider), pos in zip(parsed, positions): + for (path, parsed_path, provider), pos in zip(parsed, positions): if pos is None: if strict: raise ValueError(f"VARIANT path does not exist: {path}") continue replacement_value = provider.scalar_at(global_row + row) - if replacement_value is None: - return None - if (provider.fixed_size(replacement_value) - != _checked_value_size(value, pos)): - return None - patch_pos = values.bounds(row)[0] - data_start + pos - provider.patch(data, patch_pos, replacement_value) + value = _replace_path( + value, + row_metadata, + 0, + parsed_path, + provider.encode(replacement_value), + ) + if value != original: + rebuilt_rows[row] = value + if rebuilt_rows: + return _sparse_rebuilt_chunk( + chunk, values, data, data_start, rebuilt_rows) return _patched_chunk(chunk, values, data, data_start) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index 3fe30991e568..1f1223a6f2d7 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -160,6 +160,39 @@ def test_get_matches_java_string_and_binary_casts(self): with self.assertRaisesRegex(ValueError, "Invalid cast"): variant_get(column, '$.overflow', pa.int32()) + def test_get_matches_java_numeric_casts(self): + column = _variants([ + {'value': 1e20}, + {'value': float('nan')}, + {'value': float('inf')}, + {'value': float('-inf')}, + ]) + self.assertEqual( + variant_get(column, '$.value', pa.int32()).to_pylist(), + [2147483647, 0, 2147483647, -2147483648], + ) + + invalid = _variants([ + {'value': Decimal('1.0')}, + {'value': 1.5}, + ]) + with self.assertRaisesRegex(ValueError, "Invalid cast"): + variant_get(invalid.slice(0, 1), '$.value', pa.bool_()) + with self.assertRaisesRegex(ValueError, "Invalid cast"): + variant_get(invalid.slice(1, 1), '$.value', pa.timestamp('us')) + + timestamp = variant_get( + _variants([{'value': 1}]), '$.value', pa.timestamp('us')) + self.assertEqual( + timestamp.to_pylist(), [datetime.datetime(1970, 1, 1, 0, 0, 1)]) + + def test_variant_null_remains_arrow_null(self): + column = _variants([None, {'value': None}, {'value': 1.0}]) + + result = variant_get(column, '$.value', pa.float64()) + + self.assertEqual(result.to_pylist(), [None, None, 1.0]) + def test_get_decimal_is_exact(self): expected = Decimal('12345678901234567890123456789012345678') column = _variants([{'value': expected}]) @@ -213,6 +246,17 @@ def test_get_rejects_invalid_arguments(self): with self.assertRaisesRegex(TypeError, "PyArrow data type"): variant_get(column, '$.value', 'BIGINT') + invalid_metadata = pa.StructArray.from_arrays( + [ + pa.array([None], type=pa.binary()), + pa.array([None], type=pa.string()), + ], + names=['value', 'metadata'], + mask=pa.array([True]), + ) + with self.assertRaisesRegex(TypeError, "metadata field must be binary"): + variant_get(invalid_metadata, '$.value', pa.float64()) + class TestVariantReplace(unittest.TestCase): @@ -289,6 +333,52 @@ def test_sparse_layout_anomalies_keep_slow_path_bounded(self): self.assertEqual(_decode(result.slice(size - 1, 1)), [{'value': -1.0}]) + def test_sparse_value_types_keep_slow_path_bounded(self): + size = 4096 + uniform = _variants([ + {'value': float(index)} for index in range(1, size) + ]) + for exceptional in (None, 1): + with self.subTest(exceptional=exceptional): + column = pa.concat_arrays([ + _variants([{'value': exceptional}]), uniform, + ]) + with patch( + 'pypaimon.data.variant_path._path_positions', + wraps=_path_positions, + ) as slow_path: + current = variant_get( + column, '$.value', pa.float64()) + self.assertLessEqual(slow_path.call_count, 1) + self.assertEqual( + current[0].as_py(), + None if exceptional is None else 1.0, + ) + + with patch( + 'pypaimon.data.variant_path._path_positions', + wraps=_path_positions, + ) as slow_path: + result = variant_replace( + column, '$.value', pa.scalar(-1.0)) + self.assertLessEqual(slow_path.call_count, 1) + self.assertEqual(_decode(result.slice(0, 1)), + [{'value': -1.0}]) + self.assertEqual(_decode(result.slice(size - 1, 1)), + [{'value': -1.0}]) + + replacement = pa.array( + [None] + [-1.0] * (len(uniform) - 1), type=pa.float64()) + with patch( + 'pypaimon.data.variant_path._path_positions', + wraps=_path_positions, + ) as slow_path: + result = variant_replace(uniform, '$.value', replacement) + self.assertLessEqual(slow_path.call_count, 1) + self.assertEqual(_decode(result.slice(0, 1)), [{'value': None}]) + self.assertEqual(_decode(result.slice(len(uniform) - 1, 1)), + [{'value': -1.0}]) + def test_truncated_value_does_not_cross_row_boundary(self): first = GenericVariant.from_python({'value': 1.0}) second = GenericVariant.from_python({'value': 2.0}) From c9c50416d2ade949bb74ce8f47ad715131cb4d37 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 09:33:34 -0700 Subject: [PATCH 12/23] [python] Bound VARIANT object children by offsets --- paimon-python/pypaimon/data/variant_path.py | 30 ++++++++++++++----- .../pypaimon/tests/variant_path_test.py | 21 +++++++++++-- 2 files changed, 41 insertions(+), 10 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 68017f9a588a..b1461f8ddab1 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -171,6 +171,21 @@ def _checked_object_layout(value, pos, limit): ) +def _checked_object_child_bounds( + value, data_start, offsets, slot, end_by_offset=None): + child_offset = offsets[slot] + next_offset = ( + min(offset for offset in offsets if offset > child_offset) + if end_by_offset is None else end_by_offset[child_offset] + ) + child_start = data_start + child_offset + child_end = data_start + next_offset + if _checked_value_size(value, child_start, child_end) != ( + child_end - child_start): + _malformed("child size does not match container offsets") + return child_start, child_end + + def _checked_array_layout(value, pos, limit): _require_range(pos, 2, limit) type_info = (value[pos] >> 2) & 0x3F @@ -293,10 +308,8 @@ def _path_positions( if slot is None: bounds.append(None) continue - child_start = data_start + offsets[slot] - child_size = _checked_value_size( - value, child_start, data_start + offsets[-1]) - child_end = child_start + child_size + child_start, child_end = _checked_object_child_bounds( + value, data_start, offsets, slot) else: if basic_type != _ARRAY: bounds.append(None) @@ -339,7 +352,7 @@ def _replace_path( key_id = _metadata_key_ids(metadata).get(segment) if key_id is None: raise ValueError(f"VARIANT path does not exist: {segment}") - size, id_size, id_start, data_start, offsets, container_end = ( + size, id_size, id_start, data_start, offsets, _ = ( _checked_object_layout(value, pos, value_end)) ids = [ _read_unsigned(value, id_start + i * id_size, id_size) @@ -349,11 +362,12 @@ def _replace_path( slot = ids.index(key_id) except ValueError: raise ValueError(f"VARIANT path does not exist: {segment}") + ordered_offsets = sorted(offsets) + end_by_offset = dict(zip(ordered_offsets, ordered_offsets[1:])) children = [] for i in range(size): - child_pos = data_start + offsets[i] - child_end = child_pos + _checked_value_size( - value, child_pos, container_end) + child_pos, child_end = _checked_object_child_bounds( + value, data_start, offsets, i, end_by_offset) child = value[child_pos:child_end] if i == slot: child = _replace_path( diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index 1f1223a6f2d7..e6b6d2f2757b 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -23,13 +23,16 @@ import pyarrow.compute as pc from pypaimon.data._variant_binary import _primitive_header -from pypaimon.data.generic_variant import _DECIMAL4, GenericVariant +from pypaimon.data.generic_variant import _DECIMAL4, _DOUBLE, GenericVariant from pypaimon.data.variant_path import ( _path_positions, variant_get, variant_replace, ) -from pypaimon.data.variant_shredding import _encode_scalar_to_value_bytes +from pypaimon.data.variant_shredding import ( + _build_object_value, + _encode_scalar_to_value_bytes, +) def _variants(values): @@ -399,6 +402,20 @@ def test_truncated_value_does_not_cross_row_boundary(self): {'value': 2.0}, ) + def test_truncated_object_child_does_not_cross_sibling_boundary(self): + valid = GenericVariant.from_python({'a': 1.0, 'b': 2.0}) + value = _build_object_value([ + (0, bytes([_primitive_header(_DOUBLE)])), + (1, _encode_scalar_to_value_bytes(2.0, pa.float64())), + ]) + column = GenericVariant.to_arrow_array([ + GenericVariant(value, valid.metadata())]) + + with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"): + variant_get(column, '$.a', pa.float64()) + with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"): + variant_replace(column, '$.a', pa.scalar(3.0)) + def test_equal_length_uses_copy_on_write(self): column = _variants([ {'number': 1.0, 'text': 'keep'}, From 57e1a463616b5c5a8369f66afe2c28e8d6590a35 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 17:48:02 -0700 Subject: [PATCH 13/23] [python] Preserve VARIANT timestamp and cast semantics --- paimon-python/pypaimon/data/variant_path.py | 65 +++++++++++++++++-- .../pypaimon/data/variant_shredding.py | 15 ++++- .../pypaimon/tests/variant_path_test.py | 45 +++++++++++++ 3 files changed, 119 insertions(+), 6 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index b1461f8ddab1..9afc3737dee4 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -835,9 +835,23 @@ def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): def _decimal_text(value): text = format(value, 'f') - if '.' in text: - text = text.rstrip('0').rstrip('.') - return '0' if text in ('', '-0') else text + return text[1:] if value == 0 and text.startswith('-') else text + + +def _floating_text(value): + if math.isnan(value): + return 'NaN' + if math.isinf(value): + return 'Infinity' if value > 0 else '-Infinity' + number = decimal.Decimal(repr(value)) + exponent = number.adjusted() + if -3 <= exponent < 7: + text = format(number, 'f') + return text if '.' in text else text + '.0' + mantissa, exponent = format(number, 'E').split('E') + if '.' not in mantissa: + mantissa += '.0' + return mantissa + 'E' + str(int(exponent)) def _json_text(value): @@ -926,6 +940,45 @@ def _cast_string_integer(value, target_type): return number +def _parse_date(value): + value = value.strip() + if value and (value.isdigit() + or value[0] == '-' and value[1:].isdigit()): + return datetime.date(1970, 1, 1) + datetime.timedelta(days=int(value)) + return datetime.datetime.strptime(value, '%Y-%m-%d').date() + + +def _parse_timestamp(value, target_type): + value = value.strip() + if value and (value.isdigit() + or value[0] == '-' and value[1:].isdigit()): + raw_value = int(value) + if target_type.unit == 'ns': + if raw_value % 1000: + raise ValueError("timestamp is not microsecond-aligned") + micros = raw_value // 1000 + else: + micros = raw_value * { + 's': 1_000_000, 'ms': 1000, 'us': 1, + }[target_type.unit] + return datetime.datetime(1970, 1, 1) + datetime.timedelta( + microseconds=int(micros)) + formats = ( + '%Y-%m-%d', + '%Y-%m-%d %H:%M:%S', + '%Y-%m-%d %H:%M:%S.%f', + '%Y-%m-%dT%H:%M', + '%Y-%m-%dT%H:%M:%S', + '%Y-%m-%dT%H:%M:%S.%f', + ) + for date_format in formats: + try: + return datetime.datetime.strptime(value, date_format) + except ValueError: + pass + raise ValueError + + def _cast_python(value, target_type): if value is None or pa.types.is_null(target_type): return None @@ -972,6 +1025,8 @@ def _cast_python(value, target_type): return str(value).lower() if isinstance(value, decimal.Decimal): return _decimal_text(value) + if isinstance(value, float): + return _floating_text(value) if isinstance(value, (datetime.date, datetime.datetime)): return value.isoformat() if isinstance(value, bytes): @@ -1016,7 +1071,7 @@ def _cast_python(value, target_type): if isinstance(value, datetime.date): return value if isinstance(value, str): - return datetime.date.fromisoformat(value) + return _parse_date(value) raise TypeError if pa.types.is_timestamp(target_type): if isinstance(value, datetime.datetime): @@ -1024,7 +1079,7 @@ def _cast_python(value, target_type): if isinstance(value, datetime.date): return datetime.datetime.combine(value, datetime.time()) if isinstance(value, str): - return datetime.datetime.fromisoformat(value) + return _parse_timestamp(value, target_type) if isinstance(value, int) and not isinstance(value, bool): result = datetime.datetime.fromtimestamp( value, tz=datetime.timezone.utc) diff --git a/paimon-python/pypaimon/data/variant_shredding.py b/paimon-python/pypaimon/data/variant_shredding.py index d04eb6974704..18692cbbe602 100644 --- a/paimon-python/pypaimon/data/variant_shredding.py +++ b/paimon-python/pypaimon/data/variant_shredding.py @@ -250,6 +250,10 @@ def _append_scalar(builder, value, arrow_type: pa.DataType) -> None: elif pa.types.is_timestamp(arrow_type): # PyArrow converts timestamp to datetime.datetime if isinstance(value, datetime.datetime): + if (arrow_type.unit == 'ns' + and getattr(value, 'nanosecond', 0) != 0): + raise ValueError( + "VARIANT timestamps require microsecond-aligned values") if value.tzinfo is not None: epoch = datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) delta = value - epoch @@ -267,7 +271,16 @@ def _append_scalar(builder, value, arrow_type: pa.DataType) -> None: ) builder.append_timestamp_ntz(micros) else: - builder.append_timestamp_ntz(int(value)) + raw_value = int(value) + if arrow_type.unit == 'ns': + if raw_value % 1000: + raise ValueError( + "VARIANT timestamps require microsecond-aligned values") + micros = raw_value // 1000 + else: + micros = raw_value * {'s': 1_000_000, 'ms': 1000, 'us': 1}[ + arrow_type.unit] + builder.append_timestamp_ntz(micros) elif pa.types.is_decimal(arrow_type): decimal = ( value if isinstance(value, _decimal.Decimal) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index e6b6d2f2757b..51aae617d636 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -163,6 +163,29 @@ def test_get_matches_java_string_and_binary_casts(self): with self.assertRaisesRegex(ValueError, "Invalid cast"): variant_get(column, '$.overflow', pa.int32()) + temporal = _variants(['1', '2026-08-10']) + self.assertEqual( + variant_get(temporal, '$', pa.date32()).to_pylist(), + [datetime.date(1970, 1, 2), datetime.date(2026, 8, 10)], + ) + self.assertEqual( + variant_get(temporal.slice(0, 1), '$', pa.timestamp('us')) + .to_pylist(), + [datetime.datetime(1970, 1, 1, 0, 0, 0, 1)], + ) + + def test_get_matches_java_numeric_strings(self): + column = _variants([ + Decimal('100.00'), 1e20, 1e-4, + float('inf'), float('-inf'), float('nan'), + ]) + + self.assertEqual( + variant_get(column, '$', pa.string()).to_pylist(), + ['100.00', '1.0E20', '1.0E-4', + 'Infinity', '-Infinity', 'NaN'], + ) + def test_get_matches_java_numeric_casts(self): column = _variants([ {'value': 1e20}, @@ -277,6 +300,28 @@ def test_timestamp_replacement_is_exact(self): column, '$.value', pa.scalar(value, type=arrow_type)) self.assertEqual(_decode(result), [{'value': value}]) + column = _variants([{'value': 0}]) + for nanos in (1, -1, 1001, -1001): + with self.subTest(nanos=nanos): + with self.assertRaisesRegex(ValueError, "microsecond-aligned"): + variant_replace( + column, + '$.value', + pa.scalar(nanos, type=pa.timestamp('ns')), + ) + for nanos in (1000, -1000): + with self.subTest(nanos=nanos): + result = variant_replace( + column, + '$.value', + pa.scalar(nanos, type=pa.timestamp('ns')), + ) + self.assertEqual( + _decode(result), + [{'value': datetime.datetime(1970, 1, 1) + + datetime.timedelta(microseconds=nanos // 1000)}], + ) + def test_nullable_fast_path_does_not_devectorize_chunk(self): size = 50000 column = _variants( From 637224595f928897783d67cd29e1ba39683fc2bd Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 18:16:01 -0700 Subject: [PATCH 14/23] [python] Preserve VARIANT float string precision --- paimon-python/pypaimon/data/variant_path.py | 25 ++++++++++++------- .../pypaimon/tests/variant_path_test.py | 16 ++++++++++-- 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 9afc3737dee4..7abf04ab9f32 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -838,19 +838,21 @@ def _decimal_text(value): return text[1:] if value == 0 and text.startswith('-') else text -def _floating_text(value): +def _floating_text(value, single_precision=False): + if single_precision: + value = np.float32(value) if math.isnan(value): return 'NaN' if math.isinf(value): return 'Infinity' if value > 0 else '-Infinity' - number = decimal.Decimal(repr(value)) - exponent = number.adjusted() - if -3 <= exponent < 7: - text = format(number, 'f') - return text if '.' in text else text + '.0' - mantissa, exponent = format(number, 'E').split('E') - if '.' not in mantissa: - mantissa += '.0' + absolute = abs(value) + if value == 0 or 1e-3 <= absolute < 1e7: + text = np.format_float_positional(value, unique=True, trim='k') + return text + '0' if text.endswith('.') else text + text = np.format_float_scientific(value, unique=True, trim='k') + mantissa, exponent = text.split('e') + if mantissa.endswith('.'): + mantissa += '0' return mantissa + 'E' + str(int(exponent)) @@ -1094,6 +1096,11 @@ def _decode_scalar(value, metadata: bytes, pos: int, target_type): size = _checked_value_size(value, pos) selected = bytes(value[pos:pos + size]) decoded = GenericVariant(selected, metadata).to_python() + type_info = (value[pos] >> 2) & 0x3F + if ((pa.types.is_string(target_type) + or pa.types.is_large_string(target_type)) + and type_info in (_FLOAT, _DOUBLE)): + return _floating_text(decoded, type_info == _FLOAT) return _cast_python(decoded, target_type) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index 51aae617d636..2c2cfc621583 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -175,17 +175,29 @@ def test_get_matches_java_string_and_binary_casts(self): ) def test_get_matches_java_numeric_strings(self): - column = _variants([ + doubles = _variants([ Decimal('100.00'), 1e20, 1e-4, float('inf'), float('-inf'), float('nan'), ]) self.assertEqual( - variant_get(column, '$', pa.string()).to_pylist(), + variant_get(doubles, '$', pa.string()).to_pylist(), ['100.00', '1.0E20', '1.0E-4', 'Infinity', '-Infinity', 'NaN'], ) + floats = GenericVariant.to_arrow_array([ + GenericVariant( + _encode_scalar_to_value_bytes(value, pa.float32()), + b'\x01\x00', + ) + for value in (1.2, 1e20, 1e-4) + ]) + self.assertEqual( + variant_get(floats, '$', pa.string()).to_pylist(), + ['1.2', '1.0E20', '1.0E-4'], + ) + def test_get_matches_java_numeric_casts(self): column = _variants([ {'value': 1e20}, From b66d7edd567c07015ae1bc3b6b8f3bd3f0ef5245 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 18:25:49 -0700 Subject: [PATCH 15/23] [python] Preserve nested VARIANT numeric semantics --- paimon-python/pypaimon/data/variant_path.py | 71 +++++++++++++++++-- .../pypaimon/tests/variant_path_test.py | 37 ++++++++++ 2 files changed, 103 insertions(+), 5 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 7abf04ab9f32..da63e0e75753 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -32,6 +32,7 @@ from pypaimon.data._variant_binary import ( _ARRAY, _OBJECT, + _PRIMITIVE, _SHORT_STR, _U32_SIZE, _VERSION, @@ -883,6 +884,57 @@ def _json_text(value): return json.dumps(str(value), ensure_ascii=False) +def _variant_json(value, metadata, pos, limit, keys=None): + size = _checked_value_size(value, pos, limit) + if pos + size != limit: + _malformed("child size does not match container offsets") + header = value[pos] + basic_type = header & 0x3 + type_info = (header >> 2) & 0x3F + if basic_type == _OBJECT: + object_size, id_size, id_start, data_start, offsets, _ = ( + _checked_object_layout(value, pos, limit)) + if keys is None: + keys = { + key_id: key + for key, key_id in _metadata_key_ids(metadata).items() + } + fields = [] + for slot in range(object_size): + key_id = _read_unsigned( + value, id_start + slot * id_size, id_size) + if key_id not in keys: + _malformed("object key is missing from metadata") + child_start, child_end = _checked_object_child_bounds( + value, data_start, offsets, slot) + fields.append( + json.dumps(keys[key_id], ensure_ascii=False) + + ':' + + _variant_json( + value, metadata, child_start, child_end, keys) + ) + return '{' + ','.join(fields) + '}' + if basic_type == _ARRAY: + array_size, data_start, offsets, _ = _checked_array_layout( + value, pos, limit) + return '[' + ','.join( + _variant_json( + value, + metadata, + data_start + offsets[slot], + data_start + offsets[slot + 1], + keys, + ) + for slot in range(array_size) + ) + ']' + + decoded = GenericVariant(bytes(value[pos:limit]), metadata).to_python() + if basic_type == _PRIMITIVE and type_info in (_FLOAT, _DOUBLE): + text = _floating_text(decoded, type_info == _FLOAT) + return text if math.isfinite(decoded) else json.dumps(text) + return _json_text(decoded) + + def _cast_decimal(value, target_type): if isinstance(value, bool): value = decimal.Decimal(1 if value else 0) @@ -956,15 +1008,17 @@ def _parse_timestamp(value, target_type): or value[0] == '-' and value[1:].isdigit()): raw_value = int(value) if target_type.unit == 'ns': - if raw_value % 1000: - raise ValueError("timestamp is not microsecond-aligned") - micros = raw_value // 1000 + return pa.scalar(raw_value, type=target_type).as_py() else: micros = raw_value * { 's': 1_000_000, 'ms': 1000, 'us': 1, }[target_type.unit] return datetime.datetime(1970, 1, 1) + datetime.timedelta( microseconds=int(micros)) + if target_type.unit == 'ns': + nanos = int(np.datetime64(value.replace(' ', 'T'), 'ns').astype( + np.int64)) + return pa.scalar(nanos, type=target_type).as_py() formats = ( '%Y-%m-%d', '%Y-%m-%d %H:%M:%S', @@ -1094,11 +1148,18 @@ def _cast_python(value, target_type): def _decode_scalar(value, metadata: bytes, pos: int, target_type): size = _checked_value_size(value, pos) - selected = bytes(value[pos:pos + size]) - decoded = GenericVariant(selected, metadata).to_python() + end = pos + size type_info = (value[pos] >> 2) & 0x3F + if (pa.types.is_string(target_type) + or pa.types.is_large_string(target_type)): + basic_type = value[pos] & 0x3 + if basic_type in (_OBJECT, _ARRAY): + return _variant_json(value, metadata, pos, end) + selected = bytes(value[pos:end]) + decoded = GenericVariant(selected, metadata).to_python() if ((pa.types.is_string(target_type) or pa.types.is_large_string(target_type)) + and (value[pos] & 0x3) == _PRIMITIVE and type_info in (_FLOAT, _DOUBLE)): return _floating_text(decoded, type_info == _FLOAT) return _cast_python(decoded, target_type) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index 2c2cfc621583..b4f030da1a29 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -19,6 +19,7 @@ from decimal import Decimal from unittest.mock import patch +import numpy as np import pyarrow as pa import pyarrow.compute as pc @@ -30,6 +31,7 @@ variant_replace, ) from pypaimon.data.variant_shredding import ( + _build_array_value, _build_object_value, _encode_scalar_to_value_bytes, ) @@ -198,6 +200,26 @@ def test_get_matches_java_numeric_strings(self): ['1.2', '1.0E20', '1.0E-4'], ) + nested = GenericVariant.to_arrow_array([ + GenericVariant( + _build_object_value([(0, floats[0].as_py()['value'])]), + GenericVariant.from_python({'value': 0}).metadata(), + ), + GenericVariant( + _build_array_value([floats[0].as_py()['value']]), + b'\x01\x00', + ), + ]) + self.assertEqual( + variant_get(nested, '$', pa.string()).to_pylist(), + ['{"value":1.2}', '[1.2]'], + ) + strings = _variants(['1234567', '12345678901234']) + self.assertEqual( + variant_get(strings, '$', pa.string()).to_pylist(), + ['1234567', '12345678901234'], + ) + def test_get_matches_java_numeric_casts(self): column = _variants([ {'value': 1e20}, @@ -224,6 +246,21 @@ def test_get_matches_java_numeric_casts(self): self.assertEqual( timestamp.to_pylist(), [datetime.datetime(1970, 1, 1, 0, 0, 1)]) + nanos = variant_get( + _variants(['1', '-1', '2026-08-10 12:34:56.123456789']), + '$', + pa.timestamp('ns'), + ) + self.assertEqual( + nanos.cast(pa.int64()).to_pylist(), + [ + 1, + -1, + int(np.datetime64( + '2026-08-10T12:34:56.123456789', 'ns').astype(np.int64)), + ], + ) + def test_variant_null_remains_arrow_null(self): column = _variants([None, {'value': None}, {'value': 1.0}]) From 7e5af8e23d3e22e7d9759cac41eed63c3b10083d Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 19:31:58 -0700 Subject: [PATCH 16/23] [python] Match legacy Java floating strings --- paimon-python/pypaimon/data/_java_floating.py | 226 ++++++++++++++++++ paimon-python/pypaimon/data/variant_path.py | 17 +- .../pypaimon/tests/variant_path_test.py | 19 ++ 3 files changed, 247 insertions(+), 15 deletions(-) create mode 100644 paimon-python/pypaimon/data/_java_floating.py diff --git a/paimon-python/pypaimon/data/_java_floating.py b/paimon-python/pypaimon/data/_java_floating.py new file mode 100644 index 000000000000..ac7d0f9b1f88 --- /dev/null +++ b/paimon-python/pypaimon/data/_java_floating.py @@ -0,0 +1,226 @@ +# 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. + +"""Format IEEE floats like ``Float/Double.toString`` on the JDK 8 baseline.""" + +import struct + + +def java_floating_text(value, single_precision=False): + """Return the legacy Java decimal representation of an IEEE value.""" + if single_precision: + bits = struct.unpack('>I', struct.pack('>f', value))[0] + width, fraction_width, exponent_width, bias = 32, 23, 8, 127 + else: + bits = struct.unpack('>Q', struct.pack('>d', value))[0] + width, fraction_width, exponent_width, bias = 64, 52, 11, 1023 + + negative = bool(bits >> (width - 1)) + exponent = ((bits >> fraction_width) + & ((1 << exponent_width) - 1)) + fraction = bits & ((1 << fraction_width) - 1) + if exponent == (1 << exponent_width) - 1: + if fraction: + return 'NaN' + return '-Infinity' if negative else 'Infinity' + if exponent == 0: + if fraction == 0: + return '-0.0' if negative else '0.0' + highest_bit = fraction.bit_length() - 1 + shift = fraction_width - highest_bit + significand = fraction << shift + binary_exponent = 1 - shift - bias + significant_bits = highest_bit + 1 + else: + significand = (1 << fraction_width) | fraction + binary_exponent = exponent - bias + significant_bits = fraction_width + 1 + + trailing_zeros = (significand & -significand).bit_length() - 1 + fraction_bits = fraction_width + 1 - trailing_zeros + tiny_bits = max(0, fraction_bits - binary_exponent - 1) + if tiny_bits == 0 and -21 <= binary_exponent <= 62: + return _format_small_integer( + negative, + significand, + fraction_width, + binary_exponent, + significant_bits, + ) + + decimal_exponent = _estimate_decimal_exponent( + significand, binary_exponent, fraction_width) + base5 = max(0, -decimal_exponent) + base2 = base5 + tiny_bits + binary_exponent + scale5 = max(0, decimal_exponent) + scale2 = scale5 + tiny_bits + margin2 = base2 - significant_bits + reduced = significand >> trailing_zeros + base2 -= fraction_bits - 1 + common2 = min(base2, scale2) + base2 -= common2 + scale2 -= common2 + margin2 -= common2 + if fraction_bits == 1: + margin2 -= 1 + if margin2 < 0: + base2 -= margin2 + scale2 -= margin2 + margin2 = 0 + + base = reduced * 5 ** base5 << base2 + scale = 5 ** scale5 << scale2 + margin = 5 ** base5 << margin2 + base_bits = fraction_bits + base2 + _five_bits(base5) + ten_scale_bits = scale2 + 1 + _five_bits(scale5 + 1) + arithmetic_width = ( + 32 if base_bits < 32 and ten_scale_bits < 32 + else 64 if base_bits < 64 and ten_scale_bits < 64 + else None + ) + ten_scale = scale * 10 + + digits = [] + iteration = 0 + while True: + digit, base = divmod(base, scale) + base *= 10 + margin *= 10 + if arithmetic_width is not None: + margin = _signed(margin, arithmetic_width) + if arithmetic_width is None: + low = base < margin + high = base + margin >= ten_scale + elif iteration == 0 or margin > 0: + low = base < margin + high = _signed(base + margin, arithmetic_width) > ten_scale + else: + low = high = True + + if iteration == 0 and digit == 0 and not high: + decimal_exponent -= 1 + else: + digits.append(digit) + if iteration == 0 and ( + decimal_exponent < -3 or decimal_exponent >= 8): + low = high = False + iteration += 1 + if low or high: + break + + if high: + round_up = not low + if low: + if arithmetic_width is None: + difference = 2 * base - ten_scale + else: + difference = _signed( + _signed(base << 1, arithmetic_width) - ten_scale, + arithmetic_width, + ) + round_up = difference > 0 or ( + difference == 0 and digits[-1] & 1) + if round_up: + decimal_exponent = _round_up(digits, decimal_exponent) + return _format(negative, digits, decimal_exponent + 1) + + +def _format_small_integer( + negative, + significand, + fraction_width, + binary_exponent, + significant_bits, +): + insignificant = 0 + if binary_exponent > significant_bits: + power = binary_exponent - significant_bits - 1 + insignificant = len(str(1 << power)) - 1 + if binary_exponent >= fraction_width: + integer = significand << (binary_exponent - fraction_width) + else: + integer = significand >> (fraction_width - binary_exponent) + if insignificant: + power10 = 10 ** insignificant + integer, residue = divmod(integer, power10) + if residue >= power10 // 2: + integer += 1 + raw_digits = str(integer) + digits = [int(digit) for digit in raw_digits.rstrip('0')] + return _format(negative, digits, len(raw_digits) + insignificant) + + +def _estimate_decimal_exponent( + significand, binary_exponent, fraction_width): + normalized = significand << (52 - fraction_width) + bits = (1023 << 52) | (normalized & ((1 << 52) - 1)) + scaled = struct.unpack('>d', struct.pack('>Q', bits))[0] + estimate = ((scaled - 1.5) * 0.289529654 + 0.176091259 + + binary_exponent * 0.301029995663981) + estimate_bits = struct.unpack('>Q', struct.pack('>d', estimate))[0] + exponent = ((estimate_bits >> 52) & 0x7FF) - 1023 + negative = bool(estimate_bits >> 63) + fraction = estimate_bits & ((1 << 52) - 1) + if 0 <= exponent < 52: + mask = ((1 << 52) - 1) >> exponent + integer = ((fraction | (1 << 52)) >> (52 - exponent)) + if negative: + return -integer if fraction & mask == 0 else -integer - 1 + return integer + if exponent < 0: + magnitude = estimate_bits & ((1 << 63) - 1) + return 0 if magnitude == 0 else (-1 if negative else 0) + return int(estimate) + + +def _five_bits(power): + if power == 0: + return 0 + return (5 ** power).bit_length() if power <= 26 else power * 3 + + +def _signed(value, width): + value &= (1 << width) - 1 + sign = 1 << (width - 1) + return value - (1 << width) if value & sign else value + + +def _round_up(digits, decimal_exponent): + index = len(digits) - 1 + while index >= 0 and digits[index] == 9: + digits[index] = 0 + index -= 1 + if index < 0: + digits[0] = 1 + return decimal_exponent + 1 + digits[index] += 1 + return decimal_exponent + + +def _format(negative, digits, decimal_point): + text = ''.join(str(digit) for digit in digits) + if 0 < decimal_point < 8: + if len(text) < decimal_point: + result = text + '0' * (decimal_point - len(text)) + '.0' + else: + result = text[:decimal_point] + '.' + ( + text[decimal_point:] or '0') + elif -3 < decimal_point <= 0: + result = '0.' + '0' * -decimal_point + text + else: + result = (text[0] + '.' + (text[1:] or '0') + + 'E' + str(decimal_point - 1)) + return ('-' if negative else '') + result diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index da63e0e75753..f2010a68b9ae 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -29,6 +29,7 @@ import numpy as np import pyarrow as pa +from pypaimon.data._java_floating import java_floating_text from pypaimon.data._variant_binary import ( _ARRAY, _OBJECT, @@ -840,21 +841,7 @@ def _decimal_text(value): def _floating_text(value, single_precision=False): - if single_precision: - value = np.float32(value) - if math.isnan(value): - return 'NaN' - if math.isinf(value): - return 'Infinity' if value > 0 else '-Infinity' - absolute = abs(value) - if value == 0 or 1e-3 <= absolute < 1e7: - text = np.format_float_positional(value, unique=True, trim='k') - return text + '0' if text.endswith('.') else text - text = np.format_float_scientific(value, unique=True, trim='k') - mantissa, exponent = text.split('e') - if mantissa.endswith('.'): - mantissa += '0' - return mantissa + 'E' + str(int(exponent)) + return java_floating_text(value, single_precision) def _json_text(value): diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index b4f030da1a29..1ce64c1ce5a1 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -15,6 +15,7 @@ # limitations under the License. import datetime +import struct import unittest from decimal import Decimal from unittest.mock import patch @@ -220,6 +221,24 @@ def test_get_matches_java_numeric_strings(self): ['1234567', '12345678901234'], ) + legacy_floats = _float_variants([ + struct.unpack('>f', struct.pack('>I', bits))[0] + for bits in (0xD44F9F82, 0xEA75E34D, 0xE8FDF7D7) + ]) + self.assertEqual( + variant_get(legacy_floats, '$', pa.string()).to_pylist(), + ['-3.56693731E12', '-7.4315055E25', '-9.5946444E24'], + ) + + legacy_doubles = _variants([ + struct.unpack('>d', struct.pack('>Q', bits))[0] + for bits in (0x439F4B86CD6A5E0C, 0x439DDD7467AF36D9) + ]) + self.assertEqual( + variant_get(legacy_doubles, '$', pa.string()).to_pylist(), + ['5.6376106381000781E17', '5.3800104640575648E17'], + ) + def test_get_matches_java_numeric_casts(self): column = _variants([ {'value': 1e20}, From fe2929ec83a938269bba6bfa560d32779ada12bc Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 19:58:22 -0700 Subject: [PATCH 17/23] [python] Validate VARIANT timestamp nanosecond range --- paimon-python/pypaimon/data/variant_path.py | 60 ++++++++++++------- .../pypaimon/tests/variant_path_test.py | 24 ++++++++ 2 files changed, 62 insertions(+), 22 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index f2010a68b9ae..533577f4f1da 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -64,6 +64,11 @@ _INDEX_PATTERN = re.compile(r"\[(\d+)]") _KEY_PATTERN = re.compile(r"\.([^\.\[]+)|\['([^']+)']|\[\"([^\"]+)\"]") +_TIMESTAMP_PATTERN = re.compile( + r'^(\d{4})-(\d{1,2})-(\d{1,2})' + r'(?:(?: (\d{1,2}):(\d{1,2}):(\d{1,2})(?:\.(\d{1,9}))?)' + r'|(?:T(\d{1,2}):(\d{1,2})(?::(\d{1,2})(?:\.(\d{1,9}))?)?))?$' +) _Path = Tuple[Tuple[str, object], ...] _SLOW_PATH_ROWS = 64 @@ -995,31 +1000,42 @@ def _parse_timestamp(value, target_type): or value[0] == '-' and value[1:].isdigit()): raw_value = int(value) if target_type.unit == 'ns': - return pa.scalar(raw_value, type=target_type).as_py() - else: - micros = raw_value * { - 's': 1_000_000, 'ms': 1000, 'us': 1, - }[target_type.unit] + if raw_value < -(1 << 63) or raw_value > (1 << 63) - 1: + raise ValueError + return raw_value + micros = raw_value * { + 's': 1_000_000, 'ms': 1000, 'us': 1, + }[target_type.unit] return datetime.datetime(1970, 1, 1) + datetime.timedelta( microseconds=int(micros)) + + match = _TIMESTAMP_PATTERN.match(value) + if match is None: + raise ValueError + groups = match.groups() + year, month, day = (int(part) for part in groups[:3]) + if groups[3] is not None: + hour, minute, second, fraction = groups[3:7] + else: + hour, minute, second, fraction = groups[7:11] + hour = int(hour or 0) + minute = int(minute or 0) + second = int(second or 0) + nanos = int((fraction or '').ljust(9, '0') or 0) + base = datetime.datetime(year, month, day, hour, minute, second) + if target_type.unit == 'ns': - nanos = int(np.datetime64(value.replace(' ', 'T'), 'ns').astype( - np.int64)) - return pa.scalar(nanos, type=target_type).as_py() - formats = ( - '%Y-%m-%d', - '%Y-%m-%d %H:%M:%S', - '%Y-%m-%d %H:%M:%S.%f', - '%Y-%m-%dT%H:%M', - '%Y-%m-%dT%H:%M:%S', - '%Y-%m-%dT%H:%M:%S.%f', - ) - for date_format in formats: - try: - return datetime.datetime.strptime(value, date_format) - except ValueError: - pass - raise ValueError + delta = base - datetime.datetime(1970, 1, 1) + epoch_nanos = ( + (delta.days * 86400 + delta.seconds) * 1_000_000_000 + nanos + ) + if epoch_nanos < -(1 << 63) or epoch_nanos > (1 << 63) - 1: + raise ValueError + return epoch_nanos + + precision = {'s': 0, 'ms': 3, 'us': 6}[target_type.unit] + nanos = nanos // (10 ** (9 - precision)) * (10 ** (9 - precision)) + return base + datetime.timedelta(microseconds=nanos // 1000) def _cast_python(value, target_type): diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index 1ce64c1ce5a1..61404cd38cd8 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -280,6 +280,30 @@ def test_get_matches_java_numeric_casts(self): ], ) + def test_get_timestamp_ns_checks_range_and_format(self): + boundaries = _variants([ + '1677-09-21 00:12:43.145224192', + '2262-04-11 23:47:16.854775807', + ]) + self.assertEqual( + variant_get(boundaries, '$', pa.timestamp('ns')) + .cast(pa.int64()).to_pylist(), + [-(1 << 63), (1 << 63) - 1], + ) + + invalid = ( + '1677-09-21 00:12:43.145224191', + '2262-04-11 23:47:16.854775808', + '9999-12-31 23:59:59.999999999', + 'today', + 'now', + ) + for value in invalid: + with self.subTest(value=value): + with self.assertRaisesRegex(ValueError, "Invalid cast"): + variant_get( + _variants([value]), '$', pa.timestamp('ns')) + def test_variant_null_remains_arrow_null(self): column = _variants([None, {'value': None}, {'value': 1.0}]) From d47c1bbfb9430c3d85ca34244c1abc03698cf4a1 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 20:10:23 -0700 Subject: [PATCH 18/23] [python] Clarify JDK 8 floating-point formatting --- ...java_floating.py => _java_float_format.py} | 28 ++++--- paimon-python/pypaimon/data/variant_path.py | 21 ++++-- .../pypaimon/tests/java_float_format_test.py | 75 +++++++++++++++++++ .../pypaimon/tests/variant_path_test.py | 19 ----- 4 files changed, 110 insertions(+), 33 deletions(-) rename paimon-python/pypaimon/data/{_java_floating.py => _java_float_format.py} (89%) create mode 100644 paimon-python/pypaimon/tests/java_float_format_test.py diff --git a/paimon-python/pypaimon/data/_java_floating.py b/paimon-python/pypaimon/data/_java_float_format.py similarity index 89% rename from paimon-python/pypaimon/data/_java_floating.py rename to paimon-python/pypaimon/data/_java_float_format.py index ac7d0f9b1f88..461f724d5b53 100644 --- a/paimon-python/pypaimon/data/_java_floating.py +++ b/paimon-python/pypaimon/data/_java_float_format.py @@ -14,20 +14,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Format IEEE floats like ``Float/Double.toString`` on the JDK 8 baseline.""" +"""Format IEEE values like ``Float/Double.toString`` on JDK 8. + +Paimon VARIANT casts follow Java cast semantics, but Python's floating-point +formatting differs for some IEEE values. This module preserves the digit +generation behavior of OpenJDK 8 ``sun.misc.FloatingDecimal``. Do not replace +it with ``str`` or ``repr`` without cross-checking the JDK 8 output. +""" import struct -def java_floating_text(value, single_precision=False): - """Return the legacy Java decimal representation of an IEEE value.""" - if single_precision: - bits = struct.unpack('>I', struct.pack('>f', value))[0] - width, fraction_width, exponent_width, bias = 32, 23, 8, 127 - else: - bits = struct.unpack('>Q', struct.pack('>d', value))[0] - width, fraction_width, exponent_width, bias = 64, 52, 11, 1023 +def java_float_to_string(value): + """Return the JDK 8 ``Float.toString`` representation.""" + bits = struct.unpack('>I', struct.pack('>f', value))[0] + return _format_bits(bits, 32, 23, 8, 127) + + +def java_double_to_string(value): + """Return the JDK 8 ``Double.toString`` representation.""" + bits = struct.unpack('>Q', struct.pack('>d', value))[0] + return _format_bits(bits, 64, 52, 11, 1023) + +def _format_bits(bits, width, fraction_width, exponent_width, bias): negative = bool(bits >> (width - 1)) exponent = ((bits >> fraction_width) & ((1 << exponent_width) - 1)) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 533577f4f1da..1108ce080323 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -29,7 +29,10 @@ import numpy as np import pyarrow as pa -from pypaimon.data._java_floating import java_floating_text +from pypaimon.data._java_float_format import ( + java_double_to_string, + java_float_to_string, +) from pypaimon.data._variant_binary import ( _ARRAY, _OBJECT, @@ -845,8 +848,16 @@ def _decimal_text(value): return text[1:] if value == 0 and text.startswith('-') else text -def _floating_text(value, single_precision=False): - return java_floating_text(value, single_precision) +def _floating_text(value): + return java_double_to_string(value) + + +def _variant_floating_text(value, type_info): + if type_info == _FLOAT: + return java_float_to_string(value) + if type_info == _DOUBLE: + return java_double_to_string(value) + raise ValueError("not a floating-point VARIANT") def _json_text(value): @@ -922,7 +933,7 @@ def _variant_json(value, metadata, pos, limit, keys=None): decoded = GenericVariant(bytes(value[pos:limit]), metadata).to_python() if basic_type == _PRIMITIVE and type_info in (_FLOAT, _DOUBLE): - text = _floating_text(decoded, type_info == _FLOAT) + text = _variant_floating_text(decoded, type_info) return text if math.isfinite(decoded) else json.dumps(text) return _json_text(decoded) @@ -1164,7 +1175,7 @@ def _decode_scalar(value, metadata: bytes, pos: int, target_type): or pa.types.is_large_string(target_type)) and (value[pos] & 0x3) == _PRIMITIVE and type_info in (_FLOAT, _DOUBLE)): - return _floating_text(decoded, type_info == _FLOAT) + return _variant_floating_text(decoded, type_info) return _cast_python(decoded, target_type) diff --git a/paimon-python/pypaimon/tests/java_float_format_test.py b/paimon-python/pypaimon/tests/java_float_format_test.py new file mode 100644 index 000000000000..7f437a30b3ca --- /dev/null +++ b/paimon-python/pypaimon/tests/java_float_format_test.py @@ -0,0 +1,75 @@ +# 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. + +import struct +import unittest + +from pypaimon.data._java_float_format import ( + java_double_to_string, + java_float_to_string, +) + + +def _float_from_bits(bits): + return struct.unpack('>f', struct.pack('>I', bits))[0] + + +def _double_from_bits(bits): + return struct.unpack('>d', struct.pack('>Q', bits))[0] + + +class TestJavaFloatFormat(unittest.TestCase): + + def test_float_to_string_matches_jdk8(self): + # Generated with OpenJDK 8u492 Float.toString. + samples = ( + (0x00000000, '0.0'), + (0x80000000, '-0.0'), + (0x00000001, '1.4E-45'), + (0x007FFFFF, '1.1754942E-38'), + (0x00800000, '1.17549435E-38'), + (0xD44F9F82, '-3.56693731E12'), + (0xEA75E34D, '-7.4315055E25'), + (0xE8FDF7D7, '-9.5946444E24'), + (0x7F7FFFFF, '3.4028235E38'), + (0x7F800000, 'Infinity'), + (0xFF800000, '-Infinity'), + (0x7FC00000, 'NaN'), + ) + for bits, expected in samples: + with self.subTest(bits=hex(bits)): + self.assertEqual( + java_float_to_string(_float_from_bits(bits)), expected) + + def test_double_to_string_matches_jdk8(self): + # Generated with OpenJDK 8u492 Double.toString. + samples = ( + (0x0000000000000000, '0.0'), + (0x8000000000000000, '-0.0'), + (0x0000000000000001, '4.9E-324'), + (0x000FFFFFFFFFFFFF, '2.225073858507201E-308'), + (0x0010000000000000, '2.2250738585072014E-308'), + (0x439F4B86CD6A5E0C, '5.6376106381000781E17'), + (0x439DDD7467AF36D9, '5.3800104640575648E17'), + (0x7FEFFFFFFFFFFFFF, '1.7976931348623157E308'), + (0x7FF0000000000000, 'Infinity'), + (0xFFF0000000000000, '-Infinity'), + (0x7FF8000000000000, 'NaN'), + ) + for bits, expected in samples: + with self.subTest(bits=hex(bits)): + self.assertEqual( + java_double_to_string(_double_from_bits(bits)), expected) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index 61404cd38cd8..5b2fc967b2c0 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -15,7 +15,6 @@ # limitations under the License. import datetime -import struct import unittest from decimal import Decimal from unittest.mock import patch @@ -221,24 +220,6 @@ def test_get_matches_java_numeric_strings(self): ['1234567', '12345678901234'], ) - legacy_floats = _float_variants([ - struct.unpack('>f', struct.pack('>I', bits))[0] - for bits in (0xD44F9F82, 0xEA75E34D, 0xE8FDF7D7) - ]) - self.assertEqual( - variant_get(legacy_floats, '$', pa.string()).to_pylist(), - ['-3.56693731E12', '-7.4315055E25', '-9.5946444E24'], - ) - - legacy_doubles = _variants([ - struct.unpack('>d', struct.pack('>Q', bits))[0] - for bits in (0x439F4B86CD6A5E0C, 0x439DDD7467AF36D9) - ]) - self.assertEqual( - variant_get(legacy_doubles, '$', pa.string()).to_pylist(), - ['5.6376106381000781E17', '5.3800104640575648E17'], - ) - def test_get_matches_java_numeric_casts(self): column = _variants([ {'value': 1e20}, From b02316dd0d070ed7d29d619500dc5d4219aa7ac9 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 21:21:57 -0700 Subject: [PATCH 19/23] [python] Scope floating formatting to VARIANT --- ...java_float_format.py => _variant_float_format.py} | 6 +++--- paimon-python/pypaimon/data/variant_path.py | 12 ++++++------ ...t_format_test.py => variant_float_format_test.py} | 12 ++++++------ 3 files changed, 15 insertions(+), 15 deletions(-) rename paimon-python/pypaimon/data/{_java_float_format.py => _variant_float_format.py} (98%) rename paimon-python/pypaimon/tests/{java_float_format_test.py => variant_float_format_test.py} (89%) diff --git a/paimon-python/pypaimon/data/_java_float_format.py b/paimon-python/pypaimon/data/_variant_float_format.py similarity index 98% rename from paimon-python/pypaimon/data/_java_float_format.py rename to paimon-python/pypaimon/data/_variant_float_format.py index 461f724d5b53..9df730cea9f3 100644 --- a/paimon-python/pypaimon/data/_java_float_format.py +++ b/paimon-python/pypaimon/data/_variant_float_format.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Format IEEE values like ``Float/Double.toString`` on JDK 8. +"""Format IEEE values for Paimon VARIANT casts. Paimon VARIANT casts follow Java cast semantics, but Python's floating-point formatting differs for some IEEE values. This module preserves the digit @@ -25,13 +25,13 @@ import struct -def java_float_to_string(value): +def format_float32(value): """Return the JDK 8 ``Float.toString`` representation.""" bits = struct.unpack('>I', struct.pack('>f', value))[0] return _format_bits(bits, 32, 23, 8, 127) -def java_double_to_string(value): +def format_float64(value): """Return the JDK 8 ``Double.toString`` representation.""" bits = struct.unpack('>Q', struct.pack('>d', value))[0] return _format_bits(bits, 64, 52, 11, 1023) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 1108ce080323..f8f0f2d52233 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -29,9 +29,9 @@ import numpy as np import pyarrow as pa -from pypaimon.data._java_float_format import ( - java_double_to_string, - java_float_to_string, +from pypaimon.data._variant_float_format import ( + format_float32, + format_float64, ) from pypaimon.data._variant_binary import ( _ARRAY, @@ -849,14 +849,14 @@ def _decimal_text(value): def _floating_text(value): - return java_double_to_string(value) + return format_float64(value) def _variant_floating_text(value, type_info): if type_info == _FLOAT: - return java_float_to_string(value) + return format_float32(value) if type_info == _DOUBLE: - return java_double_to_string(value) + return format_float64(value) raise ValueError("not a floating-point VARIANT") diff --git a/paimon-python/pypaimon/tests/java_float_format_test.py b/paimon-python/pypaimon/tests/variant_float_format_test.py similarity index 89% rename from paimon-python/pypaimon/tests/java_float_format_test.py rename to paimon-python/pypaimon/tests/variant_float_format_test.py index 7f437a30b3ca..061cf3e5d8e1 100644 --- a/paimon-python/pypaimon/tests/java_float_format_test.py +++ b/paimon-python/pypaimon/tests/variant_float_format_test.py @@ -17,9 +17,9 @@ import struct import unittest -from pypaimon.data._java_float_format import ( - java_double_to_string, - java_float_to_string, +from pypaimon.data._variant_float_format import ( + format_float32, + format_float64, ) @@ -31,7 +31,7 @@ def _double_from_bits(bits): return struct.unpack('>d', struct.pack('>Q', bits))[0] -class TestJavaFloatFormat(unittest.TestCase): +class TestVariantFloatFormat(unittest.TestCase): def test_float_to_string_matches_jdk8(self): # Generated with OpenJDK 8u492 Float.toString. @@ -52,7 +52,7 @@ def test_float_to_string_matches_jdk8(self): for bits, expected in samples: with self.subTest(bits=hex(bits)): self.assertEqual( - java_float_to_string(_float_from_bits(bits)), expected) + format_float32(_float_from_bits(bits)), expected) def test_double_to_string_matches_jdk8(self): # Generated with OpenJDK 8u492 Double.toString. @@ -72,4 +72,4 @@ def test_double_to_string_matches_jdk8(self): for bits, expected in samples: with self.subTest(bits=hex(bits)): self.assertEqual( - java_double_to_string(_double_from_bits(bits)), expected) + format_float64(_double_from_bits(bits)), expected) From a5ddb1e309868196d6b7e5ff8b6f01be5eeaf864 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 21:34:32 -0700 Subject: [PATCH 20/23] [python] Harden VARIANT path edge cases --- .../pypaimon/data/_variant_float_format.py | 3 +- .../pypaimon/data/generic_variant.py | 5 +- paimon-python/pypaimon/data/variant_path.py | 70 ++++++++++++------- .../pypaimon/tests/variant_path_test.py | 37 ++++++++++ paimon-python/pypaimon/tests/variant_test.py | 7 +- 5 files changed, 93 insertions(+), 29 deletions(-) diff --git a/paimon-python/pypaimon/data/_variant_float_format.py b/paimon-python/pypaimon/data/_variant_float_format.py index 9df730cea9f3..5ab5b7707d3e 100644 --- a/paimon-python/pypaimon/data/_variant_float_format.py +++ b/paimon-python/pypaimon/data/_variant_float_format.py @@ -19,7 +19,8 @@ Paimon VARIANT casts follow Java cast semantics, but Python's floating-point formatting differs for some IEEE values. This module preserves the digit generation behavior of OpenJDK 8 ``sun.misc.FloatingDecimal``. Do not replace -it with ``str`` or ``repr`` without cross-checking the JDK 8 output. +it with ``str``, ``repr``, or newer JDK output: JDK 8 is the compatibility +contract. """ import struct diff --git a/paimon-python/pypaimon/data/generic_variant.py b/paimon-python/pypaimon/data/generic_variant.py index a7edcd048848..030bc9ab5ccc 100644 --- a/paimon-python/pypaimon/data/generic_variant.py +++ b/paimon-python/pypaimon/data/generic_variant.py @@ -327,10 +327,11 @@ def append_decimal(self, d): def append_decimal_unscaled(self, unscaled, precision, scale): if not 0 <= scale <= _MAX_DECIMAL16_PRECISION: raise ValueError(f'Unsupported VARIANT decimal scale: {scale}') - if (not 0 < precision <= _MAX_DECIMAL16_PRECISION - or not -(1 << 127) <= unscaled < (1 << 127)): + if not 0 < precision <= _MAX_DECIMAL16_PRECISION: raise ValueError( f'Unsupported VARIANT decimal precision: {precision}') + if not -(1 << 127) <= unscaled < (1 << 127): + raise ValueError('VARIANT decimal value exceeds 128 bits') if scale <= _MAX_DECIMAL4_PRECISION and precision <= _MAX_DECIMAL4_PRECISION: self._write_byte(_primitive_header(_DECIMAL4)) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index f8f0f2d52233..c9fdcc1f7bdc 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -286,9 +286,7 @@ def _path_positions( value: bytes, metadata: bytes, paths: Sequence[_Path], - plans, ) -> Sequence[Optional[int]]: - del plans root_size = _checked_value_size(value, 0) if root_size != len(value): _malformed("trailing bytes after root value") @@ -349,6 +347,7 @@ def _replace_path( path: _Path, replacement: bytes, limit=None, + key_ids=None, ) -> bytes: limit = len(value) if limit is None else limit value_end = pos + _checked_value_size(value, pos, limit) @@ -359,7 +358,9 @@ def _replace_path( if kind == 'key': if (value[pos] & 0x3) != _OBJECT: raise ValueError("VARIANT path expects an object") - key_id = _metadata_key_ids(metadata).get(segment) + if key_ids is None: + key_ids = _metadata_key_ids(metadata) + key_id = key_ids.get(segment) if key_id is None: raise ValueError(f"VARIANT path does not exist: {segment}") size, id_size, id_start, data_start, offsets, _ = ( @@ -382,7 +383,7 @@ def _replace_path( if i == slot: child = _replace_path( value, metadata, child_pos, path[1:], replacement, - child_end) + child_end, key_ids) children.append(child) return _build_object_value(list(zip(ids, children))) @@ -400,7 +401,7 @@ def _replace_path( if i == segment: child = _replace_path( value, metadata, child_pos, path[1:], replacement, - child_end) + child_end, key_ids) children.append(child) return _build_array_value(children) @@ -827,7 +828,7 @@ def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): for row in set().union(*slow_by_path): value = values.view(row) row_metadata = bytes(metadata.view(row)) - positions = _path_positions(value, row_metadata, parsed_paths, None) + positions = _path_positions(value, row_metadata, parsed_paths) for index, (pos, target_type) in enumerate( zip(positions, target_types)): if row not in slow_by_path[index] or pos is None: @@ -868,7 +869,8 @@ def _json_text(value): if isinstance(value, decimal.Decimal): return _decimal_text(value) if isinstance(value, float): - return str(value) if math.isfinite(value) else json.dumps(str(value)) + text = _floating_text(value) + return text if math.isfinite(value) else json.dumps(text) if isinstance(value, int): return str(value) if isinstance(value, str): @@ -1213,15 +1215,27 @@ def _rebuilt_chunk( ) +def _rebuilt_offsets(lengths, value_format): + total = sum(int(length) for length in lengths) + maximum = np.iinfo(np.dtype(value_format)).max + if total > maximum: + kind = 'Binary' if value_format == 'd', struct.pack('>Q', 0x439F4B86CD6A5E0C))[0] + column = _variants([{ + 'nested': {'finite': value, 'infinite': float('inf')}, + }]) + + result = variant_get( + column, '$', pa.struct([('nested', pa.string())])) + + self.assertEqual(result.to_pylist(), [{ + 'nested': ( + '{"finite":5.6376106381000781E17,' + '"infinite":"Infinity"}' + ), + }]) + def test_get_matches_java_numeric_casts(self): column = _variants([ {'value': 1e20}, @@ -241,6 +260,8 @@ def test_get_matches_java_numeric_casts(self): with self.assertRaisesRegex(ValueError, "Invalid cast"): variant_get(invalid.slice(1, 1), '$.value', pa.timestamp('us')) + # Java casts LONG as epoch seconds, but parses numeric STRING using + # the target timestamp precision. timestamp = variant_get( _variants([{'value': 1}]), '$.value', pa.timestamp('us')) self.assertEqual( @@ -785,6 +806,22 @@ def test_missing_path_is_noop_or_strict_error(self): variant_replace( column, '$.value', replacement, strict=True) + def test_all_missing_paths_reuse_input_buffers(self): + column = _variants([{'other': 1.0}, {'other': 2.0}]) + + result = variant_replace( + column, '$.missing', pa.scalar(3.0, type=pa.float64())) + + self.assertIs(result, column) + + def test_rebuilt_binary_offsets_reject_overflow(self): + self.assertEqual( + _rebuilt_offsets(np.array([2, 3]), ' Date: Mon, 10 Aug 2026 22:50:30 -0700 Subject: [PATCH 21/23] [python] Align VARIANT casts and missing paths --- paimon-python/pypaimon/data/variant_path.py | 177 ++++++++++++++---- .../pypaimon/tests/variant_path_test.py | 86 ++++++++- 2 files changed, 221 insertions(+), 42 deletions(-) diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index c9fdcc1f7bdc..3bf7456b6ece 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -72,6 +72,14 @@ r'(?:(?: (\d{1,2}):(\d{1,2}):(\d{1,2})(?:\.(\d{1,9}))?)' r'|(?:T(\d{1,2}):(\d{1,2})(?::(\d{1,2})(?:\.(\d{1,9}))?)?))?$' ) +_JAVA_DECIMAL_FLOAT_PATTERN = re.compile( + r'^[+-]?(?:(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+)' + r'(?:[eE][+-]?[0-9]+)?)[fFdD]?$' +) +_JAVA_HEX_FLOAT_PATTERN = re.compile( + r'^[+-]?0[xX](?:[0-9a-fA-F]+(?:\.[0-9a-fA-F]*)?' + r'|\.[0-9a-fA-F]+)[pP][+-]?[0-9]+[fFdD]?$' +) _Path = Tuple[Tuple[str, object], ...] _SLOW_PATH_ROWS = 64 @@ -625,6 +633,10 @@ def _vectorized_path_positions( try: for parent_node, kind, segment in nodes[1:]: parent = positions[parent_node] + if parent is None: + positions.append(None) + limits.append(None) + continue parent_ends = row_starts + limits[parent_node] absolute_parent = row_starts + parent if (np.any(absolute_parent < row_starts) @@ -638,7 +650,9 @@ def _vectorized_path_positions( return None key_id = key_ids.get(segment) if key_id is None: - return None + positions.append(None) + limits.append(None) + continue first_layout = _checked_object_layout( first_value, int(parent[0]), int(limits[parent_node][0])) size, id_size, id_start, _, first_offsets, _ = first_layout @@ -767,6 +781,9 @@ def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): uniform = True for pos, limit, target_type in zip( positions, limits, target_types): + if pos is None: + results.append(pa.nulls(len(chunk), type=target_type)) + continue absolute = row_starts + pos headers = data[absolute] if np.all(headers == _primitive_header(_DOUBLE)): @@ -799,6 +816,8 @@ def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): for planned in plans: rows, row_starts, data, positions, limits = planned for index, (pos, limit) in enumerate(zip(positions, limits)): + if pos is None: + continue absolute = row_starts + pos headers = data[absolute] handled = np.zeros(len(rows), dtype=bool) @@ -861,6 +880,22 @@ def _variant_floating_text(value, type_info): raise ValueError("not a floating-point VARIANT") +def _timestamp_text(value): + text = ( + f'{value.year:04d}-{value.month:02d}-{value.day:02d} ' + f'{value.hour:02d}:{value.minute:02d}:{value.second:02d}' + ) + if value.microsecond: + text += f'.{value.microsecond:06d}'.rstrip('0') + offset = value.utcoffset() + if offset is not None: + seconds = int(offset.total_seconds()) + sign = '+' if seconds >= 0 else '-' + minutes = abs(seconds) // 60 + text += f'{sign}{minutes // 60:02d}:{minutes % 60:02d}' + return text + + def _json_text(value): if value is None: return 'null' @@ -877,7 +912,9 @@ def _json_text(value): return json.dumps(value, ensure_ascii=False, separators=(',', ':')) if isinstance(value, bytes): return json.dumps(base64.b64encode(value).decode('ascii')) - if isinstance(value, (datetime.date, datetime.datetime)): + if isinstance(value, datetime.datetime): + return json.dumps(_timestamp_text(value)) + if isinstance(value, datetime.date): return json.dumps(value.isoformat()) if isinstance(value, dict): return '{' + ','.join( @@ -999,6 +1036,22 @@ def _cast_string_integer(value, target_type): return number +def _cast_string_floating(value): + text = value.strip() + if re.match(r'^[+-]?(?:NaN|Infinity)$', text): + if text.endswith('NaN'): + return math.copysign(float('nan'), -1.0 if text[0:1] == '-' else 1.0) + return float('-inf') if text.startswith('-') else float('inf') + if (_JAVA_DECIMAL_FLOAT_PATTERN.match(text) + or _JAVA_HEX_FLOAT_PATTERN.match(text)): + if text[-1:] in 'fFdD': + text = text[:-1] + if '0x' in text.lower(): + return float.fromhex(text) + return float(text) + raise ValueError + + def _parse_date(value): value = value.strip() if value and (value.isdigit() @@ -1055,40 +1108,6 @@ def _cast_python(value, target_type): if value is None or pa.types.is_null(target_type): return None try: - if pa.types.is_struct(target_type): - if isinstance(value, str): - value = json.loads(value) - if not isinstance(value, dict): - raise TypeError - return { - field.name: ( - None if field.name not in value - else _cast_python(value[field.name], field.type) - ) - for field in target_type - } - if (pa.types.is_list(target_type) - or pa.types.is_large_list(target_type) - or pa.types.is_fixed_size_list(target_type)): - if isinstance(value, str): - value = json.loads(value) - if not isinstance(value, list): - raise TypeError - return [ - _cast_python(child, target_type.value_type) - for child in value - ] - if pa.types.is_map(target_type): - if isinstance(value, str): - value = json.loads(value) - if not isinstance(value, dict) or not ( - pa.types.is_string(target_type.key_type) - or pa.types.is_large_string(target_type.key_type)): - raise TypeError - return [ - (key, _cast_python(child, target_type.item_type)) - for key, child in value.items() - ] if pa.types.is_string(target_type) or pa.types.is_large_string( target_type): if isinstance(value, (dict, list)): @@ -1125,6 +1144,8 @@ def _cast_python(value, target_type): if pa.types.is_floating(target_type): if not isinstance(value, (bool, int, float, decimal.Decimal, str)): raise TypeError + if isinstance(value, str): + return _cast_string_floating(value) return float(value) if pa.types.is_decimal(target_type): if not isinstance(value, (bool, int, float, decimal.Decimal, str)): @@ -1162,13 +1183,78 @@ def _cast_python(value, target_type): raise ValueError(f"Invalid cast {value!r} to {target_type}") +def _variant_object_children(value, metadata, pos, end): + size, id_size, id_start, data_start, offsets, _ = ( + _checked_object_layout(value, pos, end)) + keys = { + key_id: key for key, key_id in _metadata_key_ids(metadata).items() + } + children = {} + for slot in range(size): + key_id = _read_unsigned(value, id_start + slot * id_size, id_size) + if key_id not in keys: + _malformed("object key is missing from metadata") + children[keys[key_id]] = _checked_object_child_bounds( + value, data_start, offsets, slot) + return children + + +def _variant_array_children(value, pos, end): + size, data_start, offsets, _ = _checked_array_layout( + value, pos, end) + children = [] + for index in range(size): + child_start = data_start + offsets[index] + child_end = data_start + offsets[index + 1] + if _checked_value_size(value, child_start, child_end) != ( + child_end - child_start): + _malformed("child size does not match container offsets") + children.append((child_start, child_end)) + return children + + def _decode_scalar(value, metadata: bytes, pos: int, target_type): size = _checked_value_size(value, pos) end = pos + size + basic_type = value[pos] & 0x3 type_info = (value[pos] >> 2) & 0x3F + if pa.types.is_struct(target_type): + if basic_type != _OBJECT: + raise ValueError(f"Invalid cast VARIANT to {target_type}") + children = _variant_object_children( + value, metadata, pos, end) + return { + field.name: ( + None if field.name not in children + else _decode_scalar( + value, metadata, children[field.name][0], field.type) + ) + for field in target_type + } + if (pa.types.is_list(target_type) + or pa.types.is_large_list(target_type) + or pa.types.is_fixed_size_list(target_type)): + if basic_type != _ARRAY: + raise ValueError(f"Invalid cast VARIANT to {target_type}") + return [ + _decode_scalar(value, metadata, child_pos, target_type.value_type) + for child_pos, _ in _variant_array_children( + value, pos, end) + ] + if pa.types.is_map(target_type): + if (basic_type != _OBJECT + or not (pa.types.is_string(target_type.key_type) + or pa.types.is_large_string(target_type.key_type))): + raise ValueError(f"Invalid cast VARIANT to {target_type}") + children = _variant_object_children( + value, metadata, pos, end) + return [ + (key, _decode_scalar( + value, metadata, child_pos, target_type.item_type)) + for key, (child_pos, _) in children.items() + ] if (pa.types.is_string(target_type) or pa.types.is_large_string(target_type)): - basic_type = value[pos] & 0x3 if basic_type in (_OBJECT, _ARRAY): return _variant_json(value, metadata, pos, end) selected = bytes(value[pos:end]) @@ -1382,8 +1468,16 @@ def _vectorized_replace_chunk( rows, row_starts, _, positions, limits = planned replacements = [] compatible = np.ones(len(rows), dtype=bool) - for (_, _, provider), pos, limit in zip( + has_replacement = False + for (path, _, provider), pos, limit in zip( parsed, positions, limits): + if pos is None: + if strict: + raise ValueError( + f"VARIANT path does not exist: {path}") + replacements.append(None) + continue + has_replacement = True replacement, replacement_valid = provider.numpy_values( global_row, len(chunk), @@ -1395,6 +1489,8 @@ def _vectorized_replace_chunk( & (absolute + provider._fixed_size == row_starts + limit) ) replacements.append((pos, provider, replacement)) + if not has_replacement: + continue slow_rows.update(int(row) for row in rows[~compatible]) if not np.any(compatible): continue @@ -1403,7 +1499,10 @@ def _vectorized_replace_chunk( output_data = np.frombuffer(data, dtype=np.uint8) relative_starts = row_starts - data_start compatible_rows = rows[compatible] - for pos, provider, replacement in replacements: + for item in replacements: + if item is None: + continue + pos, provider, replacement = item absolute = (relative_starts + pos)[compatible] output_data[absolute] = provider._type_header value_size = provider._fixed_size - 1 diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index 6e06a5816a04..e3760481f6ae 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -239,6 +239,59 @@ def test_nested_json_uses_variant_float_text(self): ), }]) + metadata = GenericVariant.from_python({'value': 0}).metadata() + nested_float = GenericVariant.to_arrow_array([ + GenericVariant( + _build_object_value([( + 0, _encode_scalar_to_value_bytes(1.2, pa.float32()), + )]), + metadata, + ), + ]) + self.assertEqual( + variant_get( + nested_float, + '$', + pa.struct([('value', pa.string())]), + ).to_pylist(), + [{'value': '1.2'}], + ) + + def test_string_floating_cast_matches_java_parser(self): + for target_type in (pa.float32(), pa.float64()): + for value in ('inf', 'nan', '1_0'): + with self.subTest(target_type=target_type, value=value): + with self.assertRaisesRegex(ValueError, "Invalid cast"): + variant_get(_variants([value]), '$', target_type) + + parsed = variant_get( + _variants(['NaN', 'Infinity', '-Infinity', '1.25f', '0x1.0p0']), + '$', + pa.float64(), + ).to_pylist() + self.assertTrue(np.isnan(parsed[0])) + self.assertEqual( + parsed[1:], [float('inf'), float('-inf'), 1.25, 1.0]) + + def test_container_timestamp_json_uses_space(self): + metadata = GenericVariant.from_python({'ts': 0}).metadata() + timestamp = datetime.datetime(2026, 8, 10, 12, 34, 56, 123000) + column = GenericVariant.to_arrow_array([ + GenericVariant( + _build_object_value([( + 0, + _encode_scalar_to_value_bytes( + timestamp, pa.timestamp('us')), + )]), + metadata, + ), + ]) + + self.assertEqual( + variant_get(column, '$', pa.string()).to_pylist(), + ['{"ts":"2026-08-10 12:34:56.123"}'], + ) + def test_get_matches_java_numeric_casts(self): column = _variants([ {'value': 1e20}, @@ -807,12 +860,39 @@ def test_missing_path_is_noop_or_strict_error(self): column, '$.value', replacement, strict=True) def test_all_missing_paths_reuse_input_buffers(self): - column = _variants([{'other': 1.0}, {'other': 2.0}]) + column = _variants([ + {'other': float(index)} for index in range(4096) + ]) - result = variant_replace( - column, '$.missing', pa.scalar(3.0, type=pa.float64())) + with patch( + 'pypaimon.data.variant_path._path_positions', + wraps=_path_positions, + ) as slow_path: + current = variant_get(column, '$.missing', pa.float64()) + result = variant_replace( + column, '$.missing', pa.scalar(3.0, type=pa.float64())) + self.assertEqual(current.null_count, len(column)) self.assertIs(result, column) + slow_path.assert_not_called() + + current = variant_get(column, { + '$.other': pa.float64(), + '$.missing': pa.float64(), + }) + result = variant_replace(column, { + '$.other': pa.scalar(-1.0), + '$.missing': pa.scalar(3.0), + }) + self.assertEqual(current['$.missing'].null_count, len(column)) + self.assertEqual(_decode(result.slice(0, 1)), [{'other': -1.0}]) + with self.assertRaisesRegex(ValueError, "path does not exist"): + variant_replace( + column, + '$.missing', + pa.scalar(3.0, type=pa.float64()), + strict=True, + ) def test_rebuilt_binary_offsets_reject_overflow(self): self.assertEqual( From 09158322d22590a37d02ac60d873f1511dd3c32f Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Mon, 10 Aug 2026 23:54:24 -0700 Subject: [PATCH 22/23] [python] Scope VARIANT path updates to floating types --- docs/docs/pypaimon/python-api.mdx | 14 +- .../pypaimon/data/_variant_float_format.py | 237 ----- .../pypaimon/data/generic_variant.py | 32 +- paimon-python/pypaimon/data/variant_path.py | 707 ++------------ .../pypaimon/data/variant_shredding.py | 57 +- .../tests/variant_float_format_test.py | 75 -- .../pypaimon/tests/variant_path_test.py | 920 ++++-------------- paimon-python/pypaimon/tests/variant_test.py | 7 +- 8 files changed, 273 insertions(+), 1776 deletions(-) delete mode 100644 paimon-python/pypaimon/data/_variant_float_format.py delete mode 100644 paimon-python/pypaimon/tests/variant_float_format_test.py diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index 2e782adc9646..98be214eaa25 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -1126,10 +1126,10 @@ Supported Paimon type strings for shredded sub-fields: `BOOLEAN`, `INT`, `BIGINT -### VARIANT Path Updates +### VARIANT FLOAT/DOUBLE Path Updates -Read a path as an Arrow array, use Arrow compute, and replace its existing -values without decoding unrelated fields: +Read existing FLOAT or DOUBLE paths as Arrow arrays, use Arrow compute, and +replace them without decoding unrelated fields: ```python import pyarrow as pa @@ -1142,10 +1142,10 @@ updated_payload = variant_replace( payload, '$.velocity.y', pc.negate(current)) ``` -`variant_get` returns NULL for a missing path. `variant_replace` accepts an -Arrow Scalar, Array, or ChunkedArray. A missing path is unchanged unless -`strict=True` is specified. Pass `{path: type}` to `variant_get` and -`{path: values}` to `variant_replace` to process multiple paths in one pass. +The Arrow type passed to `variant_get` and `variant_replace` must match the +stored FLOAT or DOUBLE type; these APIs do not cast VARIANT values. Missing +paths read as NULL and remain unchanged unless `strict=True` is specified. +Pass mappings to process multiple paths in one pass. **`GenericVariant` API:** diff --git a/paimon-python/pypaimon/data/_variant_float_format.py b/paimon-python/pypaimon/data/_variant_float_format.py deleted file mode 100644 index 5ab5b7707d3e..000000000000 --- a/paimon-python/pypaimon/data/_variant_float_format.py +++ /dev/null @@ -1,237 +0,0 @@ -# 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. - -"""Format IEEE values for Paimon VARIANT casts. - -Paimon VARIANT casts follow Java cast semantics, but Python's floating-point -formatting differs for some IEEE values. This module preserves the digit -generation behavior of OpenJDK 8 ``sun.misc.FloatingDecimal``. Do not replace -it with ``str``, ``repr``, or newer JDK output: JDK 8 is the compatibility -contract. -""" - -import struct - - -def format_float32(value): - """Return the JDK 8 ``Float.toString`` representation.""" - bits = struct.unpack('>I', struct.pack('>f', value))[0] - return _format_bits(bits, 32, 23, 8, 127) - - -def format_float64(value): - """Return the JDK 8 ``Double.toString`` representation.""" - bits = struct.unpack('>Q', struct.pack('>d', value))[0] - return _format_bits(bits, 64, 52, 11, 1023) - - -def _format_bits(bits, width, fraction_width, exponent_width, bias): - negative = bool(bits >> (width - 1)) - exponent = ((bits >> fraction_width) - & ((1 << exponent_width) - 1)) - fraction = bits & ((1 << fraction_width) - 1) - if exponent == (1 << exponent_width) - 1: - if fraction: - return 'NaN' - return '-Infinity' if negative else 'Infinity' - if exponent == 0: - if fraction == 0: - return '-0.0' if negative else '0.0' - highest_bit = fraction.bit_length() - 1 - shift = fraction_width - highest_bit - significand = fraction << shift - binary_exponent = 1 - shift - bias - significant_bits = highest_bit + 1 - else: - significand = (1 << fraction_width) | fraction - binary_exponent = exponent - bias - significant_bits = fraction_width + 1 - - trailing_zeros = (significand & -significand).bit_length() - 1 - fraction_bits = fraction_width + 1 - trailing_zeros - tiny_bits = max(0, fraction_bits - binary_exponent - 1) - if tiny_bits == 0 and -21 <= binary_exponent <= 62: - return _format_small_integer( - negative, - significand, - fraction_width, - binary_exponent, - significant_bits, - ) - - decimal_exponent = _estimate_decimal_exponent( - significand, binary_exponent, fraction_width) - base5 = max(0, -decimal_exponent) - base2 = base5 + tiny_bits + binary_exponent - scale5 = max(0, decimal_exponent) - scale2 = scale5 + tiny_bits - margin2 = base2 - significant_bits - reduced = significand >> trailing_zeros - base2 -= fraction_bits - 1 - common2 = min(base2, scale2) - base2 -= common2 - scale2 -= common2 - margin2 -= common2 - if fraction_bits == 1: - margin2 -= 1 - if margin2 < 0: - base2 -= margin2 - scale2 -= margin2 - margin2 = 0 - - base = reduced * 5 ** base5 << base2 - scale = 5 ** scale5 << scale2 - margin = 5 ** base5 << margin2 - base_bits = fraction_bits + base2 + _five_bits(base5) - ten_scale_bits = scale2 + 1 + _five_bits(scale5 + 1) - arithmetic_width = ( - 32 if base_bits < 32 and ten_scale_bits < 32 - else 64 if base_bits < 64 and ten_scale_bits < 64 - else None - ) - ten_scale = scale * 10 - - digits = [] - iteration = 0 - while True: - digit, base = divmod(base, scale) - base *= 10 - margin *= 10 - if arithmetic_width is not None: - margin = _signed(margin, arithmetic_width) - if arithmetic_width is None: - low = base < margin - high = base + margin >= ten_scale - elif iteration == 0 or margin > 0: - low = base < margin - high = _signed(base + margin, arithmetic_width) > ten_scale - else: - low = high = True - - if iteration == 0 and digit == 0 and not high: - decimal_exponent -= 1 - else: - digits.append(digit) - if iteration == 0 and ( - decimal_exponent < -3 or decimal_exponent >= 8): - low = high = False - iteration += 1 - if low or high: - break - - if high: - round_up = not low - if low: - if arithmetic_width is None: - difference = 2 * base - ten_scale - else: - difference = _signed( - _signed(base << 1, arithmetic_width) - ten_scale, - arithmetic_width, - ) - round_up = difference > 0 or ( - difference == 0 and digits[-1] & 1) - if round_up: - decimal_exponent = _round_up(digits, decimal_exponent) - return _format(negative, digits, decimal_exponent + 1) - - -def _format_small_integer( - negative, - significand, - fraction_width, - binary_exponent, - significant_bits, -): - insignificant = 0 - if binary_exponent > significant_bits: - power = binary_exponent - significant_bits - 1 - insignificant = len(str(1 << power)) - 1 - if binary_exponent >= fraction_width: - integer = significand << (binary_exponent - fraction_width) - else: - integer = significand >> (fraction_width - binary_exponent) - if insignificant: - power10 = 10 ** insignificant - integer, residue = divmod(integer, power10) - if residue >= power10 // 2: - integer += 1 - raw_digits = str(integer) - digits = [int(digit) for digit in raw_digits.rstrip('0')] - return _format(negative, digits, len(raw_digits) + insignificant) - - -def _estimate_decimal_exponent( - significand, binary_exponent, fraction_width): - normalized = significand << (52 - fraction_width) - bits = (1023 << 52) | (normalized & ((1 << 52) - 1)) - scaled = struct.unpack('>d', struct.pack('>Q', bits))[0] - estimate = ((scaled - 1.5) * 0.289529654 + 0.176091259 - + binary_exponent * 0.301029995663981) - estimate_bits = struct.unpack('>Q', struct.pack('>d', estimate))[0] - exponent = ((estimate_bits >> 52) & 0x7FF) - 1023 - negative = bool(estimate_bits >> 63) - fraction = estimate_bits & ((1 << 52) - 1) - if 0 <= exponent < 52: - mask = ((1 << 52) - 1) >> exponent - integer = ((fraction | (1 << 52)) >> (52 - exponent)) - if negative: - return -integer if fraction & mask == 0 else -integer - 1 - return integer - if exponent < 0: - magnitude = estimate_bits & ((1 << 63) - 1) - return 0 if magnitude == 0 else (-1 if negative else 0) - return int(estimate) - - -def _five_bits(power): - if power == 0: - return 0 - return (5 ** power).bit_length() if power <= 26 else power * 3 - - -def _signed(value, width): - value &= (1 << width) - 1 - sign = 1 << (width - 1) - return value - (1 << width) if value & sign else value - - -def _round_up(digits, decimal_exponent): - index = len(digits) - 1 - while index >= 0 and digits[index] == 9: - digits[index] = 0 - index -= 1 - if index < 0: - digits[0] = 1 - return decimal_exponent + 1 - digits[index] += 1 - return decimal_exponent - - -def _format(negative, digits, decimal_point): - text = ''.join(str(digit) for digit in digits) - if 0 < decimal_point < 8: - if len(text) < decimal_point: - result = text + '0' * (decimal_point - len(text)) + '.0' - else: - result = text[:decimal_point] + '.' + ( - text[decimal_point:] or '0') - elif -3 < decimal_point <= 0: - result = '0.' + '0' * -decimal_point + text - else: - result = (text[0] + '.' + (text[1:] or '0') - + 'E' + str(decimal_point - 1)) - return ('-' if negative else '') + result diff --git a/paimon-python/pypaimon/data/generic_variant.py b/paimon-python/pypaimon/data/generic_variant.py index 030bc9ab5ccc..9d0e0a0b0c06 100644 --- a/paimon-python/pypaimon/data/generic_variant.py +++ b/paimon-python/pypaimon/data/generic_variant.py @@ -148,12 +148,6 @@ def _write_le(buf, pos, value, n): buf[pos:pos + n] = value.to_bytes(n, 'little') -def _decimal_from_unscaled(unscaled, scale): - sign = 1 if unscaled < 0 else 0 - digits = tuple(int(digit) for digit in str(abs(unscaled))) or (0,) - return _decimal.Decimal((sign, digits, -scale)) - - def _short_str_header(size): return (size << 2) | _SHORT_STR @@ -312,26 +306,18 @@ def append_float(self, f): self._pos += 4 def append_decimal(self, d): + d = d.normalize() sign, digits, exponent = d.as_tuple() + if exponent > 0: + raise ValueError( + f'append_decimal requires a non-positive exponent (got {d!r}); ' + 'use append_double() for Decimal values with positive exponents' + ) unscaled = int(''.join(str(x) for x in digits)) if sign: unscaled = -unscaled - if exponent > 0: - unscaled *= 10 ** exponent - scale = 0 - else: - scale = -exponent - self.append_decimal_unscaled( - unscaled, max(1, len(str(abs(unscaled)))), scale) - - def append_decimal_unscaled(self, unscaled, precision, scale): - if not 0 <= scale <= _MAX_DECIMAL16_PRECISION: - raise ValueError(f'Unsupported VARIANT decimal scale: {scale}') - if not 0 < precision <= _MAX_DECIMAL16_PRECISION: - raise ValueError( - f'Unsupported VARIANT decimal precision: {precision}') - if not -(1 << 127) <= unscaled < (1 << 127): - raise ValueError('VARIANT decimal value exceeds 128 bits') + scale = -exponent if exponent < 0 else 0 + precision = len(digits) if scale <= _MAX_DECIMAL4_PRECISION and precision <= _MAX_DECIMAL4_PRECISION: self._write_byte(_primitive_header(_DECIMAL4)) @@ -682,7 +668,7 @@ def _to_python_impl(self, value, metadata, pos): else: raw = bytes(value[pos + 2:pos + 18]) unscaled = int.from_bytes(raw, 'little', signed=True) - return _decimal_from_unscaled(unscaled, scale) + return _decimal.Decimal(unscaled) / (_decimal.Decimal(10) ** scale) if vtype == _Type.STRING: if basic_type == _SHORT_STR: return value[pos + 1:pos + 1 + type_info].decode('utf-8') diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 3bf7456b6ece..348213099ae1 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -16,12 +16,7 @@ """Read and replace paths in Arrow VARIANT columns.""" -import base64 -import datetime -import decimal import functools -import json -import math import re import struct from typing import Dict, Mapping, Optional, Sequence, Tuple @@ -29,10 +24,6 @@ import numpy as np import pyarrow as pa -from pypaimon.data._variant_float_format import ( - format_float32, - format_float64, -) from pypaimon.data._variant_binary import ( _ARRAY, _OBJECT, @@ -49,10 +40,10 @@ _DECIMAL4, _DECIMAL8, _DECIMAL16, - GenericVariant, _DOUBLE, _FLOAT, _LONG_STR, + _NULL, _MAX_DECIMAL4_PRECISION, _MAX_DECIMAL8_PRECISION, _MAX_DECIMAL16_PRECISION, @@ -61,25 +52,11 @@ from pypaimon.data.variant_shredding import ( _build_array_value, _build_object_value, - _encode_scalar_to_value_bytes, ) _INDEX_PATTERN = re.compile(r"\[(\d+)]") _KEY_PATTERN = re.compile(r"\.([^\.\[]+)|\['([^']+)']|\[\"([^\"]+)\"]") -_TIMESTAMP_PATTERN = re.compile( - r'^(\d{4})-(\d{1,2})-(\d{1,2})' - r'(?:(?: (\d{1,2}):(\d{1,2}):(\d{1,2})(?:\.(\d{1,9}))?)' - r'|(?:T(\d{1,2}):(\d{1,2})(?::(\d{1,2})(?:\.(\d{1,9}))?)?))?$' -) -_JAVA_DECIMAL_FLOAT_PATTERN = re.compile( - r'^[+-]?(?:(?:[0-9]+(?:\.[0-9]*)?|\.[0-9]+)' - r'(?:[eE][+-]?[0-9]+)?)[fFdD]?$' -) -_JAVA_HEX_FLOAT_PATTERN = re.compile( - r'^[+-]?0[xX](?:[0-9a-fA-F]+(?:\.[0-9a-fA-F]*)?' - r'|\.[0-9a-fA-F]+)[pP][+-]?[0-9]+[fFdD]?$' -) _Path = Tuple[Tuple[str, object], ...] _SLOW_PATH_ROWS = 64 @@ -765,10 +742,6 @@ def _partition_path_plans(values, metadata, valid_rows, parsed_paths): def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): - if not all(pa.types.is_float32(target_type) - or pa.types.is_float64(target_type) - for target_type in target_types): - return None valid_rows = _valid_row_indices(chunk, values, chunk.field(1)) if not len(valid_rows): return [pa.nulls(len(chunk), type=target_type) @@ -786,11 +759,12 @@ def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): continue absolute = row_starts + pos headers = data[absolute] - if np.all(headers == _primitive_header(_DOUBLE)): - value_size, data_type = 8, np.dtype('= 0 else '-' - minutes = abs(seconds) // 60 - text += f'{sign}{minutes // 60:02d}:{minutes % 60:02d}' - return text - - -def _json_text(value): - if value is None: - return 'null' - if isinstance(value, bool): - return 'true' if value else 'false' - if isinstance(value, decimal.Decimal): - return _decimal_text(value) - if isinstance(value, float): - text = _floating_text(value) - return text if math.isfinite(value) else json.dumps(text) - if isinstance(value, int): - return str(value) - if isinstance(value, str): - return json.dumps(value, ensure_ascii=False, separators=(',', ':')) - if isinstance(value, bytes): - return json.dumps(base64.b64encode(value).decode('ascii')) - if isinstance(value, datetime.datetime): - return json.dumps(_timestamp_text(value)) - if isinstance(value, datetime.date): - return json.dumps(value.isoformat()) - if isinstance(value, dict): - return '{' + ','.join( - f'{json.dumps(key, ensure_ascii=False)}:{_json_text(child)}' - for key, child in value.items() - ) + '}' - if isinstance(value, (list, tuple)): - return '[' + ','.join(_json_text(child) for child in value) + ']' - return json.dumps(str(value), ensure_ascii=False) - - -def _variant_json(value, metadata, pos, limit, keys=None): - size = _checked_value_size(value, pos, limit) - if pos + size != limit: - _malformed("child size does not match container offsets") +def _decode_floating(value, pos, target_type): + size = _checked_value_size(value, pos) header = value[pos] - basic_type = header & 0x3 + if (header & 0x3) != _PRIMITIVE: + raise TypeError("VARIANT path is not FLOAT or DOUBLE") type_info = (header >> 2) & 0x3F - if basic_type == _OBJECT: - object_size, id_size, id_start, data_start, offsets, _ = ( - _checked_object_layout(value, pos, limit)) - if keys is None: - keys = { - key_id: key - for key, key_id in _metadata_key_ids(metadata).items() - } - fields = [] - for slot in range(object_size): - key_id = _read_unsigned( - value, id_start + slot * id_size, id_size) - if key_id not in keys: - _malformed("object key is missing from metadata") - child_start, child_end = _checked_object_child_bounds( - value, data_start, offsets, slot) - fields.append( - json.dumps(keys[key_id], ensure_ascii=False) - + ':' - + _variant_json( - value, metadata, child_start, child_end, keys) - ) - return '{' + ','.join(fields) + '}' - if basic_type == _ARRAY: - array_size, data_start, offsets, _ = _checked_array_layout( - value, pos, limit) - return '[' + ','.join( - _variant_json( - value, - metadata, - data_start + offsets[slot], - data_start + offsets[slot + 1], - keys, - ) - for slot in range(array_size) - ) + ']' - - decoded = GenericVariant(bytes(value[pos:limit]), metadata).to_python() - if basic_type == _PRIMITIVE and type_info in (_FLOAT, _DOUBLE): - text = _variant_floating_text(decoded, type_info) - return text if math.isfinite(decoded) else json.dumps(text) - return _json_text(decoded) - - -def _cast_decimal(value, target_type): - if isinstance(value, bool): - value = decimal.Decimal(1 if value else 0) - elif not isinstance(value, decimal.Decimal): - value = decimal.Decimal(str(value)) - quantum = decimal.Decimal((0, (1,), -target_type.scale)) - with decimal.localcontext() as context: - context.prec = max(50, target_type.precision + abs(target_type.scale)) - result = value.quantize(quantum, rounding=decimal.ROUND_HALF_UP) - digits = len(result.as_tuple().digits) - if digits > target_type.precision: - raise ValueError("decimal precision overflow") - return result - - -def _cast_integer(value, target_type): - bits = target_type.bit_width - if isinstance(value, float): - base_bits = 64 if bits == 64 else 32 - minimum = -(1 << (base_bits - 1)) - maximum = (1 << (base_bits - 1)) - 1 - if math.isnan(value): - number = 0 - elif value <= minimum: - number = minimum - elif value >= maximum: - number = maximum - else: - number = int(value) - else: - number = int(value) - return ((number + (1 << (bits - 1))) % (1 << bits)) - (1 << (bits - 1)) - - -def _cast_string_integer(value, target_type): - if not value: - raise ValueError - negative = value[0] == '-' - if value[0] in ('-', '+'): - value = value[1:] - if not value: - raise ValueError - parts = value.split('.') - if len(parts) > 2 or any( - character < '0' or character > '9' - for part in parts for character in part): - raise ValueError - integral = parts[0] - number = int(integral) if integral else 0 - if negative: - number = -number - bits = target_type.bit_width - minimum = -(1 << (bits - 1)) - maximum = (1 << (bits - 1)) - 1 - if number < minimum or number > maximum: - raise ValueError - return number - - -def _cast_string_floating(value): - text = value.strip() - if re.match(r'^[+-]?(?:NaN|Infinity)$', text): - if text.endswith('NaN'): - return math.copysign(float('nan'), -1.0 if text[0:1] == '-' else 1.0) - return float('-inf') if text.startswith('-') else float('inf') - if (_JAVA_DECIMAL_FLOAT_PATTERN.match(text) - or _JAVA_HEX_FLOAT_PATTERN.match(text)): - if text[-1:] in 'fFdD': - text = text[:-1] - if '0x' in text.lower(): - return float.fromhex(text) - return float(text) - raise ValueError - - -def _parse_date(value): - value = value.strip() - if value and (value.isdigit() - or value[0] == '-' and value[1:].isdigit()): - return datetime.date(1970, 1, 1) + datetime.timedelta(days=int(value)) - return datetime.datetime.strptime(value, '%Y-%m-%d').date() - - -def _parse_timestamp(value, target_type): - value = value.strip() - if value and (value.isdigit() - or value[0] == '-' and value[1:].isdigit()): - raw_value = int(value) - if target_type.unit == 'ns': - if raw_value < -(1 << 63) or raw_value > (1 << 63) - 1: - raise ValueError - return raw_value - micros = raw_value * { - 's': 1_000_000, 'ms': 1000, 'us': 1, - }[target_type.unit] - return datetime.datetime(1970, 1, 1) + datetime.timedelta( - microseconds=int(micros)) - - match = _TIMESTAMP_PATTERN.match(value) - if match is None: - raise ValueError - groups = match.groups() - year, month, day = (int(part) for part in groups[:3]) - if groups[3] is not None: - hour, minute, second, fraction = groups[3:7] - else: - hour, minute, second, fraction = groups[7:11] - hour = int(hour or 0) - minute = int(minute or 0) - second = int(second or 0) - nanos = int((fraction or '').ljust(9, '0') or 0) - base = datetime.datetime(year, month, day, hour, minute, second) - - if target_type.unit == 'ns': - delta = base - datetime.datetime(1970, 1, 1) - epoch_nanos = ( - (delta.days * 86400 + delta.seconds) * 1_000_000_000 + nanos - ) - if epoch_nanos < -(1 << 63) or epoch_nanos > (1 << 63) - 1: - raise ValueError - return epoch_nanos - - precision = {'s': 0, 'ms': 3, 'us': 6}[target_type.unit] - nanos = nanos // (10 ** (9 - precision)) * (10 ** (9 - precision)) - return base + datetime.timedelta(microseconds=nanos // 1000) - - -def _cast_python(value, target_type): - if value is None or pa.types.is_null(target_type): + if type_info == _FLOAT and pa.types.is_float32(target_type): + return struct.unpack_from('> 2) & 0x3F - if pa.types.is_struct(target_type): - if basic_type != _OBJECT: - raise ValueError(f"Invalid cast VARIANT to {target_type}") - children = _variant_object_children( - value, metadata, pos, end) - return { - field.name: ( - None if field.name not in children - else _decode_scalar( - value, metadata, children[field.name][0], field.type) - ) - for field in target_type - } - if (pa.types.is_list(target_type) - or pa.types.is_large_list(target_type) - or pa.types.is_fixed_size_list(target_type)): - if basic_type != _ARRAY: - raise ValueError(f"Invalid cast VARIANT to {target_type}") - return [ - _decode_scalar(value, metadata, child_pos, target_type.value_type) - for child_pos, _ in _variant_array_children( - value, pos, end) - ] - if pa.types.is_map(target_type): - if (basic_type != _OBJECT - or not (pa.types.is_string(target_type.key_type) - or pa.types.is_large_string(target_type.key_type))): - raise ValueError(f"Invalid cast VARIANT to {target_type}") - children = _variant_object_children( - value, metadata, pos, end) - return [ - (key, _decode_scalar( - value, metadata, child_pos, target_type.item_type)) - for key, (child_pos, _) in children.items() - ] - if (pa.types.is_string(target_type) - or pa.types.is_large_string(target_type)): - if basic_type in (_OBJECT, _ARRAY): - return _variant_json(value, metadata, pos, end) - selected = bytes(value[pos:end]) - decoded = GenericVariant(selected, metadata).to_python() - if ((pa.types.is_string(target_type) - or pa.types.is_large_string(target_type)) - and (value[pos] & 0x3) == _PRIMITIVE - and type_info in (_FLOAT, _DOUBLE)): - return _variant_floating_text(decoded, type_info) - return _cast_python(decoded, target_type) + raise TypeError( + f"VARIANT path type does not match {target_type}") def _patched_chunk( @@ -1290,17 +875,6 @@ def _patched_chunk( ) -def _rebuilt_chunk( - chunk: pa.StructArray, - values: Sequence[bytes], -) -> pa.StructArray: - return pa.StructArray.from_arrays( - [pa.array(values, type=chunk.type[0].type), chunk.field(1)], - fields=list(chunk.type), - mask=chunk.is_null(), - ) - - def _rebuilt_offsets(lengths, value_format): total = sum(int(length) for length in lengths) maximum = np.iinfo(np.dtype(value_format)).max @@ -1373,41 +947,20 @@ def __init__(self, value, length: int): "VARIANT replacement must be an Arrow Scalar or Array") if not _supported_replacement_type(self.type): raise TypeError( - f"Unsupported VARIANT replacement type: {self.type}") + "VARIANT replacement type must be float32 or float64") if pa.types.is_float64(self.type): self._value_format = ' Optional[int]: - return self._fixed_size if value is not None else None - - def patch(self, data: bytearray, pos: int, value) -> None: - struct.pack_into( - self._value_format, data, pos, self._type_header, value) - def numpy_values(self, offset: int, length: int, rows=None): if self._value_format is None: return None @@ -1438,10 +991,17 @@ def numpy_values(self, offset: int, length: int, rows=None): ) def encode(self, value) -> bytes: - if value is not None and self._value_format is not None: + if value is not None: return struct.pack( self._value_format, self._type_header, value) - return _encode_scalar_to_value_bytes(value, self.type) + return bytes([_primitive_header(_NULL)]) + + def validate_source(self, value, pos) -> None: + header = value[pos] + if header not in ( + self._type_header, _primitive_header(_NULL)): + raise TypeError( + f"VARIANT path type does not match {self.type}") def _vectorized_replace_chunk( @@ -1452,9 +1012,6 @@ def _vectorized_replace_chunk( global_row, strict, ): - if not all(provider._fixed_size is not None - for _, _, provider in parsed): - return None valid_rows = _valid_row_indices(chunk, values, chunk.field(1)) if not len(valid_rows): return chunk @@ -1465,7 +1022,7 @@ def _vectorized_replace_chunk( data_start = 0 output_data = None for planned in plans: - rows, row_starts, _, positions, limits = planned + rows, row_starts, source_data, positions, limits = planned replacements = [] compatible = np.ones(len(rows), dtype=bool) has_replacement = False @@ -1486,6 +1043,7 @@ def _vectorized_replace_chunk( absolute = row_starts + pos compatible &= ( replacement_valid + & (source_data[absolute] == provider._type_header) & (absolute + provider._fixed_size == row_starts + limit) ) replacements.append((pos, provider, replacement)) @@ -1515,15 +1073,25 @@ def _vectorized_replace_chunk( metadata = _BinaryValues(chunk.field(1)) rebuilt_rows = {} for row in slow_rows: - original = bytes(values.view(row)) - value = original + value = values.view(row) row_metadata = bytes(metadata.view(row)) positions = _path_positions(value, row_metadata, parsed_paths) + if not any(pos is not None for pos in positions): + if strict: + missing = next( + path for (path, _, _), pos in zip(parsed, positions) + if pos is None) + raise ValueError( + f"VARIANT path does not exist: {missing}") + continue + original = bytes(value) + value = original for (path, parsed_path, provider), pos in zip(parsed, positions): if pos is None: if strict: raise ValueError(f"VARIANT path does not exist: {path}") continue + provider.validate_source(value, pos) replacement_value = provider.scalar_at(global_row + row) value = _replace_path( value, @@ -1546,67 +1114,31 @@ def _vectorized_replace_chunk( def _supported_replacement_type(data_type: pa.DataType) -> bool: - return ( - pa.types.is_null(data_type) - or pa.types.is_boolean(data_type) - or pa.types.is_signed_integer(data_type) - or pa.types.is_float32(data_type) - or pa.types.is_float64(data_type) - or pa.types.is_string(data_type) - or pa.types.is_large_string(data_type) - or pa.types.is_binary(data_type) - or pa.types.is_large_binary(data_type) - or pa.types.is_date32(data_type) - or pa.types.is_timestamp(data_type) - or pa.types.is_decimal128(data_type) - ) + return pa.types.is_float32(data_type) or pa.types.is_float64(data_type) def _variant_get(column, paths: Mapping[str, pa.DataType]): parsed = [] for path, target_type in paths.items(): if not isinstance(target_type, pa.DataType): - raise TypeError("VARIANT target_type must be a PyArrow data type") + raise TypeError("VARIANT data_type must be a PyArrow data type") + if not (pa.types.is_float32(target_type) + or pa.types.is_float64(target_type)): + raise TypeError("VARIANT data_type must be float32 or float64") parsed.append((path, _parse_path(path), target_type)) parsed_paths = [parsed_path for _, parsed_path, _ in parsed] chunks, chunked, _ = _variant_chunks(column) result_chunks = {path: [] for path in paths} for chunk in chunks: values = _BinaryValues(chunk.field(0)) - vectorized = _vectorized_get_chunk( + results = _vectorized_get_chunk( chunk, values, parsed_paths, [target_type for _, _, target_type in parsed], ) - if vectorized is not None: - for (path, _, _), result in zip(parsed, vectorized): - result_chunks[path].append(result) - continue - metadata = _BinaryValues(chunk.field(1)) - valid = chunk.is_valid().to_pylist() - results = {path: [] for path in paths} - for row in range(len(chunk)): - if not valid[row]: - for path in paths: - results[path].append(None) - continue - value = values.view(row) - row_metadata = bytes(metadata.view(row)) - positions = _path_positions( - value, - row_metadata, - parsed_paths, - ) - for (path, _, target_type), pos in zip(parsed, positions): - results[path].append( - None if pos is None - else _decode_scalar( - value, row_metadata, pos, target_type) - ) - for path, _, target_type in parsed: - result_chunks[path].append( - pa.array(results[path], type=target_type)) + for (path, _, _), result in zip(parsed, results): + result_chunks[path].append(result) if not chunked: return {path: chunks[0] for path, chunks in result_chunks.items()} return { @@ -1615,16 +1147,16 @@ def _variant_get(column, paths: Mapping[str, pa.DataType]): } -def variant_get(column, path, target_type=None): - """Read one or more VARIANT paths into Arrow arrays.""" +def variant_get(column, path, data_type=None): + """Read FLOAT or DOUBLE paths into matching Arrow arrays.""" if isinstance(path, Mapping): - if target_type is not None: + if data_type is not None: raise TypeError( - "VARIANT target_type must be omitted for path mappings") + "VARIANT data_type must be omitted for path mappings") return _variant_get(column, path) - if target_type is None: - raise TypeError("VARIANT target_type must be a PyArrow data type") - return _variant_get(column, {path: target_type})[path] + if data_type is None: + raise TypeError("VARIANT data_type must be a PyArrow data type") + return _variant_get(column, {path: data_type})[path] def _paths_overlap(first: _Path, second: _Path) -> bool: @@ -1646,7 +1178,7 @@ def variant_replace( replacement=None, strict: bool = False, ): - """Replace one or more existing VARIANT paths with Arrow values.""" + """Replace existing FLOAT or DOUBLE paths with matching Arrow values.""" if not isinstance(strict, bool): raise TypeError("VARIANT strict must be a boolean") if isinstance(path, Mapping): @@ -1670,99 +1202,14 @@ def variant_replace( global_row = 0 for chunk in chunks: values = _BinaryValues(chunk.field(0)) - vectorized = _vectorized_replace_chunk( + result_chunks.append(_vectorized_replace_chunk( chunk, values, parsed, parsed_paths, global_row, strict, - ) - if vectorized is not None: - result_chunks.append(vectorized) - global_row += len(chunk) - continue - metadata = _BinaryValues(chunk.field(1)) - valid = chunk.is_valid().to_pylist() - chunk_replacements = { - path: provider.values(global_row, len(chunk)) - for path, _, provider in parsed - } - patched_data = None - data_start = 0 - rebuild = False - for row in range(len(chunk)): - if not valid[row]: - continue - row_start, value = values.row(row) - row_metadata = bytes(metadata.view(row)) - positions = _path_positions( - value, - row_metadata, - parsed_paths, - ) - for (path, parsed_path, provider), pos in zip( - parsed, positions): - if pos is None: - if strict: - raise ValueError( - f"VARIANT path does not exist: {path}") - continue - replacement_value = provider.value_at( - chunk_replacements[path], row) - new_size = provider.fixed_size(replacement_value) - encoded = ( - None if new_size is not None - else provider.encode(replacement_value) - ) - if new_size is None: - new_size = len(encoded) - if new_size != _checked_value_size(value, pos): - rebuild = True - break - if patched_data is None: - patched_data, data_start = values.copy_used_data() - patch_pos = row_start - data_start + pos - if encoded is None: - provider.patch( - patched_data, patch_pos, replacement_value) - else: - patched_data[patch_pos:patch_pos + new_size] = encoded - if rebuild: - break - - if rebuild: - rebuilt_values = [] - for row in range(len(chunk)): - value = bytes(values.view(row)) - if not valid[row]: - rebuilt_values.append(value) - continue - row_metadata = bytes(metadata.view(row)) - positions = _path_positions( - value, - row_metadata, - parsed_paths, - ) - for (path, parsed_path, provider), pos in zip( - parsed, positions): - if pos is None: - if strict: - raise ValueError( - f"VARIANT path does not exist: {path}") - continue - replacement_value = provider.value_at( - chunk_replacements[path], row) - encoded = provider.encode(replacement_value) - value = _replace_path( - value, row_metadata, 0, parsed_path, encoded) - rebuilt_values.append(value) - result_chunks.append(_rebuilt_chunk(chunk, rebuilt_values)) - elif patched_data is not None: - result_chunks.append(_patched_chunk( - chunk, values, patched_data, data_start)) - else: - result_chunks.append(chunk) + )) global_row += len(chunk) if not chunked: diff --git a/paimon-python/pypaimon/data/variant_shredding.py b/paimon-python/pypaimon/data/variant_shredding.py index 18692cbbe602..72a4508dd32a 100644 --- a/paimon-python/pypaimon/data/variant_shredding.py +++ b/paimon-python/pypaimon/data/variant_shredding.py @@ -250,64 +250,21 @@ def _append_scalar(builder, value, arrow_type: pa.DataType) -> None: elif pa.types.is_timestamp(arrow_type): # PyArrow converts timestamp to datetime.datetime if isinstance(value, datetime.datetime): - if (arrow_type.unit == 'ns' - and getattr(value, 'nanosecond', 0) != 0): - raise ValueError( - "VARIANT timestamps require microsecond-aligned values") if value.tzinfo is not None: epoch = datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) - delta = value - epoch - micros = ( - (delta.days * 86400 + delta.seconds) * 1_000_000 - + delta.microseconds - ) + micros = int((value - epoch).total_seconds() * 1_000_000) builder.append_timestamp(micros) else: epoch = datetime.datetime(1970, 1, 1) - delta = value - epoch - micros = ( - (delta.days * 86400 + delta.seconds) * 1_000_000 - + delta.microseconds - ) + micros = int((value - epoch).total_seconds() * 1_000_000) builder.append_timestamp_ntz(micros) else: - raw_value = int(value) - if arrow_type.unit == 'ns': - if raw_value % 1000: - raise ValueError( - "VARIANT timestamps require microsecond-aligned values") - micros = raw_value // 1000 - else: - micros = raw_value * {'s': 1_000_000, 'ms': 1000, 'us': 1}[ - arrow_type.unit] - builder.append_timestamp_ntz(micros) + builder.append_timestamp_ntz(int(value)) elif pa.types.is_decimal(arrow_type): - decimal = ( - value if isinstance(value, _decimal.Decimal) - else _decimal.Decimal(str(value)) - ) - sign, digits, exponent = decimal.as_tuple() - unscaled = int(''.join(str(digit) for digit in digits)) - if sign: - unscaled = -unscaled - shift = exponent + arrow_type.scale - if shift < 0 and unscaled % (10 ** -shift): - raise ValueError( - f'{decimal} does not have Arrow scale {arrow_type.scale}') - unscaled = ( - unscaled * (10 ** shift) if shift >= 0 - else unscaled // (10 ** -shift) - ) - scale = arrow_type.scale - precision = max(1, len(str(abs(unscaled)))) - if precision > arrow_type.precision: - raise ValueError( - f'{decimal} exceeds Arrow precision {arrow_type.precision}') - if scale < 0: - unscaled *= 10 ** -scale - scale = 0 - precision = max(1, len(str(abs(unscaled)))) - builder.append_decimal_unscaled(unscaled, precision, scale) + if isinstance(value, _decimal.Decimal): + builder.append_decimal(value) + else: + builder.append_decimal(_decimal.Decimal(str(value))) else: # Fallback: encode as string builder.append_string(str(value)) diff --git a/paimon-python/pypaimon/tests/variant_float_format_test.py b/paimon-python/pypaimon/tests/variant_float_format_test.py deleted file mode 100644 index 061cf3e5d8e1..000000000000 --- a/paimon-python/pypaimon/tests/variant_float_format_test.py +++ /dev/null @@ -1,75 +0,0 @@ -# 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. - -import struct -import unittest - -from pypaimon.data._variant_float_format import ( - format_float32, - format_float64, -) - - -def _float_from_bits(bits): - return struct.unpack('>f', struct.pack('>I', bits))[0] - - -def _double_from_bits(bits): - return struct.unpack('>d', struct.pack('>Q', bits))[0] - - -class TestVariantFloatFormat(unittest.TestCase): - - def test_float_to_string_matches_jdk8(self): - # Generated with OpenJDK 8u492 Float.toString. - samples = ( - (0x00000000, '0.0'), - (0x80000000, '-0.0'), - (0x00000001, '1.4E-45'), - (0x007FFFFF, '1.1754942E-38'), - (0x00800000, '1.17549435E-38'), - (0xD44F9F82, '-3.56693731E12'), - (0xEA75E34D, '-7.4315055E25'), - (0xE8FDF7D7, '-9.5946444E24'), - (0x7F7FFFFF, '3.4028235E38'), - (0x7F800000, 'Infinity'), - (0xFF800000, '-Infinity'), - (0x7FC00000, 'NaN'), - ) - for bits, expected in samples: - with self.subTest(bits=hex(bits)): - self.assertEqual( - format_float32(_float_from_bits(bits)), expected) - - def test_double_to_string_matches_jdk8(self): - # Generated with OpenJDK 8u492 Double.toString. - samples = ( - (0x0000000000000000, '0.0'), - (0x8000000000000000, '-0.0'), - (0x0000000000000001, '4.9E-324'), - (0x000FFFFFFFFFFFFF, '2.225073858507201E-308'), - (0x0010000000000000, '2.2250738585072014E-308'), - (0x439F4B86CD6A5E0C, '5.6376106381000781E17'), - (0x439DDD7467AF36D9, '5.3800104640575648E17'), - (0x7FEFFFFFFFFFFFFF, '1.7976931348623157E308'), - (0x7FF0000000000000, 'Infinity'), - (0xFFF0000000000000, '-Infinity'), - (0x7FF8000000000000, 'NaN'), - ) - for bits, expected in samples: - with self.subTest(bits=hex(bits)): - self.assertEqual( - format_float64(_double_from_bits(bits)), expected) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index e3760481f6ae..1455e4cb2881 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -14,10 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import datetime -import struct import unittest -from decimal import Decimal from unittest.mock import patch import numpy as np @@ -25,7 +22,7 @@ import pyarrow.compute as pc from pypaimon.data._variant_binary import _primitive_header -from pypaimon.data.generic_variant import _DECIMAL4, _DOUBLE, GenericVariant +from pypaimon.data.generic_variant import _DOUBLE, GenericVariant from pypaimon.data.variant_path import ( _path_positions, _rebuilt_offsets, @@ -33,7 +30,6 @@ variant_replace, ) from pypaimon.data.variant_shredding import ( - _build_array_value, _build_object_value, _encode_scalar_to_value_bytes, ) @@ -46,14 +42,6 @@ def _variants(values): ]) -def _decode(column): - return [ - None if value is None - else GenericVariant.from_arrow_struct(value).to_python() - for value in column.to_pylist() - ] - - def _float_variants(values): metadata = b'\x01\x00' return GenericVariant.to_arrow_array([ @@ -63,9 +51,17 @@ def _float_variants(values): ]) +def _decode(column): + return [ + None if value is None + else GenericVariant.from_arrow_struct(value).to_python() + for value in column.to_pylist() + ] + + class TestVariantGet(unittest.TestCase): - def test_get_nested_path_and_missing_values(self): + def test_nested_paths_and_missing_values(self): column = pa.chunked_array([ _variants([{'a.b': [{'value': 1.5}]}, None]), _variants([{'other': 2.0}, {'a.b': [{'value': -3.5}]}]), @@ -78,7 +74,7 @@ def test_get_nested_path_and_missing_values(self): self.assertEqual(result.num_chunks, 2) self.assertEqual(result.to_pylist(), [1.5, None, None, -3.5]) - def test_get_float_without_full_decode(self): + def test_reads_float_without_full_decode(self): column = _float_variants([1.25, -2.5]) with patch.object( @@ -88,336 +84,72 @@ def test_get_float_without_full_decode(self): self.assertEqual(result.to_pylist(), [1.25, -2.5]) - def test_get_path_mapping_in_one_pass(self): + def test_reads_multiple_paths_in_one_pass(self): column = _variants([ - {'velocity': {'y': 1.0, 'z': -2.0}}, - {'velocity': {'y': 3.0, 'z': -4.0}}, + {'velocity': {'x': 1.0, 'y': -2.0}}, + {'velocity': {'x': 3.0, 'y': -4.0}}, ]) result = variant_get(column, { + '$.velocity.x': pa.float64(), '$.velocity.y': pa.float64(), - '$.velocity.z': pa.float64(), }) - self.assertEqual(result['$.velocity.y'].to_pylist(), [1.0, 3.0]) - self.assertEqual(result['$.velocity.z'].to_pylist(), [-2.0, -4.0]) - - def test_get_matches_java_cast_semantics(self): - column = _variants([{ - 'long': 123, - 'object': {'age': 2}, - 'array': [1, '2'], - }]) - - self.assertEqual( - variant_get(column, '$.long', pa.string()).to_pylist(), ['123']) - self.assertEqual( - variant_get(column, '$.object', pa.string()).to_pylist(), - ['{"age":2}'], - ) - self.assertEqual( - variant_get(column, '$.array', pa.string()).to_pylist(), - ['[1,"2"]'], - ) - self.assertEqual( - variant_get( - column, - '$.object', - pa.struct([('age', pa.int32()), ('name', pa.string())]), - ).to_pylist(), - [{'age': 2, 'name': None}], - ) - self.assertEqual( - variant_get( - column, '$.array', pa.list_(pa.int32())).to_pylist(), - [[1, 2]], - ) - self.assertEqual( - variant_get( - column, '$.object', - pa.map_(pa.string(), pa.int32())).to_pylist(), - [[('age', 2)]], - ) - with self.assertRaisesRegex(ValueError, "Invalid cast"): - variant_get(column, '$.object', pa.int32()) - - def test_get_matches_java_string_and_binary_casts(self): - column = _variants([{ - 'decimal': '1.9', - 'overflow': '2147483648', - 'truthy': 'yes', - 'falsey': '0', - 'binary': b'abc', - }]) - - self.assertEqual( - variant_get(column, '$.decimal', pa.int32()).to_pylist(), [1]) - self.assertEqual( - variant_get(column, '$.truthy', pa.bool_()).to_pylist(), [True]) - self.assertEqual( - variant_get(column, '$.falsey', pa.bool_()).to_pylist(), [False]) - self.assertEqual( - variant_get(column, '$.binary', pa.binary()).to_pylist(), - [b'abc'], - ) - self.assertEqual( - variant_get(column, '$.binary', pa.string()).to_pylist(), - ['abc'], - ) - with self.assertRaisesRegex(ValueError, "Invalid cast"): - variant_get(column, '$.overflow', pa.int32()) + self.assertEqual(result['$.velocity.x'].to_pylist(), [1.0, 3.0]) + self.assertEqual(result['$.velocity.y'].to_pylist(), [-2.0, -4.0]) - temporal = _variants(['1', '2026-08-10']) - self.assertEqual( - variant_get(temporal, '$', pa.date32()).to_pylist(), - [datetime.date(1970, 1, 2), datetime.date(2026, 8, 10)], - ) - self.assertEqual( - variant_get(temporal.slice(0, 1), '$', pa.timestamp('us')) - .to_pylist(), - [datetime.datetime(1970, 1, 1, 0, 0, 0, 1)], - ) - - def test_get_matches_java_numeric_strings(self): - doubles = _variants([ - Decimal('100.00'), 1e20, 1e-4, - float('inf'), float('-inf'), float('nan'), - ]) - - self.assertEqual( - variant_get(doubles, '$', pa.string()).to_pylist(), - ['100.00', '1.0E20', '1.0E-4', - 'Infinity', '-Infinity', 'NaN'], - ) - - floats = GenericVariant.to_arrow_array([ - GenericVariant( - _encode_scalar_to_value_bytes(value, pa.float32()), - b'\x01\x00', - ) - for value in (1.2, 1e20, 1e-4) - ]) - self.assertEqual( - variant_get(floats, '$', pa.string()).to_pylist(), - ['1.2', '1.0E20', '1.0E-4'], - ) - - nested = GenericVariant.to_arrow_array([ - GenericVariant( - _build_object_value([(0, floats[0].as_py()['value'])]), - GenericVariant.from_python({'value': 0}).metadata(), - ), - GenericVariant( - _build_array_value([floats[0].as_py()['value']]), - b'\x01\x00', - ), - ]) - self.assertEqual( - variant_get(nested, '$', pa.string()).to_pylist(), - ['{"value":1.2}', '[1.2]'], - ) - strings = _variants(['1234567', '12345678901234']) - self.assertEqual( - variant_get(strings, '$', pa.string()).to_pylist(), - ['1234567', '12345678901234'], - ) - - def test_nested_json_uses_variant_float_text(self): - value = struct.unpack( - '>d', struct.pack('>Q', 0x439F4B86CD6A5E0C))[0] - column = _variants([{ - 'nested': {'finite': value, 'infinite': float('inf')}, - }]) - - result = variant_get( - column, '$', pa.struct([('nested', pa.string())])) - - self.assertEqual(result.to_pylist(), [{ - 'nested': ( - '{"finite":5.6376106381000781E17,' - '"infinite":"Infinity"}' - ), - }]) - - metadata = GenericVariant.from_python({'value': 0}).metadata() - nested_float = GenericVariant.to_arrow_array([ - GenericVariant( - _build_object_value([( - 0, _encode_scalar_to_value_bytes(1.2, pa.float32()), - )]), - metadata, - ), - ]) - self.assertEqual( - variant_get( - nested_float, - '$', - pa.struct([('value', pa.string())]), - ).to_pylist(), - [{'value': '1.2'}], - ) - - def test_string_floating_cast_matches_java_parser(self): - for target_type in (pa.float32(), pa.float64()): - for value in ('inf', 'nan', '1_0'): - with self.subTest(target_type=target_type, value=value): - with self.assertRaisesRegex(ValueError, "Invalid cast"): - variant_get(_variants([value]), '$', target_type) - - parsed = variant_get( - _variants(['NaN', 'Infinity', '-Infinity', '1.25f', '0x1.0p0']), - '$', - pa.float64(), - ).to_pylist() - self.assertTrue(np.isnan(parsed[0])) - self.assertEqual( - parsed[1:], [float('inf'), float('-inf'), 1.25, 1.0]) - - def test_container_timestamp_json_uses_space(self): - metadata = GenericVariant.from_python({'ts': 0}).metadata() - timestamp = datetime.datetime(2026, 8, 10, 12, 34, 56, 123000) - column = GenericVariant.to_arrow_array([ - GenericVariant( - _build_object_value([( - 0, - _encode_scalar_to_value_bytes( - timestamp, pa.timestamp('us')), - )]), - metadata, - ), - ]) - - self.assertEqual( - variant_get(column, '$', pa.string()).to_pylist(), - ['{"ts":"2026-08-10 12:34:56.123"}'], - ) - - def test_get_matches_java_numeric_casts(self): - column = _variants([ - {'value': 1e20}, - {'value': float('nan')}, - {'value': float('inf')}, - {'value': float('-inf')}, - ]) - self.assertEqual( - variant_get(column, '$.value', pa.int32()).to_pylist(), - [2147483647, 0, 2147483647, -2147483648], - ) - - invalid = _variants([ - {'value': Decimal('1.0')}, - {'value': 1.5}, - ]) - with self.assertRaisesRegex(ValueError, "Invalid cast"): - variant_get(invalid.slice(0, 1), '$.value', pa.bool_()) - with self.assertRaisesRegex(ValueError, "Invalid cast"): - variant_get(invalid.slice(1, 1), '$.value', pa.timestamp('us')) - - # Java casts LONG as epoch seconds, but parses numeric STRING using - # the target timestamp precision. - timestamp = variant_get( - _variants([{'value': 1}]), '$.value', pa.timestamp('us')) - self.assertEqual( - timestamp.to_pylist(), [datetime.datetime(1970, 1, 1, 0, 0, 1)]) - - nanos = variant_get( - _variants(['1', '-1', '2026-08-10 12:34:56.123456789']), - '$', - pa.timestamp('ns'), - ) - self.assertEqual( - nanos.cast(pa.int64()).to_pylist(), - [ - 1, - -1, - int(np.datetime64( - '2026-08-10T12:34:56.123456789', 'ns').astype(np.int64)), - ], - ) - - def test_get_timestamp_ns_checks_range_and_format(self): - boundaries = _variants([ - '1677-09-21 00:12:43.145224192', - '2262-04-11 23:47:16.854775807', - ]) - self.assertEqual( - variant_get(boundaries, '$', pa.timestamp('ns')) - .cast(pa.int64()).to_pylist(), - [-(1 << 63), (1 << 63) - 1], + def test_requires_exact_float_type(self): + cases = ( + (_float_variants([1.25]), pa.float64()), + (_variants([1.25]), pa.float32()), + (_variants([1]), pa.float64()), ) + for column, data_type in cases: + with self.subTest(data_type=data_type): + with self.assertRaisesRegex(TypeError, "does not match"): + variant_get(column, '$', data_type) - invalid = ( - '1677-09-21 00:12:43.145224191', - '2262-04-11 23:47:16.854775808', - '9999-12-31 23:59:59.999999999', - 'today', - 'now', - ) - for value in invalid: - with self.subTest(value=value): - with self.assertRaisesRegex(ValueError, "Invalid cast"): - variant_get( - _variants([value]), '$', pa.timestamp('ns')) + with self.assertRaisesRegex(TypeError, "float32 or float64"): + variant_get(_variants([1.0]), '$', pa.string()) - def test_variant_null_remains_arrow_null(self): + def test_variant_null_is_arrow_null(self): column = _variants([None, {'value': None}, {'value': 1.0}]) result = variant_get(column, '$.value', pa.float64()) self.assertEqual(result.to_pylist(), [None, None, 1.0]) - def test_get_decimal_is_exact(self): - expected = Decimal('12345678901234567890123456789012345678') - column = _variants([{'value': expected}]) - - result = variant_get( - column, '$.value', pa.decimal128(38, 0)) - - self.assertEqual(result.to_pylist(), [expected]) - - def test_get_rejects_malformed_metadata_and_decimal(self): + def test_rejects_malformed_rows(self): valid = GenericVariant.from_python({'value': 1.0}) - bad_metadata = bytes([2]) + valid.metadata()[1:] column = pa.StructArray.from_arrays([ - pa.array([valid.value()]), - pa.array([bad_metadata]), + pa.array([valid.value()[:-8]]), + pa.array([valid.metadata()]), ], names=['value', 'metadata']) - with self.assertRaisesRegex(ValueError, "metadata version"): + with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"): variant_get(column, '$.value', pa.float64()) - for scale, unscaled in ((10, 1), (0, 2147483647)): - with self.subTest(scale=scale, unscaled=unscaled): - value = ( - bytes([_primitive_header(_DECIMAL4), scale]) - + unscaled.to_bytes(4, 'little', signed=True) - ) - malformed = GenericVariant.to_arrow_array([ - GenericVariant(value, b'\x01\x00')]) - with self.assertRaisesRegex( - ValueError, "decimal precision or scale"): - variant_get(malformed, '$', pa.decimal128(38, 0)) - - def test_get_copies_only_selected_subtree(self): - column = _variants([{'small': 'x', 'large': b'x' * (2 * 1024 * 1024)}]) - decoded_sizes = [] - original = GenericVariant.to_python - - def decode(selected): - decoded_sizes.append(len(selected.value())) - return original(selected) - - with patch.object(GenericVariant, 'to_python', decode): - result = variant_get(column, '$.small', pa.string()) - - self.assertEqual(result.to_pylist(), ['x']) - self.assertEqual(decoded_sizes, [2]) - - def test_get_rejects_invalid_arguments(self): - column = _variants([{'value': 1}]) + value = _build_object_value([ + (0, bytes([_primitive_header(_DOUBLE)])), + (1, _encode_scalar_to_value_bytes(2.0, pa.float64())), + ]) + siblings = GenericVariant.to_arrow_array([ + GenericVariant( + value, + GenericVariant.from_python({'a': 0, 'b': 0}).metadata(), + ) + ]) + with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"): + variant_get(siblings, '$.a', pa.float64()) + + def test_rejects_invalid_arguments(self): + column = _variants([{'value': 1.0}]) with self.assertRaisesRegex(ValueError, "Invalid VARIANT path"): - variant_get(column, 'value', pa.int64()) + variant_get(column, 'value', pa.float64()) with self.assertRaisesRegex(TypeError, "PyArrow data type"): - variant_get(column, '$.value', 'BIGINT') + variant_get(column, '$.value', 'DOUBLE') + with self.assertRaisesRegex(TypeError, "must be omitted"): + variant_get( + column, {'$.value': pa.float64()}, pa.float64()) invalid_metadata = pa.StructArray.from_arrays( [ @@ -427,472 +159,183 @@ def test_get_rejects_invalid_arguments(self): names=['value', 'metadata'], mask=pa.array([True]), ) - with self.assertRaisesRegex(TypeError, "metadata field must be binary"): + with self.assertRaisesRegex( + TypeError, "metadata field must be binary"): variant_get(invalid_metadata, '$.value', pa.float64()) class TestVariantReplace(unittest.TestCase): - def test_timestamp_replacement_is_exact(self): - for arrow_type, value in ( - (pa.timestamp('us'), - datetime.datetime(9999, 12, 31, 23, 59, 59, 999999)), - (pa.timestamp('us', tz='UTC'), - datetime.datetime( - 2500, 1, 1, 0, 0, 0, 1, - tzinfo=datetime.timezone.utc))): - with self.subTest(arrow_type=arrow_type): - column = _variants([{'value': 0}]) - result = variant_replace( - column, '$.value', pa.scalar(value, type=arrow_type)) - self.assertEqual(_decode(result), [{'value': value}]) - - column = _variants([{'value': 0}]) - for nanos in (1, -1, 1001, -1001): - with self.subTest(nanos=nanos): - with self.assertRaisesRegex(ValueError, "microsecond-aligned"): - variant_replace( - column, - '$.value', - pa.scalar(nanos, type=pa.timestamp('ns')), - ) - for nanos in (1000, -1000): - with self.subTest(nanos=nanos): - result = variant_replace( - column, - '$.value', - pa.scalar(nanos, type=pa.timestamp('ns')), - ) - self.assertEqual( - _decode(result), - [{'value': datetime.datetime(1970, 1, 1) - + datetime.timedelta(microseconds=nanos // 1000)}], - ) - - def test_nullable_fast_path_does_not_devectorize_chunk(self): - size = 50000 - column = _variants( - [None] + [{'value': float(index)} for index in range(1, size)]) - replacement = pa.scalar(-1.0) - - with patch( - 'pypaimon.data.variant_path._path_positions', - wraps=_path_positions, - ) as slow_path: - current = variant_get(column, '$.value', pa.float64()) - result = variant_replace(column, '$.value', replacement) - - self.assertEqual(current[0].as_py(), None) - self.assertEqual(current[-1].as_py(), float(size - 1)) - self.assertIsNone(result[0].as_py()) - self.assertEqual(_decode(result.slice(size - 1, 1)), - [{'value': -1.0}]) - slow_path.assert_not_called() - - def test_sparse_layout_anomalies_keep_slow_path_bounded(self): - size = 4096 - uniform = _variants([ - {'value': float(index)} for index in range(1, size) - ]) - cases = ( - (_variants([{'extra': 1, 'value': 0.0}]), False), - (_variants([{'other': 0.0}]), True), - ) - for first, missing in cases: - with self.subTest(missing=missing): - column = pa.concat_arrays([first, uniform]) - with patch( - 'pypaimon.data.variant_path._path_positions', - wraps=_path_positions, - ) as slow_path: - current = variant_get( - column, '$.value', pa.float64()) - self.assertLessEqual( - slow_path.call_count, 64) - self.assertEqual(current[0].as_py(), None if missing else 0.0) - self.assertEqual(current[-1].as_py(), float(size - 1)) - - with patch( - 'pypaimon.data.variant_path._path_positions', - wraps=_path_positions, - ) as slow_path: - result = variant_replace( - column, '$.value', pa.scalar(-1.0)) - self.assertLessEqual( - slow_path.call_count, 64) - expected_first = {'other': 0.0} if missing else { - 'extra': 1, 'value': -1.0, - } - self.assertEqual(_decode(result.slice(0, 1)), - [expected_first]) - self.assertEqual(_decode(result.slice(size - 1, 1)), - [{'value': -1.0}]) - - def test_sparse_value_types_keep_slow_path_bounded(self): - size = 4096 - uniform = _variants([ - {'value': float(index)} for index in range(1, size) - ]) - for exceptional in (None, 1): - with self.subTest(exceptional=exceptional): - column = pa.concat_arrays([ - _variants([{'value': exceptional}]), uniform, - ]) - with patch( - 'pypaimon.data.variant_path._path_positions', - wraps=_path_positions, - ) as slow_path: - current = variant_get( - column, '$.value', pa.float64()) - self.assertLessEqual(slow_path.call_count, 1) - self.assertEqual( - current[0].as_py(), - None if exceptional is None else 1.0, - ) - - with patch( - 'pypaimon.data.variant_path._path_positions', - wraps=_path_positions, - ) as slow_path: - result = variant_replace( - column, '$.value', pa.scalar(-1.0)) - self.assertLessEqual(slow_path.call_count, 1) - self.assertEqual(_decode(result.slice(0, 1)), - [{'value': -1.0}]) - self.assertEqual(_decode(result.slice(size - 1, 1)), - [{'value': -1.0}]) - - replacement = pa.array( - [None] + [-1.0] * (len(uniform) - 1), type=pa.float64()) - with patch( - 'pypaimon.data.variant_path._path_positions', - wraps=_path_positions, - ) as slow_path: - result = variant_replace(uniform, '$.value', replacement) - self.assertLessEqual(slow_path.call_count, 1) - self.assertEqual(_decode(result.slice(0, 1)), [{'value': None}]) - self.assertEqual(_decode(result.slice(len(uniform) - 1, 1)), - [{'value': -1.0}]) - - def test_truncated_value_does_not_cross_row_boundary(self): - first = GenericVariant.from_python({'value': 1.0}) - second = GenericVariant.from_python({'value': 2.0}) - column = pa.StructArray.from_arrays( - [ - pa.array([first.value()[:-8], second.value()]), - pa.array([first.metadata(), second.metadata()]), - ], - names=['value', 'metadata'], - ) - - with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"): - variant_replace( - column, '$.value', - pa.array([3.0, 4.0], type=pa.float64())) - self.assertEqual( - GenericVariant.from_arrow_struct(column[1].as_py()).to_python(), - {'value': 2.0}, - ) - - def test_truncated_object_child_does_not_cross_sibling_boundary(self): - valid = GenericVariant.from_python({'a': 1.0, 'b': 2.0}) - value = _build_object_value([ - (0, bytes([_primitive_header(_DOUBLE)])), - (1, _encode_scalar_to_value_bytes(2.0, pa.float64())), - ]) - column = GenericVariant.to_arrow_array([ - GenericVariant(value, valid.metadata())]) - - with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"): - variant_get(column, '$.a', pa.float64()) - with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"): - variant_replace(column, '$.a', pa.scalar(3.0)) - - def test_equal_length_uses_copy_on_write(self): - column = _variants([ - {'number': 1.0, 'text': 'keep'}, - None, - {'number': -2.0, 'text': 'also keep'}, - ]) - original = column.to_pylist() - replacement = pa.array([-1.0, 0.0, 2.0], type=pa.float64()) - - with patch.object( - GenericVariant, 'to_python', - side_effect=AssertionError("full decode is not allowed")): - result = variant_replace( - column, '$.number', replacement) - - self.assertEqual(column.to_pylist(), original) - self.assertEqual(_decode(result), [ - {'number': -1.0, 'text': 'keep'}, - None, - {'number': 2.0, 'text': 'also keep'}, - ]) - self.assertEqual( - column.field('value').buffers()[1].address, - result.field('value').buffers()[1].address, - ) - self.assertNotEqual( - column.field('value').buffers()[2].address, - result.field('value').buffers()[2].address, - ) - self.assertEqual( - column.field('metadata').buffers()[2].address, - result.field('metadata').buffers()[2].address, - ) - self.assertEqual( - column.buffers()[0].address, - result.buffers()[0].address, - ) - - def test_sliced_input_copies_only_visible_values(self): - base = _variants([ - {'number': float(i), 'padding': 'x' * 1000} - for i in range(100) - ]) - - for binary_type in (pa.binary(), pa.large_binary()): - with self.subTest(binary_type=binary_type): - values = base.field('value').cast(binary_type) - converted = pa.StructArray.from_arrays( - [values, base.field('metadata')], - names=['value', 'metadata'], - ) - column = converted.slice(50, 3) - - result = variant_replace( - column, - '$.number', - pa.scalar(-1.0, type=pa.float64()), - ) - - expected_size = sum( - len(value) - for value in column.field('value').to_pylist() - ) - self.assertEqual( - result.field('value').buffers()[2].size, - expected_size, - ) - self.assertEqual(result.field('value').offset, 0) - self.assertEqual( - [row['number'] for row in _decode(result)], - [-1.0, -1.0, -1.0], - ) - def test_get_compute_replace_pipeline(self): column = pa.chunked_array([ - _variants([{'y': 1.0, 'z': -2.0}, None]), - _variants([{'y': -3.0, 'z': 4.0}]), + _variants([{'x': 1.0, 'y': -2.0}, None]), + _variants([{'x': -3.0, 'y': 4.0}]), ]) current = variant_get(column, { + '$.x': pa.float64(), '$.y': pa.float64(), - '$.z': pa.float64(), }) result = variant_replace(column, { - '$.y': pc.negate(current['$.y']), - '$.z': pc.negate(current['$.z']), + path: pc.negate(values) + for path, values in current.items() }) self.assertIsInstance(result, pa.ChunkedArray) - self.assertEqual(result.num_chunks, 2) self.assertEqual(_decode(result), [ - {'y': -1.0, 'z': 2.0}, None, - {'y': 3.0, 'z': -4.0}, + {'x': -1.0, 'y': 2.0}, None, + {'x': 3.0, 'y': -4.0}, ]) - def test_vectorized_paths_support_varying_offsets(self): + def test_updates_four_double_paths(self): column = _variants([ - { - 'prefix': 'x' * index, - 'items': ['y' * (64 - index), { - 'value': float(index), - 'other': float(-index), - }], - } - for index in range(1, 64) + {'a': 1.0, 'b': 2.0, 'nested': {'c': 3.0, 'd': 4.0}}, + {'a': -1.0, 'b': -2.0, 'nested': {'c': -3.0, 'd': -4.0}}, ]) paths = { - '$.items[1].value': pa.float64(), - '$.items[1].other': pa.float64(), + '$.a': pa.float64(), + '$.b': pa.float64(), + '$.nested.c': pa.float64(), + '$.nested.d': pa.float64(), } - with patch( - 'pypaimon.data.variant_path._path_positions', - side_effect=AssertionError("slow path is not allowed")): - current = variant_get(column, paths) - result = variant_replace(column, { - path: pc.negate(values) - for path, values in current.items() - }) + current = variant_get(column, paths) + result = variant_replace(column, { + path: pc.negate(value) for path, value in current.items() + }) - decoded = _decode(result) - self.assertEqual( - [row['items'][1]['value'] for row in decoded], - [float(-index) for index in range(1, 64)], - ) + self.assertEqual(_decode(result), [ + {'a': -1.0, 'b': -2.0, 'nested': {'c': -3.0, 'd': -4.0}}, + {'a': 1.0, 'b': 2.0, 'nested': {'c': 3.0, 'd': 4.0}}, + ]) + + def test_float_and_double_are_distinct(self): + floats = _float_variants([1.0, 2.0]) + result = variant_replace( + floats, '$', pa.array([-1.0, -2.0], type=pa.float32())) self.assertEqual( - [row['items'][1]['other'] for row in decoded], - [float(index) for index in range(1, 64)], + variant_get(result, '$', pa.float32()).to_pylist(), + [-1.0, -2.0], ) - def test_vectorized_paths_support_wide_containers(self): - rows = [] - for row in range(3): - value = {'field_%03d' % i: i for i in range(300)} - value['field_299'] = float(row) - value['items'] = list(range(299)) + [float(row + 10)] - rows.append(value) - column = _variants(rows) - paths = { - '$.field_299': pa.float64(), - '$.items[299]': pa.float64(), - } + with self.assertRaisesRegex(TypeError, "does not match"): + variant_replace(floats, '$', pa.scalar(1.0, type=pa.float64())) + with self.assertRaisesRegex(TypeError, "does not match"): + variant_replace( + _variants([1.0]), '$', pa.scalar(1.0, type=pa.float32())) + + def test_nullable_rows_stay_vectorized(self): + size = 4096 + column = _variants( + [None] + [{'value': float(index)} for index in range(1, size)]) with patch( 'pypaimon.data.variant_path._path_positions', - side_effect=AssertionError("slow path is not allowed")): - current = variant_get(column, paths) - result = variant_replace(column, { - path: pc.negate(values) - for path, values in current.items() - }) + wraps=_path_positions, + ) as slow_path: + current = variant_get(column, '$.value', pa.float64()) + result = variant_replace(column, '$.value', pa.scalar(-1.0)) - self.assertEqual(current['$.field_299'].to_pylist(), [0.0, 1.0, 2.0]) - self.assertEqual( - current['$.items[299]'].to_pylist(), [10.0, 11.0, 12.0]) - decoded = _decode(result) - self.assertEqual( - [row['field_299'] for row in decoded], [0.0, -1.0, -2.0]) - self.assertEqual( - [row['items'][299] for row in decoded], [-10.0, -11.0, -12.0]) + self.assertIsNone(current[0].as_py()) + self.assertEqual(current[-1].as_py(), float(size - 1)) + self.assertIsNone(result[0].as_py()) + self.assertEqual(_decode(result.slice(size - 1, 1)), + [{'value': -1.0}]) + slow_path.assert_not_called() - def test_scalar_and_same_length_string_replacement(self): - column = _variants([{'text': 'aa'}, {'text': 'bb'}]) + def test_sparse_layout_fallback_is_bounded(self): + size = 4096 + column = pa.concat_arrays([ + _variants([{'extra': 1, 'value': 0.0}]), + _variants([{'value': float(index)} for index in range(1, size)]), + ]) - result = variant_replace( - column, '$.text', pa.scalar('xy', type=pa.string())) + with patch( + 'pypaimon.data.variant_path._path_positions', + wraps=_path_positions, + ) as slow_path: + current = variant_get(column, '$.value', pa.float64()) + result = variant_replace(column, '$.value', pa.scalar(-1.0)) - self.assertEqual(_decode(result), [{'text': 'xy'}, {'text': 'xy'}]) - self.assertEqual( - column.field('value').buffers()[1].address, - result.field('value').buffers()[1].address, - ) + self.assertLessEqual(slow_path.call_count, 128) + self.assertEqual(current[0].as_py(), 0.0) + self.assertEqual(_decode(result.slice(0, 1)), + [{'extra': 1, 'value': -1.0}]) - def test_decimal_replacement_preserves_arrow_value(self): - for data_type, value, expected_exponent in ( - (pa.decimal128(10, 2), Decimal('100.00'), -2), - (pa.decimal128(10, -2), Decimal('1E+2'), 0)): - with self.subTest(data_type=data_type): - column = _variants([{'value': 0}]) + def test_missing_path_is_noop_or_strict_error(self): + column = _variants([ + {'other': float(index)} for index in range(4096) + ]) - result = variant_replace( - column, - '$.value', - pa.array([value], type=data_type), - ) + with patch( + 'pypaimon.data.variant_path._path_positions', + wraps=_path_positions, + ) as slow_path: + current = variant_get(column, '$.missing', pa.float64()) + result = variant_replace( + column, '$.missing', pa.scalar(3.0, type=pa.float64())) - decoded = _decode(result)[0]['value'] - self.assertEqual(decoded, Decimal('100.00')) - self.assertEqual( - decoded.as_tuple().exponent, expected_exponent) + self.assertEqual(current.null_count, len(column)) + self.assertIs(result, column) + slow_path.assert_not_called() + with self.assertRaisesRegex(ValueError, "path does not exist"): + variant_replace( + column, '$.missing', pa.scalar(3.0), strict=True) - def test_different_length_rebuilds_offsets(self): + def test_null_replacement_rebuilds_only_affected_row(self): column = _variants([ - { - 'before': 1, - 'nested': {'items': [0, {'text': 'a'}]}, - 'after': 2, - }, - { - 'before': 3, - 'nested': {'items': [0, {'text': 'bb'}]}, - 'after': 4, - }, + {'value': 1.0, 'padding': 'x' * 1000}, + {'value': 2.0, 'padding': 'y' * 1000}, ]) result = variant_replace( column, - '$.nested.items[1].text', - pa.array(['a much longer value', 'x'], type=pa.string()), + '$.value', + pa.array([None, -2.0], type=pa.float64()), ) self.assertEqual(_decode(result), [ - { - 'before': 1, - 'nested': {'items': [0, {'text': 'a much longer value'}]}, - 'after': 2, - }, - { - 'before': 3, - 'nested': {'items': [0, {'text': 'x'}]}, - 'after': 4, - }, + {'value': None, 'padding': 'x' * 1000}, + {'value': -2.0, 'padding': 'y' * 1000}, ]) self.assertEqual( column.field('metadata').buffers()[2].address, result.field('metadata').buffers()[2].address, ) - def test_same_length_type_change_uses_fast_path(self): - column = _variants([1000000, -1000000]) + def test_copy_on_write_and_sliced_input(self): + base = _variants([ + {'value': float(index), 'padding': 'x' * 1000} + for index in range(100) + ]) + column = base.slice(50, 3) - result = variant_replace( - column, '$', pa.array([1.25, -2.5], type=pa.float32())) + result = variant_replace(column, '$.value', pa.scalar(-1.0)) - self.assertEqual(_decode(result), [1.25, -2.5]) self.assertEqual( - column.field('value').buffers()[1].address, - result.field('value').buffers()[1].address, + [row['value'] for row in _decode(result)], [-1.0, -1.0, -1.0]) + self.assertEqual( + result.field('value').buffers()[2].size, + sum(len(value) for value in column.field('value').to_pylist()), + ) + self.assertEqual( + column.field('metadata').buffers()[2].address, + result.field('metadata').buffers()[2].address, ) - def test_missing_path_is_noop_or_strict_error(self): - column = _variants([{'value': 1}, {'other': 2}]) - replacement = pa.array([10, 20], type=pa.int64()) - - result = variant_replace(column, '$.value', replacement) - - self.assertEqual(_decode(result), [{'value': 10}, {'other': 2}]) - with self.assertRaisesRegex(ValueError, "path does not exist"): - variant_replace( - column, '$.value', replacement, strict=True) - - def test_all_missing_paths_reuse_input_buffers(self): - column = _variants([ - {'other': float(index)} for index in range(4096) + def test_rejects_truncated_child_without_touching_sibling(self): + valid = GenericVariant.from_python({'a': 1.0, 'b': 2.0}) + value = _build_object_value([ + (0, bytes([_primitive_header(_DOUBLE)])), + (1, _encode_scalar_to_value_bytes(2.0, pa.float64())), ]) + column = GenericVariant.to_arrow_array([ + GenericVariant(value, valid.metadata())]) + original = column.to_pylist() - with patch( - 'pypaimon.data.variant_path._path_positions', - wraps=_path_positions, - ) as slow_path: - current = variant_get(column, '$.missing', pa.float64()) - result = variant_replace( - column, '$.missing', pa.scalar(3.0, type=pa.float64())) - - self.assertEqual(current.null_count, len(column)) - self.assertIs(result, column) - slow_path.assert_not_called() - - current = variant_get(column, { - '$.other': pa.float64(), - '$.missing': pa.float64(), - }) - result = variant_replace(column, { - '$.other': pa.scalar(-1.0), - '$.missing': pa.scalar(3.0), - }) - self.assertEqual(current['$.missing'].null_count, len(column)) - self.assertEqual(_decode(result.slice(0, 1)), [{'other': -1.0}]) - with self.assertRaisesRegex(ValueError, "path does not exist"): - variant_replace( - column, - '$.missing', - pa.scalar(3.0, type=pa.float64()), - strict=True, - ) + with self.assertRaisesRegex(ValueError, "MALFORMED_VARIANT"): + variant_replace(column, '$.a', pa.scalar(3.0)) + self.assertEqual(column.to_pylist(), original) def test_rebuilt_binary_offsets_reject_overflow(self): self.assertEqual( @@ -902,22 +345,6 @@ def test_rebuilt_binary_offsets_reject_overflow(self): with self.assertRaisesRegex(ValueError, "use LargeBinary"): _rebuilt_offsets(np.array([(1 << 31) - 1, 1]), ' Date: Tue, 11 Aug 2026 00:21:28 -0700 Subject: [PATCH 23/23] [python] Support exact VARIANT path types --- docs/docs/pypaimon/python-api.mdx | 15 +- .../pypaimon/data/generic_variant.py | 32 +- paimon-python/pypaimon/data/variant_path.py | 278 ++++++++++++++++-- .../pypaimon/data/variant_shredding.py | 57 +++- .../pypaimon/tests/variant_path_test.py | 154 +++++++++- 5 files changed, 492 insertions(+), 44 deletions(-) diff --git a/docs/docs/pypaimon/python-api.mdx b/docs/docs/pypaimon/python-api.mdx index 98be214eaa25..be95a60a0ebc 100644 --- a/docs/docs/pypaimon/python-api.mdx +++ b/docs/docs/pypaimon/python-api.mdx @@ -1126,10 +1126,10 @@ Supported Paimon type strings for shredded sub-fields: `BOOLEAN`, `INT`, `BIGINT -### VARIANT FLOAT/DOUBLE Path Updates +### VARIANT Path Updates -Read existing FLOAT or DOUBLE paths as Arrow arrays, use Arrow compute, and -replace them without decoding unrelated fields: +Read existing paths as Arrow arrays, use Arrow compute, and replace them +without decoding unrelated fields: ```python import pyarrow as pa @@ -1142,10 +1142,11 @@ updated_payload = variant_replace( payload, '$.velocity.y', pc.negate(current)) ``` -The Arrow type passed to `variant_get` and `variant_replace` must match the -stored FLOAT or DOUBLE type; these APIs do not cast VARIANT values. Missing -paths read as NULL and remain unchanged unless `strict=True` is specified. -Pass mappings to process multiple paths in one pass. +The Arrow type must match the stored VARIANT type; these APIs do not cast +values. Exact extraction supports scalar types and nested struct, list, and +string-keyed map types. Replacement supports scalar types. Missing paths read +as NULL and remain unchanged unless `strict=True` is specified. Pass mappings +to process multiple paths in one pass. **`GenericVariant` API:** diff --git a/paimon-python/pypaimon/data/generic_variant.py b/paimon-python/pypaimon/data/generic_variant.py index 9d0e0a0b0c06..030bc9ab5ccc 100644 --- a/paimon-python/pypaimon/data/generic_variant.py +++ b/paimon-python/pypaimon/data/generic_variant.py @@ -148,6 +148,12 @@ def _write_le(buf, pos, value, n): buf[pos:pos + n] = value.to_bytes(n, 'little') +def _decimal_from_unscaled(unscaled, scale): + sign = 1 if unscaled < 0 else 0 + digits = tuple(int(digit) for digit in str(abs(unscaled))) or (0,) + return _decimal.Decimal((sign, digits, -scale)) + + def _short_str_header(size): return (size << 2) | _SHORT_STR @@ -306,18 +312,26 @@ def append_float(self, f): self._pos += 4 def append_decimal(self, d): - d = d.normalize() sign, digits, exponent = d.as_tuple() - if exponent > 0: - raise ValueError( - f'append_decimal requires a non-positive exponent (got {d!r}); ' - 'use append_double() for Decimal values with positive exponents' - ) unscaled = int(''.join(str(x) for x in digits)) if sign: unscaled = -unscaled - scale = -exponent if exponent < 0 else 0 - precision = len(digits) + if exponent > 0: + unscaled *= 10 ** exponent + scale = 0 + else: + scale = -exponent + self.append_decimal_unscaled( + unscaled, max(1, len(str(abs(unscaled)))), scale) + + def append_decimal_unscaled(self, unscaled, precision, scale): + if not 0 <= scale <= _MAX_DECIMAL16_PRECISION: + raise ValueError(f'Unsupported VARIANT decimal scale: {scale}') + if not 0 < precision <= _MAX_DECIMAL16_PRECISION: + raise ValueError( + f'Unsupported VARIANT decimal precision: {precision}') + if not -(1 << 127) <= unscaled < (1 << 127): + raise ValueError('VARIANT decimal value exceeds 128 bits') if scale <= _MAX_DECIMAL4_PRECISION and precision <= _MAX_DECIMAL4_PRECISION: self._write_byte(_primitive_header(_DECIMAL4)) @@ -668,7 +682,7 @@ def _to_python_impl(self, value, metadata, pos): else: raw = bytes(value[pos + 2:pos + 18]) unscaled = int.from_bytes(raw, 'little', signed=True) - return _decimal.Decimal(unscaled) / (_decimal.Decimal(10) ** scale) + return _decimal_from_unscaled(unscaled, scale) if vtype == _Type.STRING: if basic_type == _SHORT_STR: return value[pos + 1:pos + 1 + type_info].decode('utf-8') diff --git a/paimon-python/pypaimon/data/variant_path.py b/paimon-python/pypaimon/data/variant_path.py index 348213099ae1..9a958c335998 100644 --- a/paimon-python/pypaimon/data/variant_path.py +++ b/paimon-python/pypaimon/data/variant_path.py @@ -48,10 +48,14 @@ _MAX_DECIMAL8_PRECISION, _MAX_DECIMAL16_PRECISION, _PRIMITIVE_FIXED_SIZES, + GenericVariant, + _Type, + _variant_get_type, ) from pypaimon.data.variant_shredding import ( _build_array_value, _build_object_value, + _encode_scalar_to_value_bytes, ) @@ -742,6 +746,10 @@ def _partition_path_plans(values, metadata, valid_rows, parsed_paths): def _vectorized_get_chunk(chunk, values, parsed_paths, target_types): + if not all(pa.types.is_float32(target_type) + or pa.types.is_float64(target_type) + for target_type in target_types): + return None valid_rows = _valid_row_indices(chunk, values, chunk.field(1)) if not len(valid_rows): return [pa.nulls(len(chunk), type=target_type) @@ -852,6 +860,137 @@ def _decode_floating(value, pos, target_type): f"VARIANT path type does not match {target_type}") +def _variant_object_children(value, metadata, pos, end): + size, id_size, id_start, data_start, offsets, _ = ( + _checked_object_layout(value, pos, end)) + keys = { + key_id: key for key, key_id in _metadata_key_ids(metadata).items() + } + children = {} + for slot in range(size): + key_id = _read_unsigned(value, id_start + slot * id_size, id_size) + if key_id not in keys: + _malformed("object key is missing from metadata") + children[keys[key_id]] = _checked_object_child_bounds( + value, data_start, offsets, slot) + return children + + +def _variant_array_children(value, pos, end): + size, data_start, offsets, _ = _checked_array_layout( + value, pos, end) + children = [] + for index in range(size): + child_start = data_start + offsets[index] + child_end = data_start + offsets[index + 1] + if _checked_value_size(value, child_start, child_end) != ( + child_end - child_start): + _malformed("child size does not match container offsets") + children.append((child_start, child_end)) + return children + + +def _supports_exact_get(data_type): + if (pa.types.is_boolean(data_type) + or pa.types.is_int64(data_type) + or pa.types.is_float32(data_type) + or pa.types.is_float64(data_type) + or pa.types.is_string(data_type) + or pa.types.is_large_string(data_type) + or pa.types.is_binary(data_type) + or pa.types.is_large_binary(data_type) + or pa.types.is_date32(data_type) + or pa.types.is_decimal128(data_type)): + return True + if pa.types.is_timestamp(data_type): + return data_type.unit == 'us' + if pa.types.is_struct(data_type): + return all(_supports_exact_get(field.type) for field in data_type) + if (pa.types.is_list(data_type) + or pa.types.is_large_list(data_type) + or pa.types.is_fixed_size_list(data_type)): + return _supports_exact_get(data_type.value_type) + if pa.types.is_map(data_type): + return ((pa.types.is_string(data_type.key_type) + or pa.types.is_large_string(data_type.key_type)) + and _supports_exact_get(data_type.item_type)) + return False + + +def _exact_primitive_matches(value, pos, data_type): + variant_type = _variant_get_type(value, pos) + if variant_type == _Type.NULL: + return True + if pa.types.is_boolean(data_type): + return variant_type == _Type.BOOLEAN + if pa.types.is_int64(data_type): + return variant_type == _Type.LONG + if pa.types.is_float32(data_type): + return variant_type == _Type.FLOAT + if pa.types.is_float64(data_type): + return variant_type == _Type.DOUBLE + if pa.types.is_string(data_type) or pa.types.is_large_string(data_type): + return variant_type == _Type.STRING + if pa.types.is_binary(data_type) or pa.types.is_large_binary(data_type): + return variant_type == _Type.BINARY + if pa.types.is_date32(data_type): + return variant_type == _Type.DATE + if pa.types.is_timestamp(data_type): + expected = _Type.TIMESTAMP if data_type.tz else _Type.TIMESTAMP_NTZ + return variant_type == expected + if pa.types.is_decimal128(data_type): + if variant_type != _Type.DECIMAL: + return False + scale = value[pos + 1] + return scale == data_type.scale + return False + + +def _decode_exact(value, metadata, pos, data_type): + size = _checked_value_size(value, pos) + end = pos + size + variant_type = _variant_get_type(value, pos) + if variant_type == _Type.NULL: + return None + if pa.types.is_struct(data_type): + if variant_type != _Type.OBJECT: + raise TypeError(f"VARIANT path type does not match {data_type}") + children = _variant_object_children(value, metadata, pos, end) + return { + field.name: ( + None if field.name not in children + else _decode_exact( + value, metadata, children[field.name][0], field.type) + ) + for field in data_type + } + if (pa.types.is_list(data_type) + or pa.types.is_large_list(data_type) + or pa.types.is_fixed_size_list(data_type)): + if variant_type != _Type.ARRAY: + raise TypeError(f"VARIANT path type does not match {data_type}") + children = _variant_array_children(value, pos, end) + if (pa.types.is_fixed_size_list(data_type) + and len(children) != data_type.list_size): + raise TypeError(f"VARIANT path type does not match {data_type}") + return [ + _decode_exact(value, metadata, child_pos, data_type.value_type) + for child_pos, _ in children + ] + if pa.types.is_map(data_type): + if variant_type != _Type.OBJECT: + raise TypeError(f"VARIANT path type does not match {data_type}") + return [ + (key, _decode_exact( + value, metadata, child_pos, data_type.item_type)) + for key, (child_pos, _) in _variant_object_children( + value, metadata, pos, end).items() + ] + if not _exact_primitive_matches(value, pos, data_type): + raise TypeError(f"VARIANT path type does not match {data_type}") + return GenericVariant(bytes(value[pos:end]), metadata).to_python() + + def _patched_chunk( chunk: pa.StructArray, values: _BinaryValues, @@ -947,14 +1086,20 @@ def __init__(self, value, length: int): "VARIANT replacement must be an Arrow Scalar or Array") if not _supported_replacement_type(self.type): raise TypeError( - "VARIANT replacement type must be float32 or float64") + f"Unsupported exact VARIANT replacement type: {self.type}") if pa.types.is_float64(self.type): self._value_format = ' bytes: - if value is not None: + if value is not None and self._value_format is not None: return struct.pack( self._value_format, self._type_header, value) - return bytes([_primitive_header(_NULL)]) + return _encode_scalar_to_value_bytes(value, self.type) def validate_source(self, value, pos) -> None: - header = value[pos] - if header not in ( - self._type_header, _primitive_header(_NULL)): + if not _replacement_type_matches(value, pos, self.type): raise TypeError( f"VARIANT path type does not match {self.type}") @@ -1012,6 +1155,9 @@ def _vectorized_replace_chunk( global_row, strict, ): + if not all(provider._fixed_size is not None + for _, _, provider in parsed): + return None valid_rows = _valid_row_indices(chunk, values, chunk.field(1)) if not len(valid_rows): return chunk @@ -1086,12 +1232,14 @@ def _vectorized_replace_chunk( continue original = bytes(value) value = original + for (_, _, provider), pos in zip(parsed, positions): + if pos is not None: + provider.validate_source(value, pos) for (path, parsed_path, provider), pos in zip(parsed, positions): if pos is None: if strict: raise ValueError(f"VARIANT path does not exist: {path}") continue - provider.validate_source(value, pos) replacement_value = provider.scalar_at(global_row + row) value = _replace_path( value, @@ -1114,7 +1262,68 @@ def _vectorized_replace_chunk( def _supported_replacement_type(data_type: pa.DataType) -> bool: - return pa.types.is_float32(data_type) or pa.types.is_float64(data_type) + return ( + pa.types.is_boolean(data_type) + or pa.types.is_signed_integer(data_type) + or pa.types.is_float32(data_type) + or pa.types.is_float64(data_type) + or pa.types.is_string(data_type) + or pa.types.is_large_string(data_type) + or pa.types.is_binary(data_type) + or pa.types.is_large_binary(data_type) + or pa.types.is_date32(data_type) + or (pa.types.is_timestamp(data_type) and data_type.unit == 'us') + or pa.types.is_decimal128(data_type) + ) + + +def _replacement_type_matches(value, pos, data_type): + variant_type = _variant_get_type(value, pos) + if variant_type == _Type.NULL: + return True + if pa.types.is_signed_integer(data_type): + return variant_type == _Type.LONG + return _exact_primitive_matches(value, pos, data_type) + + +def _rowwise_replace_chunk( + chunk, values, parsed, parsed_paths, global_row, strict): + metadata = _BinaryValues(chunk.field(1)) + valid = chunk.is_valid().to_pylist() + rebuilt_rows = {} + for row in range(len(chunk)): + if not valid[row]: + continue + original = values.view(row) + row_metadata = bytes(metadata.view(row)) + positions = _path_positions(original, row_metadata, parsed_paths) + for (path, _, provider), pos in zip(parsed, positions): + if pos is None: + if strict: + raise ValueError( + f"VARIANT path does not exist: {path}") + continue + provider.validate_source(original, pos) + value = None + for (path, parsed_path, provider), pos in zip(parsed, positions): + if pos is None: + continue + if value is None: + value = bytes(original) + value = _replace_path( + value, + row_metadata, + 0, + parsed_path, + provider.encode(provider.scalar_at(global_row + row)), + ) + if value is not None and value != original: + rebuilt_rows[row] = value + if not rebuilt_rows: + return chunk + data, data_start = values.copy_used_data() + return _sparse_rebuilt_chunk( + chunk, values, data, data_start, rebuilt_rows) def _variant_get(column, paths: Mapping[str, pa.DataType]): @@ -1122,9 +1331,9 @@ def _variant_get(column, paths: Mapping[str, pa.DataType]): for path, target_type in paths.items(): if not isinstance(target_type, pa.DataType): raise TypeError("VARIANT data_type must be a PyArrow data type") - if not (pa.types.is_float32(target_type) - or pa.types.is_float64(target_type)): - raise TypeError("VARIANT data_type must be float32 or float64") + if not _supports_exact_get(target_type): + raise TypeError( + f"Unsupported exact VARIANT data type: {target_type}") parsed.append((path, _parse_path(path), target_type)) parsed_paths = [parsed_path for _, parsed_path, _ in parsed] chunks, chunked, _ = _variant_chunks(column) @@ -1137,8 +1346,31 @@ def _variant_get(column, paths: Mapping[str, pa.DataType]): parsed_paths, [target_type for _, _, target_type in parsed], ) - for (path, _, _), result in zip(parsed, results): - result_chunks[path].append(result) + if results is not None: + for (path, _, _), result in zip(parsed, results): + result_chunks[path].append(result) + continue + metadata = _BinaryValues(chunk.field(1)) + valid = chunk.is_valid().to_pylist() + decoded = {path: [] for path in paths} + for row in range(len(chunk)): + if not valid[row]: + for path in paths: + decoded[path].append(None) + continue + value = values.view(row) + row_metadata = bytes(metadata.view(row)) + positions = _path_positions( + value, row_metadata, parsed_paths) + for (path, _, data_type), pos in zip(parsed, positions): + decoded[path].append( + None if pos is None + else _decode_exact( + value, row_metadata, pos, data_type) + ) + for path, _, data_type in parsed: + result_chunks[path].append( + pa.array(decoded[path], type=data_type)) if not chunked: return {path: chunks[0] for path, chunks in result_chunks.items()} return { @@ -1148,7 +1380,7 @@ def _variant_get(column, paths: Mapping[str, pa.DataType]): def variant_get(column, path, data_type=None): - """Read FLOAT or DOUBLE paths into matching Arrow arrays.""" + """Read one or more VARIANT paths without implicit casts.""" if isinstance(path, Mapping): if data_type is not None: raise TypeError( @@ -1178,7 +1410,7 @@ def variant_replace( replacement=None, strict: bool = False, ): - """Replace existing FLOAT or DOUBLE paths with matching Arrow values.""" + """Replace one or more existing VARIANT paths without implicit casts.""" if not isinstance(strict, bool): raise TypeError("VARIANT strict must be a boolean") if isinstance(path, Mapping): @@ -1202,14 +1434,24 @@ def variant_replace( global_row = 0 for chunk in chunks: values = _BinaryValues(chunk.field(0)) - result_chunks.append(_vectorized_replace_chunk( + result = _vectorized_replace_chunk( chunk, values, parsed, parsed_paths, global_row, strict, - )) + ) + if result is None: + result = _rowwise_replace_chunk( + chunk, + values, + parsed, + parsed_paths, + global_row, + strict, + ) + result_chunks.append(result) global_row += len(chunk) if not chunked: diff --git a/paimon-python/pypaimon/data/variant_shredding.py b/paimon-python/pypaimon/data/variant_shredding.py index 72a4508dd32a..18692cbbe602 100644 --- a/paimon-python/pypaimon/data/variant_shredding.py +++ b/paimon-python/pypaimon/data/variant_shredding.py @@ -250,21 +250,64 @@ def _append_scalar(builder, value, arrow_type: pa.DataType) -> None: elif pa.types.is_timestamp(arrow_type): # PyArrow converts timestamp to datetime.datetime if isinstance(value, datetime.datetime): + if (arrow_type.unit == 'ns' + and getattr(value, 'nanosecond', 0) != 0): + raise ValueError( + "VARIANT timestamps require microsecond-aligned values") if value.tzinfo is not None: epoch = datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) - micros = int((value - epoch).total_seconds() * 1_000_000) + delta = value - epoch + micros = ( + (delta.days * 86400 + delta.seconds) * 1_000_000 + + delta.microseconds + ) builder.append_timestamp(micros) else: epoch = datetime.datetime(1970, 1, 1) - micros = int((value - epoch).total_seconds() * 1_000_000) + delta = value - epoch + micros = ( + (delta.days * 86400 + delta.seconds) * 1_000_000 + + delta.microseconds + ) builder.append_timestamp_ntz(micros) else: - builder.append_timestamp_ntz(int(value)) + raw_value = int(value) + if arrow_type.unit == 'ns': + if raw_value % 1000: + raise ValueError( + "VARIANT timestamps require microsecond-aligned values") + micros = raw_value // 1000 + else: + micros = raw_value * {'s': 1_000_000, 'ms': 1000, 'us': 1}[ + arrow_type.unit] + builder.append_timestamp_ntz(micros) elif pa.types.is_decimal(arrow_type): - if isinstance(value, _decimal.Decimal): - builder.append_decimal(value) - else: - builder.append_decimal(_decimal.Decimal(str(value))) + decimal = ( + value if isinstance(value, _decimal.Decimal) + else _decimal.Decimal(str(value)) + ) + sign, digits, exponent = decimal.as_tuple() + unscaled = int(''.join(str(digit) for digit in digits)) + if sign: + unscaled = -unscaled + shift = exponent + arrow_type.scale + if shift < 0 and unscaled % (10 ** -shift): + raise ValueError( + f'{decimal} does not have Arrow scale {arrow_type.scale}') + unscaled = ( + unscaled * (10 ** shift) if shift >= 0 + else unscaled // (10 ** -shift) + ) + scale = arrow_type.scale + precision = max(1, len(str(abs(unscaled)))) + if precision > arrow_type.precision: + raise ValueError( + f'{decimal} exceeds Arrow precision {arrow_type.precision}') + if scale < 0: + unscaled *= 10 ** -scale + scale = 0 + precision = max(1, len(str(abs(unscaled)))) + builder.append_decimal_unscaled(unscaled, precision, scale) else: # Fallback: encode as string builder.append_string(str(value)) diff --git a/paimon-python/pypaimon/tests/variant_path_test.py b/paimon-python/pypaimon/tests/variant_path_test.py index 1455e4cb2881..d5797c7f2dfc 100644 --- a/paimon-python/pypaimon/tests/variant_path_test.py +++ b/paimon-python/pypaimon/tests/variant_path_test.py @@ -14,7 +14,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +import datetime import unittest +from decimal import Decimal from unittest.mock import patch import numpy as np @@ -24,6 +26,7 @@ from pypaimon.data._variant_binary import _primitive_header from pypaimon.data.generic_variant import _DOUBLE, GenericVariant from pypaimon.data.variant_path import ( + _metadata_key_ids, _path_positions, _rebuilt_offsets, variant_get, @@ -59,6 +62,19 @@ def _decode(column): ] +def _typed_object(fields): + metadata = GenericVariant.from_python({ + name: 0 for name in fields + }).metadata() + key_ids = _metadata_key_ids(metadata) + value = _build_object_value([ + (key_ids[name], _encode_scalar_to_value_bytes(item, data_type)) + for name, (item, data_type) in fields.items() + ]) + return GenericVariant.to_arrow_array([ + GenericVariant(value, metadata)]) + + class TestVariantGet(unittest.TestCase): def test_nested_paths_and_missing_values(self): @@ -98,7 +114,7 @@ def test_reads_multiple_paths_in_one_pass(self): self.assertEqual(result['$.velocity.x'].to_pylist(), [1.0, 3.0]) self.assertEqual(result['$.velocity.y'].to_pylist(), [-2.0, -4.0]) - def test_requires_exact_float_type(self): + def test_requires_exact_type(self): cases = ( (_float_variants([1.25]), pa.float64()), (_variants([1.25]), pa.float32()), @@ -109,8 +125,100 @@ def test_requires_exact_float_type(self): with self.assertRaisesRegex(TypeError, "does not match"): variant_get(column, '$', data_type) - with self.assertRaisesRegex(TypeError, "float32 or float64"): + with self.assertRaisesRegex(TypeError, "does not match"): variant_get(_variants([1.0]), '$', pa.string()) + with self.assertRaisesRegex(TypeError, "Unsupported exact"): + variant_get(_variants([1]), '$', pa.int32()) + + def test_reads_exact_primitive_types(self): + timestamp = datetime.datetime(2026, 8, 11, 1, 2, 3, 456000) + column = _typed_object({ + 'flag': (True, pa.bool_()), + 'count': (123, pa.int64()), + 'text': ('hello', pa.string()), + 'binary': (b'abc', pa.binary()), + 'decimal': (Decimal('12.30'), pa.decimal128(4, 2)), + 'date': (datetime.date(2026, 8, 11), pa.date32()), + 'timestamp': (timestamp, pa.timestamp('us')), + }) + result = variant_get(column, { + '$.flag': pa.bool_(), + '$.count': pa.int64(), + '$.text': pa.string(), + '$.binary': pa.binary(), + '$.decimal': pa.decimal128(4, 2), + '$.date': pa.date32(), + '$.timestamp': pa.timestamp('us'), + }) + + self.assertEqual( + {path: array[0].as_py() for path, array in result.items()}, + { + '$.flag': True, + '$.count': 123, + '$.text': 'hello', + '$.binary': b'abc', + '$.decimal': Decimal('12.30'), + '$.date': datetime.date(2026, 8, 11), + '$.timestamp': timestamp, + }, + ) + + def test_reads_exact_complex_types(self): + column = _variants([{ + 'object': {'count': 2, 'flag': True}, + 'array': [1, 2], + 'map': {'left': 1, 'right': 2}, + }]) + struct_type = pa.struct([ + ('count', pa.int64()), + ('flag', pa.bool_()), + ('missing', pa.string()), + ]) + + self.assertEqual( + variant_get(column, '$.object', struct_type).to_pylist(), + [{'count': 2, 'flag': True, 'missing': None}], + ) + self.assertEqual( + variant_get( + column, '$.array', pa.list_(pa.int64())).to_pylist(), + [[1, 2]], + ) + self.assertEqual( + variant_get( + column, + '$.map', + pa.map_(pa.string(), pa.int64()), + ).to_pylist(), + [[('left', 1), ('right', 2)]], + ) + + def test_decimal_extraction_preserves_38_digits(self): + expected = Decimal('12345678901234567890123456789012345678') + column = _typed_object({ + 'value': (expected, pa.decimal128(38, 0)), + }) + + result = variant_get( + column, '$.value', pa.decimal128(38, 0)) + + self.assertEqual(result.to_pylist(), [expected]) + + def test_rejects_cross_type_casts(self): + column = _variants([{ + 'count': 123, + 'object': {'value': 1}, + 'array': [1], + }]) + for path, data_type in ( + ('$.count', pa.string()), + ('$.object', pa.string()), + ('$.array', pa.string()), + ('$.array', pa.list_(pa.string()))): + with self.subTest(path=path, data_type=data_type): + with self.assertRaisesRegex(TypeError, "does not match"): + variant_get(column, path, data_type) def test_variant_null_is_arrow_null(self): column = _variants([None, {'value': None}, {'value': 1.0}]) @@ -166,6 +274,42 @@ def test_rejects_invalid_arguments(self): class TestVariantReplace(unittest.TestCase): + def test_replaces_exact_primitive_types(self): + original_timestamp = datetime.datetime(2026, 8, 11) + column = _typed_object({ + 'flag': (True, pa.bool_()), + 'count': (1, pa.int64()), + 'text': ('old', pa.string()), + 'binary': (b'old', pa.binary()), + 'decimal': (Decimal('1.00'), pa.decimal128(3, 2)), + 'date': (datetime.date(2026, 8, 10), pa.date32()), + 'timestamp': (original_timestamp, pa.timestamp('us')), + }) + new_timestamp = datetime.datetime(2026, 8, 11, 1, 2, 3, 4) + + result = variant_replace(column, { + '$.flag': pa.scalar(False), + '$.count': pa.scalar(2, type=pa.int64()), + '$.text': pa.scalar('new'), + '$.binary': pa.scalar(b'new'), + '$.decimal': pa.scalar( + Decimal('2.50'), type=pa.decimal128(3, 2)), + '$.date': pa.scalar( + datetime.date(2026, 8, 11), type=pa.date32()), + '$.timestamp': pa.scalar( + new_timestamp, type=pa.timestamp('us')), + }) + + self.assertEqual(_decode(result), [{ + 'flag': False, + 'count': 2, + 'text': 'new', + 'binary': b'new', + 'decimal': Decimal('2.50'), + 'date': datetime.date(2026, 8, 11), + 'timestamp': new_timestamp, + }]) + def test_get_compute_replace_pipeline(self): column = pa.chunked_array([ _variants([{'x': 1.0, 'y': -2.0}, None]), @@ -224,6 +368,10 @@ def test_float_and_double_are_distinct(self): variant_replace( _variants([1.0]), '$', pa.scalar(1.0, type=pa.float32())) + with self.assertRaisesRegex(TypeError, "does not match"): + variant_replace( + _variants([1.0]), '$', pa.scalar('1.0', type=pa.string())) + def test_nullable_rows_stay_vectorized(self): size = 4096 column = _variants( @@ -355,7 +503,7 @@ def test_rejects_invalid_arguments(self): ('$.value', 1.0, False, TypeError, "Arrow Scalar or Array"), ('$.value', pa.array([1, 2]), False, - TypeError, "float32 or float64"), + TypeError, "does not match"), ('$.value', pa.scalar(1.0), 'yes', TypeError, "strict must be a boolean"), ]