Bug fix for query that doesn't produce rows (#136)

When executing a statement that doesn't produce rows (ex. CREATE TABLE
....) using DB.Query() instead of DB.Exec() the driver would panic. When
DB.Query() is used we try to create a row object. In this case there was
a mismatch in the schema metadata returned by thrift and the arrow
schema metadata. The thrift schema was indicating a single column named
'Result' while the arrow schema metadata was indicating no columns.
Moved the code creating the arrow schema and column info into separate
functions. Updated creation of the column info to handle mismatch
between the arrow and thrift schemas.
Signed-off-by: Raymond Cypher <[email protected]>

Author: rcypher-databricks

internal/rows/arrowbased/arrowRows.go

@@ -103,41 +103,13 @@ func NewArrowRowScanner(resultSetMetadata *cli_service.TGetResultSetMetadataResp
 		arrowConfig = cfg.ArrowConfig
 	}
 
-	var arrowSchema *arrow.Schema
-	schemaBytes := resultSetMetadata.ArrowSchema
-	if schemaBytes == nil {
-		var err error
-		// convert the TTableSchema to an arrow Schema
-		arrowSchema, err = tTableSchemaToArrowSchema(resultSetMetadata.Schema, &arrowConfig)
-		if err != nil {
-			logger.Err(err).Msg(errArrowRowsConvertSchema)
-			return nil, dbsqlerrint.NewDriverError(ctx, errArrowRowsConvertSchema, err)
-		}
-
-		// serialize the arrow schema
-		schemaBytes, err = getArrowSchemaBytes(arrowSchema, ctx)
-		if err != nil {
-			logger.Err(err).Msg(errArrowRowsSerializeSchema)
-			return nil, dbsqlerrint.NewDriverError(ctx, errArrowRowsSerializeSchema, err)
-		}
-	} else {
-		br := &chunkedByteReader{chunks: [][]byte{schemaBytes}}
-		rdr, err := ipc.NewReader(br)
-		if err != nil {
-			return nil, dbsqlerrint.NewDriverError(ctx, errArrowRowsUnableToReadBatch, err)
-		}
-		defer rdr.Release()
-
-		arrowSchema = rdr.Schema()
+	schemaBytes, arrowSchema, metadataErr := tGetResultSetMetadataRespToArrowSchema(resultSetMetadata, arrowConfig, ctx, logger)
+	if metadataErr != nil {
+		return nil, metadataErr
 	}
 
-	// get the database type names for each column
-	colInfos := make([]colInfo, len(resultSetMetadata.Schema.Columns))
-	for i := range resultSetMetadata.Schema.Columns {
-		col := resultSetMetadata.Schema.Columns[i]
-		field := arrowSchema.Field(i)
-		colInfos[i] = colInfo{name: field.Name, arrowType: field.Type, dbType: rowscanner.GetDBType(col)}
-	}
+	// Create column info
+	colInfos := getColumnInfo(arrowSchema, resultSetMetadata.Schema)
 
 	// get the function for converting arrow timestamps to a time.Time
 	// time values from the server are returned as UTC with microsecond precision
@@ -553,6 +525,61 @@ func tColumnDescToArrowField(columnDesc *cli_service.TColumnDesc, arrowConfig *c
 	return arrowField, nil
 }
 
+// Build a slice of columnInfo using the arrow schema and the thrift schema
+func getColumnInfo(arrowSchema *arrow.Schema, schema *cli_service.TTableSchema) []colInfo {
+	if arrowSchema == nil || schema == nil {
+		return []colInfo{}
+	}
+
+	nFields := len(arrowSchema.Fields())
+	if len(schema.Columns) < nFields {
+		nFields = len(schema.Columns)
+	}
+
+	colInfos := make([]colInfo, nFields)
+	for i := 0; i < nFields; i++ {
+		col := schema.Columns[i]
+		field := arrowSchema.Field(i)
+		colInfos[i] = colInfo{name: field.Name, arrowType: field.Type, dbType: rowscanner.GetDBType(col)}
+	}
+
+	return colInfos
+}
+
+// Derive an arrow.Schema object and the corresponding serialized bytes from TGetResultSetMetadataResp
+func tGetResultSetMetadataRespToArrowSchema(resultSetMetadata *cli_service.TGetResultSetMetadataResp, arrowConfig config.ArrowConfig, ctx context.Context, logger *dbsqllog.DBSQLLogger) ([]byte, *arrow.Schema, dbsqlerr.DBError) {
+
+	var arrowSchema *arrow.Schema
+	schemaBytes := resultSetMetadata.ArrowSchema
+	if schemaBytes == nil {
+		var err error
+		// convert the TTableSchema to an arrow Schema
+		arrowSchema, err = tTableSchemaToArrowSchema(resultSetMetadata.Schema, &arrowConfig)
+		if err != nil {
+			logger.Err(err).Msg(errArrowRowsConvertSchema)
+			return nil, nil, dbsqlerrint.NewDriverError(ctx, errArrowRowsConvertSchema, err)
+		}
+
+		// serialize the arrow schema
+		schemaBytes, err = getArrowSchemaBytes(arrowSchema, ctx)
+		if err != nil {
+			logger.Err(err).Msg(errArrowRowsSerializeSchema)
+			return nil, nil, dbsqlerrint.NewDriverError(ctx, errArrowRowsSerializeSchema, err)
+		}
+	} else {
+		br := &chunkedByteReader{chunks: [][]byte{schemaBytes}}
+		rdr, err := ipc.NewReader(br)
+		if err != nil {
+			return nil, nil, dbsqlerrint.NewDriverError(ctx, errArrowRowsUnableToReadBatch, err)
+		}
+		defer rdr.Release()
+
+		arrowSchema = rdr.Schema()
+	}
+
+	return schemaBytes, arrowSchema, nil
+}
+
 type sparkRecordReader struct{}
 
 // Make sure sparkRecordReader fulfills the recordReader interface

internal/rows/arrowbased/arrowRows_test.go

@@ -1349,6 +1349,75 @@ func TestArrowRowScanner(t *testing.T) {
 			assert.Equal(t, expected[i], dest[i])
 		}
 	})
+
+	t.Run("Mismatched schemas", func(t *testing.T) {
+		// Test for
+		var arrowSchema *arrow.Schema
+		var schema *cli_service.TTableSchema
+		colInfos := getColumnInfo(arrowSchema, schema)
+		assert.NotNil(t, colInfos)
+		assert.Zero(t, len(colInfos))
+
+		arrowSchema = &arrow.Schema{}
+		colInfos = getColumnInfo(arrowSchema, schema)
+		assert.NotNil(t, colInfos)
+		assert.Zero(t, len(colInfos))
+
+		arrowSchema = nil
+		schema = &cli_service.TTableSchema{}
+		colInfos = getColumnInfo(arrowSchema, schema)
+		assert.NotNil(t, colInfos)
+		assert.Zero(t, len(colInfos))
+
+		arrowSchema = &arrow.Schema{}
+		schema.Columns = []*cli_service.TColumnDesc{{ColumnName: "Result"}}
+		colInfos = getColumnInfo(arrowSchema, schema)
+		assert.NotNil(t, colInfos)
+		assert.Zero(t, len(colInfos))
+
+		schema.Columns = nil
+		arrowSchema = arrow.NewSchema([]arrow.Field{{Name: "Result", Type: arrow.PrimitiveTypes.Int16}}, nil)
+		colInfos = getColumnInfo(arrowSchema, schema)
+		assert.NotNil(t, colInfos)
+		assert.Zero(t, len(colInfos))
+
+		schema.Columns = []*cli_service.TColumnDesc{
+			{
+				ColumnName: "Result",
+				TypeDesc: &cli_service.TTypeDesc{
+					Types: []*cli_service.TTypeEntry{
+						{
+							PrimitiveEntry: &cli_service.TPrimitiveTypeEntry{
+								Type: cli_service.TTypeId_BOOLEAN_TYPE,
+							},
+						},
+					},
+				},
+			},
+			{
+				ColumnName: "Result2",
+			},
+		}
+		colInfos = getColumnInfo(arrowSchema, schema)
+		assert.NotNil(t, colInfos)
+		assert.Equal(t, 1, len(colInfos))
+		assert.Equal(t, "Result", colInfos[0].name)
+		assert.Equal(t, cli_service.TTypeId_BOOLEAN_TYPE, colInfos[0].dbType)
+		assert.Equal(t, arrow.PrimitiveTypes.Int16, colInfos[0].arrowType)
+
+		// results of executing query:
+		// "create or replace view hive_metastore.databricks_sql_go.test as select 1"
+		// using DB.Query() instead of DB.Exec()
+		executeStatementResp := cli_service.TExecuteStatementResp{}
+		loadTestData(t, "queryVExec.json", &executeStatementResp)
+		config := config.WithDefaults()
+		config.UseArrowNativeTimestamp = true
+		config.UseArrowNativeComplexTypes = true
+		config.UseArrowNativeDecimal = false
+		config.UseArrowNativeIntervalTypes = false
+		_, err := NewArrowRowScanner(executeStatementResp.DirectResults.ResultSetMetadata, executeStatementResp.DirectResults.ResultSet.Results, config, nil, context.Background())
+		assert.Nil(t, err)
+	})
 }
 
 type fakeColumnValues struct {

internal/rows/arrowbased/testdata/queryVExec.json

@@ -0,0 +1,72 @@
+{
+ "status": {
+  "statusCode": "SUCCESS_STATUS"
+ },
+ "operationHandle": {
+  "operationId": {
+   "guid": "Ae4RVsJdEj6Egpavp/LAJA==",
+   "secret": "M41SnYJyRuuEgstBlGaDnQ=="
+  },
+  "operationType": "EXECUTE_STATEMENT",
+  "hasResultSet": true
+ },
+ "directResults": {
+  "operationStatus": {
+   "status": {
+    "statusCode": "SUCCESS_STATUS"
+   },
+   "operationState": "FINISHED_STATE",
+   "operationStarted": 1687477482139,
+   "operationCompleted": 1687477484937
+  },
+  "resultSetMetadata": {
+   "status": {
+    "statusCode": "SUCCESS_STATUS"
+   },
+   "schema": {
+    "columns": [
+     {
+      "columnName": "Result",
+      "typeDesc": {
+       "types": [
+        {
+         "primitiveEntry": {
+          "type": "STRING_TYPE"
+         }
+        }
+       ]
+      },
+      "position": 1,
+      "comment": ""
+     }
+    ]
+   },
+   "resultFormat": "ARROW_BASED_SET",
+   "lz4Compressed": false,
+   "arrowSchema": "/////0AAAAAQAAAAAAAKAA4ABgANAAgACgAAAAAABAAQAAAAAAEKAAwAAAAIAAQACgAAAAgAAAAIAAAAAAAAAAAAAAAAAAAA",
+   "cacheLookupResult": "CACHE_INELIGIBLE",
+   "uncompressedBytes": 0,
+   "compressedBytes": 0
+  },
+  "resultSet": {
+   "status": {
+    "statusCode": "SUCCESS_STATUS"
+   },
+   "hasMoreRows": false,
+   "results": {
+    "startRowOffset": 0,
+    "rows": []
+   }
+  },
+  "closeOperation": {
+   "status": {
+    "statusCode": "SUCCESS_STATUS"
+   }
+  }
+ },
+ "executionRejected": false,
+ "maxClusterCapacity": 10,
+ "queryCost": 0,
+ "currentClusterLoad": 0,
+ "idempotencyType": "NON_IDEMPOTENT"
+}