[FLINK-40370][python] Frame bytes-backed decimals as a bytes field

flink-python/pyflink/fn_execution/formats/avro.py

@@ -17,7 +17,7 @@
 ################################################################################
 import struct
 
-from avro.errors import AvroTypeException, SchemaResolutionException
+from avro.errors import AvroOutOfScaleException, AvroTypeException, SchemaResolutionException
 from avro.io import (
     BinaryDecoder,
     BinaryEncoder,
@@ -87,6 +87,13 @@ def read_bytes(self):
         assert (nbytes >= 0), nbytes
         return self.read(nbytes)
 
+    def read_decimal_from_bytes(self, precision, scale):
+        # avro's implementation sizes the payload with read_long, which is an Avro zig-zag long
+        # upstream but a fixed 8-byte long here. On the JVM a bytes-backed decimal is framed like
+        # any other bytes field, so the size is a 4-byte int.
+        size = self.read_int()
+        return self.read_decimal_from_fixed(precision, scale, size)
+
     def skip_int(self):
         self.skip(4)
 
@@ -198,6 +205,28 @@ def write_bytes(self, datum):
         self.write_int(len(datum))
         self.write(datum)
 
+    def write_decimal_bytes(self, datum, scale):
+        # avro's implementation sizes the payload with write_long, which is an Avro zig-zag long
+        # upstream but a fixed 8-byte long here, so the JVM reader consumed the size as the whole
+        # bytes field and every field after it shifted. Frame it as a bytes field instead; the
+        # payload itself is the same two's-complement big-endian unscaled value.
+        sign, digits, exp = datum.as_tuple()
+        if (-1 * int(exp)) > scale:
+            raise AvroOutOfScaleException(scale, datum, exp)
+
+        unscaled_datum = 0
+        for digit in digits:
+            unscaled_datum = (unscaled_datum * 10) + digit
+
+        bits_req = unscaled_datum.bit_length() + 1
+        if sign:
+            unscaled_datum = -unscaled_datum
+
+        bytes_req = bits_req // 8
+        bytes_req += 1 if (bytes_req << 3) < bits_req else 0
+
+        self.write_bytes(unscaled_datum.to_bytes(bytes_req, 'big', signed=True))
+
 
 class FlinkAvroDatumWriter(DatumWriter):
 

flink-python/pyflink/fn_execution/tests/test_avro_format.py

@@ -0,0 +1,103 @@
+################################################################################
+#  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.
+################################################################################
+"""Tests for the JVM-compatible Avro encoder and decoder."""
+import io
+import logging
+import unittest
+from decimal import Decimal
+
+import avro.schema
+
+from pyflink.fn_execution.formats.avro import (
+    FlinkAvroDatumReader,
+    FlinkAvroDatumWriter,
+    FlinkAvroDecoder,
+    FlinkAvroEncoder,
+)
+from pyflink.testing.test_case_utils import PyFlinkTestCase
+
+DECIMAL_THEN_INT = """
+{
+  "type": "record",
+  "name": "DecimalRecord",
+  "fields": [
+    {
+      "name": "amount",
+      "type": {"type": "bytes", "logicalType": "decimal", "precision": 8, "scale": 2}
+    },
+    {"name": "tail", "type": "int"}
+  ]
+}
+"""
+
+
+class AvroFormatTests(PyFlinkTestCase):
+
+    @staticmethod
+    def _encode(schema, record):
+        buffer = io.BytesIO()
+        FlinkAvroDatumWriter(schema).write(record, FlinkAvroEncoder(buffer))
+        return buffer.getvalue()
+
+    def test_decimal_bytes_are_framed_like_any_other_bytes_field(self):
+        schema = avro.schema.parse(DECIMAL_THEN_INT)
+        encoded = self._encode(schema, {"amount": Decimal("12.34"), "tail": 7})
+
+        # 4-byte length, the two's-complement unscaled value, then the next field. Sizing the
+        # payload as an 8-byte long instead shifts every following field.
+        self.assertEqual("0000000204d200000007", encoded.hex())
+
+    def test_decimal_bytes_are_readable_field_by_field(self):
+        schema = avro.schema.parse(DECIMAL_THEN_INT)
+        encoded = self._encode(schema, {"amount": Decimal("12.34"), "tail": 7})
+
+        decoder = FlinkAvroDecoder(io.BytesIO(encoded))
+        self.assertEqual(b"\x04\xd2", decoder.read_bytes())
+        self.assertEqual(7, decoder.read_int())
+
+    def test_decimal_payload_is_unchanged_two_s_complement(self):
+        schema = avro.schema.parse(DECIMAL_THEN_INT)
+        expected = {
+            "0.00": "00",
+            "12.34": "04d2",
+            "-12.34": "fb2e",
+            "1.28": "0080",
+            "-1.28": "ff80",
+            "-1.29": "ff7f",
+        }
+        for value, payload in expected.items():
+            with self.subTest(value=value):
+                encoded = self._encode(schema, {"amount": Decimal(value), "tail": 0})
+                length = len(bytes.fromhex(payload))
+                self.assertEqual(length.to_bytes(4, "big").hex(), encoded[:4].hex())
+                self.assertEqual(payload, encoded[4:4 + length].hex())
+
+    def test_decimal_round_trip(self):
+        schema = avro.schema.parse(DECIMAL_THEN_INT)
+        for value in ("0.00", "12.34", "-12.34", "1.27", "-1.28", "999999.99", "-999999.99"):
+            with self.subTest(value=value):
+                record = {"amount": Decimal(value), "tail": 7}
+                encoded = self._encode(schema, record)
+                decoded = FlinkAvroDatumReader(schema, schema).read(
+                    FlinkAvroDecoder(io.BytesIO(encoded)))
+                self.assertEqual(record, decoded)
+
+
+if __name__ == '__main__':
+    logging.getLogger().setLevel(logging.INFO)
+    unittest.main()