diff --git a/plugin/append_from_string.go b/plugin/append_from_string.go new file mode 100644 index 0000000000..3ed150c612 --- /dev/null +++ b/plugin/append_from_string.go @@ -0,0 +1,23 @@ +package plugin + +import ( + "strings" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/goccy/go-json" +) + +// appendFromString appends a value from its string representation. Nested types are +// JSON-encoded, so they are decoded with UseNumber to keep int64/uint64 values that +// exceed float64 precision intact. +func appendFromString(b array.Builder, s string) error { + switch b.Type().ID() { + case arrow.LIST, arrow.LARGE_LIST, arrow.LIST_VIEW, arrow.LARGE_LIST_VIEW, arrow.FIXED_SIZE_LIST, arrow.MAP, arrow.STRUCT: + dec := json.NewDecoder(strings.NewReader(s)) + dec.UseNumber() + return b.UnmarshalOne(dec) + default: + return b.AppendValueFromString(s) + } +} diff --git a/plugin/append_from_string_test.go b/plugin/append_from_string_test.go new file mode 100644 index 0000000000..c4774356f7 --- /dev/null +++ b/plugin/append_from_string_test.go @@ -0,0 +1,37 @@ +package plugin + +import ( + "testing" + + "github.com/apache/arrow-go/v18/arrow" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/apache/arrow-go/v18/arrow/memory" + "github.com/stretchr/testify/require" +) + +func TestAppendFromStringPreservesInt64Precision(t *testing.T) { + for _, tc := range []struct { + name string + dt arrow.DataType + str string + }{ + {name: "int64", dt: arrow.PrimitiveTypes.Int64, str: "-8717895732742165505"}, + {name: "list_of_int64", dt: arrow.ListOf(arrow.PrimitiveTypes.Int64), str: "[-8717895732742165505]"}, + {name: "large_list_of_int64", dt: arrow.LargeListOf(arrow.PrimitiveTypes.Int64), str: "[-8717895732742165505]"}, + {name: "struct_with_int64", dt: arrow.StructOf(arrow.Field{Name: "v", Type: arrow.PrimitiveTypes.Int64, Nullable: true}), str: `{"v":-8717895732742165505}`}, + {name: "uint64", dt: arrow.PrimitiveTypes.Uint64, str: "18428615660272232523"}, + {name: "list_of_uint64", dt: arrow.ListOf(arrow.PrimitiveTypes.Uint64), str: "[18428615660272232523]"}, + } { + t.Run(tc.name, func(t *testing.T) { + bldr := array.NewBuilder(memory.DefaultAllocator, tc.dt) + defer bldr.Release() + + require.NoError(t, appendFromString(bldr, tc.str)) + + arr := bldr.NewArray() + defer arr.Release() + + require.Equal(t, tc.str, arr.ValueStr(0)) + }) + } +} diff --git a/plugin/nulls.go b/plugin/nulls.go index 961a971037..7666b1a9e5 100644 --- a/plugin/nulls.go +++ b/plugin/nulls.go @@ -22,7 +22,7 @@ func stripNullsFromLists(list array.ListLike) array.ListLike { if slc.IsNull(k) { continue } - err := vBldr.AppendValueFromString(slc.ValueStr(k)) + err := appendFromString(vBldr, slc.ValueStr(k)) if err != nil { panic(err) } @@ -47,7 +47,7 @@ func (s *WriterTestSuite) replaceNullsByEmpty(arr arrow.Array) arrow.Array { continue } - if err := builder.AppendValueFromString(arr.ValueStr(j)); err != nil { + if err := appendFromString(builder, arr.ValueStr(j)); err != nil { panic(err) } } diff --git a/plugin/testing_write_upsert.go b/plugin/testing_write_upsert.go index f5013aca82..2997a6190e 100644 --- a/plugin/testing_write_upsert.go +++ b/plugin/testing_write_upsert.go @@ -208,7 +208,7 @@ func extractLastRowFromRecord(table *schema.Table, existingRecord arrow.RecordBa for i, c := range table.Columns { col := existingRecord.Column(i) lastRow := int(existingRecord.NumRows()) - 1 - err := bldr.Field(i).AppendValueFromString(col.ValueStr(lastRow)) + err := appendFromString(bldr.Field(i), col.ValueStr(lastRow)) if err != nil { return nil, fmt.Errorf("failed to unmarshal json `%v` for column %v: %v", col.ValueStr(lastRow), c.Name, err) }