easy_ml 0.2.0.pre.rc76 → 0.2.0.pre.rc77

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 1eebc157e0f33c3da40ef2b1bdb7cc0ed1c2b6f73615cdf26a6898cb60e60d2d
-  data.tar.gz: a12b441fe0736f251de773574858316346ba19c5b3784d73f3db200af0e619e4
+  metadata.gz: ee980703e3a768458e43d54a878bfa712d4f026967f3ccd8fa5bb2d1df50304c
+  data.tar.gz: eb5eb31b580e9112886527f416d4f360ffe9b0ee73d9e2e7dd70d9a48528ea09
 SHA512:
-  metadata.gz: 4aabb816a9d02a6f2bd870cde3db3eaaf00a314cf5e0d50a11bf707534b9d93eddee648d62304f48976916ea9d5942269dbeded81d49df23199ffcc13d6ae0eb
-  data.tar.gz: 284973f49424ac622ceb3e44071e88336ea316154dee788b0e7c865441eeb01939192289deea84283b691bf8f5a3b79f708d3d62ab9fcec3d596f67ff4c093a9
+  metadata.gz: ddc8a0005b22caf186c13790e9209d7b843181d62be8e70710e33bd6e244e3f31e6cc02efc4effa82d9f38674ff8ba2d8abdc1e43db8ba339d712f0e12d10ec4
+  data.tar.gz: 52193a2c0da5c0aca86bb627afff7efd9cbdfcf80ba2db8ae4e9baec061b716244e14c49f30451f52a221700b2b1a160a71935cb76cd77732a4e8fdcf25bd3a1

data/app/models/easy_ml/column/imputers/ordinal_encoder.rb

@@ -50,11 +50,7 @@ module EasyML
         end
 
         def cast_encoder(encoder)
-          begin
-            encoder.transform_keys { |k| column.cast(k) }
-          rescue => e
-            binding.pry
-          end
+          encoder.transform_keys { |k| column.cast(k) }
         end
 
         def cast_decoder(decoder)

data/app/models/easy_ml/dataset/refresh_reasons.rb

@@ -0,0 +1,12 @@
+module EasyML
+  class Dataset
+    class RefreshReasons < EasyML::Reasons
+      add_reason "Not split", -> { not_split? }
+      add_reason "Refreshed at is nil", -> { refreshed_at.nil? }
+      add_reason "Columns need refresh", -> { columns_need_refresh? }
+      add_reason "Features need refresh", -> { features_need_fit? }
+      add_reason "Datasource needs refresh", -> { datasource_needs_refresh? }
+      add_reason "Datasource was refreshed", -> { datasource_was_refreshed? }
+    end
+  end
+end

data/app/models/easy_ml/dataset.rb

@@ -265,9 +265,7 @@ module EasyML
 
     def refresh!(async: false)
       refreshing do
-        puts "Prepare..."
         prepare!
-        puts "Fit features..."
         fit_features!(async: async)
       end
     end
@@ -276,9 +274,7 @@ module EasyML
       return refresh_async if async
 
       refreshing do
-        puts "prepare.."
         prepare
-        puts "fit features..."
         fit_features(async: async)
       end
     end
@@ -299,7 +295,6 @@ module EasyML
     measure_method_timing :fit_features
 
     def after_fit_features
-      puts "after fit features..."
       unlock!
       reload
       return if failed?
@@ -338,45 +333,12 @@ module EasyML
     #
     # So yes this is an annoying way to structure a method, but it's helpful for performance
     #
-    def refresh_reasons(exclude: [])
-      {
-        not_split: {
-          name: "Not split",
-          check: -> { not_split? },
-        },
-        refreshed_at_is_nil: {
-          name: "Refreshed at is nil",
-          check: -> { refreshed_at.nil? },
-        },
-        columns_need_refresh: {
-          name: "Columns need refresh",
-          check: -> { columns_need_refresh? },
-        },
-        features_need_fit: {
-          name: "Features need refresh",
-          check: -> { features_need_fit? },
-        },
-        datasource_needs_refresh: {
-          name: "Datasource needs refresh",
-          check: -> { datasource_needs_refresh? },
-        },
-        refreshed_datasource: {
-          name: "Refreshed datasource",
-          check: -> { refreshed_datasource? },
-        },
-        datasource_was_refreshed: {
-          name: "Datasource was refreshed",
-          check: -> { datasource_was_refreshed? },
-        },
-      }.except(*exclude).select do |k, config|
-        config[:check].call
-      end.map do |k, config|
-        config[:name]
-      end
+    def refresh_reasons(except: [])
+      RefreshReasons.new(self).check(except: except)
     end
 
-    def needs_refresh?(exclude: [])
-      refresh_reasons(exclude: exclude).any?
+    def needs_refresh?(except: [])
+      refresh_reasons(except: except).any?
     end
 
     def processed?
@@ -518,19 +480,12 @@ module EasyML
     end
 
     def normalize(df = nil, split_ys: false, inference: false, all_columns: false, features: self.features)
-      puts "Apply missing features..."
       df = apply_missing_columns(df, inference: inference)
-      puts "Transform columns..."
       df = columns.transform(df, inference: inference)
-      puts "Apply features..."
       df = apply_features(df, features)
-      puts "Transform columns..."
       df = columns.transform(df, inference: inference, computed: true)
-      puts "Apply column mask..."
       df = apply_column_mask(df, inference: inference) unless all_columns
-      puts "Drop nulls..."
       df = drop_nulls(df) unless inference
-      puts "Split features and targets..."
       df, = processed.split_features_targets(df, true, target) if split_ys
       df
     end

data/app/models/easy_ml/feature.rb

@@ -190,31 +190,21 @@ 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),
+        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,
@@ -231,6 +221,8 @@ module EasyML
       jobs = ordered_features.map(&:build_batches)
       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 +317,7 @@ module EasyML
         params = {
           select: select,
           filter: filter,
+          sort: primary_key,
         }.compact
       else
         params = {}
@@ -438,24 +431,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
+      EasyML::FeatureStore.new(self)
     end
 
-    def query(**kwargs)
-      feature_store.query(**kwargs)
-    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 +445,7 @@ module EasyML
     def after_fit
       update_sha
 
+      feature_store.compact
       updates = {
         fit_at: Time.current,
         needs_fit: false,

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/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"

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