Module: ArrowActiveRecord::Arrowable

Defined in:
lib/arrow-activerecord/arrowable.rb

Instance Method Summary collapse

Instance Method Details

#to_arrow(batch_size: 10000) ⇒ Object 



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
 
# File 'lib/arrow-activerecord/arrowable.rb', line 5

def to_arrow(batch_size: 10000)
  target_column_names = select_values
  target_column_names = column_names if select_values.empty?

  fields = []
  target_column_names.each do |name|
    name = name.to_s
    target_column = columns.find do |column|
      column.name == name
    end
    fields << {name: name, data_type: extract_arrow_data_type(target_column)}
  end
  schema = Arrow::Schema.new(fields)

  record_batches = []
  record_batch_builder = Arrow::RecordBatchBuilder.new(schema)
  in_batches(of: batch_size).each do |relation|
    records = relation.pluck(*target_column_names)
    if target_column_names.size == 1
      # [1, 2, 3] ->
      # [[1], [2], [3]]
      record_batch_builder.append(records.collect {|value| [value]})
    else
      record_batch_builder.append(records)
    end
    record_batches << record_batch_builder.flush
  end
  Arrow::Table.new(schema, record_batches)
end