fluent-plugin-s3-arrow 0.1.0 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8d4d40af870606fbd11cf131fa371e29e45013281f7b23deae6e30b4f260ecf4
-  data.tar.gz: dc911408cac26e521a32eb50a56e544d18b1f36e2ea83b20fbf7d08c357b7a3e
+  metadata.gz: 896ac4c81632aa2283e4e3d299138bc1687529f550ea525fd6665dc3af8ddc4e
+  data.tar.gz: 6d9bd3eb2077f25992b5409ea2ca070f6eb4b1337c63ee305d83de50d910f8f9
 SHA512:
-  metadata.gz: 26a33661cab8c5d524d160c727eab8da58c43973ea708d8e1aad45eade08a220612a1620404f9647628d9e6e81cb94e8ccb2a08a59019317d1a582e5254831fc
-  data.tar.gz: c24341e83a739023ab2bfc9b151ab9d1d6cedbba1f25bce63ae1455c426d724e20807cbed89a6663d18840f96196b266446772946176a0a5fdd7d4629f31b028
+  metadata.gz: e6a29a073d6495fe0a1e038eaeebf694d40866d904bfa892581b31f05a4d02d8a7d15409b1bffe321dd9b1ddfdc65515283b195c53ba034491efd88b42569d85
+  data.tar.gz: e9fff86cae0c8a00003d1ea5b45ab7ee53293f7b5338c82b2bad93257ab2844c6a4a0f8ab785d5fbb5d8f707a8cafb7b83e3e62c156555ab66dfed6c59dd127e

data/.github/workflows/bench.yml

@@ -13,7 +13,11 @@ jobs:
     - name: 'bench'
       id: 'bench'
       run: |
-        docker run fluent-plugin-s3-arrow /bin/bash -c "benchmark-driver benchmark/compress.yml"
+        docker run fluent-plugin-s3-arrow /bin/bash -c "benchmark-driver benchmark/compress.yml -o markdown" | tee benchmark.md
+    - uses: actions/upload-artifact@v2
+      with:
+        name: benchmark.md
+        path: benchmark.md
     - name: 'bench memory'
       id: 'bench_mem'
       run: |

data/README.md

@@ -1,6 +1,6 @@
 # fluent-plugin-s3-arrow
 
-Extends the [fluent-plugin-s3](https://github.com/fluent/fluent-plugin-s3) (de)compression algorithm to enable red-arrow compression.
+Extends the [fluent-plugin-s3](https://github.com/fluent/fluent-plugin-s3) compression algorithm to enable [red-arrow](https://github.com/apache/arrow/tree/master/ruby/red-arrow) compression.
 
 ## Installation
 
@@ -33,7 +33,7 @@ $ bundle
 
 ## Configuration
 
-Example of fluent-plugin-s3 configuration.
+Example of fluent-plugin-s3-arrow configuration.
 
 ```
 <match pattern>
@@ -47,14 +47,74 @@ Example of fluent-plugin-s3 configuration.
 
   store_as arrow
   <arrow>
-    schema [
-      {"name": "test_string", "type": "string"},
-      {"name": "test_uint64", "type": "uint64"}
-    ]
+    format parquet
+    compression gzip
+
+    schema_from static
+    <static>
+      schema [
+        {"name": "test_string", "type": "string"},
+        {"name": "test_uint64", "type": "uint64"}
+      ]
+    </static>
   </arrow>
 </match>
 ```
 
+### format and compression
+
+This plugin supports multiple columnar formats and compressions by using red-arrow. Valid settings are below.
+
+|  format  |  compression |
+| ---- | ---- |
+|  arrow | gzip, zstd |
+|  feather | zstd |
+|  parquet   | gzip, snappy, zstd |
+
+### schema
+
+Schema of columnar formats.
+#### schema_from static
+
+Set the schema statically.
+
+```
+schema_from static
+<static>
+  schema [
+    {"name": "test_string", "type": "string"},
+    {"name": "test_uint64", "type": "uint64"}
+  ]
+</static>
+```
+
+##### schema (required)
+
+An array containing the names and types of the fields.
+#### schema_from glue
+
+Retrieve the schema from the AWS Glue Data Catalog.
+
+```
+schema_from glue
+<glue>
+  catalog test_catalog
+  database test_db
+  table test_table
+</glue>
+```
+
+##### catalog
+
+The name of the data catalog for which to retrieve the definition. The default value is the same as the [AWS API CatalogId](https://docs.aws.amazon.com/glue/latest/webapi/API_GetTable.html).
+
+##### database
+
+The name of the database for which to retrieve the definition.  The default value is `default`.
+##### table (required)
+
+The name of the table for which to retrieve the definition.
+
 ## License
 
 Apache License, Version 2.0

data/benchmark/prelude.rb

@@ -18,11 +18,14 @@ ARROW_CONFIG = %[
   <arrow>
     format parquet
     compression gzip
-    schema [
-      {"name": "test_string",  "type": "string"},
-      {"name": "test_uint64",  "type": "uint64"},
-      {"name": "test_boolean", "type": "boolean"}
-    ]
+    schema_from static
+    <static>
+      schema [
+        {"name": "test_string",  "type": "string"},
+        {"name": "test_uint64",  "type": "uint64"},
+        {"name": "test_boolean", "type": "boolean"}
+      ]
+    </static>
   </arrow>
 ]
 

data/fluent-plugin-s3-arrow.gemspec

@@ -3,12 +3,12 @@ $LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
 
 Gem::Specification.new do |spec|
   spec.name    = "fluent-plugin-s3-arrow"
-  spec.version = "0.1.0"
+  spec.version = "0.2.0"
   spec.authors = ["kanga333"]
   spec.email   = ["e411z7t40w@gmail.com"]
 
-  spec.summary       = %q{Extends the fluent-plugin-s3 (de)compression algorithm to enable red-arrow compression.}
-  spec.description   = %q{Extends the fluent-plugin-s3 (de)compression algorithm to enable red-arrow compression.}
+  spec.summary       = %q{Extends the fluent-plugin-s3 compression algorithm to enable red-arrow compression.}
+  spec.description   = %q{Extends the fluent-plugin-s3 compression algorithm to enable red-arrow compression.}
   spec.homepage      = "https://github.com/red-data-tools/fluent-plugin-s3-arrow"
   spec.license       = "Apache-2.0"
 
@@ -24,8 +24,10 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "faker"
   spec.add_development_dependency "rake", "~> 12.0"
   spec.add_development_dependency "test-unit", "~> 3.0"
+  spec.add_development_dependency "test-unit-rr", "~> 1.0"
   spec.add_runtime_dependency "fluentd", [">= 0.14.10", "< 2"]
   spec.add_runtime_dependency "fluent-plugin-s3", ">= 1.0"
-  spec.add_runtime_dependency "red-arrow", ">= 1.0"
-  spec.add_runtime_dependency "red-parquet", ">= 1.0"
+  spec.add_runtime_dependency "red-arrow", ">= 2.0"
+  spec.add_runtime_dependency "red-parquet", ">= 2.0"
+  spec.add_runtime_dependency "aws-sdk-glue", ">= 1.20.0"
 end

data/lib/fluent-plugin-s3-arrow/schemas.rb

@@ -0,0 +1,7 @@
+require "arrow"
+require "fluent-plugin-s3-arrow/schemas/aws_glue"
+
+module FluentPluginS3Arrow
+  module Schemas
+  end
+end

data/lib/fluent-plugin-s3-arrow/schemas/aws_glue.rb

@@ -0,0 +1,115 @@
+require 'aws-sdk-glue'
+
+module FluentPluginS3Arrow
+  module Schemas
+    class AWSGlue
+      class Error < RuntimeError; end
+      class ConvertError < Error; end
+      class Field < Struct.new(:name, :type); end
+
+      def initialize(table_name, **options)
+        @table_name = table_name
+        @database_name = options.delete(:database_name) || "default"
+        @catalog_id = options.delete(:catalog_id)
+        @client = Aws::Glue::Client.new(options)
+      end 
+
+      def to_arrow
+        glue_schema = fetch_glue_schema
+        convert_to_arrow_schema(glue_schema)
+      end
+
+      private
+      def fetch_glue_schema
+        glue_table = @client.get_table({
+          catalog_id: @catalog_id,
+          database_name: @database_name,
+          name: @table_name
+        })
+        glue_table.table.storage_descriptor.columns
+      end
+
+      def convert_to_arrow_schema(glue_schema)
+        arrow_schema_description = glue_schema.map do |glue_field|
+          convert_to_arrow_field_description(glue_field)
+        end
+        Arrow::Schema.new(arrow_schema_description)
+      end
+
+      def convert_to_arrow_field_description(glue_field)
+        arrow_field = {name: glue_field.name}
+        case glue_field.type
+        when "boolean", "float", "double"
+          arrow_field[:type] = glue_field.type
+        when "tinyint"
+          arrow_field[:type] = "int8"
+        when "smallint"
+          arrow_field[:type] = "int16"
+        when "int"
+          arrow_field[:type] = "int32"
+        when "bigint"
+          arrow_field[:type] = "int64"
+        when /\Achar/,/\Avarchar/,"string"
+          arrow_field[:type] = "string"
+        when "binary"
+          arrow_field[:type] = "binary"
+        when "date"
+          arrow_field[:type] = "date32"
+        when /\Aarray/
+          arrow_field[:type] = "list"
+          arrow_field[:field] = parse_array(glue_field.type)
+        when /\Astruct/
+          arrow_field[:type] = "struct"
+          arrow_field[:fields] = parse_struct(glue_field.type)
+        else
+          # TODO: Need support for MAP, DECIMAL, TIMESTAMP type.
+          raise ConvertError, "Input type is not supported: #{glue_field.type}"
+        end
+        arrow_field
+      end
+
+      def parse_array(str)
+        matched = str.match(/\Aarray<(.*)>\z/)
+        raise ConvertError, "Parse error on array type: #{str}" if matched.nil?
+        convert_to_arrow_field_description(Field.new("", matched[1]))
+      end
+
+      def parse_struct(str)
+        fields = []
+        matched = str.match(/\Astruct<(.*)>\z/)
+        raise ConvertError, "Parse error on struct type: #{str}" if matched.nil?
+        each_struct_fields(matched[1]) do |name, type|
+          fields << convert_to_arrow_field_description(Field.new(name, type))
+        end
+        fields
+      end
+
+      def each_struct_fields(str)
+        start, nest = 0, 0
+        name = ""
+        str.each_char.with_index do |c, i|
+          case c
+          when ':'
+            if nest == 0
+              name = str[start...i]
+              start = i + 1
+            end
+          when '<'
+            nest += 1
+          when '>'
+            nest -= 1
+          when ','
+            if nest == 0
+              type = str[start...i]
+              yield(name, type)
+              start = i + 1
+            end
+          end
+        end
+        type = str[start..]
+        yield(name, type)
+      end
+
+    end
+  end
+end

data/lib/fluent/plugin/s3_compressor_arrow.rb

@@ -1,5 +1,6 @@
 require 'arrow'
 require 'parquet'
+require 'fluent-plugin-s3-arrow/schemas'
 
 module Fluent::Plugin
   class S3Output
@@ -12,11 +13,21 @@ module Fluent::Plugin
       }
 
       config_section :arrow, multi: false do
-        config_param :schema, :array
         config_param :format, :enum, list: [:arrow, :feather, :parquet], default: :arrow
         SUPPORTED_COMPRESSION = [:gzip, :snappy, :zstd]
         config_param :compression, :enum, list: SUPPORTED_COMPRESSION, default: nil
-        config_param :chunk_size, :integer, default: 1024
+        config_param :chunk_size, :integer, default: nil
+        config_param :schema_from, :enum, list: [:static, :glue], default: :static
+
+        config_section :static, multi: false do
+          config_param :schema, :array, default: nil
+        end
+
+        config_section :glue, multi: false do
+          config_param :catalog, :string, default: nil
+          config_param :database, :string, default: "default"
+          config_param :table, :string, default: nil
+        end
       end
 
       def configure(conf)
@@ -26,9 +37,8 @@ module Fluent::Plugin
           raise Fluent::ConfigError, "#{@arrow.format} unsupported with #{@arrow.format}"
         end
 
-        @schema = Arrow::Schema.new(@arrow.schema)
         @options = Arrow::JSONReadOptions.new
-        @options.schema = @schema
+        @options.schema = resolve_schema
         @options.unexpected_field_behavior = :ignore
       end
 
@@ -51,6 +61,21 @@ module Fluent::Plugin
           compression: @arrow.compression,
         )
       end
+
+      private
+
+      def resolve_schema
+        case @arrow.schema_from
+        when :static
+          Arrow::Schema.new(@arrow.static.schema)
+        when :glue
+          glue_schema = FluentPluginS3Arrow::Schemas::AWSGlue.new(@arrow.glue.table, {
+            catalog_id: @arrow.glue.catalog,
+            database_name: @arrow.glue.database,
+          })
+          glue_schema.to_arrow
+        end
+      end
     end
   end
 end

data/test/fluent-plugin-s3-arrow/schemas/test_aws_glue.rb

@@ -0,0 +1,51 @@
+require "helper"
+require "fluent-plugin-s3-arrow/schemas"
+
+class AWSGlueTest < Test::Unit::TestCase
+    def setup
+      stub(Aws::Glue::Client).new
+      @schema = FluentPluginS3Arrow::Schemas::AWSGlue.new('test')
+    end
+
+    
+    def test_to_arrow
+      stub(@schema).fetch_glue_schema{
+        [
+          Aws::Glue::Types::Column.new({name: "a", type: "boolean"}),
+          Aws::Glue::Types::Column.new({name: "b", type: "tinyint"}),
+          Aws::Glue::Types::Column.new({name: "c", type: "smallint"}),
+          Aws::Glue::Types::Column.new({name: "d", type: "int"}),
+          Aws::Glue::Types::Column.new({name: "e", type: "bigint"}),
+          Aws::Glue::Types::Column.new({name: "f", type: "float"}),
+          Aws::Glue::Types::Column.new({name: "g", type: "double"}),
+          Aws::Glue::Types::Column.new({name: "h", type: "char(1)"}),
+          Aws::Glue::Types::Column.new({name: "i", type: "varchar(1)"}),
+          Aws::Glue::Types::Column.new({name: "j", type: "string"}),
+          Aws::Glue::Types::Column.new({name: "k", type: "binary"}),
+          Aws::Glue::Types::Column.new({name: "l", type: "date"}),
+          Aws::Glue::Types::Column.new({name: "m", type: "array<array<string>>"}),
+          Aws::Glue::Types::Column.new({name: "n", type: "struct<p1:string,p2:struct<c1:string,c2:string>,p3:string>"})
+        ]
+      }
+      actual = @schema.to_arrow
+      expect = Arrow::Schema.new([
+        {name: "a", type: "boolean"},
+        {name: "b", type: "int8"},
+        {name: "c", type: "int16"},
+        {name: "d", type: "int32"},
+        {name: "e", type: "int64"},
+        {name: "f", type: "float"},
+        {name: "g", type: "double"},
+        {name: "h", type: "string"},
+        {name: "i", type: "string"},
+        {name: "j", type: "string"},
+        {name: "k", type: "binary"},
+        {name: "l", type: "date32"},
+        {name: "m", type: "list", field: {name: "", type: "list", field: {name: "", type: "string"}}},
+        {name: "n", type: "struct", fields: [{name: "p1", type: "string"},{name: "p2", type: "struct", fields: [{name: "c1", type: "string"},{name: "c2", type: "string"}]},{name: "p3", type: "string"}]}
+      ])
+      
+      assert_equal expect, actual
+      
+    end
+end

data/test/helper.rb

@@ -1,5 +1,6 @@
 $LOAD_PATH.unshift(File.expand_path("../../", __FILE__))
 require "test-unit"
+require "test/unit/rr"
 require "fluent/test"
 require "fluent/test/driver/output"
 require "fluent/test/helpers"

data/test/plugin/test_s3_compressor_arrow.rb

@@ -10,19 +10,19 @@ class S3OutputTest < Test::Unit::TestCase
     end
   
     S3_CONFIG = {"s3_bucket" => "test", "store_as" => "arrow"}
-    SCHEMA = [
+    SCHEMA = config_element("static", "", {"schema" => [
       {"name": "test_string", "type": "string"},
       {"name": "test_uint64", "type": "uint64"},
-    ]
-    CONFIG = config_element("ROOT", "", S3_CONFIG, [config_element("arrow", "", {"schema" => SCHEMA})])
+    ]})
+    ARROW_CONFIG = config_element("arrow", "", {"schema_from" => "static"}, [SCHEMA])
+    CONFIG = config_element("ROOT", "", S3_CONFIG, [ARROW_CONFIG])
 
     def test_configure
       d = create_driver
       c = d.instance.instance_variable_get(:@compressor)
       assert_equal :arrow, c.ext
       assert_equal 'application/x-apache-arrow-file', c.content_type
-      assert c.instance_variable_get(:@schema).is_a?(Arrow::Schema)
-      assert_equal 1024, c.instance_variable_get(:@arrow).chunk_size
+      assert c.instance_variable_get(:@options).schema.is_a?(Arrow::Schema)
     end
 
     data(
@@ -32,10 +32,10 @@ class S3OutputTest < Test::Unit::TestCase
     )
     def test_invalid_configure
       format, compression = data
-      arrow_config = config_element("arrow", "", { "schema" => SCHEMA,
+      arrow_config = config_element("arrow", "", { "schema_from" => "static",
         "format" => format,
         "compression" => compression,
-      })
+      }, [SCHEMA])
       config = config_element("ROOT", "", S3_CONFIG, [arrow_config])
       assert_raise Fluent::ConfigError do
         create_driver(config)
@@ -66,9 +66,9 @@ class S3OutputTest < Test::Unit::TestCase
 
     data(gzip: "gzip", zstd: "zstd")
     def test_compress_with_compression
-      arrow_config = config_element("arrow", "", { "schema" => SCHEMA,
+      arrow_config = config_element("arrow", "", { "schema_from" => "static",
         "compression" => data,
-      })
+      },[SCHEMA])
       config = config_element("ROOT", "", S3_CONFIG, [arrow_config])
 
       d = create_driver(conf=config)
@@ -100,10 +100,10 @@ class S3OutputTest < Test::Unit::TestCase
     )
     def test_compress_with_format
       format, compression = data
-      arrow_config = config_element("arrow", "", { "schema" => SCHEMA,
+      arrow_config = config_element("arrow", "", { "schema_from" => "static",
         "format" => format,
         "compression" => compression,
-      })
+      },[SCHEMA])
       config = config_element("ROOT", "", S3_CONFIG, [arrow_config])
 
       d = create_driver(conf=config)

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: fluent-plugin-s3-arrow
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.2.0
 platform: ruby
 authors:
 - kanga333
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2020-09-30 00:00:00.000000000 Z
+date: 2020-12-27 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: benchmark-driver
@@ -66,6 +66,20 @@ dependencies:
     - - "~>"
       - !ruby/object:Gem::Version
         version: '3.0'
+- !ruby/object:Gem::Dependency
+  name: test-unit-rr
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.0'
 - !ruby/object:Gem::Dependency
   name: fluentd
   requirement: !ruby/object:Gem::Requirement
@@ -106,29 +120,43 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '1.0'
+        version: '2.0'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '1.0'
+        version: '2.0'
 - !ruby/object:Gem::Dependency
   name: red-parquet
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '1.0'
+        version: '2.0'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '1.0'
-description: Extends the fluent-plugin-s3 (de)compression algorithm to enable red-arrow
+        version: '2.0'
+- !ruby/object:Gem::Dependency
+  name: aws-sdk-glue
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 1.20.0
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 1.20.0
+description: Extends the fluent-plugin-s3 compression algorithm to enable red-arrow
   compression.
 email:
 - e411z7t40w@gmail.com
@@ -149,7 +177,10 @@ files:
 - benchmark/s3_compressor_parquet.rb
 - benchmark/schema.bq.json
 - fluent-plugin-s3-arrow.gemspec
+- lib/fluent-plugin-s3-arrow/schemas.rb
+- lib/fluent-plugin-s3-arrow/schemas/aws_glue.rb
 - lib/fluent/plugin/s3_compressor_arrow.rb
+- test/fluent-plugin-s3-arrow/schemas/test_aws_glue.rb
 - test/helper.rb
 - test/plugin/test_s3_compressor_arrow.rb
 homepage: https://github.com/red-data-tools/fluent-plugin-s3-arrow
@@ -174,8 +205,8 @@ requirements: []
 rubygems_version: 3.0.3
 signing_key:
 specification_version: 4
-summary: Extends the fluent-plugin-s3 (de)compression algorithm to enable red-arrow
-  compression.
+summary: Extends the fluent-plugin-s3 compression algorithm to enable red-arrow compression.
 test_files:
+- test/fluent-plugin-s3-arrow/schemas/test_aws_glue.rb
 - test/helper.rb
 - test/plugin/test_s3_compressor_arrow.rb