easy_ml 0.2.0.pre.rc75 → 0.2.0.pre.rc77

data/lib/easy_ml/data/dataset_manager/writer/base.rb

@@ -0,0 +1,122 @@
+module EasyML
+  module Data
+    class DatasetManager
+      class Writer
+        class Base
+          attr_accessor :filenames, :root_dir, :options, :append_only, :df
+
+          def initialize(options)
+            @root_dir = options.dig(:root_dir)
+            @filenames = options.dig(:filenames)
+            @append_only = options.dig(:append_only)
+            @options = options
+            @df = options.dig(:df)
+          end
+
+          def wipe
+            clear_unique_id
+            FileUtils.rm_rf(root_dir)
+          end
+
+          def store
+            store_to_unique_file
+          end
+
+          def compact
+            files = self.files
+
+            clear_unique_id
+            File.join(root_dir, "compacted.parquet").tap do |target_file|
+              safe_write(
+                query(lazy: true),
+                target_file
+              )
+              FileUtils.rm(files)
+            end
+            clear_unique_id
+          end
+
+          private
+
+          def files
+            DatasetManager.new(options).files
+          end
+
+          def query(**kwargs, &block)
+            DatasetManager.new(options).query(root_dir, **kwargs, &block)
+          end
+
+          def store_to_unique_file(subdir: nil)
+            safe_write(df, unique_path(subdir: subdir))
+          end
+
+          def unique_path(subdir: nil)
+            filename = [filenames, unique_id(subdir: subdir), "parquet"].compact.join(".")
+
+            File.join(root_dir, subdir.to_s, filename)
+          end
+
+          def safe_write(df, path)
+            FileUtils.mkdir_p(File.dirname(path))
+            df.is_a?(Polars::LazyFrame) ? df.sink_parquet(path) : df.write_parquet(path)
+            path
+          end
+
+          def clear_all_keys
+            keys = list_keys
+            Support::Lockable.with_lock(keys, wait_timeout: 2) do |suo|
+              suo.client.del(keys)
+            end
+          end
+
+          def clear_unique_id(subdir: nil)
+            key = unique_id_key(subdir: subdir)
+            Support::Lockable.with_lock(key, wait_timeout: 2) do |suo|
+              suo.client.del(key)
+            end
+          end
+
+          def unique_id_key(subdir: nil)
+            File.join("dataset_managers", root_dir, subdir.to_s, "sequence")
+          end
+
+          def add_key(key)
+            keylist = unique_id_key(subdir: "keylist")
+
+            Support::Lockable.with_lock(keylist, wait_timeout: 2) do |suo|
+              suo.client.sadd(keylist, key)
+            end
+          end
+
+          def list_keys
+            keylist = unique_id_key(subdir: "keylist")
+
+            Support::Lockable.with_lock(keylist, wait_timeout: 2) do |suo|
+              suo.client.smembers(keylist)
+            end
+          end
+
+          def key_exists?(key)
+            keylist = unique_id_key(subdir: "keylist")
+            Support::Lockable.with_lock(keylist, wait_timeout: 2) do |suo|
+              suo.client.sismember(keylist, key)
+            end
+          end
+
+          def unique_id(subdir: nil)
+            key = unique_id_key(subdir: subdir)
+            add_key(key)
+
+            Support::Lockable.with_lock(key, wait_timeout: 2) do |suo|
+              redis = suo.client
+
+              seq = (redis.get(key) || "0").to_i
+              redis.set(key, (seq + 1).to_s)
+              seq + 1
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/easy_ml/data/dataset_manager/writer/named.rb

@@ -0,0 +1,14 @@
+module EasyML
+  module Data
+    class DatasetManager
+      class Writer
+        class Named < Base
+          def store(name)
+            clear_unique_id(subdir: name)
+            store_to_unique_file(subdir: name)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/easy_ml/data/dataset_manager/writer/partitioned/partition_reasons.rb

@@ -0,0 +1,15 @@
+module EasyML
+  module Data
+    class DatasetManager
+      class Writer
+        class Partitioned < Base
+          class PartitionReasons < EasyML::Reasons
+            add_reason "Missing primary key", -> { primary_key.nil? }
+            add_reason "Df does not contain primary key", -> { df.columns.exclude?(primary_key) }
+            add_reason "Primary key is not numeric", -> { !numeric_primary_key? }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/easy_ml/data/dataset_manager/writer/partitioned.rb

@@ -0,0 +1,150 @@
+module EasyML
+  module Data
+    class DatasetManager
+      class Writer
+        class Partitioned < Base
+          require_relative "partitioned/partition_reasons"
+
+          attr_accessor :partition_size, :partition, :primary_key, :df
+
+          def initialize(options)
+            super
+            @partition_size = options.dig(:partition_size)
+            @partition = options.dig(:partition)
+            @primary_key = options.dig(:primary_key)
+
+            raise "filenames required: specify the prefix to use for unique new files" unless filenames.present?
+          end
+
+          def wipe
+            partitions.each do |partition|
+              FileUtils.rm_rf(File.join(root_dir, partition))
+            end
+            clear_all_keys
+          end
+
+          def store
+            unless can_partition?
+              puts cannot_partition_reasons.explain
+              return Base.new(options).store
+            end
+
+            store_each_partition
+          end
+
+          def compact
+            files = self.files
+            @df = query(lazy: true)
+
+            clear_unique_id(subdir: "compacted")
+            compact_each_partition.tap do
+              FileUtils.rm(files)
+              clear_unique_id
+            end
+          end
+
+          private
+
+          def partitions
+            Dir.glob(File.join(root_dir, "**/*")).map { |f| f.split("/").last }
+          end
+
+          def compact_each_partition
+            with_each_partition do |partition_df, _|
+              safe_write(
+                partition_df.sort(Polars.col(primary_key)),
+                unique_path(subdir: "compacted")
+              )
+            end
+          end
+
+          def with_each_partition(&block)
+            partition_boundaries.map do |partition|
+              partition_start = partition[:partition_start]
+              partition_end = partition[:partition_end]
+              partition_df = df.filter(Polars.col(primary_key).is_between(partition_start, partition_end))
+              num_rows = lazy? ? partition_df.select(Polars.length).collect[0, 0] : partition_df.shape[0]
+
+              binding.pry if num_rows == 0
+              next if num_rows == 0
+              yield partition_df, partition
+            end
+          end
+
+          def store_each_partition
+            with_each_partition do |partition_df, partition|
+              safe_write(
+                partition_df,
+                unique_path(subdir: partition[:partition])
+              )
+            end
+          end
+
+          def partition_boundaries
+            EasyML::Data::Partition::Boundaries.new(df, primary_key, partition_size).to_a
+          end
+
+          def cannot_partition_reasons
+            @cannot_partition_reasons ||= PartitionReasons.new(self)
+          end
+
+          def can_partition?
+            @partitioned ||= cannot_partition_reasons.none?
+          end
+
+          def lazy?
+            df.is_a?(Polars::LazyFrame)
+          end
+
+          def cast_primary_key
+            case dtype_primary_key
+            when Polars::Categorical
+              Polars.col(primary_key).cast(Polars::String)
+            else
+              Polars.col(primary_key)
+            end
+          end
+
+          def dtype_primary_key
+            @dtype_primary_key ||= schema[primary_key]
+          end
+
+          def schema
+            @schema ||= df.schema
+          end
+
+          def min_key
+            return @min_key if @min_key
+
+            if lazy?
+              @min_key = df.select(cast_primary_key).min.collect.to_a[0].dig(primary_key)
+            else
+              @min_key = df[primary_key].min
+            end
+          end
+
+          def max_key
+            return @max_key if @max_key
+
+            if lazy?
+              @max_key = df.select(cast_primary_key).max.collect.to_a[0].dig(primary_key)
+            else
+              @max_key = df[primary_key].max
+            end
+          end
+
+          def numeric_primary_key?
+            begin
+              # We are intentionally not using to_i, so it will raise an error for keys like "A1"
+              min = min_key.is_a?(String) ? Integer(min_key) : min_key
+              max = max_key.is_a?(String) ? Integer(max_key) : max_key
+              min.is_a?(Integer) && max.is_a?(Integer)
+            rescue ArgumentError
+              false
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/easy_ml/data/dataset_manager/writer.rb

@@ -0,0 +1,76 @@
+module EasyML
+  module Data
+    class DatasetManager
+      class Writer
+        require_relative "writer/base"
+        require_relative "writer/partitioned"
+        require_relative "writer/append_only"
+        require_relative "writer/named"
+
+        ADAPTERS = [
+          Base,
+          Partitioned,
+          AppendOnly,
+          Named,
+        ]
+
+        attr_accessor :filenames, :root_dir, :partition,
+                      :append_only, :primary_key, :options
+
+        def initialize(options)
+          @root_dir = options.dig(:root_dir)
+          @filenames = options.dig(:filenames)
+          @partition = options.dig(:partition) || (options.dig(:partition_size).present? && options.dig(:primary_key).present?)
+          @append_only = options.dig(:append_only)
+          @primary_key = options.dig(:primary_key)
+          @named = options.dig(:named) || false
+          @options = options
+        end
+
+        def store(df, *args)
+          adapter_class.new(options.merge!(df: df)).store(*args)
+        end
+
+        def wipe
+          adapter_class.new(options).wipe
+        end
+
+        def compact
+          adapter_class.new(options).compact
+        end
+
+        def inspect
+          keys = %w(root_dir append_only partition primary_key)
+          attrs = keys.map { |k| "#{k}=#{send(k)}" unless send(k).nil? }.compact
+          "#<#{self.class.name} #{attrs.join(" ")}>"
+        end
+
+        private
+
+        def adapter_class
+          if partition?
+            Partitioned
+          elsif append_only?
+            AppendOnly
+          elsif named?
+            Named
+          else
+            Base
+          end
+        end
+
+        def named?
+          @named
+        end
+
+        def partition?
+          @partition
+        end
+
+        def append_only?
+          @append_only
+        end
+      end
+    end
+  end
+end

data/lib/easy_ml/data/dataset_manager.rb

@@ -0,0 +1,134 @@
+module EasyML
+  module Data
+    class DatasetManager
+      require_relative "dataset_manager/writer"
+      require_relative "dataset_manager/reader"
+
+      attr_accessor :root_dir, :partition, :append_only, :filenames, :primary_key,
+                    :partition_size, :s3_bucket, :s3_prefix, :s3_access_key_id,
+                    :s3_secret_access_key, :polars_args, :source_of_truth,
+                    :options
+
+      def initialize(options = {})
+        @root_dir = options.dig(:root_dir)
+        @partition = options.dig(:partition) || (options.dig(:partition_size).present? && options.dig(:primary_key).present?)
+        @append_only = options.dig(:append_only) || false
+        @filenames = options.dig(:filenames) || "file"
+        @primary_key = options.dig(:primary_key)
+        @partition_size = options.dig(:partition_size) || nil
+        @s3_bucket = options.dig(:s3_bucket) || EasyML::Configuration.s3_bucket
+        @s3_prefix = options.dig(:s3_prefix) || nil
+        @s3_access_key_id = options.dig(:s3_access_key_id) || EasyML::Configuration.s3_access_key_id
+        @s3_secret_access_key = options.dig(:s3_secret_access_key) || EasyML::Configuration.s3_secret_access_key
+        @polars_args = options.dig(:polars_args) || {}
+        @source_of_truth = options.dig(:source_of_truth) || :local
+        @options = options
+
+        raise "primary_key required: how should we divide partitions?" if partition && primary_key.nil?
+        raise "partition_size required: specify number of rows in each partition" if partition && partition_size.nil?
+        raise "root_dir required: specify the root_dir of the dataset" unless root_dir.present?
+      end
+
+      def inspect
+        keys = %w(root append_only partition primary_key)
+        attrs = keys.map { |k| "#{k}=#{send(k)}" unless send(k).nil? }.compact
+        "#<#{self.class.name} #{attrs.join("\n\t")}>"
+      end
+
+      class << self
+        def query(input = nil, **kwargs, &block)
+          Reader.query(input, **kwargs, &block)
+        end
+
+        def schema(input = nil, **kwargs, &block)
+          Reader.schema(input, **kwargs, &block)
+        end
+
+        def num_rows
+          Reader.num_rows
+        end
+      end
+
+      def num_rows
+        Reader.num_rows(root_dir)
+      end
+
+      def query(input = nil, **kwargs, &block)
+        input = root_dir if input.nil?
+        DatasetManager.query(input, **kwargs, &block)
+      end
+
+      def schema(input = nil, **kwargs, &block)
+        input = root_dir if input.nil?
+        DatasetManager.schema(input, **kwargs, &block)
+      end
+
+      def sha
+        Reader.sha(root_dir)
+      end
+
+      def normalize
+        Normalizer.normalize(root_dir)
+      end
+
+      def data
+        query
+      end
+
+      def store(df, *args)
+        writer.store(df, *args)
+      end
+
+      def compact
+        writer.compact
+      end
+
+      def cp(from, to)
+        writer.cp(from, to)
+      end
+
+      def empty?
+        files.empty? || query(limit: 1).empty?
+      end
+
+      def files
+        Reader.files(root_dir)
+      end
+
+      def wipe
+        writer.wipe
+      end
+
+      def upload
+        synced_directory.upload
+      end
+
+      def download
+        synced_directory.download
+      end
+
+      private
+
+      def root
+        root_dir.gsub(/^#{Rails.root.to_s}/, "")
+      end
+
+      def writer
+        Writer.new(options)
+      end
+
+      def synced_directory
+        @synced_dir ||= EasyML::Data::SyncedDirectory.new(
+          root_dir: root_dir,
+          source_of_truth: source_of_truth,
+          s3_bucket: s3_bucket,
+          s3_prefix: s3_prefix,
+          s3_access_key_id: s3_access_key_id,
+          s3_secret_access_key: s3_secret_access_key,
+          polars_args: polars_args,
+          cache_for: 0,
+        )
+      end
+    end
+  end
+end

data/lib/easy_ml/data/partition/boundaries.rb

@@ -0,0 +1,60 @@
+module EasyML
+  module Data
+    module Partition
+      class Boundaries
+        attr_reader :df, :primary_key, :partition_size
+
+        def initialize(df, primary_key, partition_size)
+          @df = df
+          @primary_key = primary_key.is_a?(Array) ? primary_key.first : primary_key
+          @partition_size = partition_size
+        end
+
+        def inspect
+          "#<#{self.class.name.split("::").last} partition_size=#{partition_size} primary_key=#{primary_key}>"
+        end
+
+        def boundaries
+          return @boundaries if @boundaries
+
+          @boundaries = df.with_columns(
+            Polars.col(primary_key)
+              .truediv(partition_size)
+              .floor
+              .add(1)
+              .cast(Polars::Int64)
+              .alias("partition")
+          )
+          @boundaries = @boundaries.with_columns(
+            Polars.col("partition")
+                  .sub(1)
+                  .mul(partition_size)
+                  .cast(Polars::Int64)
+                  .alias("partition_start"),
+            Polars.col("partition")
+                  .mul(partition_size)
+                  .sub(1)
+                  .cast(Polars::Int64)
+                  .alias("partition_end")
+          )
+          # @boundaries = @boundaries.with_columns(
+          #   Polars.col(primary_key).is_between(Polars.col("partition_start"), Polars.col("partition_end")).select("partition")
+          # )
+        end
+
+        def to_a
+          is_lazy = df.is_a?(Polars::LazyFrame)
+          empty = is_lazy ? df.limit(1).collect.empty? : df.shape[0] == 0
+          return [] if empty
+
+          sorted = boundaries.select(["partition", "partition_start", "partition_end"]).unique.sort("partition")
+          array = (is_lazy ? sorted.collect.to_a : sorted.to_a).map(&:with_indifferent_access)
+          # For the last partition, set the end to the total number of rows (so we read the last row with is_between queries)
+          last_idx = array.size - 1
+          array[last_idx]["partition_end"] = is_lazy ? df.select(Polars.col(primary_key)).max.collect.to_a.first.dig(primary_key) : df[primary_key].max
+          array
+        end
+      end
+    end
+  end
+end

data/lib/easy_ml/data/partition.rb

@@ -0,0 +1,7 @@
+module EasyML
+  module Data
+    module Partition
+      require_relative "partition/boundaries"
+    end
+  end
+end

data/lib/easy_ml/data/synced_directory.rb

@@ -1,5 +1,3 @@
-require_relative "polars_reader"
-
 module EasyML
   module Data
     class SyncedDirectory
@@ -16,6 +14,7 @@ module EasyML
         @s3_region = options.dig(:s3_region) || EasyML::Configuration.s3_region
         @cache_for = options.dig(:cache_for)
         @polars_args = options.dig(:polars_args)
+        @source_of_truth = options.dig(:source_of_truth) || :remote
       end
 
       delegate :query, :data, :all_files, :files, :sha, to: :reader

data/lib/easy_ml/data.rb

@@ -8,5 +8,7 @@ module EasyML
     require_relative "data/polars_column"
     require_relative "data/polars_schema"
     require_relative "data/date_converter"
+    require_relative "data/dataset_manager"
+    require_relative "data/partition"
   end
 end