|
3 | 3 | require 'duckdb' |
4 | 4 | require 'polars-df' |
5 | 5 |
|
6 | | -df = Polars::DataFrame.new( |
7 | | - { |
8 | | - a: [1, 2, 3], |
9 | | - b: %w[one two three] |
10 | | - } |
11 | | -) |
12 | | - |
13 | | -module DuckDB |
14 | | - class TableFunction |
15 | | - @table_adapters = {} |
16 | | - class << self |
17 | | - def add_table_adapter(klass, adapter) |
18 | | - @table_adapters[klass] = adapter |
19 | | - end |
20 | | - |
21 | | - def table_adapter_for(klass) |
22 | | - @table_adapters[klass] |
23 | | - end |
24 | | - end |
| 6 | +class PolarsDataFrameTableAdapter |
| 7 | + def call(data_frame, name, columns: nil) |
| 8 | + columns ||= infer_columns(data_frame) |
| 9 | + DuckDB::TableFunction.create(name:, columns:, &execute_block(data_frame)) |
25 | 10 | end |
26 | 11 |
|
27 | | - class Connection |
28 | | - def create_table_function(object, name) |
29 | | - adapter = TableFunction.table_adapter_for(object.class) |
30 | | - raise ArgumentError, "No table adapter registered for #{object.class}" if adapter.nil? |
| 12 | + private |
31 | 13 |
|
32 | | - tf = adapter.call(object, name) |
33 | | - register_table_function(tf) |
| 14 | + def execute_block(data_frame) |
| 15 | + counter = 0 |
| 16 | + height = data_frame.height |
| 17 | + width = data_frame.width |
| 18 | + proc do |_func_info, output| |
| 19 | + next counter = 0 if counter >= height |
| 20 | + |
| 21 | + write_row(data_frame, output, counter, width) |
| 22 | + counter += 1 |
| 23 | + 1 |
34 | 24 | end |
35 | 25 | end |
36 | 26 |
|
37 | | - module Polars |
38 | | - module DataFrame |
39 | | - class TableAdapter |
40 | | - def call(df, name) # rubocop:disable Metrics/MethodLength, Naming/MethodParameterName |
41 | | - columns = df.columns.to_h { |header| [header, LogicalType::VARCHAR] } |
42 | | - counter = 0 |
43 | | - height = df.height |
44 | | - width = df.columns.length |
| 27 | + def write_row(data_frame, output, counter, width) |
| 28 | + width.times { |index| output.set_value(index, 0, data_frame[counter, index]) } |
| 29 | + end |
45 | 30 |
|
46 | | - DuckDB::TableFunction.create( |
47 | | - name:, |
48 | | - columns: |
49 | | - ) do |_func_info, output| |
50 | | - if counter < height |
51 | | - width.times do |index| |
52 | | - output.set_value(index, 0, df[counter, index]) |
53 | | - end |
54 | | - counter += 1 |
55 | | - 1 |
56 | | - else |
57 | | - counter = 0 |
58 | | - 0 |
59 | | - end |
60 | | - end |
61 | | - end |
62 | | - end |
63 | | - end |
| 31 | + def infer_columns(data_frame) |
| 32 | + data_frame.columns.to_h { |header| [header, DuckDB::LogicalType::VARCHAR] } |
64 | 33 | end |
65 | | - TableFunction.add_table_adapter(::Polars::DataFrame, DuckDB::Polars::DataFrame::TableAdapter.new) |
66 | 34 | end |
67 | 35 |
|
68 | | -db = DuckDB::Database.open |
| 36 | +DuckDB::TableFunction.add_table_adapter(Polars::DataFrame, PolarsDataFrameTableAdapter.new) |
| 37 | + |
| 38 | +df = Polars::DataFrame.new( |
| 39 | + { |
| 40 | + a: [1, 2, 3], |
| 41 | + b: %w[one two three] |
| 42 | + } |
| 43 | +) |
69 | 44 |
|
| 45 | +db = DuckDB::Database.open |
70 | 46 | con = db.connect |
71 | 47 | con.query('SET threads=1') |
72 | | -con.create_table_function(df, 'polars_df') |
| 48 | +con.expose_as_table(df, 'polars_df') |
73 | 49 | result = con.query('SELECT * FROM polars_df()').to_a |
74 | 50 | p result |
75 | 51 | puts result.first.first == '1' |
|
0 commit comments