polars-df 0.23.0 → 0.24.0

data/lib/polars/selectors.rb

@@ -39,7 +39,7 @@ module Polars
     #
     # @return [Selector]
     #
-    # @example
+    # @example Select all columns, casting them to string:
     #   df = Polars::DataFrame.new(
     #     {
     #       "dt" => [Date.new(1999, 12, 31), Date.new(2024, 1, 1)],
@@ -47,8 +47,6 @@ module Polars
     #     },
     #     schema_overrides: {"value" => Polars::Int32}
     #   )
-    #
-    # @example Select all columns, casting them to string:
     #   df.select(Polars.cs.all.cast(Polars::String))
     #   # =>
     #   # shape: (2, 2)
@@ -1036,13 +1034,31 @@ module Polars
 
     # Select all datetime columns, optionally filtering by time unit/zone.
     #
+    # @param time_unit ['ms', 'us', 'ns']
+    #   One (or more) of the allowed timeunit precision strings, "ms", "us", and "ns".
+    #   Omit to select columns with any valid timeunit.
+    # @param time_zone [String]
+    #   * One or more timezone strings, as defined in zoneinfo (to see valid options
+    #     run `import zoneinfo; zoneinfo.available_timezones()` for a full list).
+    #   * Set `nil` to select Datetime columns that do not have a timezone.
+    #   * Set "*" to select Datetime columns that have *any* timezone.
+    #
     # @return [Selector]
-    def self.datetime
-      time_unit = ["ms", "us", "ns"]
+    def self.datetime(time_unit = nil, time_zone: ["*", nil])
+      if time_unit.nil?
+        time_unit_lst = ["ms", "us", "ns"]
+      else
+        time_unit_lst = time_unit.is_a?(::String) ? [time_unit] : time_unit.to_a
+      end
 
-      time_zone = [nil]
+      if time_zone.nil?
+        time_zone_lst = [nil]
+      elsif time_zone
+        # TODO improve
+        time_zone_lst = time_zone.to_a
+      end
 
-      Selector._from_rbselector(RbSelector.datetime(time_unit, time_zone))
+      Selector._from_rbselector(RbSelector.datetime(time_unit_lst, time_zone_lst))
     end
 
     # Select all decimal columns.
@@ -1180,9 +1196,17 @@ module Polars
 
     # Select all duration columns, optionally filtering by time unit.
     #
+    # @param time_unit ['ms', 'us', 'ns']
+    #   One (or more) of the allowed timeunit precision strings, "ms", "us", and "ns".
+    #   Omit to select columns with any valid timeunit.
+    #
     # @return [Selector]
-    def self.duration
-      time_unit = ["ms", "us", "ns"]
+    def self.duration(time_unit = nil)
+      if time_unit.nil?
+        time_unit = ["ms", "us", "ns"]
+      else
+        time_unit = time_unit.is_a?(::String) ? [time_unit] : time_unit.to_a
+      end
 
       Selector._from_rbselector(RbSelector.duration(time_unit))
     end
@@ -1693,6 +1717,18 @@ module Polars
     # Select all object columns.
     #
     # @return [Selector]
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "idx" => [0, 1],
+    #       "uuid_obj" => ["6be063cf-c9c6-43be-878e-e446cfd42981", "7849d8f9-2cac-48e7-96d3-63cf81c14869"],
+    #       "uuid_str" => ["acab9fea-c05d-4b91-b639-418004a63f33", "28c65415-8b7d-4857-a4ce-300dca14b12b"]
+    #     },
+    #     schema_overrides: {"idx" => Polars::Int32, "uuid_obj" => Polars::Object}
+    #   )
+    #   df.select(Polars.cs.object).to_h(as_series: false)
+    #   # => {"uuid_obj"=>["6be063cf-c9c6-43be-878e-e446cfd42981", "7849d8f9-2cac-48e7-96d3-63cf81c14869"]}
     def self.object
       Selector._from_rbselector(RbSelector.object)
     end