Skip to content

Commit 66b7905

Browse files
mattdeekaymattdeekay
committed
Fix arrowRows_test.go
Signed-off-by: Matthew Kim <[email protected]>
1 parent fc9d9cb commit 66b7905

File tree

1 file changed

+46
-20
lines changed

1 file changed

+46
-20
lines changed

‎internal/rows/arrowbased/arrowRows_test.go‎

Lines changed: 46 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,7 @@ func TestArrowRowScanner(t *testing.T) {
208208

209209
rowSet := &cli_service.TRowSet{}
210210
schema := &cli_service.TTableSchema{}
211-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
211+
metadataResp := getMetadataResp(schema)
212212

213213
ars, err := NewArrowRowScanner(metadataResp, rowSet, nil, nil, context.Background())
214214
assert.NotNil(t, ars)
@@ -232,7 +232,7 @@ func TestArrowRowScanner(t *testing.T) {
232232
// holders
233233
rowSet := &cli_service.TRowSet{ArrowBatches: []*cli_service.TSparkArrowBatch{{RowCount: 2}, {RowCount: 3}}}
234234
schema := getAllTypesSchema()
235-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
235+
metadataResp := getMetadataResp(schema)
236236

237237
d, _ := NewArrowRowScanner(metadataResp, rowSet, nil, nil, context.Background())
238238

@@ -300,7 +300,7 @@ func TestArrowRowScanner(t *testing.T) {
300300
// from the default if when using native types is specified
301301
rowSet := &cli_service.TRowSet{ArrowBatches: []*cli_service.TSparkArrowBatch{{RowCount: 2}, {RowCount: 3}}}
302302
schema := getAllTypesSchema()
303-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
303+
metadataResp := getMetadataResp(schema)
304304

305305
cfg := config.Config{}
306306
cfg.UseArrowBatches = true
@@ -328,7 +328,7 @@ func TestArrowRowScanner(t *testing.T) {
328328
t.Run("Fail creating arrow row scanner on invalid native decimal type", func(t *testing.T) {
329329
rowSet := &cli_service.TRowSet{ArrowBatches: []*cli_service.TSparkArrowBatch{{RowCount: 2}, {RowCount: 3}}}
330330
schema := getAllTypesSchema()
331-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
331+
metadataResp := getMetadataResp(schema)
332332

333333
cfg := config.Config{}
334334
cfg.UseArrowBatches = true
@@ -401,7 +401,7 @@ func TestArrowRowScanner(t *testing.T) {
401401
t.Run("Fail to scan row when no batches are present", func(t *testing.T) {
402402
rowSet := &cli_service.TRowSet{}
403403
schema := getAllTypesSchema()
404-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
404+
metadataResp := getMetadataResp(schema)
405405

406406
cfg := config.Config{}
407407
cfg.UseArrowBatches = true
@@ -430,7 +430,7 @@ func TestArrowRowScanner(t *testing.T) {
430430
// The following is the code under test
431431
rowSet := &cli_service.TRowSet{ArrowBatches: []*cli_service.TSparkArrowBatch{{RowCount: 2}, {RowCount: 3}}}
432432
schema := getAllTypesSchema()
433-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
433+
metadataResp := getMetadataResp(schema)
434434

435435
cfg := config.Config{}
436436
cfg.UseArrowBatches = true
@@ -470,9 +470,12 @@ func TestArrowRowScanner(t *testing.T) {
470470
},
471471
}
472472
schema := getAllTypesSchema()
473-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
473+
metadataResp := getMetadataResp(schema)
474474

475-
d, _ := NewArrowRowScanner(metadataResp, rowSet, nil, nil, context.Background())
475+
cfg := config.Config{}
476+
cfg.UseLz4Compression = false
477+
478+
d, _ := NewArrowRowScanner(metadataResp, rowSet, &cfg, nil, context.Background())
476479

477480
var ars *arrowRowScanner = d.(*arrowRowScanner)
478481

@@ -513,9 +516,12 @@ func TestArrowRowScanner(t *testing.T) {
513516
},
514517
}
515518
schema := getAllTypesSchema()
516-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
519+
metadataResp := getMetadataResp(schema)
520+
521+
cfg := config.Config{}
522+
cfg.UseLz4Compression = false
517523

518-
d, _ := NewArrowRowScanner(metadataResp, rowSet, nil, nil, nil)
524+
d, _ := NewArrowRowScanner(metadataResp, rowSet, &cfg, nil, nil)
519525

520526
var ars *arrowRowScanner = d.(*arrowRowScanner)
521527

@@ -543,9 +549,12 @@ func TestArrowRowScanner(t *testing.T) {
543549
},
544550
}
545551
schema := getAllTypesSchema()
546-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
552+
metadataResp := getMetadataResp(schema)
553+
554+
cfg := config.Config{}
555+
cfg.UseLz4Compression = false
547556

548-
d, _ := NewArrowRowScanner(metadataResp, rowSet, nil, nil, nil)
557+
d, _ := NewArrowRowScanner(metadataResp, rowSet, &cfg, nil, nil)
549558

550559
var ars *arrowRowScanner = d.(*arrowRowScanner)
551560

@@ -574,9 +583,12 @@ func TestArrowRowScanner(t *testing.T) {
574583
}
575584

576585
schema := getAllTypesSchema()
577-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
586+
metadataResp := getMetadataResp(schema)
587+
588+
cfg := config.Config{}
589+
cfg.UseLz4Compression = false
578590

579-
d, _ := NewArrowRowScanner(metadataResp, rowSet, nil, nil, nil)
591+
d, _ := NewArrowRowScanner(metadataResp, rowSet, &cfg, nil, nil)
580592

581593
var ars *arrowRowScanner = d.(*arrowRowScanner)
582594
ars.recordReader = fakeRecordReader{fnNewRecord: func(arrowSchemaBytes []byte, sparkArrowBatch sparkArrowBatch) (arrow.Record, dbsqlerr.DBError) {
@@ -602,9 +614,12 @@ func TestArrowRowScanner(t *testing.T) {
602614
}
603615

604616
schema := getAllTypesSchema()
605-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
617+
metadataResp := getMetadataResp(schema)
606618

607-
d, _ := NewArrowRowScanner(metadataResp, rowSet, nil, nil, nil)
619+
cfg := config.Config{}
620+
cfg.UseLz4Compression = false
621+
622+
d, _ := NewArrowRowScanner(metadataResp, rowSet, &cfg, nil, nil)
608623

609624
var ars *arrowRowScanner = d.(*arrowRowScanner)
610625
ars.recordReader = fakeRecordReader{fnNewRecord: func(arrowSchemaBytes []byte, sparkArrowBatch sparkArrowBatch) (arrow.Record, dbsqlerr.DBError) {
@@ -626,9 +641,12 @@ func TestArrowRowScanner(t *testing.T) {
626641
},
627642
}
628643
schema := getAllTypesSchema()
629-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
644+
metadataResp := getMetadataResp(schema)
630645

631-
d, _ := NewArrowRowScanner(metadataResp, rowSet, nil, nil, context.Background())
646+
cfg := config.Config{}
647+
cfg.UseLz4Compression = false
648+
649+
d, _ := NewArrowRowScanner(metadataResp, rowSet, &cfg, nil, context.Background())
632650

633651
var ars *arrowRowScanner = d.(*arrowRowScanner)
634652

@@ -769,9 +787,12 @@ func TestArrowRowScanner(t *testing.T) {
769787
schema := &cli_service.TTableSchema{
770788
Columns: []*cli_service.TColumnDesc{columns[i]},
771789
}
772-
metadataResp := &cli_service.TGetResultSetMetadataResp{Schema: schema}
790+
metadataResp := getMetadataResp(schema)
791+
792+
cfg := config.Config{}
793+
cfg.UseLz4Compression = false
773794

774-
d, _ := NewArrowRowScanner(metadataResp, rowSet, nil, nil, context.Background())
795+
d, _ := NewArrowRowScanner(metadataResp, rowSet, &cfg, nil, context.Background())
775796

776797
var ars *arrowRowScanner = d.(*arrowRowScanner)
777798
ars.UseArrowNativeComplexTypes = true
@@ -1783,3 +1804,8 @@ func loadTestData(t *testing.T, name string, v any) {
17831804
}
17841805
}
17851806
}
1807+
1808+
func getMetadataResp(schema *cli_service.TTableSchema) *cli_service.TGetResultSetMetadataResp {
1809+
rowSetType := cli_service.TSparkRowSetType_ARROW_BASED_SET
1810+
return &cli_service.TGetResultSetMetadataResp{Schema: schema, ResultFormat: &rowSetType}
1811+
}

0 commit comments

Comments
 (0)