polars-df 0.17.0 → 0.17.1

data/lib/polars/data_frame.rb

@@ -495,6 +495,45 @@ module Polars
       _df.arrow_c_stream
     end
 
+    # Get an ordered mapping of column names to their data type.
+    #
+    # @return [Schema]
+    #
+    # @note
+    #   This method is included to facilitate writing code that is generic for both
+    #   DataFrame and LazyFrame.
+    #
+    # @example Determine the schema.
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "foo" => [1, 2, 3],
+    #       "bar" => [6.0, 7.0, 8.0],
+    #       "ham" => ["a", "b", "c"]
+    #     }
+    #   )
+    #   df.collect_schema
+    #   # => Polars::Schema({"foo"=>Polars::Int64, "bar"=>Polars::Float64, "ham"=>Polars::String})
+    #
+    # @example Access various properties of the schema using the `Schema` object.
+    #   schema = df.collect_schema
+    #   schema["bar"]
+    #   # => Polars::Float64
+    #
+    # @example
+    #   schema.names
+    #   # => ["foo", "bar", "ham"]
+    #
+    # @example
+    #   schema.dtypes
+    #   # => [Polars::Int64, Polars::Float64, Polars::String]
+    #
+    # @example
+    #   schema.length
+    #   # => 3
+    def collect_schema
+      Schema.new(columns.zip(dtypes), check_dtypes: false)
+    end
+
     # Return the dataframe as a scalar.
     #
     # Equivalent to `df[0,0]`, with a check that the shape is (1,1).
@@ -961,6 +1000,139 @@ module Polars
       )
     end
 
+    # Write the data in a Polars DataFrame to a database.
+    #
+    # @param table_name [String]
+    #   Schema-qualified name of the table to create or append to in the target
+    #   SQL database.
+    # @param connection [Object]
+    #   An existing Active Record connection against the target database.
+    # @param if_table_exists ['append', 'replace', 'fail']
+    #   The insert mode:
+    #
+    #   * 'replace' will create a new database table, overwriting an existing one.
+    #   * 'append' will append to an existing table.
+    #   * 'fail' will fail if table already exists.
+    #
+    # @return [Integer]
+    #
+    # @note
+    #   This functionality is experimental. It may be changed at any point without it being considered a breaking change.
+    def write_database(table_name, connection = nil, if_table_exists: "fail")
+      if !defined?(ActiveRecord)
+        raise Error, "Active Record not available"
+      elsif ActiveRecord::VERSION::MAJOR < 7
+        raise Error, "Requires Active Record 7+"
+      end
+
+      valid_write_modes = ["append", "replace", "fail"]
+      if !valid_write_modes.include?(if_table_exists)
+        msg = "write_database `if_table_exists` must be one of #{valid_write_modes.inspect}, got #{if_table_exists.inspect}"
+        raise ArgumentError, msg
+      end
+
+      with_connection(connection) do |connection|
+        table_exists = connection.table_exists?(table_name)
+        if table_exists && if_table_exists == "fail"
+          raise ArgumentError, "Table already exists"
+        end
+
+        create_table = !table_exists || if_table_exists == "replace"
+        maybe_transaction(connection, create_table) do
+          if create_table
+            mysql = connection.adapter_name.match?(/mysql|trilogy/i)
+            force = if_table_exists == "replace"
+            connection.create_table(table_name, id: false, force: force) do |t|
+              schema.each do |c, dtype|
+                options = {}
+                column_type =
+                  case dtype
+                  when Binary
+                    :binary
+                  when Boolean
+                    :boolean
+                  when Date
+                    :date
+                  when Datetime
+                    :datetime
+                  when Decimal
+                    if mysql
+                      options[:precision] = dtype.precision || 65
+                      options[:scale] = dtype.scale || 30
+                    end
+                    :decimal
+                  when Float32
+                    options[:limit] = 24
+                    :float
+                  when Float64
+                    options[:limit] = 53
+                    :float
+                  when Int8
+                    options[:limit] = 1
+                    :integer
+                  when Int16
+                    options[:limit] = 2
+                    :integer
+                  when Int32
+                    options[:limit] = 4
+                    :integer
+                  when Int64
+                    options[:limit] = 8
+                    :integer
+                  when UInt8
+                    if mysql
+                      options[:limit] = 1
+                      options[:unsigned] = true
+                    else
+                      options[:limit] = 2
+                    end
+                    :integer
+                  when UInt16
+                    if mysql
+                      options[:limit] = 2
+                      options[:unsigned] = true
+                    else
+                      options[:limit] = 4
+                    end
+                    :integer
+                  when UInt32
+                    if mysql
+                      options[:limit] = 4
+                      options[:unsigned] = true
+                    else
+                      options[:limit] = 8
+                    end
+                    :integer
+                  when UInt64
+                    if mysql
+                      options[:limit] = 8
+                      options[:unsigned] = true
+                      :integer
+                    else
+                      options[:precision] = 20
+                      options[:scale] = 0
+                      :decimal
+                    end
+                  when String
+                    :text
+                  when Time
+                    :time
+                  else
+                    raise ArgumentError, "column type not supported yet: #{dtype}"
+                  end
+                t.column c, column_type, **options
+              end
+            end
+          end
+
+          quoted_table = connection.quote_table_name(table_name)
+          quoted_columns = columns.map { |c| connection.quote_column_name(c) }
+          rows = cast({Polars::UInt64 => Polars::String}).rows(named: false).map { |row| "(#{row.map { |v| connection.quote(v) }.join(", ")})" }
+          connection.exec_update("INSERT INTO #{quoted_table} (#{quoted_columns.join(", ")}) VALUES #{rows.join(", ")}")
+        end
+      end
+    end
+
     # Write DataFrame as delta table.
     #
     # @param target [Object]
@@ -2424,6 +2596,24 @@ module Polars
     #     - true: -> Always coalesce join columns.
     #     - false: -> Never coalesce join columns.
     #   Note that joining on any other expressions than `col` will turn off coalescing.
+    # @param maintain_order ['none', 'left', 'right', 'left_right', 'right_left']
+    #   Which DataFrame row order to preserve, if any.
+    #   Do not rely on any observed ordering without explicitly
+    #   setting this parameter, as your code may break in a future release.
+    #   Not specifying any ordering can improve performance
+    #   Supported for inner, left, right and full joins
+    #
+    #   * *none*
+    #       No specific ordering is desired. The ordering might differ across
+    #       Polars versions or even between different runs.
+    #   * *left*
+    #       Preserves the order of the left DataFrame.
+    #   * *right*
+    #       Preserves the order of the right DataFrame.
+    #   * *left_right*
+    #       First preserves the order of the left DataFrame, then the right.
+    #   * *right_left*
+    #       First preserves the order of the right DataFrame, then the left.
     #
     # @return [DataFrame]
     #
@@ -2506,7 +2696,8 @@ module Polars
     #   # ╞═════╪═════╪═════╡
     #   # │ 3   ┆ 8.0 ┆ c   │
     #   # └─────┴─────┴─────┘
-    def join(other,
+    def join(
+      other,
       left_on: nil,
       right_on: nil,
       on: nil,
@@ -2514,7 +2705,8 @@ module Polars
       suffix: "_right",
       validate: "m:m",
       join_nulls: false,
-      coalesce: nil
+      coalesce: nil,
+      maintain_order: nil
     )
       lazy
         .join(
@@ -2526,7 +2718,8 @@ module Polars
           suffix: suffix,
           validate: validate,
           join_nulls: join_nulls,
-          coalesce: coalesce
+          coalesce: coalesce,
+          maintain_order: maintain_order
         )
         .collect(no_optimization: true)
     end
@@ -4865,6 +5058,90 @@ module Polars
       iter_rows(named: named, buffer_size: buffer_size, &block)
     end
 
+    # Returns an iterator over the columns of this DataFrame.
+    #
+    # @return [Object]
+    #
+    # @note
+    #   Consider whether you can use `all` instead.
+    #   If you can, it will be more efficient.
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, 3, 5],
+    #       "b" => [2, 4, 6]
+    #     }
+    #   )
+    #   df.iter_columns.map { |s| s.name }
+    #   # => ["a", "b"]
+    #
+    # @example If you're using this to modify a dataframe's columns, e.g.
+    #   # Do NOT do this
+    #   Polars::DataFrame.new(df.iter_columns.map { |column| column * 2 })
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ 2   ┆ 4   │
+    #   # │ 6   ┆ 8   │
+    #   # │ 10  ┆ 12  │
+    #   # └─────┴─────┘
+    #
+    # @example then consider whether you can use `all` instead:
+    #   df.select(Polars.all * 2)
+    #   # =>
+    #   # shape: (3, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ i64 ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ 2   ┆ 4   │
+    #   # │ 6   ┆ 8   │
+    #   # │ 10  ┆ 12  │
+    #   # └─────┴─────┘
+    def iter_columns
+      return to_enum(:iter_columns) unless block_given?
+
+      _df.get_columns.each do |s|
+        yield Utils.wrap_s(s)
+      end
+    end
+
+    # Returns a non-copying iterator of slices over the underlying DataFrame.
+    #
+    # @param n_rows [Integer]
+    #   Determines the number of rows contained in each DataFrame slice.
+    #
+    # @return [Object]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a" => 0...17_500,
+    #       "b" => Date.new(2023, 1, 1),
+    #       "c" => "klmnoopqrstuvwxyz"
+    #     },
+    #     schema_overrides: {"a" => Polars::Int32}
+    #   )
+    #   df.iter_slices.map.with_index do |frame, idx|
+    #     "#{frame.class.name}:[#{idx}]:#{frame.length}"
+    #   end
+    #   # => ["Polars::DataFrame:[0]:10000", "Polars::DataFrame:[1]:7500"]
+    def iter_slices(n_rows: 10_000)
+      return to_enum(:iter_slices, n_rows: n_rows) unless block_given?
+
+      offset = 0
+      while offset < height
+        yield slice(offset, n_rows)
+        offset += n_rows
+      end
+    end
+
     # Shrink DataFrame memory usage.
     #
     # Shrinks to fit the exact capacity needed to hold the data.
@@ -5101,12 +5378,17 @@ module Polars
       lazy.merge_sorted(other.lazy, key).collect(_eager: true)
     end
 
-    # Indicate that one or multiple columns are sorted.
+    # Flag a column as sorted.
+    #
+    # This can speed up future operations.
+    #
+    # @note
+    #   This can lead to incorrect results if the data is NOT sorted! Use with care!
     #
     # @param column [Object]
-    #   Columns that are sorted
+    #   Column that is sorted.
     # @param descending [Boolean]
-    #   Whether the columns are sorted in descending order.
+    #   Whether the column is sorted in descending order.
     #
     # @return [DataFrame]
     def set_sorted(
@@ -5538,5 +5820,21 @@ module Polars
       end
       other
     end
+
+    def with_connection(connection, &block)
+      if !connection.nil?
+        yield connection
+      else
+        ActiveRecord::Base.connection_pool.with_connection(&block)
+      end
+    end
+
+    def maybe_transaction(connection, create_table, &block)
+      if create_table && connection.adapter_name.match?(/postg|sqlite/i) && connection.open_transactions == 0
+        connection.transaction(&block)
+      else
+        yield
+      end
+    end
   end
 end

data/lib/polars/expr.rb

@@ -3994,6 +3994,37 @@ module Polars
       _from_rbexpr(_rbexpr.interpolate(method))
     end
 
+    # Fill null values using interpolation based on another column.
+    #
+    # @param by [Expr] Column to interpolate values based on.
+    #
+    # @return [Expr]
+    #
+    # @example Fill null values using linear interpolation.
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "a" => [1, nil, nil, 3],
+    #       "b" => [1, 2, 7, 8]
+    #     }
+    #   )
+    #   df.with_columns(a_interpolated: Polars.col("a").interpolate_by("b"))
+    #   # =>
+    #   # shape: (4, 3)
+    #   # ┌──────┬─────┬────────────────┐
+    #   # │ a    ┆ b   ┆ a_interpolated │
+    #   # │ ---  ┆ --- ┆ ---            │
+    #   # │ i64  ┆ i64 ┆ f64            │
+    #   # ╞══════╪═════╪════════════════╡
+    #   # │ 1    ┆ 1   ┆ 1.0            │
+    #   # │ null ┆ 2   ┆ 1.285714       │
+    #   # │ null ┆ 7   ┆ 2.714286       │
+    #   # │ 3    ┆ 8   ┆ 3.0            │
+    #   # └──────┴─────┴────────────────┘
+    def interpolate_by(by)
+      by = Utils.parse_into_expression(by)
+      _from_rbexpr(_rbexpr.interpolate_by(by))
+    end
+
     # Apply a rolling min based on another column.
     #
     # @param by [String]

data/lib/polars/functions/eager.rb

@@ -6,8 +6,7 @@ module Polars
     #   DataFrames/Series/LazyFrames to concatenate.
     # @param rechunk [Boolean]
     #   Make sure that all data is in contiguous memory.
-    # @param how ["vertical", "vertical_relaxed", "diagonal", "horizontal"]
-    #   LazyFrames do not support the `horizontal` strategy.
+    # @param how ["vertical", "vertical_relaxed", "diagonal", "diagonal_relaxed", "horizontal"]
     #
     #   - Vertical: applies multiple `vstack` operations.
     #   - Diagonal: finds a union between the column schemas and fills missing column values with null.
@@ -21,7 +20,7 @@ module Polars
     # @example
     #   df1 = Polars::DataFrame.new({"a" => [1], "b" => [3]})
     #   df2 = Polars::DataFrame.new({"a" => [2], "b" => [4]})
-    #   Polars.concat([df1, df2])
+    #   Polars.concat([df1, df2])  # default is 'vertical' strategy
     #   # =>
     #   # shape: (2, 2)
     #   # ┌─────┬─────┐
@@ -32,38 +31,168 @@ module Polars
     #   # │ 1   ┆ 3   │
     #   # │ 2   ┆ 4   │
     #   # └─────┴─────┘
+    #
+    # @example
+    #   df1 = Polars::DataFrame.new({"a" => [1], "b" => [3]})
+    #   df2 = Polars::DataFrame.new({"a" => [2.5], "b" => [4]})
+    #   Polars.concat([df1, df2], how: "vertical_relaxed")  # 'a' coerced into f64
+    #   # =>
+    #   # shape: (2, 2)
+    #   # ┌─────┬─────┐
+    #   # │ a   ┆ b   │
+    #   # │ --- ┆ --- │
+    #   # │ f64 ┆ i64 │
+    #   # ╞═════╪═════╡
+    #   # │ 1.0 ┆ 3   │
+    #   # │ 2.5 ┆ 4   │
+    #   # └─────┴─────┘
+    #
+    # @example
+    #   df_h1 = Polars::DataFrame.new({"l1" => [1, 2], "l2" => [3, 4]})
+    #   df_h2 = Polars::DataFrame.new({"r1" => [5, 6], "r2" => [7, 8], "r3" => [9, 10]})
+    #   Polars.concat([df_h1, df_h2], how: "horizontal")
+    #   # =>
+    #   # shape: (2, 5)
+    #   # ┌─────┬─────┬─────┬─────┬─────┐
+    #   # │ l1  ┆ l2  ┆ r1  ┆ r2  ┆ r3  │
+    #   # │ --- ┆ --- ┆ --- ┆ --- ┆ --- │
+    #   # │ i64 ┆ i64 ┆ i64 ┆ i64 ┆ i64 │
+    #   # ╞═════╪═════╪═════╪═════╪═════╡
+    #   # │ 1   ┆ 3   ┆ 5   ┆ 7   ┆ 9   │
+    #   # │ 2   ┆ 4   ┆ 6   ┆ 8   ┆ 10  │
+    #   # └─────┴─────┴─────┴─────┴─────┘
+    #
+    # @example
+    #   df_d1 = Polars::DataFrame.new({"a" => [1], "b" => [3]})
+    #   df_d2 = Polars::DataFrame.new({"a" => [2], "c" => [4]})
+    #   Polars.concat([df_d1, df_d2], how: "diagonal")
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌─────┬──────┬──────┐
+    #   # │ a   ┆ b    ┆ c    │
+    #   # │ --- ┆ ---  ┆ ---  │
+    #   # │ i64 ┆ i64  ┆ i64  │
+    #   # ╞═════╪══════╪══════╡
+    #   # │ 1   ┆ 3    ┆ null │
+    #   # │ 2   ┆ null ┆ 4    │
+    #   # └─────┴──────┴──────┘
+    #
+    # @example
+    #   df_a1 = Polars::DataFrame.new({"id" => [1, 2], "x" => [3, 4]})
+    #   df_a2 = Polars::DataFrame.new({"id" => [2, 3], "y" => [5, 6]})
+    #   df_a3 = Polars::DataFrame.new({"id" => [1, 3], "z" => [7, 8]})
+    #   Polars.concat([df_a1, df_a2, df_a3], how: "align")
+    #   # =>
+    #   # shape: (3, 4)
+    #   # ┌─────┬──────┬──────┬──────┐
+    #   # │ id  ┆ x    ┆ y    ┆ z    │
+    #   # │ --- ┆ ---  ┆ ---  ┆ ---  │
+    #   # │ i64 ┆ i64  ┆ i64  ┆ i64  │
+    #   # ╞═════╪══════╪══════╪══════╡
+    #   # │ 1   ┆ 3    ┆ null ┆ 7    │
+    #   # │ 2   ┆ 4    ┆ 5    ┆ null │
+    #   # │ 3   ┆ null ┆ 6    ┆ 8    │
+    #   # └─────┴──────┴──────┴──────┘
     def concat(items, rechunk: true, how: "vertical", parallel: true)
-      if items.empty?
+      elems = items.to_a
+
+      if elems.empty?
         raise ArgumentError, "cannot concat empty list"
       end
 
-      first = items[0]
+      if how == "align"
+        if !elems[0].is_a?(DataFrame) && !elems[0].is_a?(LazyFrame)
+          msg = "'align' strategy is not supported for #{elems[0].class.name}"
+          raise TypeError, msg
+        end
+
+        # establish common columns, maintaining the order in which they appear
+        all_columns = elems.flat_map { |e| e.collect_schema.names }
+        key = all_columns.uniq.map.with_index.to_h
+        common_cols = elems.map { |e| e.collect_schema.names }
+          .reduce { |x, y| Set.new(x) & Set.new(y) }
+          .sort_by { |k| key[k] }
+        # we require at least one key column for 'align'
+        if common_cols.empty?
+          msg = "'align' strategy requires at least one common column"
+          raise InvalidOperationError, msg
+        end
+
+        # align the frame data using a full outer join with no suffix-resolution
+        # (so we raise an error in case of column collision, like "horizontal")
+        lf = elems.map { |df| df.lazy }.reduce do |x, y|
+          x.join(
+            y,
+            how: "full",
+            on: common_cols,
+            suffix: "_PL_CONCAT_RIGHT",
+            maintain_order: "right_left"
+          )
+          # Coalesce full outer join columns
+          .with_columns(
+            common_cols.map { |name| F.coalesce([name, "#{name}_PL_CONCAT_RIGHT"]) }
+          )
+          .drop(common_cols.map { |name| "#{name}_PL_CONCAT_RIGHT" })
+        end.sort(common_cols)
+
+        eager = elems[0].is_a?(DataFrame)
+        return eager ? lf.collect : lf
+      end
+
+      first = elems[0]
+
       if first.is_a?(DataFrame)
         if how == "vertical"
-          out = Utils.wrap_df(Plr.concat_df(items))
+          out = Utils.wrap_df(Plr.concat_df(elems))
+        elsif how == "vertical_relaxed"
+          out = Utils.wrap_ldf(
+            Plr.concat_lf(
+              elems.map { |df| df.lazy },
+              rechunk,
+              parallel,
+              true
+            )
+          ).collect(no_optimization: true)
         elsif how == "diagonal"
-          out = Utils.wrap_df(Plr.concat_df_diagonal(items))
+          out = Utils.wrap_df(Plr.concat_df_diagonal(elems))
+        elsif how == "diagonal_relaxed"
+          out = Utils.wrap_ldf(
+            Plr.concat_lf_diagonal(
+              elems.map { |df| df.lazy },
+              rechunk,
+              parallel,
+              true
+            )
+          ).collect(no_optimization: true)
         elsif how == "horizontal"
-          out = Utils.wrap_df(Plr.concat_df_horizontal(items))
+          out = Utils.wrap_df(Plr.concat_df_horizontal(elems))
         else
-          raise ArgumentError, "how must be one of {{'vertical', 'diagonal', 'horizontal'}}, got #{how}"
+          raise ArgumentError, "how must be one of {{'vertical', 'vertical_relaxed', 'diagonal', 'diagonal_relaxed', 'horizontal'}}, got #{how}"
         end
       elsif first.is_a?(LazyFrame)
         if how == "vertical"
-          return Utils.wrap_ldf(Plr.concat_lf(items, rechunk, parallel, false))
+          return Utils.wrap_ldf(Plr.concat_lf(elems, rechunk, parallel, false))
         elsif how == "vertical_relaxed"
-          return Utils.wrap_ldf(Plr.concat_lf(items, rechunk, parallel, true))
+          return Utils.wrap_ldf(Plr.concat_lf(elems, rechunk, parallel, true))
         elsif how == "diagonal"
-          return Utils.wrap_ldf(Plr.concat_lf_diagonal(items, rechunk, parallel, false))
+          return Utils.wrap_ldf(Plr.concat_lf_diagonal(elems, rechunk, parallel, false))
+        elsif how == "diagonal_relaxed"
+          return Utils.wrap_ldf(Plr.concat_lf_diagonal(elems, rechunk, parallel, true))
+        elsif how == "horizontal"
+          return Utils.wrap_ldf(Plr.concat_lf_horizontal(elems, parallel))
         else
-          raise ArgumentError, "Lazy only allows 'vertical', 'vertical_relaxed', and 'diagonal' concat strategy."
+          raise ArgumentError, "Lazy only allows 'vertical', 'vertical_relaxed', 'diagonal', and 'diagonal_relaxed' concat strategy."
         end
       elsif first.is_a?(Series)
-        # TODO
-        out = Utils.wrap_s(Plr.concat_series(items))
+        if how == "vertical"
+          out = Utils.wrap_s(Plr.concat_series(elems))
+        else
+          msg = "Series only supports 'vertical' concat strategy"
+          raise ArgumentError, msg
+        end
       elsif first.is_a?(Expr)
         out = first
-        items[1..-1].each do |e|
+        elems[1..-1].each do |e|
           out = out.append(e)
         end
       else

data/lib/polars/io/database.rb

@@ -51,8 +51,25 @@ module Polars
           when :decimal
             Decimal
           when :float
+            # TODO uncomment in 0.18.0
+            # if column_type.limit && column_type.limit <= 24
+            #   Float32
+            # else
+            #   Float64
+            # end
             Float64
           when :integer
+            # TODO uncomment in 0.18.0
+            # case column_type.limit
+            # when 1
+            #   Int8
+            # when 2
+            #   Int16
+            # when 4
+            #   Int32
+            # else
+            #   Int64
+            # end
             Int64
           when :string, :text
             String