Repository: snowflakedb/gosnowflake
Pull request #1137 (2024-05-21 23:43:03)

[bug fix] objectType should work w/ or wo/ ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE 
### Description

If we get a snowflake object type, but we haven't set the flag `ENABLE_STRUCTURED_TYPES_IN_CLIENT_RESPONSE`, the code will panic because we are only setting the record on the batch if the snowflake type matches the expected arrow type. This will add support for the underlying type to be either structured or string for snowflake objectType, and if its neither, we will return an error with the underlying type rather than panic so we can fix it

### Checklist
- [X] Created tests which fail without the change (if possible)
   - `TestStructuredTypeInArrowBatchesWithoutEnablingStructureTypes`: check that we can have an object type without setting structured types flag
   - `TestNullObject`: ensure that we can have a null object. (This is what was failing in my test suite; ultimately this isn't the actual trigger of the panic, but its still good to test because it could be one day)
- [NA] Extended the README / documentation, if necessary



Patch:
diff --git a/converter.go b/converter.go
index 65c0343..7d1d502 100644
--- a/converter.go
+++ b/converter.go
@@ -2084,25 +2084,8 @@ func arrowToRecordSingleColumn(ctx context.Context, field arrow.Field, col arrow
 			newCol = tb.NewArray()
 		}
 	case textType:
-		if arrowBatchesUtf8ValidationEnabled(ctx) && col.DataType().ID() == arrow.STRING {
-			tb := array.NewStringBuilder(pool)
-			defer tb.Release()
-
-			for i := 0; i < int(numRows); i++ {
-				if col.(*array.String).IsValid(i) {
-					stringValue := col.(*array.String).Value(i)
-					if !utf8.ValidString(stringValue) {
-						logger.WithContext(ctx).Error("Invalid UTF-8 characters detected while reading query response, column: ", fieldMetadata.Name)
-						stringValue = strings.ToValidUTF8(stringValue, "�")
-					}
-					tb.Append(stringValue)
-				} else {
-					tb.AppendNull()
-				}
-			}
-			newCol = tb.NewArray()
-		} else {
-			col.Retain()
+		if stringCol, ok := col.(*array.String); ok {
+			newCol = arrowStringRecordToColumn(ctx, stringCol, pool, numRows, fieldMetadata)
 		}
 	case objectType:
 		if structCol, ok := col.(*array.Struct); ok {
@@ -2123,6 +2106,8 @@ func arrowToRecordSingleColumn(ctx context.Context, field arrow.Field, col arrow
 			nullBitmap := memory.NewBufferBytes(structCol.NullBitmapBytes())
 			numberOfNulls := structCol.NullN()
 			return array.NewStructArrayWithNulls(internalCols, fieldNames, nullBitmap, numberOfNulls, 0)
+		} else if stringCol, ok := col.(*array.String); ok {
+			newCol = arrowStringRecordToColumn(ctx, stringCol, pool, numRows, fieldMetadata)
 		}
 	case arrayType:
 		if listCol, ok := col.(*array.List); ok {
@@ -2134,34 +2119,71 @@ func arrowToRecordSingleColumn(ctx context.Context, field arrow.Field, col arrow
 			newData := array.NewData(arrow.ListOf(newCol.DataType()), listCol.Len(), listCol.Data().Buffers(), []arrow.ArrayData{newCol.Data()}, listCol.NullN(), 0)
 			defer newData.Release()
 			return array.NewListData(newData), nil
+		} else if stringCol, ok := col.(*array.String); ok {
+			newCol = arrowStringRecordToColumn(ctx, stringCol, pool, numRows, fieldMetadata)
 		}
 	case mapType:
-		mapCol := col.(*array.Map)
-		keyCol, err := arrowToRecordSingleColumn(ctx, field.Type.(*arrow.MapType).KeyField(), mapCol.Keys(), fieldMetadata.Fields[0], higherPrecisionEnabled, timestampOption, pool, loc, numRows)
-		if err != nil {
-			return nil, err
-		}
-		defer keyCol.Release()
-		valueCol, err := arrowToRecordSingleColumn(ctx, field.Type.(*arrow.MapType).ItemField(), mapCol.Items(), fieldMetadata.Fields[1], higherPrecisionEnabled, timestampOption, pool, loc, numRows)
-		if err != nil {
-			return nil, err
-		}
-		defer valueCol.Release()
+		if mapCol, ok := col.(*array.Map); ok {
+			keyCol, err := arrowToRecordSingleColumn(ctx, field.Type.(*arrow.MapType).KeyField(), mapCol.Keys(), fieldMetadata.Fields[0], higherPrecisionEnabled, timestampOption, pool, loc, numRows)
+			if err != nil {
+				return nil, err
+			}
+			defer keyCol.Release()
+			valueCol, err := arrowToRecordSingleColumn(ctx, field.Type.(*arrow.MapType).ItemField(), mapCol.Items(), fieldMetadata.Fields[1], higherPrecisionEnabled, timestampOption, pool, loc, numRows)
+			if err != nil {
+				return nil, err
+			}
+			defer valueCol.Release()
 
-		structArr, err := array.NewStructArray([]arrow.Array{keyCol, valueCol}, []string{"k", "v"})
-		if err != nil {
-			return nil, err
+			structArr, err := array.NewStructArray([]arrow.Array{keyCol, valueCol}, []string{"k", "v"})
+			if err != nil {
+				return nil, err
+			}
+			defer structArr.Release()
+			newData := array.NewData(arrow.MapOf(keyCol.DataType(), valueCol.DataType()), mapCol.Len(), mapCol.Data().Buffers(), []arrow.ArrayData{structArr.Data()}, mapCol.NullN(), 0)
+			defer newData.Release()
+			return array.NewMapData(newData), nil
+		} else if stringCol, ok := col.(*array.String); ok {
+			newCol = arrowStringRecordToColumn(ctx, stringCol, pool, numRows, fieldMetadata)
 		}
-		defer structArr.Release()
-		newData := array.NewData(arrow.MapOf(keyCol.DataType(), valueCol.DataType()), mapCol.Len(), mapCol.Data().Buffers(), []arrow.ArrayData{structArr.Data()}, mapCol.NullN(), 0)
-		defer newData.Release()
-		return array.NewMapData(newData), nil
 	default:
 		col.Retain()
 	}
 	return newCol, nil
 }
 
+// returns n arrow array which will be new and populated if we converted the array to valid utf8
+// or if we didn't covnert it, it will return the original column.
+func arrowStringRecordToColumn(
+	ctx context.Context,
+	stringCol *array.String,
+	mem memory.Allocator,
+	numRows int64,
+	fieldMetadata fieldMetadata,
+) arrow.Array {
+	if arrowBatchesUtf8ValidationEnabled(ctx) && stringCol.DataType().ID() == arrow.STRING {
+		tb := array.NewStringBuilder(mem)
+		defer tb.Release()
+
+		for i := 0; i < int(numRows); i++ {
+			if stringCol.IsValid(i) {
+				stringValue := stringCol.Value(i)
+				if !utf8.ValidString(stringValue) {
+					logger.WithContext(ctx).Error("Invalid UTF-8 characters detected while reading query response, column: ", fieldMetadata.Name)
+					stringValue = strings.ToValidUTF8(stringValue, "�")
+				}
+				tb.Append(stringValue)
+			} else {
+				tb.AppendNull()
+			}
+		}
+		arr := tb.NewArray()
+		return arr
+	}
+	stringCol.Retain()
+	return stringCol
+}
+
 func recordToSchema(sc *arrow.Schema, rowType []execResponseRowType, loc *time.Location, timestampOption snowflakeArrowBatchesTimestampOption, withHigherPrecision bool) (*arrow.Schema, error) {
 	fields := recordToSchemaRecursive(sc.Fields(), rowType, loc, timestampOption, withHigherPrecision)
 	meta := sc.Metadata()


Test patch:
diff --git a/structured_type_read_test.go b/structured_type_read_test.go
index 4963502..0003610 100644
--- a/structured_type_read_test.go
+++ b/structured_type_read_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"database/sql/driver"
+	"fmt"
 	"github.com/apache/arrow/go/v15/arrow"
 	"github.com/apache/arrow/go/v15/arrow/array"
 	"github.com/apache/arrow/go/v15/arrow/memory"
@@ -1621,3 +1622,143 @@ func forAllStructureTypeFormats(dbt *DBTest, f func(t *testing.T, format string)
 		})
 	}
 }
+
+func TestSelectingNullObjectsInArrowBatches(t *testing.T) {
+	testcases := []string{
+		"select null::object(v VARCHAR)",
+		"select null::object",
+	}
+	skipStructuredTypesTestsOnGHActions(t)
+	runDBTest(t, func(dbt *DBTest) {
+		for _, tc := range testcases {
+			t.Run(tc, func(t *testing.T) {
+				pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
+				defer pool.AssertSize(t, 0)
+				ctx := WithArrowBatches(WithArrowAllocator(context.Background(), pool))
+
+				var err error
+				var rows driver.Rows
+				err = dbt.conn.Raw(func(sc any) error {
+					queryer, implementsQueryContext := sc.(driver.QueryerContext)
+					assertTrueF(t, implementsQueryContext, "snowflake conn does not implement QueryerContext but needs to")
+					rows, err = queryer.QueryContext(ctx, tc, nil)
+					return err
+				})
+				assertNilF(t, err, fmt.Sprintf("test failed to run the following query: %s", tc))
+				defer rows.Close()
+
+				sfRows, isSfRows := rows.(SnowflakeRows)
+				assertTrueF(t, isSfRows, "rows are not snowflakeRows")
+
+				batches, err := sfRows.GetArrowBatches()
+				assertNilF(t, err)
+				assertNotEqualF(t, len(batches), 0)
+				batch, err := batches[0].Fetch()
+				assertNilF(t, err)
+				assertNotEqualF(t, len(*batch), 0)
+				for _, record := range *batch {
+					// check number of cols/rows so we dont get index out of range
+					assertEqualF(t, record.NumRows(), int64(1), "wrong number of rows")
+					assertEqualF(t, record.NumCols(), int64(1), "wrong number of cols")
+
+					colIndex := 0
+					rowIndex := 0
+
+					srtCol, isStrCol := record.Column(colIndex).(*array.String)
+					assertTrueF(t, isStrCol, "column is not type string")
+
+					assertEqualE(t, srtCol.Value(rowIndex), "")
+					isNull := srtCol.IsNull(rowIndex)
+					assertTrueF(t, isNull)
+					record.Release()
+				}
+
+			})
+		}
+	})
+}
+
+func TestSelectingSemistructuredTypesInArrowBatches(t *testing.T) {
+	testcases := []struct {
+		name               string
+		query              string
+		expected           string
+		withUtf8Validation bool
+	}{
+		{
+			name:               "test semistructured object in arrow batch with utf8 validation, snowflakeType = objectType",
+			withUtf8Validation: true,
+			expected:           `{"s":"someString"}`,
+			query:              "SELECT {'s':'someString'}::OBJECT",
+		},
+		{
+			name:               "test semistructured object in arrow batch without utf8 validation, snowflakeType = objectType",
+			withUtf8Validation: false,
+			expected:           `{"s":"someString"}`,
+			query:              "SELECT {'s':'someString'}::OBJECT",
+		},
+		{
+			name:               "test semistructured object in arrow batch without utf8 validation, snowflakeType = arrayType",
+			withUtf8Validation: false,
+			expected:           `[1,2,3]`,
+			query:              "SELECT [1, 2, 3]::ARRAY",
+		},
+		{
+			name:               "test semistructured object in arrow batch with utf8 validation, snowflakeType = arrayType",
+			withUtf8Validation: true,
+			expected:           `[1,2,3]`,
+			query:              "SELECT [1, 2, 3]::ARRAY",
+		},
+	}
+	skipStructuredTypesTestsOnGHActions(t)
+	runDBTest(t, func(dbt *DBTest) {
+		for _, tc := range testcases {
+			t.Run(tc.name, func(t *testing.T) {
+
+				pool := memory.NewCheckedAllocator(memory.DefaultAllocator)
+				defer pool.AssertSize(t, 0)
+				ctx := WithArrowBatches(WithArrowAllocator(context.Background(), pool))
+				if tc.withUtf8Validation {
+					ctx = WithArrowBatchesUtf8Validation(ctx)
+				}
+
+				var err error
+				var rows driver.Rows
+				err = dbt.conn.Raw(func(sc any) error {
+					queryer, implementsQueryContext := sc.(driver.QueryerContext)
+					assertTrueF(t, implementsQueryContext, "snowflake conn does not implement QueryerContext but needs to")
+					rows, err = queryer.QueryContext(ctx, tc.query, nil)
+					return err
+				})
+				assertNilF(t, err)
+				defer rows.Close()
+
+				sfRows, isSfRows := rows.(SnowflakeRows)
+				assertTrueF(t, isSfRows, "rows are not snowflakeRows")
+
+				batches, err := sfRows.GetArrowBatches()
+				assertNilF(t, err)
+				assertNotEqualF(t, len(batches), 0)
+				batch, err := batches[0].Fetch()
+				assertNilF(t, err)
+				assertNotEqualF(t, len(*batch), 0)
+				for _, record := range *batch {
+					assertEqualF(t, record.NumCols(), int64(1), "unexpected number of columns")
+					assertEqualF(t, record.NumRows(), int64(1), "unexpected number of rows")
+
+					curColIndex := 0
+					rowIndex := 0
+
+					// The underlying data may be struct, or it could be string. Either way is ok but lets not fail
+					stringCol, isString := record.Column(curColIndex).(*array.String)
+
+					assertTrueF(t, isString, "wrong type for column, expected string")
+					assertEqualIgnoringWhitespaceE(t, stringCol.Value(rowIndex), tc.expected)
+
+					record.Release()
+				}
+			})
+		}
+
+	})
+}