diff --git a/Gemfile b/Gemfile index 0368652841..8727efecf4 100644 --- a/Gemfile +++ b/Gemfile @@ -6,9 +6,14 @@ gem "rake" gem "rake-compiler" gem "minitest" gem "activerecord" -gem "sqlite3" gem "numo-narray" gem "vega" +if ENV["ADAPTER"] == "postgresql" + gem "pg" +else + gem "sqlite3" +end + # https://github.com/lsegal/yard/issues/1321 gem "yard", require: false diff --git a/lib/polars/data_frame.rb b/lib/polars/data_frame.rb index fd1b548f1b..66449e2e4c 100644 --- a/lib/polars/data_frame.rb +++ b/lib/polars/data_frame.rb @@ -20,7 +20,6 @@ class DataFrame # this does not yield conclusive results, column orientation is used. def initialize(data = nil, schema: nil, columns: nil, schema_overrides: nil, orient: nil, infer_schema_length: 100, nan_to_null: false) schema ||= columns - raise Todo if schema_overrides # TODO deprecate in favor of read_sql if defined?(ActiveRecord) && (data.is_a?(ActiveRecord::Relation) || data.is_a?(ActiveRecord::Result)) @@ -4946,7 +4945,7 @@ def self._unpack_schema(schema, schema_overrides: nil, n_expected: nil, lookup_n end if schema_overrides && schema_overrides.any? - raise Todo + column_dtypes.merge!(schema_overrides) end column_dtypes.each do |col, dtype| diff --git a/lib/polars/io.rb b/lib/polars/io.rb index 27dbb24879..369d5b0c81 100644 --- a/lib/polars/io.rb +++ b/lib/polars/io.rb @@ -622,10 +622,22 @@ def read_database(query) raise ArgumentError, "Expected ActiveRecord::Relation, ActiveRecord::Result, or String" end data = {} + schema_overrides = {} result.columns.each_with_index do |k, i| data[k] = result.rows.map { |r| r[i] } + column_type = result.column_types[i]&.type + polars_type = + case column_type + when :datetime + Datetime + when :string + Utf8 + when :integer + Int64 + end + schema_overrides[k] = polars_type if polars_type end - DataFrame.new(data) + DataFrame.new(data, schema_overrides: schema_overrides) end alias_method :read_sql, :read_database diff --git a/test/database_test.rb b/test/database_test.rb index ea23c798d1..dc613d3efb 100644 --- a/test/database_test.rb +++ b/test/database_test.rb @@ -8,7 +8,7 @@ def setup def test_relation users = create_users df = Polars::DataFrame.new(User.order(:id)) - assert_equal ["id", "name", "number"], df.columns + assert_equal ["id", "name", "number", "joined_at"], df.columns assert_series users.map(&:id), df["id"] assert_series users.map(&:name), df["name"] assert_equal Polars::Int64, df["number"].dtype @@ -17,7 +17,7 @@ def test_relation def test_result users = create_users df = Polars::DataFrame.new(User.connection.select_all("SELECT * FROM users ORDER BY id")) - assert_equal ["id", "name", "number"], df.columns + assert_equal ["id", "name", "number", "joined_at"], df.columns assert_series users.map(&:id), df["id"] assert_series users.map(&:name), df["name"] assert_equal Polars::Int64, df["number"].dtype @@ -26,7 +26,7 @@ def test_result def test_read_database_relation users = create_users df = Polars.read_database(User.order(:id)) - assert_equal ["id", "name", "number"], df.columns + assert_equal ["id", "name", "number", "joined_at"], df.columns assert_series users.map(&:id), df["id"] assert_series users.map(&:name), df["name"] assert_equal Polars::Int64, df["number"].dtype @@ -35,7 +35,7 @@ def test_read_database_relation def test_read_database_result users = create_users df = Polars.read_database(User.connection.select_all("SELECT * FROM users ORDER BY id")) - assert_equal ["id", "name", "number"], df.columns + assert_equal ["id", "name", "number", "joined_at"], df.columns assert_series users.map(&:id), df["id"] assert_series users.map(&:name), df["name"] assert_equal Polars::Int64, df["number"].dtype @@ -44,12 +44,20 @@ def test_read_database_result def test_read_database_string users = create_users df = Polars.read_database("SELECT * FROM users ORDER BY id") - assert_equal ["id", "name", "number"], df.columns + assert_equal ["id", "name", "number", "joined_at"], df.columns assert_series users.map(&:id), df["id"] assert_series users.map(&:name), df["name"] assert_equal Polars::Int64, df["number"].dtype end + def test_read_database_null + skip unless ENV["ADAPTER"] == "postgresql" + + User.create! + df = Polars.read_database("SELECT * FROM users ORDER BY id") + assert_equal [Polars::Int64, Polars::Utf8, Polars::Int64, Polars::Datetime], df.dtypes + end + def test_read_database_unsupported error = assert_raises(ArgumentError) do Polars.read_database(Object.new) diff --git a/test/test_helper.rb b/test/test_helper.rb index eee4068762..f520cad9a8 100644 --- a/test/test_helper.rb +++ b/test/test_helper.rb @@ -9,12 +9,17 @@ ActiveRecord::Base.logger = logger ActiveRecord::Migration.verbose = ENV["VERBOSE"] -ActiveRecord::Base.establish_connection adapter: "sqlite3", database: ":memory:" +if ENV["ADAPTER"] == "postgresql" + ActiveRecord::Base.establish_connection adapter: "postgresql", database: "polars_ruby_test" +else + ActiveRecord::Base.establish_connection adapter: "sqlite3", database: ":memory:" +end ActiveRecord::Schema.define do - create_table :users do |t| + create_table :users, force: true do |t| t.string :name t.integer :number + t.datetime :joined_at end end