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-go.git
The following commit(s) were added to refs/heads/main by this push:
new e117ecf8 perf(arrow/csv): avoid materializing rows before writing
(#1190)
e117ecf8 is described below
commit e117ecf8e863a7a5fcba11eb2ba11cda4d157ddf
Author: Minh Vu <[email protected]>
AuthorDate: Wed Aug 26 22:16:14 2026 +0200
perf(arrow/csv): avoid materializing rows before writing (#1190)
### Rationale for this change
`Writer.Write` currently converts columns to strings, then copies every
string header into a row-major `[][]string` before calling `WriteAll`.
This allocates one `[]string` for every row. Large record batches spend
a lot of memory on this temporary matrix.
Decimal formatting also used `math/big.Float`, which adds extra work and
allocations for normal decimal values.
### What changes are included in this PR?
- Keep the converted strings in their existing column slices.
- Reuse one row-sized `[]string` while writing each CSV record.
- Format decimal values directly from their exact integer representation
when the value fits the supported precision.
- Keep the existing `math/big.Float` path for values outside the
supported precision.
- Preserve the existing flush and error behavior.
For the existing 1,000-row, 16-column benchmark, medians from 6 runs on
an Apple M1 Pro were:
| | main | this PR | change |
|---|---:|---:|---:|
| time/op | 4.78 ms | 0.84 ms | -82.4% |
| B/op | 2,292,279 | 432,927 | -81.1% |
| allocs/op | 41,514 | 19,568 | -52.9% |
Compared with the earlier version of this PR, the decimal formatting
change reduces `B/op` by 78.5% and allocations by 51.7%:
| | before decimal formatting | after decimal formatting | change |
|---|---:|---:|---:|
| B/op | 2,012,183 | 432,927 | -78.5% |
| allocs/op | 40,515 | 19,568 | -51.7% |
### Are these changes tested?
- `go test ./arrow/... -count=1`
- `go test ./arrow/csv -run "^$" -bench "^BenchmarkWrite$" -benchmem
-count=6`
### Are there any user-facing changes?
No. The generated CSV output is unchanged.
---
arrow/csv/transformer.go | 2 +-
arrow/csv/writer.go | 21 +++++++-----
arrow/csv/writer_test.go | 68 +++++++++++++++++++++++++++++++++++++
arrow/decimal128/decimal128.go | 39 ++++++++++++++++++---
arrow/decimal128/decimal128_test.go | 31 +++++++++++++++++
arrow/decimal256/decimal256.go | 39 ++++++++++++++++++---
arrow/decimal256/decimal256_test.go | 31 +++++++++++++++++
7 files changed, 212 insertions(+), 19 deletions(-)
diff --git a/arrow/csv/transformer.go b/arrow/csv/transformer.go
index 9321a490..8073874d 100644
--- a/arrow/csv/transformer.go
+++ b/arrow/csv/transformer.go
@@ -34,7 +34,7 @@ func (w *Writer) transformColToStringArr(typ arrow.DataType,
col arrow.Array, st
if len(result) != col.Len() {
return nil, fmt.Errorf("%w: custom type
converter returned %d values for column with %d rows", arrow.ErrInvalid,
len(result), col.Len())
}
- return result, nil
+ return append([]string(nil), result...), nil
}
}
diff --git a/arrow/csv/writer.go b/arrow/csv/writer.go
index 5165d517..e14c3274 100644
--- a/arrow/csv/writer.go
+++ b/arrow/csv/writer.go
@@ -78,22 +78,27 @@ func (w *Writer) Write(record arrow.RecordBatch) error {
}
}
- recs := make([][]string, record.NumRows())
- for i := range recs {
- recs[i] = make([]string, record.NumCols())
- }
-
+ columns := make([][]string, record.NumCols())
for j, col := range record.Columns() {
rows, err := w.transformColToStringArr(w.schema.Field(j).Type,
col, w.stringReplacer)
if err != nil {
return err
}
- for i, row := range rows {
- recs[i][j] = row
+ columns[j] = rows
+ }
+
+ row := make([]string, record.NumCols())
+ for i := 0; i < int(record.NumRows()); i++ {
+ for j := range columns {
+ row[j] = columns[j][i]
+ }
+ if err := w.w.Write(row); err != nil {
+ return err
}
}
- return w.w.WriteAll(recs)
+ w.w.Flush()
+ return w.w.Error()
}
// Flush writes any buffered data to the underlying csv Writer.
diff --git a/arrow/csv/writer_test.go b/arrow/csv/writer_test.go
index 17a09bcf..ba3af2dc 100644
--- a/arrow/csv/writer_test.go
+++ b/arrow/csv/writer_test.go
@@ -194,6 +194,39 @@ func TestCSVWriter(t *testing.T) {
}
}
+func TestCSVWriterWritesMultipleRecordBatches(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "id", Type: arrow.PrimitiveTypes.Int32},
+ {Name: "label", Type: arrow.BinaryTypes.String},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+
+ ids := builder.Field(0).(*array.Int32Builder)
+ labels := builder.Field(1).(*array.StringBuilder)
+ ids.AppendValues([]int32{1, 2}, nil)
+ labels.AppendValues([]string{"first, row", "line\nbreak"}, nil)
+ first := builder.NewRecordBatch()
+ defer first.Release()
+
+ ids.Append(3)
+ labels.Append("last")
+ second := builder.NewRecordBatch()
+ defer second.Release()
+
+ var output bytes.Buffer
+ writer := csv.NewWriter(&output, schema, csv.WithHeader(true))
+ require.NoError(t, writer.Write(first))
+ require.NoError(t, writer.Write(second))
+ require.NoError(t, writer.Flush())
+ require.NoError(t, writer.Error())
+
+ assert.Equal(t, "id,label\n1,\"first,
row\"\n2,\"line\nbreak\"\n3,last\n", output.String())
+}
+
func genTimestamps(unit arrow.TimeUnit) []arrow.Timestamp {
out := []arrow.Timestamp{}
for _, input := range []string{"2014-07-28 15:04:05", "2016-09-08
15:04:05", "2021-09-18 15:04:05"} {
@@ -487,6 +520,41 @@ func TestCustomTypeConverterValidatesRowCount(t
*testing.T) {
}
}
+func TestCSVWriterSupportsReusedCustomTypeConverterResult(t *testing.T) {
+ mem := memory.NewCheckedAllocator(memory.NewGoAllocator())
+ defer mem.AssertSize(t, 0)
+
+ schema := arrow.NewSchema([]arrow.Field{
+ {Name: "first", Type: arrow.PrimitiveTypes.Int32},
+ {Name: "second", Type: arrow.PrimitiveTypes.Int32},
+ }, nil)
+ builder := array.NewRecordBuilder(mem, schema)
+ defer builder.Release()
+ builder.Field(0).(*array.Int32Builder).AppendValues([]int32{1, 2}, nil)
+ builder.Field(1).(*array.Int32Builder).AppendValues([]int32{10, 20},
nil)
+ record := builder.NewRecordBatch()
+ defer record.Release()
+
+ result := make([]string, 2)
+ var output bytes.Buffer
+ writer := csv.NewWriter(&output, schema,
csv.WithCustomTypeConverter(func(typ arrow.DataType, col arrow.Array)
([]string, bool) {
+ if typ.ID() != arrow.INT32 {
+ return nil, false
+ }
+
+ arr := col.(*array.Int32)
+ for i := 0; i < arr.Len(); i++ {
+ result[i] = fmt.Sprintf("%d", arr.Value(i))
+ }
+ return result, true
+ }))
+
+ require.NoError(t, writer.Write(record))
+ require.NoError(t, writer.Flush())
+ require.NoError(t, writer.Error())
+ assert.Equal(t, "1,10\n2,20\n", output.String())
+}
+
// TestParquetTestingCSVWriter tests that the CSV writer successfully convert
arrow/parquet-testing files to CSV
func TestParquetTestingCSVWriter(t *testing.T) {
dir := os.Getenv("PARQUET_TEST_DATA")
diff --git a/arrow/decimal128/decimal128.go b/arrow/decimal128/decimal128.go
index 660c4131..fca8c799 100644
--- a/arrow/decimal128/decimal128.go
+++ b/arrow/decimal128/decimal128.go
@@ -22,6 +22,7 @@ import (
"math"
"math/big"
"math/bits"
+ "strings"
"github.com/apache/arrow-go/v18/arrow/internal/debug"
)
@@ -526,13 +527,41 @@ func (n Num) FitsInPrecision(prec int32) bool {
}
func (n Num) ToString(scale int32) string {
- f := (&big.Float{}).SetInt(n.BigInt())
- if scale < 0 {
- f.SetPrec(128).Mul(f,
(&big.Float{}).SetInt(scaleMultipliers[-scale].BigInt()))
+ if scale < -MaxScale || scale > MaxScale {
+ panic("arrow/decimal128: scale out of range")
+ }
+
+ value := n.BigInt().String()
+ digits := len(value)
+ if value[0] == '-' {
+ digits--
+ }
+ if digits > MaxPrecision {
+ return n.ToBigFloat(scale).Text('f', int(scale))
+ }
+ if scale <= 0 {
+ if scale < 0 && value != "0" {
+ value += strings.Repeat("0", int(-scale))
+ }
+ return value
+ }
+
+ negative := value[0] == '-'
+ if negative {
+ value = value[1:]
+ }
+ places := int(scale)
+ var result string
+ if len(value) <= places {
+ result = "0." + strings.Repeat("0", places-len(value)) + value
} else {
- f.SetPrec(128).Quo(f,
(&big.Float{}).SetInt(scaleMultipliers[scale].BigInt()))
+ at := len(value) - places
+ result = value[:at] + "." + value[at:]
+ }
+ if negative {
+ return "-" + result
}
- return f.Text('f', int(scale))
+ return result
}
func GetScaleMultiplier(pow int) Num { return scaleMultipliers[pow] }
diff --git a/arrow/decimal128/decimal128_test.go
b/arrow/decimal128/decimal128_test.go
index f9758988..ba5539f1 100644
--- a/arrow/decimal128/decimal128_test.go
+++ b/arrow/decimal128/decimal128_test.go
@@ -707,3 +707,34 @@ func TestFromStringDecimal128b(t *testing.T) {
require.NoError(t, err)
assert.Equal(t, decStr, num.ToString(19))
}
+
+func TestToStringMatchesBigFloat(t *testing.T) {
+ values := []string{
+ "0",
+ "1",
+ "-1",
+ "9",
+ "-9",
+ "12345678901234567890123456789012345678",
+ "-12345678901234567890123456789012345678",
+ "170141183460469231731687303715884105727",
+ "-170141183460469231731687303715884105727",
+ }
+ scales := []int32{-38, -10, -1, 0, 1, 2, 10, 38}
+
+ for _, value := range values {
+ integer, ok := new(big.Int).SetString(value, 10)
+ require.True(t, ok)
+ number := decimal128.FromBigInt(integer)
+ for _, scale := range scales {
+ want := number.ToBigFloat(scale).Text('f', int(scale))
+ assert.Equalf(t, want, number.ToString(scale),
"value=%s scale=%d", value, scale)
+ }
+ }
+}
+
+func TestToStringRejectsInvalidScale(t *testing.T) {
+ for _, scale := range []int32{-39, 39} {
+ assert.Panics(t, func() { decimal128.FromI64(1).ToString(scale)
})
+ }
+}
diff --git a/arrow/decimal256/decimal256.go b/arrow/decimal256/decimal256.go
index 82c52a65..a473505f 100644
--- a/arrow/decimal256/decimal256.go
+++ b/arrow/decimal256/decimal256.go
@@ -22,6 +22,7 @@ import (
"math"
"math/big"
"math/bits"
+ "strings"
"github.com/apache/arrow-go/v18/arrow/decimal128"
"github.com/apache/arrow-go/v18/arrow/internal/debug"
@@ -526,13 +527,41 @@ func (n Num) FitsInPrecision(prec int32) bool {
}
func (n Num) ToString(scale int32) string {
- f := (&big.Float{}).SetInt(n.BigInt())
- if scale < 0 {
- f.SetPrec(256).Mul(f,
(&big.Float{}).SetInt(scaleMultipliers[-scale].BigInt()))
+ if scale < -MaxScale || scale > MaxScale {
+ panic("arrow/decimal256: scale out of range")
+ }
+
+ value := n.BigInt().String()
+ digits := len(value)
+ if value[0] == '-' {
+ digits--
+ }
+ if digits > MaxPrecision {
+ return n.ToBigFloat(scale).Text('f', int(scale))
+ }
+ if scale <= 0 {
+ if scale < 0 && value != "0" {
+ value += strings.Repeat("0", int(-scale))
+ }
+ return value
+ }
+
+ negative := value[0] == '-'
+ if negative {
+ value = value[1:]
+ }
+ places := int(scale)
+ var result string
+ if len(value) <= places {
+ result = "0." + strings.Repeat("0", places-len(value)) + value
} else {
- f.SetPrec(256).Quo(f,
(&big.Float{}).SetInt(scaleMultipliers[scale].BigInt()))
+ at := len(value) - places
+ result = value[:at] + "." + value[at:]
+ }
+ if negative {
+ return "-" + result
}
- return f.Text('f', int(scale))
+ return result
}
func GetScaleMultiplier(pow int) Num { return scaleMultipliers[pow] }
diff --git a/arrow/decimal256/decimal256_test.go
b/arrow/decimal256/decimal256_test.go
index a93cdfe4..0d1f14f2 100644
--- a/arrow/decimal256/decimal256_test.go
+++ b/arrow/decimal256/decimal256_test.go
@@ -590,6 +590,37 @@ func TestToString(t *testing.T) {
assert.Equal(t, decStr+"0000", dec.ToString(-4))
}
+func TestToStringMatchesBigFloat(t *testing.T) {
+ values := []string{
+ "0",
+ "1",
+ "-1",
+ "9",
+ "-9",
+
"1234567890123456789012345678901234567890123456789012345678901234567890123456",
+
"-1234567890123456789012345678901234567890123456789012345678901234567890123456",
+
"57896044618658097711785492504343953926634992332820282019728792003956564819967",
+
"-57896044618658097711785492504343953926634992332820282019728792003956564819967",
+ }
+ scales := []int32{-76, -10, -1, 0, 1, 2, 10, 38, 76}
+
+ for _, value := range values {
+ integer, ok := new(big.Int).SetString(value, 10)
+ assert.True(t, ok)
+ number := decimal256.FromBigInt(integer)
+ for _, scale := range scales {
+ want := number.ToBigFloat(scale).Text('f', int(scale))
+ assert.Equalf(t, want, number.ToString(scale),
"value=%s scale=%d", value, scale)
+ }
+ }
+}
+
+func TestToStringRejectsInvalidScale(t *testing.T) {
+ for _, scale := range []int32{-77, 77} {
+ assert.Panics(t, func() { decimal256.FromI64(1).ToString(scale)
})
+ }
+}
+
// Test issues from GH-38395
func TestHexFromString(t *testing.T) {
const decStr =
"11111111111111111111111111111111111111.00000000000000000000000000000000000000"
