| # 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. | |
| """Unit tests for VARIANT type support in pypaimon. | |
| Covers the user-facing features described in the Python API documentation: | |
| - Plain VARIANT: ordinary write (GenericVariant → Parquet) and read. | |
| - Shredded VARIANT: write-side decomposition (shred_variant_column) and | |
| read-side transparent assembly (assemble_shredded_column). | |
| Sections | |
| -------- | |
| 1. Type-system layer – schema parsing, Paimon ↔ Arrow type mapping. | |
| 2. GenericVariant container – construction, inspection, JSON round-trip. | |
| 3. Plain VARIANT I/O – PyArrow-level Parquet write/read sanity check. | |
| 4. Shredded schema utils – is_shredded_variant, parse_metadata_dict, | |
| build_variant_schema. | |
| 5. Binary encoding helpers – _encode_scalar_to_value_bytes, _build_object_value, | |
| _build_array_value. | |
| 6. Read-path assembly – rebuild_value, rebuild, assemble_shredded_column. | |
| 7. Write-path shredding – parse_shredding_schema_option, | |
| shredding_schema_to_arrow_type, decompose_variant. | |
| 8. Full-chain via Paimon API – write then read through CatalogFactory / Schema / | |
| write_arrow / to_arrow for both plain VARIANT and | |
| shredded VARIANT. | |
| """ | |
| import asyncio | |
| import io | |
| import json | |
| import os | |
| import shutil | |
| import struct as _struct | |
| import tempfile | |
| import unittest | |
| import datetime | |
| import uuid | |
| from decimal import Decimal | |
| import pyarrow as pa | |
| import pyarrow.parquet as pq | |
| from pypaimon import CatalogFactory, Schema | |
| from pypaimon.data.generic_variant import GenericVariant | |
| from pypaimon.data.variant_shredding import ( | |
| VARIANT_ARROW_TYPE, | |
| VariantSchema, | |
| _NULL_VALUE_BYTES, | |
| _build_array_value, | |
| _build_object_value, | |
| _extract_overflow_fields, | |
| _encode_scalar_to_value_bytes, | |
| assemble_shredded_column, | |
| build_variant_schema, | |
| decompose_variant, | |
| is_shredded_variant, | |
| parse_metadata_dict, | |
| parse_shredding_schema_option, | |
| rebuild, | |
| rebuild_value, | |
| shredding_schema_to_arrow_type, | |
| ) | |
| from pypaimon.read.native_plan import native_split_bridge_available | |
| from pypaimon.schema.data_types import ( | |
| AtomicType, | |
| DataField, | |
| DataTypeParser, | |
| PyarrowFieldParser, | |
| RowType, | |
| is_variant_struct, | |
| ) | |
| from pypaimon.table.row.generic_row import ( | |
| GenericRow, | |
| GenericRowDeserializer, | |
| GenericRowSerializer, | |
| ) | |
| # --------------------------------------------------------------------------- | |
| # Shared helpers | |
| # --------------------------------------------------------------------------- | |
| def _variant_arrow_type() -> pa.StructType: | |
| """The canonical Arrow representation of a plain VARIANT column.""" | |
| return pa.struct([ | |
| pa.field('value', pa.binary(), nullable=False), | |
| pa.field('metadata', pa.binary(), nullable=False), | |
| ]) | |
| def _native_variant_projection_available() -> bool: | |
| if not native_split_bridge_available(): | |
| return False | |
| try: | |
| from pypaimon_rust.datafusion import ReadBuilder as NativeReadBuilder | |
| return callable(getattr(NativeReadBuilder, 'with_read_type', None)) | |
| except (ImportError, AttributeError, TypeError, ValueError): | |
| return False | |
| def _make_metadata(*keys: str) -> bytes: | |
| """Build variant metadata bytes. With no keys returns the minimal header.""" | |
| if keys: | |
| return GenericVariant.from_python({k: None for k in keys}).metadata() | |
| return b'\x01\x00' # version=1, dict_size=0 | |
| def _scalar_sub_struct(arrow_type: pa.DataType) -> pa.StructType: | |
| """Return struct<value: binary, typed_value: arrow_type> for one scalar sub-field.""" | |
| return pa.struct([ | |
| pa.field('value', pa.binary(), nullable=True), | |
| pa.field('typed_value', arrow_type, nullable=True), | |
| ]) | |
| def _object_shredded_type(*key_type_pairs) -> pa.StructType: | |
| """Build a top-level shredded VARIANT struct with the given (key, arrow_type) sub-fields.""" | |
| tv_fields = [ | |
| pa.field(name, _scalar_sub_struct(atype), nullable=True) | |
| for name, atype in key_type_pairs | |
| ] | |
| return pa.struct([ | |
| pa.field('metadata', pa.binary(), nullable=False), | |
| pa.field('value', pa.binary(), nullable=True), | |
| pa.field('typed_value', pa.struct(tv_fields), nullable=True), | |
| ]) | |
| def _schema_json(col_name: str, sub_fields) -> str: | |
| """Build a variant.shreddingSchema JSON string for the given column and sub-fields. | |
| Args: | |
| col_name: VARIANT column name. | |
| sub_fields: list of (field_name, paimon_type_str) pairs, e.g. | |
| [('age', 'BIGINT'), ('city', 'VARCHAR')]. | |
| """ | |
| return json.dumps({ | |
| 'type': 'ROW', | |
| 'fields': [ | |
| { | |
| 'id': 0, | |
| 'name': col_name, | |
| 'type': { | |
| 'type': 'ROW', | |
| 'fields': [ | |
| {'id': i, 'name': name, 'type': dtype} | |
| for i, (name, dtype) in enumerate(sub_fields) | |
| ], | |
| }, | |
| } | |
| ], | |
| }) | |
| # =========================================================================== | |
| # 1. Type-system layer | |
| # =========================================================================== | |
| class TestVariantSchemaParsing(unittest.TestCase): | |
| def test_parse_variant_keyword(self): | |
| dt = DataTypeParser.parse_atomic_type_sql_string('VARIANT') | |
| self.assertIsInstance(dt, AtomicType) | |
| self.assertEqual(dt.type, 'VARIANT') | |
| self.assertTrue(dt.nullable) | |
| def test_parse_variant_not_null(self): | |
| dt = DataTypeParser.parse_atomic_type_sql_string('VARIANT NOT NULL') | |
| self.assertIsInstance(dt, AtomicType) | |
| self.assertFalse(dt.nullable) | |
| def test_variant_to_dict_roundtrip(self): | |
| dt = AtomicType('VARIANT') | |
| restored = DataTypeParser.parse_data_type(dt.to_dict()) | |
| self.assertEqual(dt, restored) | |
| def test_variant_str(self): | |
| self.assertEqual(str(AtomicType('VARIANT')), 'VARIANT') | |
| self.assertEqual(str(AtomicType('VARIANT', nullable=False)), 'VARIANT NOT NULL') | |
| class TestVariantFromPaimonType(unittest.TestCase): | |
| def _arrow_type(self): | |
| return PyarrowFieldParser.from_paimon_type(AtomicType('VARIANT')) | |
| def test_returns_struct(self): | |
| self.assertTrue(pa.types.is_struct(self._arrow_type())) | |
| self.assertEqual(self._arrow_type().num_fields, 2) | |
| def test_field_names(self): | |
| t = self._arrow_type() | |
| self.assertEqual(t.field(0).name, 'value') | |
| self.assertEqual(t.field(1).name, 'metadata') | |
| def test_field_types_are_binary(self): | |
| t = self._arrow_type() | |
| self.assertTrue(pa.types.is_binary(t.field(0).type)) | |
| self.assertTrue(pa.types.is_binary(t.field(1).type)) | |
| def test_fields_not_nullable(self): | |
| t = self._arrow_type() | |
| self.assertFalse(t.field(0).nullable) | |
| self.assertFalse(t.field(1).nullable) | |
| def test_from_paimon_field(self): | |
| df = DataField(id=0, name='payload', type=AtomicType('VARIANT')) | |
| pa_field = PyarrowFieldParser.from_paimon_field(df) | |
| self.assertEqual(pa_field.name, 'payload') | |
| self.assertTrue(pa.types.is_struct(pa_field.type)) | |
| self.assertTrue(pa_field.nullable) | |
| def test_from_paimon_schema(self): | |
| fields = [ | |
| DataField(id=0, name='id', type=AtomicType('BIGINT')), | |
| DataField(id=1, name='payload', type=AtomicType('VARIANT')), | |
| ] | |
| schema = PyarrowFieldParser.from_paimon_schema(fields) | |
| self.assertEqual(schema.field('payload').type, _variant_arrow_type()) | |
| class TestVariantToPaimonType(unittest.TestCase): | |
| def test_is_variant_struct_positive(self): | |
| self.assertTrue(is_variant_struct(_variant_arrow_type())) | |
| def test_is_variant_struct_wrong_names(self): | |
| st = pa.struct([ | |
| pa.field('val', pa.binary(), nullable=False), | |
| pa.field('meta', pa.binary(), nullable=False), | |
| ]) | |
| self.assertFalse(is_variant_struct(st)) | |
| def test_is_variant_struct_nullable_fields(self): | |
| st = pa.struct([ | |
| pa.field('value', pa.binary(), nullable=True), | |
| pa.field('metadata', pa.binary(), nullable=False), | |
| ]) | |
| self.assertFalse(is_variant_struct(st)) | |
| def test_is_variant_struct_wrong_types(self): | |
| st = pa.struct([ | |
| pa.field('value', pa.string(), nullable=False), | |
| pa.field('metadata', pa.binary(), nullable=False), | |
| ]) | |
| self.assertFalse(is_variant_struct(st)) | |
| def test_is_variant_struct_extra_fields(self): | |
| st = pa.struct([ | |
| pa.field('value', pa.binary(), nullable=False), | |
| pa.field('metadata', pa.binary(), nullable=False), | |
| pa.field('typed_value', pa.int64(), nullable=True), | |
| ]) | |
| self.assertFalse(is_variant_struct(st)) | |
| def test_to_paimon_type_variant(self): | |
| result = PyarrowFieldParser.to_paimon_type(_variant_arrow_type(), nullable=True) | |
| self.assertIsInstance(result, AtomicType) | |
| self.assertEqual(result.type, 'VARIANT') | |
| self.assertTrue(result.nullable) | |
| def test_to_paimon_type_variant_not_null(self): | |
| result = PyarrowFieldParser.to_paimon_type(_variant_arrow_type(), nullable=False) | |
| self.assertFalse(result.nullable) | |
| def test_ordinary_struct_maps_to_row_type(self): | |
| st = pa.struct([pa.field('a', pa.int32()), pa.field('b', pa.string())]) | |
| result = PyarrowFieldParser.to_paimon_type(st, nullable=True) | |
| self.assertIsInstance(result, RowType) | |
| def test_struct_same_names_different_types_is_row_type(self): | |
| st = pa.struct([ | |
| pa.field('value', pa.string(), nullable=False), | |
| pa.field('metadata', pa.string(), nullable=False), | |
| ]) | |
| result = PyarrowFieldParser.to_paimon_type(st, nullable=True) | |
| self.assertIsInstance(result, RowType) | |
| class TestVariantSchemaRoundTrip(unittest.TestCase): | |
| def test_paimon_to_arrow_to_paimon(self): | |
| original = DataField(id=0, name='v', type=AtomicType('VARIANT')) | |
| pa_field = PyarrowFieldParser.from_paimon_field(original) | |
| restored = PyarrowFieldParser.to_paimon_type(pa_field.type, pa_field.nullable) | |
| self.assertIsInstance(restored, AtomicType) | |
| self.assertEqual(restored.type, 'VARIANT') | |
| def test_mixed_schema_round_trip(self): | |
| fields = [ | |
| DataField(id=0, name='id', type=AtomicType('BIGINT')), | |
| DataField(id=1, name='payload', type=AtomicType('VARIANT')), | |
| DataField(id=2, name='ts', type=AtomicType('TIMESTAMP(6)')), | |
| ] | |
| pa_schema = PyarrowFieldParser.from_paimon_schema(fields) | |
| restored = PyarrowFieldParser.to_paimon_schema(pa_schema) | |
| self.assertEqual(restored[1].name, 'payload') | |
| self.assertIsInstance(restored[1].type, AtomicType) | |
| self.assertEqual(restored[1].type.type, 'VARIANT') | |
| self.assertEqual(restored[2].name, 'ts') | |
| # =========================================================================== | |
| # 2. GenericVariant container | |
| # =========================================================================== | |
| class TestGenericVariantContainer(unittest.TestCase): | |
| def test_from_python_returns_instance(self): | |
| gv = GenericVariant.from_python({'a': 1, 'b': 'hello'}) | |
| self.assertIsInstance(gv, GenericVariant) | |
| self.assertIsInstance(gv.value(), bytes) | |
| def test_from_arrow_struct_roundtrip(self): | |
| original = GenericVariant.from_python({'x': 1, 'y': 2}) | |
| restored = GenericVariant.from_arrow_struct( | |
| {'value': original.value(), 'metadata': original.metadata()}) | |
| self.assertEqual(restored.value(), original.value()) | |
| self.assertEqual(restored.metadata(), original.metadata()) | |
| def test_to_python_object(self): | |
| gv = GenericVariant.from_python({'age': 30, 'city': 'Beijing'}) | |
| result = gv.to_python() | |
| self.assertEqual(result, {'age': 30, 'city': 'Beijing'}) | |
| def test_to_python_array(self): | |
| gv = GenericVariant.from_python([1, 2, 3]) | |
| self.assertEqual(gv.to_python(), [1, 2, 3]) | |
| def test_to_python_null(self): | |
| gv = GenericVariant.from_python(None) | |
| self.assertIsNone(gv.to_python()) | |
| def test_to_python_string(self): | |
| gv = GenericVariant.from_python('hello') | |
| self.assertEqual(gv.to_python(), 'hello') | |
| def test_to_python_number(self): | |
| gv = GenericVariant.from_python(42) | |
| self.assertEqual(gv.to_python(), 42) | |
| def test_from_python_none(self): | |
| gv = GenericVariant.from_python(None) | |
| self.assertIsNone(gv.to_python()) | |
| def test_from_python_nested(self): | |
| obj = {'user': {'name': 'Alice', 'scores': [10, 20, 30]}, 'active': True} | |
| gv = GenericVariant.from_python(obj) | |
| result = gv.to_python() | |
| self.assertEqual(result['user']['name'], 'Alice') | |
| self.assertEqual(result['user']['scores'], [10, 20, 30]) | |
| self.assertTrue(result['active']) | |
| def test_equality(self): | |
| gv1 = GenericVariant.from_python({'a': 1}) | |
| gv2 = GenericVariant.from_python({'a': 1}) | |
| self.assertEqual(gv1, gv2) | |
| def test_repr_and_str(self): | |
| gv = GenericVariant.from_python('hello') | |
| self.assertIn('hello', repr(gv)) | |
| self.assertIn('hello', str(gv)) | |
| def test_from_python_date(self): | |
| value = datetime.date(2024, 1, 15) | |
| gv = GenericVariant.from_python(value) | |
| self.assertEqual(gv.to_python(), value) | |
| def test_from_python_timestamp_ntz(self): | |
| cases = [ | |
| datetime.datetime(2024, 1, 15, 12, 30, 45, 123456), | |
| datetime.datetime(1600, 6, 15, 12, 30, 45, 123456), | |
| datetime.datetime(1900, 1, 1, 0, 0, 0, 654321), | |
| datetime.datetime(2500, 12, 31, 23, 59, 59, 111111), | |
| datetime.datetime(5000, 3, 3, 3, 3, 3, 222222), | |
| datetime.datetime(9998, 7, 8, 12, 34, 56, 654321), | |
| ] | |
| for value in cases: | |
| gv = GenericVariant.from_python(value) | |
| self.assertEqual(gv.to_python(), value) | |
| def test_from_python_timestamp_ltz(self): | |
| cases = [ | |
| datetime.datetime(2024, 1, 15, 12, 30, 45, 123456, tzinfo=datetime.timezone.utc), | |
| datetime.datetime( | |
| 1600, 6, 15, 12, 30, 45, 123456, tzinfo=datetime.timezone(datetime.timedelta(hours=8))), | |
| datetime.datetime( | |
| 9998, 7, 8, 12, 34, 56, 654321, tzinfo=datetime.timezone(datetime.timedelta(hours=8))), | |
| ] | |
| for value in cases: | |
| gv = GenericVariant.from_python(value) | |
| self.assertEqual(gv.to_python(), value) | |
| def test_from_python_nested_datetime(self): | |
| obj = {'created_at': datetime.datetime(2024, 1, 15, 12, 0)} | |
| gv = GenericVariant.from_python(obj) | |
| self.assertEqual(gv.to_python(), obj) | |
| def test_from_python_uuid(self): | |
| value = uuid.UUID('12345678-1234-5678-1234-567812345678') | |
| gv = GenericVariant.from_python(value) | |
| self.assertEqual(gv.to_python(), value) | |
| def test_from_python_nested_uuid(self): | |
| obj = {'id': uuid.UUID('12345678-1234-5678-1234-567812345678')} | |
| gv = GenericVariant.from_python(obj) | |
| self.assertEqual(gv.to_python(), obj) | |
| class TestToArrowArray(unittest.TestCase): | |
| def test_basic(self): | |
| gv1 = GenericVariant.from_python({'a': 1}) | |
| gv2 = GenericVariant.from_python([1, 2]) | |
| arr = GenericVariant.to_arrow_array([gv1, gv2]) | |
| self.assertIsInstance(arr, pa.StructArray) | |
| self.assertEqual(len(arr), 2) | |
| restored = GenericVariant.from_arrow_struct(arr[0].as_py()) | |
| self.assertEqual(restored.to_python(), {'a': 1}) | |
| def test_with_nulls(self): | |
| arr = GenericVariant.to_arrow_array([GenericVariant.from_python(42), None]) | |
| self.assertEqual(len(arr), 2) | |
| self.assertIsNotNone(arr[0].as_py()) | |
| self.assertIsNone(arr[1].as_py()) | |
| def test_empty(self): | |
| self.assertEqual(len(GenericVariant.to_arrow_array([])), 0) | |
| def test_arrow_type(self): | |
| gv = GenericVariant.from_python(True) | |
| arr = GenericVariant.to_arrow_array([gv]) | |
| self.assertEqual(arr.type, _variant_arrow_type()) | |
| class TestJsonRoundtrip(unittest.TestCase): | |
| def _check(self, json_str): | |
| gv = GenericVariant.from_python(json.loads(json_str)) | |
| self.assertEqual(gv.to_python(), json.loads(json_str)) | |
| def test_nested_object_array(self): | |
| self._check('{"users":[{"name":"Alice","age":30},{"name":"Bob","age":25}]}') | |
| def test_deep_nesting(self): | |
| self._check('{"a":{"b":{"c":{"d":42}}}}') | |
| def test_array_of_objects(self): | |
| self._check('[{"x":1},{"x":2},{"x":3}]') | |
| def test_all_primitive_types(self): | |
| self._check('{"n":null,"b":true,"i":42,"s":"hello","f":1.5}') | |
| def test_empty_object(self): | |
| gv = GenericVariant.from_python({}) | |
| self.assertEqual(gv.to_python(), {}) | |
| def test_empty_array(self): | |
| gv = GenericVariant.from_python([]) | |
| self.assertEqual(gv.to_python(), []) | |
| # =========================================================================== | |
| # 3. Plain VARIANT — PyArrow-level Parquet sanity check | |
| # (verifies the physical struct<value, metadata> layout works in Parquet) | |
| # =========================================================================== | |
| class TestVariantBinaryRow(unittest.TestCase): | |
| def test_round_trip_non_utf8_payload(self): | |
| field = DataField(0, 'v', AtomicType('VARIANT')) | |
| variant = GenericVariant.from_python(-1) | |
| serialized = GenericRowSerializer.to_bytes( | |
| GenericRow([variant], [field])) | |
| result = GenericRowDeserializer.from_bytes( | |
| serialized, [field]).values[0] | |
| self.assertEqual(variant.value(), result.value()) | |
| self.assertEqual(variant.metadata(), result.metadata()) | |
| self.assertEqual(-1, result.to_python()) | |
| def _make_variant_bytes(json_str: str) -> bytes: | |
| """Produce a minimal VARIANT value payload encoding a long-string primitive.""" | |
| payload = json_str.encode('utf-8') | |
| return _struct.pack('<B', 0x15) + _struct.pack('<I', len(payload)) + payload | |
| class TestVariantParquetCycle(unittest.TestCase): | |
| def test_write_and_read_parquet_values(self): | |
| """Write struct<value, metadata> bytes to Parquet and verify schema + values.""" | |
| meta = _make_metadata() | |
| payload_col = pa.array( | |
| [{'value': _make_variant_bytes('{"key":"hello"}'), 'metadata': meta}, | |
| {'value': _make_variant_bytes('42'), 'metadata': meta}], | |
| type=_variant_arrow_type(), | |
| ) | |
| original = pa.table({'id': [1, 2], 'payload': payload_col}) | |
| buf = io.BytesIO() | |
| pq.write_table(original, buf) | |
| buf.seek(0) | |
| restored = pq.read_table(buf) | |
| self.assertEqual(restored.schema.field('payload').type, _variant_arrow_type()) | |
| self.assertEqual(restored.num_rows, 2) | |
| row0 = restored.column('payload')[0].as_py() | |
| self.assertEqual(row0['value'], _make_variant_bytes('{"key":"hello"}')) | |
| self.assertEqual(row0['metadata'], meta) | |
| def test_null_variant_row(self): | |
| """SQL-NULL VARIANT row survives Parquet round-trip as None.""" | |
| payload_col = pa.array( | |
| [None, {'value': _make_variant_bytes('true'), 'metadata': _make_metadata()}], | |
| type=_variant_arrow_type(), | |
| ) | |
| buf = io.BytesIO() | |
| pq.write_table(pa.table({'id': [1, 2], 'payload': payload_col}), buf) | |
| buf.seek(0) | |
| restored = pq.read_table(buf) | |
| self.assertIsNone(restored.column('payload')[0].as_py()) | |
| self.assertIsNotNone(restored.column('payload')[1].as_py()) | |
| def test_to_arrow_array_roundtrip_parquet(self): | |
| """GenericVariant.to_arrow_array() produces bytes that survive a Parquet cycle.""" | |
| gvs = [ | |
| GenericVariant.from_python({'score': 1}), | |
| GenericVariant.from_python({'score': 2}), | |
| None, | |
| ] | |
| col = GenericVariant.to_arrow_array(gvs) | |
| buf = io.BytesIO() | |
| pq.write_table(pa.table({'v': col}), buf) | |
| buf.seek(0) | |
| restored = pq.read_table(buf).column('v') | |
| self.assertEqual( | |
| GenericVariant.from_arrow_struct(restored[0].as_py()).to_python(), {'score': 1}) | |
| self.assertEqual( | |
| GenericVariant.from_arrow_struct(restored[1].as_py()).to_python(), {'score': 2}) | |
| self.assertIsNone(restored[2].as_py()) | |
| # =========================================================================== | |
| # 4. Shredded VARIANT — schema detection and parsing | |
| # =========================================================================== | |
| class TestIsShredded(unittest.TestCase): | |
| def test_shredded_is_detected(self): | |
| t = _object_shredded_type(('age', pa.int64())) | |
| self.assertTrue(is_shredded_variant(t)) | |
| def test_plain_variant_not_shredded(self): | |
| plain = pa.struct([ | |
| pa.field('value', pa.binary(), nullable=False), | |
| pa.field('metadata', pa.binary(), nullable=False), | |
| ]) | |
| self.assertFalse(is_shredded_variant(plain)) | |
| def test_arbitrary_struct_not_shredded(self): | |
| t = pa.struct([pa.field('a', pa.int32()), pa.field('b', pa.string())]) | |
| self.assertFalse(is_shredded_variant(t)) | |
| def test_non_struct_not_shredded(self): | |
| self.assertFalse(is_shredded_variant(pa.binary())) | |
| self.assertFalse(is_shredded_variant(pa.int64())) | |
| def test_missing_typed_value_not_shredded(self): | |
| t = pa.struct([ | |
| pa.field('metadata', pa.binary()), | |
| pa.field('value', pa.binary()), | |
| ]) | |
| self.assertFalse(is_shredded_variant(t)) | |
| class TestParseMetadataDict(unittest.TestCase): | |
| def test_single_key(self): | |
| meta = _make_metadata('age') | |
| d = parse_metadata_dict(meta) | |
| self.assertIn('age', d) | |
| self.assertEqual(d['age'], 0) | |
| def test_multiple_keys(self): | |
| meta = _make_metadata('age', 'name', 'city') | |
| d = parse_metadata_dict(meta) | |
| self.assertEqual(set(d.keys()), {'age', 'name', 'city'}) | |
| self.assertEqual(d['age'], 0) | |
| self.assertEqual(d['name'], 1) | |
| self.assertEqual(d['city'], 2) | |
| def test_empty_metadata(self): | |
| meta = b'\x01\x00' | |
| d = parse_metadata_dict(meta) | |
| self.assertEqual(d, {}) | |
| def test_empty_bytes(self): | |
| d = parse_metadata_dict(b'') | |
| self.assertEqual(d, {}) | |
| class TestBuildVariantSchema(unittest.TestCase): | |
| def test_simple_scalar(self): | |
| t = _object_shredded_type(('age', pa.int64()), ('name', pa.string())) | |
| schema = build_variant_schema(t) | |
| self.assertEqual(schema.metadata_idx, 0) | |
| self.assertEqual(schema.value_idx, 1) | |
| self.assertEqual(schema.typed_idx, 2) | |
| self.assertIsNotNone(schema.object_fields) | |
| self.assertEqual(len(schema.object_fields), 2) | |
| self.assertEqual(schema.object_fields[0].field_name, 'age') | |
| self.assertEqual(schema.object_fields[1].field_name, 'name') | |
| def test_sub_field_scalar_type(self): | |
| t = _object_shredded_type(('age', pa.int64())) | |
| schema = build_variant_schema(t) | |
| age_schema = schema.object_fields[0].schema | |
| self.assertEqual(age_schema.scalar_arrow_type, pa.int64()) | |
| self.assertEqual(age_schema.value_idx, 0) | |
| self.assertEqual(age_schema.typed_idx, 1) | |
| def test_object_schema_map(self): | |
| t = _object_shredded_type(('age', pa.int64()), ('name', pa.string())) | |
| schema = build_variant_schema(t) | |
| self.assertIn('age', schema.object_schema_map) | |
| self.assertIn('name', schema.object_schema_map) | |
| def test_is_unshredded_for_plain_variant(self): | |
| schema = VariantSchema(metadata_idx=0, value_idx=1) | |
| self.assertTrue(schema.is_unshredded()) | |
| def test_not_unshredded_for_shredded_schema(self): | |
| t = _object_shredded_type(('age', pa.int64())) | |
| schema = build_variant_schema(t) | |
| self.assertFalse(schema.is_unshredded()) | |
| def test_no_typed_value_field(self): | |
| plain = pa.struct([ | |
| pa.field('metadata', pa.binary(), nullable=False), | |
| pa.field('value', pa.binary(), nullable=True), | |
| ]) | |
| schema = build_variant_schema(plain) | |
| self.assertEqual(schema.metadata_idx, 0) | |
| self.assertEqual(schema.value_idx, 1) | |
| self.assertEqual(schema.typed_idx, -1) | |
| self.assertTrue(schema.is_unshredded()) | |
| # =========================================================================== | |
| # 5. Binary encoding helpers | |
| # =========================================================================== | |
| class TestEncodeScalar(unittest.TestCase): | |
| def _roundtrip(self, json_str: str, arrow_type: pa.DataType): | |
| """Encode a scalar to bytes via the Arrow type, then decode via GenericVariant.""" | |
| gv_orig = GenericVariant.from_python(json.loads(json_str)) | |
| typed_value = pa.array([gv_orig.to_python()], type=arrow_type).to_pylist()[0] | |
| value_bytes = _encode_scalar_to_value_bytes(typed_value, arrow_type) | |
| gv = GenericVariant(value_bytes, b'\x01\x00') | |
| return gv.to_python() | |
| def test_tinyint(self): | |
| self.assertEqual(self._roundtrip('42', pa.int8()), 42) | |
| def test_smallint(self): | |
| self.assertEqual(self._roundtrip('1000', pa.int16()), 1000) | |
| def test_int(self): | |
| self.assertEqual(self._roundtrip('1000000', pa.int32()), 1000000) | |
| def test_bigint(self): | |
| self.assertEqual(self._roundtrip('12345678901234', pa.int64()), 12345678901234) | |
| def test_float(self): | |
| self.assertAlmostEqual(self._roundtrip('3.14159', pa.float32()), 3.14159, places=5) | |
| def test_double(self): | |
| self.assertEqual(self._roundtrip('1.012345678901234567890123456789012345678', | |
| pa.float64()), 1.012345678901234567890123456789012345678) | |
| def test_decimal(self): | |
| value = Decimal('12345.6789') | |
| value_bytes = _encode_scalar_to_value_bytes(value, pa.decimal128(10, 4)) | |
| gv = GenericVariant(value_bytes, b'\x01\x00') | |
| self.assertEqual(gv.to_python(), value) | |
| def test_date(self): | |
| value = datetime.date(2024, 1, 15) | |
| value_bytes = _encode_scalar_to_value_bytes(value, pa.date32()) | |
| gv = GenericVariant(value_bytes, b'\x01\x00') | |
| self.assertEqual(gv.to_python(), value) | |
| def test_timestamp_ntz(self): | |
| value = datetime.datetime(2024, 1, 15, 12, 30, 45, 123456) | |
| value_bytes = _encode_scalar_to_value_bytes(value, pa.timestamp('us')) | |
| gv = GenericVariant(value_bytes, b'\x01\x00') | |
| self.assertEqual(gv.to_python(), value) | |
| def test_timestamp_ltz(self): | |
| value = datetime.datetime( | |
| 2024, 1, 15, 12, 30, 45, 123456, tzinfo=datetime.timezone.utc | |
| ) | |
| value_bytes = _encode_scalar_to_value_bytes(value, pa.timestamp('us', tz='UTC')) | |
| gv = GenericVariant(value_bytes, b'\x01\x00') | |
| self.assertEqual(gv.to_python(), value) | |
| def test_bool_true(self): | |
| self.assertEqual(self._roundtrip('true', pa.bool_()), True) | |
| def test_bool_false(self): | |
| self.assertEqual(self._roundtrip('false', pa.bool_()), False) | |
| def test_string(self): | |
| self.assertEqual(self._roundtrip('"hello"', pa.string()), 'hello') | |
| def test_binary(self): | |
| value_bytes = _encode_scalar_to_value_bytes(b'\x00\x01\x02', pa.binary()) | |
| gv = GenericVariant(value_bytes, b'\x01\x00') | |
| self.assertEqual(gv.to_python(), b'\x00\x01\x02') | |
| def test_null(self): | |
| value_bytes = _encode_scalar_to_value_bytes(None, pa.int64()) | |
| gv = GenericVariant(value_bytes, b'\x01\x00') | |
| self.assertIsNone(gv.to_python()) | |
| class TestBuildBinary(unittest.TestCase): | |
| def test_build_object_empty(self): | |
| obj_bytes = _build_object_value([]) | |
| gv = GenericVariant(obj_bytes, b'\x01\x00') | |
| self.assertEqual(gv.to_python(), {}) | |
| def test_build_object_one_field(self): | |
| meta = _make_metadata('age') | |
| key_dict = parse_metadata_dict(meta) | |
| age_val = _encode_scalar_to_value_bytes(30, pa.int64()) | |
| obj_bytes = _build_object_value([(key_dict['age'], age_val)]) | |
| gv = GenericVariant(obj_bytes, meta) | |
| self.assertEqual(gv.to_python(), {'age': 30}) | |
| def test_build_object_orders_keys_not_metadata_ids(self): | |
| metadata = _make_metadata('z', 'a') | |
| key_dict = parse_metadata_dict(metadata) | |
| scalar = _encode_scalar_to_value_bytes(1, pa.int64()) | |
| value = _build_object_value( | |
| [(key_dict['z'], scalar), (key_dict['a'], scalar)], key_dict) | |
| self.assertEqual( | |
| [key_id for key_id, _ in _extract_overflow_fields(value)], | |
| [key_dict['a'], key_dict['z']], | |
| ) | |
| def test_build_array_empty(self): | |
| arr_bytes = _build_array_value([]) | |
| gv = GenericVariant(arr_bytes, b'\x01\x00') | |
| self.assertEqual(gv.to_python(), []) | |
| def test_build_array_three_ints(self): | |
| elem1 = _encode_scalar_to_value_bytes(1, pa.int64()) | |
| elem2 = _encode_scalar_to_value_bytes(2, pa.int64()) | |
| elem3 = _encode_scalar_to_value_bytes(3, pa.int64()) | |
| arr_bytes = _build_array_value([elem1, elem2, elem3]) | |
| gv = GenericVariant(arr_bytes, b'\x01\x00') | |
| self.assertEqual(gv.to_python(), [1, 2, 3]) | |
| def test_build_array_with_null(self): | |
| elem1 = _encode_scalar_to_value_bytes(1, pa.int64()) | |
| arr_bytes = _build_array_value([elem1, _NULL_VALUE_BYTES]) | |
| gv = GenericVariant(arr_bytes, b'\x01\x00') | |
| result = gv.to_python() | |
| self.assertEqual(result[0], 1) | |
| self.assertIsNone(result[1]) | |
| # =========================================================================== | |
| # 6. Read-path assembly | |
| # =========================================================================== | |
| class TestRebuildValue(unittest.TestCase): | |
| """Test the core rebuild_value() function for various shredding scenarios.""" | |
| def _age_schema(self) -> VariantSchema: | |
| """Schema for a sub-field struct {value: binary, typed_value: int64}.""" | |
| return build_variant_schema(_scalar_sub_struct(pa.int64())) | |
| def _metadata_and_dict(self, *keys: str): | |
| meta = _make_metadata(*keys) | |
| return meta, parse_metadata_dict(meta) | |
| def test_scalar_typed_value(self): | |
| schema = self._age_schema() | |
| row = {'value': None, 'typed_value': 30} | |
| result = rebuild_value(row, schema, {}) | |
| gv = GenericVariant(result, b'\x01\x00') | |
| self.assertEqual(gv.to_python(), 30) | |
| def test_scalar_overflow(self): | |
| """When typed_value is None, fall back to overflow value bytes.""" | |
| age_val = _encode_scalar_to_value_bytes(42, pa.int64()) | |
| schema = self._age_schema() | |
| row = {'value': age_val, 'typed_value': None} | |
| result = rebuild_value(row, schema, {}) | |
| gv = GenericVariant(result, b'\x01\x00') | |
| self.assertEqual(gv.to_python(), 42) | |
| def test_both_null_returns_none(self): | |
| """Both value and typed_value being None signals an absent field.""" | |
| schema = self._age_schema() | |
| row = {'value': None, 'typed_value': None} | |
| self.assertIsNone(rebuild_value(row, schema, {})) | |
| def test_object_with_shredded_fields(self): | |
| meta, key_dict = self._metadata_and_dict('age', 'name') | |
| t = _object_shredded_type(('age', pa.int64()), ('name', pa.string())) | |
| schema = build_variant_schema(t) | |
| row = { | |
| 'metadata': meta, | |
| 'value': None, | |
| 'typed_value': { | |
| 'age': {'value': None, 'typed_value': 30}, | |
| 'name': {'value': None, 'typed_value': 'Alice'}, | |
| }, | |
| } | |
| value_bytes, _ = rebuild(row, schema, key_dict) | |
| gv = GenericVariant(value_bytes, meta) | |
| result = gv.to_python() | |
| self.assertEqual(result['age'], 30) | |
| self.assertEqual(result['name'], 'Alice') | |
| def test_object_absent_field_skipped(self): | |
| """A sub-field with both value=None and typed_value=None is omitted from output.""" | |
| meta, key_dict = self._metadata_and_dict('age', 'name') | |
| t = _object_shredded_type(('age', pa.int64()), ('name', pa.string())) | |
| schema = build_variant_schema(t) | |
| row = { | |
| 'metadata': meta, | |
| 'value': None, | |
| 'typed_value': { | |
| 'age': {'value': None, 'typed_value': 25}, | |
| 'name': {'value': None, 'typed_value': None}, # absent | |
| }, | |
| } | |
| value_bytes, _ = rebuild(row, schema, key_dict) | |
| gv = GenericVariant(value_bytes, meta) | |
| result = gv.to_python() | |
| self.assertEqual(result['age'], 25) | |
| self.assertNotIn('name', result) | |
| def test_object_with_overflow(self): | |
| """Fields not in typed_value are preserved from overflow bytes.""" | |
| original = GenericVariant.from_python({'age': 30, 'extra': 'overflow_val'}) | |
| overflow_bytes = original.value() | |
| meta = original.metadata() | |
| key_dict = parse_metadata_dict(meta) | |
| t = _object_shredded_type(('age', pa.int64())) | |
| schema = build_variant_schema(t) | |
| row = { | |
| 'metadata': meta, | |
| 'value': overflow_bytes, | |
| 'typed_value': { | |
| 'age': {'value': None, 'typed_value': 30}, | |
| }, | |
| } | |
| value_bytes = rebuild_value(row, schema, key_dict) | |
| gv = GenericVariant(value_bytes, meta) | |
| result = gv.to_python() | |
| self.assertEqual(result['age'], 30) | |
| self.assertIn('extra', result) | |
| def test_typed_value_null_uses_overflow(self): | |
| """If typed_value struct is None for the whole row, full overflow bytes are used.""" | |
| original = GenericVariant.from_python({'age': 99}) | |
| meta = original.metadata() | |
| key_dict = parse_metadata_dict(meta) | |
| t = _object_shredded_type(('age', pa.int64())) | |
| schema = build_variant_schema(t) | |
| row = { | |
| 'metadata': meta, | |
| 'value': original.value(), | |
| 'typed_value': None, # typed_value absent for this row | |
| } | |
| value_bytes, _ = rebuild(row, schema, key_dict) | |
| gv = GenericVariant(value_bytes, meta) | |
| self.assertEqual(gv.to_python()['age'], 99) | |
| def test_rebuild_matches_direct_variant(self): | |
| """Bytes rebuilt from shredded form must equal bytes from GenericVariant.from_python.""" | |
| original_json = '{"score": 42, "tag": "test"}' | |
| original_gv = GenericVariant.from_python(json.loads(original_json)) | |
| meta = original_gv.metadata() | |
| key_dict = parse_metadata_dict(meta) | |
| t = _object_shredded_type(('score', pa.int64()), ('tag', pa.string())) | |
| schema = build_variant_schema(t) | |
| row = { | |
| 'metadata': meta, | |
| 'value': None, | |
| 'typed_value': { | |
| 'score': {'value': None, 'typed_value': 42}, | |
| 'tag': {'value': None, 'typed_value': 'test'}, | |
| }, | |
| } | |
| value_bytes, _ = rebuild(row, schema, key_dict) | |
| gv_rebuilt = GenericVariant(value_bytes, meta) | |
| self.assertEqual(gv_rebuilt.to_python(), original_gv.to_python()) | |
| class TestRebuild(unittest.TestCase): | |
| def test_missing_metadata_raises(self): | |
| schema = VariantSchema() | |
| with self.assertRaises(ValueError): | |
| rebuild({'value': None, 'typed_value': None}, schema) | |
| def test_basic_roundtrip(self): | |
| meta = _make_metadata('x') | |
| key_dict = parse_metadata_dict(meta) | |
| t = _object_shredded_type(('x', pa.int64())) | |
| schema = build_variant_schema(t) | |
| row = { | |
| 'metadata': meta, | |
| 'value': None, | |
| 'typed_value': {'x': {'value': None, 'typed_value': 7}}, | |
| } | |
| value_bytes, ret_meta = rebuild(row, schema, key_dict) | |
| self.assertEqual(ret_meta, meta) | |
| gv = GenericVariant(value_bytes, ret_meta) | |
| self.assertEqual(gv.to_python(), {'x': 7}) | |
| def test_auto_parse_key_dict(self): | |
| """When key_dict is None, rebuild() parses it automatically from the metadata field.""" | |
| meta = _make_metadata('y') | |
| t = _object_shredded_type(('y', pa.string())) | |
| schema = build_variant_schema(t) | |
| row = { | |
| 'metadata': meta, | |
| 'value': None, | |
| 'typed_value': {'y': {'value': None, 'typed_value': 'hi'}}, | |
| } | |
| value_bytes, _ = rebuild(row, schema) | |
| gv = GenericVariant(value_bytes, meta) | |
| self.assertEqual(gv.to_python(), {'y': 'hi'}) | |
| class TestAssembleShreddedColumn(unittest.TestCase): | |
| def test_basic_two_rows(self): | |
| """assemble_shredded_column converts a shredded Arrow column to plain VARIANT.""" | |
| meta = _make_metadata('age', 'name') | |
| t = _object_shredded_type(('age', pa.int64()), ('name', pa.string())) | |
| rows = [ | |
| {'metadata': meta, 'value': None, 'typed_value': { | |
| 'age': {'value': None, 'typed_value': 30}, | |
| 'name': {'value': None, 'typed_value': 'Alice'}, | |
| }}, | |
| {'metadata': meta, 'value': None, 'typed_value': { | |
| 'age': {'value': None, 'typed_value': 25}, | |
| 'name': {'value': None, 'typed_value': 'Bob'}, | |
| }}, | |
| ] | |
| col = pa.array(rows, type=t) | |
| schema = build_variant_schema(col.type) | |
| result = assemble_shredded_column(col, schema) | |
| self.assertEqual(result.type, VARIANT_ARROW_TYPE) | |
| self.assertEqual(len(result), 2) | |
| gv0 = GenericVariant.from_arrow_struct(result[0].as_py()) | |
| self.assertEqual(gv0.to_python(), {'age': 30, 'name': 'Alice'}) | |
| gv1 = GenericVariant.from_arrow_struct(result[1].as_py()) | |
| self.assertEqual(gv1.to_python(), {'age': 25, 'name': 'Bob'}) | |
| # =========================================================================== | |
| # 7. Write-path shredding | |
| # =========================================================================== | |
| class TestShreddingWrite(unittest.TestCase): | |
| """Tests for the static-mode write shredding API.""" | |
| def _obj_fields_for(self, col_name, sub_fields): | |
| """Helper: build obj_fields for a single column from (name, type) pairs.""" | |
| return parse_shredding_schema_option(_schema_json(col_name, sub_fields))[col_name] | |
| # ----------------------------------------------------------------------- | |
| # parse_shredding_schema_option | |
| # ----------------------------------------------------------------------- | |
| def test_parse_single_scalar_field(self): | |
| result = parse_shredding_schema_option(json.dumps({ | |
| 'type': 'ROW', | |
| 'fields': [ | |
| {'id': 0, 'name': 'payload', 'type': { | |
| 'type': 'ROW', | |
| 'fields': [{'id': 0, 'name': 'age', 'type': 'INT'}], | |
| }}, | |
| ], | |
| })) | |
| self.assertIn('payload', result) | |
| obj_fields = result['payload'] | |
| self.assertEqual(len(obj_fields), 1) | |
| self.assertEqual(obj_fields[0].field_name, 'age') | |
| self.assertIsNotNone(obj_fields[0].schema.scalar_arrow_type) | |
| def test_parse_multiple_fields_multiple_cols(self): | |
| result = parse_shredding_schema_option(json.dumps({ | |
| 'type': 'ROW', | |
| 'fields': [ | |
| {'id': 0, 'name': 'col_a', 'type': {'type': 'ROW', 'fields': [ | |
| {'id': 0, 'name': 'x', 'type': 'BIGINT'}, | |
| {'id': 1, 'name': 'y', 'type': 'VARCHAR'}, | |
| ]}}, | |
| {'id': 1, 'name': 'col_b', 'type': {'type': 'ROW', 'fields': [ | |
| {'id': 0, 'name': 'flag', 'type': 'BOOLEAN'}, | |
| ]}}, | |
| ], | |
| })) | |
| self.assertIn('col_a', result) | |
| self.assertIn('col_b', result) | |
| self.assertEqual(len(result['col_a']), 2) | |
| self.assertEqual(result['col_a'][0].field_name, 'x') | |
| self.assertEqual(result['col_a'][1].field_name, 'y') | |
| self.assertEqual(result['col_b'][0].field_name, 'flag') | |
| def test_parse_invalid_top_level_type_raises(self): | |
| with self.assertRaises(ValueError): | |
| parse_shredding_schema_option(json.dumps({'type': 'ARRAY', 'element': 'INT'})) | |
| # ----------------------------------------------------------------------- | |
| # shredding_schema_to_arrow_type | |
| # ----------------------------------------------------------------------- | |
| def test_arrow_type_structure_single_scalar(self): | |
| obj_fields = self._obj_fields_for('col', [('score', 'BIGINT')]) | |
| arrow_type = shredding_schema_to_arrow_type(obj_fields) | |
| self.assertTrue(pa.types.is_struct(arrow_type)) | |
| names = {arrow_type.field(i).name for i in range(arrow_type.num_fields)} | |
| self.assertIn('metadata', names) | |
| self.assertIn('value', names) | |
| self.assertIn('typed_value', names) | |
| tv = arrow_type.field('typed_value').type | |
| self.assertTrue(pa.types.is_struct(tv)) | |
| self.assertEqual(tv.num_fields, 1) | |
| self.assertEqual(tv.field(0).name, 'score') | |
| sub = tv.field('score').type | |
| self.assertTrue(pa.types.is_struct(sub)) | |
| sub_names = {sub.field(i).name for i in range(sub.num_fields)} | |
| self.assertIn('value', sub_names) | |
| self.assertIn('typed_value', sub_names) | |
| # ----------------------------------------------------------------------- | |
| # decompose_variant | |
| # ----------------------------------------------------------------------- | |
| def test_decompose_scalar_field_extracted(self): | |
| """A shredded scalar field ends up in typed_value; non-shredded fields go to overflow.""" | |
| obj_fields = self._obj_fields_for('col', [('age', 'BIGINT')]) | |
| gv = GenericVariant.from_python({'age': 25, 'name': 'alice'}) | |
| result = decompose_variant(gv, obj_fields) | |
| self.assertIn('metadata', result) | |
| self.assertIn('value', result) | |
| self.assertIn('typed_value', result) | |
| tv = result['typed_value'] | |
| self.assertIn('age', tv) | |
| self.assertIsNone(tv['age']['value']) | |
| self.assertEqual(tv['age']['typed_value'], 25) | |
| # 'name' not in schema → must appear in overflow | |
| overflow = result['value'] | |
| if overflow: | |
| overflow_gv = GenericVariant(overflow, result['metadata']) | |
| self.assertIn('name', overflow_gv.to_python()) | |
| def test_unicode_overflow_and_rebuild_use_utf8_key_order(self): | |
| bmp_key = '\uff21' | |
| supplementary_key = '\U0001f600' | |
| obj_fields = self._obj_fields_for('col', [('keep', 'BIGINT')]) | |
| original = GenericVariant.from_python({ | |
| 'keep': 0, bmp_key: 1, supplementary_key: 2}) | |
| key_dict = parse_metadata_dict(original.metadata()) | |
| shredded = decompose_variant(original, obj_fields) | |
| self.assertEqual( | |
| [key_id for key_id, _ in _extract_overflow_fields(shredded['value'])], | |
| [key_dict[bmp_key], key_dict[supplementary_key]], | |
| ) | |
| schema = build_variant_schema(shredding_schema_to_arrow_type(obj_fields)) | |
| rebuilt_value, metadata = rebuild(shredded, schema, key_dict) | |
| self.assertEqual( | |
| [key_id for key_id, _ in _extract_overflow_fields(rebuilt_value)], | |
| [key_dict['keep'], key_dict[bmp_key], key_dict[supplementary_key]], | |
| ) | |
| self.assertEqual( | |
| GenericVariant(rebuilt_value, metadata).to_python(), | |
| original.to_python(), | |
| ) | |
| def test_decompose_absent_field_is_null(self): | |
| """A shredded field absent from the variant yields {value: None, typed_value: None}.""" | |
| obj_fields = self._obj_fields_for('col', [('missing_field', 'BIGINT')]) | |
| gv = GenericVariant.from_python({'other': 99}) | |
| result = decompose_variant(gv, obj_fields) | |
| tv = result['typed_value'] | |
| self.assertIsNone(tv['missing_field']['value']) | |
| self.assertIsNone(tv['missing_field']['typed_value']) | |
| def test_decompose_non_object_variant_all_overflow(self): | |
| """A non-object variant goes fully into overflow; typed_value is null (not a struct). | |
| Java's ShreddingUtils.rebuild requires typed_value to be null (not a non-null struct | |
| with all-null sub-fields) when the variant is not an object, so that the overflow | |
| binary can be consumed directly without expecting it to be an OBJECT variant. | |
| """ | |
| obj_fields = self._obj_fields_for('col', [('x', 'BIGINT')]) | |
| gv = GenericVariant.from_python(42) | |
| result = decompose_variant(gv, obj_fields) | |
| self.assertIsNotNone(result['value']) | |
| self.assertIsNone(result['typed_value']) | |
| def test_decompose_multiple_typed_fields(self): | |
| """Multiple shredded fields are all correctly split from the variant.""" | |
| obj_fields = self._obj_fields_for('col', [('score', 'BIGINT'), ('tag', 'VARCHAR')]) | |
| gv = GenericVariant.from_python({'score': 100, 'tag': 'gold', 'extra': True}) | |
| result = decompose_variant(gv, obj_fields) | |
| tv = result['typed_value'] | |
| self.assertEqual(tv['score']['typed_value'], 100) | |
| self.assertEqual(tv['tag']['typed_value'], 'gold') | |
| # 'extra' must be in overflow | |
| overflow = result['value'] | |
| if overflow: | |
| overflow_py = GenericVariant(overflow, result['metadata']).to_python() | |
| self.assertIn('extra', overflow_py) | |
| # =========================================================================== | |
| # 8. Full-chain via Paimon API | |
| # Write and read using CatalogFactory / Schema / write_arrow / to_arrow — | |
| # the same code path used by real applications. | |
| # =========================================================================== | |
| class TestVariantPaimonTable(unittest.TestCase): | |
| """End-to-end tests that use the real Paimon Python API. | |
| Each test creates its own table (unique name) within a shared in-process | |
| warehouse so that tests are fully independent. | |
| """ | |
| @classmethod | |
| def setUpClass(cls): | |
| cls.tmpdir = tempfile.mkdtemp() | |
| warehouse = os.path.join(cls.tmpdir, 'warehouse') | |
| cls.catalog = CatalogFactory.create({'warehouse': warehouse}) | |
| cls.catalog.create_database('default', True) | |
| @classmethod | |
| def tearDownClass(cls): | |
| shutil.rmtree(cls.tmpdir, ignore_errors=True) | |
| def _write_and_read(self, table, pa_data: pa.Table) -> pa.Table: | |
| """Write pa_data to the table and read all rows back.""" | |
| write_builder = table.new_batch_write_builder() | |
| table_write = write_builder.new_write() | |
| table_commit = write_builder.new_commit() | |
| table_write.write_arrow(pa_data) | |
| table_commit.commit(table_write.prepare_commit()) | |
| table_write.close() | |
| table_commit.close() | |
| read_builder = table.new_read_builder() | |
| splits = read_builder.new_scan().plan().splits() | |
| return read_builder.new_read().to_arrow(splits) | |
| def _pa_schema(self): | |
| return pa.schema([ | |
| pa.field('id', pa.int64()), | |
| pa.field('payload', _variant_arrow_type()), | |
| ]) | |
| def test_native_variant_projection_rejects_row_readers(self): | |
| for mode, options in ( | |
| ('append', {}), | |
| ('lazy', { | |
| 'data-evolution.enabled': 'true', | |
| 'row-tracking.enabled': 'true', | |
| }), | |
| ): | |
| with self.subTest(mode=mode): | |
| identifier = 'default.variant_projection_row_' + mode | |
| schema = Schema.from_pyarrow_schema( | |
| self._pa_schema(), options=options) | |
| self.catalog.create_table(identifier, schema, False) | |
| table = self.catalog.get_table(identifier) | |
| data = pa.table({ | |
| 'id': [1], | |
| 'payload': GenericVariant.to_arrow_array([ | |
| GenericVariant.from_python({'ratio': 1.25})]), | |
| }, schema=self._pa_schema()) | |
| write_builder = table.new_batch_write_builder() | |
| writer = write_builder.new_write() | |
| commit = write_builder.new_commit() | |
| writer.write_arrow(data) | |
| commit.commit(writer.prepare_commit()) | |
| writer.close() | |
| commit.close() | |
| builder = table.new_read_builder().with_projection( | |
| {'ratio': "try_variant_get(payload, '$.ratio', 'float')"}) | |
| splits = builder.new_scan().plan().splits() | |
| self.assertTrue(splits) | |
| read = builder.new_read() | |
| with self.assertRaisesRegex(RuntimeError, 'to_iterator'): | |
| read.to_iterator(splits) | |
| with self.assertRaisesRegex(RuntimeError, 'Torch row format'): | |
| read.to_torch(splits, streaming=True) | |
| with self.assertRaisesRegex(RuntimeError, 'Torch row format'): | |
| read.to_torch(splits, streaming=False) | |
| def _check_variant_projection_literal_top_level_names( | |
| self, suffix, native_read): | |
| pa_schema = pa.schema([ | |
| pa.field('id.dot', pa.int64()), | |
| pa.field('payload', _variant_arrow_type()), | |
| pa.field('payload.dot', _variant_arrow_type()), | |
| pa.field('payload[raw]', _variant_arrow_type()), | |
| ]) | |
| schema = Schema.from_pyarrow_schema( | |
| pa_schema, options={'read.native.enabled': 'true'}) | |
| identifier = 'default.variant_projection_literal_names_' + suffix | |
| self.catalog.create_table(identifier, schema, False) | |
| table = self.catalog.get_table(identifier) | |
| values = GenericVariant.to_arrow_array([ | |
| GenericVariant.from_python({'ratio': 1.25})]) | |
| data = pa.Table.from_arrays([ | |
| pa.array([1], type=pa.int64()), values, values, values, | |
| ], schema=pa_schema) | |
| write_builder = table.new_batch_write_builder() | |
| writer = write_builder.new_write() | |
| commit = write_builder.new_commit() | |
| writer.write_arrow(data) | |
| commit.commit(writer.prepare_commit()) | |
| writer.close() | |
| commit.close() | |
| for expressions, projection, column in ( | |
| ({'ratio': 'try_variant_get("payload.dot", "$.ratio", "float")'}, | |
| ['payload.dot'], 'payload.dot'), | |
| ({'ratio': 'try_variant_get("payload[raw]", "$.ratio", "float")'}, | |
| ['payload[raw]'], 'payload[raw]'), | |
| ({'id': 'id.dot', | |
| 'ratio': 'try_variant_get(payload, "$.ratio", "float")'}, | |
| ['id.dot', 'payload'], 'payload'), | |
| ): | |
| with self.subTest(expressions=expressions): | |
| builder = table.new_read_builder().with_projection( | |
| expressions) | |
| splits = builder.new_scan().plan().splits() | |
| self.assertTrue(splits) | |
| read = builder.new_read() | |
| kwargs = read._native_read_kwargs() | |
| self.assertEqual( | |
| [field.name for field in kwargs['read_type']], projection) | |
| self.assertNotIn('nested_projection', kwargs) | |
| field = next(field for field in kwargs['read_type'] | |
| if field.name == column) | |
| self.assertEqual( | |
| field.type.fields[0].description, | |
| '__VARIANT_METADATA$.ratio;false;UTC') | |
| self.assertEqual( | |
| read._output_arrow_schema().field('ratio').type, | |
| pa.float32()) | |
| if native_read: | |
| projected = read.to_arrow(splits) | |
| self.assertEqual(projected.column_names, | |
| list(expressions)) | |
| self.assertEqual( | |
| projected['ratio'].to_pylist(), [1.25]) | |
| plain = table.new_read_builder().with_projection(projection) | |
| self.assertEqual( | |
| plain.new_read().to_arrow(splits).column_names, | |
| projection) | |
| def test_variant_projection_literal_top_level_names_bridge(self): | |
| self._check_variant_projection_literal_top_level_names( | |
| 'bridge', native_read=False) | |
| @unittest.skipUnless(_native_variant_projection_available(), | |
| 'compatible native Variant projection API not installed') | |
| def test_variant_projection_literal_top_level_names_native(self): | |
| self._check_variant_projection_literal_top_level_names( | |
| 'native', native_read=True) | |
| @unittest.skipUnless(_native_variant_projection_available(), | |
| 'compatible native Variant projection API not installed') | |
| def test_named_variant_projection_sql_escaped_path(self): | |
| schema = Schema.from_pyarrow_schema( | |
| self._pa_schema(), options={'read.native.enabled': 'true'}) | |
| identifier = 'default.named_variant_escaped_path' | |
| self.catalog.create_table(identifier, schema, False) | |
| table = self.catalog.get_table(identifier) | |
| data = pa.table({ | |
| 'id': [1], | |
| 'payload': GenericVariant.to_arrow_array([ | |
| GenericVariant.from_python({"it's": 1.0, 'its': 2.0})]), | |
| }, schema=self._pa_schema()) | |
| write_builder = table.new_batch_write_builder() | |
| writer = write_builder.new_write() | |
| commit = write_builder.new_commit() | |
| writer.write_arrow(data) | |
| commit.commit(writer.prepare_commit()) | |
| writer.close() | |
| commit.close() | |
| from pypaimon_rust.datafusion import SQLContext | |
| context = SQLContext() | |
| context.register_catalog( | |
| 'paimon', {'warehouse': os.path.join(self.tmpdir, 'warehouse')}) | |
| for function in ('variant_get', 'try_variant_get'): | |
| expression = "%s(payload, '$[\"it''s\"]', 'float')" % function | |
| builder = table.new_read_builder().with_projection({'x': expression}) | |
| splits = builder.new_scan().plan().splits() | |
| actual = builder.new_read().to_arrow(splits)['x'].to_pylist() | |
| sql = pa.Table.from_batches(context.sql( | |
| 'SELECT %s AS x FROM %s' % (expression, identifier))) | |
| self.assertEqual([1.0], actual) | |
| self.assertEqual(sql['x'].to_pylist(), actual) | |
| @unittest.skipUnless(_native_variant_projection_available(), | |
| 'compatible native Variant projection API not installed') | |
| def test_named_variant_projection_batch_and_stream(self): | |
| schema = Schema.from_pyarrow_schema( | |
| self._pa_schema(), options={'read.native.enabled': 'true'}) | |
| identifier = 'default.named_variant_projection' | |
| self.catalog.create_table(identifier, schema, False) | |
| table = self.catalog.get_table(identifier) | |
| data = pa.table({ | |
| 'id': [1, 2], | |
| 'payload': GenericVariant.to_arrow_array([ | |
| GenericVariant.from_python({'x': 1.25, 'y': 2}), None]), | |
| }, schema=self._pa_schema()) | |
| writer_builder = table.new_batch_write_builder() | |
| writer = writer_builder.new_write() | |
| commit = writer_builder.new_commit() | |
| writer.write_arrow(data) | |
| commit.commit(writer.prepare_commit()) | |
| writer.close() | |
| commit.close() | |
| expressions = { | |
| 'identifier': 'id', | |
| 'x': "try_variant_get(payload, '$.x', 'float')", | |
| 'y': "variant_get(payload, '$.y', 'float')", | |
| } | |
| batch_builder = table.new_read_builder().with_projection(expressions) | |
| expected = {'identifier': [1, 2], 'x': [1.25, None], 'y': [2.0, None]} | |
| result = batch_builder.new_read().to_arrow( | |
| batch_builder.new_scan().plan().splits()) | |
| self.assertEqual(result.column_names, list(expressions)) | |
| self.assertEqual(result.to_pydict(), expected) | |
| stream_builder = table.new_stream_read_builder().with_projection( | |
| expressions) | |
| async def first_plan(): | |
| stream = stream_builder.new_streaming_scan().stream() | |
| try: | |
| return await stream.__anext__() | |
| finally: | |
| await stream.aclose() | |
| stream_result = stream_builder.new_read().to_arrow( | |
| asyncio.run(first_plan()).splits()) | |
| self.assertEqual(stream_result.to_pydict(), expected) | |
| bad = pa.table({ | |
| 'id': [3], | |
| 'payload': GenericVariant.to_arrow_array([ | |
| GenericVariant.from_python({'x': 'bad'})]), | |
| }, schema=self._pa_schema()) | |
| writer = writer_builder.new_write() | |
| commit = writer_builder.new_commit() | |
| writer.write_arrow(bad) | |
| commit.commit(writer.prepare_commit()) | |
| writer.close() | |
| commit.close() | |
| tolerant = table.new_read_builder().with_projection({ | |
| 'x': "try_variant_get(payload, '$.x', 'float')"}) | |
| result = tolerant.new_read().to_arrow( | |
| tolerant.new_scan().plan().splits()) | |
| self.assertEqual(result['x'].to_pylist(), [1.25, None, None]) | |
| strict = table.new_read_builder().with_projection({ | |
| 'x': "variant_get(payload, '$.x', 'float')"}) | |
| with self.assertRaisesRegex(Exception, 'Cannot cast Variant'): | |
| strict.new_read().to_arrow(strict.new_scan().plan().splits()) | |
| def test_plain_variant_write_and_read(self): | |
| """Plain VARIANT: GenericVariant → write_arrow → to_arrow → GenericVariant.""" | |
| schema = Schema.from_pyarrow_schema(self._pa_schema()) | |
| self.catalog.create_table('default.plain_variant', schema, False) | |
| table = self.catalog.get_table('default.plain_variant') | |
| test_uuid = uuid.UUID('12345678-1234-5678-1234-567812345678') | |
| gvs = [ | |
| GenericVariant.from_python({'age': 30, 'city': 'Beijing'}), | |
| GenericVariant.from_python({'score': 99, 'active': True}), | |
| GenericVariant.from_python([1, 2, 3]), | |
| GenericVariant.from_python({'dt': datetime.date(2024, 1, 15)}), | |
| GenericVariant.from_python( | |
| {'ts': datetime.datetime(2024, 1, 15, 12, 30, 45, 123456)} | |
| ), | |
| GenericVariant.from_python( | |
| {'ts_ltz': datetime.datetime( | |
| 2024, 1, 15, 12, 30, 45, 123456, tzinfo=datetime.timezone.utc | |
| )} | |
| ), | |
| GenericVariant.from_python({'id': test_uuid}), | |
| GenericVariant.from_python(test_uuid), | |
| ] | |
| data = pa.table( | |
| {'id': list(range(1, 9)), 'payload': GenericVariant.to_arrow_array(gvs)}, | |
| schema=self._pa_schema(), | |
| ) | |
| result = self._write_and_read(table, data) | |
| self.assertEqual(result.num_rows, 8) | |
| payload_col = result.column('payload') | |
| self.assertEqual( | |
| GenericVariant.from_arrow_struct(payload_col[0].as_py()).to_python(), | |
| {'age': 30, 'city': 'Beijing'}, | |
| ) | |
| self.assertEqual( | |
| GenericVariant.from_arrow_struct(payload_col[1].as_py()).to_python(), | |
| {'score': 99, 'active': True}, | |
| ) | |
| self.assertEqual( | |
| GenericVariant.from_arrow_struct(payload_col[2].as_py()).to_python(), | |
| [1, 2, 3], | |
| ) | |
| py3 = GenericVariant.from_arrow_struct(payload_col[3].as_py()).to_python() | |
| self.assertEqual(py3['dt'], datetime.date(2024, 1, 15)) | |
| py4 = GenericVariant.from_arrow_struct(payload_col[4].as_py()).to_python() | |
| self.assertEqual(py4['ts'], datetime.datetime(2024, 1, 15, 12, 30, 45, 123456)) | |
| py5 = GenericVariant.from_arrow_struct(payload_col[5].as_py()).to_python() | |
| self.assertEqual( | |
| py5['ts_ltz'], | |
| datetime.datetime(2024, 1, 15, 12, 30, 45, 123456, tzinfo=datetime.timezone.utc), | |
| ) | |
| py6 = GenericVariant.from_arrow_struct(payload_col[6].as_py()).to_python() | |
| self.assertEqual(py6['id'], test_uuid) | |
| py7 = GenericVariant.from_arrow_struct(payload_col[7].as_py()).to_python() | |
| self.assertEqual(py7, test_uuid) | |
| def test_plain_variant_null_row(self): | |
| """SQL-NULL VARIANT rows are stored and retrieved as None.""" | |
| schema = Schema.from_pyarrow_schema(self._pa_schema()) | |
| self.catalog.create_table('default.plain_variant_null', schema, False) | |
| table = self.catalog.get_table('default.plain_variant_null') | |
| gvs = [GenericVariant.from_python({'x': 1}), None, GenericVariant.from_python({'x': 3})] | |
| data = pa.table( | |
| {'id': [1, 2, 3], 'payload': GenericVariant.to_arrow_array(gvs)}, | |
| schema=self._pa_schema(), | |
| ) | |
| result = self._write_and_read(table, data) | |
| payload_col = result.column('payload') | |
| self.assertIsNotNone(payload_col[0].as_py()) | |
| self.assertIsNone(payload_col[1].as_py()) | |
| self.assertIsNotNone(payload_col[2].as_py()) | |
| def test_shredded_variant_write_and_read(self): | |
| """Shredded VARIANT: writer shreds automatically, reader assembles transparently.""" | |
| # Build shredding schema manually because nested ARRAY/ROW require JSON objects, | |
| # not just atomic type strings. | |
| shredding_json = json.dumps({ | |
| 'type': 'ROW', | |
| 'fields': [ | |
| { | |
| 'id': 0, | |
| 'name': 'payload', | |
| 'type': { | |
| 'type': 'ROW', | |
| 'fields': [ | |
| {'name': 'name', 'type': 'VARCHAR'}, | |
| {'name': 'active', 'type': 'BOOLEAN'}, | |
| {'name': 'age', 'type': 'TINYINT'}, | |
| {'name': 'score', 'type': 'SMALLINT'}, | |
| {'name': 'count', 'type': 'INT'}, | |
| {'name': 'id', 'type': 'BIGINT'}, | |
| {'name': 'ratio', 'type': 'DOUBLE'}, | |
| {'name': 'amount', 'type': 'DECIMAL(10,2)'}, | |
| {'name': 'fixed', 'type': 'BINARY'}, | |
| {'name': 'raw', 'type': 'VARBINARY'}, | |
| { | |
| 'name': 'tags', | |
| 'type': {'type': 'ARRAY', 'element': 'INT'}, | |
| }, | |
| { | |
| 'name': 'address', | |
| 'type': { | |
| 'type': 'ROW', | |
| 'fields': [ | |
| {'name': 'city', 'type': 'VARCHAR'}, | |
| {'name': 'zip', 'type': 'INT'}, | |
| ], | |
| }, | |
| }, | |
| ], | |
| }, | |
| } | |
| ], | |
| }) | |
| schema = Schema.from_pyarrow_schema( | |
| self._pa_schema(), | |
| options={'variant.shreddingSchema': shredding_json}, | |
| ) | |
| self.catalog.create_table('default.shredded_variant', schema, False) | |
| table = self.catalog.get_table('default.shredded_variant') | |
| gvs = [ | |
| GenericVariant.from_python( | |
| { | |
| 'name': 'Apache Paimon', | |
| 'active': True, | |
| 'age': 3, | |
| 'score': 3000, | |
| 'count': 400000, | |
| 'id': 12345678901234, | |
| 'ratio': 1.012345678901234567890123456789, | |
| 'amount': Decimal('100.99'), | |
| # BINARY stores raw bytes; fixed-length semantics are Parquet-level. | |
| 'fixed': 'Apache Paimon'.encode('utf-8'), | |
| # VARBINARY stores raw bytes without padding. | |
| 'raw': b'\x01\x02\x03\x04\x05', | |
| 'tags': [1, 2, 3], | |
| 'address': {'city': 'Beijing', 'zip': 100000}, | |
| } | |
| ), | |
| GenericVariant.from_python( | |
| { | |
| 'name': 'Pypaimon', | |
| 'active': False, | |
| 'age': 1, | |
| 'score': 100, | |
| 'count': 42, | |
| 'id': 98765432109876, | |
| 'ratio': 2.718281828459045, | |
| 'amount': Decimal('42.50'), | |
| 'fixed': b'\x00\x01\x02\x03', | |
| 'raw': 'hello'.encode('utf-8'), | |
| 'tags': [10, 20], | |
| 'address': {'city': 'Shanghai', 'zip': 200000}, | |
| } | |
| ), | |
| ] | |
| data = pa.table( | |
| {'id': [1, 2], 'payload': GenericVariant.to_arrow_array(gvs)}, | |
| schema=self._pa_schema(), | |
| ) | |
| result = self._write_and_read(table, data) | |
| self.assertEqual(result.num_rows, 2) | |
| payload_col = result.column('payload') | |
| py0 = GenericVariant.from_arrow_struct(payload_col[0].as_py()).to_python() | |
| self.assertEqual(py0['name'], 'Apache Paimon') | |
| self.assertEqual(py0['active'], True) | |
| self.assertEqual(py0['age'], 3) | |
| self.assertEqual(py0['score'], 3000) | |
| self.assertEqual(py0['count'], 400000) | |
| self.assertEqual(py0['id'], 12345678901234) | |
| self.assertAlmostEqual(py0['ratio'], 1.012345678901234567890123456789) | |
| self.assertEqual(py0['amount'], Decimal('100.99')) | |
| self.assertEqual(py0['fixed'], 'Apache Paimon'.encode('utf-8')) | |
| self.assertEqual(py0['raw'], b'\x01\x02\x03\x04\x05') | |
| self.assertEqual(py0['tags'], [1, 2, 3]) | |
| self.assertEqual(py0['address'], {'city': 'Beijing', 'zip': 100000}) | |
| py1 = GenericVariant.from_arrow_struct(payload_col[1].as_py()).to_python() | |
| self.assertEqual(py1['name'], 'Pypaimon') | |
| self.assertEqual(py1['active'], False) | |
| self.assertEqual(py1['age'], 1) | |
| self.assertEqual(py1['score'], 100) | |
| self.assertEqual(py1['count'], 42) | |
| self.assertEqual(py1['id'], 98765432109876) | |
| self.assertAlmostEqual(py1['ratio'], 2.718281828459045) | |
| self.assertEqual(py1['amount'], Decimal('42.50')) | |
| self.assertEqual(py1['fixed'], b'\x00\x01\x02\x03') | |
| self.assertEqual(py1['raw'], 'hello'.encode('utf-8')) | |
| self.assertEqual(py1['tags'], [10, 20]) | |
| self.assertEqual(py1['address'], {'city': 'Shanghai', 'zip': 200000}) | |
| def test_shredded_variant_overflow_preserved(self): | |
| """Fields outside the shredding schema survive in overflow bytes end-to-end.""" | |
| shredding_json = _schema_json('payload', [('age', 'BIGINT')]) | |
| schema = Schema.from_pyarrow_schema( | |
| self._pa_schema(), | |
| options={'variant.shreddingSchema': shredding_json}, | |
| ) | |
| self.catalog.create_table('default.shredded_overflow', schema, False) | |
| table = self.catalog.get_table('default.shredded_overflow') | |
| # 'city' and 'active' are NOT in the shredding schema → stored as overflow | |
| gv = GenericVariant.from_python({'age': 30, 'city': 'Beijing', 'active': True}) | |
| data = pa.table( | |
| {'id': [1], 'payload': GenericVariant.to_arrow_array([gv])}, | |
| schema=self._pa_schema(), | |
| ) | |
| result = self._write_and_read(table, data) | |
| py = GenericVariant.from_arrow_struct(result.column('payload')[0].as_py()).to_python() | |
| self.assertEqual(py['age'], 30) | |
| self.assertEqual(py['city'], 'Beijing') | |
| self.assertEqual(py['active'], True) | |
| def test_shredded_variant_null_row(self): | |
| """SQL-NULL VARIANT rows survive the shred → assemble cycle as None.""" | |
| shredding_json = _schema_json('payload', [('x', 'BIGINT')]) | |
| schema = Schema.from_pyarrow_schema( | |
| self._pa_schema(), | |
| options={'variant.shreddingSchema': shredding_json}, | |
| ) | |
| self.catalog.create_table('default.shredded_null', schema, False) | |
| table = self.catalog.get_table('default.shredded_null') | |
| gvs = [GenericVariant.from_python({'x': 7}), None] | |
| data = pa.table( | |
| {'id': [1, 2], 'payload': GenericVariant.to_arrow_array(gvs)}, | |
| schema=self._pa_schema(), | |
| ) | |
| result = self._write_and_read(table, data) | |
| payload_col = result.column('payload') | |
| gv = GenericVariant.from_arrow_struct(payload_col[0].as_py()) | |
| self.assertEqual(gv.to_python(), {'x': 7}) | |
| self.assertIsNone(payload_col[1].as_py()) | |
| if __name__ == '__main__': | |
| unittest.main() |