athena-udf 0.1.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: be808b4f155b8cdfa5c52bcc3c486d77cb9d056099f592b304bc8aa6426694b4
+  data.tar.gz: c41565a48ec8702756c48f01112b6d7dc881723a26a521c09cbb8bf506666a24
+SHA512:
+  metadata.gz: 971c50fb5eee0df1c0f4522b19a1583fca1ad94144f38766f0a6ff42b82151ee8501d7b303cc361f70a151c4e29e57aaca360bdd2e57646fb8509c5d4fbd74bd
+  data.tar.gz: af0deed14deb22d1e382adfe9bc4097f07dff0cea835f9188e847df7bd27017e91b21c8c84bdf56ebbf6ac4dc729a3e6fe9dfbff43847f6eca1ee423389d3824

data/.dockerignore

@@ -0,0 +1,2 @@
+.git
+Dockerfile.example

data/.rspec

@@ -0,0 +1,3 @@
+--format documentation
+--color
+--require spec_helper

data/.rubocop.yml

@@ -0,0 +1,36 @@
+AllCops:
+  TargetRubyVersion: 3.2
+
+Style/StringLiterals:
+  Enabled: true
+  EnforcedStyle: single_quotes
+
+Style/StringLiteralsInInterpolation:
+  Enabled: true
+  EnforcedStyle: single_quotes
+
+Layout/LineLength:
+  Max: 120
+
+Style/Documentation:
+  Enabled: false
+
+Metrics/MethodLength:
+  Max: 30
+
+Metrics/AbcSize:
+  Max: 40
+
+Metrics/BlockLength:
+  Exclude:
+    - "spec/**/*"
+
+Style/TrailingCommaInArguments:
+  EnforcedStyleForMultiline: comma
+
+Style/TrailingCommaInArrayLiteral:
+  EnforcedStyleForMultiline: comma
+
+Style/TrailingCommaInHashLiteral:
+  EnforcedStyleForMultiline: comma
+

data/Dockerfile.example

@@ -0,0 +1,31 @@
+FROM public.ecr.aws/lambda/ruby:3.2
+
+# General packages
+RUN yum update -y \
+  && yum install -y \
+    amazon-linux-extras \
+    gcc-c++ \
+    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
+
+# Update bundler
+RUN gem update bundler
+
+# Install gems
+COPY Gemfile Gemfile.lock athena-udf.gemspec ${LAMBDA_TASK_ROOT}
+COPY lib/athena-udf/version.rb ${LAMBDA_TASK_ROOT}/lib/athena-udf/version.rb
+ENV GEM_HOME=${LAMBDA_TASK_ROOT}
+RUN bundle config set --local without development \
+  && bundle install
+
+COPY . ${LAMBDA_TASK_ROOT}
+
+CMD ["example.SimpleVarcharUDF.lambda_handler"]

data/Gemfile

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+source 'https://rubygems.org'
+
+gemspec
+
+group :development, :test do
+  gem 'rake', '~> 13.0'
+
+  gem 'rspec', '~> 3.0'
+
+  gem 'rubocop', '~> 1.21'
+
+  gem 'gem-release', '~> 2.2'
+end

data/Gemfile.lock

@@ -0,0 +1,90 @@
+PATH
+  remote: .
+  specs:
+    athena-udf (0.1.0)
+      base64
+      csv
+      red-arrow (~> 12.0.1)
+
+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)
+    extpp (0.1.1)
+    fiddle (1.1.2)
+    gem-release (2.2.2)
+    gio2 (4.2.2)
+      fiddle
+      gobject-introspection (= 4.2.2)
+    glib2 (4.2.2)
+      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)
+    native-package-installer (1.1.9)
+    parallel (1.26.3)
+    parser (3.3.4.2)
+      ast (~> 2.4.1)
+      racc
+    pkg-config (1.5.6)
+    racc (1.8.1)
+    rainbow (3.1.1)
+    rake (13.2.1)
+    red-arrow (12.0.1)
+      bigdecimal (>= 3.1.0)
+      extpp (>= 0.1.1)
+      gio2 (>= 3.5.0)
+      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)
+      rspec-mocks (~> 3.13.0)
+    rspec-core (3.13.0)
+      rspec-support (~> 3.13.0)
+    rspec-expectations (3.13.1)
+      diff-lcs (>= 1.2.0, < 2.0)
+      rspec-support (~> 3.13.0)
+    rspec-mocks (3.13.1)
+      diff-lcs (>= 1.2.0, < 2.0)
+      rspec-support (~> 3.13.0)
+    rspec-support (3.13.1)
+    rubocop (1.65.1)
+      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)
+      ruby-progressbar (~> 1.7)
+      unicode-display_width (>= 2.4.0, < 3.0)
+    rubocop-ast (1.32.1)
+      parser (>= 3.3.1.0)
+    ruby-progressbar (1.13.0)
+    strscan (3.1.0)
+    unicode-display_width (2.5.0)
+
+PLATFORMS
+  arm64-darwin-23
+  ruby
+
+DEPENDENCIES
+  athena-udf!
+  gem-release (~> 2.2)
+  rake (~> 13.0)
+  rspec (~> 3.0)
+  rubocop (~> 1.21)
+
+BUNDLED WITH
+   2.5.17

data/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2024 Daisuke Taniwaki
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

data/README.md

@@ -0,0 +1,90 @@
+# AthenaUDF
+
+Ruby-version Athena User Defined Function (UDF).
+
+This gem is highly inspired by [the Python-version Athena UDF](https://github.com/dmarkey/python-athena-udf).
+
+See [an official example implementation](https://github.com/awslabs/aws-athena-query-federation/blob/fc2e4e9cdcb71ec7f7c7d44cbda7f56c5835811e/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/handlers/UserDefinedFunctionHandler.java) for more detail of a lambda function for Athena UDF.
+
+## Installation
+
+Install the gem and add to the application's Gemfile by executing:
+
+```sh
+$ bundle add athena-udf
+```
+
+If bundler is not being used to manage dependencies, install the gem by executing:
+
+```sh
+$ gem install athena-udf
+```
+
+## Usage
+
+Just make a subclass of `AthenaUDF::BaseUDF` and implement a concrete function logic.
+
+```rb
+require "athena-udf"
+
+class SimpleVarcharUDF < AthenaUDF::BaseUDF
+  def self.handle_athena_record(_input_schema, _output_schema, record)
+    [record[0].downcase]
+  end
+end
+```
+
+Then, it can be called as `SimpleVarcharUDF.lambda_handler` in your lambda function for Athena UDF workloads.
+
+After pushing an image to Amazon ECR, you can call the function like the following SQL.
+
+```sql
+USING EXTERNAL FUNCTION my_udf(col1 varchar) RETURNS varchar LAMBDA 'athena-udf-simple-varchar'
+
+SELECT my_udf('FooBar');
+```
+
+See [the official document](https://docs.aws.amazon.com/athena/latest/ug/querying-udf.html) for the UDF usage.
+
+## Development
+
+To contribute to this library, first checkout the code. Then, install the dependent gems.
+
+```sh
+$ bundle install
+```
+
+To run the tests:
+
+```sh
+$ bundle exec rspec
+```
+
+## Deployment
+
+You can try the example with the following steps.
+
+First, push a container image to Amazon ECR:
+
+```sh
+$ aws ecr get-login-password | docker login --username AWS --password-stdin https://<ACCOUNT_ID>.dkr.ecr.<AWS_REGION>.amazonaws.com
+$ docker build --platform=linux/amd64 -t <ACCOUNT_ID>.dkr.ecr.<AWS_REGION>.amazonaws.com/athena-udf-test -f Dockerfile.example .
+$ docker push <ACCOUNT_ID>.dkr.ecr.<AWS_REGION>.amazonaws.com/athena-udf-test
+```
+
+Then, create a lambda function with the CLI:
+
+```sh
+$ aws iam create-role --role-name athena-udf-simple-varchar --assume-role-policy-document '{"Version": "2012-10-17","Statement": [{ "Effect": "Allow", "Principal": {"Service": "lambda.amazonaws.com"}, "Action": "sts:AssumeRole"}]}'
+$ aws iam attach-role-policy --role-name athena-udf-simple-varchar --policy-arn arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
+$ 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
+```
+
+## Contributing
+
+Bug reports and pull requests are welcome on GitHub at https://github.com/dtaniwaki/ruby-athena-udf.
+
+## License
+
+The gem is available as open source under the terms of the [MIT License](https://opensource.org/licenses/MIT).
+

data/Rakefile

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require 'bundler/gem_tasks'
+require 'rspec/core/rake_task'
+
+RSpec::Core::RakeTask.new(:spec)
+
+require 'rubocop/rake_task'
+
+RuboCop::RakeTask.new
+
+task default: %i[spec rubocop]

data/example.rb

@@ -0,0 +1,9 @@
+# frozen_string_literal: true
+
+require_relative 'lib/athena_udf'
+
+class SimpleVarcharUDF < AthenaUDF::BaseUDF
+  def self.handle_athena_record(_input_schema, _output_schema, record)
+    [record[0].downcase]
+  end
+end

data/lib/athena-udf/base_udf.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+require 'securerandom'
+require 'base64'
+require 'tempfile'
+require 'arrow'
+require_relative 'utils'
+
+module AthenaUDF
+  class BaseUDF
+    extend AthenaUDF::Utils
+
+    def self.lambda_handler(event:, context:) # rubocop:disable Lint/UnusedMethodArgument
+      incoming_type = event['@type']
+      if incoming_type == 'PingRequest'
+        return handle_ping(event)
+      elsif incoming_type == 'UserDefinedFunctionRequest'
+        return handle_udf_request(event)
+      end
+
+      raise "Unknown event type #{incoming_type} from Athena"
+    end
+
+    def self.handle_ping(event)
+      {
+        '@type' => 'PingResponse',
+        'catalogName' => 'event',
+        'queryId' => event['queryId'],
+        'sourceType' => 'athena_udf',
+        'capabilities' => capabilities,
+      }
+    end
+
+    def self.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']))
+      output_builder = Arrow::RecordBatchBuilder.new(output_schema)
+
+      input_schema_data = Base64.decode64(event['inputRecords']['schema'])
+      input_records_data = Base64.decode64(event['inputRecords']['records'])
+      read_record_batches(input_schema_data, input_records_data) do |input_schema, record_batch|
+        output_builder.append_records(
+          record_batch.each_record.map do |record|
+            handle_athena_record(input_schema, output_schema, record)
+          end,
+        )
+      end
+
+      output_record_batch = output_builder.flush
+      output_records_bytes = get_record_batch_bytes(output_schema, output_record_batch)
+
+      {
+        '@type' => 'UserDefinedFunctionResponse',
+        'methodName' => event['methodName'],
+        'records' => {
+          'aId' => SecureRandom.uuid,
+          'schema' => event['outputSchema']['schema'],
+          'records' => Base64.strict_encode64(output_records_bytes),
+        },
+      }
+    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)
+      raise NotImplementedError
+    end
+  end
+end

data/lib/athena-udf/utils.rb

@@ -0,0 +1,82 @@
+# frozen_string_literal: true
+
+module AthenaUDF
+  module Utils
+    def read_record_batches(schema_data, record_batch_data)
+      Tempfile.create do |in_f|
+        in_f.write(schema_data)
+        in_f.write(record_batch_data)
+        in_f.flush
+
+        Arrow::MemoryMappedInputStream.open(in_f.path) do |inp|
+          reader = Arrow::RecordBatchStreamReader.new(inp)
+          input_schema = reader.schema
+          reader.each do |record_batch|
+            yield input_schema, record_batch
+          end
+        end
+      end
+    end
+
+    def read_schema(schema_data)
+      # schema_buf = Arrow::Buffer.try_convert(schema_data)
+      Tempfile.create do |f|
+        f.write(schema_data)
+        f.flush
+
+        Arrow::MemoryMappedInputStream.open(f.path) do |inp|
+          reader = Arrow::RecordBatchStreamReader.new(inp)
+          reader.schema
+        end
+      end
+    end
+
+    def get_schema_bytes(schema, record_batch)
+      Tempfile.create do |f|
+        Arrow::FileOutputStream.open(f.path, false) do |oup|
+          Arrow::RecordBatchStreamWriter.open(oup, schema) do |writer|
+            writer.write_record_batch(record_batch)
+          end
+        end
+        f.flush
+
+        data = File.binread(f.path)
+        start_index = get_record_batch_index(data)
+        data.bytes[4..start_index - 5].pack('C*')
+      end
+    end
+
+    def get_record_batch_bytes(schema, record_batch)
+      Tempfile.create do |f|
+        Arrow::FileOutputStream.open(f.path, false) do |oup|
+          Arrow::RecordBatchStreamWriter.open(oup, schema) do |writer|
+            writer.write_record_batch(record_batch)
+          end
+        end
+        f.flush
+
+        data = File.binread(f.path)
+        start_index = get_record_batch_index(data)
+        data.bytes[start_index..].pack('C*')
+      end
+    end
+
+    def get_record_batch_index(raw)
+      size = raw.bytes.size
+      found_count = 0
+      start_index = 0
+      0.upto(size - 4).each do |i|
+        has_ffff = 4.times.all? do |n|
+          raw.bytes[i + n] == 255
+        end
+        found_count += 1 if has_ffff
+        if found_count == 2
+          start_index = i + 4
+          break
+        end
+      end
+
+      start_index
+    end
+  end
+end

data/lib/athena-udf/version.rb

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

data/lib/athena_udf.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+require_relative 'athena-udf/version'
+require_relative 'athena-udf/base_udf'
+
+module AthenaUDF
+end

metadata

@@ -0,0 +1,101 @@
+--- !ruby/object:Gem::Specification
+name: athena-udf
+version: !ruby/object:Gem::Version
+  version: 0.1.0
+platform: ruby
+authors:
+- Daisuke Taniwaki
+autorequire:
+bindir: exe
+cert_chain: []
+date: 2024-08-23 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: base64
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: csv
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: red-arrow
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 12.0.1
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 12.0.1
+description: ''
+email:
+- daisuketaniwaki@gmail.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- ".dockerignore"
+- ".rspec"
+- ".rubocop.yml"
+- Dockerfile.example
+- Gemfile
+- Gemfile.lock
+- LICENSE.txt
+- README.md
+- Rakefile
+- example.rb
+- lib/athena-udf/base_udf.rb
+- lib/athena-udf/utils.rb
+- lib/athena-udf/version.rb
+- lib/athena_udf.rb
+homepage: https://github.com/dtaniwaki/ruby-athena-udf
+licenses:
+- MIT
+metadata:
+  homepage_uri: https://github.com/dtaniwaki/ruby-athena-udf
+  source_code_uri: https://github.com/dtaniwaki/ruby-athena-udf
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: 3.2.0
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.5.11
+signing_key:
+specification_version: 4
+summary: Ruby-version Athena UDF
+test_files: []