fluent-plugin-arrow 0.0.2 → 0.0.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: a1ed6b1b40e55501097a3a0455b2f2be93b3cf55d7e8cb831d4dc742f4007b97
-  data.tar.gz: fdcaed4462a919076c29f5947c3c81744b887e0099d2e5ad7a317e6845391da1
+  metadata.gz: 9325c65f679d81a53d202a11a842a6c3a82b1cc3d600203064fc1329eeb3ded2
+  data.tar.gz: 134a1bb09054b3feeed480efb0d7df558ad8a8131dbe3951d73395511ed2e843
 SHA512:
-  metadata.gz: 5f219215b1f8b8a57b2f9d79a6d1fdec39e10652c6d674a8d1390f7f3ce75d8752c15e89b268e36fb781bed44cfed2e128fd840b0978eb508c75dec1fc2b82a1
-  data.tar.gz: 2d473fbead6b6c12e361b045e7de325652b36c22d94685be289bc8a076c03f4dcb8b2c5270d922f76c0f462475d6846c96176abf4085abf24affaad0bc4b04fd
+  metadata.gz: c72b02f95969d7dd8cdbfe525b579afc23a16f8fae1fdd72ac095479f94514c890287f7b3c41822888965a4f9554b33b5ed49e23d92401cefbfee3f8f8feb9b0
+  data.tar.gz: 18eaeb2135ecc0339e8986db79545a772aeeef0a2499119bfd3b2194beb125c5be518d6bb86dd0b5caaae1c1d842f98926557219c8fd6586fa04e0ac51b34e2c

data/fluent-plugin-arrow.gemspec

@@ -3,7 +3,7 @@ $LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
 
 Gem::Specification.new do |spec|
   spec.name    = "fluent-plugin-arrow"
-  spec.version = "0.0.2"
+  spec.version = "0.0.3"
   spec.authors = ["joker1007"]
   spec.email   = ["kakyoin.hierophant@gmail.com"]
 

data/lib/fluent/plugin/buf_arrow_file.rb

@@ -0,0 +1,62 @@
+
+#
+# Copyright 2018- joker1007
+#
+# Licensed 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.
+
+require "arrow"
+require 'fluent/plugin/buf_file'
+require 'fluent/plugin/buffer/arrow_file_chunk'
+require 'fluent/plugin/arrow/field_wrapper'
+
+module Fluent
+  module Plugin
+    class ArrowFileBuffer < Fluent::Plugin::FileBuffer
+      Plugin.register_buffer('arrow_file', self)
+
+      config_param :schema, :array
+      config_param :arrow_format, :enum, list: [:arrow, :parquet], default: :arrow
+      config_param :row_group_chunk_size, :integer, default: 1024
+
+      attr_reader :arrow_schema
+
+      def configure(conf)
+        super
+
+        # [{"name" => foo1, "type" => "uint64"}, {"name" => foo2, "type" => "struct", "fields" => [{"name" => bar1, "type" => "string"}]}
+        @field_wrappers = @schema.each_with_object({}) do |field, h|
+          h[field["name"]] = Fluent::Plugin::Arrow::FieldWrapper.build(field)
+        end
+
+        @arrow_schema = ::Arrow::Schema.new(@field_wrappers.values.map(&:arrow_field))
+      end
+
+      def resume
+        return {}, []
+      end
+
+      def generate_chunk(metadata)
+        # FileChunk generates real path with unique_id
+        if @file_permission
+          chunk = Fluent::Plugin::Buffer::ArrowFileChunk.new(metadata, @path, :create, @arrow_schema, @field_wrappers, perm: @file_permission, chunk_size: @row_group_chunk_size, format: @arrow_format)
+        else
+          chunk = Fluent::Plugin::Buffer::ArrowFileChunk.new(metadata, @path, :create, @arrow_schema, @field_wrappers, chunk_size: @row_group_chunk_size, format: @arrow_format)
+        end
+
+        log.debug "Created new chunk", chunk_id: dump_unique_id_hex(chunk.unique_id), metadata: metadata
+
+        return chunk
+      end
+    end
+  end
+end

data/lib/fluent/plugin/buffer/arrow_buffer_string_builder.rb

@@ -0,0 +1,44 @@
+module Fluent
+  module Plugin
+    class Buffer
+      module ArrowBufferStringBuilder
+
+        private
+
+        def each_record(&block)
+          raise NotImplementedError
+        end
+
+        def build_arrow_buffer_string
+          count = 0
+          each_record do |record|
+            count += 1
+            record.each do |k, v|
+              @field_wrappers[k].append(v)
+            end
+          end
+          arrow_buf = ::Arrow::ResizableBuffer.new(bytesize * 1.2)
+
+          ::Arrow::BufferOutputStream.open(arrow_buf) do |output|
+            if @format == :parquet
+              Parquet::ArrowFileWriter.open(@schema, output) do |writer|
+                columns = @schema.fields.map do |f|
+                  ::Arrow::Column.new(f, @field_wrappers[f.name].finish)
+                end
+                table = ::Arrow::Table.new(@schema, columns)
+                writer.write_table(table, @chunk_size)
+              end
+            else
+              ::Arrow::RecordBatchFileWriter.open(output, @schema) do |writer|
+                record_batch = ::Arrow::RecordBatch.new(@schema, count, @field_wrappers.values.map(&:finish))
+                writer.write_record_batch(record_batch)
+              end
+            end
+          end
+
+          arrow_buf.data.to_s
+        end
+      end
+    end
+  end
+end

data/lib/fluent/plugin/buffer/arrow_file_chunk.rb

@@ -0,0 +1,63 @@
+#
+# Copyright 2018- joker1007
+#
+# Licensed 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.
+
+require 'arrow'
+require 'parquet'
+require 'fluent/msgpack_factory'
+require 'fluent/plugin/buffer/chunk'
+require 'fluent/plugin/buffer/file_chunk'
+require 'fluent/plugin/buffer/arrow_buffer_string_builder'
+require 'fluent/plugin/arrow/field_wrapper'
+
+module Fluent
+  module Plugin
+    class Buffer
+      class ArrowFileChunk < FileChunk
+        include ArrowBufferStringBuilder
+
+        def initialize(metadata, path, mode, schema, field_wrappers, perm: system_config.file_permission || FILE_PERMISSION, chunk_size: 1024, format: :arrow)
+          super(metadata, path, mode, perm: perm, compress: :text)
+          @schema = schema
+          @field_wrappers = field_wrappers
+          @chunk_size = chunk_size
+          @format = format
+        end
+
+        def read(**kwargs)
+          @chunk.seek(0, IO::SEEK_SET)
+          build_arrow_buffer_string
+        end
+
+        def open(**kwargs, &block)
+          @chunk.seek(0, IO::SEEK_SET)
+          val = StringIO.open(build_arrow_buffer_string, &block)
+          @chunk.seek(0, IO::SEEK_END) if self.staged?
+          val
+        end
+
+        def write_to(io, **kwargs)
+          @chunk.seek(0, IO::SEEK_SET)
+          io.write build_arrow_buffer_string
+        end
+
+        private
+
+        def each_record(&block)
+          Fluent::MessagePackFactory.engine_factory.unpacker(@chunk).each(&block)
+        end
+      end
+    end
+  end
+end

data/lib/fluent/plugin/buffer/arrow_memory_chunk.rb

@@ -18,19 +18,21 @@ require 'parquet'
 require 'fluent/msgpack_factory'
 require 'fluent/plugin/buffer/chunk'
 require 'fluent/plugin/buffer/memory_chunk'
+require 'fluent/plugin/buffer/arrow_buffer_string_builder'
 require 'fluent/plugin/arrow/field_wrapper'
 
 module Fluent
   module Plugin
     class Buffer
       class ArrowMemoryChunk < MemoryChunk
+        include ArrowBufferStringBuilder
+
         def initialize(metadata, schema, field_wrappers, chunk_size: 1024, format: :arrow)
           super(metadata, compress: :text)
           @schema = schema
           @field_wrappers = field_wrappers
           @chunk_size = chunk_size
           @format = format
-          @unpacker = Fluent::MessagePackFactory.engine_factory.unpacker
         end
 
         def read(**kwargs)
@@ -48,34 +50,8 @@ module Fluent
 
         private
 
-        def build_arrow_buffer_string
-          count = 0
-          @unpacker.feed_each(@chunk) do |record|
-            count += 1
-            record.each do |k, v|
-              @field_wrappers[k].append(v)
-            end
-          end
-          arrow_buf = ::Arrow::ResizableBuffer.new(@chunk_bytes * 1.2)
-
-          ::Arrow::BufferOutputStream.open(arrow_buf) do |output|
-            if @format == :parquet
-              Parquet::ArrowFileWriter.open(@schema, output) do |writer|
-                columns = @schema.fields.map do |f|
-                  ::Arrow::Column.new(f, @field_wrappers[f.name].finish)
-                end
-                table = ::Arrow::Table.new(@schema, columns)
-                writer.write_table(table, @chunk_size)
-              end
-            else
-              ::Arrow::RecordBatchFileWriter.open(output, @schema) do |writer|
-                record_batch = ::Arrow::RecordBatch.new(@schema, count, @field_wrappers.values.map(&:finish))
-                writer.write_record_batch(record_batch)
-              end
-            end
-          end
-
-          arrow_buf.data.to_s
+        def each_record(&block)
+          Fluent::MessagePackFactory.engine_factory.unpacker.feed_each(@chunk, &block)
         end
       end
     end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: fluent-plugin-arrow
 version: !ruby/object:Gem::Version
-  version: 0.0.2
+  version: 0.0.3
 platform: ruby
 authors:
 - joker1007
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2018-12-01 00:00:00.000000000 Z
+date: 2018-12-17 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: bundler
@@ -114,7 +114,10 @@ files:
 - Rakefile
 - fluent-plugin-arrow.gemspec
 - lib/fluent/plugin/arrow/field_wrapper.rb
+- lib/fluent/plugin/buf_arrow_file.rb
 - lib/fluent/plugin/buf_arrow_memory.rb
+- lib/fluent/plugin/buffer/arrow_buffer_string_builder.rb
+- lib/fluent/plugin/buffer/arrow_file_chunk.rb
 - lib/fluent/plugin/buffer/arrow_memory_chunk.rb
 - lib/fluent/plugin/formatter_arrow.rb
 - test/helper.rb