# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

module WriterHelper
  def convert_time_unit(red_arrow_time_unit)
    if red_arrow_time_unit.nick == "second"
      red_arrow_time_unit.nick.to_sym
    else
      :"#{red_arrow_time_unit.nick}second"
    end
  end

  def convert_type(red_arrow_type)
    case red_arrow_type
    when Arrow::NullDataType
      ArrowFormat::NullType.singleton
    when Arrow::BooleanDataType
      ArrowFormat::BooleanType.singleton
    when Arrow::Int8DataType
      ArrowFormat::Int8Type.singleton
    when Arrow::UInt8DataType
      ArrowFormat::UInt8Type.singleton
    when Arrow::Int16DataType
      ArrowFormat::Int16Type.singleton
    when Arrow::UInt16DataType
      ArrowFormat::UInt16Type.singleton
    when Arrow::Int32DataType
      ArrowFormat::Int32Type.singleton
    when Arrow::UInt32DataType
      ArrowFormat::UInt32Type.singleton
    when Arrow::Int64DataType
      ArrowFormat::Int64Type.singleton
    when Arrow::UInt64DataType
      ArrowFormat::UInt64Type.singleton
    when Arrow::FloatDataType
      ArrowFormat::Float32Type.singleton
    when Arrow::DoubleDataType
      ArrowFormat::Float64Type.singleton
    when Arrow::Date32DataType
      ArrowFormat::Date32Type.singleton
    when Arrow::Date64DataType
      ArrowFormat::Date64Type.singleton
    when Arrow::Time32DataType
      ArrowFormat::Time32Type.new(convert_time_unit(red_arrow_type.unit))
    when Arrow::Time64DataType
      ArrowFormat::Time64Type.new(convert_time_unit(red_arrow_type.unit))
    when Arrow::TimestampDataType
      ArrowFormat::TimestampType.new(convert_time_unit(red_arrow_type.unit),
                                     red_arrow_type.time_zone&.identifier)
    when Arrow::MonthIntervalDataType
      ArrowFormat::YearMonthIntervalType.singleton
    when Arrow::DayTimeIntervalDataType
      ArrowFormat::DayTimeIntervalType.singleton
    when Arrow::MonthDayNanoIntervalDataType
      ArrowFormat::MonthDayNanoIntervalType.singleton
    when Arrow::DurationDataType
      ArrowFormat::DurationType.new(convert_time_unit(red_arrow_type.unit))
    when Arrow::BinaryDataType
      ArrowFormat::BinaryType.singleton
    when Arrow::LargeBinaryDataType
      ArrowFormat::LargeBinaryType.singleton
    when Arrow::StringDataType
      ArrowFormat::UTF8Type.singleton
    when Arrow::LargeStringDataType
      ArrowFormat::LargeUTF8Type.singleton
    when Arrow::Decimal128DataType
      ArrowFormat::Decimal128Type.new(red_arrow_type.precision,
                                      red_arrow_type.scale)
    when Arrow::Decimal256DataType
      ArrowFormat::Decimal256Type.new(red_arrow_type.precision,
                                      red_arrow_type.scale)
    when Arrow::FixedSizeBinaryDataType
      ArrowFormat::FixedSizeBinaryType.new(red_arrow_type.byte_width)
    when Arrow::MapDataType
      ArrowFormat::MapType.new(convert_field(red_arrow_type.field),
                               red_arrow_type.keys_sorted?)
    when Arrow::ListDataType
      ArrowFormat::ListType.new(convert_field(red_arrow_type.field))
    when Arrow::LargeListDataType
      ArrowFormat::LargeListType.new(convert_field(red_arrow_type.field))
    when Arrow::FixedSizeListDataType
      ArrowFormat::FixedSizeListType.new(convert_field(red_arrow_type.field),
                                         red_arrow_type.list_size)
    when Arrow::StructDataType
      fields = red_arrow_type.fields.collect do |field|
        convert_field(field)
      end
      ArrowFormat::StructType.new(fields)
    when Arrow::DenseUnionDataType
      fields = red_arrow_type.fields.collect do |field|
        convert_field(field)
      end
      ArrowFormat::DenseUnionType.new(fields, red_arrow_type.type_codes)
    when Arrow::SparseUnionDataType
      fields = red_arrow_type.fields.collect do |field|
        convert_field(field)
      end
      ArrowFormat::SparseUnionType.new(fields, red_arrow_type.type_codes)
    when Arrow::DictionaryDataType
      @dictionary_id ||= 0
      dictionary_id = @dictionary_id
      @dictionary_id += 1
      index_type = convert_type(red_arrow_type.index_data_type)
      value_type = convert_type(red_arrow_type.value_data_type)
      ArrowFormat::DictionaryType.new(dictionary_id,
                                      index_type,
                                      value_type,
                                      red_arrow_type.ordered?)
    else
      raise "Unsupported type: #{red_arrow_type.inspect}"
    end
  end

  def convert_field(red_arrow_field)
    type = convert_type(red_arrow_field.data_type)
    ArrowFormat::Field.new(red_arrow_field.name,
                           type,
                           nullable: red_arrow_field.nullable?,
                           metadata: red_arrow_field.metadata)
  end

  def convert_buffer(buffer)
    return nil if buffer.nil?
    IO::Buffer.for(buffer.data.to_s.dup)
  end

  def convert_array(red_arrow_array)
    type = convert_type(red_arrow_array.value_data_type)
    case type
    when ArrowFormat::NullType
      type.build_array(red_arrow_array.size)
    when ArrowFormat::PrimitiveType
      type.build_array(red_arrow_array.size,
                       convert_buffer(red_arrow_array.null_bitmap),
                       convert_buffer(red_arrow_array.data_buffer))
    when ArrowFormat::VariableSizeBinaryType
      type.build_array(red_arrow_array.size,
                       convert_buffer(red_arrow_array.null_bitmap),
                       convert_buffer(red_arrow_array.offsets_buffer),
                       convert_buffer(red_arrow_array.data_buffer))
    when ArrowFormat::FixedSizeBinaryType
      type.build_array(red_arrow_array.size,
                       convert_buffer(red_arrow_array.null_bitmap),
                       convert_buffer(red_arrow_array.data_buffer))
    when ArrowFormat::VariableSizeListType
      type.build_array(red_arrow_array.size,
                       convert_buffer(red_arrow_array.null_bitmap),
                       convert_buffer(red_arrow_array.value_offsets_buffer),
                       convert_array(red_arrow_array.values_raw))
    when ArrowFormat::FixedSizeListType
      type.build_array(red_arrow_array.size,
                       convert_buffer(red_arrow_array.null_bitmap),
                       convert_array(red_arrow_array.values_raw))
    when ArrowFormat::StructType
      children = red_arrow_array.fields.collect do |red_arrow_field|
        convert_array(red_arrow_field)
      end
      type.build_array(red_arrow_array.size,
                       convert_buffer(red_arrow_array.null_bitmap),
                       children)
    when ArrowFormat::DenseUnionType
      types_buffer = convert_buffer(red_arrow_array.type_ids.data_buffer)
      offsets_buffer = convert_buffer(red_arrow_array.value_offsets.data_buffer)
      children = red_arrow_array.fields.collect do |red_arrow_field|
        convert_array(red_arrow_field)
      end
      type.build_array(red_arrow_array.size,
                       types_buffer,
                       offsets_buffer,
                       children)
    when ArrowFormat::SparseUnionType
      types_buffer = convert_buffer(red_arrow_array.type_ids.data_buffer)
      children = red_arrow_array.fields.collect do |red_arrow_field|
        convert_array(red_arrow_field)
      end
      type.build_array(red_arrow_array.size,
                       types_buffer,
                       children)
    when ArrowFormat::DictionaryType
      validity_buffer = convert_buffer(red_arrow_array.null_bitmap)
      indices_buffer = convert_buffer(red_arrow_array.indices.data_buffer)
      dictionary_array = convert_array(red_arrow_array.dictionary)
      dictionary = ArrowFormat::Dictionary.new(0, dictionary_array)
      type.build_array(red_arrow_array.size,
                       validity_buffer,
                       indices_buffer,
                       [dictionary])
    else
      raise "Unsupported array #{red_arrow_array.inspect}"
    end
  end

  def write(writer, *inputs)
    inputs.each_with_index do |input, i|
      case input
      when ArrowFormat::RecordBatch
        record_batch = input
      else
        red_arrow_array = input
        array = convert_array(red_arrow_array)
        red_arrow_field = Arrow::Field.new("value",
                                           red_arrow_array.value_data_type,
                                           true)
        fields = [convert_field(red_arrow_field)]
        schema = ArrowFormat::Schema.new(fields)
        record_batch = ArrowFormat::RecordBatch.new(schema,
                                                    array.size,
                                                    [array])
      end
      writer.start(record_batch.schema) if i.zero?
      writer.write_record_batch(record_batch)
    end
    writer.finish
  end

  def roundtrip(*inputs)
    Dir.mktmpdir do |tmp_dir|
      path = File.join(tmp_dir, "data.#{file_extension}")
      File.open(path, "wb") do |output|
        writer = writer_class.new(output)
        write(writer, *inputs)
      end
      # pp(read(path)) # debug
      data = File.open(path, "rb", &:read).freeze
      case file_extension
      when "arrow"
        format = :arrow_file
      else
        format = :arrow_streaming
      end
      table = Arrow::Table.load(Arrow::Buffer.new(data), format: format)
      if inputs[0].is_a?(Arrow::Array)
        [table.value.data_type, table.value.values]
      else
        table
      end
    end
  end
end

module WriterTests
  def test_custom_metadata_field
    field = ArrowFormat::Field.new("value",
                                   ArrowFormat::BooleanType.new,
                                   metadata: {
                                     "key1" => "value1",
                                     "key2" => "value2",
                                   })
    schema = ArrowFormat::Schema.new([field])
    column = convert_array(Arrow::BooleanArray.new([true, nil, false]))
    record_batch = ArrowFormat::RecordBatch.new(schema, 3, [column])
    table = roundtrip(record_batch)
    assert_equal({
                   "key1" => "value1",
                   "key2" => "value2",
                 },
                 table.schema.fields[0].metadata)
  end

  def test_custom_metadata_schema
    field = ArrowFormat::Field.new("value", ArrowFormat::BooleanType.new)
    schema = ArrowFormat::Schema.new([field],
                                     metadata: {
                                       "key1" => "value1",
                                       "key2" => "value2",
                                     })
    column = convert_array(Arrow::BooleanArray.new([true, nil, false]))
    record_batch = ArrowFormat::RecordBatch.new(schema, 3, [column])
    table = roundtrip(record_batch)
    assert_equal({
                   "key1" => "value1",
                   "key2" => "value2",
                 },
                 table.schema.metadata)
  end

  def test_custom_metadata_message_record_batch
    field = ArrowFormat::Field.new("value", ArrowFormat::BooleanType.new)
    schema = ArrowFormat::Schema.new([field])
    column = convert_array(Arrow::BooleanArray.new([true, nil, false]))
    record_batch = ArrowFormat::RecordBatch.new(schema, 3, [column],
                                                message_metadata: {
                                                  "key1" => "value1",
                                                  "key2" => "value2",
                                                })
    output = StringIO.new(+"".b)
    writer = writer_class.new(output)
    write(writer, record_batch)
    writer.finish
    reader = reader_class.new(output.string)
    assert_equal({
                   "key1" => "value1",
                   "key2" => "value2",
                 },
                 reader.first.message_metadata)
  end

  def test_null
    array = Arrow::NullArray.new(3)
    type, values = roundtrip(array)
    assert_equal(["null", [nil, nil, nil]],
                 [type.to_s, values])
  end

  def test_boolean
    array = Arrow::BooleanArray.new([true, nil, false])
    type, values = roundtrip(array)
    assert_equal(["bool", [true, nil, false]],
                 [type.to_s, values])
  end

  def test_int8
    array = Arrow::Int8Array.new([-128, nil, 127])
    type, values = roundtrip(array)
    assert_equal(["int8", [-128, nil, 127]],
                 [type.to_s, values])
  end

  def test_uint8
    array = Arrow::UInt8Array.new([0, nil, 255])
    type, values = roundtrip(array)
    assert_equal(["uint8", [0, nil, 255]],
                 [type.to_s, values])
  end

  def test_int16
    array = Arrow::Int16Array.new([-32768, nil, 32767])
    type, values = roundtrip(array)
    assert_equal(["int16", [-32768, nil, 32767]],
                 [type.to_s, values])
  end

  def test_uint16
    array = Arrow::UInt16Array.new([0, nil, 65535])
    type, values = roundtrip(array)
    assert_equal(["uint16", [0, nil, 65535]],
                 [type.to_s, values])
  end

  def test_int32
    array = Arrow::Int32Array.new([-2147483648, nil, 2147483647])
    type, values = roundtrip(array)
    assert_equal(["int32", [-2147483648, nil, 2147483647]],
                 [type.to_s, values])
  end

  def test_uint32
    array = Arrow::UInt32Array.new([0, nil, 4294967295])
    type, values = roundtrip(array)
    assert_equal(["uint32", [0, nil, 4294967295]],
                 [type.to_s, values])
  end

  def test_int64
    array = Arrow::Int64Array.new([
                                    -9223372036854775808,
                                    nil,
                                    9223372036854775807
                                  ])
    type, values = roundtrip(array)
    assert_equal([
                   "int64",
                   [
                     -9223372036854775808,
                     nil,
                     9223372036854775807
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_uint64
    array = Arrow::UInt64Array.new([0, nil, 18446744073709551615])
    type, values = roundtrip(array)
    assert_equal(["uint64", [0, nil, 18446744073709551615]],
                 [type.to_s, values])
  end

  def test_float32
    array = Arrow::FloatArray.new([-0.5, nil, 0.5])
    type, values = roundtrip(array)
    assert_equal(["float", [-0.5, nil, 0.5]],
                 [type.to_s, values])
  end

  def test_float64
    array = Arrow::DoubleArray.new([-0.5, nil, 0.5])
    type, values = roundtrip(array)
    assert_equal(["double", [-0.5, nil, 0.5]],
                 [type.to_s, values])
  end

  def test_date32
    date_2017_08_28 = 17406
    date_2025_12_09 = 20431
    array = Arrow::Date32Array.new([
                                     date_2017_08_28,
                                     nil,
                                     date_2025_12_09,
                                   ])
    type, values = roundtrip(array)
    assert_equal([
                   "date32[day]",
                   [Date.new(2017, 8, 28), nil, Date.new(2025, 12, 9)],
                 ],
                 [type.to_s, values])
  end

  def test_date64
    date_2017_08_28_00_00_00 = 1503878400000
    date_2025_12_10_00_00_00 = 1765324800000
    array = Arrow::Date64Array.new([
                                     date_2017_08_28_00_00_00,
                                     nil,
                                     date_2025_12_10_00_00_00,
                                   ])
    type, values = roundtrip(array)
    assert_equal([
                   "date64[ms]",
                   [
                     DateTime.new(2017, 8, 28, 0, 0, 0),
                     nil,
                     DateTime.new(2025, 12, 10, 0, 0, 0),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_time32_second
    time_00_00_10 = 10
    time_00_01_10 = 60 + 10
    array = Arrow::Time32Array.new(:second,
                                   [time_00_00_10, nil, time_00_01_10])
    type, values = roundtrip(array)
    assert_equal([
                   "time32[s]",
                   [
                     Arrow::Time.new(:second, time_00_00_10),
                     nil,
                     Arrow::Time.new(:second, time_00_01_10),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_time32_millisecond
    time_00_00_10_000 = 10 * 1000
    time_00_01_10_000 = (60 + 10) * 1000
    array = Arrow::Time32Array.new(:milli,
                                   [
                                     time_00_00_10_000,
                                     nil,
                                     time_00_01_10_000,
                                   ])
    type, values = roundtrip(array)
    assert_equal([
                   "time32[ms]",
                   [
                     Arrow::Time.new(:milli, time_00_00_10_000),
                     nil,
                     Arrow::Time.new(:milli, time_00_01_10_000),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_time64_microsecond
    time_00_00_10_000_000 = 10 * 1_000_000
    time_00_01_10_000_000 = (60 + 10) * 1_000_000
    array = Arrow::Time64Array.new(:micro,
                                   [
                                     time_00_00_10_000_000,
                                     nil,
                                     time_00_01_10_000_000,
                                   ])
    type, values = roundtrip(array)
    assert_equal([
                   "time64[us]",
                   [
                     Arrow::Time.new(:micro, time_00_00_10_000_000),
                     nil,
                     Arrow::Time.new(:micro, time_00_01_10_000_000),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_time64_nanosecond
    time_00_00_10_000_000_000 = 10 * 1_000_000_000
    time_00_01_10_000_000_000 = (60 + 10) * 1_000_000_000
    array = Arrow::Time64Array.new(:nano,
                                   [
                                     time_00_00_10_000_000_000,
                                     nil,
                                     time_00_01_10_000_000_000,
                                   ])
    type, values = roundtrip(array)
    assert_equal([
                   "time64[ns]",
                   [
                     Arrow::Time.new(:nano, time_00_00_10_000_000_000),
                     nil,
                     Arrow::Time.new(:nano, time_00_01_10_000_000_000),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_timestamp_second
    timestamp_2019_11_17_15_09_11 = 1574003351
    timestamp_2025_12_16_05_33_58 = 1765863238
    array = Arrow::TimestampArray.new(:second,
                                      [
                                        timestamp_2019_11_17_15_09_11,
                                        nil,
                                        timestamp_2025_12_16_05_33_58,
                                      ])
    type, values = roundtrip(array)
    assert_equal([
                   "timestamp[s]",
                   [
                     Time.at(timestamp_2019_11_17_15_09_11),
                     nil,
                     Time.at(timestamp_2025_12_16_05_33_58),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_timestamp_millisecond
    timestamp_2019_11_17_15_09_11 = 1574003351 * 1_000
    timestamp_2025_12_16_05_33_58 = 1765863238 * 1_000
    array = Arrow::TimestampArray.new(:milli,
                                      [
                                        timestamp_2019_11_17_15_09_11,
                                        nil,
                                        timestamp_2025_12_16_05_33_58,
                                      ])
    type, values = roundtrip(array)
    assert_equal([
                   "timestamp[ms]",
                   [
                     Time.at(timestamp_2019_11_17_15_09_11 / 1_000),
                     nil,
                     Time.at(timestamp_2025_12_16_05_33_58 / 1_000),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_timestamp_microsecond
    timestamp_2019_11_17_15_09_11 = 1574003351 * 1_000_000
    timestamp_2025_12_16_05_33_58 = 1765863238 * 1_000_000
    array = Arrow::TimestampArray.new(:micro,
                                      [
                                        timestamp_2019_11_17_15_09_11,
                                        nil,
                                        timestamp_2025_12_16_05_33_58,
                                      ])
    type, values = roundtrip(array)
    assert_equal([
                   "timestamp[us]",
                   [
                     Time.at(timestamp_2019_11_17_15_09_11 / 1_000_000),
                     nil,
                     Time.at(timestamp_2025_12_16_05_33_58 / 1_000_000),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_timestamp_nanosecond
    timestamp_2019_11_17_15_09_11 = 1574003351 * 1_000_000_000
    timestamp_2025_12_16_05_33_58 = 1765863238 * 1_000_000_000
    array = Arrow::TimestampArray.new(:nano,
                                      [
                                        timestamp_2019_11_17_15_09_11,
                                        nil,
                                        timestamp_2025_12_16_05_33_58,
                                      ])
    type, values = roundtrip(array)
    assert_equal([
                   "timestamp[ns]",
                   [
                     Time.at(timestamp_2019_11_17_15_09_11 / 1_000_000_000),
                     nil,
                     Time.at(timestamp_2025_12_16_05_33_58 / 1_000_000_000),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_timestamp_time_zone
    time_zone = "UTC"
    timestamp_2019_11_17_15_09_11 = 1574003351
    timestamp_2025_12_16_05_33_58 = 1765863238
    data_type = Arrow::TimestampDataType.new(:second, time_zone)
    array = Arrow::TimestampArray.new(data_type,
                                      [
                                        timestamp_2019_11_17_15_09_11,
                                        nil,
                                        timestamp_2025_12_16_05_33_58,
                                      ])
    type, values = roundtrip(array)
    assert_equal([
                   "timestamp[s, tz=#{time_zone}]",
                   [
                     Time.at(timestamp_2019_11_17_15_09_11),
                     nil,
                     Time.at(timestamp_2025_12_16_05_33_58),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_year_month_interval
    array = Arrow::MonthIntervalArray.new([0, nil, 100])
    type, values = roundtrip(array)
    assert_equal(["month_interval", [0, nil, 100]],
                 [type.to_s, values])
  end

  def test_day_time_interval
    array =
      Arrow::DayTimeIntervalArray.new([
                                        {day: 1, millisecond: 100},
                                        nil,
                                        {day: 3, millisecond: 300},
                                      ])
    type, values = roundtrip(array)
    assert_equal([
                   "day_time_interval",
                   [
                     {day: 1, millisecond: 100},
                     nil,
                     {day: 3, millisecond: 300},
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_month_day_nano_interval
    array =
      Arrow::MonthDayNanoIntervalArray.new([
                                             {
                                               month: 1,
                                               day: 1,
                                               nanosecond: 100,
                                             },
                                             nil,
                                             {
                                               month: 3,
                                               day: 3,
                                               nanosecond: 300,
                                             },
                                           ])
    type, values = roundtrip(array)
    assert_equal([
                   "month_day_nano_interval",
                   [
                     {
                       month: 1,
                       day: 1,
                       nanosecond: 100,
                     },
                     nil,
                     {
                       month: 3,
                       day: 3,
                       nanosecond: 300,
                     },
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_duration_second
    array = Arrow::DurationArray.new(:second, [0, nil, 100])
    type, values = roundtrip(array)
    assert_equal(["duration[s]", [0, nil, 100]],
                 [type.to_s, values])
  end

  def test_duration_millisecond
    array = Arrow::DurationArray.new(:milli, [0, nil, 100])
    type, values = roundtrip(array)
    assert_equal(["duration[ms]", [0, nil, 100]],
                 [type.to_s, values])
  end

  def test_duration_microsecond
    array = Arrow::DurationArray.new(:micro, [0, nil, 100])
    type, values = roundtrip(array)
    assert_equal(["duration[us]", [0, nil, 100]],
                 [type.to_s, values])
  end

  def test_duration_nanosecond
    array = Arrow::DurationArray.new(:nano, [0, nil, 100])
    type, values = roundtrip(array)
    assert_equal(["duration[ns]", [0, nil, 100]],
                 [type.to_s, values])
  end

  def test_binary
    array = Arrow::BinaryArray.new(["Hello".b, nil, "World".b])
    type, values = roundtrip(array)
    assert_equal(["binary", ["Hello".b, nil, "World".b]],
                 [type.to_s, values])
  end

  def test_large_binary
    array = Arrow::LargeBinaryArray.new(["Hello".b, nil, "World".b])
    type, values = roundtrip(array)
    assert_equal(["large_binary", ["Hello".b, nil, "World".b]],
                 [type.to_s, values])
  end

  def test_utf8
    array = Arrow::StringArray.new(["Hello", nil, "World"])
    type, values = roundtrip(array)
    assert_equal(["string", ["Hello", nil, "World"]],
                 [type.to_s, values])
  end

  def test_large_utf8
    array = Arrow::LargeStringArray.new(["Hello", nil, "World"])
    type, values = roundtrip(array)
    assert_equal(["large_string", ["Hello", nil, "World"]],
                 [type.to_s, values])
  end

  def test_fixed_size_binary
    data_type = Arrow::FixedSizeBinaryDataType.new(4)
    array = Arrow::FixedSizeBinaryArray.new(data_type,
                                            ["0124".b, nil, "abcd".b])
    type, values = roundtrip(array)
    assert_equal(["fixed_size_binary[4]", ["0124".b, nil, "abcd".b]],
                 [type.to_s, values])
  end

  def test_decimal128
    positive_small = "1.200"
    positive_large = ("1234567890" * 3) + "12345.678"
    negative_small = "-1.200"
    negative_large = "-" + ("1234567890" * 3) + "12345.678"
    array = Arrow::Decimal128Array.new({precision: 38, scale: 3},
                                       [
                                         positive_large,
                                         positive_small,
                                         nil,
                                         negative_small,
                                         negative_large,
                                       ])
    type, values = roundtrip(array)
    assert_equal([
                   "decimal128(38, 3)",
                   [
                     BigDecimal(positive_large),
                     BigDecimal(positive_small),
                     nil,
                     BigDecimal(negative_small),
                     BigDecimal(negative_large),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_decimal256
    positive_small = "1.200"
    positive_large = ("1234567890" * 7) + "123.456"
    negative_small = "-1.200"
    negative_large = "-" + ("1234567890" * 7) + "123.456"
    array = Arrow::Decimal256Array.new({precision: 76, scale: 3},
                                       [
                                         positive_large,
                                         positive_small,
                                         nil,
                                         negative_small,
                                         negative_large,
                                       ])
    type, values = roundtrip(array)
    assert_equal([
                   "decimal256(76, 3)",
                   [
                     BigDecimal(positive_large),
                     BigDecimal(positive_small),
                     nil,
                     BigDecimal(negative_small),
                     BigDecimal(negative_large),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_list
    data_type = Arrow::ListDataType.new(name: "count", type: :int8)
    array = Arrow::ListArray.new(data_type,
                                 [[-128, 127], nil, [-1, 0, 1]])
    type, values = roundtrip(array)
    assert_equal([
                   "list<count: int8>",
                   [[-128, 127], nil, [-1, 0, 1]],
                 ],
                 [type.to_s, values])
  end

  def test_large_list
    data_type = Arrow::LargeListDataType.new(name: "count",
                                             type: :int8)
    array = Arrow::LargeListArray.new(data_type,
                                      [[-128, 127], nil, [-1, 0, 1]])
    type, values = roundtrip(array)
    assert_equal([
                   "large_list<count: int8>",
                   [[-128, 127], nil, [-1, 0, 1]],
                 ],
                 [type.to_s, values])
  end

  def test_fixed_size_list
    data_type = Arrow::FixedSizeListDataType.new({
                                                   name: "count",
                                                   type: :int8,
                                                 },
                                                 2)
    array = Arrow::FixedSizeListArray.new(data_type,
                                          [[-128, 127], nil, [-1, 1]])
    type, values = roundtrip(array)
    assert_equal([
                   "fixed_size_list<count: int8>[2]",
                   [[-128, 127], nil, [-1, 1]],
                 ],
                 [type.to_s, values])
  end

  def test_map
    data_type = Arrow::MapDataType.new(:string, :int8)
    array = Arrow::MapArray.new(data_type,
                                [
                                  {"a" => -128, "b" => 127},
                                  nil,
                                  {"c" => nil},
                                ])
    type, values = roundtrip(array)
    assert_equal([
                   "map<string, int8>",
                   [
                     {"a" => -128, "b" => 127},
                     nil,
                     {"c" => nil},
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_struct
    data_type = Arrow::StructDataType.new(count: :int8,
                                          visible: :boolean)
    array = Arrow::StructArray.new(data_type,
                                   [[-128, nil], nil, [nil, true]])
    type, values = roundtrip(array)
    assert_equal([
                   "struct<count: int8, visible: bool>",
                   [
                     {"count" => -128, "visible" => nil},
                     nil,
                     {"count" => nil, "visible" => true},
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_dense_union
    fields = [
      Arrow::Field.new("number", :int8),
      Arrow::Field.new("text", :string),
    ]
    type_ids = [11, 13]
    data_type = Arrow::DenseUnionDataType.new(fields, type_ids)
    types = Arrow::Int8Array.new([11, 13, 11, 13, 13])
    value_offsets = Arrow::Int32Array.new([0, 0, 1, 1, 2])
    children = [
      Arrow::Int8Array.new([1, nil]),
      Arrow::StringArray.new(["a", "b", "c"])
    ]
    array = Arrow::DenseUnionArray.new(data_type,
                                       types,
                                       value_offsets,
                                       children)
    type, values = roundtrip(array)
    assert_equal([
                   "dense_union<number: int8=11, text: string=13>",
                   [1, "a", nil, "b", "c"],
                 ],
                 [type.to_s, values])
  end

  def test_sparse_union
    fields = [
      Arrow::Field.new("number", :int8),
      Arrow::Field.new("text", :string),
    ]
    type_ids = [11, 13]
    data_type = Arrow::SparseUnionDataType.new(fields, type_ids)
    types = Arrow::Int8Array.new([11, 13, 11, 13, 11])
    children = [
      Arrow::Int8Array.new([1, nil, nil, nil, 5]),
      Arrow::StringArray.new([nil, "b", nil, "d", nil])
    ]
    array = Arrow::SparseUnionArray.new(data_type, types, children)
    type, values = roundtrip(array)
    assert_equal([
                   "sparse_union<number: int8=11, text: string=13>",
                   [1, "b", nil, "d", 5],
                 ],
                 [type.to_s, values])
  end

  def test_dictionary
    values = ["a", "b", "c", nil, "a"]
    string_array = Arrow::StringArray.new(values)
    array = string_array.dictionary_encode
    type, values = roundtrip(array)
    assert_equal([
                   "dictionary<values=string, " +
                   "indices=int32, " +
                   "ordered=0>",
                   ["a", "b", "c", nil, "a"],
                 ],
                 [type.to_s, values])
  end
end

module FileWriterTests
  def test_custom_metadata_footer
    output = StringIO.new(+"".b)
    writer = writer_class.new(output)
    field = ArrowFormat::Field.new("value", ArrowFormat::BooleanType.new)
    schema = ArrowFormat::Schema.new([field])
    writer.start(schema)
    metadata = {
      "key1" => "value1",
      "key2" => "value2",
    }
    writer.finish(metadata)
    buffer = Arrow::Buffer.new(output.string)
    Arrow::BufferInputStream.open(buffer) do |input|
      reader = Arrow::RecordBatchFileReader.new(input)
      assert_equal(metadata, reader.metadata)
    end
  end
end

module StreamingWriterTests
  def test_custom_metadata_message_schema
    field = ArrowFormat::Field.new("value", ArrowFormat::BooleanType.new)
    schema = ArrowFormat::Schema.new([field],
                                     message_metadata: {
                                       "key1" => "value1",
                                       "key2" => "value2",
                                     })
    column = convert_array(Arrow::BooleanArray.new([true, nil, false]))
    record_batch = ArrowFormat::RecordBatch.new(schema, 3, [column])
    output = StringIO.new(+"".b)
    writer = writer_class.new(output)
    write(writer, record_batch)
    writer.finish
    input = StringIO.new(output.string)
    reader = reader_class.new(input)
    assert_equal({
                   "key1" => "value1",
                   "key2" => "value2",
                 },
                 reader.schema.message_metadata)
  end
end

module WriterDictionaryDeltaTests
  def build_schema(value_type)
    index_type = ArrowFormat::Int32Type.singleton
    dictionary_id = 1
    ordered = false
    type = ArrowFormat::DictionaryType.new(dictionary_id,
                                           index_type,
                                           value_type,
                                           ordered)
    field = ArrowFormat::Field.new("value", type)
    ArrowFormat::Schema.new([field])
  end

  def build_dictionary_array(type, indices, dictionaries)
    indices_buffer = IO::Buffer.for(indices.pack("l<*"))
    ArrowFormat::DictionaryArray.new(type,
                                     indices.size,
                                     nil,
                                     indices_buffer,
                                     dictionaries)
  end

  def build_record_batches(red_arrow_value_type, values1, values2)
    value_type = convert_type(red_arrow_value_type)
    schema = build_schema(value_type)
    type = schema.fields[0].type

    dictionary_id = 1

    # The first record batch with new dictionary.
    raw_dictionary = values1.uniq
    red_arrow_dictionary =
      red_arrow_value_type.build_array(raw_dictionary)
    dictionary_array = convert_array(red_arrow_dictionary)
    dictionary =
      ArrowFormat::Dictionary.new(dictionary_id, dictionary_array,
                                  message_metadata: {
                                    "key1" => "value1",
                                    "key2" => "value2",
                                  })
    indices1 = values1.collect do |value|
      raw_dictionary.index(value)
    end
    array1 = build_dictionary_array(type, indices1, [dictionary])
    record_batch =
      ArrowFormat::RecordBatch.new(schema, array1.size, [array1])

    if chunked_dictionaries?
      # The second record batch with the first dictionary and
      # a delta dictionary.
      raw_dictionary_delta = (values2.uniq - raw_dictionary)
      raw_dictionary_more = raw_dictionary + raw_dictionary_delta
      red_arrow_dictionary_delta =
        red_arrow_value_type.build_array(raw_dictionary_delta)
      dictionary_array_delta = convert_array(red_arrow_dictionary_delta)
      dictionary_delta =
        ArrowFormat::Dictionary.new(dictionary_id, dictionary_array_delta)
      indices2 = values2.collect do |value|
        raw_dictionary_more.index(value)
      end
      array2 = build_dictionary_array(type,
                                      indices2,
                                      [dictionary, dictionary_delta])
    else
      # The second record batch with the combined dictionary.
      raw_dictionary_more = raw_dictionary | values2.uniq
      red_arrow_dictionary_more =
        red_arrow_value_type.build_array(raw_dictionary_more)
      dictionary_array_more = convert_array(red_arrow_dictionary_more)
      dictionary_more = ArrowFormat::Dictionary.new(dictionary_id,
                                                    dictionary_array_more)
      indices2 = values2.collect do |value|
        raw_dictionary_more.index(value)
      end
      array2 = build_dictionary_array(type,
                                      indices2,
                                      [dictionary_more])
    end
    record_batch_delta =
      ArrowFormat::RecordBatch.new(schema, array2.size, [array2])

    [record_batch, record_batch_delta]
  end

  def roundtrip(value_type, values1, values2)
    record_batches = build_record_batches(value_type, values1, values2)
    GC.start
    table = super(*record_batches)
    [table.value.data_type, table.value.values]
  end

  def test_boolean
    value_type = Arrow::BooleanDataType.new
    values1 = [true, true]
    values2 = [false, true, false]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=bool, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_int8
    value_type = Arrow::Int8DataType.new
    values1 = [-128, 0, -128]
    values2 = [127, -128, 0, 127]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=int8, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_uint8
    value_type = Arrow::UInt8DataType.new
    values1 = [1, 0, 1]
    values2 = [255, 0, 1, 255]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=uint8, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_int16
    value_type = Arrow::Int16DataType.new
    values1 = [-32768, 0, -32768]
    values2 = [32767, -32768, 0, 32767]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=int16, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_uint16
    value_type = Arrow::UInt16DataType.new
    values1 = [1, 0, 1]
    values2 = [65535, 0, 1, 65535]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=uint16, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_int32
    value_type = Arrow::Int32DataType.new
    values1 = [-2147483648, 0, -2147483648]
    values2 = [2147483647, -2147483648, 0, 2147483647]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=int32, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_uint32
    value_type = Arrow::UInt32DataType.new
    values1 = [1, 0, 1]
    values2 = [4294967295, 0, 1, 4294967295]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=uint32, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_int64
    value_type = Arrow::Int64DataType.new
    values1 = [
      -9223372036854775808,
      0,
      -9223372036854775808,
    ]
    values2 = [
      9223372036854775807,
      -9223372036854775808,
      0,
      9223372036854775807,
    ]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=int64, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_uint64
    value_type = Arrow::UInt64DataType.new
    values1 = [1, 0, 1]
    values2 = [
      18446744073709551615,
      0,
      1,
      18446744073709551615,
    ]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=uint64, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_float32
    value_type = Arrow::FloatDataType.new
    values1 = [-0.5, 0.0, -0.5]
    values2 = [0.5, -0.5, 0.0, 0.5]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=float, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_float64
    value_type = Arrow::DoubleDataType.new
    values1 = [-0.5, 0.0, -0.5]
    values2 = [0.5, -0.5, 0.0, 0.5]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=double, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_date32
    date_2017_08_28 = 17406
    date_2025_12_09 = 20431
    value_type = Arrow::Date32DataType.new
    values1 = [date_2017_08_28, date_2017_08_28]
    values2 = [date_2025_12_09, date_2017_08_28, date_2025_12_09]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=date32[day], " +
                   "indices=int32, " +
                   "ordered=0>",
                   [
                     Date.new(2017, 8, 28),
                     Date.new(2017, 8, 28),
                     Date.new(2025, 12, 9),
                     Date.new(2017, 8, 28),
                     Date.new(2025, 12, 9),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_date64
    date_2017_08_28_00_00_00 = 1503878400000
    date_2025_12_10_00_00_00 = 1765324800000
    value_type = Arrow::Date64DataType.new
    values1 = [date_2017_08_28_00_00_00, date_2017_08_28_00_00_00]
    values2 = [
      date_2025_12_10_00_00_00,
      date_2017_08_28_00_00_00,
      date_2025_12_10_00_00_00,
    ]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=date64[ms], " +
                   "indices=int32, " +
                   "ordered=0>",
                   [
                     DateTime.new(2017, 8, 28),
                     DateTime.new(2017, 8, 28),
                     DateTime.new(2025, 12, 10),
                     DateTime.new(2017, 8, 28),
                     DateTime.new(2025, 12, 10),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_time32
    time_00_00_10 = 10
    time_00_01_10 = 60 + 10
    value_type = Arrow::Time32DataType.new(:second)
    values1 = [time_00_00_10, time_00_00_10]
    values2 = [time_00_01_10, time_00_00_10, time_00_01_10]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=time32[s], " +
                   "indices=int32, " +
                   "ordered=0>",
                   [
                     Arrow::Time.new(:second, time_00_00_10),
                     Arrow::Time.new(:second, time_00_00_10),
                     Arrow::Time.new(:second, time_00_01_10),
                     Arrow::Time.new(:second, time_00_00_10),
                     Arrow::Time.new(:second, time_00_01_10),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_time64
    time_00_00_10_000_000 = 10 * 1_000_000
    time_00_01_10_000_000 = (60 + 10) * 1_000_000
    value_type = Arrow::Time64DataType.new(:micro)
    values1 = [time_00_00_10_000_000, time_00_00_10_000_000]
    values2 = [
      time_00_01_10_000_000,
      time_00_00_10_000_000,
      time_00_01_10_000_000,
    ]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=time64[us], " +
                   "indices=int32, " +
                   "ordered=0>",
                   [
                     Arrow::Time.new(:micro, time_00_00_10_000_000),
                     Arrow::Time.new(:micro, time_00_00_10_000_000),
                     Arrow::Time.new(:micro, time_00_01_10_000_000),
                     Arrow::Time.new(:micro, time_00_00_10_000_000),
                     Arrow::Time.new(:micro, time_00_01_10_000_000),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_timestamp
    timestamp_2019_11_17_15_09_11 = 1574003351
    timestamp_2025_12_16_05_33_58 = 1765863238
    value_type = Arrow::TimestampDataType.new(:second)
    values1 = [
      timestamp_2019_11_17_15_09_11,
      timestamp_2019_11_17_15_09_11,
    ]
    values2 = [
      timestamp_2025_12_16_05_33_58,
      timestamp_2019_11_17_15_09_11,
      timestamp_2025_12_16_05_33_58,
    ]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=timestamp[s], " +
                   "indices=int32, " +
                   "ordered=0>",
                   [
                     Time.at(timestamp_2019_11_17_15_09_11),
                     Time.at(timestamp_2019_11_17_15_09_11),
                     Time.at(timestamp_2025_12_16_05_33_58),
                     Time.at(timestamp_2019_11_17_15_09_11),
                     Time.at(timestamp_2025_12_16_05_33_58),
                   ],
                 ],
                 [type.to_s, values])
  end

  def test_year_month_interval
    value_type = Arrow::MonthIntervalDataType.new
    values1 = [100, 0, 100]
    values2 = [1000, 100, 0, 1000]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=month_interval, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_day_time_interval
    value_type = Arrow::DayTimeIntervalDataType.new
    values1 = [
      {day: 1, millisecond: 100},
      {day: 1, millisecond: 100},
    ]
    values2 = [
      {day: 3, millisecond: 300},
      {day: 1, millisecond: 100},
      {day: 3, millisecond: 300},
    ]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=day_time_interval, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_month_day_nano_interval
    value_type = Arrow::MonthDayNanoIntervalDataType.new
    values1 = [
      {month: 1, day: 1, nanosecond: 100},
      {month: 1, day: 1, nanosecond: 100},
    ]
    values2 = [
      {month: 3, day: 3, nanosecond: 300},
      {month: 1, day: 1, nanosecond: 100},
      {month: 3, day: 3, nanosecond: 300},
    ]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=month_day_nano_interval, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_duration
    value_type = Arrow::DurationDataType.new(:second)
    values1 = [100, 0, 100]
    values2 = [1000, 100, 0, 1000]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=duration[s], " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_binary
    value_type = Arrow::BinaryDataType.new
    values1 = ["ab".b, "c".b, "ab".b]
    values2 = ["c".b, "de".b, "ab".b, "de".b]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=binary, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_large_binary
    value_type = Arrow::LargeBinaryDataType.new
    values1 = ["ab".b, "c".b, "ab".b]
    values2 = ["c".b, "de".b, "ab".b, "de".b]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=large_binary, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_utf8
    value_type = Arrow::StringDataType.new
    values1 = ["ab", "c", "ab"]
    values2 = ["c", "de", "ab", "de"]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=string, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_large_utf8
    value_type = Arrow::LargeStringDataType.new
    values1 = ["ab", "c", "ab"]
    values2 = ["c", "de", "ab", "de"]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=large_string, " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_fixed_size_binary
    value_type = Arrow::FixedSizeBinaryDataType.new(2)
    values1 = ["ab".b, "cd".b, "ab".b]
    values2 = ["ef".b, "cd".b, "ab".b, "ef".b]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=fixed_size_binary[2], " +
                   "indices=int32, " +
                   "ordered=0>",
                   values1 + values2,
                 ],
                 [type.to_s, values])
  end

  def test_decimal128
    positive_small = "1.200"
    positive_large = ("1234567890" * 3) + "12345.678"
    negative_small = "-1.200"
    negative_large = "-" + ("1234567890" * 3) + "12345.678"
    value_type = Arrow::Decimal128DataType.new(precision: 38,
                                               scale: 3)
    values1 = [positive_small, negative_small, positive_small]
    values2 = [
      positive_large,
      positive_small,
      negative_small,
      positive_large,
      negative_large,
    ]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=decimal128(38, 3), " +
                   "indices=int32, " +
                   "ordered=0>",
                   (values1 + values2).collect {|v| BigDecimal(v)},
                 ],
                 [type.to_s, values])
  end

  def test_decimal256
    positive_small = "1.200"
    positive_large = ("1234567890" * 7) + "123.456"
    negative_small = "-1.200"
    negative_large = "-" + ("1234567890" * 7) + "123.456"
    value_type = Arrow::Decimal256DataType.new(precision: 76,
                                               scale: 3)
    values1 = [positive_small, negative_small, positive_small]
    values2 = [
      positive_large,
      positive_small,
      negative_small,
      positive_large,
      negative_large,
    ]
    type, values = roundtrip(value_type, values1, values2)
    assert_equal([
                   "dictionary<values=decimal256(76, 3), " +
                   "indices=int32, " +
                   "ordered=0>",
                   (values1 + values2).collect {|v| BigDecimal(v)},
                 ],
                 [type.to_s, values])
  end
end

class TestFileWriter < Test::Unit::TestCase
  include WriterHelper

  def file_extension
    "arrow"
  end

  def writer_class
    ArrowFormat::FileWriter
  end

  def reader_class
    ArrowFormat::FileReader
  end

  def read(path)
    File.open(path, "rb") do |input|
      reader = reader_class.new(input)
      reader.collect do |record_batch|
        record_batch.to_h.tap do |hash|
          hash.each do |key, value|
            hash[key] = value.to_a
          end
        end
      end
    end
  end

  sub_test_case("Basic") do
    include WriterTests
    include FileWriterTests
  end

  sub_test_case("Dictionary: delta") do
    include WriterDictionaryDeltaTests

    def chunked_dictionaries?
      true
    end
  end

  sub_test_case("Dictionary: delta: slice") do
    include WriterDictionaryDeltaTests

    def chunked_dictionaries?
      false
    end
  end
end

class TestStreamingWriter < Test::Unit::TestCase
  include WriterHelper

  def file_extension
    "arrows"
  end

  def writer_class
    ArrowFormat::StreamingWriter
  end

  def reader_class
    ArrowFormat::StreamingReader
  end

  def read(path)
    File.open(path, "rb") do |input|
      reader = reader_class.new(input)
      reader.collect do |record_batch|
        record_batch.to_h.tap do |hash|
          hash.each do |key, value|
            hash[key] = value.to_a
          end
        end
      end
    end
  end

  sub_test_case("Basic") do
    include WriterTests
    include StreamingWriterTests
  end

  sub_test_case("Dictionary: delta") do
    include WriterDictionaryDeltaTests

    def chunked_dictionaries?
      true
    end
  end

  sub_test_case("Dictionary: delta: slice") do
    include WriterDictionaryDeltaTests

    def chunked_dictionaries?
      false
    end
  end
end
