forked from DuanWeiFan/parquet-go
/
parquet_go18.go
87 lines (81 loc) · 2.48 KB
/
parquet_go18.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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
//go:build go1.18
package parquet
import (
"io"
"os"
)
// Read reads and returns rows from the parquet file in the given reader.
//
// The type T defines the type of rows read from r. T must be compatible with
// the file's schema or an error will be returned. The row type might represent
// a subset of the full schema, in which case only a subset of the columns will
// be loaded from r.
//
// This function is provided for convenience to facilitate reading of parquet
// files from arbitrary locations in cases where the data set fit in memory.
func Read[T any](r io.ReaderAt, size int64, options ...ReaderOption) (rows []T, err error) {
config, err := NewReaderConfig(options...)
if err != nil {
return nil, err
}
file, err := OpenFile(r, size)
if err != nil {
return nil, err
}
rows = make([]T, file.NumRows())
reader := NewGenericReader[T](file, config)
n, err := reader.Read(rows)
if err == io.EOF {
err = nil
}
reader.Close()
return rows[:n], err
}
// ReadFile reads rows of the parquet file at the given path.
//
// The type T defines the type of rows read from r. T must be compatible with
// the file's schema or an error will be returned. The row type might represent
// a subset of the full schema, in which case only a subset of the columns will
// be loaded from the file.
//
// This function is provided for convenience to facilitate reading of parquet
// files from the file system in cases where the data set fit in memory.
func ReadFile[T any](path string, options ...ReaderOption) (rows []T, err error) {
f, err := os.Open(path)
if err != nil {
return nil, err
}
defer f.Close()
s, err := f.Stat()
if err != nil {
return nil, err
}
return Read[T](f, s.Size())
}
// Write writes the given list of rows to a parquet file written to w.
//
// This function is provided for convenience to facilitate the creation of
// parquet files.
func Write[T any](w io.Writer, rows []T, options ...WriterOption) error {
config, err := NewWriterConfig(options...)
if err != nil {
return err
}
writer := NewGenericWriter[T](w, config)
if _, err := writer.Write(rows); err != nil {
return err
}
return writer.Close()
}
// Write writes the given list of rows to a parquet file written to w.
//
// This function is provided for convenience to facilitate writing parquet
// files to the file system.
func WriteFile[T any](path string, rows []T, options ...WriterOption) error {
f, err := os.Create(path)
if err != nil {
return err
}
defer f.Close()
return Write(f, rows, options...)
}