Method: Polars::IO#read_database

Defined in:
lib/polars/io/database.rb

#read_database(query, schema_overrides: nil) ⇒ DataFrame

Read a SQL query into a DataFrame.



12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
 
# File 'lib/polars/io/database.rb', line 12

def read_database(query, schema_overrides: nil)
  if !defined?(ActiveRecord)
    raise Error, "Active Record not available"
  end

  result =
    if query.is_a?(ActiveRecord::Result)
      query
    elsif query.is_a?(ActiveRecord::Relation)
      query.connection_pool.with_connection { |c| c.select_all(query.to_sql) }
    elsif query.is_a?(::String)
      ActiveRecord::Base.connection_pool.with_connection { |c| c.select_all(query) }
    else
      raise ArgumentError, "Expected ActiveRecord::Relation, ActiveRecord::Result, or String"
    end

  data = {}
  schema_overrides = (schema_overrides || {}).transform_keys(&:to_s)

  result.columns.each_with_index do |k, i|
    column_type = result.column_types[i]

    data[k] =
      if column_type
        result.rows.map { |r| column_type.deserialize(r[i]) }
      else
        result.rows.map { |r| r[i] }
      end

    polars_type =
      case column_type&.type
      when :binary
        Binary
      when :boolean
        Boolean
      when :date
        Date
      when :datetime, :timestamp
        Datetime
      when :decimal
        Decimal
      when :float
        # TODO uncomment in future release
        # if column_type.limit && column_type.limit <= 24
        #   Float32
        # else
        #   Float64
        # end
        Float64
      when :integer
        # TODO uncomment in future release
        # case column_type.limit
        # when 1
        #   Int8
        # when 2
        #   Int16
        # when 4
        #   Int32
        # else
        #   Int64
        # end
        Int64
      when :string, :text
        String
      when :time
        Time
      # TODO fix issue with null
      # when :json, :jsonb
      #   Struct
      end

    schema_overrides[k] ||= polars_type if polars_type
  end

  DataFrame.new(data, schema_overrides: schema_overrides)
end