athena-udf 0.1.1 → 0.1.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 1ddf95a3316e3f1c9f3a6b21d067c4412aef09cc34982a5be141428c64974674
-  data.tar.gz: a9a189b6a6a1d2651f96a198f75829b22cd1289bfd41f3a2543ec408531ae064
+  metadata.gz: ab79db7e49eadd3d4448779185ecfc592f497e0a81209964ae7c2b8e4fdbeed5
+  data.tar.gz: 3c2b577e2506f67b9c0e455d8e3988b844fb5762dba69fa91c508d75d5de18fe
 SHA512:
-  metadata.gz: 83ce24ec41d03b9cf909c1a1f7b34cf952a79af775eabc00b2769cb6f862fd44c255a74961b9a99207f8581789b05f8202a3b09fb7b03f3785dde385c90f78d1
-  data.tar.gz: 16a47b3b3504e204c65bd71590a8c63618ad5fd25f0453504874f91ba1aed92624a77b3e660ece6815eb31d096cfab6fcfc680e2bbd6ff5943e235a2c24c838b
+  metadata.gz: 03b34a9703a30047eaeaf03c2bf8b975f190bf1f7ad08c30f1f06f8dcf2b6a7eca90b29621437699ce9acd22579bc44f59d03941493fb27adcb173caa7488f7b
+  data.tar.gz: c9dc25db956961249e40e4909ec284b5234622c05933c0fa78347431344afd5e5ca09cc6c36fa6b480ff2b5e4893d095e94369e6a07fc72917d8df279d9791f8

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    athena-udf (0.1.0)
+    athena-udf (0.1.3)
       base64
       csv
       red-arrow (~> 12.0.1)

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,61 +19,55 @@ 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) == "\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.1'
+  VERSION = '0.1.3'
 end

metadata

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