polars-df 0.10.0-x86_64-linux-musl

data/lib/polars/slice.rb

@@ -0,0 +1,104 @@
+module Polars
+  # @private
+  class Slice
+    def initialize(obj)
+      @obj = obj
+    end
+
+    # Apply a slice operation, taking advantage of any potential fast paths.
+    def apply(s)
+      # normalize slice
+      _slice_setup(s)
+
+      # check for fast-paths / single-operation calls
+      if @slice_length == 0
+        @obj.cleared
+      elsif @is_unbounded && [-1, 1].include?(@stride)
+        @stride < 0 ? @obj.reverse : @obj.clone
+      elsif @start >= 0 && @stop >= 0 && @stride == 1
+        @obj.slice(@start, @slice_length)
+      elsif @stride < 0 && @slice_length == 1
+        @obj.slice(@stop + 1, 1)
+      else
+        # multi-operation calls; make lazy
+        lazyobj = _lazify(@obj)
+        sliced = @stride > 0 ? _slice_positive(lazyobj) : _slice_negative(lazyobj)
+        _as_original(sliced, @obj)
+      end
+    end
+
+    private
+
+    # Return lazy variant back to its original type.
+    def _as_original(lazy, original)
+      frame = lazy.collect
+      original.is_a?(DataFrame) ? frame : frame.to_series
+    end
+
+    # Make lazy to ensure efficient/consistent handling.
+    def _lazify(obj)
+      obj.is_a?(DataFrame) ? obj.lazy : obj.to_frame.lazy
+    end
+
+    # Logic for slices with positive stride.
+    def _slice_positive(obj)
+      # note: at this point stride is guaranteed to be > 1
+      obj.slice(@start, @slice_length).take_every(@stride)
+    end
+
+    # Logic for slices with negative stride.
+    def _slice_negative(obj)
+      stride = @stride.abs
+      lazyslice = obj.slice(@stop + 1, @slice_length).reverse
+      stride > 1 ? lazyslice.take_every(stride) : lazyslice
+    end
+
+    # Normalize slice bounds, identify unbounded and/or zero-length slices.
+    def _slice_setup(s)
+      # can normalize slice indices as we know object size
+      obj_len = @obj.length
+      start = if s.begin
+        if s.begin < 0
+          [s.begin + obj_len, 0].max
+        else
+          s.begin
+        end
+      else
+        0
+      end
+      stop = if s.end
+        if s.end < 0
+          s.end + (s.exclude_end? ? 0 : 1) + obj_len
+        else
+          s.end + (s.exclude_end? ? 0 : 1)
+        end
+      else
+        obj_len
+      end
+      stride = 1
+
+      # check if slice is actually unbounded
+      if stride >= 1
+        @is_unbounded = start <= 0 && stop >= obj_len
+      else
+        @is_unbounded = stop == -1 && start >= obj_len - 1
+      end
+
+      # determine slice length
+      if @obj.is_empty
+        @slice_length = 0
+      elsif @is_unbounded
+        @slice_length = obj_len
+      else
+        @slice_length = if start == stop || (stride > 0 && start > stop) || (stride < 0 && start < stop)
+          0
+        else
+          (stop - start).abs
+        end
+      end
+      @start = start
+      @stop = stop
+      @stride = stride
+    end
+  end
+end

data/lib/polars/sql_context.rb

@@ -0,0 +1,194 @@
+module Polars
+  # Run SQL queries against DataFrame/LazyFrame data.
+  class SQLContext
+    # @private
+    attr_accessor :_ctxt, :_eager_execution
+
+    # Initialize a new `SQLContext`.
+    def initialize(frames = nil, eager_execution: false, **named_frames)
+      self._ctxt = RbSQLContext.new
+      self._eager_execution = eager_execution
+
+      frames = (frames || {}).to_h
+
+      if frames.any? || named_frames.any?
+        register_many(frames, **named_frames)
+      end
+    end
+
+    # Parse the given SQL query and execute it against the registered frame data.
+    #
+    # @param query [String]
+    #   A valid string SQL query.
+    # @param eager [Boolean]
+    #   Apply the query eagerly, returning `DataFrame` instead of `LazyFrame`.
+    #   If unset, the value of the init-time parameter "eager_execution" will be
+    #   used. (Note that the query itself is always executed in lazy-mode; this
+    #   parameter only impacts the type of the returned frame).
+    #
+    # @return [Object]
+    #
+    # @example Execute a SQL query against the registered frame data:
+    #   df = Polars::DataFrame.new(
+    #     [
+    #       ["The Godfather", 1972, 6_000_000, 134_821_952, 9.2],
+    #       ["The Dark Knight", 2008, 185_000_000, 533_316_061, 9.0],
+    #       ["Schindler's List", 1993, 22_000_000, 96_067_179, 8.9],
+    #       ["Pulp Fiction", 1994, 8_000_000, 107_930_000, 8.9],
+    #       ["The Shawshank Redemption", 1994, 25_000_000, 28_341_469, 9.3],
+    #     ],
+    #     schema: ["title", "release_year", "budget", "gross", "imdb_score"]
+    #   )
+    #   ctx = Polars::SQLContext.new(films: df)
+    #   ctx.execute(
+    #     "
+    #     SELECT title, release_year, imdb_score
+    #     FROM films
+    #     WHERE release_year > 1990
+    #     ORDER BY imdb_score DESC
+    #     ",
+    #     eager: true
+    #   )
+    #   # =>
+    #   # shape: (4, 3)
+    #   # ┌──────────────────────────┬──────────────┬────────────┐
+    #   # │ title                    ┆ release_year ┆ imdb_score │
+    #   # │ ---                      ┆ ---          ┆ ---        │
+    #   # │ str                      ┆ i64          ┆ f64        │
+    #   # ╞══════════════════════════╪══════════════╪════════════╡
+    #   # │ The Shawshank Redemption ┆ 1994         ┆ 9.3        │
+    #   # │ The Dark Knight          ┆ 2008         ┆ 9.0        │
+    #   # │ Schindler's List         ┆ 1993         ┆ 8.9        │
+    #   # │ Pulp Fiction             ┆ 1994         ┆ 8.9        │
+    #   # └──────────────────────────┴──────────────┴────────────┘
+    #
+    # @example Execute a GROUP BY query:
+    #   ctx.execute(
+    #     "
+    #     SELECT
+    #         MAX(release_year / 10) * 10 AS decade,
+    #         SUM(gross) AS total_gross,
+    #         COUNT(title) AS n_films,
+    #     FROM films
+    #     GROUP BY (release_year / 10) -- decade
+    #     ORDER BY total_gross DESC
+    #     ",
+    #     eager: true
+    #   )
+    #   # =>
+    #   # shape: (3, 3)
+    #   # ┌────────┬─────────────┬─────────┐
+    #   # │ decade ┆ total_gross ┆ n_films │
+    #   # │ ---    ┆ ---         ┆ ---     │
+    #   # │ i64    ┆ i64         ┆ u32     │
+    #   # ╞════════╪═════════════╪═════════╡
+    #   # │ 2000   ┆ 533316061   ┆ 1       │
+    #   # │ 1990   ┆ 232338648   ┆ 3       │
+    #   # │ 1970   ┆ 134821952   ┆ 1       │
+    #   # └────────┴─────────────┴─────────┘
+    def execute(query, eager: nil)
+      res = Utils.wrap_ldf(_ctxt.execute(query))
+      eager || _eager_execution ? res.collect : res
+    end
+
+    # Register a single frame as a table, using the given name.
+    #
+    # @param name [String]
+    #   Name of the table.
+    # @param frame [Object]
+    #   eager/lazy frame to associate with this table name.
+    #
+    # @return [SQLContext]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"hello" => ["world"]})
+    #   ctx = Polars::SQLContext.new
+    #   ctx.register("frame_data", df).execute("SELECT * FROM frame_data").collect
+    #   # =>
+    #   # shape: (1, 1)
+    #   # ┌───────┐
+    #   # │ hello │
+    #   # │ ---   │
+    #   # │ str   │
+    #   # ╞═══════╡
+    #   # │ world │
+    #   # └───────┘
+    def register(name, frame)
+      if frame.is_a?(DataFrame)
+        frame = frame.lazy
+      end
+      _ctxt.register(name.to_s, frame._ldf)
+      self
+    end
+
+    # Register multiple eager/lazy frames as tables, using the associated names.
+    #
+    # @param frames [Hash]
+    #   A `{name:frame, ...}` mapping.
+    # @param named_frames [Object]
+    #   Named eager/lazy frames, provided as kwargs.
+    #
+    # @return [SQLContext]
+    def register_many(frames, **named_frames)
+      frames = (frames || {}).to_h
+      frames = frames.merge(named_frames)
+      frames.each do |name, frame|
+        register(name, frame)
+      end
+      self
+    end
+
+    # Unregister one or more eager/lazy frames by name.
+    #
+    # @param names [Object]
+    #   Names of the tables to unregister.
+    #
+    # @return [SQLContext]
+    #
+    # @example Register with a SQLContext object:
+    #   df0 = Polars::DataFrame.new({"ints" => [9, 8, 7, 6, 5]})
+    #   lf1 = Polars::LazyFrame.new({"text" => ["a", "b", "c"]})
+    #   lf2 = Polars::LazyFrame.new({"misc" => ["testing1234"]})
+    #   ctx = Polars::SQLContext.new(test1: df0, test2: lf1, test3: lf2)
+    #   ctx.tables
+    #   # => ["test1", "test2", "test3"]
+    #
+    # @example Unregister one or more of the tables:
+    #   ctx.unregister(["test1", "test3"]).tables
+    #   # => ["test2"]
+    def unregister(names)
+      if names.is_a?(::String)
+        names = [names]
+      end
+      names.each do |nm|
+        _ctxt.unregister(nm)
+      end
+      self
+    end
+
+    # Return a list of the registered table names.
+    #
+    # @return [Array]
+    #
+    # @example Executing as SQL:
+    #   frame_data = Polars::DataFrame.new({"hello" => ["world"]})
+    #   ctx = Polars::SQLContext.new(hello_world: frame_data)
+    #   ctx.execute("SHOW TABLES", eager: true)
+    #   # =>
+    #   # shape: (1, 1)
+    #   # ┌─────────────┐
+    #   # │ name        │
+    #   # │ ---         │
+    #   # │ str         │
+    #   # ╞═════════════╡
+    #   # │ hello_world │
+    #   # └─────────────┘
+    #
+    # @example Calling the method:
+    #   ctx.tables
+    #   # => ["hello_world"]
+    def tables
+      _ctxt.get_tables.sort
+    end
+  end
+end

data/lib/polars/string_cache.rb

@@ -0,0 +1,75 @@
+module Polars
+  # Context manager for enabling and disabling the global string cache.
+  class StringCache
+    def initialize(&block)
+      RbStringCacheHolder.hold(&block)
+    end
+  end
+
+  module Functions
+    # Enable the global string cache.
+    #
+    # `Categorical` columns created under the same global string cache have
+    # the same underlying physical value when string values are equal. This allows the
+    # columns to be concatenated or used in a join operation, for example.
+    #
+    # @return [nil]
+    #
+    # @example Construct two Series using the same global string cache.
+    #   Polars.enable_string_cache
+    #   s1 = Polars::Series.new("color", ["red", "green", "red"], dtype: Polars::Categorical)
+    #   s2 = Polars::Series.new("color", ["blue", "red", "green"], dtype: Polars::Categorical)
+    #   Polars.disable_string_cache
+    #
+    # @example As both Series are constructed under the same global string cache, they can be concatenated.
+    #   Polars.concat([s1, s2])
+    #   # =>
+    #   # shape: (6,)
+    #   # Series: 'color' [cat]
+    #   # [
+    #   #         "red"
+    #   #         "green"
+    #   #         "red"
+    #   #         "blue"
+    #   #         "red"
+    #   #         "green"
+    #   # ]
+    def enable_string_cache
+      Plr.enable_string_cache
+    end
+
+    # Disable and clear the global string cache.
+    #
+    # @return [nil]
+    #
+    # @example Construct two Series using the same global string cache.
+    #   Polars.enable_string_cache
+    #   s1 = Polars::Series.new("color", ["red", "green", "red"], dtype: Polars::Categorical)
+    #   s2 = Polars::Series.new("color", ["blue", "red", "green"], dtype: Polars::Categorical)
+    #   Polars.disable_string_cache
+    #
+    # @example As both Series are constructed under the same global string cache, they can be concatenated.
+    #   Polars.concat([s1, s2])
+    #   # =>
+    #   # shape: (6,)
+    #   # Series: 'color' [cat]
+    #   # [
+    #   #         "red"
+    #   #         "green"
+    #   #         "red"
+    #   #         "blue"
+    #   #         "red"
+    #   #         "green"
+    #   # ]
+    def disable_string_cache
+      Plr.disable_string_cache
+    end
+
+    # Check whether the global string cache is enabled.
+    #
+    # @return [Boolean]
+    def using_string_cache
+      Plr.using_string_cache
+    end
+  end
+end