1# Licensed to the Apache Software Foundation (ASF) under one 2# or more contributor license agreements. See the NOTICE file 3# distributed with this work for additional information 4# regarding copyright ownership. The ASF licenses this file 5# to you under the Apache License, Version 2.0 (the 6# "License"); you may not use this file except in compliance 7# with the License. You may obtain a copy of the License at 8# 9# http://www.apache.org/licenses/LICENSE-2.0 10# 11# Unless required by applicable law or agreed to in writing, 12# software distributed under the License is distributed on an 13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 14# KIND, either express or implied. See the License for the 15# specific language governing permissions and limitations 16# under the License. 17 18class TestDatasetFileWriter < Test::Unit::TestCase 19 include Helper::Buildable 20 include Helper::Readable 21 22 def setup 23 omit("Arrow Dataset is required") unless defined?(ArrowDataset) 24 Dir.mktmpdir do |tmpdir| 25 @dir = tmpdir 26 @format = ArrowDataset::IPCFileFormat.new 27 @file_system = Arrow::LocalFileSystem.new 28 @path = File.join(@dir, "data.arrow") 29 @output = @file_system.open_output_stream(@path) 30 @schema = build_schema(visible: Arrow::BooleanDataType.new, 31 point: Arrow::UInt8DataType.new) 32 @writer = @format.open_writer(@output, 33 @file_system, 34 @path, 35 @schema, 36 @format.default_write_options) 37 yield 38 end 39 end 40 41 def test_write_record_batch 42 record_batch = build_record_batch( 43 visible: build_boolean_array([true, false, true]), 44 point: build_uint8_array([1, 2, 3])) 45 @writer.write_record_batch(record_batch) 46 @writer.finish 47 @output.close 48 read_table(@path) do |written_table| 49 assert_equal(Arrow::Table.new(record_batch.schema, 50 [record_batch]), 51 written_table) 52 end 53 end 54 55 def test_write_record_batch_reader 56 table = build_table(visible: build_boolean_array([true, false, true]), 57 point: build_uint8_array([1, 2, 3])) 58 @writer.write_record_batch_reader(Arrow::TableBatchReader.new(table)) 59 @writer.finish 60 @output.close 61 read_table(@path) do |written_table| 62 assert_equal(table, written_table) 63 end 64 end 65end 66