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 <[email protected]>
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 <[email protected]>
Signed-off-by: Matt Topol <[email protected]>
---
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].
