From f248f2b209e71d127d70e8358a835c7dff3ded74 Mon Sep 17 00:00:00 2001
From: Alvaro Viebrantz <aviebrantz@google.com>
Date: Wed, 17 Apr 2024 17:00:13 -0400
Subject: [PATCH] feat(bigquery): RANGE support when reading Arrow format

---
 bigquery/arrow.go                    | 15 +++++
 bigquery/integration_test.go         | 89 +++++++++++++---------------
 bigquery/storage_integration_test.go | 32 +++++-----
 3 files changed, 74 insertions(+), 62 deletions(-)

diff --git a/bigquery/arrow.go b/bigquery/arrow.go
index 0d9efeff8c9d..825256dec0b4 100644
--- a/bigquery/arrow.go
+++ b/bigquery/arrow.go
@@ -272,6 +272,21 @@ func convertArrowValue(col arrow.Array, i int, ft arrow.DataType, fs *FieldSchem
 		arr := col.(*array.Struct)
 		nestedValues := []Value{}
 		fields := ft.(*arrow.StructType).Fields()
+		if fs.Type == RangeFieldType {
+			rangeFieldSchema := &FieldSchema{
+				Type: fs.RangeElementType.Type,
+			}
+			start, err := convertArrowValue(arr.Field(0), i, fields[0].Type, rangeFieldSchema)
+			if err != nil {
+				return nil, err
+			}
+			end, err := convertArrowValue(arr.Field(1), i, fields[1].Type, rangeFieldSchema)
+			if err != nil {
+				return nil, err
+			}
+			rangeValue := &RangeValue{Start: start, End: end}
+			return Value(rangeValue), nil
+		}
 		for fIndex, f := range fields {
 			v, err := convertArrowValue(arr.Field(fIndex), i, f.Type, fs.Schema[fIndex])
 			if err != nil {
diff --git a/bigquery/integration_test.go b/bigquery/integration_test.go
index 0af0cc77805d..95c3a5180470 100644
--- a/bigquery/integration_test.go
+++ b/bigquery/integration_test.go
@@ -2154,7 +2154,7 @@ var (
 	queryParameterTestCases = []queryParameterTestCase{}
 )
 
-func initQueryParameterTestCases(includeRangeCases bool) {
+func initQueryParameterTestCases() {
 	d := civil.Date{Year: 2016, Month: 3, Day: 20}
 	tm := civil.Time{Hour: 15, Minute: 04, Second: 05, Nanosecond: 3008}
 	rtm := tm
@@ -2276,6 +2276,46 @@ func initQueryParameterTestCases(includeRangeCases bool) {
 			[]Value{"{\"alpha\":\"beta\"}"},
 			"{\"alpha\":\"beta\"}",
 		},
+		{
+			"RangeUnboundedStart",
+			"SELECT @val",
+			[]QueryParameter{
+				{
+					Name: "val",
+					Value: &QueryParameterValue{
+						Type: StandardSQLDataType{
+							TypeKind: "RANGE",
+							RangeElementType: &StandardSQLDataType{
+								TypeKind: "TIMESTAMP",
+							},
+						},
+						Value: rangeTimestamp1,
+					},
+				},
+			},
+			[]Value{rangeTimestamp1},
+			rangeTimestamp1,
+		},
+		{
+			"RangeUnboundedEnd",
+			"SELECT @val",
+			[]QueryParameter{
+				{
+					Name: "val",
+					Value: &QueryParameterValue{
+						Type: StandardSQLDataType{
+							TypeKind: "RANGE",
+							RangeElementType: &StandardSQLDataType{
+								TypeKind: "TIMESTAMP",
+							},
+						},
+						Value: rangeTimestamp2,
+					},
+				},
+			},
+			[]Value{rangeTimestamp2},
+			rangeTimestamp2,
+		},
 		{
 			"NestedStructParam",
 			"SELECT @val",
@@ -2448,51 +2488,6 @@ func initQueryParameterTestCases(includeRangeCases bool) {
 			},
 		},
 	}
-
-	if includeRangeCases {
-		queryParameterTestCases = append(queryParameterTestCases, []queryParameterTestCase{
-			{
-				"RangeUnboundedEnd",
-				"SELECT @val",
-				[]QueryParameter{
-					{
-						Name: "val",
-						Value: &QueryParameterValue{
-							Type: StandardSQLDataType{
-								TypeKind: "RANGE",
-								RangeElementType: &StandardSQLDataType{
-									TypeKind: "TIMESTAMP",
-								},
-							},
-							Value: rangeTimestamp1,
-						},
-					},
-				},
-				[]Value{rangeTimestamp1},
-				rangeTimestamp1,
-			},
-			{
-				"RangeUnboundedStart",
-				"SELECT @val",
-				[]QueryParameter{
-					{
-						Name: "val",
-						Value: &QueryParameterValue{
-							Type: StandardSQLDataType{
-								TypeKind: "RANGE",
-								RangeElementType: &StandardSQLDataType{
-									TypeKind: "TIMESTAMP",
-								},
-							},
-							Value: rangeTimestamp2,
-						},
-					},
-				},
-				[]Value{rangeTimestamp2},
-				rangeTimestamp2,
-			},
-		}...)
-	}
 }
 
 func TestIntegration_QueryParameters(t *testing.T) {
@@ -2501,7 +2496,7 @@ func TestIntegration_QueryParameters(t *testing.T) {
 	}
 	ctx := context.Background()
 
-	initQueryParameterTestCases(true)
+	initQueryParameterTestCases()
 
 	for _, tc := range queryParameterTestCases {
 		t.Run(tc.name, func(t *testing.T) {
diff --git a/bigquery/storage_integration_test.go b/bigquery/storage_integration_test.go
index 6f275927c6bc..81d68aee8318 100644
--- a/bigquery/storage_integration_test.go
+++ b/bigquery/storage_integration_test.go
@@ -37,23 +37,25 @@ func TestIntegration_StorageReadBasicTypes(t *testing.T) {
 	}
 	ctx := context.Background()
 
-	initQueryParameterTestCases(false)
+	initQueryParameterTestCases()
 
 	for _, c := range queryParameterTestCases {
-		q := storageOptimizedClient.Query(c.query)
-		q.Parameters = c.parameters
-		q.forceStorageAPI = true
-		it, err := q.Read(ctx)
-		if err != nil {
-			t.Fatal(err)
-		}
-		err = checkIteratorRead(it, c.wantRow)
-		if err != nil {
-			t.Fatalf("%s: error on query `%s`[%v]: %v", it.SourceJob().ID(), c.query, c.parameters, err)
-		}
-		if !it.IsAccelerated() {
-			t.Fatalf("%s: expected storage api to be used", it.SourceJob().ID())
-		}
+		t.Run(c.name, func(t *testing.T) {
+			q := storageOptimizedClient.Query(c.query)
+			q.Parameters = c.parameters
+			q.forceStorageAPI = true
+			it, err := q.Read(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+			err = checkIteratorRead(it, c.wantRow)
+			if err != nil {
+				t.Fatalf("%s: error on query `%s`[%v]: %v", it.SourceJob().ID(), c.query, c.parameters, err)
+			}
+			if !it.IsAccelerated() {
+				t.Fatalf("%s: expected storage api to be used", it.SourceJob().ID())
+			}
+		})
 	}
 }