wip -- fix things

Author: TomAugspurger

python/cudf_polars/cudf_polars/dsl/ir.py

@@ -460,6 +460,7 @@ def __init__(
             schema,
             typ,
             reader_options,
+            cloud_options,
             paths,
             with_columns,
             skip_rows,
@@ -928,7 +929,14 @@ def __init__(
         self.options = options
         self.cloud_options = cloud_options
         self.children = (df,)
-        self._non_child_args = (schema, kind, path, parquet_options, options)
+        self._non_child_args = (
+            schema,
+            kind,
+            path,
+            parquet_options,
+            options,
+            cloud_options,
+        )
         if self.cloud_options is not None and any(
             self.cloud_options.get(k) is not None
             for k in ("config", "credential_provider")
@@ -1209,7 +1217,7 @@ def __init__(self, schema: Schema, key: int, refcount: int | None, value: IR):
         self.key = key
         self.refcount = refcount
         self.children = (value,)
-        self._non_child_args = (key, refcount)
+        self._non_child_args = (schema, key, refcount)
 
     def get_hashable(self) -> Hashable:  # noqa: D102
         # Polars arranges that the keys are unique across all cache
@@ -1355,7 +1363,7 @@ def __init__(
         self.exprs = tuple(exprs)
         self.should_broadcast = should_broadcast
         self.children = (df,)
-        self._non_child_args = (self.exprs, should_broadcast)
+        self._non_child_args = (schema, self.exprs, should_broadcast)
         if (
             Select._is_len_expr(self.exprs)
             and isinstance(df, Scan)
@@ -1465,7 +1473,7 @@ def __init__(
         self.schema = schema
         self.exprs = tuple(exprs)
         self.children = (df,)
-        self._non_child_args = (self.exprs,)
+        self._non_child_args = (schema, self.exprs)
 
     @classmethod
     @log_do_evaluate
@@ -2009,7 +2017,7 @@ def __init__(
         assert not nulls_equal
         assert not coalesce
         assert maintain_order == "none"
-        self._non_child_args = (predicate_wrapper, options)
+        self._non_child_args = (schema, predicate_wrapper, options)
 
     @classmethod
     @log_do_evaluate
@@ -2129,7 +2137,7 @@ def __init__(
         self.right_on = tuple(right_on)
         self.options = options
         self.children = (left, right)
-        self._non_child_args = (self.left_on, self.right_on, self.options)
+        self._non_child_args = (schema, self.left_on, self.right_on, self.options)
 
     @staticmethod
     @cache
@@ -2500,7 +2508,7 @@ def __init__(
         self.schema = schema
         self.columns = tuple(columns)
         self.should_broadcast = should_broadcast
-        self._non_child_args = (self.columns, self.should_broadcast)
+        self._non_child_args = (schema, self.columns, self.should_broadcast)
         self.children = (df,)
 
     @classmethod
@@ -2564,7 +2572,7 @@ def __init__(
         self.subset = subset
         self.zlice = zlice
         self.stable = stable
-        self._non_child_args = (keep, subset, zlice, stable)
+        self._non_child_args = (schema, keep, subset, zlice, stable)
         self.children = (df,)
 
     _KEEP_MAP: ClassVar[dict[str, plc.stream_compaction.DuplicateKeepOption]] = {
@@ -2663,6 +2671,7 @@ def __init__(
         self.stable = stable
         self.zlice = zlice
         self._non_child_args = (
+            schema,
             self.by,
             self.order,
             self.null_order,
@@ -2727,7 +2736,7 @@ def __init__(self, schema: Schema, offset: int, length: int | None, df: IR):
         self.schema = schema
         self.offset = offset
         self.length = length
-        self._non_child_args = (offset, length)
+        self._non_child_args = (schema, offset, length)
         self.children = (df,)
 
     @classmethod
@@ -2757,7 +2766,7 @@ class Filter(IR):
     def __init__(self, schema: Schema, mask: expr.NamedExpr, df: IR):
         self.schema = schema
         self.mask = mask
-        self._non_child_args = (mask,)
+        self._non_child_args = (schema, mask)
         self.children = (df,)
 
     @classmethod
@@ -2825,7 +2834,7 @@ def __init__(self, schema: Schema, key: str, left: IR, right: IR):
         self.schema = schema
         self.key = key
         self.children = (left, right)
-        self._non_child_args = (key,)
+        self._non_child_args = (schema, key)
 
     @classmethod
     @log_do_evaluate
@@ -3075,7 +3084,7 @@ class Union(IR):
     def __init__(self, schema: Schema, zlice: Zlice | None, *children: IR):
         self.schema = schema
         self.zlice = zlice
-        self._non_child_args = (zlice,)
+        self._non_child_args = (schema, zlice)
         self.children = children
         schema = self.children[0].schema
 
@@ -3126,7 +3135,7 @@ def __init__(
     ):
         self.schema = schema
         self.should_broadcast = should_broadcast
-        self._non_child_args = (should_broadcast,)
+        self._non_child_args = (schema, should_broadcast)
         self.children = children
 
     @staticmethod

python/cudf_polars/cudf_polars/experimental/io.py

@@ -417,13 +417,14 @@ def _sink_to_directory(
     path: str,
     parquet_options: ParquetOptions,
     options: dict[str, Any],
+    cloud_options: dict[str, Any] | None,
     df: DataFrame,
     ready: None,
     context: IRExecutionContext,
 ) -> DataFrame:
     """Sink a partition to a new file."""
     return Sink.do_evaluate(
-        schema, kind, path, parquet_options, options, df, context=context
+        schema, kind, path, parquet_options, options, cloud_options, df, context=context
     )
 
 
@@ -566,6 +567,7 @@ def _directory_sink_graph(
             f"{sink.path}/part.{str(i).zfill(width)}.{suffix}",
             sink.parquet_options,
             sink.options,
+            sink.cloud_options,
             (child_name, i),
             setup_name,
             context,

python/cudf_polars/cudf_polars/experimental/join.py

@@ -144,6 +144,7 @@ def _make_bcast_join(
     left: IR,
     right: IR,
     shuffle_method: ShuffleMethod,
+    *,
     streaming_runtime: str,
 ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]:
     if ir.options[0] != "Inner":
@@ -290,7 +291,7 @@ def _(
             left,
             right,
             config_options.executor.shuffle_method,
-            config_options.executor.runtime,
+            streaming_runtime=config_options.executor.runtime,
         )
     else:
         # Create a hash join
@@ -391,6 +392,7 @@ def _(
                 inter_key = (inter_name, part_out, j)
                 graph[(inter_name, part_out, j)] = (
                     partial(ir.do_evaluate, context=context),
+                    ir.schema,
                     ir.left_on,
                     ir.right_on,
                     ir.options,

python/cudf_polars/cudf_polars/experimental/utils.py

@@ -28,7 +28,11 @@
 
 def _concat(*dfs: DataFrame, context: IRExecutionContext) -> DataFrame:
     # Concatenate a sequence of DataFrames vertically
-    return dfs[0] if len(dfs) == 1 else Union.do_evaluate(None, *dfs, context=context)
+    return (
+        dfs[0]
+        if len(dfs) == 1
+        else Union.do_evaluate(None, None, *dfs, context=context)
+    )
 
 
 def _fallback_inform(msg: str, config_options: ConfigOptions) -> None: