diff --git a/ruby/red-arrow/lib/arrow/csv-loader.rb b/ruby/red-arrow/lib/arrow/csv-loader.rb index c63ab89eb099b..5e244e2334afb 100644 --- a/ruby/red-arrow/lib/arrow/csv-loader.rb +++ b/ruby/red-arrow/lib/arrow/csv-loader.rb @@ -93,10 +93,17 @@ def reader_options @options.each do |key, value| case key when :headers - if value - options.n_header_rows = 1 + case value + when ::Array + options.column_names = value + when String + return nil else - options.n_header_rows = 0 + if value + options.generate_column_names = false + else + options.generate_column_names = true + end end when :column_types value.each do |name, type| diff --git a/ruby/red-arrow/test/test-csv-loader.rb b/ruby/red-arrow/test/test-csv-loader.rb index 96de9c8634540..42eef7f67ba57 100644 --- a/ruby/red-arrow/test/test-csv-loader.rb +++ b/ruby/red-arrow/test/test-csv-loader.rb @@ -121,6 +121,81 @@ def load_csv(data, options) Arrow::CSVLoader.load(data, options) end + sub_test_case(":headers") do + test("true") do + values = Arrow::StringArray.new(["a", "b", "c"]) + assert_equal(Arrow::Table.new(value: values), + load_csv(<<-CSV, headers: true)) +value +a +b +c + CSV + end + + test(":first_line") do + values = Arrow::StringArray.new(["a", "b", "c"]) + assert_equal(Arrow::Table.new(value: values), + load_csv(<<-CSV, headers: :first_line)) +value +a +b +c + CSV + end + + test("truthy") do + values = Arrow::StringArray.new(["a", "b", "c"]) + assert_equal(Arrow::Table.new(value: values), + load_csv(<<-CSV, headers: 0)) +value +a +b +c + CSV + end + + test("Array of column names") do + values = Arrow::StringArray.new(["a", "b", "c"]) + assert_equal(Arrow::Table.new(column: values), + load_csv(<<-CSV, headers: ["column"])) +a +b +c + CSV + end + + test("false") do + values = Arrow::StringArray.new(["a", "b", "c"]) + assert_equal(Arrow::Table.new(f0: values), + load_csv(<<-CSV, headers: false)) +a +b +c + CSV + end + + test("nil") do + values = Arrow::StringArray.new(["a", "b", "c"]) + assert_equal(Arrow::Table.new(f0: values), + load_csv(<<-CSV, headers: nil)) +a +b +c + CSV + end + + test("string") do + values = Arrow::StringArray.new(["a", "b", "c"]) + assert_equal(Arrow::Table.new(column: values), + load_csv(<<-CSV, headers: "column")) +a +b +c + CSV + end + end + test(":column_types") do assert_equal(Arrow::Table.new(:count => Arrow::UInt16Array.new([1, 2, 4])), load_csv(<<-CSV, column_types: {count: :uint16}))