Add support for one column case

Author: kou

lib/arrow-activerecord/arrowable.rb

@@ -20,7 +20,11 @@ def to_arrow(batch_size: 10000)
       record_batch_builder = Arrow::RecordBatchBuilder.new(schema)
       in_batches(of: batch_size).each do |relation|
         records = relation.pluck(*target_column_names)
-        record_batch_builder.append(records)
+        if target_column_names.size == 1
+          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)

test/arrowable_test.rb

@@ -23,32 +23,58 @@ class Data < ActiveRecord::Base
         table.boolean :boolean_column
         table.bigint :bigint_column
       end
-      Data.create(:string_column => "Hello",
-                  :date_column => @date_value,
-                  :datetime_column => @datetime_value,
-                  :boolean_column => false,
-                  :bigint_column => @bigint_value)
+      Data.create(string_column: "Hello",
+                  date_column: @date_value,
+                  datetime_column: @datetime_value,
+                  boolean_column: false,
+                  bigint_column: @bigint_value)
+      Data.create(string_column: "Hello2",
+                  date_column: @date_value + 1,
+                  datetime_column: @datetime_value + 1,
+                  boolean_column: true,
+                  bigint_column: -@bigint_value - 1)
+    end
+
+    teardown do
+      ActiveRecord::Base.connection.drop_table(:data)
+    end
+
+    test "one" do
+      arrow = Data.all.select(:id).to_arrow
+      assert_equal(<<-RECORD_BATCH, arrow.each_record_batch.first.to_s)
+id:   [
+    1,
+    2
+  ]
+      RECORD_BATCH
     end
 
     test "all" do
-      assert_equal(<<-RECORD_BATCH, Data.all.to_arrow.each_record_batch.first.to_s)
+      arrow = Data.all.to_arrow
+      assert_equal(<<-RECORD_BATCH, arrow.each_record_batch.first.to_s)
 id:   [
-    1
+    1,
+    2
   ]
 string_column:   [
-    "Hello"
+    "Hello",
+    "Hello2"
   ]
 date_column:   [
-    #{@date_value}
+    #{@date_value},
+    #{@date_value + 1}
   ]
 datetime_column:   [
-    #{@datetime_value.strftime("%Y-%m-%d %H:%M:%S.%9N")}
+    #{@datetime_value.strftime("%Y-%m-%d %H:%M:%S.%9N")},
+    #{(@datetime_value + 1).strftime("%Y-%m-%d %H:%M:%S.%9N")}
   ]
 boolean_column:   [
-    false
+    false,
+    true
   ]
 bigint_column:   [
-    #{@bigint_value}
+    #{@bigint_value},
+    #{-@bigint_value - 1}
   ]
       RECORD_BATCH
     end