paimon-python/pypaimon/data/variant_shredding.py

# 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.

"""Shredded Parquet VARIANT support for pypaimon.

Paimon stores VARIANT columns in Parquet using a "shredded" format that enables
efficient sub-field reading and predicate pushdown.  This module provides:

  1. VariantSchema / build_variant_schema()
     Parse the PyArrow struct type of a shredded VARIANT column into a tree
     that mirrors the Java ``VariantSchema`` class.

  2. rebuild_value() / rebuild()
     Reconstruct standard ``(value: bytes, metadata: bytes)`` VARIANT binary
     from a shredded row dict.  Mirrors ``ShreddingUtils.rebuild()`` in Java.

  3. assemble_shredded_column()
     High-level helper used by ``FormatPyArrowReader`` to post-process batches
     that contain shredded VARIANT columns.

  4. is_shredded_variant()
     Detect shredded VARIANT columns in a file schema.

Shredded column layout (Parquet GROUP → PyArrow struct):
  - ``metadata``: binary — the top-level key dictionary
  - ``value``:    binary (optional) — overflow bytes for un-shredded fields
  - ``typed_value``: struct — per-field typed sub-columns (shredded fields)

Each field inside ``typed_value`` has the same ``{value, typed_value}``
structure recursively.  No ``field_`` prefix is used; sub-column names are
the exact variant key names.
"""

import datetime
import decimal as _decimal
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple

import pyarrow as pa

from pypaimon.data._variant_binary import (
    _OBJECT,
    _U8_MAX, _U32_SIZE,
    _VERSION, _VERSION_MASK,
    _read_unsigned, _get_int_size,
    _primitive_header, _object_header, _array_header,
)

# ---------------------------------------------------------------------------
# Local constants for null/true/false type IDs
# ---------------------------------------------------------------------------

_NULL_TYPE_ID = 0
_TRUE_TYPE_ID = 1
_FALSE_TYPE_ID = 2

_NULL_VALUE_BYTES: bytes = bytes([_primitive_header(_NULL_TYPE_ID)])


# ---------------------------------------------------------------------------
# Low-level binary helpers (shredding-specific)
# ---------------------------------------------------------------------------

def _append_le(buf: bytearray, value: int, n: int) -> None:
    buf.extend(value.to_bytes(n, 'little'))


# ---------------------------------------------------------------------------
# Metadata parsing
# ---------------------------------------------------------------------------

def parse_metadata_dict(metadata: bytes) -> Dict[str, int]:
    """Parse variant metadata bytes into a ``{key_name: key_id}`` mapping.

    The top-level metadata is shared across all shredded sub-fields.  We parse
    it once and pass the dict down to ``rebuild_value()`` so every recursive
    call can look up key IDs without re-parsing.
    """
    if not metadata or len(metadata) < 1:
        return {}
    if (metadata[0] & _VERSION_MASK) != _VERSION:
        raise ValueError('MALFORMED_VARIANT: invalid metadata version')
    offset_size = ((metadata[0] >> 6) & 0x3) + 1
    if len(metadata) < 1 + offset_size:
        return {}
    dict_size = _read_unsigned(metadata, 1, offset_size)
    if dict_size == 0:
        return {}
    string_start = 1 + (dict_size + 2) * offset_size
    result: Dict[str, int] = {}
    for key_id in range(dict_size):
        off = _read_unsigned(metadata, 1 + (key_id + 1) * offset_size, offset_size)
        next_off = _read_unsigned(metadata, 1 + (key_id + 2) * offset_size, offset_size)
        key = metadata[string_start + off:string_start + next_off].decode('utf-8')
        result[key] = key_id
    return result


# ---------------------------------------------------------------------------
# VariantSchema / build_variant_schema
# ---------------------------------------------------------------------------

@dataclass
class ObjectField:
    """One shredded field inside a ``typed_value`` object group."""
    field_name: str         # exact variant key name, no "field_" prefix
    schema: 'VariantSchema'


@dataclass
class VariantSchema:
    """Describes the shredding layout of a VARIANT column or sub-column.

    Mirrors the Java ``VariantSchema`` class.  Indices are positions within
    the PyArrow struct type that was parsed via ``build_variant_schema()``.

    For a plain (un-shredded) VARIANT:
      ``metadata_idx >= 0``, ``value_idx >= 0``, ``typed_idx < 0``.
    """
    typed_idx: int = -1
    value_idx: int = -1
    metadata_idx: int = -1
    num_fields: int = 0
    scalar_arrow_type: Optional[pa.DataType] = None
    object_fields: Optional[List[ObjectField]] = None
    object_schema_map: Optional[Dict[str, int]] = None
    array_schema: Optional['VariantSchema'] = None

    def is_unshredded(self) -> bool:
        """Return True if this is a plain (non-shredded) VARIANT layout."""
        return self.metadata_idx >= 0 and self.typed_idx < 0


def is_shredded_variant(pa_type: pa.DataType) -> bool:
    """Return True if *pa_type* is a shredded Parquet VARIANT struct.

    A shredded VARIANT column has three top-level fields: ``metadata``,
    ``value`` (overflow), and ``typed_value`` (shredded sub-columns).
    """
    if not pa.types.is_struct(pa_type):
        return False
    names = {pa_type[i].name for i in range(pa_type.num_fields)}
    return 'metadata' in names and 'value' in names and 'typed_value' in names


def build_variant_schema(pa_type: pa.StructType) -> VariantSchema:
    """Parse a PyArrow struct type into a ``VariantSchema`` tree.

    Works for both the top-level VARIANT column struct (which has ``metadata``
    in addition to ``value`` / ``typed_value``) and sub-field structs (which
    only have ``value`` / ``typed_value``).
    """
    schema = VariantSchema(num_fields=pa_type.num_fields)
    for i in range(pa_type.num_fields):
        f = pa_type[i]
        if f.name == 'metadata':
            schema.metadata_idx = i
        elif f.name == 'value':
            schema.value_idx = i
        elif f.name == 'typed_value':
            schema.typed_idx = i
            schema = _parse_typed_value_field(schema, f.type)
    return schema


def _parse_typed_value_field(schema: VariantSchema, tv_type: pa.DataType) -> VariantSchema:
    """Fill in the shredding details for a ``typed_value`` field."""
    if pa.types.is_struct(tv_type):
        object_fields: List[ObjectField] = []
        for j in range(tv_type.num_fields):
            sub_f = tv_type[j]
            if pa.types.is_struct(sub_f.type):
                sub_schema = build_variant_schema(sub_f.type)
            else:
                # Scalar typed_value embedded directly (no surrounding struct)
                sub_schema = VariantSchema(
                    typed_idx=0, num_fields=1, scalar_arrow_type=sub_f.type
                )
            object_fields.append(ObjectField(sub_f.name, sub_schema))
        schema.object_fields = object_fields
        schema.object_schema_map = {
            of.field_name: idx for idx, of in enumerate(object_fields)
        }
    elif pa.types.is_list(tv_type) or pa.types.is_large_list(tv_type):
        elem_type = tv_type.value_type
        if pa.types.is_struct(elem_type):
            schema.array_schema = build_variant_schema(elem_type)
        else:
            schema.array_schema = VariantSchema(
                typed_idx=0, num_fields=1, scalar_arrow_type=elem_type
            )
    else:
        schema.scalar_arrow_type = tv_type
    return schema


# ---------------------------------------------------------------------------
# Scalar encoding: Arrow typed value → variant binary bytes
# ---------------------------------------------------------------------------

def _encode_scalar_to_value_bytes(typed_value, arrow_type: pa.DataType) -> bytes:
    """Encode a typed Python scalar (from PyArrow .as_py()) to variant value bytes."""
    from pypaimon.data.generic_variant import _GenericVariantBuilder  # local import avoids circular
    builder = _GenericVariantBuilder()
    _append_scalar(builder, typed_value, arrow_type)
    gv = builder.result()
    # _pos == 0 so value() returns all bytes
    return gv.value()


def _append_scalar(builder, value, arrow_type: pa.DataType) -> None:
    """Dispatch a typed Python scalar into the appropriate builder method."""
    if pa.types.is_decimal(arrow_type) and arrow_type.scale < 0:
        raise ValueError('VARIANT decimal scale must be non-negative')
    if value is None:
        builder.append_null()
        return

    if pa.types.is_boolean(arrow_type):
        builder.append_boolean(bool(value))
    elif pa.types.is_integer(arrow_type):
        builder.append_long(int(value))
    elif pa.types.is_float64(arrow_type):
        builder.append_double(float(value))
    elif pa.types.is_float32(arrow_type):
        builder.append_float(float(value))
    elif pa.types.is_string(arrow_type) or pa.types.is_large_string(arrow_type):
        builder.append_string(str(value))
    elif pa.types.is_binary(arrow_type) or pa.types.is_large_binary(arrow_type):
        builder.append_binary(bytes(value))
    elif pa.types.is_date32(arrow_type):
        # PyArrow converts date32 to datetime.date
        if isinstance(value, datetime.date):
            days = (value - datetime.date(1970, 1, 1)).days
        else:
            days = int(value)
        builder.append_date(days)
    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
                )
                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
                )
                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)
    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}')
        builder.append_decimal_unscaled(unscaled, precision, scale)
    else:
        # Fallback: encode as string
        builder.append_string(str(value))


# ---------------------------------------------------------------------------
# Object / array binary construction
# ---------------------------------------------------------------------------

def _build_object_value(
    fields: List[Tuple[int, bytes]],
    key_dict: Optional[Dict[str, int]] = None,
) -> bytes:
    """Build object variant value bytes from ``(key_id, value_bytes)`` pairs.

    Variant object fields are ordered by key name, not metadata key ID. The
    metadata dictionary may assign IDs in a different order.
    """
    if not fields:
        # Empty object: header + size=0 + one zero-offset sentinel
        buf = bytearray()
        buf.append(_object_header(False, 1, 1))
        buf.append(0)       # size = 0
        buf.append(0)       # offset[0] = 0 (sentinel)
        return bytes(buf)

    if key_dict is None:
        fields = sorted(fields, key=lambda f: f[0])
    else:
        id_to_name = {key_id: name for name, key_id in key_dict.items()}
        fields = sorted(
            fields,
            key=lambda f: id_to_name[f[0]].encode('utf-8'),
        )
    size = len(fields)
    data = b''.join(vb for _, vb in fields)
    data_size = len(data)

    large_size = size > _U8_MAX
    size_bytes = _U32_SIZE if large_size else 1
    max_id = max(kid for kid, _ in fields)
    id_size = _get_int_size(max_id)
    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))
    _append_le(buf, size, size_bytes)

    for kid, _ in fields:
        _append_le(buf, kid, id_size)

    offset = 0
    for _, vb in fields:
        _append_le(buf, offset, offset_size)
        offset += len(vb)
    _append_le(buf, offset, offset_size)   # sentinel = total data size

    buf.extend(data)
    return bytes(buf)


def _build_array_value(element_bytes_list: List[bytes]) -> bytes:
    """Build array variant value bytes from per-element value bytes."""
    size = len(element_bytes_list)
    data = b''.join(element_bytes_list)
    data_size = len(data)

    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))
    _append_le(buf, size, size_bytes)

    offset = 0
    for eb in element_bytes_list:
        _append_le(buf, offset, offset_size)
        offset += len(eb)
    _append_le(buf, offset, offset_size)   # sentinel

    buf.extend(data)
    return bytes(buf)


def _extract_overflow_fields(overflow_bytes: bytes) -> List[Tuple[int, bytes]]:
    """Parse an overflow binary (a variant object) into ``(key_id, value_bytes)`` pairs.

    The overflow binary contains fields that were NOT shredded — they remain
    encoded as a compact variant object.
    """
    if not overflow_bytes:
        return []

    b = overflow_bytes[0]
    basic_type = b & 0x3
    if basic_type != _OBJECT:
        return []

    type_info = (b >> 2) & 0x3F
    large_size = bool((type_info >> 4) & 0x1)
    size_bytes = _U32_SIZE if large_size else 1
    size = _read_unsigned(overflow_bytes, 1, size_bytes)
    if size == 0:
        return []

    id_size = ((type_info >> 2) & 0x3) + 1
    offset_size = (type_info & 0x3) + 1
    id_start = 1 + size_bytes
    offset_start = id_start + size * id_size
    data_start = offset_start + (size + 1) * offset_size

    # The Parquet variant spec stores field offsets in the same order as the id_table,
    # but the DATA section may be laid out in a different order (e.g. GenericVariantBuilder
    # sorts the id/offset tables alphabetically while writing data in insertion order).
    # We must sort by offset to determine each field's byte boundaries correctly.
    #
    # Example: variant {"b": 2, "a": 1} written by GenericVariantBuilder
    #   id_table:     [id_a=0,  id_b=1]    (alphabetical order)
    #   offset_table: [off_a=2, off_b=0]   (id-order: a first, then b)
    #   data section: [enc(2),  enc(1)]    (insertion order: b first, then a)
    #   So pairs = [(id_a=0, off=2), (id_b=1, off=0)]
    #   Sorted by offset: [(id_b=1, off=0), (id_a=0, off=2)]
    #   Boundaries:        [0, 2, sentinel] → enc(2)=data[0:2], enc(1)=data[2:sentinel]
    pairs: List[Tuple[int, int]] = []
    for i in range(size):
        key_id = _read_unsigned(overflow_bytes, id_start + i * id_size, id_size)
        off = _read_unsigned(overflow_bytes, offset_start + i * offset_size, offset_size)
        pairs.append((key_id, off))

    sentinel = _read_unsigned(overflow_bytes, offset_start + size * offset_size, offset_size)

    # Sort by offset so that adjacent entries define contiguous data boundaries.
    # Track by original index (always unique) to avoid dict key collisions when
    # two fields share the same offset (malformed data).
    indexed_pairs = sorted(enumerate(pairs), key=lambda x: x[1][1])
    boundaries = [ip[1][1] for ip in indexed_pairs] + [sentinel]

    # end_by_orig[i] = end offset for pairs[i]
    end_by_orig = [0] * size
    for rank, (orig_idx, _) in enumerate(indexed_pairs):
        end_by_orig[orig_idx] = boundaries[rank + 1]

    fields: List[Tuple[int, bytes]] = []
    for orig_idx, (key_id, off) in enumerate(pairs):
        end = end_by_orig[orig_idx]
        field_bytes = bytes(overflow_bytes[data_start + off:data_start + end])
        fields.append((key_id, field_bytes))
    return fields


# ---------------------------------------------------------------------------
# Core rebuild algorithm (mirrors ShreddingUtils.rebuild() in Java)
# ---------------------------------------------------------------------------

def rebuild_value(
    row: dict,
    schema: VariantSchema,
    key_dict: Dict[str, int],
) -> Optional[bytes]:
    """Reconstruct variant value bytes from a shredded sub-row dict.

    Args:
        row:      Python dict from ``PyArrow StructScalar.as_py()``.  Keys are
                  ``'value'`` and/or ``'typed_value'``.
        schema:   VariantSchema for this level.
        key_dict: ``{key_name: key_id}`` parsed from the top-level metadata.

    Returns:
        Variant value bytes, or ``None`` if the field is absent (both
        ``typed_value`` and ``value`` are null — "missing from this row").
    """
    typed_value = row.get('typed_value') if schema.typed_idx >= 0 else None
    overflow = row.get('value') if schema.value_idx >= 0 else None

    # if both null → field is absent in this row
    if typed_value is None and overflow is None:
        return None

    # if typed_value is null → use overflow bytes directly
    if typed_value is None:
        return bytes(overflow)

    if schema.scalar_arrow_type is not None:
        return _encode_scalar_to_value_bytes(typed_value, schema.scalar_arrow_type)

    if schema.object_fields is not None:
        return _rebuild_object(typed_value, schema, key_dict, overflow)

    if schema.array_schema is not None:
        return _rebuild_array(typed_value, schema.array_schema, key_dict)

    # No sub-schema for typed_value; fall back to overflow
    return bytes(overflow) if overflow is not None else None


def _rebuild_object(
    typed_value: dict,
    schema: VariantSchema,
    key_dict: Dict[str, int],
    overflow_bytes: Optional[bytes],
) -> bytes:
    """Rebuild an object variant from shredded object sub-fields."""
    fields: List[Tuple[int, bytes]] = []

    for obj_field in schema.object_fields:
        fname = obj_field.field_name
        sub_row = typed_value.get(fname) if isinstance(typed_value, dict) else None
        if sub_row is None:
            continue

        field_value = rebuild_value(sub_row, obj_field.schema, key_dict)
        if field_value is None:
            continue

        key_id = key_dict.get(fname)
        if key_id is None:
            # Key not found in metadata — data integrity issue, skip
            continue
        fields.append((key_id, field_value))

    # Merge with overflow (fields not shredded)
    if overflow_bytes:
        fields.extend(_extract_overflow_fields(bytes(overflow_bytes)))

    return _build_object_value(fields, key_dict)


def _rebuild_array(
    typed_value,
    element_schema: VariantSchema,
    key_dict: Dict[str, int],
) -> bytes:
    """Rebuild an array variant from a shredded list."""
    if typed_value is None:
        typed_value = []

    element_bytes_list: List[bytes] = []
    for element_row in typed_value:
        if element_row is None:
            element_bytes_list.append(_NULL_VALUE_BYTES)
        elif isinstance(element_row, dict):
            eb = rebuild_value(element_row, element_schema, key_dict)
            element_bytes_list.append(eb if eb is not None else _NULL_VALUE_BYTES)
        else:
            # Scalar element directly (no surrounding {value, typed_value} struct)
            if element_schema.scalar_arrow_type is not None:
                eb = _encode_scalar_to_value_bytes(element_row, element_schema.scalar_arrow_type)
            else:
                eb = _NULL_VALUE_BYTES
            element_bytes_list.append(eb)

    return _build_array_value(element_bytes_list)


def rebuild(
    row: dict,
    schema: VariantSchema,
    key_dict: Optional[Dict[str, int]] = None,
) -> Tuple[bytes, bytes]:
    """Reconstruct ``(value_bytes, metadata_bytes)`` from a top-level shredded row.

    Args:
        row:      Top-level shredded row dict (contains ``'metadata'``,
                  ``'value'``, and/or ``'typed_value'`` keys).
        schema:   Top-level VariantSchema (from ``build_variant_schema()``).
        key_dict: Optional pre-parsed ``{key_name: key_id}`` mapping.  If None,
                  it is parsed from ``row['metadata']`` automatically.

    Returns:
        ``(value_bytes, metadata_bytes)`` suitable for constructing a standard
        VARIANT ``struct<value: binary, metadata: binary>`` row.
    """
    raw_metadata = row.get('metadata')
    if raw_metadata is None:
        raise ValueError("Shredded VARIANT row missing 'metadata' field")
    metadata = bytes(raw_metadata)

    if key_dict is None:
        key_dict = parse_metadata_dict(metadata)

    value_bytes = rebuild_value(row, schema, key_dict)
    if value_bytes is None:
        value_bytes = _NULL_VALUE_BYTES

    return value_bytes, metadata


# ---------------------------------------------------------------------------
# High-level column assembly
# ---------------------------------------------------------------------------

#: The canonical VARIANT Arrow type (struct<value: binary NOT NULL,
#: metadata: binary NOT NULL>).
VARIANT_ARROW_TYPE = pa.struct([
    pa.field('value', pa.binary(), nullable=False),
    pa.field('metadata', pa.binary(), nullable=False),
])


def assemble_shredded_column(column: pa.Array, schema: VariantSchema) -> pa.Array:
    """Convert a shredded VARIANT column to standard ``struct<value, metadata>``.

    Args:
        column: A PyArrow Array whose type is a shredded VARIANT struct.
        schema: VariantSchema built from the column's type.

    Returns:
        A ``pa.StructArray`` with type ``struct<value: binary, metadata: binary>``.
    """
    rows = column.to_pylist()
    assembled = []
    # cache parsed key dicts keyed by metadata bytes; most files share one metadata
    key_dict_cache: Dict[bytes, Dict[str, int]] = {}

    for row in rows:
        if row is None:
            assembled.append(None)
            continue

        raw_meta = row.get('metadata')
        if raw_meta is None:
            assembled.append(None)
            continue

        metadata = bytes(raw_meta)
        if metadata not in key_dict_cache:
            key_dict_cache[metadata] = parse_metadata_dict(metadata)

        value_bytes = rebuild_value(row, schema, key_dict_cache[metadata])
        if value_bytes is None:
            value_bytes = _NULL_VALUE_BYTES

        assembled.append({'value': value_bytes, 'metadata': metadata})

    return pa.array(assembled, type=VARIANT_ARROW_TYPE)


# ---------------------------------------------------------------------------
# Write-side shredding
# ---------------------------------------------------------------------------

def _paimon_type_str_to_arrow(type_str: str) -> pa.DataType:
    """Map a Paimon SQL type string to the Arrow type used in ``typed_value``."""
    from pypaimon.schema.data_types import AtomicType, PyarrowFieldParser
    try:
        return PyarrowFieldParser.from_paimon_type(AtomicType(type_str.upper()))
    except Exception:
        return pa.binary()


def _parse_sub_schema(type_def) -> VariantSchema:
    """Recursively parse a Paimon type definition (string or dict) into a VariantSchema.

    The resulting schema is for a sub-field struct, so ``value_idx=0``,
    ``typed_idx=1``, no ``metadata_idx``.
    """
    if isinstance(type_def, str):
        # Scalar leaf: e.g. "BIGINT", "VARCHAR", "DOUBLE"
        arrow_type = _paimon_type_str_to_arrow(type_def)
        return VariantSchema(
            value_idx=0,
            typed_idx=1,
            num_fields=2,
            scalar_arrow_type=arrow_type,
        )

    if not isinstance(type_def, dict):
        return VariantSchema(value_idx=0, num_fields=1)

    kind = type_def.get('type', '').upper()

    if kind == 'ROW':
        sub_fields_def = type_def.get('fields', [])
        object_fields: List[ObjectField] = []
        for f in sub_fields_def:
            fname = f.get('name', '')
            ftype = f.get('type', 'BINARY')
            sub_schema = _parse_sub_schema(ftype)
            object_fields.append(ObjectField(fname, sub_schema))
        schema = VariantSchema(
            value_idx=0,
            typed_idx=1,
            num_fields=2,
            object_fields=object_fields,
            object_schema_map={of.field_name: i for i, of in enumerate(object_fields)},
        )
        return schema

    if kind == 'ARRAY':
        elem_type = type_def.get('element', 'BINARY')
        elem_schema = _parse_sub_schema(elem_type)
        return VariantSchema(
            value_idx=0,
            typed_idx=1,
            num_fields=2,
            array_schema=elem_schema,
        )

    # Fallback: treat as scalar using the type string
    arrow_type = _paimon_type_str_to_arrow(kind)
    return VariantSchema(
        value_idx=0,
        typed_idx=1,
        num_fields=2,
        scalar_arrow_type=arrow_type,
    )


def parse_shredding_schema_option(json_str: str) -> Dict[str, List[ObjectField]]:
    """Parse the ``variant.shreddingSchema`` option value.

    Args:
        json_str: JSON-encoded Paimon ROW type where each top-level field
                  corresponds to a VARIANT column name, and its type is a ROW
                  listing the sub-fields to shred.

    Returns:
        ``{variant_col_name: [ObjectField, ...]}`` mapping.

    Raises:
        ValueError: if the JSON is invalid or the top-level type is not ROW.
    """
    import json as _json
    data = _json.loads(json_str)
    if data.get('type', '').upper() != 'ROW':
        raise ValueError(
            f"variant.shreddingSchema must be a JSON-encoded ROW type, got: {data.get('type')}"
        )

    result: Dict[str, List[ObjectField]] = {}
    for field_def in data.get('fields', []):
        col_name = field_def.get('name', '')
        col_type = field_def.get('type', {})

        if isinstance(col_type, dict) and col_type.get('type', '').upper() == 'ROW':
            sub_fields_def = col_type.get('fields', [])
            obj_fields: List[ObjectField] = []
            for sf in sub_fields_def:
                fname = sf.get('name', '')
                ftype = sf.get('type', 'BINARY')
                sub_schema = _parse_sub_schema(ftype)
                obj_fields.append(ObjectField(fname, sub_schema))
            result[col_name] = obj_fields
        else:
            # Top-level field type is not a ROW — skip (can't shred a non-object schema)
            pass

    return result


def _fid(field_id: int) -> dict:
    """Return PyArrow field metadata dict that sets the Parquet field ID.

    PyArrow respects the ``PARQUET:field_id`` key when writing Parquet, which
    ensures ``parquetType.getId()`` returns a non-null value on the Java reader
    side (``ParquetSchemaConverter.convertToPaimonField`` calls
    ``parquetType.getId().intValue()`` unconditionally).
    """
    return {b'PARQUET:field_id': str(field_id).encode()}


def _leaf_arrow_type_for_write(schema: VariantSchema) -> pa.DataType:
    """Return the Arrow type for the ``typed_value`` leaf of a sub-field struct.

    Used by ``sub_field_output_type``; field IDs are NOT embedded here since the
    result describes output column types, not Parquet-serialised fields.
    """
    if schema.scalar_arrow_type is not None:
        return schema.scalar_arrow_type
    if schema.object_fields is not None:
        return pa.struct([
            pa.field(of.field_name, pa.struct([
                pa.field('value', pa.binary(), nullable=True),
                pa.field('typed_value', _leaf_arrow_type_for_write(of.schema), nullable=True),
            ]), nullable=True)
            for of in schema.object_fields
        ])
    return pa.binary()


def _leaf_arrow_type_for_write_with_ids(schema: VariantSchema) -> pa.DataType:
    """Like ``_leaf_arrow_type_for_write`` but embeds ``PARQUET:field_id`` metadata.

    Used by ``shredding_schema_to_arrow_type`` so that Parquet field IDs are
    present in every nested field of the written file.
    """
    if schema.scalar_arrow_type is not None:
        return schema.scalar_arrow_type
    if schema.object_fields is not None:
        inner_fields = []
        for i, of in enumerate(schema.object_fields):
            inner_sub = pa.struct([
                pa.field('value', pa.binary(), nullable=True, metadata=_fid(0)),
                pa.field(
                    'typed_value',
                    _leaf_arrow_type_for_write_with_ids(of.schema),
                    nullable=True,
                    metadata=_fid(1),
                ),
            ])
            inner_fields.append(
                pa.field(of.field_name, inner_sub, nullable=False, metadata=_fid(i + 1))
            )
        return pa.struct(inner_fields)
    return pa.binary()


def shredding_schema_to_arrow_type(obj_fields: List[ObjectField]) -> pa.StructType:
    """Convert an ``[ObjectField]`` list into the PyArrow struct type for a shredded column.

    The produced type is the canonical Parquet shredding layout::

        struct<
            metadata: binary NOT NULL,   (field_id=0)
            value:    binary,             (field_id=1)
            typed_value: struct<          (field_id=2)
                field_a: struct<          (field_id=1, NOT NULL)
                    value: binary,        (field_id=0)
                    typed_value: <type_a> (field_id=1)
                >,
                ...
            >
        >

    ``PARQUET:field_id`` metadata is embedded on every field so that the Java
    ``ParquetSchemaConverter.convertToPaimonField`` can call
    ``parquetType.getId().intValue()`` without a NullPointerException.

    Sub-field structs within ``typed_value`` are marked NOT NULL (``nullable=False``)
    to match the Java shredding schema where each named sub-field carries a
    ``.notNull()`` annotation.
    """
    sub_field_defs = []
    for i, of in enumerate(obj_fields):
        typed_val_type = _leaf_arrow_type_for_write_with_ids(of.schema)
        sub_struct = pa.struct([
            pa.field('value', pa.binary(), nullable=True, metadata=_fid(0)),
            pa.field('typed_value', typed_val_type, nullable=True, metadata=_fid(1)),
        ])
        # Java's variantShreddingSchema marks each named sub-field as .notNull()
        sub_field_defs.append(pa.field(of.field_name, sub_struct, nullable=False, metadata=_fid(i + 1)))

    return pa.struct([
        pa.field('metadata', pa.binary(), nullable=False, metadata=_fid(0)),
        pa.field('value', pa.binary(), nullable=True, metadata=_fid(1)),
        pa.field('typed_value', pa.struct(sub_field_defs), nullable=True, metadata=_fid(2)),
    ])


def _decompose_field_bytes(
    field_bytes: bytes,
    schema: VariantSchema,
    metadata: bytes,
) -> dict:
    """Decompose a variant field's value bytes into a ``{value, typed_value}`` sub-struct dict.

    This is the write-direction counterpart of ``rebuild_value()``.

    Args:
        field_bytes: Variant value bytes for a single field extracted from the
                     original variant binary.
        schema:      VariantSchema describing how this field should be shredded.
        metadata:    Top-level metadata bytes (used for key ID lookups in nested objects).

    Returns:
        A Python dict ``{'value': bytes_or_none, 'typed_value': val_or_dict_or_none}``
        suitable for building a PyArrow struct array row.
    """
    from pypaimon.data.generic_variant import GenericVariant

    if schema.scalar_arrow_type is not None:
        try:
            py_val = GenericVariant(field_bytes, metadata).to_python()
        except Exception:
            return {'value': field_bytes, 'typed_value': None}
        return {'value': None, 'typed_value': py_val}

    if schema.object_fields is not None:
        if not field_bytes or (field_bytes[0] & 0x3) != _OBJECT:
            return {'value': field_bytes, 'typed_value': None}

        all_sub = _extract_overflow_fields(field_bytes)
        key_dict = parse_metadata_dict(metadata)
        id_to_name = {v: k for k, v in key_dict.items()}
        shredded_names = {of.field_name for of in schema.object_fields}

        overflow_pairs = [
            (kid, fb) for kid, fb in all_sub
            if id_to_name.get(kid) not in shredded_names
        ]
        shredded_by_name = {
            id_to_name[kid]: fb for kid, fb in all_sub
            if kid in id_to_name and id_to_name[kid] in shredded_names
        }

        typed_value: Dict[str, dict] = {}
        for of in schema.object_fields:
            fname = of.field_name
            if fname in shredded_by_name:
                typed_value[fname] = _decompose_field_bytes(
                    shredded_by_name[fname], of.schema, metadata
                )
            else:
                typed_value[fname] = {'value': None, 'typed_value': None}

        overflow_bytes = (_build_object_value(overflow_pairs, key_dict)
                          if overflow_pairs else None)
        return {'value': overflow_bytes, 'typed_value': typed_value}

    # No shredding sub-schema: treat field bytes as overflow
    return {'value': field_bytes, 'typed_value': None}


def decompose_variant(
    gv: 'GenericVariant',
    obj_fields: List[ObjectField],
) -> dict:
    """Decompose a ``GenericVariant`` into a shredded row dict for writing.

    This is the inverse of ``rebuild()`` / ``rebuild_value()`` — it takes a
    fully-encoded variant and splits it into the shredded ``{metadata, value,
    typed_value}`` structure that Parquet shredding expects.

    Args:
        gv:         The ``GenericVariant`` to decompose.
        obj_fields: List of ``ObjectField`` from ``parse_shredding_schema_option()``,
                    describing which top-level object keys to shred.

    Returns:
        A Python dict matching the Arrow type produced by
        ``shredding_schema_to_arrow_type(obj_fields)``.
    """
    metadata = gv.metadata()
    value_bytes = gv.value()

    # Non-object variants cannot be shredded: put everything in overflow and set
    # typed_value to NULL so the Java reader falls through to the overflow path.
    if not value_bytes or (value_bytes[0] & 0x3) != _OBJECT:
        return {'metadata': metadata, 'value': value_bytes, 'typed_value': None}

    all_fields = _extract_overflow_fields(value_bytes)
    key_dict = parse_metadata_dict(metadata)
    id_to_name = {v: k for k, v in key_dict.items()}
    shredded_names = {of.field_name for of in obj_fields}

    overflow_pairs = [
        (kid, fb) for kid, fb in all_fields
        if id_to_name.get(kid) not in shredded_names
    ]
    shredded_by_name = {
        id_to_name[kid]: fb for kid, fb in all_fields
        if kid in id_to_name and id_to_name[kid] in shredded_names
    }

    typed_value = {}
    for of in obj_fields:
        fname = of.field_name
        if fname in shredded_by_name:
            typed_value[fname] = _decompose_field_bytes(
                shredded_by_name[fname], of.schema, metadata
            )
        else:
            typed_value[fname] = {'value': None, 'typed_value': None}

    overflow_bytes = (_build_object_value(overflow_pairs, key_dict)
                      if overflow_pairs else None)
    return {'metadata': metadata, 'value': overflow_bytes, 'typed_value': typed_value}


def shred_variant_column(
    column: pa.Array,
    obj_fields: List[ObjectField],
    target_type: pa.StructType,
) -> pa.Array:
    """Convert a standard VARIANT column to its shredded representation.

    Args:
        column:      A ``pa.Array`` of type ``struct<value: binary, metadata: binary>``
                     (the standard Paimon VARIANT layout).
        obj_fields:  ``[ObjectField, ...]`` from ``parse_shredding_schema_option()``.
        target_type: The Arrow struct type from ``shredding_schema_to_arrow_type(obj_fields)``.

    Returns:
        A ``pa.StructArray`` with ``target_type`` suitable for writing to Parquet
        in the shredded format.
    """
    from pypaimon.data.generic_variant import GenericVariant

    rows = column.to_pylist()
    result = []
    for row in rows:
        if row is None:
            result.append(None)
        else:
            gv = GenericVariant.from_arrow_struct(row)
            result.append(decompose_variant(gv, obj_fields))
    return pa.array(result, type=target_type)