-
Notifications
You must be signed in to change notification settings - Fork 67
/
writer.go
46 lines (38 loc) · 1017 Bytes
/
writer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
package parquetio
import (
"fmt"
"io"
"strings"
"github.com/apache/arrow/go/v14/arrow"
"github.com/apache/arrow/go/v14/parquet/pqarrow"
"github.com/brimdata/zed"
"github.com/brimdata/zed/zio"
"github.com/brimdata/zed/zio/arrowio"
)
type Writer struct {
*arrowio.Writer
}
func NewWriter(wc io.WriteCloser) *Writer {
w := arrowio.NewWriter(wc)
w.NewWriterFunc = func(w io.Writer, s *arrow.Schema) (arrowio.WriteCloser, error) {
fw, err := pqarrow.NewFileWriter(s, zio.NopCloser(w), nil, pqarrow.DefaultWriterProps())
if err != nil {
return nil, fmt.Errorf("%w: %s", arrowio.ErrUnsupportedType, err)
}
return fw, nil
}
return &Writer{w}
}
func (w *Writer) Write(val zed.Value) error {
if err := w.Writer.Write(val); err != nil {
return parquetioError{err}
}
return nil
}
type parquetioError struct {
err error
}
func (p parquetioError) Error() string {
return "parquetio: " + strings.TrimPrefix(p.err.Error(), "arrowio: ")
}
func (p parquetioError) Unwrap() error { return p.err }