Added support for SQL querying - #33

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.7.0 (unreleased)
 
 - Updated Polars to 0.34.2
+- Added support for SQL querying
 - Added `!` for `Expr`
 - Added `Config` module
 - Added `none?` method to `Series`

ext/polars/Cargo.toml

@@ -82,6 +82,7 @@ features = [
     "semi_anti_join",
     "serde-lazy",
     "sign",
+    "sql",
     "string_encoding",
     "string_from_radix",
     "string_pad",

ext/polars/src/lib.rs

@@ -12,6 +12,7 @@ mod object;
 mod prelude;
 pub(crate) mod rb_modules;
 mod series;
+mod sql;
 mod utils;
 
 use batched_csv::RbBatchedCsv;
@@ -25,6 +26,7 @@ use lazyframe::RbLazyFrame;
 use lazygroupby::RbLazyGroupBy;
 use magnus::{define_module, function, method, prelude::*, Error, Ruby};
 use series::RbSeries;
+use sql::RbSQLContext;
 
 #[cfg(target_os = "linux")]
 use jemallocator::Jemalloc;
@@ -983,5 +985,13 @@ fn init(ruby: &Ruby) -> RbResult<()> {
     let class = module.define_class("RbWhenThen", ruby.class_object())?;
     class.define_method("otherwise", method!(RbThen::overwise, 1))?;
 
+    // sql
+    let class = module.define_class("RbSQLContext", ruby.class_object())?;
+    class.define_singleton_method("new", function!(RbSQLContext::new, 0))?;
+    class.define_method("execute", method!(RbSQLContext::execute, 1))?;
+    class.define_method("get_tables", method!(RbSQLContext::get_tables, 0))?;
+    class.define_method("register", method!(RbSQLContext::register, 2))?;
+    class.define_method("unregister", method!(RbSQLContext::unregister, 1))?;
+
     Ok(())
 }

ext/polars/src/sql.rs

@@ -0,0 +1,46 @@
+use polars::sql::SQLContext;
+use std::cell::RefCell;
+
+use crate::{RbLazyFrame, RbPolarsErr, RbResult};
+
+#[magnus::wrap(class = "Polars::RbSQLContext")]
+#[repr(transparent)]
+#[derive(Clone)]
+pub struct RbSQLContext {
+    pub context: RefCell<SQLContext>,
+}
+
+#[allow(
+    clippy::wrong_self_convention,
+    clippy::should_implement_trait,
+    clippy::len_without_is_empty
+)]
+impl RbSQLContext {
+    #[allow(clippy::new_without_default)]
+    pub fn new() -> RbSQLContext {
+        RbSQLContext {
+            context: SQLContext::new().into(),
+        }
+    }
+
+    pub fn execute(&self, query: String) -> RbResult<RbLazyFrame> {
+        Ok(self
+            .context
+            .borrow_mut()
+            .execute(&query)
+            .map_err(RbPolarsErr::from)?
+            .into())
+    }
+
+    pub fn get_tables(&self) -> RbResult<Vec<String>> {
+        Ok(self.context.borrow().get_tables())
+    }
+
+    pub fn register(&self, name: String, lf: &RbLazyFrame) {
+        self.context.borrow_mut().register(&name, lf.ldf.clone())
+    }
+
+    pub fn unregister(&self, name: String) {
+        self.context.borrow_mut().unregister(&name)
+    }
+}

lib/polars.rb

@@ -42,6 +42,7 @@
 require_relative "polars/rolling_group_by"
 require_relative "polars/series"
 require_relative "polars/slice"
+require_relative "polars/sql_context"
 require_relative "polars/string_expr"
 require_relative "polars/string_name_space"
 require_relative "polars/struct_expr"

lib/polars/sql_context.rb

@@ -0,0 +1,39 @@
+module Polars
+  class SQLContext
+    # @private
+    attr_accessor :_ctxt, :_eager_execution
+
+    def initialize(frames = nil, eager_execution: false, **named_frames)
+      self._ctxt = RbSQLContext.new
+      self._eager_execution = eager_execution
+
+      frames = (frames || {}).to_h
+
+      if frames.any? || named_frames.any?
+        register_many(frames, **named_frames)
+      end
+    end
+
+    def execute(query, eager: nil)
+      res = Utils.wrap_ldf(_ctxt.execute(query))
+      eager || _eager_execution ? res.collect : res
+    end
+
+    def register(name, frame)
+      if frame.is_a?(DataFrame)
+        frame = frame.lazy
+      end
+      _ctxt.register(name.to_s, frame._ldf)
+      self
+    end
+
+    def register_many(frames, **named_frames)
+      frames = (frames || {}).to_h
+      frames = frames.merge(named_frames)
+      frames.each do |name, frame|
+        register(name, frame)
+      end
+      self
+    end
+  end
+end

test/sql_test.rb

@@ -0,0 +1,12 @@
+require_relative "test_helper"
+
+class SqlTest < Minitest::Test
+  def test_works
+    lf = Polars::LazyFrame.new({"a" => [1, 2, 3], "b" => ["x", nil, "z"]})
+    res = Polars::SQLContext.new(frame: lf).execute(
+      "SELECT b, a*2 AS two_a FROM frame WHERE b IS NOT NULL"
+    )
+    expected = Polars::DataFrame.new({"b" => ["x", "z"], "two_a" => [2, 6]})
+    assert_frame expected, res.collect
+  end
+end