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

data/app/models/easy_ml/datasource.rb

@@ -22,7 +22,6 @@ module EasyML
     DATASOURCE_OPTIONS = {
       "s3" => "EasyML::Datasources::S3Datasource",
       "file" => "EasyML::Datasources::FileDatasource",
-      "polars" => "EasyML::Datasources::PolarsDatasource",
     }
     DATASOURCE_TYPES = [
       {
@@ -35,11 +34,6 @@ module EasyML
         label: "Local Files",
         description: "Connect to data stored in local files",
       },
-      {
-        value: "polars",
-        label: "Polars DataFrame",
-        description: "In-memory dataframe storage using Polars",
-      },
     ].freeze
     DATASOURCE_NAMES = DATASOURCE_OPTIONS.keys.freeze
     DATASOURCE_CONSTANTS = DATASOURCE_OPTIONS.values.map(&:constantize)

data/app/models/easy_ml/feature.rb

@@ -88,6 +88,7 @@ module EasyML
     before_save :update_sha
     after_find :update_from_feature_class
     before_save :update_from_feature_class
+    before_destroy :wipe
 
     def feature_klass
       feature_class.constantize
@@ -190,34 +191,23 @@ module EasyML
       reader = dataset.raw
 
       if adapter.respond_to?(:batch)
-        array = adapter.batch(reader, self)
-        min_id = array.min
-        max_id = array.max
+        series = adapter.batch(reader, self)
+        primary_key = series.name
       else
-        # Get all primary keys
-        begin
-          unless primary_key.present?
-            raise "Couldn't find primary key for feature #{feature_class}, check your feature class"
-          end
-          df = reader.query(select: primary_key)
-        rescue => e
-          raise "Couldn't find primary key #{primary_key.first} for feature #{feature_class}: #{e.message}"
-        end
-        return [] if df.nil?
-
-        min_id = df[primary_key.first].min
-        max_id = df[primary_key.last].max
+        primary_key = self.primary_key
       end
 
-      (min_id..max_id).step(batch_size).map.with_index do |batch_start, idx|
-        batch_end = [batch_start + batch_size, max_id + 1].min - 1
+      EasyML::Data::Partition::Boundaries.new(
+        reader.data(lazy: true, all_columns: true),
+        primary_key,
+        batch_size
+      ).to_a.map.with_index do |partition, idx|
         {
           feature_id: id,
-          batch_start: batch_start,
-          batch_end: batch_end,
+          batch_start: partition[:partition_start],
+          batch_end: partition[:partition_end],
           batch_number: feature_position,
           subbatch_number: idx,
-          parent_batch_id: Random.uuid,
         }
       end
     end
@@ -228,9 +218,16 @@ module EasyML
 
     def fit(features: [self], async: false)
       ordered_features = features.sort_by(&:feature_position)
-      jobs = ordered_features.map(&:build_batches)
+      parent_batch_id = Random.uuid
+      jobs = ordered_features.map do |feature|
+        feature.build_batches.map do |batch_args|
+          batch_args.merge(parent_batch_id: parent_batch_id)
+        end
+      end
       job_count = jobs.dup.flatten.size
 
+      ordered_features.each(&:wipe)
+
       # This is very important! For whatever reason, Resque BatchJob does not properly
       # handle batch finished callbacks for batch size = 1
       if async && job_count > 1
@@ -325,6 +322,7 @@ module EasyML
         params = {
           select: select,
           filter: filter,
+          sort: primary_key,
         }.compact
       else
         params = {}
@@ -438,24 +436,10 @@ module EasyML
     end
 
     def feature_store
-      @feature_store ||= EasyML::FeatureStore.new(self)
-    end
-
-    def upload_remote_files
-      feature_store.upload_remote_files
-    end
-
-    def files
-      feature_store.list_partitions
-    end
-
-    def query(**kwargs)
-      feature_store.query(**kwargs)
+      EasyML::FeatureStore.new(self)
     end
 
-    def store(df)
-      feature_store.store(df)
-    end
+    delegate :files, :query, :store, :compact, to: :feature_store
 
     def batch_size
       read_attribute(:batch_size) ||
@@ -466,6 +450,7 @@ module EasyML
     def after_fit
       update_sha
 
+      feature_store.compact
       updates = {
         fit_at: Time.current,
         needs_fit: false,
@@ -474,6 +459,10 @@ module EasyML
       update!(updates)
     end
 
+    def unlock!
+      feature_store.unlock!
+    end
+
     UNCONFIGURABLE_COLUMNS = %w(
       id
       dataset_id

data/app/models/easy_ml/model.rb

@@ -45,7 +45,7 @@ module EasyML
     MODEL_NAMES = MODEL_OPTIONS.keys.freeze
     MODEL_CONSTANTS = MODEL_OPTIONS.values.map(&:constantize)
 
-    add_configuration_attributes :task, :objective, :hyperparameters, :callbacks, :metrics
+    add_configuration_attributes :task, :objective, :hyperparameters, :callbacks, :metrics, :weights_column
     MODEL_CONSTANTS.flat_map(&:configuration_attributes).each do |attribute|
       add_configuration_attributes attribute
     end
@@ -179,6 +179,8 @@ module EasyML
     end
 
     def actually_train(&progress_block)
+      raise untrainable_error unless trainable?
+
       lock_model do
         run = pending_run
         run.wrap_training do
@@ -258,7 +260,7 @@ module EasyML
 
     def formatted_version
       return nil unless version
-      Time.strptime(version, "%Y%m%d%H%M%S").strftime("%B %-d, %Y at %-l:%M %p")
+      UTC.parse(version).in_time_zone(EasyML::Configuration.timezone).strftime("%B %-d, %Y at %-l:%M %p")
     end
 
     def last_run_at
@@ -277,6 +279,22 @@ module EasyML
     alias_method :latest_version, :inference_version
     alias_method :deployed, :inference_version
 
+    def trainable?
+      adapter.trainable?
+    end
+
+    def untrainable_columns
+      adapter.untrainable_columns
+    end
+
+    def untrainable_error
+      %Q(
+        Cannot train dataset containing null values!
+        Apply preprocessing to the following columns:
+        #{untrainable_columns.join(", ")}
+      )
+    end
+
     def predict(xs)
       load_model!
       unless xs.is_a?(XGBoost::DMatrix)

data/app/models/easy_ml/models/xgboost/evals_callback.rb

@@ -36,7 +36,7 @@ module EasyML
           if tuner.present?
             [tuner.x_valid, tuner.y_valid]
           else
-            model.dataset.valid(split_ys: true)
+            model.dataset.valid(split_ys: true, lazy: true)
           end
         end
 
@@ -47,7 +47,8 @@ module EasyML
           if epoch % log_frequency == 0
             model.adapter.external_model = booster
             x_valid, y_valid = valid_dataset
-            @preprocessed ||= model.preprocess(x_valid)
+            x_valid = x_valid.select(model.dataset.col_order(inference: true))
+            @preprocessed ||= model.preprocess(x_valid, y_valid)
             y_pred = model.predict(@preprocessed)
             dataset = model.dataset.valid(all_columns: true)
 

data/app/models/easy_ml/models/xgboost.rb

@@ -421,11 +421,11 @@ module EasyML
       def prepare_data
         if @d_train.nil?
           col_order = dataset.col_order
-          x_sample, y_sample = dataset.train(split_ys: true, limit: 5, select: col_order)
+          x_sample, y_sample = dataset.train(split_ys: true, limit: 5, select: col_order, lazy: true)
           preprocess(x_sample, y_sample) # Ensure we fail fast if the dataset is misconfigured
-          x_train, y_train = dataset.train(split_ys: true, select: col_order)
-          x_valid, y_valid = dataset.valid(split_ys: true, select: col_order)
-          x_test, y_test = dataset.test(split_ys: true, select: col_order)
+          x_train, y_train = dataset.train(split_ys: true, select: col_order, lazy: true)
+          x_valid, y_valid = dataset.valid(split_ys: true, select: col_order, lazy: true)
+          x_test, y_test = dataset.test(split_ys: true, select: col_order, lazy: true)
           @d_train = preprocess(x_train, y_train)
           @d_valid = preprocess(x_valid, y_valid)
           @d_test = preprocess(x_test, y_test)
@@ -434,21 +434,60 @@ module EasyML
         [@d_train, @d_valid, @d_test]
       end
 
+      def trainable?
+        untrainable_columns.empty?
+      end
+
+      def untrainable_columns
+        df = model.dataset.processed.data(lazy: true)
+
+        columns = df.columns
+        selects = columns.map do |col|
+          Polars.col(col).null_count.alias(col)
+        end
+        null_info = df.select(selects).collect
+        null_info.to_hashes.first.compact
+        col_list = null_info.to_hashes.first.transform_values { |v| v > 0 ? v : nil }.compact.keys
+
+        model.dataset.regular_columns(col_list)
+      end
+
       def preprocess(xs, ys = nil)
         return xs if xs.is_a?(::XGBoost::DMatrix)
+        weights_col = model.weights_column || nil
+
+        if weights_col == model.dataset.target
+          raise ArgumentError, "Weight column cannot be the target column"
+        end
+
+        # Extract feature columns (all columns except label and weight)
+        feature_cols = xs.columns
+        feature_cols -= [weights_col] if weights_col
+        lazy = xs.is_a?(Polars::LazyFrame)
+
+        # Get features, labels and weights
+        features = lazy ? xs.select(feature_cols).collect.to_numo : xs.select(feature_cols).to_numo
+        weights = weights_col ? (lazy ? xs.select(weights_col).collect.to_numo : xs.select(weights_col).to_numo) : nil
+        weights = weights.flatten if weights
+        if ys.present?
+          ys = ys.is_a?(Array) ? Polars::Series.new(ys) : ys
+          labels = lazy ? ys.collect.to_numo.flatten : ys.to_numo.flatten
+        else
+          labels = nil
+        end
+
+        kwargs = {
+          label: labels,
+          weight: weights,
+        }.compact
 
-        orig_xs = xs.dup
-        column_names = xs.columns
-        xs = _preprocess(xs)
-        ys = ys.nil? ? nil : _preprocess(ys).flatten
-        kwargs = { label: ys }.compact
         begin
-          ::XGBoost::DMatrix.new(xs, **kwargs).tap do |dmat|
-            dmat.feature_names = column_names
+          ::XGBoost::DMatrix.new(features, **kwargs).tap do |dmatrix|
+            dmatrix.feature_names = feature_cols
           end
         rescue StandardError => e
-          problematic_columns = orig_xs.schema.select { |k, v| [Polars::Categorical, Polars::String].include?(v) }
-          problematic_xs = orig_xs.select(problematic_columns.keys)
+          problematic_columns = xs.schema.select { |k, v| [Polars::Categorical, Polars::String].include?(v) }
+          problematic_xs = lazy ? xs.lazy.select(problematic_columns.keys).collect : xs.select(problematic_columns.keys)
           raise %(
             Error building data for XGBoost.
             Apply preprocessing to columns 
@@ -501,29 +540,6 @@ module EasyML
         cb_container.after_iteration(@booster, current_iteration, d_train, evals)
       end
 
-      def _preprocess(df)
-        return df if df.is_a?(Array)
-
-        df.to_a.map do |row|
-          row.values.map do |value|
-            case value
-            when Time
-              value.to_i # Convert Time to Unix timestamp
-            when Date
-              value.to_time.to_i # Convert Date to Unix timestamp
-            when String
-              value
-            when TrueClass, FalseClass
-              value ? 1.0 : 0.0 # Convert booleans to 1.0 and 0.0
-            when Integer
-              value
-            else
-              value.to_f # Ensure everything else is converted to a float
-            end
-          end
-        end
-      end
-
       def initialize_model
         @xgboost_model = model_class.new(n_estimators: @hyperparameters.to_h.dig(:n_estimators))
         if block_given?

data/app/models/easy_ml/retraining_run.rb

@@ -150,7 +150,7 @@ module EasyML
 
       training_model.dataset.refresh
       evaluator = retraining_job.evaluator.symbolize_keys
-      x_test, y_test = training_model.dataset.test(split_ys: true)
+      x_test, y_test = training_model.dataset.test(split_ys: true, all_columns: true)
       y_pred = training_model.predict(x_test)
 
       metric = evaluator[:metric].to_sym

data/app/serializers/easy_ml/dataset_serializer.rb

@@ -84,7 +84,7 @@ module EasyML
     end
 
     attribute :needs_refresh do |dataset|
-      dataset.needs_refresh?(exclude: [:datasource_needs_refresh])
+      dataset.needs_refresh?(except: [:datasource_needs_refresh])
     end
 
     attribute :stacktrace do |object|

data/app/serializers/easy_ml/model_serializer.rb

@@ -27,6 +27,7 @@ module EasyML
                :model_type,
                :task,
                :objective,
+               :weights_column,
                :metrics,
                :dataset_id,
                :status,

data/lib/easy_ml/core/tuner.rb

@@ -8,7 +8,7 @@ module EasyML
                     :metrics, :objective, :n_trials, :direction, :evaluator,
                     :study, :results, :adapter, :tune_started_at, :x_valid, :y_valid,
                     :project_name, :job, :current_run, :trial_enumerator, :progress_block,
-                    :tuner_job, :dataset
+                    :tuner_job, :dataset, :x_normalized
 
       def initialize(options = {})
         @model = options[:model]
@@ -73,9 +73,12 @@ module EasyML
         model.task = task
 
         model.dataset.refresh if model.dataset.needs_refresh?
-        x_valid, y_valid = model.dataset.valid(split_ys: true, select: model.dataset.col_order)
+        x_valid, y_valid = model.dataset.valid(split_ys: true, all_columns: true)
+        x_normalized = model.dataset.normalize(x_valid, inference: true)
+        x_normalized = model.preprocess(x_normalized)
         self.x_valid = x_valid
         self.y_valid = y_valid
+        self.x_normalized = x_normalized
         self.dataset = model.dataset.valid(all_columns: true)
         adapter.tune_started_at = tune_started_at
         adapter.x_valid = x_valid
@@ -99,7 +102,7 @@ module EasyML
             @study.tell(@current_trial, result)
           rescue StandardError => e
             puts EasyML::Event.easy_ml_context(e.backtrace)
-            @tuner_run.update!(status: :failed, hyperparameters: {})
+            @tuner_run.update!(status: :failed, hyperparameters: model.hyperparameters.to_h)
             puts "Optuna failed with: #{e.message}"
             raise e
           end
@@ -138,7 +141,7 @@ module EasyML
           end
         end
 
-        y_pred = model.predict(x_valid)
+        y_pred = model.predict(x_normalized)
         model.metrics = metrics
         metrics = model.evaluate(y_pred: y_pred, y_true: y_valid, x_true: x_valid, dataset: dataset)
         metric = metrics.symbolize_keys.dig(model.evaluator[:metric].to_sym)

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

@@ -0,0 +1,80 @@
+module EasyML
+  module Data
+    class DatasetManager
+      class Reader
+        class Base
+          DEFAULTS = {
+            drop_cols: [],
+            filter: nil,
+            limit: nil,
+            select: nil,
+            unique: nil,
+            sort: nil,
+            descending: false,
+            batch_size: nil,
+            batch_start: nil,
+            batch_key: nil,
+            lazy: false,
+          }
+
+          DEFAULTS.each do |k, _|
+            attr_accessor k
+          end
+          attr_accessor :block, :options, :input
+          attr_accessor :options
+
+          def initialize(options, &block)
+            options = apply_defaults(options)
+            @block = block
+            @options = options
+          end
+
+          def query
+            raise "Not implemented"
+          end
+
+          private
+
+          def apply_defaults(kwargs)
+            options = kwargs.dup
+
+            DEFAULTS.each do |k, default|
+              unless options.key?(k)
+                options[k] = default
+              end
+            end
+
+            options.each do |k, v|
+              send("#{k}=", v)
+            end
+
+            options
+          end
+
+          def query_dataframes(df, schema)
+            num_rows = df.is_a?(Polars::LazyFrame) ? df.select(Polars.length).collect[0, 0] : df.shape[0]
+            return df if num_rows == 0
+
+            # Apply the predicate filter if given
+            df = df.filter(filter) if filter
+            # Apply select columns if provided
+            df = df.select(select) if select.present?
+            df = df.unique if unique
+
+            # Apply sorting if provided
+            df = df.sort(sort, reverse: descending) if sort
+
+            # Apply drop columns
+            drop_cols = self.drop_cols
+            drop_cols &= schema.keys
+            df = df.drop(drop_cols) unless drop_cols.empty?
+
+            # Collect the DataFrame (execute the lazy operations)
+            df = df.limit(limit) if limit
+            lazy ? df : df.collect
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,106 @@
+module EasyML
+  module Data
+    class DatasetManager
+      class Reader
+        class Batch < File
+          def query
+            return batch_enumerator unless block.present?
+            return process_batches
+          end
+
+          private
+
+          def batch_enumerator
+            Enumerator.new do |yielder|
+              process_batches do |batch|
+                yielder << batch
+              end
+            end
+          end
+
+          def process_batches(&b)
+            raise "When using batch_size, sort must match primary key (#{batch_key})" if sort.present? && batch_key != sort
+            block = b || self.block
+
+            sort = batch_key
+
+            current_start = get_batch_start
+            final_value = get_final_value
+
+            while current_start < final_value
+              filter = Polars.col(sort) >= current_start
+              batch = query_files(filter: filter, limit: batch_size, lazy: true, sort: sort, descending: descending)
+              block.yield(batch)
+              current_start = File.new(input: input, lazy: true)
+                                  .query
+                                  .filter(filter)
+                                  .sort(sort, reverse: descending)
+                                  .limit(batch_size + 1)
+                                  .sort(sort, reverse: !descending)
+                                  .limit(1)
+                                  .select(sort)
+                                  .collect
+                                  .to_a.first&.dig(sort) || final_value
+            end
+          end
+
+          def query_files(overrides = {})
+            query = options.deep_dup.merge!(overrides).except(:batch_size, :batch_start, :batch_key)
+            File.new(query).query
+          end
+
+          def get_batch_start
+            if batch_start.present?
+              batch_start
+            else
+              get_sorted_batch_keys(descending)
+            end
+          end
+
+          def get_final_value
+            get_sorted_batch_keys(!descending)
+          end
+
+          def get_sorted_batch_keys(descending, filter: nil)
+            query = query_files(lazy: true)
+            query = query.filter(filter) if filter
+            query.sort(batch_key, reverse: descending).limit(1).select(batch_key).collect.to_a.last.dig(batch_key)
+          end
+
+          def batch_key
+            return @batch_key if @batch_key
+
+            lazy_df = lazy_frames([files.first]).first
+            if select
+              # Lazily filter only the selected columns
+              lazy_df = lazy_df.select(select)
+
+              # Lazily compute the unique count for each column and compare with total row count
+              primary_keys = select.select do |col|
+                lazy_df.select(col).unique.collect.height == lazy_df.collect.height
+              end
+            else
+              primary_keys = lazy_df.collect.columns.select do |col|
+                # Lazily count unique values and compare with the total row count
+                lazy_df.select(col).unique.collect.height == lazy_df.collect.height
+              end
+            end
+
+            if primary_keys.count > 1
+              key = primary_keys.detect { |key| key.underscore.split("_").any? { |k| k.match?(/id/) } }
+              if key
+                primary_keys = [key]
+              end
+            end
+
+            if primary_keys.count != 1
+              raise "Unable to determine primary key for dataset"
+            end
+
+            @batch_key = primary_keys.first
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,23 @@
+
+module EasyML
+  module Data
+    class DatasetManager
+      class Reader
+        class DataFrame < File
+          def query
+            return query_dataframes(lazy_frames, schema)
+          end
+
+          def schema
+            input.schema
+          end
+
+        private
+          def lazy_frames
+            input.lazy
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,75 @@
+module EasyML
+  module Data
+    class DatasetManager
+      class Reader
+        class File < Base
+          attr_accessor :file_filter
+
+          def initialize(options = {})
+            super
+            @file_filter = options.dig(:file_filter) || ->(file) { true }
+          end
+
+          def query
+            return query_dataframes(dataframe, schema) unless batch_size.present?
+            return Batch.new(options, &block).query
+          end
+
+          def schema
+            @schema ||= files.any? ? Polars.read_parquet_schema(files.first) : nil
+          end
+
+          def files
+            filter_files do
+              if is_file?
+                @files ||= [input]
+              elsif is_dir?
+                @files ||= Dir.glob(::File.join(root_dir, "**/*.{parquet}"))
+              else
+                @files ||= []
+              end
+            end
+          end
+
+          private
+
+          def filter_files(&block)
+            yield
+            @files = @files.select(&file_filter)
+          end
+
+          def is_dir?
+            path.directory?
+          end
+
+          def is_file?
+            path.file?
+          end
+
+          def root_dir
+            path if is_dir?
+          end
+
+          def path
+            @path ||= input.is_a?(Pathname) ? input : Pathname.new(input)
+          end
+
+          def dataframe
+            @dataframe = lazy_frames.any? ? Polars.concat(lazy_frames) : Polars::LazyFrame.new
+          end
+
+          def lazy_frames(files = nil)
+            return @lazy_frames if @lazy_frames
+
+            files ||= self.files
+            @lazy_frames = files.map do |file|
+              Polars.scan_parquet(file)
+            end
+          end
+        end
+      end
+    end
+  end
+end
+
+require_relative "batch"