parqueteur 1.0.3 → 1.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 51c58dad3007e9c6a08b4673b378975c7839158240c27fb0a883d0e59c05d5d5
-  data.tar.gz: fa63fe9cd3cfa71d587400f3672ca2697c87a77fc574d8e2463d89f0090cd686
+  metadata.gz: 90ffcadf5b78e4ffc3329eac9be1be34af9209a77a39a6395eefc1c56afa7ce6
+  data.tar.gz: 1d7f5257d3f86443e0d13b789d7449565198f6cda1563d111a36ad3728264044
 SHA512:
-  metadata.gz: 58c98870d35b4af03bf52fea1572d26b4073d16696637efe8992eba9ddfbe8931e3b015976cd997ee85fc1a0388b031ecafe1b165b911a36d6748a90ce8af88d
-  data.tar.gz: 288418ec8a410f94fe38e76d94bb2da4173a4760ed8a9825132a41f3d808370ff60c8975ea8e6b8b72c1198dd36e5b4ceb1682570edf4ae9fa60be93ddc3cba8
+  metadata.gz: 262039094dd3aa5890f9d1a87d836eb64004c51fc57456d7880a36332295ac8a447fcd045c6bd2e053986ad4e9981021448b39e20010d0e41fef6b7e233d91ca
+  data.tar.gz: 73988e1f836acbe22e26c20b8559d8f79e9eeda55f9612fbefae7aa5bef131187b7a0716a3bd7804db5321b29b19f54b383e8594830814ae765d2c7b30986f59

data/.gitignore

@@ -7,3 +7,5 @@
 /spec/reports/
 /tmp/
 /*.gem
+/tmp/*
+!/tmp/.keep

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    parqueteur (1.0.3)
+    parqueteur (1.1.0)
       red-parquet (~> 5.0)
 
 GEM

data/example.rb

@@ -2,19 +2,38 @@ require 'bundler/setup'
 require 'parqueteur'
 
 class Foo < Parqueteur::Converter
-  column :id, :long
+  column :id, :bigint
   column :reference, :string
   column :hash, :map, key: :string, value: :string
   column :valid, :boolean
   column :total, :integer
+  column :numbers, :array, elements: :integer
+  column :my_struct, :struct do
+    field :test, :string
+    field :mon_nombre, :integer
+  end
 end
 
 LETTERS = ('a'..'z').to_a
 
 data = 1000.times.collect do |i|
-  { 'id' => i + 1, 'reference' => "coucou:#{i}", 'hash' => { 'a' => LETTERS.sample }, 'valid' => rand < 0.5, 'total' => rand(100..500) }
+  {
+    'id' => i + 1,
+    'reference' => "coucou:#{i}",
+    'hash' => { 'a' => LETTERS.sample },
+    'valid' => rand < 0.5,
+    'total' => rand(100..500),
+    'numbers' => [1, 2, 3],
+    'my_struct' => {
+      'test' => 'super'
+    }
+  }
 end
 
-chunked_converter = Parqueteur::ChunkedConverter.new(data, Foo)
-pp chunked_converter.write_files('test')
-# puts Foo.convert(data, output: 'test.parquet')
+# chunked_converter = Parqueteur::ChunkedConverter.new(data, Foo)
+# pp chunked_converter.write_files('test')
+puts Foo.convert(data, output: 'tmp/test.parquet')
+table = Arrow::Table.load('tmp/test.parquet')
+table.each_record do |record|
+  puts record.to_h
+end

data/lib/parqueteur/column.rb

@@ -4,31 +4,14 @@ module Parqueteur
   class Column
     attr_reader :name, :type, :options
 
-    def initialize(name, type, options = {})
+    def initialize(name, type, options = {}, &block)
       @name = name.to_s
-      @type = type
+      @type = Parqueteur::TypeResolver.resolve(type, options, &block)
       @options = options
     end
 
     def arrow_type
-      @arrow_type ||= Parqueteur::TypeResolver.resolve(@type, @options)
-    end
-
-    def cast(value)
-      case @type
-      when :string then value.to_s
-      when :boolean then value == true
-      when :integer then value.to_i
-      when :long then value.to_i
-      when :timestamp
-        case value
-        when String then Time.parse(value).to_i
-        when Integer then value
-        else
-          raise ArgumentError, "Unable to cast '#{value}' to timestamp"
-        end
-      when :map then value
-      end
+      @type.arrow_type
     end
 
     def to_arrow_field

data/lib/parqueteur/column_collection.rb

@@ -4,11 +4,18 @@ module Parqueteur
   class ColumnCollection
     include Enumerable
 
+    attr_reader :column_names
+
     def initialize
       @columns = []
+      @column_names = []
       @columns_idx = {}
     end
 
+    def key?(key)
+      @columns_idx.key?(key)
+    end
+
     def each(&block)
       @columns.each(&block)
     end
@@ -17,6 +24,7 @@ module Parqueteur
       unless @columns_idx.key?(column.name)
         @columns_idx[column.name] = column
         @columns << column
+        @column_names << column.name
       end
 
       true

data/lib/parqueteur/converter.rb

@@ -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
@@ -69,7 +69,10 @@ module Parqueteur
     def to_arrow_table
       transforms = self.class.transforms
 
-      chunks = {}
+      chunks = self.class.columns.each_with_object({}) do |column, hash|
+        hash[column.name] = []
+      end
+      items_count = 0
       @input.each_slice(100) do |items|
         values = self.class.columns.each_with_object({}) do |column, hash|
           hash[column.name] = []
@@ -98,20 +101,27 @@ module Parqueteur
 
         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
-            )
+            column.type.build_value_array(item[1])
           )
         end
+
+        items_count += items.length
       end
 
-      Arrow::Table.new(
-        chunks.transform_values! do |value|
-          Arrow::ChunkedArray.new(value)
-        end
-      )
+      if items_count > 0
+        Arrow::Table.new(
+          chunks.transform_values! do |value|
+            Arrow::ChunkedArray.new(value)
+          end
+        )
+      else
+        Arrow::Table.new(
+          self.class.columns.each_with_object({}) do |column, hash|
+            hash[column.name] = column.type.build_value_array([])
+          end
+        )
+      end
     end
   end
 end

data/lib/parqueteur/input.rb

@@ -13,7 +13,7 @@ module Parqueteur
           else
             arg.split("\n")
           end
-        when Array, Enumerator
+        when Enumerable
           arg
         end,
         options
@@ -36,7 +36,7 @@ module Parqueteur
           JSON.parse(@source.read).each(&block)
         end
         @source.rewind
-      when Array, Enumerator
+      when Enumerable
         @source.each(&block)
       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,46 @@
 
 module Parqueteur
   class TypeResolver
-    def self.resolve(*args)
-      new.resolve(*args)
+    include Singleton
+
+    def self.registered_types
+      @registered_types ||= {
+        array: Parqueteur::Types::ArrayType,
+        bigint: Parqueteur::Types::Int64Type,
+        boolean: Parqueteur::Types::BooleanType,
+        int32: Parqueteur::Types::Int32Type,
+        int64: Parqueteur::Types::Int64Type,
+        integer: Parqueteur::Types::Int32Type,
+        map: Parqueteur::Types::MapType,
+        string: Parqueteur::Types::StringType,
+        struct: Parqueteur::Types::StructType,
+        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/int32_type.rb

@@ -0,0 +1,25 @@
+# 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
+
+# when :integer

data/lib/parqueteur/types/int64_type.rb

@@ -0,0 +1,25 @@
+# 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
+
+# when :integer

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

data/lib/parqueteur/types/timestamp_type.rb

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

data/lib/parqueteur/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Parqueteur
-  VERSION = '1.0.3'
+  VERSION = '1.1.0'
 end

data/lib/parqueteur.rb

@@ -1,17 +1,29 @@
 # frozen_string_literal: true
 
+require 'json'
+require 'singleton'
+
 require_relative "parqueteur/version"
-require 'parqueteur/type_resolver'
+require 'parqueteur/chunked_converter'
 require 'parqueteur/column'
 require 'parqueteur/column_collection'
 require 'parqueteur/converter'
-require 'parqueteur/chunked_converter'
 require 'parqueteur/input'
-require 'parqueteur/value_array_builder'
-require 'json'
+require 'parqueteur/struct'
+require 'parqueteur/type'
+require 'parqueteur/type_resolver'
+require 'parqueteur/types/array_type'
+require 'parqueteur/types/boolean_type'
+require 'parqueteur/types/int32_type'
+require 'parqueteur/types/int64_type'
+require 'parqueteur/types/map_type'
+require 'parqueteur/types/string_type'
+require 'parqueteur/types/struct_type'
+require 'parqueteur/types/timestamp_type'
 require 'parquet'
 
 module Parqueteur
   class Error < StandardError; end
+  class TypeNotFound < Error; end
   # Your code goes here...
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: parqueteur
 version: !ruby/object:Gem::Version
-  version: 1.0.3
+  version: 1.1.0
 platform: ruby
 authors:
 - Julien D.
@@ -45,8 +45,17 @@ files:
 - lib/parqueteur/column_collection.rb
 - lib/parqueteur/converter.rb
 - lib/parqueteur/input.rb
+- lib/parqueteur/struct.rb
+- lib/parqueteur/type.rb
 - lib/parqueteur/type_resolver.rb
-- lib/parqueteur/value_array_builder.rb
+- lib/parqueteur/types/array_type.rb
+- lib/parqueteur/types/boolean_type.rb
+- lib/parqueteur/types/int32_type.rb
+- lib/parqueteur/types/int64_type.rb
+- lib/parqueteur/types/map_type.rb
+- lib/parqueteur/types/string_type.rb
+- lib/parqueteur/types/struct_type.rb
+- lib/parqueteur/types/timestamp_type.rb
 - lib/parqueteur/version.rb
 - parqueteur.gemspec
 - test.json

data/lib/parqueteur/value_array_builder.rb

@@ -1,59 +0,0 @@
-# frozen_string_literal: true
-
-module Parqueteur
-  class ValueArrayBuilder
-    attr_reader :type, :options, :arrow_type
-
-    def self.build(input, type, options)
-      new(type, options).build(input)
-    end
-
-    def initialize(type, options)
-      @type = type
-      @options = options
-      @arrow_type = Parqueteur::TypeResolver.resolve(type, options)
-    end
-
-    def build(input)
-      return if input.nil?
-
-      case type
-      when :array
-        Arrow::ListArrayBuilder.build(arrow_type, input)
-      when :map
-        builder = Arrow::MapArrayBuilder.new(arrow_type)
-        input.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
-      when :boolean
-        Arrow::BooleanArray.new(input)
-      when :integer
-        if options.fetch(:unsigned, false) == true
-          Arrow::UInt32Array.new(input)
-        else
-          Arrow::Int32Array.new(input)
-        end
-      when :long
-        if options.fetch(:unsigned, false) == true
-          Arrow::UInt64Array.new(input)
-        else
-          Arrow::Int64Array.new(input)
-        end
-      when :string
-        Arrow::StringArray.new(input)
-      when :timestamp
-        Arrow::TimestampArray.new(input)
-      else
-        raise Error, "unknown type: #{type}"
-      end
-    end
-  end
-end