athena-udf 0.1.3 → 0.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: ab79db7e49eadd3d4448779185ecfc592f497e0a81209964ae7c2b8e4fdbeed5
-  data.tar.gz: 3c2b577e2506f67b9c0e455d8e3988b844fb5762dba69fa91c508d75d5de18fe
+  metadata.gz: 9e5ff4cb6e2f166491b1bb18cc2be70b9de602852ff9535eb7b888de86d46569
+  data.tar.gz: df66a5f63ea2ac6a036d4f47befd0dc55ae0f8487de80a8d7d1da8d38606307d
 SHA512:
-  metadata.gz: 03b34a9703a30047eaeaf03c2bf8b975f190bf1f7ad08c30f1f06f8dcf2b6a7eca90b29621437699ce9acd22579bc44f59d03941493fb27adcb173caa7488f7b
-  data.tar.gz: c9dc25db956961249e40e4909ec284b5234622c05933c0fa78347431344afd5e5ca09cc6c36fa6b480ff2b5e4893d095e94369e6a07fc72917d8df279d9791f8
+  metadata.gz: 18e43fce15698f0fc1dd0160cb0098887cf78085ca13774e72819cc89369b3fd466b6bf27a216e1ef1ff863caa412881f0fab4ba5df9af62b6cc0ddfa77f2e85
+  data.tar.gz: eba9d3b48c2adabe2c02c6a22686e0f8b7a4ee35f535738d9a72a7384948765c14a4d978e3061b5d43e75610bab7df8580dd8cbcb9a270bd025a7bb8f0392d5a

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,32 @@
+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
+
+COPY Gemfile Gemfile.lock athena-udf.gemspec /src
+COPY lib/athena-udf/version.rb /src/lib/athena-udf/version.rb
+RUN bundle install
+
+CMD ["/bin/bash"]

data/Dockerfile.example

@@ -1,20 +1,36 @@
-FROM public.ecr.aws/lambda/ruby:3.2
+# Can NOT install apache-arrow on the amazonlinux:2023-minimal image,
+# so install the ruby directly on amazonlinux:2023.
+FROM public.ecr.aws/amazonlinux/amazonlinux:2023 AS builder
+
+# Apache Arrow
+RUN dnf upgrade -y --releasever=latest \
+  && dnf install -y https://apache.jfrog.io/artifactory/arrow/amazon-linux/$(cut -d: -f6 /etc/system-release-cpe)/apache-arrow-release-latest.rpm \
+  && dnf install -y arrow-devel arrow-glib-devel arrow-dataset-devel arrow-dataset-glib-devel
+
+############
+
+FROM public.ecr.aws/lambda/ruby:3.3
+
+COPY --from=builder /usr/include/parquet/arrow/ /usr/include/parquet/arrow/
+COPY --from=builder /usr/include/arrow* /usr/include/
+COPY --from=builder /usr/share/gdb/auto-load/usr/lib64/libarrow.so* /usr/share/gdb/auto-load/usr/lib64/
+COPY --from=builder /usr/share/licenses/arrow* /usr/share/licenses/
+COPY --from=builder /usr/share/doc/arrow* /usr/share/doc/
+COPY --from=builder /usr/share/arrow* /usr/share/
+COPY --from=builder /usr/share/vala/vapi/arrow* /usr/share/vala/vapi/
+COPY --from=builder /usr/lib64/libarrow* /usr/lib64/
+COPY --from=builder /usr/lib64/pkgconfig/arrow* /usr/lib64/pkgconfig/
+COPY --from=builder /usr/lib64/libarrow* /usr/lib64/libarrow/
+COPY --from=builder /usr/lib64/cmake/Arrow/ /usr/lib64/cmake/Arrow/
 
 # General packages
-RUN yum update -y \
-  && yum install -y \
-    amazon-linux-extras \
+RUN dnf upgrade -y --releasever=latest \
+  && dnf install -y \
     gcc-c++ \
+    glib2-devel \
     make \
     git \
-  && amazon-linux-extras install -y epel \
-  && yum clean all
-
-# Apache Arrow
-RUN yum update -y \
-  && yum install -y https://apache.jfrog.io/artifactory/arrow/amazon-linux/2/apache-arrow-release-latest.rpm \
-  && yum install -y --enablerepo=epel arrow-devel arrow-glib-devel arrow-dataset-devel arrow-dataset-glib-devel \
-  && yum clean all
+  && dnf clean all
 
 # Update bundler
 RUN gem update bundler
@@ -29,3 +45,4 @@ RUN bundle config set --local without development \
 COPY . ${LAMBDA_TASK_ROOT}
 
 CMD ["example.SimpleVarcharUDF.lambda_handler"]
+

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,39 +1,41 @@
 PATH
   remote: .
   specs:
-    athena-udf (0.1.3)
+    athena-udf (0.2.0)
       base64
       csv
-      red-arrow (~> 12.0.1)
+      red-arrow (< 20.0.0)
 
 GEM
   remote: https://rubygems.org/
   specs:
     ast (2.4.2)
     base64 (0.2.0)
-    bigdecimal (3.1.8)
-    csv (3.3.0)
-    diff-lcs (1.5.1)
+    benchmark (0.3.0)
+    bigdecimal (3.1.9)
+    csv (3.3.2)
+    diff-lcs (1.6.0)
     docile (1.4.1)
     extpp (0.1.1)
-    fiddle (1.1.2)
-    gem-release (2.2.2)
-    gio2 (4.2.2)
+    fiddle (1.1.6)
+    gem-release (2.2.4)
+    gio2 (4.2.7)
       fiddle
-      gobject-introspection (= 4.2.2)
-    glib2 (4.2.2)
+      gobject-introspection (= 4.2.7)
+    glib2 (4.2.7)
       native-package-installer (>= 1.0.3)
       pkg-config (>= 1.3.5)
-    gobject-introspection (4.2.2)
-      glib2 (= 4.2.2)
-    json (2.7.2)
-    language_server-protocol (3.17.0.3)
+    gobject-introspection (4.2.7)
+      glib2 (= 4.2.7)
+    json (2.10.2)
+    language_server-protocol (3.17.0.4)
+    lint_roller (1.1.0)
     native-package-installer (1.1.9)
     parallel (1.26.3)
-    parser (3.3.4.2)
+    parser (3.3.7.1)
       ast (~> 2.4.1)
       racc
-    pkg-config (1.5.6)
+    pkg-config (1.6.0)
     racc (1.8.1)
     rainbow (3.1.1)
     rake (13.2.1)
@@ -43,45 +45,44 @@ GEM
       gio2 (>= 3.5.0)
       native-package-installer
       pkg-config
-    regexp_parser (2.9.2)
-    rexml (3.3.5)
-      strscan
+    regexp_parser (2.10.0)
     rspec (3.13.0)
       rspec-core (~> 3.13.0)
       rspec-expectations (~> 3.13.0)
       rspec-mocks (~> 3.13.0)
-    rspec-core (3.13.0)
+    rspec-core (3.13.3)
       rspec-support (~> 3.13.0)
-    rspec-expectations (3.13.1)
+    rspec-expectations (3.13.3)
       diff-lcs (>= 1.2.0, < 2.0)
       rspec-support (~> 3.13.0)
-    rspec-mocks (3.13.1)
+    rspec-mocks (3.13.2)
       diff-lcs (>= 1.2.0, < 2.0)
       rspec-support (~> 3.13.0)
-    rspec-support (3.13.1)
-    rubocop (1.65.1)
+    rspec-support (3.13.2)
+    rubocop (1.74.0)
       json (~> 2.3)
-      language_server-protocol (>= 3.17.0)
+      language_server-protocol (~> 3.17.0.2)
+      lint_roller (~> 1.1.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)
+      regexp_parser (>= 2.9.3, < 3.0)
+      rubocop-ast (>= 1.38.0, < 2.0)
       ruby-progressbar (~> 1.7)
-      unicode-display_width (>= 2.4.0, < 3.0)
-    rubocop-ast (1.32.1)
+      unicode-display_width (>= 2.4.0, < 4.0)
+    rubocop-ast (1.38.1)
       parser (>= 3.3.1.0)
     ruby-progressbar (1.13.0)
     simplecov (0.22.0)
       docile (~> 1.1)
       simplecov-html (~> 0.11)
       simplecov_json_formatter (~> 0.1)
-    simplecov-html (0.12.3)
+    simplecov-html (0.13.1)
     simplecov-lcov (0.8.0)
     simplecov_json_formatter (0.1.4)
-    strscan (3.1.0)
-    unicode-display_width (2.5.0)
+    unicode-display_width (3.1.4)
+      unicode-emoji (~> 4.0, >= 4.0.4)
+    unicode-emoji (4.0.4)
 
 PLATFORMS
   arm64-darwin-23
@@ -89,6 +90,7 @@ PLATFORMS
 
 DEPENDENCIES
   athena-udf!
+  benchmark (~> 0.3.0)
   gem-release (~> 2.2)
   rake (~> 13.0)
   rspec (~> 3.0)

data/README.md

@@ -31,7 +31,7 @@ Just make a subclass of `AthenaUDF::BaseUDF` and implement a concrete function l
 require "athena-udf"
 
 class SimpleVarcharUDF < AthenaUDF::BaseUDF
-  def self.handle_athena_record(_input_schema, _output_schema, record)
+  def handle_athena_record(_input_schema, _output_schema, record)
     [record[0].downcase]
   end
 end
@@ -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

@@ -2,6 +2,9 @@
 
 module AthenaUDF
   module Utils
+    SEPARATOR = "\xFF\xFF\xFF\xFF".b
+    SEPARATOR_SIZE = SEPARATOR.bytesize
+
     def read_record_batches(schema_data, record_batch_data)
       buffer = Arrow::ResizableBuffer.new(schema_data.bytes.size + record_batch_data.bytes.size)
       Arrow::BufferOutputStream.open(buffer) do |output|
@@ -30,16 +33,16 @@ 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
-        start_index = get_record_batch_index(bytes)
-        bytes[4..start_index - 5]
+        last_index = bytes.index(SEPARATOR, SEPARATOR_SIZE)
+        bytes[SEPARATOR_SIZE...last_index]
       end
     end
 
@@ -51,26 +54,9 @@ module AthenaUDF
         end
 
         bytes = buffer.data.to_s
-        start_index = get_record_batch_index(bytes)
+        start_index = bytes.index(SEPARATOR, SEPARATOR_SIZE) + SEPARATOR_SIZE
         bytes[start_index..]
       end
     end
-
-    def get_record_batch_index(bytes)
-      size = bytes.size
-      found_count = 0
-      start_index = 0
-      0.upto(size - 4).each do |i|
-        has_ffff = bytes.slice(i, 4) == "\xFF\xFF\xFF\xFF".b
-
-        found_count += 1 if has_ffff
-        next unless found_count == 2
-
-        start_index = i + 4
-        break
-      end
-
-      start_index
-    end
   end
 end

data/lib/athena-udf/version.rb

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

data/scripts/benchmark.rb

@@ -0,0 +1,105 @@
+# frozen_string_literal: true
+
+require 'benchmark'
+require 'athena_udf'
+
+Benchmark.bm 20 do |r|
+  include AthenaUDF::Utils
+
+  instance = Class.new(AthenaUDF::BaseUDF) do
+    def handle_athena_record(_input_schema, _output_schema, record)
+      [record[0]]
+    end
+  end.new(event: {}, context: {})
+
+  input_schema_1 = Arrow::Schema.new("0": :string)
+  input_schema_bytes_1 = get_schema_bytes(input_schema_1)
+  input_schema_100 = Arrow::Schema.new(0.upto(100).map { |n| [n.to_s, :string] }.to_h)
+  input_schema_bytes_100 = get_schema_bytes(input_schema_100)
+  output_schema = Arrow::Schema.new("0": :string)
+  output_schema_bytes = get_schema_bytes(output_schema)
+
+  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(output_schema_bytes),
+    },
+    'functionType' => 'SCALAR',
+  }
+
+  r.report '1 record 1 column' do
+    100.times do
+      instance.handle_udf_request(event_1_1)
+    end
+  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(output_schema_bytes),
+    },
+    'functionType' => 'SCALAR',
+  }
+
+  r.report '100 records 1 column' do
+    100.times do
+      instance.handle_udf_request(event_100)
+    end
+  end
+
+  input_table1_100 = Arrow::Table.new(input_schema_100, [['FooBar'] * 100])
+  input_records_bytes_1_100 = get_record_batch_bytes(input_schema_100, 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(output_schema_bytes),
+    },
+    'functionType' => 'SCALAR',
+  }
+
+  r.report '1 record 100 column' do
+    100.times do
+      instance.handle_udf_request(event_1_100)
+    end
+  end
+
+  input_table_100_100 = Arrow::Table.new(input_schema_100, [['FooBar'] * 100] * 100)
+  input_records_bytes_100_100 = get_record_batch_bytes(input_schema_100, 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(output_schema_bytes),
+    },
+    'functionType' => 'SCALAR',
+  }
+
+  r.report '100 record 100 column' do
+    100.times do
+      instance.handle_udf_request(event_100_100)
+    end
+  end
+end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: athena-udf
 version: !ruby/object:Gem::Version
-  version: 0.1.3
+  version: 0.2.1
 platform: ruby
 authors:
 - Daisuke Taniwaki
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2024-09-03 00:00:00.000000000 Z
+date: 2025-03-13 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: base64
@@ -42,16 +42,22 @@ dependencies:
   name: red-arrow
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 12.0.0
+    - - "<"
       - !ruby/object:Gem::Version
-        version: 12.0.1
+        version: 20.0.0
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 12.0.0
+    - - "<"
       - !ruby/object:Gem::Version
-        version: 12.0.1
+        version: 20.0.0
 description: ''
 email:
 - daisuketaniwaki@gmail.com
@@ -62,6 +68,7 @@ files:
 - ".dockerignore"
 - ".rspec"
 - ".rubocop.yml"
+- Dockerfile.dev
 - Dockerfile.example
 - Gemfile
 - Gemfile.lock
@@ -73,6 +80,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