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
18module Helper
19  module Buildable
20    def build_null_array(values)
21      build_array(Arrow::NullArrayBuilder.new, values)
22    end
23
24    def build_boolean_array(values)
25      build_array(Arrow::BooleanArrayBuilder.new, values)
26    end
27
28    def build_int_array(values)
29      build_array(Arrow::IntArrayBuilder.new, values)
30    end
31
32    def build_uint_array(values)
33      build_array(Arrow::UIntArrayBuilder.new, values)
34    end
35
36    def build_int8_array(values)
37      build_array(Arrow::Int8ArrayBuilder.new, values)
38    end
39
40    def build_uint8_array(values)
41      build_array(Arrow::UInt8ArrayBuilder.new, values)
42    end
43
44    def build_int16_array(values)
45      build_array(Arrow::Int16ArrayBuilder.new, values)
46    end
47
48    def build_uint16_array(values)
49      build_array(Arrow::UInt16ArrayBuilder.new, values)
50    end
51
52    def build_int32_array(values)
53      build_array(Arrow::Int32ArrayBuilder.new, values)
54    end
55
56    def build_uint32_array(values)
57      build_array(Arrow::UInt32ArrayBuilder.new, values)
58    end
59
60    def build_int64_array(values)
61      build_array(Arrow::Int64ArrayBuilder.new, values)
62    end
63
64    def build_uint64_array(values)
65      build_array(Arrow::UInt64ArrayBuilder.new, values)
66    end
67
68    def build_float_array(values)
69      build_array(Arrow::FloatArrayBuilder.new, values)
70    end
71
72    def build_double_array(values)
73      build_array(Arrow::DoubleArrayBuilder.new, values)
74    end
75
76    def build_date32_array(values)
77      build_array(Arrow::Date32ArrayBuilder.new, values)
78    end
79
80    def build_date64_array(values)
81      build_array(Arrow::Date64ArrayBuilder.new, values)
82    end
83
84    def build_timestamp_array(unit, values)
85      data_type = Arrow::TimestampDataType.new(unit)
86      build_array(Arrow::TimestampArrayBuilder.new(data_type),
87                  values)
88    end
89
90    def build_time32_array(unit, values)
91      build_array(Arrow::Time32ArrayBuilder.new(Arrow::Time32DataType.new(unit)),
92                  values)
93    end
94
95    def build_time64_array(unit, values)
96      build_array(Arrow::Time64ArrayBuilder.new(Arrow::Time64DataType.new(unit)),
97                  values)
98    end
99
100    def build_binary_array(values)
101      build_array(Arrow::BinaryArrayBuilder.new, values)
102    end
103
104    def build_large_binary_array(values)
105      build_array(Arrow::LargeBinaryArrayBuilder.new, values)
106    end
107
108    def build_string_array(values)
109      build_array(Arrow::StringArrayBuilder.new, values)
110    end
111
112    def build_large_string_array(values)
113      build_array(Arrow::LargeStringArrayBuilder.new, values)
114    end
115
116    def build_list_array(value_data_type, values_list, field_name: "value")
117      value_field = Arrow::Field.new(field_name, value_data_type)
118      data_type = Arrow::ListDataType.new(value_field)
119      builder = Arrow::ListArrayBuilder.new(data_type)
120      values_list.each do |values|
121        if values.nil?
122          builder.append_null
123        else
124          append_to_builder(builder, values)
125        end
126      end
127      builder.finish
128    end
129
130    def build_large_list_array(value_data_type, values_list, field_name: "value")
131      value_field = Arrow::Field.new(field_name, value_data_type)
132      data_type = Arrow::LargeListDataType.new(value_field)
133      builder = Arrow::LargeListArrayBuilder.new(data_type)
134      values_list.each do |values|
135        if values.nil?
136          builder.append_null
137        else
138          append_to_builder(builder, values)
139        end
140      end
141      builder.finish
142    end
143
144    def build_struct_array(fields, structs)
145      data_type = Arrow::StructDataType.new(fields)
146      builder = Arrow::StructArrayBuilder.new(data_type)
147      structs.each do |struct|
148        if struct.nil?
149          builder.append_null
150        else
151          append_to_builder(builder, struct)
152        end
153      end
154      builder.finish
155    end
156
157    def append_to_builder(builder, value)
158      if value.nil?
159        builder.append_null
160      else
161        data_type = builder.value_data_type
162        case data_type
163        when Arrow::ListDataType, Arrow::LargeListDataType
164          builder.append_value
165          value_builder = builder.value_builder
166          value.each do |v|
167            append_to_builder(value_builder, v)
168          end
169        when Arrow::StructDataType
170          builder.append_value
171          value.each do |name, v|
172            field_index = data_type.get_field_index(name)
173            field_builder = builder.get_field_builder(field_index)
174            append_to_builder(field_builder, v)
175          end
176        else
177          builder.append_value(value)
178        end
179      end
180    end
181
182    def build_table(columns)
183      fields = []
184      arrays = []
185      columns.each do |name, array|
186        fields << Arrow::Field.new(name, array.value_data_type)
187        arrays << array
188      end
189      schema = Arrow::Schema.new(fields)
190      Arrow::Table.new(schema, arrays)
191    end
192
193    def build_record_batch(arrays)
194      n_rows = arrays.collect {|_, array| array.length}.min || 0
195      fields = arrays.collect do |name, array|
196        Arrow::Field.new(name, array.value_data_type)
197      end
198      schema = Arrow::Schema.new(fields)
199      Arrow::RecordBatch.new(schema, n_rows, arrays.values)
200    end
201
202    private
203    def build_array(builder, values)
204      values.each do |value|
205        if value.nil?
206          builder.append_null
207        elsif builder.respond_to?(:append_string)
208          builder.append_string(value)
209        else
210          builder.append_value(value)
211        end
212      end
213      builder.finish
214    end
215  end
216end
217