Module: ArrowActiveRecord::Arrowable

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

Instance Method Summary collapse

Instance Method Details

#each_record_batch(batch_size: 10000, &block) ⇒ Object 



10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
 
# File 'lib/arrow-activerecord/arrowable.rb', line 10

def each_record_batch(batch_size: 10000, &block)
  return to_enum(__method__, batch_size:) unless block_given?

  schema = build_arrow_schema
  target_column_names = schema.fields.collect(&:name)
  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
    yield(record_batch_builder.flush)
  end
end

#to_arrow(batch_size: 10000) ⇒ Object 



5
6
7
8
 
# File 'lib/arrow-activerecord/arrowable.rb', line 5

def to_arrow(batch_size: 10000)
  record_batches = each_record_batch(batch_size:).to_a
  Arrow::Table.new(record_batches.first.schema, record_batches)
end