Google Git Google Git
Sign in
apache/paimon/refs/heads/codex/append-unknown-row-count/./paimon-python/pypaimon/data/variant_path.py
blob: fea3291d5a37d70473463659bb6537ba0fe42ecb [file]
# 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.
"""Read, replace, and upsert paths in Arrow VARIANT columns."""
import functools
import re
import struct
import threading
from typing import Dict, Mapping, Optional, Sequence, Tuple
import numpy as np
import pyarrow as pa
from cachetools import LRUCache
from pypaimon.data._variant_binary import (
_ARRAY,
_OBJECT,
_PRIMITIVE,
_SHORT_STR,
_U8_MAX,
_U32_SIZE,
_VERSION,
_VERSION_MASK,
_array_header,
_get_int_size,
_object_header,
_primitive_header,
_read_unsigned,
)
from pypaimon.data.generic_variant import (
_BINARY,
_DECIMAL4,
_DECIMAL8,
_DECIMAL16,
_DOUBLE,
_FLOAT,
_LONG_STR,
_NULL,
_MAX_DECIMAL4_PRECISION,
_MAX_DECIMAL8_PRECISION,
_MAX_DECIMAL16_PRECISION,
_PRIMITIVE_FIXED_SIZES,
GenericVariant,
_Type,
_check_variant_sizes,
_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+)]")
_KEY_PATTERN = re.compile(r"\.([^\.\[]+)|\['([^']+)']|\[\"([^\"]+)\"]")
_Path = Tuple[Tuple[str, object], ...]
_SLOW_PATH_ROWS = 64
# Bound the dominant temporary allocation during batch structure matching.
_STRUCTURE_MATCH_INDEX_BUDGET = 8 * 1024 * 1024
# Bound encoded variable-width payloads retained while rebuilt rows accumulate.
_ROOT_INSERT_SPLICE_PAYLOAD_BUDGET = 8 * 1024 * 1024
# Bound per-row Python and NumPy temporaries for tiny payloads.
_ROOT_INSERT_SPLICE_MAX_BATCH_ROWS = 64 * 1024
_FAST_CHILD_SIZES = np.asarray([
_PRIMITIVE_FIXED_SIZES.get(type_info, 0)
if type_info not in (_DECIMAL4, _DECIMAL8, _DECIMAL16) else 0
for type_info in range(64)
], dtype=np.uint8)
@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))))
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()
return tuple(segments)
def _metadata_key_ids(metadata: bytes) -> Dict[str, int]:
if not 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)
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")
_metadata_cache = threading.local()
_NO_METADATA_CACHE = object()
def _with_metadata_cache(func):
@functools.wraps(func)
def wrapper(*args, **kwargs):
previous = getattr(_metadata_cache, 'value', _NO_METADATA_CACHE)
_metadata_cache.value = LRUCache(maxsize=256)
try:
return func(*args, **kwargs)
finally:
if previous is _NO_METADATA_CACHE:
del _metadata_cache.value
else:
_metadata_cache.value = previous
return wrapper
def _cached_metadata_key_ids(metadata: bytes) -> Dict[str, int]:
cache = getattr(_metadata_cache, 'value', None)
if cache is not None:
hit = cache.get(metadata)
if hit is not None:
return hit
key_ids = _metadata_key_ids(metadata)
if cache is not None:
cache[metadata] = key_ids
return key_ids
def _metadata_with_keys(metadata: bytes, new_keys: Tuple[str, ...]):
"""Append missing keys to the metadata dictionary, keeping ids.
Returns read-only ``(new_metadata, key_ids, names_by_id)``;
``new_metadata`` is None when every key already exists.
"""
cache_key = (metadata, new_keys)
cache = getattr(_metadata_cache, 'value', None)
if cache is not None:
hit = cache.get(cache_key)
if hit is not None:
return hit
key_ids = dict(_cached_metadata_key_ids(metadata))
names = [None] * len(key_ids)
for key, key_id in key_ids.items():
names[key_id] = key
missing = list(dict.fromkeys(
key for key in new_keys if key not in key_ids))
for key in missing:
key_ids[key] = len(names)
names.append(key)
names_by_id = {key_id: key for key, key_id in key_ids.items()}
if not missing:
result = (None, key_ids, names_by_id)
else:
encoded = [name.encode('utf-8') for name in names]
total_size = sum(len(name) for name in encoded)
max_size = max(total_size, len(encoded))
offset_size = _get_int_size(max_size) if max_size > 0 else 1
offset_start = 1 + offset_size
string_start = offset_start + (len(encoded) + 1) * offset_size
metadata_size = string_start + total_size
_check_variant_sizes(0, metadata_size)
rebuilt = bytearray(metadata_size)
rebuilt[0] = _VERSION | ((offset_size - 1) << 6)
rebuilt[1:1 + offset_size] = len(encoded).to_bytes(
offset_size, 'little')
current = 0
for index, name in enumerate(encoded):
rebuilt[offset_start + index * offset_size:
offset_start + (index + 1) * offset_size] = (
current.to_bytes(offset_size, 'little'))
rebuilt[string_start + current:
string_start + current + len(name)] = name
current += len(name)
rebuilt[offset_start + len(encoded) * offset_size:string_start] = (
current.to_bytes(offset_size, 'little'))
result = (bytes(rebuilt), key_ids, names_by_id)
if cache is not None:
cache[cache_key] = result
return result
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 ((not size and sentinel != 0)
or (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")
if size and len({
_read_unsigned(value, id_start + i * id_width, id_width)
for i in range(size)}) != size:
_malformed("duplicate object field id")
_require_range(data_start, sentinel, limit)
return (
size, id_width, id_start, data_start, offsets,
data_start + offsets[-1],
)
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
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]
basic_type = header & 0x3
type_info = (header >> 2) & 0x3F
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
_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 = (
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
def _validate_value_field_ids(
value, pos, limit, metadata_size, structure_ranges=None):
"""Validate object field ids in one unedited value subtree."""
stack = [(pos, limit)]
while stack:
current_pos, current_limit = stack.pop()
_require_range(current_pos, 1, current_limit)
basic_type = value[current_pos] & 0x3
if basic_type == _OBJECT:
size, id_size, id_start, data_start, offsets, value_end = (
_checked_object_layout(
value, current_pos, current_limit))
if structure_ranges is not None:
structure_ranges.append((current_pos, data_start))
ids = [
_read_unsigned(value, id_start + i * id_size, id_size)
for i in range(size)
]
if any(field_id >= metadata_size for field_id in ids):
_malformed("object field id is missing from metadata")
ordered_offsets = sorted(offsets)
end_by_offset = dict(zip(
ordered_offsets, ordered_offsets[1:]))
for slot in range(size):
child_offset = offsets[slot]
stack.append((
data_start + child_offset,
data_start + end_by_offset[child_offset],
))
elif basic_type == _ARRAY:
size, data_start, offsets, value_end = _checked_array_layout(
value, current_pos, current_limit)
if structure_ranges is not None:
structure_ranges.append((current_pos, data_start))
for index in range(size):
stack.append((
data_start + offsets[index],
data_start + offsets[index + 1],
))
else:
value_end = current_pos + _checked_value_size(
value, current_pos, current_limit)
if structure_ranges is not None:
type_info = (value[current_pos] >> 2) & 0x3F
structure_end = current_pos + 1
if basic_type == _PRIMITIVE:
if type_info in (_BINARY, _LONG_STR):
structure_end += _U32_SIZE
elif type_info in (
_DECIMAL4, _DECIMAL8, _DECIMAL16):
structure_end = value_end
structure_ranges.append((current_pos, structure_end))
if value_end != current_limit:
_malformed("child size does not match container offsets")
def _matching_value_structures(
value, source_data, row_starts, metadata_size):
"""Match equal-length rows against one validated value structure."""
ranges = []
_validate_value_field_ids(
value, 0, len(value), metadata_size, ranges)
position_count = sum(end - start for start, end in ranges)
index_size = np.dtype(np.int64).itemsize
position_bytes = position_count * index_size
available_bytes = _STRUCTURE_MATCH_INDEX_BUDGET - position_bytes
if available_bytes < index_size:
return None
max_cells = available_bytes // index_size
positions = np.empty(position_count, dtype=np.int64)
cursor = 0
for start, end in ranges:
count = end - start
positions[cursor:cursor + count] = np.arange(
start, end, dtype=np.int64)
cursor += count
matches = np.ones(len(row_starts), dtype=bool)
expected = np.frombuffer(value, dtype=np.uint8)
width = max(1, min(position_count, max_cells))
row_chunk_size = min(len(matches), max(1, max_cells // width))
for row_start in range(0, len(matches), row_chunk_size):
row_end = min(row_start + row_chunk_size, len(matches))
batch_matches = matches[row_start:row_end]
batch_starts = row_starts[row_start:row_end]
for start in range(0, position_count, width):
offsets = positions[start:start + width]
batch_matches &= np.all(
source_data[batch_starts[:, None] + offsets]
== expected[offsets],
axis=1,
)
if (start + width < position_count
and not np.any(batch_matches)):
break
return matches
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
@functools.lru_cache(maxsize=256)
def _compile_paths(paths: Tuple[_Path, ...]):
nodes = [(None, None, None)]
node_by_edge = {}
results = []
for path in paths:
parent = 0
for segment in path:
edge = (parent, segment)
node = node_by_edge.get(edge)
if node is None:
node = len(nodes)
node_by_edge[edge] = node
nodes.append((parent,) + segment)
parent = node
results.append(parent)
return tuple(nodes), tuple(results)
def _path_positions(
value: bytes,
metadata: bytes,
paths: Sequence[_Path],
) -> Sequence[Optional[int]]:
root_size = _checked_value_size(value, 0)
if root_size != len(value):
_malformed("trailing bytes after root value")
nodes, result_nodes = _compile_paths(tuple(paths))
_validate_metadata_version(metadata)
key_ids = (
_cached_metadata_key_ids(bytes(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]
if parent is None:
bounds.append(None)
continue
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 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, child_end = _checked_object_child_bounds(
value, data_start, offsets, slot)
else:
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 _replace_path(
value: bytes,
metadata: bytes,
pos: int,
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)
if not path:
return replacement
kind, segment = path[0]
if kind == 'key':
if (value[pos] & 0x3) != _OBJECT:
raise ValueError("VARIANT path expects an object")
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, _ = (
_checked_object_layout(value, pos, value_end))
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}")
ordered_offsets = sorted(offsets)
end_by_offset = dict(zip(ordered_offsets, ordered_offsets[1:]))
children = []
for i in range(size):
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(
value, metadata, child_pos, path[1:], replacement,
child_end, key_ids)
children.append(child)
return _build_object_value(list(zip(ids, children)))
if (value[pos] & 0x3) != _ARRAY:
raise ValueError("VARIANT path expects an array")
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 + 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,
child_end, key_ids)
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
elif isinstance(column, pa.Array):
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)
or [field.name for field in data_type]
!= ['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")
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
class _BinaryValues:
def __init__(self, array: pa.Array):
self.array = array
if pa.types.is_binary(array.type):
self.width, self.value_format = 4, '<i'
elif pa.types.is_large_binary(array.type):
self.width, self.value_format = 8, '<q'
else:
raise TypeError("VARIANT binary field has an unsupported type")
self.offsets = array.buffers()[1]
data_buffer = array.buffers()[2]
self.data = (memoryview(data_buffer) if data_buffer is not None
else memoryview(b''))
def bounds(self, row: int) -> 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 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)
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 _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
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, 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(lengths) 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))
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_rows,
paths: Sequence[_Path],
):
if not len(valid_rows):
return None
metadata_values = _BinaryValues(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
nodes, result_nodes = _compile_paths(tuple(paths))
_validate_metadata_version(first_metadata)
key_ids = (
_cached_metadata_key_ids(first_metadata)
if any(kind == 'key' for _, kind, _ in nodes[1:]) else {}
)
row_offsets = values.numpy_offsets()
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(first_row)
positions = [np.zeros(len(valid_rows), dtype=np.int64)]
limits = [row_ends - row_starts]
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)
or np.any(absolute_parent >= parent_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:
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
id_table = bytes(
first_value[id_start:id_start + size * id_size])
slot = _field_slot(id_table, id_size, key_id)
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 != 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
for index in range(size):
expected_id = _read_unsigned(
first_value, id_start + index * id_size, id_size)
ids = _take_unsigned(
data, id_starts + index * id_widths, id_widths)
if np.any(ids != expected_id):
return None
if slot is None:
positions.append(None)
limits.append(None)
continue
successor_slot = min(
(
index for index in range(size + 1)
if first_offsets[index] > first_offsets[slot]
),
key=lambda index: first_offsets[index],
)
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 != size):
return None
slot = segment
successor_slot = slot + 1
offset_widths = (type_info & 0x3) + 1
offset_starts = absolute_parent + 1 + size_widths
data_starts = offset_starts + (sizes + 1) * offset_widths
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,
)
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 (
valid_rows,
row_starts,
data,
tuple(positions[node] for node in result_nodes),
tuple(limits[node] for node in result_nodes),
)
def _partition_path_plans(values, metadata, valid_rows, parsed_paths):
planned = _vectorized_path_positions(
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
valid_rows = _valid_row_indices(chunk, values, chunk.field(1))
if not len(valid_rows):
return [pa.nulls(len(chunk), type=target_type)
for target_type in target_types]
plans, slow_rows = _partition_path_plans(
values, chunk.field(1), valid_rows, parsed_paths)
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):
if pos is None:
results.append(pa.nulls(len(chunk), type=target_type))
continue
absolute = row_starts + pos
headers = data[absolute]
type_info = (
_FLOAT if pa.types.is_float32(target_type) else _DOUBLE)
value_size = 4 if type_info == _FLOAT else 8
data_type = (
np.dtype('<f4') if type_info == _FLOAT else np.dtype('<f8'))
if not np.all(headers == _primitive_header(type_info)):
uniform = False
break
if np.any(absolute + 1 + value_size != row_starts + limit):
uniform = False
break
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))
if uniform:
return results
outputs = [
np.empty(
len(chunk),
dtype=np.dtype('<f4') if pa.types.is_float32(target_type)
else np.dtype('<f8'),
)
for target_type in target_types
]
masks = [np.ones(len(chunk), dtype=bool) for _ in target_types]
slow_by_path = [set(int(row) for row in slow_rows)
for _ in 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]
target_type = target_types[index]
type_info = (
_FLOAT if pa.types.is_float32(target_type) else _DOUBLE)
value_size = 4 if type_info == _FLOAT else 8
data_type = (
np.dtype('<f4') if type_info == _FLOAT else np.dtype('<f8'))
handled = (
(headers == _primitive_header(type_info))
& (absolute + 1 + value_size == row_starts + limit)
)
if np.any(handled):
selected_rows = rows[handled]
selected_absolute = absolute[handled]
indices = (
selected_absolute[:, None] + 1 + np.arange(value_size)
)
raw = np.ascontiguousarray(data[indices])
outputs[index][selected_rows] = raw.view(
data_type).reshape(-1)
masks[index][selected_rows] = False
slow_by_path[index].update(
int(row) for row in rows[~handled])
metadata = _BinaryValues(chunk.field(1))
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)
for index, (pos, target_type) in enumerate(
zip(positions, target_types)):
if row not in slow_by_path[index] or pos is None:
continue
decoded = _decode_floating(value, pos, target_type)
if decoded is not None:
outputs[index][row] = decoded
masks[index][row] = False
return [
pa.array(output, mask=mask, type=target_type)
for output, mask, target_type in zip(outputs, masks, target_types)
]
def _decode_floating(value, pos, target_type):
size = _checked_value_size(value, pos)
header = value[pos]
if (header & 0x3) != _PRIMITIVE:
raise TypeError("VARIANT path is not FLOAT or DOUBLE")
type_info = (header >> 2) & 0x3F
if type_info == _FLOAT and pa.types.is_float32(target_type):
return struct.unpack_from('<f', value, pos + 1)[0]
if type_info == _DOUBLE and pa.types.is_float64(target_type):
return struct.unpack_from('<d', value, pos + 1)[0]
if type_info == _NULL and size == 1:
return None
raise TypeError(
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_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_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 _validate_decimal_scale(data_type):
if pa.types.is_decimal128(data_type) and data_type.scale < 0:
raise ValueError("VARIANT decimal scale must be non-negative")
if pa.types.is_struct(data_type):
for field in data_type:
_validate_decimal_scale(field.type)
elif (pa.types.is_list(data_type)
or pa.types.is_large_list(data_type)
or pa.types.is_fixed_size_list(data_type)):
_validate_decimal_scale(data_type.value_type)
elif pa.types.is_map(data_type):
_validate_decimal_scale(data_type.item_type)
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_signed_integer(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,
data: bytearray,
start: int,
) -> pa.StructArray:
patched_values = values.array_from_data(data, start)
metadata = chunk.field(1)
if chunk.offset == 0 and patched_values.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(
[patched_values, metadata],
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
if total > maximum:
kind = 'Binary' if value_format == '<i' else 'LargeBinary'
suffix = '; use LargeBinary' if kind == 'Binary' else ''
raise ValueError(
f'Rebuilt VARIANT values exceed the {kind} offset limit{suffix}')
offsets = np.empty(len(lengths) + 1, dtype=np.dtype(value_format))
offsets[0] = 0
np.cumsum(lengths, out=offsets[1:])
return offsets
def _sparse_rebuilt_binary(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 = _rebuilt_offsets(lengths, values.value_format)
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]
)
return pa.Array.from_buffers(
values.array.type,
len(values.array),
[validity, pa.py_buffer(offsets), pa.py_buffer(output)],
null_count=values.array.null_count,
)
def _sparse_rebuilt_chunk(
chunk, values, data, data_start, rebuilt_rows):
rebuilt_values = _sparse_rebuilt_binary(
values, data, data_start, rebuilt_rows)
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):
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")
_validate_decimal_scale(self.type)
if not _supported_replacement_type(self.type):
raise TypeError(
f"Unsupported exact VARIANT replacement type: {self.type}")
if pa.types.is_float64(self.type):
self._value_format = '<Bd'
self._type_header = _primitive_header(_DOUBLE)
elif pa.types.is_float32(self.type):
self._value_format = '<Bf'
self._type_header = _primitive_header(_FLOAT)
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 scalar_at(self, row: int):
if self._array is None:
return self._value.as_py()
return self._array[row].as_py()
def numpy_values(self, offset: int, length: int, rows=None):
if self._value_format is None:
return None
data_type = (
np.dtype('<f8') if pa.types.is_float64(self.type)
else np.dtype('<f4')
)
if self._array is None:
size = length if rows is None else len(rows)
if not self._value.is_valid:
return np.empty(size, dtype=data_type), np.zeros(
size, dtype=bool)
return (
np.full(size, self._value.as_py(), dtype=data_type),
np.ones(size, dtype=bool),
)
values = self._array.slice(offset, length)
if isinstance(values, pa.ChunkedArray):
values = values.combine_chunks()
if rows is not None:
values = values.take(pa.array(rows, type=pa.int64()))
valid = np.asarray(
values.is_valid().to_numpy(zero_copy_only=False), dtype=bool)
return (
np.asarray(values.to_numpy(zero_copy_only=False), dtype=data_type),
valid,
)
def encode(self, value) -> 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 validate_source(self, value, pos) -> None:
if not _replacement_type_matches(value, pos, self.type):
raise TypeError(
f"VARIANT path type does not match {self.type}")
class _PatchState:
"""Lazy copy-on-write buffer shared by per-group in-place patches."""
def __init__(self, values: _BinaryValues):
self._values = values
self.data = None
self.data_start = 0
self.output_data = None
def ensure(self):
if self.data is None:
self.data, self.data_start = self._values.copy_used_data()
self.output_data = np.frombuffer(self.data, dtype=np.uint8)
return self.output_data
def _patch_planned_group(
planned, parsed, chunk_length, global_row, state, slow_rows, strict):
rows, row_starts, source_data, positions, limits = planned
replacements = []
compatible = np.ones(len(rows), dtype=bool)
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,
chunk_length,
None if len(rows) == chunk_length else rows,
)
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))
if not has_replacement:
return
slow_rows.update(int(row) for row in rows[~compatible])
if not np.any(compatible):
return
output_data = state.ensure()
relative_starts = row_starts - state.data_start
compatible_rows = rows[compatible]
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
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
def _vectorized_replace_chunk(
chunk,
values,
parsed,
parsed_paths,
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
plans, slow_rows = _partition_path_plans(
values, chunk.field(1), valid_rows, parsed_paths)
slow_rows = set(int(row) for row in slow_rows)
state = _PatchState(values)
for planned in plans:
_patch_planned_group(
planned, parsed, len(chunk), global_row, state,
slow_rows, strict)
data = state.data
data_start = state.data_start
metadata = _BinaryValues(chunk.field(1))
rebuilt_rows = {}
for row in slow_rows:
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 (_, _, 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
replacement_value = provider.scalar_at(global_row + row)
value = _replace_path(
value,
row_metadata,
0,
parsed_path,
provider.encode(replacement_value),
)
if value != original:
rebuilt_rows[row] = value
if rebuilt_rows:
if data is None:
data, data_start = values.copy_used_data()
return _sparse_rebuilt_chunk(
chunk, values, data, data_start, rebuilt_rows)
if data is None:
return chunk
return _patched_chunk(chunk, values, data, data_start)
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) and data_type.unit == 'us')
or pa.types.is_decimal128(data_type)
)
def _replacement_type_matches(value, pos, data_type):
if pa.types.is_null(data_type):
return True
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 _flat_object_unsigneds(value, start, count, width):
"""Read an object header array without creating Python ints per field."""
if width in (1, 2, 4):
return np.frombuffer(
value, dtype=np.dtype('<u%d' % width), count=count,
offset=start)
raw = np.frombuffer(
value, dtype=np.uint8, count=count * 3,
offset=start).reshape(count, 3)
return (raw[:, 0].astype(np.uint32)
| (raw[:, 1].astype(np.uint32) << 8)
| (raw[:, 2].astype(np.uint32) << 16))
def _flat_object_layout(value):
"""Validate a root object and read its offsets as a NumPy array."""
limit = len(value)
_require_range(0, 2, limit)
type_info = (value[0] >> 2) & 0x3F
size_width = _U32_SIZE if ((type_info >> 4) & 0x1) else 1
_require_range(1, size_width, limit)
size = _read_unsigned(value, 1, size_width)
id_width = ((type_info >> 2) & 0x3) + 1
offset_width = (type_info & 0x3) + 1
id_start = 1 + size_width
offset_start = id_start + size * id_width
data_start = offset_start + (size + 1) * offset_width
_require_range(0, data_start, limit)
offsets = _flat_object_unsigneds(
value, offset_start, size + 1, offset_width)
end_offset = int(offsets[-1])
if size:
ordered = np.sort(offsets)
if (int(ordered[0]) != 0 or int(ordered[-1]) != end_offset
or np.any(np.diff(ordered) == 0)):
_malformed("invalid object offsets")
else:
if end_offset != 0:
_malformed("invalid object offsets")
ordered = offsets
_require_range(data_start, end_offset, limit)
return size, id_width, id_start, offset_start, data_start, offsets, ordered
def _flat_object_get_chunk(chunk, values, parsed):
"""Look up wide, top-level object fields together instead of per path."""
if (len(parsed) < 2
or any(len(path) != 1 or path[0][0] != 'key'
for _, path, _ in parsed)):
return None
valid_rows = _valid_row_indices(chunk, values, chunk.field(1))
if not len(valid_rows):
return [pa.nulls(len(chunk), type=data_type)
for _, _, data_type in parsed]
first_value = values.view(int(valid_rows[0]))
if not first_value or (first_value[0] & 0x3) != _OBJECT:
return None
first_size = _checked_object_layout(
first_value, 0, len(first_value))[0]
# Keep the vectorized reader for narrow objects or a few requested keys.
if first_size * len(parsed) < 1024:
return None
# For multi-row float batches, the existing reader vectorizes across rows.
# Its per-path cost wins over per-row lookup on modest object widths.
if (len(valid_rows) >= _SLOW_PATH_ROWS
and all(pa.types.is_float32(data_type)
or pa.types.is_float64(data_type)
for _, _, data_type in parsed)
and (first_size < 256 or first_size * len(parsed) < 4096)):
return None
metadata = _BinaryValues(chunk.field(1))
decoded = [[None] * len(chunk) for _ in parsed]
# Object field IDs may differ across rows, so the key includes the actual
# encoded layout as well as the metadata dictionary and requested IDs.
slot_cache = LRUCache(maxsize=64)
for row in valid_rows:
row = int(row)
value = values.view(row)
if not value or (value[0] & 0x3) != _OBJECT:
return None
row_metadata = bytes(metadata.view(row))
key_ids = _cached_metadata_key_ids(row_metadata)
targets = {}
for index, (_, path, _) in enumerate(parsed):
key_id = key_ids.get(path[0][1])
if key_id is not None:
targets.setdefault(key_id, []).append(index)
(size, id_width, id_start, offset_start, data_start,
offsets, ordered_offsets) = _flat_object_layout(value)
if data_start + int(offsets[-1]) != len(value):
_malformed("trailing bytes after root value")
layout = (row_metadata, bytes(value[id_start:offset_start]),
id_width, tuple(sorted(targets)))
slots = slot_cache.get(layout)
if slots is None:
field_ids = _flat_object_unsigneds(
value, id_start, size, id_width)
if len(np.unique(field_ids)) != size:
_malformed("duplicate object field id")
slots = {int(key_id): slot
for slot, key_id in enumerate(field_ids)
if int(key_id) in targets}
slot_cache[layout] = slots
if not targets:
continue
if not slots:
continue
selected_offsets = offsets[np.fromiter(
slots.values(), dtype=np.int64, count=len(slots))]
next_indices = np.searchsorted(
ordered_offsets, selected_offsets, side='right')
for (key_id, _), offset, next_index in zip(
slots.items(), selected_offsets, next_indices):
child_pos = data_start + int(offset)
child_end = data_start + int(ordered_offsets[next_index])
if _checked_value_size(value, child_pos, child_end) != (
child_end - child_pos):
_malformed("child size does not match container offsets")
for index in targets[key_id]:
data_type = parsed[index][2]
if (pa.types.is_float32(data_type)
or pa.types.is_float64(data_type)):
decoded[index][row] = _decode_floating(
value, child_pos, data_type)
else:
decoded[index][row] = _decode_exact(
value, row_metadata, child_pos, data_type)
return [pa.array(result, type=data_type)
for result, (_, _, data_type) in zip(decoded, parsed)]
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 data_type must be a PyArrow data type")
_validate_decimal_scale(target_type)
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)
result_chunks = {path: [] for path in paths}
for chunk in chunks:
values = _BinaryValues(chunk.field(0))
results = _flat_object_get_chunk(chunk, values, parsed)
if results is None:
results = _vectorized_get_chunk(
chunk,
values,
parsed_paths,
[target_type for _, _, target_type in parsed],
)
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 {
path: pa.chunked_array(chunks, type=paths[path])
for path, chunks in result_chunks.items()
}
@_with_metadata_cache
def variant_get(column, path, data_type=None):
"""Read one or more VARIANT paths without implicit casts."""
if isinstance(path, Mapping):
if data_type is not None:
raise TypeError(
"VARIANT data_type must be omitted for path mappings")
return _variant_get(column, 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]
@_with_metadata_cache
def variant_to_pylist(column, fields: Sequence[str]):
"""Decode selected top-level VARIANT object fields to Python values.
Values use their natural Python types, even when a field has different
types in different rows. A missing field is omitted from its row's dict;
a VARIANT null is present with value None, and a SQL null row returns None.
Field names are literal names, not VARIANT path expressions.
"""
if isinstance(fields, (str, bytes)):
raise TypeError("VARIANT fields must be a sequence of field names")
try:
names = tuple(dict.fromkeys(fields))
except (TypeError, ValueError):
raise TypeError(
"VARIANT fields must be a sequence of field names") from None
if any(not isinstance(name, str) for name in names):
raise TypeError("VARIANT field names must be strings")
chunks, _, _ = _variant_chunks(column)
result = []
slot_cache = LRUCache(maxsize=64)
for chunk in chunks:
values = _BinaryValues(chunk.field(0))
metadata = _BinaryValues(chunk.field(1))
valid = set(map(int, _valid_row_indices(
chunk, values, chunk.field(1))))
for row in range(len(chunk)):
if row not in valid:
result.append(None)
continue
value = bytes(values.view(row))
row_metadata = bytes(metadata.view(row))
if not value or (value[0] & 0x3) != _OBJECT:
raise TypeError("VARIANT root must be an object")
variant = GenericVariant(value, row_metadata)
key_ids = _cached_metadata_key_ids(row_metadata)
targets = {key_ids[name]: name for name in names
if name in key_ids}
(size, id_width, id_start, offset_start, data_start,
offsets, ordered_offsets) = _flat_object_layout(value)
if data_start + int(offsets[-1]) != len(value):
_malformed("trailing bytes after root value")
layout = (row_metadata, value[id_start:offset_start],
id_width, tuple(sorted(targets)))
selected_slots = slot_cache.get(layout)
if selected_slots is None:
field_ids = _flat_object_unsigneds(
value, id_start, size, id_width)
if len(np.unique(field_ids)) != size:
_malformed("duplicate object field id")
selected_slots = tuple(
(slot, targets[int(key_id)])
for slot, key_id in enumerate(field_ids)
if int(key_id) in targets)
slot_cache[layout] = selected_slots
if not targets:
result.append({})
continue
slots = np.fromiter(
(slot for slot, _ in selected_slots),
dtype=np.int64, count=len(selected_slots))
starts = offsets[slots]
if np.any(starts >= offsets[-1]):
_malformed("invalid object offset")
next_indices = np.searchsorted(
ordered_offsets, starts, side='right')
ends = ordered_offsets[next_indices]
# Encoded offsets may be uint8/uint16. Adding data_start before
# widening can wrap or raise for otherwise valid objects.
absolute_starts = starts.astype(np.intp, copy=False) + data_start
headers = np.frombuffer(value, dtype=np.uint8)[absolute_starts]
basic_types = headers & 0x3
type_infos = headers >> 2
fixed_sizes = _FAST_CHILD_SIZES[type_infos]
is_short_string = basic_types == _SHORT_STR
is_fast = (is_short_string
| ((basic_types == _PRIMITIVE) & (fixed_sizes != 0)))
expected_sizes = np.where(
is_short_string, type_infos + 1, fixed_sizes)
if np.any(is_fast & (ends - starts != expected_sizes)):
_malformed("child size does not match container offsets")
selected = {}
for index, ((_, name), offset, end) in enumerate(zip(
selected_slots, starts, ends)):
child_pos = data_start + int(offset)
if not is_fast[index]:
child_end = data_start + int(end)
if basic_types[index] in (_OBJECT, _ARRAY):
_validate_value_field_ids(
value, child_pos, child_end, len(key_ids))
elif _checked_value_size(
value, child_pos, child_end) != (
child_end - child_pos):
_malformed(
"child size does not match container offsets")
selected[name] = variant._to_python_impl(
value, row_metadata, child_pos)
result.append(selected)
return result
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")
@_with_metadata_cache
def variant_replace(
column,
path,
replacement=None,
strict: bool = False,
):
"""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):
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), _Replacement(value, len(column)))
for path, value in replacements.items()
]
_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)
result_chunks = []
global_row = 0
for chunk in chunks:
values = _BinaryValues(chunk.field(0))
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:
return result_chunks[0]
return pa.chunked_array(result_chunks, type=data_type)
class _ValueParts:
__slots__ = ('parts', 'size')
def __init__(self, parts, size):
_check_variant_sizes(size, 0)
self.parts = tuple(parts)
self.size = size
def __len__(self):
return self.size
def _part_size(part):
return part.size if isinstance(part, _ValueParts) else len(part)
def _materialize_value(part):
if not isinstance(part, _ValueParts):
_check_variant_sizes(len(part), 0)
return bytes(part)
output = bytearray(part.size)
output_pos = 0
stack = list(reversed(part.parts))
while stack:
current = stack.pop()
if isinstance(current, _ValueParts):
stack.extend(reversed(current.parts))
continue
size = len(current)
output[output_pos:output_pos + size] = current
output_pos += size
return bytes(output)
def _build_object_value_parts(fields):
"""Build an object without copying child values."""
size = len(fields)
data_size = sum(_part_size(child) for _, child in fields)
large_size = size > _U8_MAX
size_bytes = _U32_SIZE if large_size else 1
max_id = max((field_id for field_id, _ in fields), default=0)
id_size = _get_int_size(max_id) if max_id > 0 else 1
offset_size = _get_int_size(data_size) if data_size > 0 else 1
buf = bytearray()
buf.append(_object_header(large_size, id_size, offset_size))
buf += size.to_bytes(size_bytes, 'little')
for field_id, _ in fields:
buf += field_id.to_bytes(id_size, 'little')
offset = 0
for _, child in fields:
buf += offset.to_bytes(offset_size, 'little')
offset += _part_size(child)
buf += offset.to_bytes(offset_size, 'little')
header = bytes(buf)
return _ValueParts(
[header] + [child for _, child in fields],
len(header) + data_size,
)
def _build_array_value_parts(children):
"""Build an array without copying child values."""
size = len(children)
data_size = sum(_part_size(child) for child in children)
large_size = size > _U8_MAX
size_bytes = _U32_SIZE if large_size else 1
offset_size = _get_int_size(data_size) if data_size > 0 else 1
buf = bytearray()
buf.append(_array_header(large_size, offset_size))
buf += size.to_bytes(size_bytes, 'little')
offset = 0
for child in children:
buf += offset.to_bytes(offset_size, 'little')
offset += _part_size(child)
buf += offset.to_bytes(offset_size, 'little')
header = bytes(buf)
return _ValueParts([header] + children, len(header) + data_size)
def _build_object_value_ordered(fields):
"""Build object value bytes keeping the given field order."""
return _materialize_value(_build_object_value_parts(fields))
def _apply_edits(
value,
pos,
limit,
edits,
key_ids,
names_by_id,
source_metadata_size=None,
):
"""Apply edits and validate source ids before metadata extension."""
source = value if isinstance(value, memoryview) else memoryview(value)
results = {}
next_token = 1
stack = [('visit', 0, pos, limit, edits)]
while stack:
action = stack.pop()
kind = action[0]
if kind == 'finish_object':
_, token, ids, children, inserts, child_tokens = action
for slot, child_token in child_tokens:
children[slot] = results.pop(child_token)
fields = list(zip(ids, children)) + inserts
if len({field_id for field_id, _ in fields}) != len(fields):
_malformed("duplicate object field id")
if inserts:
try:
fields.sort(
key=lambda field: names_by_id[
field[0]].encode('utf-8'))
except KeyError:
_malformed("object key is missing from metadata")
results[token] = _build_object_value_parts(fields)
continue
if kind == 'finish_array':
_, token, children, child_tokens = action
for index, child_token in child_tokens:
children[index] = results.pop(child_token)
results[token] = _build_array_value_parts(children)
continue
_, token, node_pos, node_limit, node_edits = action
value_end = node_pos + _checked_value_size(
value, node_pos, node_limit)
if value_end != node_limit:
_malformed("child size does not match container offsets")
inserts = []
descend = {}
replacement = None
for segments, op, key_id, payload in node_edits:
if op == 'replace' and not segments:
replacement = payload
break
if op == 'insert' and len(segments) == 1:
inserts.append((key_id, payload))
else:
descend.setdefault(segments[0], []).append(
(segments[1:], op, key_id, payload))
if replacement is not None:
results[token] = replacement
continue
basic_type = value[node_pos] & 0x3
child_actions = []
if basic_type == _OBJECT:
size, id_size, id_start, data_start, offsets, _ = (
_checked_object_layout(value, node_pos, value_end))
ids = [
_read_unsigned(value, id_start + i * id_size, id_size)
for i in range(size)
]
if (source_metadata_size is not None
and any(field_id >= source_metadata_size
for field_id in ids)):
_malformed("object field id is missing from metadata")
ordered_offsets = sorted(offsets)
end_by_offset = dict(zip(
ordered_offsets, ordered_offsets[1:]))
slot_by_id = {
field_id: index for index, field_id in enumerate(ids)
}
edits_by_slot = {}
for (_, segment), child_edits in descend.items():
slot = slot_by_id[key_ids[segment]]
edits_by_slot[slot] = child_edits
children = []
child_tokens = []
for slot in range(size):
child_pos, child_end = _checked_object_child_bounds(
value, data_start, offsets, slot, end_by_offset)
child_edits = edits_by_slot.get(slot)
if child_edits is not None:
child_token = next_token
next_token += 1
children.append(None)
child_tokens.append((slot, child_token))
child_actions.append((
'visit', child_token, child_pos, child_end,
child_edits,
))
else:
if (source_metadata_size is not None
and (value[child_pos] & 0x3)
in (_OBJECT, _ARRAY)):
_validate_value_field_ids(
value, child_pos, child_end,
source_metadata_size)
children.append(source[child_pos:child_end])
stack.append((
'finish_object', token, ids, children, inserts,
child_tokens,
))
elif basic_type == _ARRAY:
size, data_start, offsets, _ = _checked_array_layout(
value, node_pos, value_end)
edits_by_index = {
segment: child_edits
for (_, segment), child_edits in descend.items()
}
children = []
child_tokens = []
for index in range(size):
child_pos = data_start + offsets[index]
child_end = data_start + offsets[index + 1]
child_edits = edits_by_index.get(index)
if child_edits is not None:
child_token = next_token
next_token += 1
children.append(None)
child_tokens.append((index, child_token))
child_actions.append((
'visit', child_token, child_pos, child_end,
child_edits,
))
else:
if source_metadata_size is not None:
_validate_value_field_ids(
value, child_pos, child_end,
source_metadata_size)
children.append(source[child_pos:child_end])
stack.append((
'finish_array', token, children, child_tokens,
))
else:
_malformed("path segment does not match the value type")
stack.extend(child_actions)
return _materialize_value(results[0])
def _root_insert_splice_layout(value, key_id, key_name, names_by_id):
"""Return a root layout that can splice the new field."""
size, id_size, id_start, data_start, first_offsets, _ = (
_checked_object_layout(value, 0, len(value)))
header = value[0]
type_info = (header >> 2) & 0x3F
large_size = ((type_info >> 4) & 0x1) != 0
size_width = _U32_SIZE if large_size else 1
offset_size = (type_info & 0x3) + 1
offset_start = id_start + size * id_size
if not large_size and size + 1 > _U8_MAX:
return None
if key_id >= 1 << (8 * id_size):
return None
ordered_offsets = sorted(first_offsets)
end_by_offset = dict(zip(ordered_offsets, ordered_offsets[1:]))
for index in range(size):
_checked_object_child_bounds(
value, data_start, first_offsets, index, end_by_offset)
ids = [
_read_unsigned(value, id_start + i * id_size, id_size)
for i in range(size)
]
names = [names_by_id.get(field_id) for field_id in ids]
if any(name is None for name in names) or names != sorted(names):
return None
slot = sum(name < key_name for name in names)
return (
header, size, size_width, id_size, id_start, data_start,
offset_size, offset_start, ids, slot,
)
def _encoded_payload_batches(rows, provider, global_row):
"""Encode array-backed splice payloads within byte and row budgets."""
batch_start = 0
payloads = []
payload_bytes = 0
for index, row in enumerate(rows):
payload = provider.encode(
provider.scalar_at(global_row + int(row)))
if (payloads
and (len(payloads)
>= _ROOT_INSERT_SPLICE_MAX_BATCH_ROWS
or payload_bytes + len(payload)
> _ROOT_INSERT_SPLICE_PAYLOAD_BUDGET)):
yield batch_start, index, payloads
batch_start = index
payloads = []
payload_bytes = 0
payloads.append(payload)
payload_bytes += len(payload)
if payloads:
yield batch_start, len(rows), payloads
def _repeated_payload_batches(row_count, payload):
"""Repeat a scalar payload without creating an unbounded row batch."""
for batch_start in range(
0, row_count, _ROOT_INSERT_SPLICE_MAX_BATCH_ROWS):
batch_end = min(
batch_start + _ROOT_INSERT_SPLICE_MAX_BATCH_ROWS,
row_count)
yield (
batch_start,
batch_end,
[payload] * (batch_end - batch_start),
)
def _plan_root_insert_splice(
values, rows, row_starts, row_lengths, source_data,
key_id, key_name, names_by_id, payloads):
"""Plan a splice and identify rows matching one root layout."""
layout = None
for row in rows:
layout = _root_insert_splice_layout(
values.view(int(row)), key_id, key_name, names_by_id)
if layout is not None:
break
if layout is None:
return None
(
header, size, size_width, id_size, id_start, data_start,
offset_size, offset_start, ids, slot,
) = layout
widths = np.full(len(rows), size_width, dtype=np.int64)
ok = source_data[row_starts] == header
ok &= row_lengths >= data_start
safe_starts = np.where(ok, row_starts, 0)
ok &= _take_unsigned(source_data, safe_starts + 1, widths) == size
widths = np.full(len(rows), id_size, dtype=np.int64)
for index in range(size):
ok &= _take_unsigned(
source_data,
safe_starts + id_start + index * id_size,
widths,
) == ids[index]
widths = np.full(len(rows), offset_size, dtype=np.int64)
sentinels = _take_unsigned(
source_data,
safe_starts + offset_start + size * offset_size,
widths,
)
ok &= data_start + sentinels == row_lengths
if size:
minimum = None
for index in range(size):
entry = _take_unsigned(
source_data,
safe_starts + offset_start + index * offset_size,
widths,
)
ok &= entry < sentinels
minimum = entry if minimum is None else np.minimum(
minimum, entry)
ok &= minimum == 0
payload_lengths = np.fromiter(
(len(payload) for payload in payloads), np.int64, len(payloads))
new_sentinels = sentinels + payload_lengths
ok &= new_sentinels < 1 << (8 * offset_size)
return (
header, size, size_width, id_size, id_start, data_start,
offset_size, offset_start, slot, sentinels, ok,
)
def _root_insert_splice(
values, state, rows, row_starts, row_lengths, source_data,
key_id, key_name, names_by_id, payloads,
source_metadata_size, output_metadata_size):
"""Splice one field into uniform root objects."""
plan = _plan_root_insert_splice(
values, rows, row_starts, row_lengths, source_data,
key_id, key_name, names_by_id, payloads)
if plan is None:
return None
(
header, size, size_width, id_size, id_start, data_start,
offset_size, offset_start, slot, sentinels, ok,
) = plan
matching_structures = np.zeros(len(rows), dtype=bool)
candidates = np.flatnonzero(ok)
if len(candidates):
lengths = row_lengths[candidates]
order = np.argsort(lengths, kind='stable')
candidates = candidates[order]
lengths = lengths[order]
boundaries = np.flatnonzero(lengths[1:] != lengths[:-1]) + 1
for group in np.split(candidates, boundaries):
exemplar = int(group[0])
value = values.view(int(rows[exemplar]))
if len(group) == 1:
_validate_value_field_ids(
value, 0, len(value), source_metadata_size)
matching_structures[exemplar] = True
continue
matches = _matching_value_structures(
value, source_data, row_starts[group],
source_metadata_size)
if matches is not None:
matching_structures[group[matches]] = True
if state.data is not None:
source_view = memoryview(state.data)
source_base = state.data_start
else:
source_view = values.data
source_base = 0
prefix = bytes([header]) + (size + 1).to_bytes(size_width, 'little')
id_bytes = key_id.to_bytes(id_size, 'little')
id_slot = id_start + slot * id_size
offset_slot = offset_start + slot * offset_size
sentinel_slot = offset_start + size * offset_size
rebuilt = {}
fallback_rows = []
for index, row in enumerate(rows):
row = int(row)
if not ok[index]:
fallback_rows.append(row)
continue
original = values.view(row)
if not matching_structures[index]:
_validate_value_field_ids(
original, 0, len(original), source_metadata_size)
base = int(row_starts[index]) - source_base
sentinel = int(sentinels[index])
_check_variant_sizes(
int(row_lengths[index]) + id_size + offset_size
+ len(payloads[index]),
output_metadata_size,
)
rebuilt[row] = b''.join((
prefix,
source_view[base + id_start:base + id_slot],
id_bytes,
source_view[base + id_slot:base + offset_slot],
sentinel.to_bytes(offset_size, 'little'),
source_view[base + offset_slot:base + sentinel_slot],
(sentinel + len(payloads[index])).to_bytes(
offset_size, 'little'),
source_view[base + data_start:base + data_start + sentinel],
payloads[index],
))
return rebuilt, fallback_rows
def _set_chunk(chunk, values, parsed, global_row):
parsed_paths = [parsed_path for _, parsed_path, _ in parsed]
parent_paths = [parsed_path[:-1] for parsed_path in parsed_paths]
query_paths = tuple(parsed_paths) + tuple(parent_paths)
count = len(parsed)
metadata_column = chunk.field(1)
valid_rows = _valid_row_indices(chunk, values, metadata_column)
if not len(valid_rows):
return chunk
plans, slow_rows = _partition_path_plans(
values, metadata_column, valid_rows, query_paths)
slow_rows = set(int(row) for row in slow_rows)
metadata_values = _BinaryValues(metadata_column)
state = _PatchState(values)
rebuilt_rows = {}
rebuilt_metadata = {}
scalar_payloads = {}
def payload_for(index, provider, row):
if provider._array is not None:
return provider.encode(provider.scalar_at(global_row + row))
if index not in scalar_payloads:
scalar_payloads[index] = provider.encode(provider.scalar_at(0))
return scalar_payloads[index]
def rebuild_row(row, view, insert_set, key_ids, names_by_id,
original_metadata, new_metadata,
source_metadata_size=None,
validated_positions=None):
edits = []
for index, (path, parsed_path, provider) in enumerate(parsed):
payload = payload_for(index, provider, row)
if index in insert_set:
edits.append((
parsed_path, 'insert',
key_ids[parsed_path[-1][1]], payload))
else:
if validated_positions is not None:
provider.validate_source(
view, validated_positions[index])
edits.append((parsed_path, 'replace', None, payload))
rebuilt = _apply_edits(
view, 0, len(view), edits, key_ids, names_by_id,
source_metadata_size)
_check_variant_sizes(
len(rebuilt),
len(new_metadata if new_metadata is not None
else original_metadata),
)
if new_metadata is not None or rebuilt != view:
rebuilt_rows[row] = rebuilt
if new_metadata is not None:
rebuilt_metadata[row] = new_metadata
for planned in plans:
rows, row_starts, source_data, positions, limits = planned
target_positions = positions[:count]
target_limits = limits[:count]
parent_positions = positions[count:]
parent_limits = limits[count:]
group_slow = set()
insert_indices = []
for index, (path, parsed_path, provider) in enumerate(parsed):
if target_positions[index] is not None:
continue
parent_pos = parent_positions[index]
if parent_pos is None:
raise ValueError(
f"VARIANT parent path does not exist: {path}")
if not parsed_path or parsed_path[-1][0] != 'key':
raise ValueError(
"VARIANT array index insertion is not supported: "
+ path)
parent_headers = source_data[row_starts + parent_pos]
if np.any((parent_headers & 0x3) != _OBJECT):
raise ValueError(
f"VARIANT parent path is not an object: {path}")
insert_indices.append(index)
if not insert_indices and all(
provider._fixed_size is not None
for _, _, provider in parsed):
_patch_planned_group(
(rows, row_starts, source_data,
target_positions, target_limits),
parsed, len(chunk), global_row, state, slow_rows, False)
continue
first_metadata = bytes(metadata_values.view(int(rows[0])))
insert_keys = tuple(
parsed[index][1][-1][1] for index in insert_indices)
metadata_key_ids = _cached_metadata_key_ids(first_metadata)
source_metadata_size = len(metadata_key_ids)
rebuild_validation_size = (
source_metadata_size
if any(key not in metadata_key_ids for key in insert_keys)
else None
)
new_metadata, key_ids, names_by_id = _metadata_with_keys(
first_metadata, insert_keys)
insert_set = set(insert_indices)
splice_eligible = (
len(insert_indices) == 1
and len(parsed[insert_indices[0]][1]) == 1
and all(
index in insert_set
or parsed[index][2]._fixed_size is not None
for index in range(count))
)
if splice_eligible:
insert_index = insert_indices[0]
replace_indices = [
index for index in range(count) if index != insert_index
]
if replace_indices:
_patch_planned_group(
(rows, row_starts, source_data,
[target_positions[index] for index in replace_indices],
[target_limits[index] for index in replace_indices]),
[parsed[index] for index in replace_indices],
len(chunk), global_row, state, group_slow, False)
slow_rows |= group_slow
if group_slow:
keep = np.fromiter(
(int(row) not in group_slow for row in rows),
bool, len(rows))
live_rows = rows[keep]
live_starts = row_starts[keep]
live_lengths = parent_limits[insert_index][keep]
else:
live_rows = rows
live_starts = row_starts
live_lengths = parent_limits[insert_index]
provider = parsed[insert_index][2]
if not len(live_rows):
continue
if provider._array is None:
payload = payload_for(insert_index, provider, 0)
batches = _repeated_payload_batches(
len(live_rows), payload)
else:
batches = _encoded_payload_batches(
live_rows, provider, global_row)
key_name = parsed[insert_index][1][-1][1]
output_metadata_size = len(
new_metadata if new_metadata is not None
else first_metadata)
for batch_start, batch_end, payloads in batches:
batch_rows = live_rows[batch_start:batch_end]
spliced = _root_insert_splice(
values, state, batch_rows,
live_starts[batch_start:batch_end],
live_lengths[batch_start:batch_end],
source_data, key_ids[key_name], key_name,
names_by_id, payloads, source_metadata_size,
output_metadata_size)
if spliced is None:
for row in batch_rows:
row = int(row)
rebuild_row(
row, values.view(row), insert_set, key_ids,
names_by_id, first_metadata, new_metadata,
rebuild_validation_size)
continue
rebuilt, fallback_rows = spliced
rebuilt_rows.update(rebuilt)
if new_metadata is not None:
for row in rebuilt:
rebuilt_metadata[row] = new_metadata
for row in fallback_rows:
rebuild_row(
row, values.view(row), insert_set, key_ids,
names_by_id, first_metadata, new_metadata,
rebuild_validation_size)
continue
for offset_index, row in enumerate(rows):
row = int(row)
if row in group_slow:
continue
rebuild_row(
row, values.view(row), insert_set, key_ids, names_by_id,
first_metadata, new_metadata, rebuild_validation_size,
[
None if target_positions[index] is None
else int(target_positions[index][offset_index])
for index in range(count)
],
)
for row in sorted(slow_rows):
view = values.view(row)
row_metadata = bytes(metadata_values.view(row))
positions = _path_positions(view, row_metadata, query_paths)
target_positions = positions[:count]
parent_positions = positions[count:]
insert_keys = []
insert_set = set()
for index, (path, parsed_path, provider) in enumerate(parsed):
if target_positions[index] is not None:
provider.validate_source(view, target_positions[index])
continue
parent_pos = parent_positions[index]
if parent_pos is None:
raise ValueError(
f"VARIANT parent path does not exist: {path}")
if not parsed_path or parsed_path[-1][0] != 'key':
raise ValueError(
"VARIANT array index insertion is not supported: "
+ path)
if (view[parent_pos] & 0x3) != _OBJECT:
raise ValueError(
f"VARIANT parent path is not an object: {path}")
insert_set.add(index)
insert_keys.append(parsed_path[-1][1])
metadata_key_ids = _cached_metadata_key_ids(row_metadata)
source_metadata_size = len(metadata_key_ids)
rebuild_validation_size = (
source_metadata_size
if any(key not in metadata_key_ids for key in insert_keys)
else None
)
new_metadata, key_ids, names_by_id = _metadata_with_keys(
row_metadata, tuple(insert_keys))
rebuild_row(
row, view, insert_set, key_ids, names_by_id, row_metadata,
new_metadata, rebuild_validation_size)
if rebuilt_rows:
state.ensure()
new_values = _sparse_rebuilt_binary(
values, state.data, state.data_start, rebuilt_rows)
elif state.data is not None:
return _patched_chunk(chunk, values, state.data, state.data_start)
else:
return chunk
if rebuilt_metadata:
new_metadata_column = _sparse_rebuilt_binary(
metadata_values, metadata_values.data, 0, rebuilt_metadata)
else:
new_metadata_column = metadata_column
return pa.StructArray.from_arrays(
[new_values, new_metadata_column],
fields=list(chunk.type),
mask=chunk.is_null(),
)
@_with_metadata_cache
def variant_set(column, path, value=None):
"""Replace existing VARIANT paths or insert missing final OBJECT keys.
Missing intermediate paths and non-OBJECT parents raise ValueError.
"""
if isinstance(path, Mapping):
if value is not None:
raise TypeError(
"VARIANT value must be omitted for path mappings")
updates = path
else:
updates = {path: value}
parsed = [
(target, _parse_path(target), _Replacement(item, len(column)))
for target, item in updates.items()
]
_validate_distinct_paths(parsed)
if not parsed:
return column
chunks, chunked, data_type = _variant_chunks(column)
result_chunks = []
global_row = 0
for chunk in chunks:
values = _BinaryValues(chunk.field(0))
result_chunks.append(_set_chunk(chunk, values, parsed, global_row))
global_row += len(chunk)
if not chunked:
return result_chunks[0]
return pa.chunked_array(result_chunks, type=data_type)