easy_ml 0.1.4 → 0.2.0.pre.rc1

data/lib/easy_ml/data/polars_reader.rb

@@ -0,0 +1,297 @@
+module EasyML
+  module Data
+    class PolarsReader
+      attr_accessor :root_dir, :polars_args, :refresh, :num_rows
+      attr_reader :schema
+
+      def initialize(options = {})
+        @root_dir = options[:root_dir]
+        @polars_args = options[:polars_args] || {}
+        @refresh = options[:refresh] || false
+        @num_rows = options[:num_rows]
+        @schema = options[:schema]
+      end
+
+      def schema=(value)
+        @schema = value
+        polars_args[:dtypes] = value
+      end
+
+      def normalize
+        learn_dataset
+        convert_to_parquet
+        files
+      end
+
+      def clean
+        FileUtils.rm(parquet_files) unless Rails.env.test?
+      end
+
+      def in_batches
+        normalize
+
+        files.each do |file|
+          yield read_file(file)
+        end
+      end
+
+      def all_files
+        (parquet_files + csv_files)
+      end
+
+      def files
+        if parquet_files.any? && !refresh
+          parquet_files
+        else
+          csv_files
+        end
+      end
+
+      def data
+        query
+      end
+
+      def convert_to_parquet(columns = nil)
+        return files if any_parquet? && columns.nil?
+
+        puts "Converting to Parquet..."
+
+        csv_files.each do |path|
+          df = read_file(path, columns)
+          df = cast(df, columns)
+          path.dup
+          filename = Pathname.new(path).basename
+          ext = Pathname.new(path).extname.gsub(/\./, "")
+          filename = filename.to_s.gsub(Regexp.new(ext), "parquet")
+          path = File.join(root_dir, filename).to_s
+          df.write_parquet(path)
+        end
+      end
+
+      def query(files = 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 ||= self.files
+        PolarsReader.query(files, drop_cols: drop_cols, filter: filter, limit: limit,
+                                  select: select, unique: unique, sort: sort, descending: descending,
+                                  batch_size: batch_size, batch_start: batch_start, batch_key: batch_key, &block)
+      end
+
+      def self.query(files, drop_cols: [], filter: nil, limit: nil, select: nil, unique: nil, sort: nil, descending: false,
+                            batch_size: nil, batch_start: nil, batch_key: nil, &block)
+        return query_files(files, drop_cols: drop_cols, filter: filter, limit: limit, select: select,
+                                  unique: unique, sort: sort, descending: descending).collect unless batch_size.present?
+
+        return batch_enumerator(files, drop_cols: drop_cols, filter: filter, limit: limit, select: select, unique: unique, sort: sort, descending: descending,
+                                       batch_size: batch_size, batch_start: batch_start, batch_key: batch_key) unless block_given?
+
+        process_batches(files, drop_cols: drop_cols, filter: filter, limit: limit, select: select, unique: unique, sort: sort, descending: descending,
+                               batch_size: batch_size, batch_start: batch_start, batch_key: batch_key, &block)
+      end
+
+      private
+
+      def self.batch_enumerator(files, drop_cols: [], filter: nil, limit: nil, select: nil, unique: nil, sort: nil, descending: false,
+                                       batch_size: nil, batch_start: nil, batch_key: nil, &block)
+        Enumerator.new do |yielder|
+          process_batches(files, drop_cols: drop_cols, filter: filter, limit: limit, select: select, unique: unique, sort: sort, descending: descending,
+                                 batch_size: batch_size, batch_start: batch_start, batch_key: batch_key) do |batch|
+            yielder << batch
+          end
+        end
+      end
+
+      def self.process_batches(files, drop_cols: [], filter: nil, limit: nil, select: nil, unique: nil, sort: nil, descending: false,
+                                      batch_size: nil, batch_start: nil, batch_key: nil, &block)
+        batch_key ||= identify_primary_key(files, select: select)
+        raise "When using batch_size, sort must match primary key (#{batch_key})" if sort.present? && batch_key != sort
+
+        sort = batch_key
+        batch_start = query_files(files, sort: sort, descending: descending, select: batch_key, limit: 1).collect[batch_key].to_a.last unless batch_start
+        final_value = query_files(files, sort: sort, descending: !descending, select: batch_key, limit: 1).collect[batch_key].to_a.last
+
+        is_first_batch = true
+        current_start = batch_start
+
+        while current_start < final_value
+          filter = is_first_batch ? Polars.col(sort) >= current_start : Polars.col(sort) > current_start
+          batch = query_files(files, drop_cols: drop_cols, filter: filter, limit: batch_size, select: select, unique: unique, sort: sort, descending: descending)
+          yield batch
+          current_start = query_files(files, sort: sort, descending: descending, limit: batch_size, filter: filter).sort(sort, reverse: !descending).limit(1).select(batch_key).collect[batch_key].to_a.last
+          is_first_batch = false
+        end
+      end
+
+      def self.query_files(files, drop_cols: [], filter: nil, limit: nil, select: nil, unique: nil, sort: nil, descending: false)
+        lazy_frames = to_lazy_frames(files)
+        combined_lazy_df = Polars.concat(lazy_frames)
+
+        # Apply the predicate filter if given
+        combined_lazy_df = combined_lazy_df.filter(filter) if filter
+        # Apply select columns if provided
+        combined_lazy_df = combined_lazy_df.select(select) if select.present?
+        combined_lazy_df = combined_lazy_df.unique if unique
+
+        # Apply sorting if provided
+        combined_lazy_df = combined_lazy_df.sort(sort, reverse: descending) if sort
+
+        # Apply drop columns
+        drop_cols &= combined_lazy_df.columns
+        combined_lazy_df = combined_lazy_df.drop(drop_cols) unless drop_cols.empty?
+
+        # Collect the DataFrame (execute the lazy operations)
+        combined_lazy_df = combined_lazy_df.limit(limit) if limit
+        combined_lazy_df
+      end
+
+      def self.identify_primary_key(files, select: nil)
+        lazy_df = to_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
+
+        return primary_keys.first
+      end
+
+      def self.lazy_schema(files)
+        to_lazy_frames([files.first]).first.schema
+      end
+
+      def self.to_lazy_frames(files)
+        files.map do |file|
+          case Pathname.new(file).extname.gsub(/\./, "")
+          when "csv"
+            Polars.scan_csv(file)
+          when "parquet"
+            Polars.scan_parquet(file)
+          end
+        end
+      end
+
+      def read_file(file, columns = nil)
+        if columns
+          dtypes = columns_to_dtypes(columns)
+          polars_args[:dtypes] ||= {}
+          polars_args[:dtypes].merge!(dtypes)
+        end
+        ext = Pathname.new(file).extname.gsub(/\./, "")
+        case ext
+        when "csv"
+          filtered_args = filter_polars_args(Polars.method(:read_csv))
+          filtered_args.merge!(infer_schema_length: 1_000_000, null_values: ["\\N", "\\\\N", "NULL"])
+          df = Polars.read_csv(file, **filtered_args)
+        when "parquet"
+          filtered_args = filter_polars_args(Polars.method(:read_parquet))
+          df = Polars.read_parquet(file, **filtered_args)
+        end
+        df
+      end
+
+      def any_parquet?
+        files.any? { |f| f.match?(/\.parquet$/) }
+      end
+
+      def filter_polars_args(method)
+        supported_params = method.parameters.map { |_, name| name }
+        polars_args.select { |k, _| supported_params.include?(k) }
+      end
+
+      def csv_files
+        Dir.glob(File.join(root_dir, "**/*.{csv}"))
+      end
+
+      def parquet_files
+        Dir.glob(File.join(root_dir, "**/*.{parquet}"))
+      end
+
+      def columns_to_dtypes(columns)
+        columns.reduce({}) { |h, c| h[c.name] = c.polars_type; h }
+      end
+
+      def cast(df, columns = [])
+        keep_cols = columns && columns.any? ? columns.map(&:name) : schema.keys
+        lookup = columns && columns.any? ? columns_to_dtypes(columns) : schema
+        cast_cols = keep_cols & df.columns
+        df = df.with_columns(
+          cast_cols.map do |column|
+            dtype = lookup[column]
+            df[column].cast(dtype).alias(column)
+          end
+        )
+      end
+
+      def existing_parquet_schema
+        return nil if parquet_files.empty?
+
+        file_path = parquet_files.first
+        df = Polars.scan_parquet(file_path)
+        df = df.limit(1)
+        df.collect.schema
+      end
+
+      def learn_dataset
+        return schema if schema.present?
+
+        existing_schema = existing_parquet_schema
+        schema = existing_schema || normalize_dataset
+
+        self.schema = schema
+        polars_args[:dtypes] = schema
+      end
+
+      def normalize_dataset
+        puts "Normalizing schema..."
+        self.num_rows = 0
+        first_file = read_file(files.first)
+
+        files.map.with_index do |path, _idx|
+          df = read_file(path)
+          self.num_rows += df.shape[0]
+          df.schema
+        end.inject({}) do |h, schema|
+          h.tap do
+            schema.each do |key, value|
+              h[key] ||= []
+              h[key] << value unless h[key].include?(value)
+            end
+          end
+        end.inject({}) do |h, (k, v)|
+          h.tap do
+            values = v.map { |klass| klass.to_s.gsub(/Polars::/, "") }
+            h[k] = if values.any? { |v| v.match?(/Float/) }
+                Polars::Float64
+              elsif values.any? { |v| v.match?(/Int/) }
+                Polars::Int64
+              else
+                type = EasyML::Data::PolarsColumn.determine_type(first_file[k], true)
+                raise "Cannot determine polars type for field #{k}" if type.nil?
+
+                type
+              end
+          end
+        end
+      end
+    end
+  end
+end