From 355beb56f91b2c6bd8cc49d9fce1826e85f72144 Mon Sep 17 00:00:00 2001
From: Nathan <148575555+nathan-artie@users.noreply.github.com>
Date: Fri, 14 Jun 2024 12:29:41 -0700
Subject: [PATCH] [bigquery] Support additional types for float columns (#723)

---
 clients/bigquery/storagewrite.go      | 13 ++++++++++++-
 clients/bigquery/storagewrite_test.go |  9 +++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/clients/bigquery/storagewrite.go b/clients/bigquery/storagewrite.go
index 2aafa9adb..3cd550bde 100644
--- a/clients/bigquery/storagewrite.go
+++ b/clients/bigquery/storagewrite.go
@@ -2,6 +2,7 @@ package bigquery
 
 import (
 	"fmt"
+	"strconv"
 	"time"
 
 	"cloud.google.com/go/bigquery/storage/apiv1/storagepb"
@@ -132,8 +133,18 @@ func rowToMessage(row map[string]any, columns []columns.Column, messageDescripto
 				message.Set(field, protoreflect.ValueOfFloat64(float64(value)))
 			case float64:
 				message.Set(field, protoreflect.ValueOfFloat64(value))
+			case int32:
+				message.Set(field, protoreflect.ValueOfFloat64(float64(value)))
+			case int64:
+				message.Set(field, protoreflect.ValueOfFloat64(float64(value)))
+			case string:
+				floatValue, err := strconv.ParseFloat(value, 64)
+				if err != nil {
+					return nil, fmt.Errorf("failed to parse string to float64: %w", err)
+				}
+				message.Set(field, protoreflect.ValueOfFloat64(floatValue))
 			default:
-				return nil, fmt.Errorf("expected float32/float64 recieved %T with value %v", value, value)
+				return nil, fmt.Errorf("expected float32/float64/int32/int64/string recieved %T with value %v", value, value)
 			}
 		case typing.EDecimal.Kind:
 			if decimalValue, ok := value.(*decimal.Decimal); ok {
diff --git a/clients/bigquery/storagewrite_test.go b/clients/bigquery/storagewrite_test.go
index 53c83a268..ceaf26168 100644
--- a/clients/bigquery/storagewrite_test.go
+++ b/clients/bigquery/storagewrite_test.go
@@ -106,6 +106,9 @@ func TestRowToMessage(t *testing.T) {
 		columns.NewColumn("c_int64", typing.Integer),
 		columns.NewColumn("c_float32", typing.Float),
 		columns.NewColumn("c_float64", typing.Float),
+		columns.NewColumn("c_float_int32", typing.Float),
+		columns.NewColumn("c_float_int64", typing.Float),
+		columns.NewColumn("c_float_string", typing.Float),
 		columns.NewColumn("c_numeric", typing.EDecimal),
 		columns.NewColumn("c_string", typing.String),
 		columns.NewColumn("c_string_decimal", typing.String),
@@ -123,6 +126,9 @@ func TestRowToMessage(t *testing.T) {
 		"c_int64":          int64(1234),
 		"c_float32":        float32(1234.567),
 		"c_float64":        float64(1234.567),
+		"c_float_int32":    int32(1234),
+		"c_float_int64":    int64(1234),
+		"c_float_string":   "4444.55555",
 		"c_numeric":        decimal.NewDecimal(nil, 5, big.NewFloat(3.1415926)),
 		"c_string":         "foo bar",
 		"c_string_decimal": decimal.NewDecimal(nil, 5, big.NewFloat(1.618033)),
@@ -149,6 +155,9 @@ func TestRowToMessage(t *testing.T) {
 		"cBool":          true,
 		"cFloat32":       1234.5670166015625,
 		"cFloat64":       1234.567,
+		"cFloatInt32":    1234.0,
+		"cFloatInt64":    1234.0,
+		"cFloatString":   4444.55555,
 		"cInt":           "1234",
 		"cInt32":         "1234",
 		"cInt64":         "1234",