diff --git a/arrow/array/table.go b/arrow/array/table.go index 540124b1f..6bd440455 100644 --- a/arrow/array/table.go +++ b/arrow/array/table.go @@ -149,10 +149,15 @@ func NewTableFromSlice(schema *arrow.Schema, data [][]arrow.Array) arrow.Table { chunked.Release() } + var rows int64 + if len(cols) > 0 { + rows = int64(cols[0].Len()) + } + tbl := simpleTable{ schema: schema, cols: cols, - rows: int64(cols[0].Len()), + rows: rows, } tbl.refCount.Add(1) @@ -178,6 +183,13 @@ func NewTableFromSlice(schema *arrow.Schema, data [][]arrow.Array) arrow.Table { func NewTableFromRecords(schema *arrow.Schema, recs []arrow.RecordBatch) arrow.Table { arrs := make([]arrow.Array, len(recs)) cols := make([]arrow.Column, schema.NumFields()) + rows := int64(-1) + if len(cols) == 0 { + rows = 0 + for _, rec := range recs { + rows += rec.NumRows() + } + } defer func(cols []arrow.Column) { for i := range cols { @@ -195,7 +207,7 @@ func NewTableFromRecords(schema *arrow.Schema, recs []arrow.RecordBatch) arrow.T chunk.Release() } - return NewTable(schema, cols, -1) + return NewTable(schema, cols, rows) } func (tbl *simpleTable) Schema() *arrow.Schema { return tbl.schema } diff --git a/arrow/array/table_test.go b/arrow/array/table_test.go index bbf4a4e66..110b1c7c5 100644 --- a/arrow/array/table_test.go +++ b/arrow/array/table_test.go @@ -113,6 +113,36 @@ func TestChunked(t *testing.T) { } } +func TestTableFromSliceWithoutColumns(t *testing.T) { + schema := arrow.NewSchema(nil, nil) + tbl := array.NewTableFromSlice(schema, nil) + defer tbl.Release() + + if got, want := tbl.NumRows(), int64(0); got != want { + t.Fatalf("NumRows = %d, want %d", got, want) + } +} + +func TestTableFromRecordsWithoutColumns(t *testing.T) { + schema := arrow.NewSchema(nil, nil) + records := []arrow.RecordBatch{ + array.NewRecordBatch(schema, nil, 2), + array.NewRecordBatch(schema, nil, 3), + } + defer func() { + for _, rec := range records { + rec.Release() + } + }() + + tbl := array.NewTableFromRecords(schema, records) + defer tbl.Release() + + if got, want := tbl.NumRows(), int64(5); got != want { + t.Fatalf("NumRows = %d, want %d", got, want) + } +} + func TestChunkedEqualDataType(t *testing.T) { mem := memory.NewCheckedAllocator(memory.NewGoAllocator()) defer mem.AssertSize(t, 0)