flat_kit 0.1.0

data/lib/flat_kit/jsonl.rb

@@ -0,0 +1,8 @@
+module FlatKit
+  module Jsonl
+  end
+end
+require 'flat_kit/jsonl/record'
+require 'flat_kit/jsonl/reader'
+require 'flat_kit/jsonl/writer'
+require 'flat_kit/jsonl/format'

data/lib/flat_kit/jsonl/format.rb

@@ -0,0 +1,25 @@
+module FlatKit
+  module Jsonl
+    class Format < ::FlatKit::Format
+      def self.format_name
+        "jsonl"
+      end
+
+      def self.handles?(filename)
+        parts = filename.split(".")
+        %w[ json jsonl ndjson ].each do |ext|
+          return true if parts.include?(ext)
+        end
+        return false
+      end
+
+      def self.reader
+        ::FlatKit::Jsonl::Reader
+      end
+
+      def self.writer
+        ::FlatKit::Jsonl::Writer
+      end
+    end
+  end
+end

data/lib/flat_kit/jsonl/reader.rb

@@ -0,0 +1,30 @@
+module FlatKit
+  module Jsonl
+    class Reader < ::FlatKit::Reader
+      attr_reader :input
+      attr_reader :count
+
+      def self.format_name
+        ::FlatKit::Jsonl::Format.format_name
+      end
+
+      def initialize(source:, compare_fields: :none)
+        super
+        @input = ::FlatKit::Input.from(source)
+        @count = 0
+      end
+
+      def each
+        while line = input.io.gets do
+          record = ::FlatKit::Jsonl::Record.new(data: line, compare_fields: compare_fields)
+          @count += 1
+          yield record
+        end
+        input.close
+      rescue => e
+        ::FlatKit.logger.error "Error reading jsonl records from #{input.name}: #{e}"
+        raise ::FlatKit::Error, e
+      end
+    end
+  end
+end

data/lib/flat_kit/jsonl/record.rb

@@ -0,0 +1,84 @@
+require 'oj'
+require 'flat_kit/record'
+
+module FlatKit
+  module Jsonl
+    class Record < ::FlatKit::Record
+      attr_reader :compare_data
+
+      def self.format_name
+        ::FlatKit::Jsonl::Format.format_name
+      end
+
+      def self.from_record(record)
+        if record.instance_of?(FlatKit::Jsonl::Record) then
+
+          structured = record.complete_structured_data? ? record.complete_structured_data : nil
+
+          new(data: record.data, compare_fields: record.compare_fields,
+              compare_data: record.compare_data,
+              complete_structured_data: structured)
+        else
+          new(data: nil, compare_fields: record.compare_fields,
+              complete_structured_data: record.to_hash)
+        end
+      end
+
+      def initialize(data:, compare_fields: :none,
+                     compare_data: Hash.new,
+                     complete_structured_data: nil)
+        super(data: data, compare_fields: compare_fields)
+
+        @complete_structured_data = complete_structured_data
+
+        if complete_structured_data? && (compare_data.nil? || compare_data.empty?) then
+          @compare_data = complete_structured_data
+        else
+          @compare_data = compare_data
+        end
+
+        # only load compare data if it dosn't exist
+        if data && compare_data.empty? then
+          quick_parse
+        end
+      end
+
+      def [](key)
+        compare_data[key]
+      end
+
+      def complete_structured_data
+        @complete_structured_data ||= Oj.load(data)
+      end
+      alias to_hash complete_structured_data
+
+      def complete_structured_data?
+        !(@complete_structured_data.nil? || @complete_structured_data.empty?)
+      end
+
+      # overriding parent accessor since we may be initialized without raw bytes
+      # to parse
+      def data
+        if @data.nil? && complete_structured_data? then
+          @data = Oj.dump(complete_structured_data)
+        end
+        @data
+      end
+      alias to_s data
+
+      private
+
+      def quick_parse
+        Oj::Doc.open(data) do |doc|
+          compare_fields.each do |field|
+            val = doc.fetch("/#{field}")
+            @compare_data[field] = val
+          end
+        end
+      end
+    end
+  end
+end
+
+
+

data/lib/flat_kit/jsonl/writer.rb

@@ -0,0 +1,45 @@
+module FlatKit
+  module Jsonl
+    class Writer < ::FlatKit::Writer
+      attr_reader :output
+      attr_reader :count
+
+      def self.format_name
+        ::FlatKit::Jsonl::Format.format_name
+      end
+
+      def initialize(destination:)
+        super
+        @output = ::FlatKit::Output.from(@destination)
+        @count = 0
+      end
+
+      def write(record)
+        case record
+        when FlatKit::Jsonl::Record
+          write_record(record)
+        when FlatKit::Record
+          converted_record = ::FlatKit::Jsonl::Record.from_record(record)
+          write_record(converted_record)
+        else
+          raise FlatKit::Error, "Unable to write records of type #{record.class}"
+        end
+      rescue FlatKit::Error => fe
+        raise fe
+      rescue => e
+        ::FlatKit.logger.error "Error reading jsonl records from #{output.name}: #{e}"
+        raise ::FlatKit::Error, e
+      end
+
+      def close
+        @output.close
+      end
+
+      def write_record(record)
+        # enforces ending in newlin if it doesn't already have one
+        output.io.puts record.to_s
+        @count += 1
+      end
+    end
+  end
+end

data/lib/flat_kit/leaf_node.rb

@@ -0,0 +1,71 @@
+module FlatKit
+  # Private: The LeafNode is a wrapper around a Reader object to enable
+  # a consistent api for use in the MergeTree
+  #
+  # The LeafNode keeps track of the head of the reader list and when its value
+  # is used up, it pulls the next value and then notifies the next level of the
+  # MergeTree that its value has changed and so should do another play.
+  #
+  # If all the data is used up from the reader, it also notifies the next level
+  # of that so the next level can remove it from the tree.
+  class LeafNode
+
+    include Comparable
+
+    attr_reader :reader
+    attr_reader :value
+
+    attr_accessor :next_level
+
+    def initialize(reader)
+      @reader     = reader
+      @enum       = @reader.to_enum
+      @value      = @enum.next
+
+      @next_level = nil
+    end
+
+    def winner
+      value
+    end
+
+    def sentinel?
+      false
+    end
+
+    def leaf?
+      true
+    end
+
+    def leaf
+      self
+    end
+
+    def update_and_replay
+      self.next
+      if finished? then
+        ::FlatKit.logger.debug "#{reader.source} has finished reading #{reader.count} records"
+        next_level.player_finished(self)
+      end
+      next_level.play
+    end
+
+    def next
+      begin
+        @value = @enum.next
+      rescue StopIteration
+        @value = nil
+      end
+      @value
+    end
+
+    def finished?
+      @enum && @value.nil?
+    end
+
+    def <=>(other)
+      return -1 if other.sentinel?
+      self.value.<=>(other.value)
+    end
+  end
+end

data/lib/flat_kit/logger.rb

@@ -0,0 +1,39 @@
+require 'logger'
+
+module FlatKit
+  class LogFormatter < ::Logger::Formatter
+    FORMAT          = "%s %5d %05s : %s\n".freeze
+    DATETIME_FORMAT = "%Y-%m-%dT%H:%M:%SZ".freeze
+    def initialize
+      super
+      self.datetime_format = DATETIME_FORMAT
+    end
+
+    def call(severity, time, progname, msg)
+      FORMAT % [format_datetime(time.utc), Process.pid, severity, msg2str(msg)]
+    end
+  end
+
+  class Logger
+    def self.for_io(io)
+      ::Logger.new(io, formatter: LogFormatter.new)
+    end
+
+    def self.for_path(path)
+      io = File.open(path.to_s, "a")
+      for_io(io)
+    end
+  end
+
+  def self.log_to(destination = $stderr)
+    if destination.kind_of?(::IO) then
+      @logger = ::FlatKit::Logger.for_io(destination)
+    else
+      @logger = ::FlatKit::Logger.for_path(destination)
+    end
+  end
+
+  def self.logger
+    @logger ||= ::FlatKit::Logger.for_io($stderr)
+  end
+end

data/lib/flat_kit/merge.rb

@@ -0,0 +1,35 @@
+module FlatKit
+  class Merge
+    attr_reader :readers
+    attr_reader :writer
+    attr_reader :compare_fields
+
+    def initialize(inputs:, input_fallback: "auto",
+                   output:, output_fallback: "auto",
+                   compare_fields:)
+      @compare_fields = compare_fields
+      @readers = ::FlatKit::Reader.create_readers_from_paths(paths: inputs, compare_fields: @compare_fields,
+                                                             fallback: input_fallback)
+      @writer  = ::FlatKit::Writer.create_writer_from_path(path: output, fallback: output_fallback,
+                                                           reader_format: @readers.first.format_name)
+    end
+
+    def call
+      ::FlatKit.logger.info "Merging the following files into #{writer.destination}"
+      ::FlatKit.logger.info "Using this key for sorting: #{compare_fields.join(", ")}"
+      readers.each do |r|
+        ::FlatKit.logger.info "  #{r.source}"
+      end
+
+      merge_tree = ::FlatKit::MergeTree.new(readers)
+      merge_tree.each do |record|
+        writer.write(record)
+      end
+      readers.each do |r|
+        ::FlatKit.logger.info "  #{r.source} produced #{r.count} records"
+      end
+      writer.close
+      ::FlatKit.logger.info "Wrote #{writer.count} records to #{writer.destination}"
+    end
+  end
+end

data/lib/flat_kit/merge_tree.rb

@@ -0,0 +1,104 @@
+module FlatKit
+  # Public: Merge a list of sorted records from Readers into a single output Writer
+  #
+  # The MergeTree implements a Tournament Tree algorightm to do a k-way merge
+  # between Reader objects:
+  #
+  #   https://en.wikipedia.org/wiki/K-way_merge_algorithm
+  #
+  # Usage:
+  #
+  #   compare_fields = %w[ key timestamp ]
+  #   format = ::FlatKit::Format.for('json')
+  #
+  #   readers = ARGV.map do |path|
+  #     format.reader.new(source: path, compare_fields: compare_fields)
+  #   end
+  #
+  #   write_path = "merged.jsonl"
+  #   writer = format.writer.new(destination: write_path.to_s)
+  #
+  #   tree = ::FlatKit::MergeTree.new(readers)
+  #
+  #   tree.each do |record|
+  #      writer.write(record)
+  #   end
+  #   writer.close
+  #
+  #
+  class MergeTree
+    include Enumerable
+
+    attr_reader :leaves
+    attr_reader :levels
+    attr_reader :readers
+
+    def initialize(readers)
+      @readers = readers
+      @leaves = []
+      @levels = []
+
+      @readers.each do |reader|
+        @leaves << LeafNode.new(reader)
+      end
+
+      # Need to pad the leaves to an even number so that the slicing by 2 for
+      # the tournament will work
+      if @leaves.size.odd? then
+        @leaves << SentinelLeafNode.new
+      end
+
+      init_tree
+    end
+
+    #
+    # Initialize the tournament tree, go in depths - bottom layer will be the
+    # winners of the 2 leaf nodes, continuing until top layer is just 1 node
+    #
+    def init_tree
+      values = @leaves.dup
+      loop do
+        break if values.size == 1
+
+        winners = []
+
+        # we alays need a left and a right node, there is the possibility of
+        # adding a single Sentinel node as the final right node in each level
+        values.each_slice(2) do |left, right|
+          right = SentinelInternalNode.new if right.nil?
+          winners << InternalNode.new(left: left, right: right)
+        end
+        values = winners
+        @levels << winners
+      end
+    end
+
+    #
+    # Root is the last level - should only have one node
+    #
+    def root
+      @levels[-1].first
+    end
+
+    #
+    # The number of levels, this shold be the logn(readers.size)
+    #
+    def depth
+      @levels.size
+    end
+
+    #
+    # Iterate over all the ements from all the readers yielding them in sorted
+    # order.
+    #
+    def each
+      loop do
+        break if root.leaf.finished?
+        yield root.value
+        # consume the yielded value and have the tournament tree replay those
+        # brackets affected
+        root.leaf.update_and_replay
+      end
+    end
+  end
+end