[arrow] branch main updated: GH-34832: [Go] Add Record SetColumn method (#34794)

[ze...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow.git


The following commit(s) were added to refs/heads/main by this push:
     new c582300007 GH-34832: [Go] Add Record SetColumn method  (#34794)
c582300007 is described below

commit c582300007c3aa261682e43bf24de77f50c2ca9e
Author: Yevgeny Pats <ye...@gmail.com>
AuthorDate: Sat Apr 1 17:21:03 2023 -0400

    GH-34832: [Go] Add Record SetColumn method  (#34794)
    
    Small [method](https://arrow.apache.org/docs/cpp/api/table.html#_CPPv4NK5arrow11RecordBatch9SetColumnEiRKNSt10shared_ptrI5FieldEERKNSt10shared_ptrI5ArrayEE) similar to the c/c++ implementation.
    
    * Closes: #34832
    
    Authored-by: Yevgeny Pats <16...@users.noreply.github.com>
    Signed-off-by: Matt Topol <zo...@gmail.com>
---
 go/arrow/array/record.go      | 26 ++++++++++++++++++++++++++
 go/arrow/array/record_test.go | 20 ++++++++++++++++++++
 go/arrow/record.go            |  1 +
 3 files changed, 47 insertions(+)

diff --git a/go/arrow/array/record.go b/go/arrow/array/record.go
index 3b55d83d4a..b00a59d16f 100644
--- a/go/arrow/array/record.go
+++ b/go/arrow/array/record.go
@@ -154,6 +154,32 @@ func NewRecord(schema *arrow.Schema, cols []arrow.Array, nrows int64) *simpleRec
 	return rec
 }
 
+func (rec *simpleRecord) SetColumn(i int, arr arrow.Array) (arrow.Record, error) {
+	if i < 0 || i >= len(rec.arrs) {
+		return nil, fmt.Errorf("arrow/array: column index out of range [0, %d): got=%d", len(rec.arrs), i)
+	}
+
+	if arr.Len() != int(rec.rows) {
+		return nil, fmt.Errorf("arrow/array: mismatch number of rows in column %q: got=%d, want=%d",
+			rec.schema.Field(i).Name,
+			arr.Len(), rec.rows,
+		)
+	}
+
+	f := rec.schema.Field(i)
+	if !arrow.TypeEqual(f.Type, arr.DataType()) {
+		return nil, fmt.Errorf("arrow/array: column %q type mismatch: got=%v, want=%v",
+			f.Name,
+			arr.DataType(), f.Type,
+		)
+	}
+	arrs := make([]arrow.Array, len(rec.arrs))
+	copy(arrs, rec.arrs)
+	arrs[i] = arr
+
+	return NewRecord(rec.schema, arrs, rec.rows), nil
+}
+
 func (rec *simpleRecord) validate() error {
 	if rec.rows == 0 && len(rec.arrs) == 0 {
 		return nil
diff --git a/go/arrow/array/record_test.go b/go/arrow/array/record_test.go
index 2724d3943c..d9e62ce2b8 100644
--- a/go/arrow/array/record_test.go
+++ b/go/arrow/array/record_test.go
@@ -56,6 +56,15 @@ func TestRecord(t *testing.T) {
 	}()
 	defer col2.Release()
 
+	col2_1 := func() arrow.Array {
+		b := array.NewFloat64Builder(mem)
+		defer b.Release()
+
+		b.AppendValues([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, nil)
+		return b.NewFloat64Array()
+	}()
+	defer col2_1.Release()
+
 	cols := []arrow.Array{col1, col2}
 	rec := array.NewRecord(schema, cols, -1)
 	defer rec.Release()
@@ -82,6 +91,17 @@ func TestRecord(t *testing.T) {
 	if got, want := rec.ColumnName(0), schema.Field(0).Name; got != want {
 		t.Fatalf("invalid column name: got=%q, want=%q", got, want)
 	}
+	if _, err := rec.SetColumn(0, col2_1); err == nil {
+		t.Fatalf("expected an error")
+	}
+	newRec, err := rec.SetColumn(1, col2_1);
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	defer newRec.Release()
+	if !reflect.DeepEqual(newRec.Column(1), col2_1) {
+		t.Fatalf("invalid column: got=%q, want=%q", rec.Column(1), col2_1)
+	}
 
 	for _, tc := range []struct {
 		i, j int64
diff --git a/go/arrow/record.go b/go/arrow/record.go
index fa63641d5a..b93f1015b9 100644
--- a/go/arrow/record.go
+++ b/go/arrow/record.go
@@ -37,6 +37,7 @@ type Record interface {
 	Columns() []Array
 	Column(i int) Array
 	ColumnName(i int) string
+	SetColumn(i int, col Array) (Record, error)
 
 	// NewSlice constructs a zero-copy slice of the record with the indicated
 	// indices i and j, corresponding to array[i:j].