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

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/polars_column.rb

@@ -124,11 +124,25 @@ module EasyML
       # @param series [Polars::Series] The string series to analyze
       # @return [Symbol] One of :datetime, :text, or :categorical
       def determine_string_type(series)
-        if EasyML::Data::DateConverter.maybe_convert_date(Polars::DataFrame.new({ temp: series }),
-                                                          :temp)[:temp].dtype.is_a?(Polars::Datetime)
-          :datetime
-        else
-          categorical_or_text?(series)
+        # Try to parse as numeric first
+        begin
+          # Try integer first
+          series.cast(Polars::Int64)
+          return :numeric
+        rescue StandardError
+          begin
+            # Try float if integer fails
+            series.cast(Polars::Float64)
+            return :numeric
+          rescue StandardError
+            # If not numeric, check for datetime or categorical
+            if EasyML::Data::DateConverter.maybe_convert_date(Polars::DataFrame.new({ temp: series }),
+                                                              :temp)[:temp].dtype.is_a?(Polars::Datetime)
+              :datetime
+            else
+              categorical_or_text?(series)
+            end
+          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

data/lib/easy_ml/engine.rb

@@ -55,7 +55,7 @@ module EasyML
       Polars.enable_string_cache
     end
 
-    if %w[db:migrate db:migrate:status db:setup db:drop assets:precompile].include?(ARGV.first)
+    if %w[db:create db:migrate db:migrate:status db:setup db:drop assets:precompile].include?(ARGV.first)
       config.eager_load_paths = config.eager_load_paths.without(config.eager_load_paths.map(&:to_s).grep(/easy_ml/).map { |p| Pathname.new(p) })
     else
       config.after_initialize do
@@ -77,6 +77,21 @@ module EasyML
       end
     end
 
+    unless %w[db:create db:migrate db:migrate:status db:setup db:drop assets:precompile].include?(ARGV.first)
+      initializer "easy_ml.configure_secrets" do
+        EasyML::Configuration.configure do |config|
+          raise "S3_ACCESS_KEY_ID is missing. Set ENV['S3_ACCESS_KEY_ID']" unless ENV["S3_ACCESS_KEY_ID"]
+          raise "S3_SECRET_ACCESS_KEY is missing. Set ENV['S3_SECRET_ACCESS_KEY']" unless ENV["S3_SECRET_ACCESS_KEY"]
+
+          config.s3_access_key_id = ENV["S3_ACCESS_KEY_ID"]
+          config.s3_secret_access_key = ENV["S3_SECRET_ACCESS_KEY"]
+          config.s3_region = ENV["S3_REGION"] ? ENV["S3_REGION"] : "us-east-1"
+          config.timezone = ENV["TIMEZONE"].present? ? ENV["TIMEZONE"] : "America/New_York"
+          config.wandb_api_key = ENV["WANDB_API_KEY"] if ENV["WANDB_API_KEY"]
+        end
+      end
+    end
+
     initializer "easy_ml.check_pending_migrations" do
       if defined?(Rails::Server)
         config.after_initialize do
@@ -96,19 +111,6 @@ module EasyML
       end
     end
 
-    initializer "easy_ml.configure_secrets" do
-      EasyML::Configuration.configure do |config|
-        raise "S3_ACCESS_KEY_ID is missing. Set ENV['S3_ACCESS_KEY_ID']" unless ENV["S3_ACCESS_KEY_ID"]
-        raise "S3_SECRET_ACCESS_KEY is missing. Set ENV['S3_SECRET_ACCESS_KEY']" unless ENV["S3_SECRET_ACCESS_KEY"]
-
-        config.s3_access_key_id = ENV["S3_ACCESS_KEY_ID"]
-        config.s3_secret_access_key = ENV["S3_SECRET_ACCESS_KEY"]
-        config.s3_region = ENV["S3_REGION"] if ENV["S3_REGION"]
-        config.timezone = ENV["TIMEZONE"].present? ? ENV["TIMEZONE"] : "America/New_York"
-        config.wandb_api_key = ENV["WANDB_API_KEY"] if ENV["WANDB_API_KEY"]
-      end
-    end
-
     initializer "easy_ml.setup_generators" do |app|
       generators_path = EasyML::Engine.root.join("lib/easy_ml/railtie/generators")
       generators_dirs = Dir[File.join(generators_path, "**", "*.rb")]

data/lib/easy_ml/feature_store.rb

@@ -1,69 +1,28 @@
 module EasyML
-  class FeatureStore
+  class FeatureStore < EasyML::Data::DatasetManager
     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
-
-      begin
-        # We are intentionally not using to_i, so it will raise an error for keys like "A1"
-        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)
+      datasource_config = feature&.dataset&.datasource&.configuration
+      if datasource_config
+        options = {
+          root_dir: feature_dir,
+          filenames: "feature",
+          append_only: false,
+          primary_key: feature.primary_key&.first,
+          partition_size: batch_size,
+          s3_bucket: datasource_config.dig("s3_bucket") || EasyML::Configuration.s3_bucket,
+          s3_prefix: s3_prefix,
+          polars_args: datasource_config.dig("polars_args"),
+        }.compact
+        super(options)
+      else
+        super({ root_dir: "" })
       end
     end
 
-    def query(**kwargs)
-      query_all_partitions(**kwargs)
-    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)
@@ -74,7 +33,7 @@ module EasyML
       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)
+        target_file = file.gsub(old_dir, new_dir)
         FileUtils.mkdir_p(File.dirname(target_file))
         FileUtils.cp(file, target_file)
       end
@@ -82,77 +41,17 @@ module EasyML
 
     private
 
-    def cleanup(type: :partitions)
-      case type
-      when :partitions
-        list_partitions.each do |partition|
-          FileUtils.rm(partition)
-        end
-      when :no_partitions
-        FileUtils.rm_rf(feature_path)
-      when :all
-        wipe
-      end
-    end
-
-    def store_without_partitioning(df)
-      lock_file do
-        cleanup(type: :partitions)
-        path = feature_path
-        safe_write(df, path)
-      end
-    end
-
-    def safe_write(df, path)
-      FileUtils.mkdir_p(File.dirname(path))
-      df.write_parquet(path)
-    end
-
-    def store_partition(partition_df, primary_key, partition_start)
-      lock_partition(partition_start) do
-        cleanup(type: :no_partitions)
-        path = partition_path(partition_start)
-
-        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
-          )
-          if preserved_records.shape[1] != partition_df.shape[1]
-            wipe
-          else
-            partition_df = Polars.concat([preserved_records, partition_df], how: "vertical")
-          end
-        end
-
-        safe_write(partition_df, path)
-      end
-    end
-
-    def query_all_partitions(**kwargs)
-      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, **kwargs)
-    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
+    def batch_size
+      @batch_size ||= feature.batch_size || 10_000
     end
 
     def feature_dir_for_version(version)
       File.join(
         Rails.root,
         "easy_ml/datasets",
-        feature.dataset.name.parameterize.gsub("-", "_"),
+        feature&.dataset&.name&.parameterize&.gsub("-", "_"),
         "features",
-        feature.name.parameterize.gsub("-", "_"),
+        feature&.name&.parameterize&.gsub("-", "_"),
         version.to_s
       )
     end
@@ -161,74 +60,8 @@ module EasyML
       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/reasons.rb

@@ -0,0 +1,41 @@
+module EasyML
+  class Reasons
+    def initialize(context)
+      @context = context
+      @reasons = {}
+    end
+
+    class << self
+      def add_reason(name, check)
+        @reasons ||= {}
+        key = name.to_s.downcase.gsub(/\s/, "_").to_sym
+        @reasons[key] = { name: name, check: check }
+      end
+
+      def reasons
+        @reasons ||= {}
+      end
+    end
+
+    def inspect
+      "#<#{self.class.name.split("::").last} checks=[#{self.class.reasons.map { |k, v| "#{v[:name]}" }.join(", ")}]>"
+    end
+
+    def none?(except: [])
+      check(except: except).none?
+    end
+
+    def check(except: [])
+      self.class.reasons.except(*except).select do |_, config|
+        @context.instance_exec(&config[:check])
+      end.map do |_, config|
+        config[:name]
+      end
+    end
+
+    def explain
+      reasons = check
+      reasons.any? ? reasons.join(", ") : :none
+    end
+  end
+end

data/lib/easy_ml/support/lockable.rb

@@ -22,7 +22,7 @@ module EasyML
       end
 
       def self.lock_client(key, wait_timeout: 0.1, stale_timeout: 60 * 10, resources: 1)
-        Suo::Client::Redis.new(prefixed_key(key), {
+        Suo::Client::Redis.new(key, {
           acquisition_timeout: wait_timeout,
           stale_lock_expiry: stale_timeout,
           resources: resources,
@@ -30,10 +30,6 @@ module EasyML
         })
       end
 
-      def self.prefixed_key(key)
-        "easy_ml:#{key}"
-      end
-
       # Execute a block with a Redis lock
       def self.with_lock(key, wait_timeout: 0.1, stale_timeout: 60 * 10, resources: 1)
         lock_key = nil

data/lib/easy_ml/version.rb

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 
 module EasyML
-  VERSION = "0.2.0-rc76"
+  VERSION = "0.2.0-rc78"
 
   module Version
   end

data/lib/easy_ml.rb

@@ -15,13 +15,13 @@ module EasyML
   class Error < StandardError; end
 
   require_relative "easy_ml/configuration"
+  require_relative "easy_ml/reasons"
   require_relative "easy_ml/deep_compact"
   require_relative "easy_ml/timing"
   require_relative "easy_ml/support"
   require_relative "easy_ml/core_ext"
   require_relative "easy_ml/logging"
   require_relative "easy_ml/data"
-  require_relative "easy_ml/data/filter_extensions"
   require_relative "easy_ml/evaluators"
   require_relative "easy_ml/features"
   require_relative "easy_ml/feature_store"

data/public/easy_ml/assets/.vite/manifest.json

@@ -1,6 +1,6 @@
 {
   "entrypoints/Application.tsx": {
-    "file": "assets/entrypoints/Application.tsx-B1qLZuyu.js",
+    "file": "assets/entrypoints/Application.tsx-Bbf3mD_b.js",
     "name": "entrypoints/Application.tsx",
     "src": "entrypoints/Application.tsx",
     "isEntry": true,