Merge pull request #50 from Recidiviz/dan/arrow-extreme

ohaibbq · web-flow · commit 431a583698fc · 2025-11-20T21:02:23.000-06:00
fix: Arrow serialization of extreme values
diff --git a/server/storage_test.go b/server/storage_test.go
@@ -1098,7 +1098,7 @@ func generateExampleMessages(numMessages int) ([][]byte, error) {
 			// NUMERIC and BIGNUMERIC can be passed as string, or more efficiently
 			// using a packed byte representation.
 			NumericCol:    proto.String("99999999999999999999999999999.999999999"),
-			BignumericCol: proto.String("578960446186580977117854925043439539266.34992332820282019728792003956564819967"),
+			BignumericCol: proto.String("-578960446186580977117854925043439539266.34992332820282019728792003956564819968"),
 
 			// TIME also uses literal format.
 			TimeCol: proto.String("12:13:14.000000"),
@@ -1394,18 +1394,17 @@ func TestStorageReadWithAPICreatedTable(t *testing.T) {
 
 			testData := []TestRow{
 				{
-					StringCol:    "hello",
-					IntCol:       42,
-					FloatCol:     3.14,
-					BoolCol:      true,
-					BytesCol:     []byte("abc"),
-					DateCol:      "2024-01-01",
-					DatetimeCol:  "2024-01-01T12:00:00",
-					TimestampCol: time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC),
-					TimeCol:      civil.Time{Hour: 12},
-					NumericCol:   "123.456",
-					// BIGNUMERIC max: 38 integer digits, 38 fractional digits
-					BignumericCol: "12345678901234567890123456789012345678.12345678901234567890123456789012345678",
+					StringCol:     "hello",
+					IntCol:        42,
+					FloatCol:      3.14,
+					BoolCol:       true,
+					BytesCol:      []byte("abc"),
+					DateCol:       "2024-01-01",
+					DatetimeCol:   "2024-01-01T12:00:00",
+					TimestampCol:  time.Date(2024, 1, 1, 12, 0, 0, 0, time.UTC),
+					TimeCol:       civil.Time{Hour: 12},
+					NumericCol:    "123.456",
+					BignumericCol: "578960446186580977117854925043439539266.34992332820282019728792003956564819967",
 					ArrayCol:      []string{"x", "y"},
 					StructCol:     map[string]interface{}{"field1": int64(1), "field2": "nested"},
 				},
diff --git a/test/python/extreme_values_test.py b/test/python/extreme_values_test.py
@@ -0,0 +1,168 @@
+"""Tests for extreme/boundary values in BigQuery types.
+
+This test suite validates that extreme values (min/max) for various BigQuery
+data types are correctly handled by both the REST API and Storage Arrow API.
+"""
+import datetime
+from decimal import Decimal
+from typing import Any, Dict, Iterable
+
+import grpc
+from google.cloud import bigquery, bigquery_storage
+from google.cloud.bigquery_storage_v1.services.big_query_read.transports import (
+    BigQueryReadGrpcTransport,
+)
+from google.api_core.client_options import ClientOptions
+
+from utils.big_query_emulator_test_case import BigQueryEmulatorTestCase
+
+
+class TestExtremeValues(BigQueryEmulatorTestCase):
+    """Tests for extreme/boundary values in BigQuery types."""
+
+    def _run_test_with_both_apis(
+        self,
+        query_str: str,
+        expected_result: Iterable[Dict[str, Any]],
+    ) -> None:
+        """Run the same query against both REST API and Storage Arrow API.
+
+        Args:
+            query_str: SQL query to execute
+            expected_result: Expected results as list of dictionaries
+        """
+        # Test 1: REST API via job.result()
+        query_job = self.client.query(query=query_str)
+        rest_contents = list(
+            {key: row.get(key) for key in row.keys()} for row in query_job.result()
+        )
+        self.assertEqual(
+            list(expected_result),
+            rest_contents,
+            f"REST API results don't match. Expected: {expected_result}, Got: {rest_contents}"
+        )
+
+        # Test 2: Storage Arrow API via query_job.to_dataframe()
+        # Create a fresh query for the Storage API test
+        query_job = self.client.query(query=query_str)
+
+        # Patch the client to use Storage API
+        host = f"localhost:{self.emulator.grpc_port}"
+        channel = grpc.insecure_channel(host)
+        transport = BigQueryReadGrpcTransport(channel=channel, host=host)
+        bqstorage_client = bigquery_storage.BigQueryReadClient(
+            transport=transport,
+            client_options=ClientOptions(api_endpoint=host),
+        )
+
+        # Get DataFrame using Storage API
+        df = query_job.to_dataframe(bqstorage_client=bqstorage_client)
+
+        # Convert DataFrame to same format as expected_result
+        arrow_contents = df.to_dict(orient='records')
+
+        self.assertEqual(
+            list(expected_result),
+            arrow_contents,
+            f"Storage Arrow API results don't match. Expected: {expected_result}, Got: {arrow_contents}"
+        )
+
+        channel.close()
+
+    def test_timestamp_min_max(self) -> None:
+        """Tests resolution of https://github.com/goccy/go-zetasqlite/issues/132
+        and https://github.com/goccy/bigquery-emulator/issues/262"""
+        self._run_test_with_both_apis(
+            """SELECT TIMESTAMP '0001-01-01 00:00:00.000000+00' AS min_ts, TIMESTAMP '9999-12-31 23:59:59.999999+00' AS max_ts""",
+            expected_result=[
+                {
+                    "min_ts": datetime.datetime(
+                        1, 1, 1, 0, 0, tzinfo=datetime.timezone.utc
+                    ),
+                    "max_ts": datetime.datetime(
+                        9999, 12, 31, 23, 59, 59, 999999, tzinfo=datetime.timezone.utc
+                    ),
+                }
+            ],
+        )
+
+    def test_int64_min_max(self) -> None:
+        """Tests extreme INT64 values (64-bit signed integer range)"""
+        self._run_test_with_both_apis(
+            """SELECT -9223372036854775808 AS min_int64, 9223372036854775807 AS max_int64""",
+            expected_result=[
+                {
+                    "min_int64": -9223372036854775808,
+                    "max_int64": 9223372036854775807,
+                }
+            ],
+        )
+
+    def test_numeric_min_max(self) -> None:
+        """Tests extreme NUMERIC values (38 digits, 9 after decimal point)"""
+        self._run_test_with_both_apis(
+            """SELECT
+                NUMERIC '-99999999999999999999999999999.999999999' AS min_numeric,
+                NUMERIC '99999999999999999999999999999.999999999' AS max_numeric
+            """,
+            expected_result=[
+                {
+                    "min_numeric": Decimal('-99999999999999999999999999999.999999999'),
+                    "max_numeric": Decimal('99999999999999999999999999999.999999999'),
+                }
+            ],
+        )
+
+    def test_bignumeric_min_max(self) -> None:
+        """Tests extreme BIGNUMERIC values"""
+        self._run_test_with_both_apis(
+            """SELECT
+                BIGNUMERIC '-578960446186580977117854925043439539266.34992332820282019728792003956564819968' AS min_bignumeric,
+                BIGNUMERIC '578960446186580977117854925043439539266.34992332820282019728792003956564819967' AS max_bignumeric
+            """,
+            expected_result=[
+                {
+                    "min_bignumeric": Decimal('-578960446186580977117854925043439539266.34992332820282019728792003956564819968'),
+                    "max_bignumeric": Decimal('578960446186580977117854925043439539266.34992332820282019728792003956564819967'),
+                }
+            ],
+        )
+
+    def test_date_min_max(self) -> None:
+        """Tests extreme DATE values"""
+        self._run_test_with_both_apis(
+            """SELECT DATE '0001-01-01' AS min_date, DATE '9999-12-31' AS max_date""",
+            expected_result=[
+                {
+                    "min_date": datetime.date(1, 1, 1),
+                    "max_date": datetime.date(9999, 12, 31),
+                }
+            ],
+        )
+
+    def test_datetime_min_max(self) -> None:
+        """Tests extreme DATETIME values"""
+        self._run_test_with_both_apis(
+            """SELECT
+                DATETIME '0001-01-01 00:00:00' AS min_datetime,
+                DATETIME '9999-12-31 23:59:59.999999' AS max_datetime
+            """,
+            expected_result=[
+                {
+                    "min_datetime": datetime.datetime(1, 1, 1, 0, 0, 0),
+                    "max_datetime": datetime.datetime(9999, 12, 31, 23, 59, 59, 999999),
+                }
+            ],
+        )
+
+    def test_time_min_max(self) -> None:
+        """Tests extreme TIME values"""
+        self._run_test_with_both_apis(
+            """SELECT TIME '00:00:00' AS min_time, TIME '23:59:59.999999' AS max_time""",
+            expected_result=[
+                {
+                    "min_time": datetime.time(0, 0, 0),
+                    "max_time": datetime.time(23, 59, 59, 999999),
+                }
+            ],
+        )
diff --git a/types/arrow.go b/types/arrow.go
@@ -1,19 +1,19 @@
 package types
 
 import (
+	"cloud.google.com/go/bigquery"
 	"encoding/base64"
+	"encoding/binary"
 	"fmt"
-	"math/big"
-	"strconv"
-	"strings"
-	"time"
-
 	"github.com/apache/arrow-go/v18/arrow"
 	"github.com/apache/arrow-go/v18/arrow/array"
 	"github.com/apache/arrow-go/v18/arrow/decimal128"
 	"github.com/apache/arrow-go/v18/arrow/decimal256"
 	"github.com/goccy/go-zetasqlite"
 	bigqueryv2 "google.golang.org/api/bigquery/v2"
+	"math/big"
+	"strconv"
+	"strings"
 )
 
 func TableToARROW(t *bigqueryv2.Table) (*arrow.Schema, error) {
@@ -59,7 +59,7 @@ func tableFieldToARROW(f *bigqueryv2.TableFieldSchema) (*arrow.Field, error) {
 	case FieldBytes:
 		return &arrow.Field{Name: f.Name, Type: arrow.BinaryTypes.Binary}, nil
 	case FieldDate:
-		return &arrow.Field{Name: f.Name, Type: arrow.PrimitiveTypes.Date32}, nil
+		return &arrow.Field{Name: f.Name, Type: arrow.PrimitiveTypes.Date64}, nil
 	case FieldDatetime:
 		return &arrow.Field{
 			Name: f.Name,
@@ -96,11 +96,21 @@ func tableFieldToARROW(f *bigqueryv2.TableFieldSchema) (*arrow.Field, error) {
 		return &arrow.Field{Name: f.Name, Type: arrow.StructOf(fields...)}, nil
 	case FieldNumeric:
 		// NUMERIC is a DECIMAL with precision 38, scale 9
-		return &arrow.Field{Name: f.Name, Type: &arrow.Decimal128Type{Precision: 38, Scale: 9}}, nil
+		return &arrow.Field{Name: f.Name, Type: &arrow.Decimal128Type{
+			Precision: bigquery.NumericPrecisionDigits,
+			Scale:     bigquery.NumericScaleDigits,
+		}}, nil
 	case FieldBignumeric:
-		// BIGNUMERIC is a DECIMAL with precision 76, scale 38
-		// BigQuery supports 76.76 digits (76 full digits, 77th is partial)
-		return &arrow.Field{Name: f.Name, Type: &arrow.Decimal256Type{Precision: 76, Scale: 38}}, nil
+		// In BigQuery, BIGNUMERIC is a DECIMAL with precision 76 (partial 77), scale 38
+		// Values requiring 77 digits when scaled by 10^38 work fine, including the maximum value (±2^255 / 10^38).
+		// These values can technically be encoded into the Arrow format, but most libraries (including arrow-go)
+		// raise validation errors when trying to build them.
+		// The values returned by the BigQuery Storage Read API raise errors when you try to validate them client side
+		// but if you only access their values, it is fine.
+		return &arrow.Field{Name: f.Name, Type: &arrow.Decimal256Type{
+			Precision: bigquery.BigNumericPrecisionDigits, // 76
+			Scale:     bigquery.BigNumericScaleDigits,     // 38
+		}}, nil
 	case FieldGeography:
 		return &arrow.Field{Name: f.Name, Type: arrow.BinaryTypes.String}, nil
 	case FieldInterval:
@@ -148,12 +158,12 @@ func AppendValueToARROWBuilder(ptrv *string, builder array.Builder) error {
 		}
 		b.Append(decoded)
 		return nil
-	case *array.Date32Builder:
+	case *array.Date64Builder:
 		t, err := parseDate(v)
 		if err != nil {
 			return err
 		}
-		b.Append(arrow.Date32(int32(t.Sub(time.Unix(0, 0)) / (24 * time.Hour))))
+		b.Append(arrow.Date64FromTime(t))
 		return nil
 	case *array.Time64Builder:
 		t, err := parseTime(v)
@@ -198,32 +208,83 @@ func AppendValueToARROWBuilder(ptrv *string, builder array.Builder) error {
 		// Convert to integer (this truncates any remaining fractional part)
 		scaledInt := new(big.Int).Div(scaled.Num(), scaled.Denom())
 
-		// Convert to decimal128.Num
-		num := decimal128.FromBigInt(scaledInt)
-		b.Append(num)
+		// Convert to decimal128.Num using Arrow's FromBigInt (handles two's complement correctly)
+		b.Append(decimal128.FromBigInt(scaledInt))
 		return nil
 	case *array.Decimal256Builder:
-		// BIGNUMERIC type: precision 77, scale 38
-		// Parse the string value to a big.Rat, then convert to scaled integer
+		// BIGNUMERIC type: precision 76, scale 38
+		// NOTE: BigQuery declares decimal256(76, 38) in the schema but doesn't enforce
+		// precision during encoding. Values requiring 77 digits when scaled work fine in BigQuery.
+		// We bypass Arrow's FromBigInt validation (bitlen > 255 check) and manually construct
+		// the Decimal256 using the same logic but without the strict check.
+
+		// Parse as rational number
 		rat := new(big.Rat)
 		if _, ok := rat.SetString(v); !ok {
-			return fmt.Errorf("failed to parse decimal value: %s", v)
+			return fmt.Errorf("failed to parse BIGNUMERIC value: %s", v)
 		}
 
-		// Scale the value by 10^scale to get the integer representation
-		scale := int32(38)
-		scaleFactor := new(big.Int).Exp(big.NewInt(10), big.NewInt(int64(scale)), nil)
-
-		// Multiply the rational by the scale factor
+		// Scale by 10^38 to get integer representation
+		scale := int64(bigquery.BigNumericScaleDigits) // 38
+		scaleFactor := new(big.Int).Exp(big.NewInt(10), big.NewInt(scale), nil)
 		scaled := new(big.Rat).Mul(rat, new(big.Rat).SetInt(scaleFactor))
 
-		// Convert to integer (this truncates any remaining fractional part)
+		// Convert to integer (truncating any remaining fractional part)
 		scaledInt := new(big.Int).Div(scaled.Num(), scaled.Denom())
 
-		// Convert to decimal256.Num
-		num := decimal256.FromBigInt(scaledInt)
+		// Replicate decimal256.FromBigInt logic without the bitlen > 255 check
+		// This matches how Arrow handles two's complement representation
+		var num decimal256.Num
+		if scaledInt.Sign() == 0 {
+			// Zero value, return default
+			b.Append(num)
+			return nil
+		}
+
+		num = decimal256FromScaledInt(scaledInt)
+
 		b.Append(num)
 		return nil
 	}
 	return fmt.Errorf("unexpected builder type %T", builder)
 }
+
+func decimal256FromScaledInt(scaledInt *big.Int) decimal256.Num {
+	b := scaledInt.FillBytes(make([]byte, 32))
+
+	var limbs [4]uint64
+
+	// Arrow Decimal256 uses little-endian uint64 limbs.
+	// BigQuery and BigInt.FillBytes produce big-endian bytes.
+	//
+	// So the 256-bit structure:
+	//   b[0] ... b[31]   (big endian)
+	// maps to Arrow limbs:
+	//   limbs[0] = low  64 bits
+	//   limbs[3] = high 64 bits
+
+	for i := 0; i < 4; i++ {
+		// Big-endian slice for limb i:
+		start := 32 - (i+1)*8
+		end := 32 - i*8
+
+		// Convert this BE 8-byte block into LE uint64
+		// Arrow stores each limb as native endian (LE)
+		limbs[i] = binary.LittleEndian.Uint64(reverse8(b[start:end]))
+	}
+
+	dec := decimal256.New(limbs[3], limbs[2], limbs[1], limbs[0])
+	// If negative, negate to get two's complement representation
+	if scaledInt.Sign() < 0 {
+		dec = dec.Negate()
+	}
+	return dec
+}
+
+func reverse8(b []byte) []byte {
+	out := make([]byte, 8)
+	for i := 0; i < 8; i++ {
+		out[i] = b[7-i]
+	}
+	return out
+}
diff --git a/types/arrow_test.go b/types/arrow_test.go
