easy_ml 0.2.0.pre.rc76 → 0.2.0.pre.rc78

data/lib/easy_ml/data/dataset_manager/reader.rb

@@ -0,0 +1,58 @@
+module EasyML
+  module Data
+    class DatasetManager
+      class Reader
+        require_relative "reader/base"
+        require_relative "reader/file"
+        require_relative "reader/data_frame"
+
+        ADAPTERS = [
+          File,
+          DataFrame,
+        ]
+
+        def self.query(input, **kwargs, &block)
+          adapter(input).new(
+            kwargs.merge!(input: input), &block
+          ).query
+        end
+
+        def self.schema(input, **kwargs, &block)
+          adapter(input).new(
+            kwargs.merge!(input: input), &block
+          ).schema
+        end
+
+        def self.files(dir)
+          Dir.glob(::File.join(dir, "**/*.{parquet}"))
+        end
+
+        def self.sha
+          files = sha.sort
+
+          file_hashes = files.map do |file|
+            meta = Polars.read_parquet_schema(file)
+            row_count = Polars.scan_parquet(file).select(Polars.col("*").count).collect[0, 0]
+
+            Digest::SHA256.hexdigest([
+              meta.to_json,
+              row_count.to_s,
+            ].join("|"))
+          end
+
+          Digest::SHA256.hexdigest(file_hashes.join)
+        end
+
+        private
+
+        def self.adapter(input)
+          if input.is_a?(Polars::DataFrame) || input.is_a?(Polars::LazyFrame)
+            DataFrame
+          else
+            File
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,67 @@
+module EasyML
+  module Data
+    class DatasetManager
+      class Writer
+        class AppendOnly < Base
+          attr_accessor :primary_key
+
+          def initialize(options)
+            super
+            @primary_key = options.dig(:primary_key)
+            raise "primary_key required for append_only writer" if primary_key.nil?
+            raise "filenames required: specify the prefix to uuse for unique new files" unless filenames.present?
+          end
+
+          def store
+            # If there are no existing files, just store as normal
+            return super if files.empty?
+
+            # Get existing data lazily
+            existing_keys = query(lazy: true)
+              .select(primary_key)
+              .collect[primary_key]
+              .to_a
+
+            # Convert input to lazy if it isn't already
+            input_data = df.is_a?(Polars::LazyFrame) ? df : df.lazy
+
+            # Filter out records that already exist
+            new_records = input_data.filter(
+              Polars.col(primary_key).is_in(existing_keys).not_
+            )
+
+            # If we have new records, store them
+            if new_records.clone.select(Polars.length).collect[0, 0] > 0
+              @df = new_records
+              store_to_unique_file
+            end
+          end
+
+          def compact
+            files = self.files
+            return if files.empty?
+
+            clear_unique_id
+
+            # Mv existing compacted parquet to a temp file, so it doesn't conflict with write,
+            # but can still be queried
+            compacted_file = File.join(root_dir, "compacted.parquet")
+            if File.exist?(compacted_file)
+              tmp_file = File.join(root_dir, "compacted.orig.parquet")
+              FileUtils.mv(compacted_file, tmp_file)
+            end
+            files = self.files
+
+            compacted_file.tap do |target_file|
+              compacted_data = query(lazy: true).sort(primary_key)
+
+              safe_write(compacted_data, target_file)
+              FileUtils.rm(files)
+              clear_unique_id
+            end
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,139 @@
+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
+
+          def unlock!
+            clear_all_keys
+          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 acquire_lock(key, &block)
+            Support::Lockable.with_lock("#{key}:lock", wait_timeout: 2, &block)
+          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
+            list_keys.each { |key| unlock_file(key) }
+          end
+
+          def unlock_file(key)
+            acquire_lock(key) do |suo|
+              suo.client.del(key)
+            end
+          end
+
+          def clear_unique_id(subdir: nil)
+            key = unique_id_key(subdir: subdir)
+            acquire_lock(key) 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")
+
+            acquire_lock(keylist) do |suo|
+              suo.client.sadd(keylist, key)
+            end
+          end
+
+          def list_keys
+            keylist = unique_id_key(subdir: "keylist")
+
+            acquire_lock(keylist) do |suo|
+              if suo.client.type(keylist) == "set"
+                suo.client.smembers(keylist)
+              else
+                suo.client.del(keylist)
+                []
+              end
+            end
+          end
+
+          def key_exists?(key)
+            keylist = unique_id_key(subdir: "keylist")
+
+            acquire_lock(keylist) do |suo|
+              suo.client.sismember(keylist, key)
+            end
+          end
+
+          def unique_id(subdir: nil)
+            key = unique_id_key(subdir: subdir)
+            add_key(key)
+
+            acquire_lock(key) 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,80 @@
+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,
+                      :primary_key, :options, :append_only, :named
+
+        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 unlock!
+          adapter_class.new(options).unlock!
+        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,140 @@
+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
+
+      # Transform CSV files into Parquet files, of all the same datatype.
+      # Learn datatypes of columns and store schema.
+      def normalize
+        Normalizer.normalize(root_dir)
+      end
+
+      def data
+        query
+      end
+
+      def unlock!
+        writer.unlock!
+      end
+
+      def compact
+        writer.compact
+      end
+
+      def store(df, *args)
+        writer.store(df, *args)
+      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