athena-udf 0.1.2 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 1fca59aa26a814ceb98c86febbff75b38a5de2cb888c7291c483ed49aafe2301
-  data.tar.gz: 1a1020ab4893a1689771c5a35049eeca8ec778d7e452de5de660790c8394f600
+  metadata.gz: 9cc2a9e47dee420f0b11442039d1873aca9ad04ecf12775dbe8ea49df3153f02
+  data.tar.gz: 54d15b93ffc2fb58793aad53d03c69608e607c00d0f8da3f5210b47860a5aecc
 SHA512:
-  metadata.gz: 9f2f349f806c505d4957674ac6c18bf326b8c2cae08d259d30501d550b4f5fa8f73f5e79c5ec5eecab8da250abc91a3bb4bf0cddcc3c04e7288a2df2df79f76a
-  data.tar.gz: efe4807a051325233719ffa97aeec66564ad5c02a2703f70ebccb725a464f0be7d8142be42fae0073706f8b56edddf94a58fa7e15bbb3734c0bd2242a9f14475
+  metadata.gz: b0f36011d582e681d575fe8d6f33a6cf9935519ed9bb21953dac980aa1c479c7d943992bbdd0e992f053f32eee6b228f473d772e4e698c3517370009fe49d3b6
+  data.tar.gz: 318b411aa1949b998c1a34176ca049cbccea541861a89931b1a3f71730ab45fec383e03549ec941f7a48305d4ad1e9c0de419d9a46b8e5469205467772df8598

data/.rubocop.yml

@@ -24,6 +24,7 @@ Metrics/AbcSize:
 Metrics/BlockLength:
   Exclude:
     - "spec/**/*"
+    - "scripts/benchmark.rb"
 
 Style/TrailingCommaInArguments:
   EnforcedStyleForMultiline: comma
@@ -37,3 +38,7 @@ Style/TrailingCommaInHashLiteral:
 Style/ClassVars:
   Enabled: false
 
+Naming/VariableNumber:
+  Exclude:
+    - "scripts/benchmark.rb"
+

data/Dockerfile.dev

@@ -0,0 +1,28 @@
+FROM ruby:3.2
+
+# General packages
+RUN apt-get update -y \
+  && apt-get install -y \
+    build-essential \
+    ca-certificates \
+    lsb-release \
+    wget \
+    vim \
+  && apt-get clean
+
+# Apache Arrow
+RUN apt-get update -y \
+  && wget https://apache.jfrog.io/artifactory/arrow/$(lsb_release --id --short | tr 'A-Z' 'a-z')/apache-arrow-apt-source-latest-$(lsb_release --codename --short).deb \
+  && apt install -y ./apache-arrow-apt-source-latest-$(lsb_release --codename --short).deb \
+  && apt-get update -y \
+  && apt-get install -y \
+    libarrow-dev \
+    libarrow-glib-dev \
+  && apt-get clean
+
+# Update bundler
+RUN gem update bundler
+
+WORKDIR /src
+
+CMD ["/bin/bash"]

data/Gemfile

@@ -5,6 +5,7 @@ source 'https://rubygems.org'
 gemspec
 
 group :development, :test do
+  gem 'benchmark', '~> 0.3.0'
   gem 'gem-release', '~> 2.2'
   gem 'rake', '~> 13.0'
   gem 'rspec', '~> 3.0'

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    athena-udf (0.1.2)
+    athena-udf (0.1.3)
       base64
       csv
       red-arrow (~> 12.0.1)
@@ -11,6 +11,7 @@ GEM
   specs:
     ast (2.4.2)
     base64 (0.2.0)
+    benchmark (0.3.0)
     bigdecimal (3.1.8)
     csv (3.3.0)
     diff-lcs (1.5.1)
@@ -44,8 +45,6 @@ GEM
       native-package-installer
       pkg-config
     regexp_parser (2.9.2)
-    rexml (3.3.5)
-      strscan
     rspec (3.13.0)
       rspec-core (~> 3.13.0)
       rspec-expectations (~> 3.13.0)
@@ -59,18 +58,17 @@ GEM
       diff-lcs (>= 1.2.0, < 2.0)
       rspec-support (~> 3.13.0)
     rspec-support (3.13.1)
-    rubocop (1.65.1)
+    rubocop (1.66.0)
       json (~> 2.3)
       language_server-protocol (>= 3.17.0)
       parallel (~> 1.10)
       parser (>= 3.3.0.2)
       rainbow (>= 2.2.2, < 4.0)
       regexp_parser (>= 2.4, < 3.0)
-      rexml (>= 3.2.5, < 4.0)
-      rubocop-ast (>= 1.31.1, < 2.0)
+      rubocop-ast (>= 1.32.1, < 2.0)
       ruby-progressbar (~> 1.7)
       unicode-display_width (>= 2.4.0, < 3.0)
-    rubocop-ast (1.32.1)
+    rubocop-ast (1.32.2)
       parser (>= 3.3.1.0)
     ruby-progressbar (1.13.0)
     simplecov (0.22.0)
@@ -80,7 +78,6 @@ GEM
     simplecov-html (0.12.3)
     simplecov-lcov (0.8.0)
     simplecov_json_formatter (0.1.4)
-    strscan (3.1.0)
     unicode-display_width (2.5.0)
 
 PLATFORMS
@@ -89,6 +86,7 @@ PLATFORMS
 
 DEPENDENCIES
   athena-udf!
+  benchmark (~> 0.3.0)
   gem-release (~> 2.2)
   rake (~> 13.0)
   rspec (~> 3.0)

data/README.md

@@ -83,6 +83,15 @@ $ aws iam attach-role-policy --role-name athena-udf-simple-varchar --policy-arn
 $ aws lambda create-function --function-name athena-udf-simple-varchar --package-type Image --role arn:aws:iam::<ACCOUNT_ID>:role/athena-udf-simple-varchar --code ImageUri=<ACCOUNT_ID>.dkr.ecr.<AWS_REGION>.amazonaws.com/athena-udf-test:latest --publish
 ```
 
+## Development
+
+You can use the dev container image, which includes necessary packages, to develop this library.
+
+```sh
+$ docker build -t ruby-athena-udf-dev -f Dockerfile.dev .
+$ docker run -v $PWD:/src -it ruby-athena-udf-dev
+```
+
 ## Contributing
 
 Bug reports and pull requests are welcome on GitHub at https://github.com/dtaniwaki/ruby-athena-udf.

data/lib/athena-udf/base_udf.rb

@@ -9,33 +9,43 @@ require_relative 'utils'
 
 module AthenaUDF
   class BaseUDF
-    extend AthenaUDF::Utils
+    include AthenaUDF::Utils
 
-    @@logger = Logger.new($stdout)
-    @@logger.level = Logger.const_get(ENV.fetch('LOG_LEVEL', 'WARN').upcase)
+    attr_reader :logger
 
-    def self.lambda_handler(event:, context:) # rubocop:disable Lint/UnusedMethodArgument
+    def self.lambda_handler(event:, context:)
+      instance = new(event:, context:)
       incoming_type = event['@type']
       if incoming_type == 'PingRequest'
-        return handle_ping(event)
+        return instance.handle_ping(event)
       elsif incoming_type == 'UserDefinedFunctionRequest'
-        return handle_udf_request(event)
+        return instance.handle_udf_request(event)
       end
 
       raise "Unknown event type #{incoming_type} from Athena"
     end
 
-    def self.handle_ping(event)
+    # About capabilities: https://github.com/awslabs/aws-athena-query-federation/blob/f52d929a109099a1e7180fa242e26331137ed84c/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/FederationCapabilities.java#L29-L32
+    def self.capabilities
+      1
+    end
+
+    def initialize(event:, context:) # rubocop:disable Lint/UnusedMethodArgument
+      @logger = Logger.new($stdout)
+      @logger.level = Logger.const_get(ENV.fetch('LOG_LEVEL', 'WARN').upcase)
+    end
+
+    def handle_ping(event)
       {
         '@type' => 'PingResponse',
         'catalogName' => 'event',
         'queryId' => event['queryId'],
         'sourceType' => 'athena_udf',
-        'capabilities' => capabilities,
+        'capabilities' => self.class.capabilities,
       }
     end
 
-    def self.handle_udf_request(event)
+    def handle_udf_request(event)
       # Cannot find a way to write Arrow::RecordBatch to a buffer directly in Ruby.
 
       output_schema = read_schema(Base64.decode64(event['outputSchema']['schema']))
@@ -66,17 +76,8 @@ module AthenaUDF
       }
     end
 
-    # About capabilities: https://github.com/awslabs/aws-athena-query-federation/blob/f52d929a109099a1e7180fa242e26331137ed84c/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/FederationCapabilities.java#L29-L32
-    def self.capabilities
-      1
-    end
-
-    def self.handle_athena_record(input_schema, output_schema, records)
+    def handle_athena_record(input_schema, output_schema, records)
       raise NotImplementedError
     end
-
-    def self.logger
-      @@logger
-    end
   end
 end

data/lib/athena-udf/utils.rb

@@ -30,11 +30,11 @@ module AthenaUDF
       end
     end
 
-    def get_schema_bytes(schema, record_batch)
+    def get_schema_bytes(schema)
       buffer = Arrow::ResizableBuffer.new(0)
       Arrow::BufferOutputStream.open(buffer) do |output|
         Arrow::RecordBatchStreamWriter.open(output, schema) do |writer|
-          writer.write_record_batch(record_batch)
+          # noop
         end
 
         bytes = buffer.data.to_s
@@ -61,12 +61,13 @@ module AthenaUDF
       found_count = 0
       start_index = 0
       0.upto(size - 4).each do |i|
-        has_ffff = bytes.slice(i, 4) == [255, 255, 255, 255]
+        has_ffff = bytes.slice(i, 4) == "\xFF\xFF\xFF\xFF".b
+
         found_count += 1 if has_ffff
-        if found_count == 2
-          start_index = i + 4
-          break
-        end
+        next unless found_count == 2
+
+        start_index = i + 4
+        break
       end
 
       start_index

data/lib/athena-udf/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module AthenaUDF
-  VERSION = '0.1.2'
+  VERSION = '0.2.0'
 end

data/scripts/benchmark.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+require 'benchmark'
+require 'athena_udf'
+
+Benchmark.bm 10 do |r|
+  include AthenaUDF::Utils
+
+  instance = Class.new(AthenaUDF::BaseUDF) do
+    def handle_athena_record(_input_schema, _output_schema, record)
+      record.to_a
+    end
+  end.new(event: {}, context: {})
+
+  input_schema_1 = Arrow::Schema.new("0": :string)
+  input_schema_bytes_1 = get_schema_bytes(input_schema_1)
+  Arrow::Schema.new(0.upto(100).map { |n| [n.to_s, :string] }.to_h)
+  input_schema_bytes_100 = get_schema_bytes(input_schema_1)
+
+  input_table1_1 = Arrow::Table.new(input_schema_1, [['FooBar']])
+  input_records_bytes_1_1 = get_record_batch_bytes(input_schema_1, input_table1_1.each_record_batch.first)
+  event_1_1 = {
+    '@type' => 'UserDefinedFunctionRequest',
+    'inputRecords' => {
+      'schema' => Base64.strict_encode64(input_schema_bytes_1),
+      'records' => Base64.strict_encode64(input_records_bytes_1_1),
+    },
+    'methodName' => 'lower',
+    'outputSchema' => {
+      'schema' => Base64.strict_encode64(input_schema_bytes_1),
+    },
+    'functionType' => 'SCALAR',
+  }
+
+  r.report '1 record 1 column' do
+    instance.handle_udf_request(event_1_1)
+  end
+
+  input_table100_1 = Arrow::Table.new(input_schema_1, [['FooBar']] * 100)
+  input_records_bytes_100_1 = get_record_batch_bytes(input_schema_1, input_table100_1.each_record_batch.first)
+  event_100 = {
+    '@type' => 'UserDefinedFunctionRequest',
+    'inputRecords' => {
+      'schema' => Base64.strict_encode64(input_schema_bytes_1),
+      'records' => Base64.strict_encode64(input_records_bytes_100_1),
+    },
+    'methodName' => 'lower',
+    'outputSchema' => {
+      'schema' => Base64.strict_encode64(input_schema_bytes_1),
+    },
+    'functionType' => 'SCALAR',
+  }
+
+  r.report '100 records 1 column' do
+    instance.handle_udf_request(event_100)
+  end
+
+  input_table1_100 = Arrow::Table.new(input_schema_1, [['FooBar']])
+  input_records_bytes_1_100 = get_record_batch_bytes(input_schema_1, input_table1_100.each_record_batch.first)
+  event_1_100 = {
+    '@type' => 'UserDefinedFunctionRequest',
+    'inputRecords' => {
+      'schema' => Base64.strict_encode64(input_schema_bytes_100),
+      'records' => Base64.strict_encode64(input_records_bytes_1_100),
+    },
+    'methodName' => 'lower',
+    'outputSchema' => {
+      'schema' => Base64.strict_encode64(input_schema_bytes_100),
+    },
+    'functionType' => 'SCALAR',
+  }
+
+  r.report '1 record 100 column' do
+    instance.handle_udf_request(event_1_100)
+  end
+
+  input_table_100_100 = Arrow::Table.new(input_schema_1, [['FooBar']])
+  input_records_bytes_100_100 = get_record_batch_bytes(input_schema_1, input_table_100_100.each_record_batch.first)
+  event_100_100 = {
+    '@type' => 'UserDefinedFunctionRequest',
+    'inputRecords' => {
+      'schema' => Base64.strict_encode64(input_schema_bytes_100),
+      'records' => Base64.strict_encode64(input_records_bytes_100_100),
+    },
+    'methodName' => 'lower',
+    'outputSchema' => {
+      'schema' => Base64.strict_encode64(input_schema_bytes_100),
+    },
+    'functionType' => 'SCALAR',
+  }
+
+  r.report '100 record 100 column' do
+    instance.handle_udf_request(event_100_100)
+  end
+end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: athena-udf
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.2.0
 platform: ruby
 authors:
 - Daisuke Taniwaki
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2024-09-03 00:00:00.000000000 Z
+date: 2024-09-07 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: base64
@@ -62,6 +62,7 @@ files:
 - ".dockerignore"
 - ".rspec"
 - ".rubocop.yml"
+- Dockerfile.dev
 - Dockerfile.example
 - Gemfile
 - Gemfile.lock
@@ -73,6 +74,7 @@ files:
 - lib/athena-udf/utils.rb
 - lib/athena-udf/version.rb
 - lib/athena_udf.rb
+- scripts/benchmark.rb
 homepage: https://github.com/dtaniwaki/ruby-athena-udf
 licenses:
 - MIT
@@ -94,7 +96,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.5.11
+rubygems_version: 3.4.19
 signing_key:
 specification_version: 4
 summary: Ruby-version Athena UDF