easy_ml 0.1.4 → 0.2.0.pre.rc1

data/lib/easy_ml/data/simple_imputer.rb

@@ -0,0 +1,255 @@
+require "active_support/core_ext/hash/deep_transform_values"
+require "numo/narray"
+require "json"
+
+module EasyML
+  module Data
+    class SimpleImputer
+      attr_reader :statistics
+      attr_accessor :path, :attribute, :strategy, :options
+
+      def initialize(strategy: "mean", path: nil, attribute: nil, options: {}, statistics: {}, &block)
+        @strategy = strategy.to_sym
+        @path = path
+        @attribute = attribute
+        @options = options || {}
+        apply_defaults
+        @statistics = statistics || {}
+        deep_symbolize_keys!
+        return unless block_given?
+
+        instance_eval(&block)
+      end
+
+      def deep_symbolize_keys!
+        @statistics = @statistics.deep_symbolize_keys
+      end
+
+      def apply_defaults
+        @options[:date_column] ||= "CREATED_DATE"
+
+        if strategy == :categorical
+          @options[:categorical_min] ||= 25
+        elsif strategy == :custom
+          itself = ->(col) { col }
+          @options[:fit] ||= itself
+          @options[:transform] ||= itself
+        end
+      end
+
+      def fit(x, df = nil)
+        x = validate_input(x)
+
+        fit_values = case @strategy
+        when :mean
+          fit_mean(x)
+        when :median
+          fit_median(x)
+        when :ffill
+          fit_ffill(x, df)
+        when :most_frequent
+          fit_most_frequent(x)
+        when :categorical
+          fit_categorical(x)
+        when :constant
+          fit_constant(x)
+        when :clip
+          fit_no_op(x)
+        when :today
+          fit_no_op(x)
+        when :one_hot
+          fit_no_op(x)
+        when :custom
+          fit_custom(x)
+        else
+          raise ArgumentError, "Invalid strategy: #{@strategy}"
+        end || {}
+
+        @statistics[attribute] ||= {}
+        @statistics[attribute][@strategy] = fit_values.merge!(original_dtype: x.dtype)
+        @statistics.deep_symbolize_keys
+      end
+
+      def transform(x)
+        check_is_fitted
+
+        if x.is_a?(Polars::Series)
+          transform_polars(x)
+        else
+          transform_dense(x)
+        end
+      end
+
+      def transform_polars(x)
+        case @strategy
+        when :mean, :median
+          x.fill_null(@statistics[@strategy])
+        when :ffill
+          x.fill_null(@statistics[:last_value])
+        when :most_frequent
+          x.fill_null(@statistics[:most_frequent_value])
+        when :constant
+          x.fill_null(@options[:constant])
+        when :categorical
+          allowed_cats = statistics[:allowed_categories]
+          df = Polars::DataFrame.new({ x: x })
+          df.with_column(
+            Polars.when(Polars.col("x").is_in(allowed_cats))
+              .then(Polars.col("x"))
+              .otherwise(Polars.lit("other"))
+              .alias("x")
+          )["x"]
+        when :clip
+          min = options["min"] || 0
+          max = options["max"] || 1_000_000_000_000
+          if x.null_count != x.len
+            x.clip(min, max)
+          else
+            x
+          end
+        when :today
+          x.fill_null(transform_today(nil))
+        when :custom
+          if x.null_count == x.len
+            x.fill_null(transform_custom(nil))
+          else
+            x.apply do |val|
+              should_transform_custom?(val) ? transform_custom(val) : val
+            end
+          end
+        else
+          raise ArgumentError, "Unsupported strategy for Polars::Series: #{@strategy}"
+        end
+      end
+
+      def file_path
+        raise "Need both attribute and path to save/load statistics" unless attribute.present? && path.to_s.present?
+
+        File.join(path, "statistics.json")
+      end
+
+      def transform_today(_val)
+        UTC.now.beginning_of_day
+      end
+
+      def fit_custom(x)
+        x
+      end
+
+      def should_transform_custom?(x)
+        if options.key?(:should_transform)
+          options[:should_transform].call(x)
+        else
+          should_transform_default?(x)
+        end
+      end
+
+      def transform_custom(x)
+        raise "transform required" unless options.key?(:transform)
+
+        options[:transform].call(x)
+      end
+
+      private
+
+      def validate_input(x)
+        raise ArgumentError, "Input must be a Polars::Series" unless x.is_a?(Polars::Series)
+
+        x
+      end
+
+      def fit_mean(x)
+        { value: x.mean }
+      end
+
+      def fit_median(x)
+        { value: x.median }
+      end
+
+      def fit_ffill(x, df)
+        values = { value: nil, max_date: nil }
+
+        date_col = df[options[:date_column]]
+        return if date_col.is_null.all
+
+        sorted_df = df.sort(options[:date_column])
+        new_max_date = sorted_df[options[:date_column]].max
+
+        current_max_date = values[:max_date]
+        return if current_max_date && current_max_date > new_max_date
+
+        values[:max_date] = [current_max_date, new_max_date].compact.max
+
+        # Get the last non-null value
+        last_non_null = sorted_df[x.name].filter(sorted_df[x.name].is_not_null).tail(1).to_a.first
+        values[:value] = last_non_null
+
+        values
+      end
+
+      def fit_most_frequent(x)
+        value_counts = x.filter(x.is_not_null).value_counts
+        column_names = value_counts.columns
+        column_names[0]
+        count_column = column_names[1]
+
+        most_frequent_value = value_counts.sort(count_column, descending: true).row(0)[0]
+        { value: most_frequent_value }
+      end
+
+      def fit_no_op(_x)
+        {}
+      end
+
+      def fit_constant(_x)
+        { value: @options[:fill_value] }
+      end
+
+      def transform_default(_val)
+        @statistics[strategy][:value]
+      end
+
+      def should_transform_default?(val)
+        checker_method = val.respond_to?(:nan?) ? :nan? : :nil?
+        val.send(checker_method)
+      end
+
+      def transform_dense(x)
+        result = x.map do |val|
+          strategy_method = respond_to?("transform_#{strategy}") ? "transform_#{strategy}" : "transform_default"
+          checker_method = respond_to?("should_transform_#{strategy}?") ? "should_transform_#{strategy}?" : "should_transform_default?"
+          send(checker_method, val) ? send(strategy_method, val) : val
+        end
+
+        # Cast the result back to the original dtype
+        original_dtype = @statistics[:original_dtype]
+        if original_dtype
+          result.map { |val| cast_to_dtype(val, original_dtype) }
+        else
+          result
+        end
+      end
+
+      def check_is_fitted
+        return if %i[clip today custom].include?(strategy)
+
+        pass_check = case strategy
+          when :mean
+            @statistics.dig(:mean).present?
+          when :median
+            @statistics.dig(:median).present?
+          when :ffill
+            @statistics.dig(:last_value).present?
+          when :most_frequent
+            @statistics.key?(:most_frequent_value)
+          when :constant
+            options.dig(:constant).present?
+          when :categorical
+            true
+          end
+
+        raise "SimpleImputer has not been fitted yet for #{attribute}##{strategy}" unless pass_check
+      end
+    end
+  end
+end

data/lib/easy_ml/data/splits/file_split.rb

@@ -0,0 +1,252 @@
+require_relative "split"
+
+module EasyML
+  module Data
+    module Splits
+      class FileSplit < Split
+        include EasyML::Data::Utils
+
+        attr_accessor :dir, :polars_args, :max_rows_per_file, :batch_size, :verbose,
+                      :dataset, :datasource
+
+        def initialize(options = {})
+          super
+          @dir = options[:dir]
+          @polars_args = options[:polars_args] || {}
+          @max_rows_per_file = options[:max_rows_per_file] || 1_000_000
+          @batch_size = options[:batch_size] || 10_000
+          @verbose = options[:verbose] || false
+          @dataset = options[:dataset]
+          @datasource = options[:datasource]
+          FileUtils.mkdir_p(dir)
+        end
+
+        def attributes
+          {
+            dir: dir,
+            polars_args: polars_args,
+            max_rows_per_file: max_rows_per_file,
+            batch_size: batch_size,
+            verbose: verbose,
+            dataset: dataset,
+            datasource: datasource,
+          }.with_indifferent_access
+        end
+
+        def s3_prefix
+          File.join("datasets", dir.split("datasets").last)
+        end
+
+        def synced_directory
+          datasource_config = datasource.configuration || {}
+          @synced_dir ||= EasyML::Data::SyncedDirectory.new(
+            root_dir: 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 should_sync?
+          synced_directory.should_sync?
+        end
+
+        def download
+          synced_directory.download
+        end
+
+        def upload
+          synced_directory.upload
+        end
+
+        # cp can receive a directory or just a version string
+        def cp(target_dir)
+          target_dir = version_to_dir(target_dir) if is_version?(target_dir)
+          return self if target_dir.nil?
+
+          target_dir = target_dir.to_s
+
+          puts "copying #{dir} to #{target_dir}"
+          FileUtils.mkdir_p(target_dir)
+
+          files_to_cp = Dir.glob(Pathname.new(dir).join("**/*")).select { |f| File.file?(f) }
+
+          files_to_cp.each do |file|
+            target_parts = target_dir.split("/")
+            file_parts = file.split("/")
+            _, not_shared = file_parts.partition.with_index { |part, index| target_parts[index] == part }
+            not_shared = not_shared[1..-1].join("/") # remove leading "processed", "raw", etc
+
+            target_file = File.join(target_dir, not_shared)
+            FileUtils.mkdir_p(File.dirname(target_file))
+            FileUtils.cp(file, target_file)
+          end
+
+          self.class.new(dir: target_dir, **attributes.except("dir"))
+        end
+
+        def save(segment, df)
+          return unless df.present?
+
+          segment_dir = File.join(dir, segment.to_s)
+          FileUtils.mkdir_p(segment_dir)
+
+          file_path = new_file_path_for_segment(segment)
+          df.write_parquet(file_path)
+          file_path
+        end
+
+        def query(**kwargs, &block)
+          read(:all, **kwargs, &block)
+        end
+
+        def read(segment, split_ys: false, target: nil, drop_cols: [], filter: nil, limit: nil, select: nil,
+                          unique: nil, sort: nil, descending: false, batch_size: nil, batch_start: nil, batch_key: nil, &block)
+          files = files_for_segment(segment)
+          return split_ys ? [nil, nil] : nil if files.empty?
+
+          query_params = {
+            filter: filter,
+            limit: limit,
+            select: select,
+            unique: unique,
+            drop_cols: drop_cols,
+            sort: sort,
+            descending: descending,
+            batch_size: batch_size,
+            batch_start: batch_start,
+            batch_key: batch_key,
+          }.compact
+
+          if batch_size.present?
+            base_enumerator = EasyML::Data::PolarsReader.query(files, **query_params)
+
+            if block_given?
+              wrap_with_block(base_enumerator, split_ys, target, &block)
+            else
+              wrap_with_split(base_enumerator, target)
+            end
+          else
+            df = EasyML::Data::PolarsReader.query(files, **query_params, &block)
+            split_features_targets(df, split_ys, target)
+          end
+        end
+
+        def cleanup
+          FileUtils.rm_rf(dir)
+          FileUtils.mkdir_p(dir)
+        end
+
+        def split_at
+          return nil if output_files.empty?
+
+          output_files.map { |file| File.mtime(file) }.max
+        end
+
+        def num_batches(segment)
+          files_for_segment(segment).count
+        end
+
+        def files
+          files_for_segment("all")
+        end
+
+        private
+
+        def wrap_with_block(base_enumerator, split_ys, target, &block)
+          Enumerator.new do |yielder|
+            base_enumerator.each do |df|
+              df = block.call(df)
+              result = process_dataframe(df, split_ys, target)
+              yielder << result
+            end
+          end
+        end
+
+        def wrap_with_split(base_enumerator, target)
+          Enumerator.new do |yielder|
+            base_enumerator.each do |df|
+              result = process_dataframe(df, true, target)
+              yielder << result
+            end
+          end
+        end
+
+        def process_dataframe(df, split_ys, target)
+          return df unless split_ys
+
+          df = df.collect if df.is_a?(Polars::LazyFrame)
+          split_features_targets(df, split_ys, target)
+        end
+
+        def version_to_dir(version)
+          relative_path = dir.gsub(Regexp.new(Rails.root.to_s), "")
+          current_path = Pathname.new(relative_path)
+
+          # Find the version component in the path
+          path_parts = current_path.each_filename.to_a
+          version_index = path_parts.find_index { |part| part.match?(version_pattern) }
+
+          return unless version_index
+
+          old_version = path_parts[version_index]
+          return if old_version == version
+
+          # Replace the version number with the new version
+          path_parts[version_index] = version
+          Rails.root.join(File.join(*path_parts))
+        end
+
+        def version_pattern
+          /^\d{14}$/
+        end
+
+        def is_version?(string)
+          string.to_s.match?(version_pattern)
+        end
+
+        def df(path)
+          filtered_args = filter_polars_args(Polars.method(:read_parquet))
+          Polars.read_parquet(path, **filtered_args)
+        end
+
+        def filter_polars_args(method)
+          supported_params = method.parameters.map { |_, name| name }
+          polars_args.select { |k, _| supported_params.include?(k) }
+        end
+
+        def output_files
+          Dir.glob("#{dir}/**/*.parquet")
+        end
+
+        def files_for_segment(segment)
+          if segment.to_s == "all"
+            files_for_segment("train") + files_for_segment("test") + files_for_segment("valid")
+          else
+            segment_dir = File.join(dir, segment.to_s)
+            Dir.glob(File.join(segment_dir, "**/*.parquet")).sort
+          end
+        end
+
+        def current_file_for_segment(segment)
+          segment_dir = File.join(dir, segment.to_s)
+          File.join(segment_dir, "#{segment}.parquet")
+        end
+
+        def new_file_path_for_segment(segment)
+          segment_dir = File.join(dir, segment.to_s)
+          file_number = Dir.glob(File.join(segment_dir, "*.parquet")).count
+          File.join(segment_dir, "#{segment}_%04d.parquet" % file_number)
+        end
+
+        def combine_dataframes(files)
+          dfs = files.map { |file| df(file) }
+          Polars.concat(dfs)
+        end
+      end
+    end
+  end
+end

data/lib/easy_ml/data/splits/in_memory_split.rb

@@ -0,0 +1,54 @@
+module EasyML
+  module Data
+    module Splits
+      class InMemorySplit < Split
+        attr_accessor :dataset
+
+        def initialize(options = {})
+          @data = {}
+          @dataset = options[:dataset]
+        end
+
+        # We don't backup in-memory splits to s3
+        def download; end
+
+        def upload; end
+
+        def files
+          []
+        end
+
+        def save(segment, df)
+          @data[segment] = df
+        end
+
+        def read(segment, split_ys: false, target: nil, drop_cols: [], filter: nil, limit: nil, select: nil,
+                          unique: nil)
+          return nil if @data.keys.none?
+
+          df = if segment.to_s == "all"
+              Polars.concat(EasyML::Dataset::SPLIT_ORDER.map { |segment| @data[segment] }.compact)
+            else
+              @data[segment]
+            end
+          return nil if df.nil?
+
+          df = df.filter(filter) if filter.present?
+          drop_cols &= df.columns
+          df = df.drop(drop_cols) unless drop_cols.empty?
+          df = df.unique if unique
+
+          split_features_targets(df, split_ys, target)
+        end
+
+        def cleanup
+          @data.clear
+        end
+
+        def split_at
+          @data.keys.empty? ? nil : Time.now
+        end
+      end
+    end
+  end
+end

data/lib/easy_ml/data/splits/split.rb

@@ -0,0 +1,95 @@
+module EasyML
+  module Data
+    module Splits
+      class Split
+        include EasyML::Data::Utils
+
+        VALID_SEGMENTS = %w[train test valid all].freeze
+
+        def initialize(options = {})
+        end
+
+        def load_data(segment, **kwargs)
+          drop_cols = dataset.drop_columns(all_columns: kwargs[:all_columns] || false)
+          kwargs.delete(:all_columns)
+          kwargs = kwargs.merge!(drop_cols: drop_cols, target: dataset.target)
+          read(segment, **kwargs)
+        end
+
+        def save(segment, _df)
+          validate_segment!(segment)
+          raise NotImplementedError, "Subclasses must implement #save"
+        end
+
+        def data(**kwargs, &block)
+          load_data(:all, **kwargs, &block)
+        end
+
+        def train(**kwargs, &block)
+          load_data(:train, **kwargs, &block)
+        end
+
+        def test(**kwargs, &block)
+          load_data(:test, **kwargs, &block)
+        end
+
+        def valid(**kwargs, &block)
+          load_data(:valid, **kwargs, &block)
+        end
+
+        def cleanup
+          raise NotImplementedError, "Subclasses must implement #cleanup"
+        end
+
+        def split_at
+          raise NotImplementedError, "Subclasses must implement #split_at"
+        end
+
+        protected
+
+        def split_features_targets(df, split_ys, target)
+          return df unless split_ys
+          raise ArgumentError, "Target column must be specified when split_ys is true" if target.nil?
+
+          xs = df.drop(target)
+          ys = df.select(target)
+          [xs, ys]
+        end
+
+        def validate_segment!(segment)
+          segment = segment.to_s
+          return if VALID_SEGMENTS.include?(segment)
+
+          raise ArgumentError, "Invalid segment: #{segment}. Must be one of: #{VALID_SEGMENTS.join(", ")}"
+        end
+
+        def validate_read_options!(options)
+          valid_options = %i[filter limit select unique sort descending batch_size batch_start batch_key]
+          invalid_options = options.keys - valid_options
+          return if invalid_options.empty?
+
+          raise ArgumentError,
+                "Invalid options: #{invalid_options.join(", ")}. Valid options are: #{valid_options.join(", ")}"
+        end
+
+        private
+
+        def process_block_with_split_ys(block, result, xs, ys)
+          case block.arity
+          when 3 then result.nil? ? [xs, ys] : block.call(result, xs, ys)
+          when 2 then block.call(xs, ys) && result
+          else raise ArgumentError, "Block must accept 2 or 3 arguments when split_ys is true"
+          end
+        end
+
+        def process_block_without_split_ys(block, result, df)
+          case block.arity
+          when 2 then result.nil? ? df : block.call(result, df)
+          when 1 then block.call(df) && result
+          else raise ArgumentError, "Block must accept 1 or 2 arguments when split_ys is false"
+          end
+        end
+      end
+    end
+  end
+end

data/lib/easy_ml/data/splits.rb

@@ -0,0 +1,9 @@
+module EasyML
+  module Data
+    module Splits
+      require_relative "splits/split"
+      require_relative "splits/file_split"
+      require_relative "splits/in_memory_split"
+    end
+  end
+end