ARROW-3625: [Go] add examples for Record and Table

Author: sbinet

go/arrow/array/table.go

@@ -246,6 +246,32 @@ func NewTable(schema *arrow.Schema, cols []Column, rows int64) *simpleTable {
 	return &tbl
 }
 
+// NewTableFromRecords returns a new basic, non-lazy in-memory table.
+//
+// NewTableFromRecords panics if the records and schema are inconsistent.
+func NewTableFromRecords(schema *arrow.Schema, recs []Record) *simpleTable {
+	arrs := make([]Interface, len(recs))
+	cols := make([]Column, len(schema.Fields()))
+
+	defer func(cols []Column) {
+		for i := range cols {
+			cols[i].Release()
+		}
+	}(cols)
+
+	for i := range cols {
+		field := schema.Field(i)
+		for j, rec := range recs {
+			arrs[j] = rec.Column(i)
+		}
+		chunk := NewChunked(field.Type, arrs)
+		cols[i] = *NewColumn(field, chunk)
+		chunk.Release()
+	}
+
+	return NewTable(schema, cols, -1)
+}
+
 func (tbl *simpleTable) Schema() *arrow.Schema { return tbl.schema }
 func (tbl *simpleTable) NumRows() int64        { return tbl.rows }
 func (tbl *simpleTable) NumCols() int64        { return int64(len(tbl.cols)) }

go/arrow/array/table_test.go

@@ -556,6 +556,52 @@ func TestTable(t *testing.T) {
 	}
 }
 
+func TestTableFromRecords(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+	defer mem.AssertSize(t, 0)
+
+	schema := arrow.NewSchema(
+		[]arrow.Field{
+			arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
+			arrow.Field{Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64},
+		},
+		nil,
+	)
+
+	b := array.NewRecordBuilder(mem, schema)
+	defer b.Release()
+
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{1, 2, 3, 4, 5, 6}, nil)
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{7, 8, 9, 10}, []bool{true, true, false, true})
+	b.Field(1).(*array.Float64Builder).AppendValues([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)
+
+	rec1 := b.NewRecord()
+	defer rec1.Release()
+
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{11, 12, 13, 14, 15, 16, 17, 18, 19, 20}, nil)
+	b.Field(1).(*array.Float64Builder).AppendValues([]float64{11, 12, 13, 14, 15, 16, 17, 18, 19, 20}, nil)
+
+	rec2 := b.NewRecord()
+	defer rec2.Release()
+
+	tbl := array.NewTableFromRecords(schema, []array.Record{rec1, rec2})
+	defer tbl.Release()
+
+	if got, want := tbl.Schema(), schema; !got.Equal(want) {
+		t.Fatalf("invalid schema: got=%#v, want=%#v", got, want)
+	}
+
+	if got, want := tbl.NumRows(), int64(20); got != want {
+		t.Fatalf("invalid number of rows: got=%d, want=%d", got, want)
+	}
+	if got, want := tbl.NumCols(), int64(2); got != want {
+		t.Fatalf("invalid number of columns: got=%d, want=%d", got, want)
+	}
+	if got, want := tbl.Column(0).Name(), schema.Field(0).Name; got != want {
+		t.Fatalf("invalid column: got=%q, want=%q", got, want)
+	}
+}
+
 func TestTableReader(t *testing.T) {
 	mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
 	defer mem.AssertSize(t, 0)

go/arrow/example_test.go

@@ -18,6 +18,7 @@ package arrow_test
 
 import (
 	"fmt"
+	"log"
 
 	"github.com/apache/arrow/go/arrow"
 	"github.com/apache/arrow/go/arrow/array"
@@ -409,3 +410,135 @@ func Example_float64Tensor2x5ColMajor() {
 	// arr[1 3] = 8
 	// arr[1 4] = 10
 }
+
+func Example_record() {
+	pool := memory.NewGoAllocator()
+
+	schema := arrow.NewSchema(
+		[]arrow.Field{
+			arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
+			arrow.Field{Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64},
+		},
+		nil,
+	)
+
+	b := array.NewRecordBuilder(pool, schema)
+	defer b.Release()
+
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{1, 2, 3, 4, 5, 6}, nil)
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{7, 8, 9, 10}, []bool{true, true, false, true})
+	b.Field(1).(*array.Float64Builder).AppendValues([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)
+
+	rec := b.NewRecord()
+	defer rec.Release()
+
+	for i, col := range rec.Columns() {
+		fmt.Printf("column[%d] %q: %v\n", i, rec.ColumnName(i), col)
+	}
+
+	// Output:
+	// column[0] "f1-i32": [1 2 3 4 5 6 7 8 (null) 10]
+	// column[1] "f2-f64": [1 2 3 4 5 6 7 8 9 10]
+}
+
+func Example_recordReader() {
+	pool := memory.NewGoAllocator()
+
+	schema := arrow.NewSchema(
+		[]arrow.Field{
+			arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
+			arrow.Field{Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64},
+		},
+		nil,
+	)
+
+	b := array.NewRecordBuilder(pool, schema)
+	defer b.Release()
+
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{1, 2, 3, 4, 5, 6}, nil)
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{7, 8, 9, 10}, []bool{true, true, false, true})
+	b.Field(1).(*array.Float64Builder).AppendValues([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)
+
+	rec1 := b.NewRecord()
+	defer rec1.Release()
+
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{11, 12, 13, 14, 15, 16, 17, 18, 19, 20}, nil)
+	b.Field(1).(*array.Float64Builder).AppendValues([]float64{11, 12, 13, 14, 15, 16, 17, 18, 19, 20}, nil)
+
+	rec2 := b.NewRecord()
+	defer rec2.Release()
+
+	itr, err := array.NewRecordReader(schema, []array.Record{rec1, rec2})
+	if err != nil {
+		log.Fatal(err)
+	}
+	defer itr.Release()
+
+	n := 0
+	for itr.Next() {
+		rec := itr.Record()
+		for i, col := range rec.Columns() {
+			fmt.Printf("rec[%d][%q]: %v\n", n, rec.ColumnName(i), col)
+		}
+		n++
+	}
+
+	// Output:
+	// rec[0]["f1-i32"]: [1 2 3 4 5 6 7 8 (null) 10]
+	// rec[0]["f2-f64"]: [1 2 3 4 5 6 7 8 9 10]
+	// rec[1]["f1-i32"]: [11 12 13 14 15 16 17 18 19 20]
+	// rec[1]["f2-f64"]: [11 12 13 14 15 16 17 18 19 20]
+}
+
+func Example_table() {
+	pool := memory.NewGoAllocator()
+
+	schema := arrow.NewSchema(
+		[]arrow.Field{
+			arrow.Field{Name: "f1-i32", Type: arrow.PrimitiveTypes.Int32},
+			arrow.Field{Name: "f2-f64", Type: arrow.PrimitiveTypes.Float64},
+		},
+		nil,
+	)
+
+	b := array.NewRecordBuilder(pool, schema)
+	defer b.Release()
+
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{1, 2, 3, 4, 5, 6}, nil)
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{7, 8, 9, 10}, []bool{true, true, false, true})
+	b.Field(1).(*array.Float64Builder).AppendValues([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)
+
+	rec1 := b.NewRecord()
+	defer rec1.Release()
+
+	b.Field(0).(*array.Int32Builder).AppendValues([]int32{11, 12, 13, 14, 15, 16, 17, 18, 19, 20}, nil)
+	b.Field(1).(*array.Float64Builder).AppendValues([]float64{11, 12, 13, 14, 15, 16, 17, 18, 19, 20}, nil)
+
+	rec2 := b.NewRecord()
+	defer rec2.Release()
+
+	tbl := array.NewTableFromRecords(schema, []array.Record{rec1, rec2})
+	defer tbl.Release()
+
+	tr := array.NewTableReader(tbl, 5)
+	defer tr.Release()
+
+	n := 0
+	for tr.Next() {
+		rec := tr.Record()
+		for i, col := range rec.Columns() {
+			fmt.Printf("rec[%d][%q]: %v\n", n, rec.ColumnName(i), col)
+		}
+		n++
+	}
+
+	// Output:
+	// rec[0]["f1-i32"]: [1 2 3 4 5]
+	// rec[0]["f2-f64"]: [1 2 3 4 5]
+	// rec[1]["f1-i32"]: [6 7 8 (null) 10]
+	// rec[1]["f2-f64"]: [6 7 8 9 10]
+	// rec[2]["f1-i32"]: [11 12 13 14 15]
+	// rec[2]["f2-f64"]: [11 12 13 14 15]
+	// rec[3]["f1-i32"]: [16 17 18 19 20]
+	// rec[3]["f2-f64"]: [16 17 18 19 20]
+}