]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
1 | -- Licensed to the Apache Software Foundation (ASF) under one |
2 | -- or more contributor license agreements. See the NOTICE file | |
3 | -- distributed with this work for additional information | |
4 | -- regarding copyright ownership. The ASF licenses this file | |
5 | -- to you under the Apache License, Version 2.0 (the | |
6 | -- "License"); you may not use this file except in compliance | |
7 | -- with the License. You may obtain a copy of the License at | |
8 | -- | |
9 | -- http://www.apache.org/licenses/LICENSE-2.0 | |
10 | -- | |
11 | -- Unless required by applicable law or agreed to in writing, | |
12 | -- software distributed under the License is distributed on an | |
13 | -- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY | |
14 | -- KIND, either express or implied. See the License for the | |
15 | -- specific language governing permissions and limitations | |
16 | -- under the License. | |
17 | ||
18 | local lgi = require 'lgi' | |
19 | local Arrow = lgi.Arrow | |
20 | ||
21 | local output_path = arg[1] or "/tmp/stream.arrow"; | |
22 | ||
23 | local fields = { | |
24 | Arrow.Field.new("uint8", Arrow.UInt8DataType.new()), | |
25 | Arrow.Field.new("uint16", Arrow.UInt16DataType.new()), | |
26 | Arrow.Field.new("uint32", Arrow.UInt32DataType.new()), | |
27 | Arrow.Field.new("uint64", Arrow.UInt64DataType.new()), | |
28 | Arrow.Field.new("int8", Arrow.Int8DataType.new()), | |
29 | Arrow.Field.new("int16", Arrow.Int16DataType.new()), | |
30 | Arrow.Field.new("int32", Arrow.Int32DataType.new()), | |
31 | Arrow.Field.new("int64", Arrow.Int64DataType.new()), | |
32 | Arrow.Field.new("float", Arrow.FloatDataType.new()), | |
33 | Arrow.Field.new("double", Arrow.DoubleDataType.new()), | |
34 | } | |
35 | local schema = Arrow.Schema.new(fields) | |
36 | ||
37 | local output = Arrow.FileOutputStream.new(output_path, false) | |
38 | local writer = Arrow.RecordBatchStreamWriter.new(output, schema) | |
39 | ||
40 | function build_array(builder, values) | |
41 | for _, value in pairs(values) do | |
42 | builder:append(value) | |
43 | end | |
44 | return builder:finish() | |
45 | end | |
46 | ||
47 | local uints = {1, 2, 4, 8} | |
48 | local ints = {1, -2, 4, -8} | |
49 | local floats = {1.1, -2.2, 4.4, -8.8} | |
50 | local columns = { | |
51 | build_array(Arrow.UInt8ArrayBuilder.new(), uints), | |
52 | build_array(Arrow.UInt16ArrayBuilder.new(), uints), | |
53 | build_array(Arrow.UInt32ArrayBuilder.new(), uints), | |
54 | build_array(Arrow.UInt64ArrayBuilder.new(), uints), | |
55 | build_array(Arrow.Int8ArrayBuilder.new(), ints), | |
56 | build_array(Arrow.Int16ArrayBuilder.new(), ints), | |
57 | build_array(Arrow.Int32ArrayBuilder.new(), ints), | |
58 | build_array(Arrow.Int64ArrayBuilder.new(), ints), | |
59 | build_array(Arrow.FloatArrayBuilder.new(), floats), | |
60 | build_array(Arrow.DoubleArrayBuilder.new(), floats), | |
61 | } | |
62 | ||
63 | local record_batch = Arrow.RecordBatch.new(schema, 4, columns) | |
64 | writer:write_record_batch(record_batch) | |
65 | ||
66 | local sliced_columns = {} | |
67 | for i, column in pairs(columns) do | |
68 | sliced_columns[i] = column:slice(1, 3) | |
69 | end | |
70 | record_batch = Arrow.RecordBatch.new(schema, 3, sliced_columns) | |
71 | writer:write_record_batch(record_batch) | |
72 | ||
73 | writer:close() | |
74 | output:close() |