parqueteur 1.0.3 → 1.3.0

data/lib/parqueteur/converter.rb

@@ -2,7 +2,7 @@
 
 module Parqueteur
   class Converter
-    attr_reader :schema
+    DEFAULT_BATCH_SIZE = 10
 
     def self.inline(&block)
       Class.new(self, &block)
@@ -12,8 +12,8 @@ module Parqueteur
       @columns ||= Parqueteur::ColumnCollection.new
     end
 
-    def self.column(name, type, options = {})
-      columns.add(Parqueteur::Column.new(name, type, options))
+    def self.column(name, type, options = {}, &block)
+      columns.add(Parqueteur::Column.new(name, type, options, &block))
     end
 
     def self.transforms
@@ -24,92 +24,135 @@ module Parqueteur
       transforms << (method_name || block)
     end
 
-    def self.convert(input, output: nil)
-      converter = new(input)
-      if !output.nil?
-        converter.write(output)
-      else
-        converter.to_blob
+    def self.convert(input, **kwargs)
+      new(input, **kwargs).to_io
+    end
+
+    def self.convert_to(input, output_path, **kwargs)
+      converter = new(input, **kwargs)
+      converter.write(output_path)
+    end
+
+    # @param [Enumerable] An enumerable object
+    # @option [Symbol] compression - :gzip
+    def initialize(input, **kwargs)
+      @input = Parqueteur::Input.from(input)
+      @batch_size = kwargs.fetch(:batch_size, DEFAULT_BATCH_SIZE)
+      @compression = kwargs.fetch(:compression, nil)&.to_sym
+    end
+
+    def split(size, batch_size: nil, compression: nil)
+      Enumerator.new do |arr|
+        options = {
+          batch_size: batch_size || @batch_size,
+          compression: compression || @compression
+        }
+        @input.each_slice(size) do |records|
+          local_converter = self.class.new(records, **options)
+          file = local_converter.to_tmpfile
+          arr << file
+          file.close
+          file.unlink
+        end
       end
     end
 
-    def initialize(input, options = {})
-      @input = Parqueteur::Input.from(input, options)
-    end
-
-    def write(output)
-      case output
-      when :io
-        to_io
-      when String
-        to_arrow_table.save(output)
-      when StringIO, IO
-        buffer = Arrow::ResizableBuffer.new(0)
-        to_arrow_table.save(buffer, format: :parquet)
-        output.write(buffer.data.to_s)
-        output.rewind
-        output
-      else
-        raise ArgumentError, "unsupported output: #{output.class}, accepted: String (filename), IO, StringIO"
+    def split_by_io(size, batch_size: nil, compression: nil)
+      Enumerator.new do |arr|
+        options = {
+          batch_size: batch_size || @batch_size,
+          compression: compression || @compression
+        }
+        @input.each_slice(size) do |records|
+          local_converter = self.class.new(records, **options)
+          arr << local_converter.to_io
+        end
+      end
+    end
+
+    def write(path, batch_size: nil, compression: nil)
+      compression = @compression if compression.nil?
+      batch_size = @batch_size if batch_size.nil?
+      arrow_schema = self.class.columns.arrow_schema
+      writer_properties = Parquet::WriterProperties.new
+      if !compression.nil? && compression != false
+        writer_properties.set_compression(compression)
+      end
+
+      Arrow::FileOutputStream.open(path, false) do |output|
+        Parquet::ArrowFileWriter.open(arrow_schema, output, writer_properties) do |writer|
+          @input.each_slice(batch_size) do |records|
+            arrow_table = build_arrow_table(records)
+            writer.write_table(arrow_table, 1024)
+          end
+        end
       end
+
+      true
     end
 
-    def to_s
-      inspect
+    def to_tmpfile(options = {})
+      tempfile = Tempfile.new
+      tempfile.binmode
+      write(tempfile.path, **options)
+      tempfile.rewind
+      tempfile
     end
 
-    def to_io
-      write(StringIO.new)
+    def to_io(options = {})
+      tmpfile = to_tmpfile(options)
+      strio = StringIO.new(tmpfile.read)
+      tmpfile.close
+      tmpfile.unlink
+      strio
     end
 
-    def to_blob
-      write(StringIO.new).read
+    def to_arrow_table(options = {})
+      file = to_tmpfile(options)
+      table = Arrow::Table.load(file.path, format: :parquet)
+      file.close
+      file.unlink
+      table
     end
 
-    def to_arrow_table
-      transforms = self.class.transforms
+    def to_blob(options = {})
+      to_tmpfile(options).read
+    end
 
-      chunks = {}
-      @input.each_slice(100) do |items|
-        values = self.class.columns.each_with_object({}) do |column, hash|
-          hash[column.name] = []
-        end
+    private
 
-        items.each do |item|
-          if transforms.length > 0
-            transforms.each do |transform|
-              item = \
-                if transform.is_a?(Symbol)
-                  __send__(transform, item)
-                else
-                  transform.call(item)
-                end
-            end
-          end
+    def build_arrow_table(records)
+      transforms = self.class.transforms
+
+      values = self.class.columns.each_with_object({}) do |column, hash|
+        hash[column.name] = []
+      end
 
-          values.each_key do |value_key|
-            if item.key?(value_key)
-              values[value_key] << item[value_key]
-            else
-              values[value_key] << nil
-            end
+      records.each do |item|
+        if transforms.length > 0
+          transforms.each do |transform|
+            item = \
+              if transform.is_a?(Symbol)
+                __send__(transform, item)
+              else
+                transform.call(item)
+              end
           end
         end
 
-        values.each_with_object(chunks) do |item, hash|
-          column = self.class.columns.find(item[0])
-          hash[item[0]] ||= []
-          hash[item[0]].push(
-            Parqueteur::ValueArrayBuilder.build(
-              item[1], column.type, column.options
-            )
-          )
+        values.each_key do |value_key|
+          if item.key?(value_key)
+            values[value_key] << item[value_key]
+          else
+            values[value_key] << nil
+          end
         end
       end
 
       Arrow::Table.new(
-        chunks.transform_values! do |value|
-          Arrow::ChunkedArray.new(value)
+        values.each_with_object({}) do |item, hash|
+          column = self.class.columns.find(item[0])
+          hash[item[0]] = column.type.build_value_array(item[1])
         end
       )
     end

data/lib/parqueteur/input.rb

@@ -4,40 +4,25 @@ module Parqueteur
   class Input
     include Enumerable
 
-    def self.from(arg, options = {})
-      new(
-        case arg
-        when String
-          if File.exist?(arg)
-            File.new(arg, 'r')
-          else
-            arg.split("\n")
-          end
-        when Array, Enumerator
-          arg
-        end,
-        options
-      )
+    def self.from(arg)
+      return arg if arg.is_a?(self)
+
+      new(arg)
     end
 
-    def initialize(source, options = {})
+    def initialize(source)
+      unless source.is_a?(Enumerable)
+        raise ArgumentError, 'Enumerable object expected'
+      end
+
       @source = source
-      @options = options
     end
 
     def each(&block)
-      case @source
-      when File
-        if @options.fetch(:json_newlines, true) == true
-          @source.each_line do |line|
-            yield(JSON.parse(line.strip))
-          end
-        else
-          JSON.parse(@source.read).each(&block)
-        end
-        @source.rewind
-      when Array, Enumerator
+      if block_given?
         @source.each(&block)
+      else
+        @source.to_enum(:each)
       end
     end
   end

data/lib/parqueteur/struct.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  class Struct
+    def initialize(&block)
+      instance_exec(&block)
+    end
+
+    def fields
+      @fields ||= Parqueteur::ColumnCollection.new
+    end
+
+    def field(name, type, options = {}, &block)
+      fields.add(Parqueteur::Column.new(name, type, options, &block))
+    end
+
+    def key?(key)
+      fields.key?(key)
+    end
+
+    def to_arrow_type
+      fields.collect(&:to_arrow_field)
+    end
+  end
+end

data/lib/parqueteur/type.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  class Type
+    attr_reader :options, :arrow_type
+
+    def initialize(options = {}, &block)
+      @options = options
+      @block = block
+      @arrow_type = arrow_type_builder
+    end
+
+    def build_value_array(values)
+      raise "#to_arrow_field must be implemented in #{self.class}"
+    end
+
+    def resolve(type, options = {})
+      Parqueteur::TypeResolver.resolve(type, options)
+    end
+  end
+end

data/lib/parqueteur/type_resolver.rb

@@ -2,60 +2,56 @@
 
 module Parqueteur
   class TypeResolver
-    def self.resolve(*args)
-      new.resolve(*args)
+    include Singleton
+
+    def self.registered_types
+      @registered_types ||= {
+        array: Parqueteur::Types::ArrayType,
+        bigdecimal: Parqueteur::Types::Decimal256Type,
+        bigint: Parqueteur::Types::Int64Type,
+        boolean: Parqueteur::Types::BooleanType,
+        date: Parqueteur::Types::Date32Type,
+        date32: Parqueteur::Types::Date64Type,
+        date64: Parqueteur::Types::Date64Type,
+        decimal: Parqueteur::Types::Decimal128Type,
+        decimal128: Parqueteur::Types::Decimal128Type,
+        decimal256: Parqueteur::Types::Decimal256Type,
+        int32: Parqueteur::Types::Int32Type,
+        int64: Parqueteur::Types::Int64Type,
+        integer: Parqueteur::Types::Int32Type,
+        map: Parqueteur::Types::MapType,
+        string: Parqueteur::Types::StringType,
+        struct: Parqueteur::Types::StructType,
+        time: Parqueteur::Types::Time32Type,
+        time32: Parqueteur::Types::Time32Type,
+        time64: Parqueteur::Types::Time64Type,
+        timestamp: Parqueteur::Types::TimestampType
+      }
+    end
+
+    def self.register_type(type, klass)
+      registered_types[type] = klass
+    end
+
+    def self.resolve(*args, &block)
+      instance.resolve(*args, &block)
     end
 
-    def resolve(type, options = {})
-      case type
-      when :array
-        elements_opt = options.fetch(:elements)
-        Arrow::ListDataType.new(
-          if elements_opt.is_a?(Hash)
-            resolve(elements_opt.fetch(:type), elements_opt)
-          else
-            resolve(elements_opt)
-          end
-        )
-      when :boolean
-        Arrow::BooleanDataType.new
-      when :integer
-        if options.fetch(:unsigned, false) == true
-          Arrow::UInt32DataType.new
-        else
-          Arrow::Int32DataType.new
-        end
-      when :long
-        if options.fetch(:unsigned, false) == true
-          Arrow::UInt64DataType.new
-        else
-          Arrow::Int64DataType.new
-        end
-      when :timestamp
-        Arrow::TimestampDataType.new(
-          options.fetch(:unit, :second)
-        )
-      when :string
-        Arrow::StringDataType.new
-      when :map
-        map_value = options.fetch(:value)
-        Arrow::MapDataType.new(
-          resolve(options.fetch(:key)),
-          if map_value.is_a?(Hash)
-            resolve(map_value.fetch(:type), map_value)
-          else
-            resolve(map_value)
-          end
-        )
+    def resolve(type, options = {}, &block)
+      if type.is_a?(Symbol)
+        resolve_from_symbol(type, options, &block)
       else
-        raise Error, "unknown type: #{type}"
+        type.new(options, &block)
       end
     end
-  end
-end
 
-private
+    private
 
-def build_arrow_type(type, options = {})
+    def resolve_from_symbol(type, options, &block)
+      type_klass = self.class.registered_types.fetch(type.to_sym, nil)
+      raise Parqueteur::TypeNotFound, type if type_klass.nil?
 
+      type_klass.new(options, &block)
+    end
+  end
 end

data/lib/parqueteur/types/array_type.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  module Types
+    class ArrayType < Parqueteur::Type
+      def build_value_array(values)
+        Arrow::ListArrayBuilder.build(arrow_type, values)
+      end
+
+      def arrow_type_builder
+        Arrow::ListDataType.new(
+          if options[:elements].is_a?(Hash)
+            resolve(options[:elements].fetch(:type), options[:elements]).arrow_type
+          else
+            resolve(options[:elements]).arrow_type
+          end
+        )
+      end
+    end
+  end
+end

data/lib/parqueteur/types/boolean_type.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  module Types
+    class BooleanType < Parqueteur::Type
+      def build_value_array(values)
+        Arrow::BooleanArray.new(values)
+      end
+
+      def arrow_type_builder
+        Arrow::BooleanDataType.new
+      end
+    end
+  end
+end

data/lib/parqueteur/types/date32_type.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  module Types
+    class Date32Type < Parqueteur::Type
+      def build_value_array(values)
+        Arrow::Date32ArrayBuilder.build(values)
+      end
+
+      def arrow_type_builder
+        Arrow::Date32DataType.new
+      end
+    end
+  end
+end

data/lib/parqueteur/types/date64_type.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  module Types
+    class Date64Type < Parqueteur::Type
+      def build_value_array(values)
+        Arrow::Date64ArrayBuilder.build([values])
+      end
+
+      def arrow_type_builder
+        Arrow::Date64DataType.new
+      end
+    end
+  end
+end

data/lib/parqueteur/types/decimal128_type.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  module Types
+    class Decimal128Type < Parqueteur::Type
+      def build_value_array(values)
+        Arrow::Decimal128ArrayBuilder.build(@arrow_type, values)
+      end
+
+      def arrow_type_builder
+        Arrow::Decimal128DataType.new(
+          precision: @options.fetch(:precision),
+          scale: @options.fetch(:scale)
+        )
+      end
+    end
+  end
+end

data/lib/parqueteur/types/decimal256_type.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  module Types
+    class Decimal256Type < Parqueteur::Type
+      def build_value_array(values)
+        Arrow::Decimal256ArrayBuilder.build(@arrow_type, values)
+      end
+
+      def arrow_type_builder
+        Arrow::Decimal256DataType.new(
+          precision: @options.fetch(:precision),
+          scale: @options.fetch(:scale)
+        )
+      end
+    end
+  end
+end

data/lib/parqueteur/types/int32_type.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  module Types
+    class Int32Type < Parqueteur::Type
+      def build_value_array(values)
+        if options.fetch(:unsigned, false) == true
+          Arrow::UInt32Array.new(values)
+        else
+          Arrow::Int32Array.new(values)
+        end
+      end
+
+      def arrow_type_builder
+        if options.fetch(:unsigned, false) == true
+          Arrow::UInt32DataType.new
+        else
+          Arrow::Int32DataType.new
+        end
+      end
+    end
+  end
+end

data/lib/parqueteur/types/int64_type.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  module Types
+    class Int64Type < Parqueteur::Type
+      def build_value_array(values)
+        if options.fetch(:unsigned, false) == true
+          Arrow::UInt64Array.new(values)
+        else
+          Arrow::Int64Array.new(values)
+        end
+      end
+
+      def arrow_type_builder
+        if options.fetch(:unsigned, false) == true
+          Arrow::UInt64DataType.new
+        else
+          Arrow::Int64DataType.new
+        end
+      end
+    end
+  end
+end

data/lib/parqueteur/types/map_type.rb

@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+module Parqueteur
+  module Types
+    class MapType < Parqueteur::Type
+      def build_value_array(values)
+        builder = Arrow::MapArrayBuilder.new(arrow_type)
+        values.each do |entry|
+          builder.append_value
+          next if entry.nil?
+
+          entry.each do |k, v|
+            builder.key_builder.append(k)
+            builder.item_builder.append(v)
+          end
+        end
+
+        builder.finish
+      end
+
+      def arrow_type_builder
+        map_value = options.fetch(:value)
+
+        Arrow::MapDataType.new(
+          resolve(options.fetch(:key)).arrow_type,
+          if map_value.is_a?(Hash)
+            resolve(map_value.fetch(:type), map_value).arrow_type
+          else
+            resolve(map_value).arrow_type
+          end
+        )
+      end
+    end
+  end
+end
+

data/lib/parqueteur/types/string_type.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+# when :timestamp
+#   Arrow::TimestampDataType.new(
+#     options.fetch(:unit, :second)
+#   )
+
+module Parqueteur
+  module Types
+    class StringType < Parqueteur::Type
+      def build_value_array(values)
+        Arrow::StringArray.new(values)
+      end
+
+      def arrow_type_builder
+        Arrow::StringDataType.new
+      end
+    end
+  end
+end

data/lib/parqueteur/types/struct_type.rb

@@ -0,0 +1,35 @@
+# frozen_string_literal: true
+
+# when :timestamp
+#   Arrow::TimestampDataType.new(
+#     options.fetch(:unit, :second)
+#   )
+
+module Parqueteur
+  module Types
+    class StructType < Parqueteur::Type
+      def build_value_array(values)
+        values.each do |value|
+          next if value.nil?
+
+          value.each_key do |key|
+            next if struct_object.key?(key)
+
+            raise Parqueteur::Error, "Struct field '#{key}' not found"
+          end
+        end
+        Arrow::StructArrayBuilder.build(arrow_type, values)
+      end
+
+      def arrow_type_builder
+        Arrow::StructDataType.new(struct_object.to_arrow_type)
+      end
+
+      private
+
+      def struct_object
+        @struct_object ||= Parqueteur::Struct.new(&@block)
+      end
+    end
+  end
+end