Improved Array type

ankane · Mar 3, 2024 · 0b401af · 0b401af
1 parent 11c7b28
commit 0b401af
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 10 deletions.
diff --git a/lib/polars/data_types.rb b/lib/polars/data_types.rb
@@ -381,30 +381,34 @@ def to_s
 
   # Nested list/array type.
   class Array < NestedType
-    attr_reader :width, :inner
+    attr_reader :inner, :width
 
-    def initialize(width, inner = nil)
+    def initialize(inner, width)
       if width.is_a?(DataType) || (width.is_a?(Class) && width < DataType)
         inner, width = width, inner
       end
-      @width = width
       @inner = Utils.rb_type_to_dtype(inner) if inner
+      @width = width
     end
 
-    # TODO check width?
     def ==(other)
       if other.eql?(Array)
         true
       elsif other.is_a?(Array)
-        @inner.nil? || other.inner.nil? || @inner == other.inner
+        if @width != other.width
+          false
+        elsif @inner.nil? || other.inner.nil?
+          true
+        else
+          @inner == other.inner
+        end
       else
         false
       end
     end
 
-    # TODO add width?
     def to_s
-      "#{self.class.name}(#{inner})"
+      "#{self.class.name}(#{inner}, width: #{width.inspect})"
     end
   end
 

diff --git a/lib/polars/series.rb b/lib/polars/series.rb
@@ -4351,7 +4351,7 @@ def sequence_to_rbseries(name, values, dtype: nil, strict: true, dtype_if_empty:
 
       if !dtype.nil? && ![List, Struct, Unknown].include?(dtype) && Utils.is_polars_dtype(dtype) && ruby_dtype.nil?
         if dtype == Array && !dtype.is_a?(Array) && value.is_a?(::Array)
-          dtype = Array.new(value.size)
+          dtype = Array.new(nil, value.size)
         end
 
         constructor = polars_type_to_constructor(dtype)

diff --git a/lib/polars/utils.rb b/lib/polars/utils.rb
@@ -187,7 +187,7 @@ def self.rb_type_to_dtype(data_type)
       begin
         map_rb_type_to_dtype(data_type)
       rescue TypeError
-        raise ArgumentError, "Conversion of Ruby data type #{data_type} to Polars data type not implemented."
+        raise ArgumentError, "Conversion of Ruby data type #{data_type.inspect} to Polars data type not implemented."
       end
     end
 

diff --git a/test/data_types_test.rb b/test/data_types_test.rb
@@ -16,7 +16,7 @@ def test_to_s
     assert_equal %!Polars::Duration(time_unit: "ns")!, Polars::Duration.new("ns").to_s
     assert_equal "Polars::List", Polars::List.to_s
     assert_equal "Polars::List(Polars::Int64)", Polars::List.new(Polars::Int64).to_s
-    assert_equal "Polars::Array(Polars::Int64)", Polars::Array.new(3, Polars::Int64).to_s
+    assert_equal "Polars::Array(Polars::Int64, width: 3)", Polars::Array.new(3, Polars::Int64).to_s
     assert_equal %!Polars::Struct([Polars::Field("a", Polars::Int64)])!, Polars::Struct.new([Polars::Field.new("a", Polars::Int64)]).to_s
   end