easy_ml 0.1.4 → 0.2.0.pre.rc1

data/lib/easy_ml/feature_store.rb

@@ -0,0 +1,227 @@
+module EasyML
+  class FeatureStore
+    attr_reader :feature
+
+    def initialize(feature)
+      @feature = feature
+    end
+
+    def store(df)
+      primary_key = feature.primary_key&.first
+      return store_without_partitioning(df) unless df.columns.include?(primary_key)
+      return store_without_partitioning(df) unless primary_key
+
+      min_key = df[primary_key].min
+      max_key = df[primary_key].max
+      batch_size = feature.batch_size || 10_000
+
+      # Try to parse as integers if they're strings
+      begin
+        min_key = Integer(min_key) if min_key.is_a?(String)
+        max_key = Integer(max_key) if max_key.is_a?(String)
+      rescue ArgumentError
+        return store_without_partitioning(df)
+      end
+
+      # Only partition if we have integer keys where we can predict boundaries
+      return store_without_partitioning(df) unless min_key.is_a?(Integer) && max_key.is_a?(Integer)
+
+      partitions = compute_partition_boundaries(min_key, max_key, batch_size)
+      partitions.each do |partition_start|
+        partition_end = partition_start + batch_size - 1
+        partition_df = df.filter(
+          (Polars.col(primary_key) >= partition_start) &
+          (Polars.col(primary_key) <= partition_end)
+        )
+
+        next if partition_df.height == 0
+
+        store_partition(partition_df, primary_key, partition_start)
+      end
+    end
+
+    def query(filter: nil)
+      query_all_partitions(filter)
+    end
+
+    def empty?
+      list_partitions.empty?
+    end
+
+    def list_partitions
+      Dir.glob(File.join(feature_dir, "feature*.parquet")).sort
+    end
+
+    def wipe
+      FileUtils.rm_rf(feature_dir)
+    end
+
+    def upload_remote_files
+      synced_directory.upload
+    end
+
+    def download
+      synced_directory&.download
+    end
+
+    def cp(old_version, new_version)
+      old_dir = feature_dir_for_version(old_version)
+      new_dir = feature_dir_for_version(new_version)
+
+      return if old_dir.nil? || !Dir.exist?(old_dir)
+
+      FileUtils.mkdir_p(new_dir)
+      files_to_cp = Dir.glob(Pathname.new(old_dir).join("**/*")).select { |f| File.file?(f) }
+
+      files_to_cp.each do |file|
+        target_file = file.gsub(old_version.to_s, new_version.to_s)
+        FileUtils.mkdir_p(File.dirname(target_file))
+        FileUtils.cp(file, target_file)
+      end
+    end
+
+    private
+
+    def store_without_partitioning(df)
+      lock_file do
+        path = feature_path
+        FileUtils.mkdir_p(File.dirname(path))
+        df.write_parquet(path)
+      end
+    end
+
+    def store_partition(partition_df, primary_key, partition_start)
+      lock_partition(partition_start) do
+        path = partition_path(partition_start)
+        FileUtils.mkdir_p(File.dirname(path))
+
+        if File.exist?(path)
+          reader = EasyML::Data::PolarsReader.new
+          existing_df = reader.query([path])
+          preserved_records = existing_df.filter(
+            Polars.col(primary_key).is_in(partition_df[primary_key]).is_not
+          )
+          partition_df = Polars.concat([preserved_records, partition_df], how: "vertical")
+        end
+
+        partition_df.write_parquet(path)
+      end
+    end
+
+    def query_partitions(filter)
+      primary_key_values = filter.extract_primary_key_values
+      batch_size = feature.batch_size || 10_000
+
+      partition_files = primary_key_values.map do |key|
+        partition_start = (key / batch_size.to_f).floor * batch_size
+        partition_path(partition_start)
+      end.uniq.select { |path| File.exist?(path) }
+
+      return Polars::DataFrame.new if partition_files.empty?
+
+      reader = EasyML::Data::PolarsReader.new
+      reader.query(partition_files, filter: filter)
+    end
+
+    def query_all_partitions(filter)
+      reader = EasyML::Data::PolarsReader.new
+      pattern = File.join(feature_dir, "feature*.parquet")
+      files = Dir.glob(pattern)
+
+      return Polars::DataFrame.new if files.empty?
+
+      reader.query(files, filter: filter)
+    end
+
+    def compute_partition_boundaries(min_key, max_key, batch_size)
+      start_partition = (min_key / batch_size.to_f).floor * batch_size
+      end_partition = (max_key / batch_size.to_f).floor * batch_size
+      (start_partition..end_partition).step(batch_size).to_a
+    end
+
+    def feature_dir_for_version(version)
+      File.join(
+        Rails.root,
+        "easy_ml/datasets",
+        feature.dataset.name.parameterize.gsub("-", "_"),
+        "features",
+        feature.name.parameterize.gsub("-", "_"),
+        version.to_s
+      )
+    end
+
+    def feature_dir
+      feature_dir_for_version(feature.version)
+    end
+
+    def feature_path
+      File.join(feature_dir, "feature.parquet")
+    end
+
+    def partition_path(partition_start)
+      File.join(feature_dir, "feature#{partition_start}.parquet")
+    end
+
+    def s3_prefix
+      File.join("datasets", feature_dir.split("datasets").last)
+    end
+
+    def synced_directory
+      return unless feature.dataset&.datasource.present?
+
+      datasource_config = feature.dataset.datasource.configuration || {}
+      @synced_dir ||= EasyML::Data::SyncedDirectory.new(
+        root_dir: feature_dir,
+        s3_bucket: datasource_config.dig("s3_bucket") || EasyML::Configuration.s3_bucket,
+        s3_prefix: s3_prefix,
+        s3_access_key_id: EasyML::Configuration.s3_access_key_id,
+        s3_secret_access_key: EasyML::Configuration.s3_secret_access_key,
+        polars_args: datasource_config.dig("polars_args"),
+        cache_for: 0,
+      )
+    end
+
+    def lock_partition(partition_start)
+      Support::Lockable.with_lock(partition_lock_key(partition_start), wait_timeout: 2, stale_timeout: 60) do |client|
+        begin
+          yield client if block_given?
+        ensure
+          unlock_partition(partition_start)
+        end
+      end
+    end
+
+    def lock_file
+      Support::Lockable.with_lock(file_lock_key, wait_timeout: 2, stale_timeout: 60) do |client|
+        begin
+          yield client if block_given?
+        ensure
+          unlock_file
+        end
+      end
+    end
+
+    def unlock_partition(partition_start)
+      Support::Lockable.unlock!(partition_lock_key(partition_start))
+    end
+
+    def unlock_file
+      Support::Lockable.unlock!(file_lock_key)
+    end
+
+    def unlock_all_partitions
+      list_partitions.each do |partition_path|
+        partition_start = partition_path.match(/feature(\d+)\.parquet/)[1].to_i
+        unlock_partition(partition_start)
+      end
+    end
+
+    def partition_lock_key(partition_start)
+      "feature_store:#{feature.id}.partition.#{partition_start}"
+    end
+
+    def file_lock_key
+      "feature_store:#{feature.id}.file"
+    end
+  end
+end

data/lib/easy_ml/features.rb

@@ -0,0 +1,61 @@
+module EasyML::Features
+  def transform(df, feature)
+    raise NotImplementedError
+  end
+
+  def self.included(base)
+    base.extend(ClassMethods)
+    Registry.register(base)
+  end
+
+  module ClassMethods
+    def features
+      @features ||= []
+    end
+
+    def feature(**kwargs)
+      features << kwargs.merge!(feature_class: self.to_s)
+    end
+  end
+
+  def missing_any?(list1, list2)
+    (list1 - list2).any?
+  end
+
+  class Registry
+    class << self
+      def register(feature_class, namespace: nil)
+        namespace = namespace&.to_sym
+        registry[namespace] ||= {}
+        registry[namespace][feature_class] = feature_class
+      end
+
+      def list(namespace: nil)
+        require_files
+        namespace ? registry[namespace.to_sym] : registry
+      end
+
+      def require_files
+        Dir.glob(Rails.root.join("app/features/**/*.rb")).each { |f| require_dependency f }
+      end
+
+      def list_flat
+        (list.try(:values) || []).flat_map(&:values).flat_map(&:features)
+      end
+
+      def find(name)
+        list_flat.detect { |feature| feature[:name] == name || feature[:feature_class] == name }
+      end
+
+      def clear
+        @registry = {}
+      end
+
+      private
+
+      def registry
+        @registry ||= {}
+      end
+    end
+  end
+end

data/lib/easy_ml/initializers/inflections.rb

@@ -1,4 +1,18 @@
-ActiveSupport::Inflector.inflections(:en) do |inflect|
-  inflect.acronym "EasyML"
-  inflect.acronym "ML"
+module EasyML
+  module Initializers
+    module Inflections
+      def self.inflect
+        ActiveSupport::Inflector.inflections(:en) do |inflect|
+          inflect.acronym "EasyML"
+          inflect.acronym "ML"
+          inflect.acronym "STI"
+          inflect.acronym "XGBoost"
+          inflect.acronym "GBLinear"
+          inflect.acronym "GBTree"
+          inflect.acronym "EST"
+          inflect.acronym "UTC"
+        end
+      end
+    end
+  end
 end

data/lib/easy_ml/logging.rb

@@ -7,9 +7,9 @@ module EasyML
     module ClassMethods
       def log_method(method_name, message, verbose: false)
         original_method = instance_method(method_name)
-        define_method(method_name) do |*args, &block|
+        define_method(method_name) do |*args, **kwargs, &block|
           log_message(message, verbose: verbose)
-          result = original_method.bind(self).call(*args, &block)
+          result = original_method.bind(self).call(*args, **kwargs, &block)
           result
         end
       end

data/lib/easy_ml/predict.rb

@@ -0,0 +1,74 @@
+require "singleton"
+
+module EasyML
+  class Predict
+    include Singleton
+
+    attr_reader :models
+
+    def initialize
+      @models = {}
+    end
+
+    def self.predict(model_name, df)
+      if df.is_a?(Hash)
+        df = Polars::DataFrame.new(df)
+      end
+      raw_input = df.to_hashes&.first
+      df = instance.normalize(model_name, df)
+      preds = instance.predict(model_name, df)
+      current_version = instance.get_model(model_name)
+
+      EasyML::Prediction.create!(
+        model: current_version.model,
+        model_history: current_version,
+        prediction_type: current_version.model.task,
+        prediction_value: {
+          value: preds.first,
+        }.compact,
+        raw_input: raw_input,
+        normalized_input: df.to_hashes&.first,
+      )
+
+      preds
+    end
+
+    def self.train(model_name, tuner: nil, evaluator: nil)
+      instance.train(model_name, tuner: tuner, evaluator: evaluator)
+    end
+
+    def predict(model_name, df)
+      get_model(model_name).predict(df)
+    end
+
+    def normalize(model_name, df)
+      get_model(model_name).dataset.normalize(df, inference: true)
+    end
+
+    def get_model(model_name)
+      load_model(model_name)
+      models[model_name]
+    end
+
+    def reset
+      @models = {}
+    end
+
+    def self.reset
+      instance.reset
+    end
+
+    private
+
+    def load_model(model_name)
+      current_model = EasyML::Model.find_by!(name: model_name).inference_version
+
+      # Load new model if not loaded or different version
+      model_not_loaded = models[model_name].nil?
+      model_is_new_version = models[model_name]&.id != current_model&.id
+      return unless model_not_loaded || model_is_new_version
+
+      models[model_name] = current_model
+    end
+  end
+end

data/lib/easy_ml/railtie/generators/migration/migration_generator.rb

@@ -2,47 +2,203 @@ require "rails/generators"
 require "rails/generators/active_record/migration"
 
 module EasyML
-  module Generators
-    module Migration
-      class MigrationGenerator < Rails::Generators::Base
-        include Rails::Generators::Migration
-        namespace "easy_ml:migration"
-
-        # Set the source directory for templates
-        source_root File.expand_path("../../templates/migration", __dir__)
-
-        # Define the migration name
-        desc "Generates migrations for EasyMLModel, Dataset, and TunerRun"
-
-        # Specify the next migration number
-        def self.next_migration_number(dirname)
-          if ActiveRecord.version < Gem::Version.new("7")
-            Time.now.utc.strftime("%Y%m%d%H%M%S")
-          elsif ActiveRecord.timestamped_migrations
-            Time.now.utc.strftime("%Y%m%d%H%M%S")
-          else
-            format("%.3d", (current_migration_number(dirname) + 1))
+  module Railtie
+    module Generators
+      module Migration
+        class MigrationGenerator < Rails::Generators::Base
+          include Rails::Generators::Migration
+          namespace "easy_ml:migration"
+
+          # Set the source directory for templates
+          source_root File.expand_path("../../templates/migration", __dir__)
+
+          # Define the migration name
+          desc "Generates migrations for EasyMLModel"
+
+          # Specify the next migration number
+          def self.next_migration_number(dirname)
+            sleep(1)
+            if ActiveRecord.version < Gem::Version.new("7")
+              Time.now.utc.strftime("%Y%m%d%H%M%S")
+            elsif ActiveRecord.timestamped_migrations
+              Time.now.utc.strftime("%Y%m%d%H%M%S")
+            else
+              format("%.3d", (current_migration_number(dirname) + 1))
+            end
           end
-        end
 
-        # Generate the migration files using the templates
-        def create_migration_files
-          create_easy_ml_models_migration
-        end
+          # Generate the migration files using the templates
+          def create_migration_files
+            create_easy_ml_datasource_migration
+            create_easy_ml_datasets_migration
+            create_easy_ml_columns_migration
+            create_easy_ml_models_migration
+            create_easy_ml_model_files_migration
+            create_easy_ml_tuner_jobs_migration
+            create_easy_ml_retraining_jobs_migration
+            create_easy_ml_settings_migration
+            create_easy_ml_events_migration
+            create_easy_ml_features_migration
+            create_easy_ml_splitters_migration
+            create_easy_ml_splitter_histories_migration
+            create_easy_ml_deploys
+
+            create_easy_ml_datasource_histories_migration
+            create_easy_ml_dataset_histories_migration
+            create_easy_ml_column_histories_migration
+            create_easy_ml_model_histories_migration
+            create_easy_ml_model_file_histories_migration
+            create_easy_ml_feature_histories_migration
+            create_easy_ml_predictions_migration
+          end
 
-        private
+          private
 
-        # Generate the migration file for EasyMLModel using the template
-        def create_easy_ml_models_migration
-          migration_template(
-            "create_easy_ml_models.rb.tt",
-            "db/migrate/create_easy_ml_models.rb"
-          )
-        end
+          # Generate the migration file for EasyMLModel using the template
+          def create_easy_ml_models_migration
+            migration_template(
+              "create_easy_ml_models.rb.tt",
+              "db/migrate/create_easy_ml_models.rb"
+            )
+          end
+
+          def create_easy_ml_model_files_migration
+            migration_template(
+              "create_easy_ml_model_files.rb.tt",
+              "db/migrate/create_easy_ml_model_files.rb"
+            )
+          end
+
+          def create_easy_ml_datasource_migration
+            migration_template(
+              "create_easy_ml_datasources.rb.tt",
+              "db/migrate/create_easy_ml_datasources.rb"
+            )
+          end
+
+          def create_easy_ml_datasets_migration
+            migration_template(
+              "create_easy_ml_datasets.rb.tt",
+              "db/migrate/create_easy_ml_datasets.rb"
+            )
+          end
+
+          def create_easy_ml_tuner_jobs_migration
+            migration_template(
+              "create_easy_ml_tuner_jobs.rb.tt",
+              "db/migrate/create_easy_ml_tuner_jobs.rb"
+            )
+          end
+
+          def create_easy_ml_retraining_jobs_migration
+            migration_template(
+              "create_easy_ml_retraining_jobs.rb.tt",
+              "db/migrate/create_easy_ml_retraining_jobs.rb"
+            )
+          end
+
+          def create_easy_ml_settings_migration
+            migration_template(
+              "create_easy_ml_settings.rb.tt",
+              "db/migrate/create_easy_ml_settings.rb"
+            )
+          end
+
+          def create_easy_ml_events_migration
+            migration_template(
+              "create_easy_ml_events.rb.tt",
+              "db/migrate/create_easy_ml_events.rb"
+            )
+          end
+
+          def create_easy_ml_columns_migration
+            migration_template(
+              "create_easy_ml_columns.rb.tt",
+              "db/migrate/create_easy_ml_columns.rb"
+            )
+          end
+
+          def create_easy_ml_features_migration
+            migration_template(
+              "create_easy_ml_features.rb.tt",
+              "db/migrate/create_easy_ml_features.rb"
+            )
+          end
+
+          def create_easy_ml_splitters_migration
+            migration_template(
+              "create_easy_ml_splitters.rb.tt",
+              "db/migrate/create_easy_ml_splitters.rb"
+            )
+          end
+
+          def create_easy_ml_splitter_histories_migration
+            migration_template(
+              "create_easy_ml_splitter_histories.rb.tt",
+              "db/migrate/create_easy_ml_splitter_histories.rb"
+            )
+          end
+
+          def create_easy_ml_datasource_histories_migration
+            migration_template(
+              "create_easy_ml_datasource_histories.rb.tt",
+              "db/migrate/create_easy_ml_datasource_histories.rb"
+            )
+          end
+
+          def create_easy_ml_dataset_histories_migration
+            migration_template(
+              "create_easy_ml_dataset_histories.rb.tt",
+              "db/migrate/create_easy_ml_dataset_histories.rb"
+            )
+          end
+
+          def create_easy_ml_column_histories_migration
+            migration_template(
+              "create_easy_ml_column_histories.rb.tt",
+              "db/migrate/create_easy_ml_column_histories.rb"
+            )
+          end
+
+          def create_easy_ml_model_histories_migration
+            migration_template(
+              "create_easy_ml_model_histories.rb.tt",
+              "db/migrate/create_easy_ml_model_histories.rb"
+            )
+          end
+
+          def create_easy_ml_feature_histories_migration
+            migration_template(
+              "create_easy_ml_feature_histories.rb.tt",
+              "db/migrate/create_easy_ml_feature_histories.rb"
+            )
+          end
+
+          def create_easy_ml_model_file_histories_migration
+            migration_template(
+              "create_easy_ml_model_file_histories.rb.tt",
+              "db/migrate/create_easy_ml_model_file_histories.rb"
+            )
+          end
+
+          def create_easy_ml_deploys
+            migration_template(
+              "create_easy_ml_deploys.rb.tt",
+              "db/migrate/create_easy_ml_deploys.rb"
+            )
+          end
 
-        # Get the next migration number
-        def next_migration_number
-          self.class.next_migration_number(Rails.root.join("db/migrate"))
+          def create_easy_ml_predictions_migration
+            migration_template(
+              "create_easy_ml_predictions.rb.tt",
+              "db/migrate/create_easy_ml_predictions.rb"
+            )
+          end
+
+          # Get the next migration number
+          def next_migration_number
+            self.class.next_migration_number(Rails.root.join("db/migrate"))
+          end
         end
       end
     end

data/lib/easy_ml/railtie/templates/migration/create_easy_ml_column_histories.rb.tt

@@ -0,0 +1,9 @@
+require "historiographer/postgres_migration"
+
+class CreateEasyMLColumnHistories < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def change
+    create_table :easy_ml_column_histories do |t|
+      t.histories(foreign_key: :column_id)
+    end
+  end
+end 

data/lib/easy_ml/railtie/templates/migration/create_easy_ml_columns.rb.tt

@@ -0,0 +1,25 @@
+class CreateEasyMLColumns < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def change
+    create_table :easy_ml_columns do |t|
+      t.bigint :dataset_id, null: false
+      t.string :name, null: false
+      t.string :description
+      t.string :datatype # The symbol representation (e.g., 'float', 'integer')
+      t.string :polars_datatype # The full Polars class name (e.g., 'Polars::Float64')
+      t.boolean :is_target
+      t.boolean :hidden, default: false
+      t.boolean :drop_if_null, default: false
+      t.json :preprocessing_steps
+      t.json :sample_values # Store up to 3 sample values
+      t.json :statistics
+
+      t.timestamps
+
+      t.index [:dataset_id, :name], unique: true
+      t.index :datatype
+      t.index :hidden
+      t.index :drop_if_null
+      t.index :is_target
+    end
+  end
+end 

data/lib/easy_ml/railtie/templates/migration/create_easy_ml_dataset_histories.rb.tt

@@ -0,0 +1,9 @@
+require "historiographer/postgres_migration"
+
+class CreateEasyMLDatasetHistories < ActiveRecord::Migration[<%= ActiveRecord::Migration.current_version %>]
+  def change
+    create_table :easy_ml_dataset_histories do |t|
+      t.histories(foreign_key: :dataset_id)
+    end
+  end
+end