athena-udf 0.1.0 → 0.1.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: be808b4f155b8cdfa5c52bcc3c486d77cb9d056099f592b304bc8aa6426694b4
-  data.tar.gz: c41565a48ec8702756c48f01112b6d7dc881723a26a521c09cbb8bf506666a24
+  metadata.gz: 1fca59aa26a814ceb98c86febbff75b38a5de2cb888c7291c483ed49aafe2301
+  data.tar.gz: 1a1020ab4893a1689771c5a35049eeca8ec778d7e452de5de660790c8394f600
 SHA512:
-  metadata.gz: 971c50fb5eee0df1c0f4522b19a1583fca1ad94144f38766f0a6ff42b82151ee8501d7b303cc361f70a151c4e29e57aaca360bdd2e57646fb8509c5d4fbd74bd
-  data.tar.gz: af0deed14deb22d1e382adfe9bc4097f07dff0cea835f9188e847df7bd27017e91b21c8c84bdf56ebbf6ac4dc729a3e6fe9dfbff43847f6eca1ee423389d3824
+  metadata.gz: 9f2f349f806c505d4957674ac6c18bf326b8c2cae08d259d30501d550b4f5fa8f73f5e79c5ec5eecab8da250abc91a3bb4bf0cddcc3c04e7288a2df2df79f76a
+  data.tar.gz: efe4807a051325233719ffa97aeec66564ad5c02a2703f70ebccb725a464f0be7d8142be42fae0073706f8b56edddf94a58fa7e15bbb3734c0bd2242a9f14475

data/.rubocop.yml

@@ -34,3 +34,6 @@ Style/TrailingCommaInArrayLiteral:
 Style/TrailingCommaInHashLiteral:
   EnforcedStyleForMultiline: comma
 
+Style/ClassVars:
+  Enabled: false
+

data/Gemfile

@@ -5,11 +5,10 @@ source 'https://rubygems.org'
 gemspec
 
 group :development, :test do
+  gem 'gem-release', '~> 2.2'
   gem 'rake', '~> 13.0'
-
   gem 'rspec', '~> 3.0'
-
   gem 'rubocop', '~> 1.21'
-
-  gem 'gem-release', '~> 2.2'
+  gem 'simplecov', '~> 0.22'
+  gem 'simplecov-lcov', '~> 0.8.0'
 end

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    athena-udf (0.1.0)
+    athena-udf (0.1.2)
       base64
       csv
       red-arrow (~> 12.0.1)
@@ -14,6 +14,7 @@ GEM
     bigdecimal (3.1.8)
     csv (3.3.0)
     diff-lcs (1.5.1)
+    docile (1.4.1)
     extpp (0.1.1)
     fiddle (1.1.2)
     gem-release (2.2.2)
@@ -72,6 +73,13 @@ GEM
     rubocop-ast (1.32.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-lcov (0.8.0)
+    simplecov_json_formatter (0.1.4)
     strscan (3.1.0)
     unicode-display_width (2.5.0)
 
@@ -85,6 +93,8 @@ DEPENDENCIES
   rake (~> 13.0)
   rspec (~> 3.0)
   rubocop (~> 1.21)
+  simplecov (~> 0.22)
+  simplecov-lcov (~> 0.8.0)
 
 BUNDLED WITH
    2.5.17

data/README.md

@@ -1,5 +1,8 @@
 # AthenaUDF
 
+[![Gem Version][gem-image]][gem-link]
+[![Coverage Status][cov-image]][cov-link]
+
 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).
@@ -88,3 +91,7 @@ Bug reports and pull requests are welcome on GitHub at https://github.com/dtaniw
 
 The gem is available as open source under the terms of the [MIT License](https://opensource.org/licenses/MIT).
 
+[gem-image]:   https://badge.fury.io/rb/athena-udf.svg
+[gem-link]:    http://badge.fury.io/rb/athena-udf
+[cov-image]:   https://coveralls.io/repos/dtaniwaki/ruby-athena-udf/badge.png
+[cov-link]:    https://coveralls.io/r/dtaniwaki/ruby-athena-udf

data/lib/athena-udf/base_udf.rb

@@ -4,12 +4,16 @@ require 'securerandom'
 require 'base64'
 require 'tempfile'
 require 'arrow'
+require 'logger'
 require_relative 'utils'
 
 module AthenaUDF
   class BaseUDF
     extend AthenaUDF::Utils
 
+    @@logger = Logger.new($stdout)
+    @@logger.level = Logger.const_get(ENV.fetch('LOG_LEVEL', 'WARN').upcase)
+
     def self.lambda_handler(event:, context:) # rubocop:disable Lint/UnusedMethodArgument
       incoming_type = event['@type']
       if incoming_type == 'PingRequest'
@@ -40,6 +44,7 @@ module AthenaUDF
       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|
+        logger.info("Processing #{record_batch.size} records")
         output_builder.append_records(
           record_batch.each_record.map do |record|
             handle_athena_record(input_schema, output_schema, record)
@@ -69,5 +74,9 @@ module AthenaUDF
     def self.handle_athena_record(input_schema, output_schema, records)
       raise NotImplementedError
     end
+
+    def self.logger
+      @@logger
+    end
   end
 end

data/lib/athena-udf/utils.rb

@@ -3,13 +3,13 @@
 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
+      buffer = Arrow::ResizableBuffer.new(schema_data.bytes.size + record_batch_data.bytes.size)
+      Arrow::BufferOutputStream.open(buffer) do |output|
+        output.write(schema_data)
+        output.write(record_batch_data)
 
-        Arrow::MemoryMappedInputStream.open(in_f.path) do |inp|
-          reader = Arrow::RecordBatchStreamReader.new(inp)
+        Arrow::BufferInputStream.open(buffer) do |input|
+          reader = Arrow::RecordBatchStreamReader.new(input)
           input_schema = reader.schema
           reader.each do |record_batch|
             yield input_schema, record_batch
@@ -19,56 +19,49 @@ module AthenaUDF
     end
 
     def read_schema(schema_data)
-      # schema_buf = Arrow::Buffer.try_convert(schema_data)
-      Tempfile.create do |f|
-        f.write(schema_data)
-        f.flush
+      buffer = Arrow::ResizableBuffer.new(schema_data.bytes.size)
+      Arrow::BufferOutputStream.open(buffer) do |output|
+        output.write(schema_data)
 
-        Arrow::MemoryMappedInputStream.open(f.path) do |inp|
-          reader = Arrow::RecordBatchStreamReader.new(inp)
+        Arrow::BufferInputStream.open(buffer) do |input|
+          reader = Arrow::RecordBatchStreamReader.new(input)
           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
+      buffer = Arrow::ResizableBuffer.new(0)
+      Arrow::BufferOutputStream.open(buffer) do |output|
+        Arrow::RecordBatchStreamWriter.open(output, schema) do |writer|
+          writer.write_record_batch(record_batch)
         end
-        f.flush
 
-        data = File.binread(f.path)
-        start_index = get_record_batch_index(data)
-        data.bytes[4..start_index - 5].pack('C*')
+        bytes = buffer.data.to_s
+        start_index = get_record_batch_index(bytes)
+        bytes[4..start_index - 5]
       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
+      buffer = Arrow::ResizableBuffer.new(0)
+      Arrow::BufferOutputStream.open(buffer) do |output|
+        Arrow::RecordBatchStreamWriter.open(output, schema) do |writer|
+          writer.write_record_batch(record_batch)
         end
-        f.flush
 
-        data = File.binread(f.path)
-        start_index = get_record_batch_index(data)
-        data.bytes[start_index..].pack('C*')
+        bytes = buffer.data.to_s
+        start_index = get_record_batch_index(bytes)
+        bytes[start_index..]
       end
     end
 
-    def get_record_batch_index(raw)
-      size = raw.bytes.size
+    def get_record_batch_index(bytes)
+      size = 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
+        has_ffff = bytes.slice(i, 4) == [255, 255, 255, 255]
         found_count += 1 if has_ffff
         if found_count == 2
           start_index = i + 4

data/lib/athena-udf/version.rb

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

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: athena-udf
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.1.2
 platform: ruby
 authors:
 - Daisuke Taniwaki
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2024-08-23 00:00:00.000000000 Z
+date: 2024-09-03 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: base64