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