-- 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.

local lgi = require 'lgi'
local Arrow = lgi.Arrow

local output_path = arg[1] or "/tmp/batch.arrow";

local fields = {
  Arrow.Field.new("uint8",  Arrow.UInt8DataType.new()),
  Arrow.Field.new("uint16", Arrow.UInt16DataType.new()),
  Arrow.Field.new("uint32", Arrow.UInt32DataType.new()),
  Arrow.Field.new("uint64", Arrow.UInt64DataType.new()),
  Arrow.Field.new("int8",   Arrow.Int8DataType.new()),
  Arrow.Field.new("int16",  Arrow.Int16DataType.new()),
  Arrow.Field.new("int32",  Arrow.Int32DataType.new()),
  Arrow.Field.new("int64",  Arrow.Int64DataType.new()),
  Arrow.Field.new("float",  Arrow.FloatDataType.new()),
  Arrow.Field.new("double", Arrow.DoubleDataType.new()),
}
local schema = Arrow.Schema.new(fields)

local output = Arrow.FileOutputStream.new(output_path, false)
local writer = Arrow.RecordBatchFileWriter.new(output, schema)

function build_array(builder, values)
   for _, value in pairs(values) do
      builder:append(value)
   end
   return builder:finish()
end

local uints = {1, 2, 4, 8}
local ints = {1, -2, 4, -8}
local floats = {1.1, -2.2, 4.4, -8.8}
local columns = {
   build_array(Arrow.UInt8ArrayBuilder.new(), uints),
   build_array(Arrow.UInt16ArrayBuilder.new(), uints),
   build_array(Arrow.UInt32ArrayBuilder.new(), uints),
   build_array(Arrow.UInt64ArrayBuilder.new(), uints),
   build_array(Arrow.Int8ArrayBuilder.new(), ints),
   build_array(Arrow.Int16ArrayBuilder.new(), ints),
   build_array(Arrow.Int32ArrayBuilder.new(), ints),
   build_array(Arrow.Int64ArrayBuilder.new(), ints),
   build_array(Arrow.FloatArrayBuilder.new(), floats),
   build_array(Arrow.DoubleArrayBuilder.new(), floats),
}

local record_batch = Arrow.RecordBatch.new(schema, 4, columns)
writer:write_record_batch(record_batch)

local sliced_columns = {}
for i, column in pairs(columns) do
   sliced_columns[i] = column:slice(1, 3)
end
record_batch = Arrow.RecordBatch.new(schema, 3, sliced_columns)
writer:write_record_batch(record_batch)

writer:close()
output:close()