polars-df 0.19.0 → 0.21.0

data/LICENSE.txt

@@ -1,4 +1,4 @@
-Copyright (c) 2020 Ritchie Vink
+Copyright (c) 2025 Ritchie Vink
 Copyright (c) 2022-2025 Andrew Kane
 Some portions Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 

data/ext/polars/Cargo.toml

@@ -1,9 +1,9 @@
 [package]
 name = "polars"
-version = "0.19.0"
+version = "0.21.0"
 license = "MIT"
 authors = ["Andrew Kane <andrew@ankane.org>"]
-edition = "2021"
+edition = "2024"
 rust-version = "1.85.0"
 publish = false
 
@@ -12,22 +12,25 @@ crate-type = ["cdylib"]
 
 [dependencies]
 ahash = "0.8"
-arrow = { package = "polars-arrow", version = "=0.48.0" }
+arrow = { package = "polars-arrow", version = "=0.50.0" }
 bytes = "1"
 chrono = "0.4"
 either = "1.8"
-magnus = "0.7"
+magnus = { git = "https://github.com/matsadler/magnus.git", rev = "4125ee2ab5a405d17aa6d7503961c1ecdbe499ef", features = ["chrono"] }
 num-traits = "0.2"
-polars-core = "=0.48.0"
-polars-plan = "=0.48.0"
-polars-parquet = "=0.48.0"
-polars-utils = "=0.48.0"
+polars-core = "=0.50.0"
+polars-dtype = "=0.50.0"
+polars-io = "=0.50.0"
+polars-lazy = { version = "=0.50.0", features = ["catalog"] }
+polars-plan = "=0.50.0"
+polars-parquet = "=0.50.0"
+polars-utils = "=0.50.0"
 rayon = "1.9"
 regex = "1"
 serde_json = "1"
 
 [dependencies.polars]
-version = "=0.48.0"
+version = "=0.50.0"
 features = [
     "abs",
     "approx_unique",
@@ -80,6 +83,7 @@ features = [
     "list_eval",
     "list_gather",
     "list_sample",
+    "list_sets",
     "list_to_struct",
     "log",
     "merge_sorted",

data/ext/polars/src/batched_csv.rs

@@ -2,10 +2,10 @@ use std::cell::RefCell;
 use std::path::PathBuf;
 use std::sync::Mutex;
 
-use magnus::{prelude::*, RArray, Value};
+use magnus::{RArray, Value, prelude::*};
+use polars::io::RowIndex;
 use polars::io::csv::read::OwnedBatchedCsvReader;
 use polars::io::mmap::MmapBytesReader;
-use polars::io::RowIndex;
 use polars::prelude::*;
 
 use crate::conversion::*;

data/ext/polars/src/catalog/mod.rs

@@ -0,0 +1 @@
+pub mod unity;

data/ext/polars/src/catalog/unity.rs

@@ -0,0 +1,450 @@
+use std::str::FromStr;
+
+use magnus::value::{Lazy, ReprValue};
+use magnus::{IntoValue, Module, RArray, RClass, RHash, RModule, Ruby, Value};
+use polars::prelude::{PlHashMap, PlSmallStr, Schema};
+use polars_io::catalog::unity::client::{CatalogClient, CatalogClientBuilder};
+use polars_io::catalog::unity::models::{
+    CatalogInfo, ColumnInfo, DataSourceFormat, NamespaceInfo, TableInfo, TableType,
+};
+use polars_io::catalog::unity::schema::parse_type_json_str;
+use polars_io::pl_async;
+
+use crate::rb_modules::polars;
+use crate::utils::to_rb_err;
+use crate::{RbResult, RbValueError, Wrap};
+
+macro_rules! rbdict_insert_keys {
+    ($dict:expr, {$a:expr}) => {
+        $dict.aset(stringify!($a), $a)?;
+    };
+
+    ($dict:expr, {$a:expr, $($args:expr),+}) => {
+        rbdict_insert_keys!($dict, { $a });
+        rbdict_insert_keys!($dict, { $($args),+ });
+    };
+
+    ($dict:expr, {$a:expr, $($args:expr),+,}) => {
+        rbdict_insert_keys!($dict, {$a, $($args),+});
+    };
+}
+
+static CATALOG_INFO_CLS: Lazy<RClass> = Lazy::new(|_| {
+    polars()
+        .const_get::<_, RClass>("Catalog")
+        .unwrap()
+        .const_get::<_, RModule>("Unity")
+        .unwrap()
+        .const_get("CatalogInfo")
+        .unwrap()
+});
+
+static NAMESPACE_INFO_CLS: Lazy<RClass> = Lazy::new(|_| {
+    polars()
+        .const_get::<_, RClass>("Catalog")
+        .unwrap()
+        .const_get::<_, RModule>("Unity")
+        .unwrap()
+        .const_get("NamespaceInfo")
+        .unwrap()
+});
+
+static TABLE_INFO_CLS: Lazy<RClass> = Lazy::new(|_| {
+    polars()
+        .const_get::<_, RClass>("Catalog")
+        .unwrap()
+        .const_get::<_, RModule>("Unity")
+        .unwrap()
+        .const_get("TableInfo")
+        .unwrap()
+});
+
+static COLUMN_INFO_CLS: Lazy<RClass> = Lazy::new(|_| {
+    polars()
+        .const_get::<_, RClass>("Catalog")
+        .unwrap()
+        .const_get::<_, RModule>("Unity")
+        .unwrap()
+        .const_get("ColumnInfo")
+        .unwrap()
+});
+
+#[magnus::wrap(class = "Polars::RbCatalogClient")]
+pub struct RbCatalogClient(CatalogClient);
+
+impl RbCatalogClient {
+    pub fn new(workspace_url: String, bearer_token: Option<String>) -> RbResult<Self> {
+        let builder = CatalogClientBuilder::new().with_workspace_url(workspace_url);
+
+        let builder = if let Some(bearer_token) = bearer_token {
+            builder.with_bearer_token(bearer_token)
+        } else {
+            builder
+        };
+
+        builder.build().map(RbCatalogClient).map_err(to_rb_err)
+    }
+
+    pub fn list_catalogs(&self) -> RbResult<Value> {
+        let v = pl_async::get_runtime()
+            .block_in_place_on(self.client().list_catalogs())
+            .map_err(to_rb_err)?;
+
+        let mut opt_err = None;
+
+        let out = RArray::from_iter(v.into_iter().map(|x| {
+            let v = catalog_info_to_rbobject(x);
+            if let Ok(v) = v {
+                Some(v)
+            } else {
+                opt_err.replace(v);
+                None
+            }
+        }));
+
+        opt_err.transpose()?;
+
+        Ok(out.into_value())
+    }
+
+    pub fn list_namespaces(&self, catalog_name: String) -> RbResult<Value> {
+        let v = pl_async::get_runtime()
+            .block_in_place_on(self.client().list_namespaces(&catalog_name))
+            .map_err(to_rb_err)?;
+
+        let mut opt_err = None;
+
+        let out = RArray::from_iter(v.into_iter().map(|x| {
+            let v = namespace_info_to_rbobject(x);
+            match v {
+                Ok(v) => Some(v),
+                Err(_) => {
+                    opt_err.replace(v);
+                    None
+                }
+            }
+        }));
+
+        opt_err.transpose()?;
+
+        Ok(out.into_value())
+    }
+
+    pub fn list_tables(&self, catalog_name: String, namespace: String) -> RbResult<Value> {
+        let v = pl_async::get_runtime()
+            .block_in_place_on(self.client().list_tables(&catalog_name, &namespace))
+            .map_err(to_rb_err)?;
+
+        let mut opt_err = None;
+
+        let out = RArray::from_iter(v.into_iter().map(|table_info| {
+            let v = table_info_to_rbobject(table_info);
+
+            if let Ok(v) = v {
+                Some(v)
+            } else {
+                opt_err.replace(v);
+                None
+            }
+        }))
+        .into_value();
+
+        opt_err.transpose()?;
+
+        Ok(out)
+    }
+
+    pub fn get_table_info(
+        &self,
+        table_name: String,
+        catalog_name: String,
+        namespace: String,
+    ) -> RbResult<Value> {
+        let table_info = pl_async::get_runtime()
+            .block_in_place_on(
+                self.client()
+                    .get_table_info(&table_name, &catalog_name, &namespace),
+            )
+            .map_err(to_rb_err)?;
+
+        table_info_to_rbobject(table_info)
+    }
+
+    pub fn create_catalog(
+        &self,
+        catalog_name: String,
+        comment: Option<String>,
+        storage_root: Option<String>,
+    ) -> RbResult<Value> {
+        let catalog_info = pl_async::get_runtime()
+            .block_in_place_on(self.client().create_catalog(
+                &catalog_name,
+                comment.as_deref(),
+                storage_root.as_deref(),
+            ))
+            .map_err(to_rb_err)?;
+
+        catalog_info_to_rbobject(catalog_info)
+    }
+
+    pub fn delete_catalog(&self, catalog_name: String, force: bool) -> RbResult<()> {
+        pl_async::get_runtime()
+            .block_in_place_on(self.client().delete_catalog(&catalog_name, force))
+            .map_err(to_rb_err)
+    }
+
+    pub fn create_namespace(
+        &self,
+        catalog_name: String,
+        namespace: String,
+        comment: Option<String>,
+        storage_root: Option<String>,
+    ) -> RbResult<Value> {
+        let namespace_info = pl_async::get_runtime()
+            .block_in_place_on(self.client().create_namespace(
+                &catalog_name,
+                &namespace,
+                comment.as_deref(),
+                storage_root.as_deref(),
+            ))
+            .map_err(to_rb_err)?;
+
+        namespace_info_to_rbobject(namespace_info)
+    }
+
+    pub fn delete_namespace(
+        &self,
+        catalog_name: String,
+        namespace: String,
+        force: bool,
+    ) -> RbResult<()> {
+        pl_async::get_runtime()
+            .block_in_place_on(
+                self.client()
+                    .delete_namespace(&catalog_name, &namespace, force),
+            )
+            .map_err(to_rb_err)
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn create_table(
+        &self,
+        catalog_name: String,
+        namespace: String,
+        table_name: String,
+        schema: Option<Wrap<Schema>>,
+        table_type: String,
+        data_source_format: Option<String>,
+        comment: Option<String>,
+        storage_root: Option<String>,
+        properties: Vec<(String, String)>,
+    ) -> RbResult<Value> {
+        let table_info = pl_async::get_runtime()
+            .block_in_place_on(
+                self.client().create_table(
+                    &catalog_name,
+                    &namespace,
+                    &table_name,
+                    schema.as_ref().map(|x| &x.0),
+                    &TableType::from_str(&table_type)
+                        .map_err(|e| RbValueError::new_err(e.to_string()))?,
+                    data_source_format
+                        .as_deref()
+                        .map(DataSourceFormat::from_str)
+                        .transpose()
+                        .map_err(|e| RbValueError::new_err(e.to_string()))?
+                        .as_ref(),
+                    comment.as_deref(),
+                    storage_root.as_deref(),
+                    &mut properties.iter().map(|(a, b)| (a.as_str(), b.as_str())),
+                ),
+            )
+            .map_err(to_rb_err)?;
+
+        table_info_to_rbobject(table_info)
+    }
+
+    pub fn delete_table(
+        &self,
+        catalog_name: String,
+        namespace: String,
+        table_name: String,
+    ) -> RbResult<()> {
+        pl_async::get_runtime()
+            .block_in_place_on(
+                self.client()
+                    .delete_table(&catalog_name, &namespace, &table_name),
+            )
+            .map_err(to_rb_err)
+    }
+
+    pub fn type_json_to_polars_type(type_json: String) -> RbResult<Value> {
+        Ok(Wrap(parse_type_json_str(&type_json).map_err(to_rb_err)?).into_value())
+    }
+}
+
+impl RbCatalogClient {
+    fn client(&self) -> &CatalogClient {
+        &self.0
+    }
+}
+
+fn catalog_info_to_rbobject(
+    CatalogInfo {
+        name,
+        comment,
+        storage_location,
+        properties,
+        options,
+        created_at,
+        created_by,
+        updated_at,
+        updated_by,
+    }: CatalogInfo,
+) -> RbResult<Value> {
+    let dict = RHash::new();
+
+    let properties = properties_to_rbobject(properties);
+    let options = properties_to_rbobject(options);
+
+    rbdict_insert_keys!(dict, {
+        name,
+        comment,
+        storage_location,
+        properties,
+        options,
+        created_at,
+        created_by,
+        updated_at,
+        updated_by
+    });
+
+    Ruby::get()
+        .unwrap()
+        .get_inner(&CATALOG_INFO_CLS)
+        .funcall("new", (dict,))
+}
+
+fn namespace_info_to_rbobject(
+    NamespaceInfo {
+        name,
+        comment,
+        properties,
+        storage_location,
+        created_at,
+        created_by,
+        updated_at,
+        updated_by,
+    }: NamespaceInfo,
+) -> RbResult<Value> {
+    let dict = RHash::new();
+
+    let properties = properties_to_rbobject(properties);
+
+    rbdict_insert_keys!(dict, {
+        name,
+        comment,
+        properties,
+        storage_location,
+        created_at,
+        created_by,
+        updated_at,
+        updated_by
+    });
+
+    Ruby::get()
+        .unwrap()
+        .get_inner(&NAMESPACE_INFO_CLS)
+        .funcall("new", (dict,))
+}
+
+fn table_info_to_rbobject(table_info: TableInfo) -> RbResult<Value> {
+    let TableInfo {
+        name,
+        table_id,
+        table_type,
+        comment,
+        storage_location,
+        data_source_format,
+        columns,
+        properties,
+        created_at,
+        created_by,
+        updated_at,
+        updated_by,
+    } = table_info;
+
+    let column_info_cls = Ruby::get().unwrap().get_inner(&COLUMN_INFO_CLS);
+
+    let columns = columns
+        .map(|columns| {
+            Ruby::get()
+                .unwrap()
+                .ary_try_from_iter(columns.into_iter().map(
+                    |ColumnInfo {
+                         name,
+                         type_name,
+                         type_text,
+                         type_json,
+                         position,
+                         comment,
+                         partition_index,
+                     }| {
+                        let dict = RHash::new();
+
+                        let name = name.as_str();
+                        let type_name = type_name.as_str();
+                        let type_text = type_text.as_str();
+
+                        rbdict_insert_keys!(dict, {
+                            name,
+                            type_name,
+                            type_text,
+                            type_json,
+                            position,
+                            comment,
+                            partition_index,
+                        });
+
+                        column_info_cls.funcall::<_, _, Value>("new", (dict,))
+                    },
+                ))
+        })
+        .transpose()?;
+
+    let dict = RHash::new();
+
+    let data_source_format = data_source_format.map(|x| x.to_string());
+    let table_type = table_type.to_string();
+    let properties = properties_to_rbobject(properties);
+
+    rbdict_insert_keys!(dict, {
+        name,
+        comment,
+        table_id,
+        table_type,
+        storage_location,
+        data_source_format,
+        columns,
+        properties,
+        created_at,
+        created_by,
+        updated_at,
+        updated_by,
+    });
+
+    Ruby::get()
+        .unwrap()
+        .get_inner(&TABLE_INFO_CLS)
+        .funcall("new", (dict,))
+}
+
+fn properties_to_rbobject(properties: PlHashMap<PlSmallStr, String>) -> RHash {
+    let dict = RHash::new();
+
+    for (key, value) in properties.into_iter() {
+        dict.aset(key.as_str(), value).unwrap();
+    }
+
+    dict
+}

data/ext/polars/src/conversion/any_value.rs

@@ -1,11 +1,11 @@
 use magnus::encoding::{EncodingCapable, Index};
 use magnus::{
-    class, prelude::*, r_hash::ForEach, IntoValue, RArray, RHash, RString, Ruby, TryConvert, Value,
+    IntoValue, RArray, RHash, RString, Ruby, TryConvert, Value, class, prelude::*, r_hash::ForEach,
 };
 use polars::prelude::*;
 use polars_core::utils::any_values_to_supertype_and_n_dtypes;
 
-use super::{struct_dict, ObjectValue, Wrap};
+use super::{ObjectValue, Wrap, struct_dict};
 
 use crate::exceptions::RbOverflowError;
 use crate::rb_modules::utils;
@@ -40,22 +40,12 @@ pub(crate) fn any_value_into_rb_object(av: AnyValue, ruby: &Ruby) -> Value {
         AnyValue::Boolean(v) => ruby.into_value(v),
         AnyValue::String(v) => ruby.into_value(v),
         AnyValue::StringOwned(v) => ruby.into_value(v.as_str()),
-        AnyValue::Categorical(idx, rev, arr) | AnyValue::Enum(idx, rev, arr) => {
-            let s = if arr.is_null() {
-                rev.get(idx)
-            } else {
-                unsafe { arr.deref_unchecked().value(idx as usize) }
-            };
-            s.into_value()
-        }
-        AnyValue::CategoricalOwned(idx, rev, arr) | AnyValue::EnumOwned(idx, rev, arr) => {
-            let s = if arr.is_null() {
-                rev.get(idx)
-            } else {
-                unsafe { arr.deref_unchecked().value(idx as usize) }
-            };
-            s.into_value()
-        }
+        AnyValue::Categorical(cat, map) | AnyValue::Enum(cat, map) => unsafe {
+            map.cat_to_str_unchecked(cat).into_value()
+        },
+        AnyValue::CategoricalOwned(cat, map) | AnyValue::EnumOwned(cat, map) => unsafe {
+            map.cat_to_str_unchecked(cat).into_value()
+        },
         AnyValue::Date(v) => utils().funcall("_to_ruby_date", (v,)).unwrap(),
         AnyValue::Datetime(v, time_unit, time_zone) => {
             datetime_to_rb_object(v, time_unit, time_zone)
@@ -273,6 +263,6 @@ pub(crate) fn rb_object_to_any_value<'s>(ob: Value, strict: bool) -> RbResult<An
     } else if ob.is_kind_of(crate::rb_modules::bigdecimal()) {
         get_decimal(ob, strict)
     } else {
-        Err(RbPolarsErr::Other(format!("object type not supported {:?}", ob)).into())
+        Err(RbPolarsErr::Other(format!("object type not supported {ob:?}")).into())
     }
 }

data/ext/polars/src/conversion/categorical.rs

@@ -0,0 +1,30 @@
+use std::sync::Arc;
+
+use polars_dtype::categorical::Categories;
+
+#[magnus::wrap(class = "Polars::RbCategories")]
+#[repr(transparent)]
+#[derive(Clone)]
+pub struct RbCategories {
+    categories: Arc<Categories>,
+}
+
+impl RbCategories {
+    pub fn categories(&self) -> &Arc<Categories> {
+        &self.categories
+    }
+}
+
+impl RbCategories {
+    pub fn global_categories() -> Self {
+        Self {
+            categories: Categories::global(),
+        }
+    }
+}
+
+impl From<Arc<Categories>> for RbCategories {
+    fn from(categories: Arc<Categories>) -> Self {
+        Self { categories }
+    }
+}

data/ext/polars/src/conversion/chunked_array.rs

@@ -1,10 +1,10 @@
-use magnus::{prelude::*, IntoValue, RArray, RString, Ruby, TryConvert, Value};
+use magnus::{IntoValue, RArray, RString, Ruby, TryConvert, Value, prelude::*};
 use polars::prelude::*;
 
-use super::{get_rbseq, struct_dict, Wrap};
+use super::{Wrap, get_rbseq, struct_dict};
 
-use crate::rb_modules::utils;
 use crate::RbResult;
+use crate::rb_modules::utils;
 
 impl TryConvert for Wrap<StringChunked> {
     fn try_convert(obj: Value) -> RbResult<Self> {
@@ -75,7 +75,7 @@ impl IntoValue for Wrap<&DurationChunked> {
     fn into_value_with(self, _: &Ruby) -> Value {
         let utils = utils();
         let time_unit = Wrap(self.0.time_unit()).into_value();
-        let iter = self.0.into_iter().map(|opt_v| {
+        let iter = self.0.physical().into_iter().map(|opt_v| {
             opt_v.map(|v| {
                 utils
                     .funcall::<_, _, Value>("_to_ruby_duration", (v, time_unit))
@@ -91,7 +91,7 @@ impl IntoValue for Wrap<&DatetimeChunked> {
         let utils = utils();
         let time_unit = Wrap(self.0.time_unit()).into_value();
         let time_zone = self.0.time_zone().as_deref().map(|v| v.into_value());
-        let iter = self.0.into_iter().map(|opt_v| {
+        let iter = self.0.physical().into_iter().map(|opt_v| {
             opt_v.map(|v| {
                 utils
                     .funcall::<_, _, Value>("_to_ruby_datetime", (v, time_unit, time_zone))
@@ -105,7 +105,7 @@ impl IntoValue for Wrap<&DatetimeChunked> {
 impl IntoValue for Wrap<&TimeChunked> {
     fn into_value_with(self, _: &Ruby) -> Value {
         let utils = utils();
-        let iter = self.0.into_iter().map(|opt_v| {
+        let iter = self.0.physical().into_iter().map(|opt_v| {
             opt_v.map(|v| utils.funcall::<_, _, Value>("_to_ruby_time", (v,)).unwrap())
         });
         RArray::from_iter(iter).into_value()
@@ -115,7 +115,7 @@ impl IntoValue for Wrap<&TimeChunked> {
 impl IntoValue for Wrap<&DateChunked> {
     fn into_value_with(self, _: &Ruby) -> Value {
         let utils = utils();
-        let iter = self.0.into_iter().map(|opt_v| {
+        let iter = self.0.physical().into_iter().map(|opt_v| {
             opt_v.map(|v| utils.funcall::<_, _, Value>("_to_ruby_date", (v,)).unwrap())
         });
         RArray::from_iter(iter).into_value()
@@ -126,7 +126,7 @@ impl IntoValue for Wrap<&DecimalChunked> {
     fn into_value_with(self, _: &Ruby) -> Value {
         let utils = utils();
         let rb_scale = (-(self.0.scale() as i32)).into_value();
-        let iter = self.0.into_iter().map(|opt_v| {
+        let iter = self.0.physical().into_iter().map(|opt_v| {
             opt_v.map(|v| {
                 utils
                     .funcall::<_, _, Value>("_to_ruby_decimal", (v.to_string(), rb_scale))