polars-df 0.15.0 → 0.26.0

data/ext/polars/src/io/scan_options.rs

@@ -0,0 +1,113 @@
+use std::sync::Arc;
+
+use magnus::{TryConvert, Value, value::ReprValue};
+use polars::prelude::default_values::DefaultFieldValues;
+use polars::prelude::deletion::DeletionFilesList;
+use polars::prelude::{
+    CastColumnsPolicy, CloudScheme, ColumnMapping, ExtraColumnsPolicy, MissingColumnsPolicy,
+    PlSmallStr, Schema, TableStatistics, UnifiedScanArgs,
+};
+use polars_io::{HiveOptions, RowIndex};
+use polars_utils::IdxSize;
+use polars_utils::slice_enum::Slice;
+
+use crate::io::cloud_options::OptRbCloudOptions;
+use crate::prelude::Wrap;
+use crate::{RbDataFrame, RbResult};
+
+/// Interface to `class ScanOptions` on the Ruby side
+pub struct RbScanOptions(Value);
+
+impl TryConvert for RbScanOptions {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        Ok(Self(ob))
+    }
+}
+
+impl TryConvert for Wrap<TableStatistics> {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        let attr: Value = ob.funcall("_df", ())?;
+        Ok(Wrap(TableStatistics(Arc::new(
+            <&RbDataFrame>::try_convert(attr)?.clone().df.into_inner(),
+        ))))
+    }
+}
+
+impl RbScanOptions {
+    pub fn extract_unified_scan_args(
+        &self,
+        // For cloud_options init
+        cloud_scheme: Option<CloudScheme>,
+    ) -> RbResult<UnifiedScanArgs> {
+        let row_index: Option<(Wrap<PlSmallStr>, IdxSize)> = self.0.funcall("row_index", ())?;
+        let pre_slice: Option<(i64, usize)> = self.0.funcall("pre_slice", ())?;
+        let cast_options: Wrap<CastColumnsPolicy> = self.0.funcall("cast_options", ())?;
+        let extra_columns: Wrap<ExtraColumnsPolicy> = self.0.funcall("extra_columns", ())?;
+        let missing_columns: Wrap<MissingColumnsPolicy> = self.0.funcall("missing_columns", ())?;
+        let include_file_paths: Option<Wrap<PlSmallStr>> =
+            self.0.funcall("include_file_paths", ())?;
+        let glob: bool = self.0.funcall("glob", ())?;
+        let hidden_file_prefix: Option<Vec<String>> = self.0.funcall("hidden_file_prefix", ())?;
+        let column_mapping: Option<Wrap<ColumnMapping>> = self.0.funcall("column_mapping", ())?;
+        let default_values: Option<Wrap<DefaultFieldValues>> =
+            self.0.funcall("default_values", ())?;
+        let hive_partitioning: Option<bool> = self.0.funcall("hive_partitioning", ())?;
+        let hive_schema: Option<Wrap<Schema>> = self.0.funcall("hive_schema", ())?;
+        let try_parse_hive_dates: bool = self.0.funcall("try_parse_hive_dates", ())?;
+        let rechunk: bool = self.0.funcall("rechunk", ())?;
+        let cache: bool = self.0.funcall("cache", ())?;
+        let storage_options: OptRbCloudOptions = self.0.funcall("storage_options", ())?;
+        let credential_provider: Option<Value> = self.0.funcall("credential_provider", ())?;
+        let deletion_files: Option<Wrap<DeletionFilesList>> =
+            self.0.funcall("deletion_files", ())?;
+        let table_statistics: Option<Wrap<TableStatistics>> =
+            self.0.funcall("table_statistics", ())?;
+        let row_count: Option<(u64, u64)> = self.0.funcall("row_count", ())?;
+
+        let cloud_options =
+            storage_options.extract_opt_cloud_options(cloud_scheme, credential_provider)?;
+
+        let hive_schema = hive_schema.map(|s| Arc::new(s.0));
+
+        let row_index = row_index.map(|(name, offset)| RowIndex {
+            name: name.0,
+            offset,
+        });
+
+        let hive_options = HiveOptions {
+            enabled: hive_partitioning,
+            hive_start_idx: 0,
+            schema: hive_schema,
+            try_parse_dates: try_parse_hive_dates,
+        };
+
+        let unified_scan_args = UnifiedScanArgs {
+            // Schema is currently still stored inside the options per scan type, but we do eventually
+            // want to put it here instead.
+            schema: None,
+            cloud_options,
+            hive_options,
+            rechunk,
+            cache,
+            glob,
+            hidden_file_prefix: hidden_file_prefix
+                .map(|x| x.into_iter().map(|x| (*x).into()).collect()),
+            projection: None,
+            column_mapping: column_mapping.map(|x| x.0),
+            default_values: default_values
+                .map(|x| x.0)
+                .filter(|DefaultFieldValues::Iceberg(v)| !v.is_empty()),
+            row_index,
+            pre_slice: pre_slice.map(Slice::from),
+            cast_columns_policy: cast_options.0,
+            missing_columns_policy: missing_columns.0,
+            extra_columns_policy: extra_columns.0,
+            include_file_paths: include_file_paths.map(|x| x.0),
+            deletion_files: DeletionFilesList::filter_empty(deletion_files.map(|x| x.0)),
+            table_statistics: table_statistics.map(|x| x.0),
+            row_count,
+        };
+
+        Ok(unified_scan_args)
+    }
+}

data/ext/polars/src/io/sink_options.rs

@@ -0,0 +1,46 @@
+use std::sync::Arc;
+
+use magnus::{TryConvert, Value, value::ReprValue};
+use polars::prelude::sync_on_close::SyncOnCloseType;
+use polars::prelude::{CloudScheme, UnifiedSinkArgs};
+
+use crate::io::cloud_options::OptRbCloudOptions;
+use crate::{RbResult, Wrap};
+
+/// Interface to `class SinkOptions` on the Ruby side
+pub struct RbSinkOptions(Value);
+
+impl TryConvert for RbSinkOptions {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        Ok(Self(ob))
+    }
+}
+
+impl RbSinkOptions {
+    pub fn extract_unified_sink_args(
+        &self,
+        cloud_scheme: Option<CloudScheme>,
+    ) -> RbResult<UnifiedSinkArgs> {
+        let mkdir: bool = self.0.funcall("mkdir", ())?;
+        let maintain_order: bool = self.0.funcall("maintain_order", ())?;
+        let sync_on_close: Option<Wrap<SyncOnCloseType>> = self.0.funcall("sync_on_close", ())?;
+        let storage_options: OptRbCloudOptions = self.0.funcall("storage_options", ())?;
+        let credential_provider: Option<Value> = self.0.funcall("credential_provider", ())?;
+        let sinked_paths_callback: Option<Value> = self.0.funcall("sinked_paths_callback", ())?;
+
+        let cloud_options =
+            storage_options.extract_opt_cloud_options(cloud_scheme, credential_provider)?;
+
+        let sync_on_close = sync_on_close.map_or(SyncOnCloseType::default(), |x| x.0);
+
+        let unified_sink_args = UnifiedSinkArgs {
+            mkdir,
+            maintain_order,
+            sync_on_close,
+            cloud_options: cloud_options.map(Arc::new),
+            sinked_paths_callback: sinked_paths_callback.map(|_| todo!()),
+        };
+
+        Ok(unified_sink_args)
+    }
+}

data/ext/polars/src/io/sink_output.rs

@@ -0,0 +1,21 @@
+use magnus::{TryConvert, Value};
+use polars::prelude::SinkDestination;
+
+use crate::RbResult;
+use crate::prelude::Wrap;
+
+pub struct RbFileSinkDestination(Value);
+
+impl TryConvert for RbFileSinkDestination {
+    fn try_convert(ob: Value) -> RbResult<Self> {
+        Ok(Self(ob))
+    }
+}
+
+impl RbFileSinkDestination {
+    pub fn extract_file_sink_destination(&self) -> RbResult<SinkDestination> {
+        let v = Wrap::<polars_plan::dsl::SinkTarget>::try_convert(self.0)?;
+
+        Ok(SinkDestination::File { target: v.0 })
+    }
+}

data/ext/polars/src/lazyframe/exitable.rs

@@ -0,0 +1,39 @@
+use magnus::Ruby;
+use polars::prelude::*;
+
+use super::RbLazyFrame;
+use crate::utils::EnterPolarsExt;
+use crate::{RbDataFrame, RbResult};
+
+impl RbLazyFrame {
+    pub fn collect_concurrently(ruby: &Ruby, self_: &Self) -> RbResult<RbInProcessQuery> {
+        let ipq = ruby.enter_polars(|| {
+            let ldf = self_.ldf.read().clone();
+            ldf.collect_concurrently()
+        })?;
+        Ok(RbInProcessQuery { ipq })
+    }
+}
+
+#[magnus::wrap(class = "Polars::RbInProcessQuery")]
+#[repr(transparent)]
+#[derive(Clone)]
+pub struct RbInProcessQuery {
+    pub ipq: InProcessQuery,
+}
+
+impl RbInProcessQuery {
+    pub fn cancel(ruby: &Ruby, self_: &Self) -> RbResult<()> {
+        ruby.enter_polars_ok(|| self_.ipq.cancel())
+    }
+
+    pub fn fetch(ruby: &Ruby, self_: &Self) -> RbResult<Option<RbDataFrame>> {
+        let out = ruby.enter_polars(|| self_.ipq.fetch().transpose())?;
+        Ok(out.map(|df| df.into()))
+    }
+
+    pub fn fetch_blocking(ruby: &Ruby, self_: &Self) -> RbResult<RbDataFrame> {
+        let out = ruby.enter_polars(|| self_.ipq.fetch_blocking())?;
+        Ok(out.into())
+    }
+}