]>
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 | class RecordBatchFileReaderTest < Test::Unit::TestCase | |
19 | test("write/read") do | |
20 | fields = [ | |
21 | Arrow::Field.new("uint8", :uint8), | |
22 | Arrow::Field.new("uint16", :uint16), | |
23 | Arrow::Field.new("uint32", :uint32), | |
24 | Arrow::Field.new("uint64", :uint64), | |
25 | Arrow::Field.new("int8", :int8), | |
26 | Arrow::Field.new("int16", :int16), | |
27 | Arrow::Field.new("int32", :int32), | |
28 | Arrow::Field.new("int64", :int64), | |
29 | Arrow::Field.new("float", :float), | |
30 | Arrow::Field.new("double", :double), | |
31 | ] | |
32 | schema = Arrow::Schema.new(fields) | |
33 | ||
34 | tempfile = Tempfile.new(["batch", ".arrow"]) | |
35 | Arrow::FileOutputStream.open(tempfile.path, false) do |output| | |
36 | Arrow::RecordBatchFileWriter.open(output, schema) do |writer| | |
37 | uints = [1, 2, 4, 8] | |
38 | ints = [1, -2, 4, -8] | |
39 | floats = [1.1, -2.2, 4.4, -8.8] | |
40 | columns = [ | |
41 | Arrow::UInt8Array.new(uints), | |
42 | Arrow::UInt16Array.new(uints), | |
43 | Arrow::UInt32Array.new(uints), | |
44 | Arrow::UInt64Array.new(uints), | |
45 | Arrow::Int8Array.new(ints), | |
46 | Arrow::Int16Array.new(ints), | |
47 | Arrow::Int32Array.new(ints), | |
48 | Arrow::Int64Array.new(ints), | |
49 | Arrow::FloatArray.new(floats), | |
50 | Arrow::DoubleArray.new(floats), | |
51 | ] | |
52 | ||
53 | record_batch = Arrow::RecordBatch.new(schema, 4, columns) | |
54 | writer.write_record_batch(record_batch) | |
55 | end | |
56 | end | |
57 | ||
58 | Arrow::MemoryMappedInputStream.open(tempfile.path) do |input| | |
59 | reader = Arrow::RecordBatchFileReader.new(input) | |
60 | reader.each do |record_batch| | |
61 | assert_equal([ | |
62 | { | |
63 | "uint8" => 1, | |
64 | "uint16" => 1, | |
65 | "uint32" => 1, | |
66 | "uint64" => 1, | |
67 | "int8" => 1, | |
68 | "int16" => 1, | |
69 | "int32" => 1, | |
70 | "int64" => 1, | |
71 | "float" => 1.100000023841858, | |
72 | "double" => 1.1, | |
73 | }, | |
74 | { | |
75 | "uint8" => 2, | |
76 | "uint16" => 2, | |
77 | "uint32" => 2, | |
78 | "uint64" => 2, | |
79 | "int8" => -2, | |
80 | "int16" => -2, | |
81 | "int32" => -2, | |
82 | "int64" => -2, | |
83 | "float" => -2.200000047683716, | |
84 | "double" => -2.2, | |
85 | }, | |
86 | { | |
87 | "uint8" => 4, | |
88 | "uint16" => 4, | |
89 | "uint32" => 4, | |
90 | "uint64" => 4, | |
91 | "int8" => 4, | |
92 | "int16" => 4, | |
93 | "int32" => 4, | |
94 | "int64" => 4, | |
95 | "float" => 4.400000095367432, | |
96 | "double" => 4.4, | |
97 | }, | |
98 | { | |
99 | "uint8" => 8, | |
100 | "uint16" => 8, | |
101 | "uint32" => 8, | |
102 | "uint64" => 8, | |
103 | "int8" => -8, | |
104 | "int16" => -8, | |
105 | "int32" => -8, | |
106 | "int64" => -8, | |
107 | "float" => -8.800000190734863, | |
108 | "double" => -8.8, | |
109 | }, | |
110 | ], | |
111 | record_batch.collect(&:to_h)) | |
112 | end | |
113 | end | |
114 | end | |
115 | end |