SNOW-988793 Do not read arrow chunk when it is empty for arrow batches

assert_test.go

@@ -34,6 +34,10 @@ func assertTrueF(t *testing.T, actual bool, descriptions ...string) {
 	fatalOnNonEmpty(t, validateEqual(actual, true, descriptions...))
 }
 
+func assertFalseF(t *testing.T, actual bool, descriptions ...string) {
+	fatalOnNonEmpty(t, validateEqual(actual, false, descriptions...))
+}
+
 func assertStringContainsE(t *testing.T, actual string, expectedToContain string, descriptions ...string) {
 	errorOnNonEmpty(t, validateStringContains(actual, expectedToContain, descriptions...))
 }

chunk_downloader.go

@@ -271,6 +271,9 @@ func getChunk(
 }
 
 func (scd *snowflakeChunkDownloader) startArrowBatches() error {
+	if scd.RowSet.RowSetBase64 == "" {
+		return nil
+	}
 	var err error
 	chunkMetaLen := len(scd.ChunkMetas)
 	var loc *time.Location

chunk_test.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"database/sql/driver"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"math/rand"
@@ -462,6 +463,101 @@ func TestWithArrowBatches(t *testing.T) {
 	})
 }
 
+func TestWithArrowBatchesWhenQueryReturnsNoRowsWhenUsingNativeGoSQLInterface(t *testing.T) {
+	runDBTest(t, func(dbt *DBTest) {
+		ctx := context.Background()
+
+		var rows driver.Rows
+		var err error
+		err = dbt.conn.Raw(func(x interface{}) error {
+			queryer, implementsQueryContext := x.(driver.QueryerContext)
+			if !implementsQueryContext {
+				return errors.New("gosnowflake driver does not implement queryerContext")
+			}
+			rows, err = queryer.QueryContext(WithArrowBatches(ctx), "SELECT 1 WHERE 0 = 1", nil)
+			return err
+		})
+		assertNilF(t, err)
+		defer rows.Close()
+	})
+}
+
+func TestWithArrowBatchesWhenQueryReturnsNoRows(t *testing.T) {
+	runDBTest(t, func(dbt *DBTest) {
+		rows := dbt.mustQueryContext(WithArrowBatches(context.Background()), "SELECT 1")
+		defer rows.Close()
+		assertFalseF(t, rows.Next())
+	})
+}
+
+func TestWithArrowBatchesWhenQueryReturnsSomeRowsInGivenFormatUsingNativeGoSQLInterface(t *testing.T) {
+	for _, tc := range []struct {
+		useJson bool
+		desc    string
+	}{
+		{
+			useJson: true,
+			desc:    "json",
+		},
+		{
+			useJson: false,
+			desc:    "arrow",
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			runDBTest(t, func(dbt *DBTest) {
+				if tc.useJson {
+					dbt.mustExec(forceJSON)
+				}
+				ctx := context.Background()
+
+				var rows driver.Rows
+				var err error
+				err = dbt.conn.Raw(func(x interface{}) error {
+					queryer, implementsQueryContext := x.(driver.QueryerContext)
+					if !implementsQueryContext {
+						return errors.New("gosnowflake driver does not implement queryerContext")
+					}
+					rows, err = queryer.QueryContext(WithArrowBatches(ctx), "SELECT 1", nil)
+					return err
+				})
+				assertNilF(t, err)
+				defer rows.Close()
+				values := make([]driver.Value, 1)
+				rows.Next(values)
+				assertEqualE(t, values[0], nil)
+			})
+		})
+	}
+}
+
+func TestWithArrowBatchesWhenQueryReturnsSomeRowsInGivenFormat(t *testing.T) {
+	for _, tc := range []struct {
+		useJson bool
+		desc    string
+	}{
+		{
+			useJson: true,
+			desc:    "json",
+		},
+		{
+			useJson: false,
+			desc:    "arrow",
+		},
+	} {
+		t.Run(tc.desc, func(t *testing.T) {
+			runDBTest(t, func(dbt *DBTest) {
+				if tc.useJson {
+					dbt.mustExec(forceJSON)
+				}
+				rows := dbt.mustQueryContext(WithArrowBatches(context.Background()), "SELECT 1")
+				defer rows.Close()
+				assertFalseF(t, rows.Next())
+			})
+		})
+	}
+}
+
 func TestWithArrowBatchesAsync(t *testing.T) {
 	runSnowflakeConnTest(t, func(sct *SCTest) {
 		ctx := WithAsyncMode(sct.sc.ctx)