|
15 | 15 | # specific language governing permissions and limitations |
16 | 16 | # under the License. |
17 | 17 |
|
| 18 | +from collections.abc import Callable |
18 | 19 | from tempfile import TemporaryDirectory |
19 | 20 |
|
20 | 21 | import pytest |
21 | 22 | from pydantic import Field |
22 | 23 |
|
| 24 | +from pyiceberg.avro.decoder import BinaryDecoder, StreamingBinaryDecoder |
| 25 | +from pyiceberg.avro.decoder_fast import CythonBinaryDecoder |
23 | 26 | from pyiceberg.avro.file import AvroFile |
24 | 27 | from pyiceberg.avro.reader import ( |
25 | 28 | DecimalReader, |
|
32 | 35 | StringReader, |
33 | 36 | StructReader, |
34 | 37 | ) |
35 | | -from pyiceberg.avro.resolver import resolve_reader, resolve_writer |
| 38 | +from pyiceberg.avro.resolver import EnumReader, resolve_reader, resolve_writer |
36 | 39 | from pyiceberg.avro.writer import ( |
37 | 40 | BinaryWriter, |
38 | 41 | DefaultWriter, |
|
46 | 49 | ) |
47 | 50 | from pyiceberg.exceptions import ResolveError |
48 | 51 | from pyiceberg.io.pyarrow import PyArrowFileIO |
49 | | -from pyiceberg.manifest import MANIFEST_ENTRY_SCHEMAS |
| 52 | +from pyiceberg.manifest import MANIFEST_ENTRY_SCHEMAS, ManifestEntryStatus |
50 | 53 | from pyiceberg.schema import Schema |
51 | 54 | from pyiceberg.typedef import Record |
52 | 55 | from pyiceberg.types import ( |
@@ -418,3 +421,30 @@ def test_writer_missing_optional_in_read_schema() -> None: |
418 | 421 | expected = StructWriter(field_writers=((None, OptionWriter(option=StringWriter())),)) |
419 | 422 |
|
420 | 423 | assert actual == expected |
| 424 | + |
| 425 | + |
| 426 | +@pytest.mark.parametrize("decoder_class", [StreamingBinaryDecoder, CythonBinaryDecoder]) |
| 427 | +def test_enum_reader_skip_advances_the_decoder(decoder_class: Callable[[bytes], BinaryDecoder]) -> None: |
| 428 | + """Skipping an enum field must consume its bytes, or the next field is read from them.""" |
| 429 | + # Two ints, zigzag encoded: the enum's ordinal 1, then the next field's value 12. |
| 430 | + decoder = decoder_class(b"\x02\x18") |
| 431 | + reader = EnumReader(ManifestEntryStatus, IntegerReader()) |
| 432 | + |
| 433 | + reader.skip(decoder) |
| 434 | + |
| 435 | + assert IntegerReader().read(decoder) == 12 |
| 436 | + |
| 437 | + |
| 438 | +@pytest.mark.parametrize("decoder_class", [StreamingBinaryDecoder, CythonBinaryDecoder]) |
| 439 | +def test_enum_reader_skip_matches_read(decoder_class: Callable[[bytes], BinaryDecoder]) -> None: |
| 440 | + """Reading and skipping must leave the decoder at the same position.""" |
| 441 | + encoded = b"\x02\x18" |
| 442 | + reader = EnumReader(ManifestEntryStatus, IntegerReader()) |
| 443 | + |
| 444 | + read_decoder = decoder_class(encoded) |
| 445 | + assert reader.read(read_decoder) == ManifestEntryStatus.ADDED |
| 446 | + |
| 447 | + skip_decoder = decoder_class(encoded) |
| 448 | + reader.skip(skip_decoder) |
| 449 | + |
| 450 | + assert IntegerReader().read(read_decoder) == IntegerReader().read(skip_decoder) |
0 commit comments