diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index df4a495cd93..935be5cbbf2 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -67,7 +67,7 @@ dependencies: - packaging - pandas>=3.0.0,<3.1.0 - pandoc -- polars>=1.35,<1.39 +- polars>=1.35,<1.40 - pre-commit - pyarrow>=19.0.0,<24 - pydata-sphinx-theme>=0.15.4 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index e95fc711b1e..20d28e469e8 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -67,7 +67,7 @@ dependencies: - packaging - pandas>=3.0.0,<3.1.0 - pandoc -- polars>=1.35,<1.39 +- polars>=1.35,<1.40 - pre-commit - pyarrow>=19.0.0,<24 - pydata-sphinx-theme>=0.15.4 diff --git a/conda/environments/all_cuda-132_arch-aarch64.yaml b/conda/environments/all_cuda-132_arch-aarch64.yaml index 74af84c5666..0b7610a49e7 100644 --- a/conda/environments/all_cuda-132_arch-aarch64.yaml +++ b/conda/environments/all_cuda-132_arch-aarch64.yaml @@ -67,7 +67,7 @@ dependencies: - packaging - pandas>=3.0.0,<3.1.0 - pandoc -- polars>=1.35,<1.39 +- polars>=1.35,<1.40 - pre-commit - pyarrow>=19.0.0,<24 - pydata-sphinx-theme>=0.15.4 diff --git a/conda/environments/all_cuda-132_arch-x86_64.yaml b/conda/environments/all_cuda-132_arch-x86_64.yaml index c0e1a5b36a1..5fcacb73088 100644 --- a/conda/environments/all_cuda-132_arch-x86_64.yaml +++ b/conda/environments/all_cuda-132_arch-x86_64.yaml @@ -67,7 +67,7 @@ dependencies: - packaging - pandas>=3.0.0,<3.1.0 - pandoc -- polars>=1.35,<1.39 +- polars>=1.35,<1.40 - pre-commit - pyarrow>=19.0.0,<24 - pydata-sphinx-theme>=0.15.4 diff --git a/conda/recipes/cudf-polars/recipe.yaml b/conda/recipes/cudf-polars/recipe.yaml index 482ea5f0b49..b3cb69835c6 100644 --- a/conda/recipes/cudf-polars/recipe.yaml +++ b/conda/recipes/cudf-polars/recipe.yaml @@ -37,7 +37,7 @@ requirements: - python - pylibcudf =${{ version }} - rapidsmpf =${{ minor_version }} - - polars>=1.35,<1.39 + - polars>=1.35,<1.40 - packaging - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} - if: cuda_major == "12" diff --git a/cpp/cmake/thirdparty/patches/arrow_rapidjson_cmake_policy_version_minimum.diff b/cpp/cmake/thirdparty/patches/arrow_rapidjson_cmake_policy_version_minimum.diff new file mode 100644 index 00000000000..153d1333b35 --- /dev/null +++ b/cpp/cmake/thirdparty/patches/arrow_rapidjson_cmake_policy_version_minimum.diff @@ -0,0 +1,15 @@ +diff --git a/cpp/cmake_modules/FindRapidJSONAlt.cmake b/cpp/cmake_modules/FindRapidJSONAlt.cmake +index babb450e20..148dd93a78 100644 +--- a/cpp/cmake_modules/FindRapidJSONAlt.cmake ++++ b/cpp/cmake_modules/FindRapidJSONAlt.cmake +@@ -26,7 +26,10 @@ endif() + if(RapidJSONAlt_FIND_QUIETLY) + list(APPEND find_package_args QUIET) + endif() ++set(_CMAKE_POLICY_VERSION_MINIMUM_OLD ${CMAKE_POLICY_VERSION_MINIMUM}) ++set(CMAKE_POLICY_VERSION_MINIMUM 3.5) + find_package(RapidJSON ${find_package_args}) ++set(CMAKE_POLICY_VERSION_MINIMUM ${_CMAKE_POLICY_VERSION_MINIMUM_OLD}) + if(RapidJSON_FOUND) + set(RapidJSONAlt_FOUND TRUE) + if(NOT TARGET RapidJSON) diff --git a/cpp/cmake/thirdparty/patches/override.json b/cpp/cmake/thirdparty/patches/override.json new file mode 100644 index 00000000000..3fb069dfce8 --- /dev/null +++ b/cpp/cmake/thirdparty/patches/override.json @@ -0,0 +1,17 @@ +{ + "packages": { + "arrow": { + "version": "21.0.0", + "git_url": "https://github.com/apache/arrow.git", + "git_tag": "apache-arrow-${version}", + "git_shallow": false, + "patches": [ + { + "file": "${current_json_dir}/arrow_rapidjson_cmake_policy_version_minimum.diff", + "issue": "https://github.com/apache/arrow/pull/49993" + } + ], + "source_subdir": "cpp" + } + } +} diff --git a/cpp/src/io/json/json_tree.cu b/cpp/src/io/json/json_tree.cu index f38c4f0c55c..5f3c9b314ac 100644 --- a/cpp/src/io/json/json_tree.cu +++ b/cpp/src/io/json/json_tree.cu @@ -130,14 +130,14 @@ struct is_nested_end { }; struct checked_token_level_output { - bool* depth_out_of_range; + int32_t* depth_out_of_range; __device__ TreeDepthT operator()(size_type level) const { static_assert(sizeof(TreeDepthT) < sizeof(size_type)); if (level < static_cast(cuda::std::numeric_limits::min()) || level > static_cast(cuda::std::numeric_limits::max())) { - cuda::atomic_ref flag{*depth_out_of_range}; + cuda::atomic_ref flag{*depth_out_of_range}; if (!flag.load(cuda::std::memory_order_relaxed)) { flag.store(true, cuda::std::memory_order_relaxed); } @@ -302,7 +302,7 @@ tree_meta_t get_tree_representation(device_span tokens, return does_push(token) - does_pop(token); })); auto depth_out_of_range = - cudf::detail::device_scalar(false, stream, cudf::get_current_device_resource_ref()); + cudf::detail::device_scalar(0, stream, cudf::get_current_device_resource_ref()); auto const token_level_output_it = thrust::make_transform_output_iterator( token_levels.begin(), checked_token_level_output{depth_out_of_range.data()}); thrust::exclusive_scan(rmm::exec_policy_nosync(stream, cudf::get_current_device_resource_ref()), diff --git a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp index 9b6953b4bd1..2a3fad68d7c 100644 --- a/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp +++ b/cpp/src/io/parquet/io_utils/parquet_io_utils.cpp @@ -6,6 +6,7 @@ #include "io/comp/common.hpp" #include "io/parquet/parquet_common.hpp" +#include #include #include #include @@ -21,6 +22,7 @@ #include #include +#include /** * @file parquet_io_utils.cpp @@ -89,6 +91,9 @@ fetch_byte_ranges_to_device_async( std::vector io_offsets; std::vector io_sizes; std::vector destinations; + io_offsets.reserve(byte_ranges.size()); + io_sizes.reserve(byte_ranges.size()); + destinations.reserve(byte_ranges.size()); for (size_t chunk = 0; chunk < byte_ranges.size();) { auto const io_offset = static_cast(byte_ranges[chunk].offset()); @@ -110,12 +115,24 @@ fetch_byte_ranges_to_device_async( CUDF_EXPECTS(io_offsets.size() == io_sizes.size() and io_sizes.size() == destinations.size(), "Unexpected number of IO offsets, sizes, or destinations"); + using host_read_buffer = std::unique_ptr; + + // Vectors to hold futures from datasource std::vector> device_read_tasks{}; - std::vector> host_read_tasks{}; - device_read_tasks.reserve(byte_ranges.size()); - host_read_tasks.reserve(byte_ranges.size()); + std::vector> host_read_tasks{}; + device_read_tasks.reserve(io_offsets.size()); + host_read_tasks.reserve(io_offsets.size()); + + // Vectors to store intermediate host read buffers and relevant pointers + std::vector copy_dsts{}; + std::vector copy_sizes{}; + copy_dsts.reserve(io_offsets.size()); + copy_sizes.reserve(io_offsets.size()); + + // Vector to store intermediate host buffers + std::vector host_buffers{}; - // device_read_async is not guaranteed to follow stream-ordering (see datasource API docs). + // `device_read_async` is not guaranteed to follow stream-ordering (see datasource API docs) stream.synchronize(); { @@ -133,35 +150,40 @@ fetch_byte_ranges_to_device_async( device_read_tasks.emplace_back( datasource.device_read_async(io_offset, io_size, dest, stream)); } else { - // Read the column chunk data to the host buffer copy it to the device buffer - host_read_tasks.emplace_back(cudf::detail::host_worker_pool().submit_task( - [&datasource, io_offset, io_size, dest, stream]() { - auto host_buffer = datasource.host_read(io_offset, io_size); - cudf::detail::cuda_memcpy_async( - cudf::device_span{dest, io_size}, - cudf::host_span{host_buffer->data(), io_size}, - stream); - return io_size; - })); + // Asynchronously read column chunk data to a host buffer + host_read_tasks.emplace_back(datasource.host_read_async(io_offset, io_size)); + copy_dsts.push_back(static_cast(dest)); + copy_sizes.push_back(io_size); } }); - } - auto sync_function = [](decltype(host_read_tasks) host_read_tasks, - decltype(device_read_tasks) device_read_tasks) { - for (auto& task : host_read_tasks) { - task.get(); - } - for (auto& task : device_read_tasks) { - task.get(); + // If there are host reads, schedule a batched memcpy to device + if (not host_read_tasks.empty()) { + std::vector copy_srcs{}; + copy_srcs.reserve(host_read_tasks.size()); + host_buffers.reserve(host_read_tasks.size()); + + for (auto& task : host_read_tasks) { + host_buffers.emplace_back(task.get()); + copy_srcs.push_back(host_buffers.back().get()->data()); + } + CUDF_CUDA_TRY(cudf::detail::memcpy_batch_async( + copy_dsts.data(), copy_srcs.data(), copy_sizes.data(), copy_dsts.size(), stream)); } - }; - return {std::move(column_chunk_buffers), - std::move(column_chunk_data), - std::async(std::launch::deferred, - sync_function, - std::move(host_read_tasks), - std::move(device_read_tasks))}; + + // Synchronize the stream if `memcpy_batch_async` was scheduled to safely discard the host + // buffers + if (not host_buffers.empty()) { stream.synchronize(); } + + auto sync_function = [](decltype(device_read_tasks) device_read_tasks) { + for (auto& task : device_read_tasks) { + task.get(); + } + }; + return {std::move(column_chunk_buffers), + std::move(column_chunk_data), + std::async(std::launch::deferred, sync_function, std::move(device_read_tasks))}; + } } } // namespace cudf::io::parquet diff --git a/dependencies.yaml b/dependencies.yaml index f50d7814eb2..b9ef783e44a 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -838,7 +838,7 @@ dependencies: # 'nvidia-ml-py' provides the 'pynvml' module - &nvidia_ml_py nvidia-ml-py>=12 - packaging - - polars>=1.35,<1.39 + - polars>=1.35,<1.40 - typing_extensions>=4.0.0 run_cudf_polars_dask: common: diff --git a/java/src/main/native/CMakeLists.txt b/java/src/main/native/CMakeLists.txt index b34c3bd82c5..0f0f5f1e9c8 100644 --- a/java/src/main/native/CMakeLists.txt +++ b/java/src/main/native/CMakeLists.txt @@ -12,7 +12,9 @@ include(rapids-cmake) include(rapids-cuda) include(rapids-find) include(rapids-cpm) -rapids_cpm_init() +rapids_cpm_init( + OVERRIDE "${CMAKE_CURRENT_SOURCE_DIR}/../../../../cpp/cmake/thirdparty/patches/override.json" +) rapids_cuda_init_architectures(CUDF_JNI) diff --git a/python/cudf_polars/cudf_polars/callback.py b/python/cudf_polars/cudf_polars/callback.py index acd0452ae1b..93dd0790bd5 100644 --- a/python/cudf_polars/cudf_polars/callback.py +++ b/python/cudf_polars/cudf_polars/callback.py @@ -307,7 +307,7 @@ def _callback( else: return df, timer.timings elif config_options.executor.name == "streaming": - from cudf_polars.experimental.parallel import evaluate_streaming + from cudf_polars.streaming.parallel import evaluate_streaming if timer is not None: msg = textwrap.dedent("""\ diff --git a/python/cudf_polars/cudf_polars/containers/column.py b/python/cudf_polars/cudf_polars/containers/column.py index c28db5a0d14..9f8514dca77 100644 --- a/python/cudf_polars/cudf_polars/containers/column.py +++ b/python/cudf_polars/cudf_polars/containers/column.py @@ -134,7 +134,7 @@ def serialize( To enable dask support, dask serializers must be registered - >>> from cudf_polars.experimental.dask_serialize import register + >>> from cudf_polars.streaming.dask_serialize import register >>> register() Returns diff --git a/python/cudf_polars/cudf_polars/containers/dataframe.py b/python/cudf_polars/cudf_polars/containers/dataframe.py index d96e96cc997..7d6c93c097f 100644 --- a/python/cudf_polars/cudf_polars/containers/dataframe.py +++ b/python/cudf_polars/cudf_polars/containers/dataframe.py @@ -299,7 +299,7 @@ def serialize( To enable dask support, dask serializers must be registered - >>> from cudf_polars.experimental.dask_serialize import register + >>> from cudf_polars.streaming.dask_serialize import register >>> register() Parameters diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py index f21451e702e..70ced1c440d 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/rolling.py @@ -22,7 +22,7 @@ offsets_to_windows, range_window_bounds, ) -from cudf_polars.utils.versions import POLARS_VERSION_LT_136 +from cudf_polars.utils.versions import POLARS_VERSION_LT_136, POLARS_VERSION_LT_139 if TYPE_CHECKING: from collections.abc import Sequence @@ -79,14 +79,14 @@ def to_request( elif isinstance(value, expr.Agg): child = value.children[0] col = child.evaluate(df, context=ExecutionContext.ROLLING) - if POLARS_VERSION_LT_136 and value.name == "var": # pragma: no cover + if (POLARS_VERSION_LT_136 or not POLARS_VERSION_LT_139) and value.name == "var": # Polars variance produces null if nvalues <= ddof # libcudf produces NaN. However, we can get the polars # behaviour by setting the minimum window size to ddof + # 1. - # - # We still need this check, polars is not hitting it because - # of https://github.com/pola-rs/polars/pull/25117 + # In polars 1.36-1.38 this code path is not hit because + # rolling goes through the IR path (not RollingWindow). + # See https://github.com/pola-rs/polars/pull/25117 min_periods = value.options + 1 else: col = value.evaluate( @@ -146,12 +146,12 @@ def __init__( raise NotImplementedError( "Incorrect handling of empty groups for list collection" ) - if POLARS_VERSION_LT_136 and not plc.rolling.is_valid_rolling_aggregation( + if ( + POLARS_VERSION_LT_136 or not POLARS_VERSION_LT_139 + ) and not plc.rolling.is_valid_rolling_aggregation( agg.dtype.plc_type, agg.agg_request ): - raise NotImplementedError( - f"Unsupported rolling aggregation {agg}" - ) # pragma: no cover; polars may raise ahead of time + raise NotImplementedError(f"Unsupported rolling aggregation {agg}") def do_evaluate( # noqa: D102 self, df: DataFrame, *, context: ExecutionContext = ExecutionContext.FRAME diff --git a/python/cudf_polars/cudf_polars/dsl/expressions/unary.py b/python/cudf_polars/cudf_polars/dsl/expressions/unary.py index 6f033ed70cb..de03a4623a4 100644 --- a/python/cudf_polars/cudf_polars/dsl/expressions/unary.py +++ b/python/cudf_polars/cudf_polars/dsl/expressions/unary.py @@ -240,15 +240,14 @@ def do_evaluate( if maintain_order: column = column.sorted_like(values) return column - elif self.name == "set_sorted": # pragma: no cover - # TODO: LazyFrame.set_sorted is proper IR concept (ie. FunctionIR::Hint) - # and is is currently not implemented. We should reimplement it as a MapFunction. + elif self.name == "set_sorted": (column,) = (child.evaluate(df, context=context) for child in self.children) - (asc,) = self.options + if isinstance(self.options[0], str): + descending = self.options[0] == "descending" # pragma: no cover + else: + descending, _ = self.options order = ( - plc.types.Order.ASCENDING - if asc == "ascending" - else plc.types.Order.DESCENDING + plc.types.Order.DESCENDING if descending else plc.types.Order.ASCENDING ) null_order = plc.types.NullOrder.BEFORE if column.null_count > 0 and (n := column.size) > 1: diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index e4a2e2fcc12..8b41e6d6dca 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -3245,19 +3245,21 @@ def do_evaluate( class HConcat(IR): """Concatenate dataframes horizontally.""" - __slots__ = ("should_broadcast",) - _non_child = ("schema", "should_broadcast") - _n_non_child_args = 2 + __slots__ = ("should_broadcast", "strict") + _non_child = ("schema", "should_broadcast", "strict") + _n_non_child_args = 3 def __init__( self, schema: Schema, should_broadcast: bool, # noqa: FBT001 + strict: bool, # noqa: FBT001 *children: IR, ): self.schema = schema self.should_broadcast = should_broadcast - self._non_child_args = (schema, should_broadcast) + self.strict = strict + self._non_child_args = (schema, should_broadcast, strict) self.children = children @staticmethod @@ -3300,6 +3302,7 @@ def do_evaluate( cls, schema: Schema, should_broadcast: bool, # noqa: FBT001 + strict: bool, # noqa: FBT001 *dfs: DataFrame, context: IRExecutionContext, ) -> DataFrame: @@ -3317,6 +3320,13 @@ def do_evaluate( result = DataFrame(ordered, stream=stream) else: max_rows = max(df.num_rows for df in dfs) + if strict and any(df.num_rows != max_rows for df in dfs): + heights = [df.num_rows for df in dfs] + msg = ( + f"cannot concat DataFrames horizontally" + f" with strict=True: height mismatch {heights}" + ) + raise pl.exceptions.ShapeError(msg) # Horizontal concatenation extends shorter tables with nulls result = DataFrame( itertools.chain.from_iterable( diff --git a/python/cudf_polars/cudf_polars/dsl/translate.py b/python/cudf_polars/cudf_polars/dsl/translate.py index 8f68b259954..b2e560c8b36 100644 --- a/python/cudf_polars/cudf_polars/dsl/translate.py +++ b/python/cudf_polars/cudf_polars/dsl/translate.py @@ -35,6 +35,7 @@ from cudf_polars.utils.versions import ( POLARS_VERSION_LT_136, POLARS_VERSION_LT_138, + POLARS_VERSION_LT_139, ) if TYPE_CHECKING: @@ -47,6 +48,29 @@ __all__ = ["Translator", "translate_named_expr"] +def _align_decimal_float_for_comparison( + *operands: expr.Expr, +) -> tuple[expr.Expr, ...]: + """ + Cast Decimal operands to Float64 when mixed with Float operands. + + libcudf does not produce correct results for Decimal vs Float + comparisons. This matches Polars' supertype resolution where + Decimal + Float yields Float64. + """ + has_decimal = any(plc.traits.is_fixed_point(op.dtype.plc_type) for op in operands) + has_float = any(plc.traits.is_floating_point(op.dtype.plc_type) for op in operands) + if has_decimal and has_float: + f64 = DataType(pl.Float64()) + return tuple( + expr.Cast(f64, False, op) # noqa: FBT003 + if plc.traits.is_fixed_point(op.dtype.plc_type) + else op + for op in operands + ) + return operands + + def _check_compression(data: bytes) -> str | None: # Vendored from Polars' SupportedCompression::check in # polars-io/src/utils/compression.rs @@ -128,7 +152,7 @@ def translate_ir(self, *, n: int | None = None) -> ir.IR: # IR is versioned with major.minor, minor is bumped for backwards # compatible changes (e.g. adding new nodes), major is bumped for # incompatible changes (e.g. renaming nodes). - if (version := self.visitor.version()) >= (12, 1): + if (version := self.visitor.version()) >= (12, 2): e = NotImplementedError( f"No support for polars IR {version=}" ) # pragma: no cover; no such version for now. @@ -250,6 +274,17 @@ def __exit__(self, *args: Any) -> None: self.translator._expr_context = self._prev +def _is_dynamic_pred(visitor: Any, expr_ir: Any) -> bool: + try: + visitor.view_expression(expr_ir.node) + except Exception as e: + if str(e) == "dynamic_pred": + return True + raise # pragma: no cover + else: + return False + + class set_internal_name_gen(AbstractContextManager[None]): """Share one internal-name generator across sibling expression translations.""" @@ -328,7 +363,7 @@ def _(node: plrs._ir_nodes.Scan, translator: Translator, schema: Schema) -> ir.I if typ in ("csv", "ndjson"): for p in paths: if plc.io.SourceInfo._is_remote_uri(p): - continue + continue # pragma: no cover data = _read_file_bytes(Path(p)) compression = _check_compression(data) if compression is not None: @@ -347,9 +382,12 @@ def _(node: plrs._ir_nodes.Scan, translator: Translator, schema: Schema) -> ir.I n_rows, row_index, include_file_paths, - translate_named_expr(translator, n=node.predicate, schema=schema) - if node.predicate is not None - else None, + ( + None + if node.predicate is None + or _is_dynamic_pred(translator.visitor, node.predicate) + else translate_named_expr(translator, n=node.predicate, schema=schema) + ), parquet_options, ) @@ -562,6 +600,8 @@ def _(node: plrs._ir_nodes.Slice, translator: Translator, schema: Schema) -> ir. def _(node: plrs._ir_nodes.Filter, translator: Translator, schema: Schema) -> ir.IR: with set_node(translator.visitor, node.input): inp = translator.translate_ir(n=None) + if _is_dynamic_pred(translator.visitor, node.predicate): + return inp mask = translate_named_expr(translator, n=node.predicate, schema=inp.schema) return ir.Filter(schema, mask, inp) @@ -610,9 +650,14 @@ def _(node: plrs._ir_nodes.Union, translator: Translator, schema: Schema) -> ir. @_translate_ir.register def _(node: plrs._ir_nodes.HConcat, translator: Translator, schema: Schema) -> ir.IR: + if POLARS_VERSION_LT_139: + strict = False # pragma: no cover + else: + _, strict, _ = node.options return ir.HConcat( schema, False, # noqa: FBT003 + strict, *(translator.translate_ir(n=n) for n in node.inputs), ) @@ -650,7 +695,7 @@ def _(node: plrs._ir_nodes.Sink, translator: Translator, schema: Schema) -> ir.I if k in {"mkdir", "maintain_order", "sync_on_close"}: options[k] = v - if sink_kind in ("Csv", "NDJson"): + if sink_kind in ("Csv", "NDJson", "Json"): compression = format_options.get("compression") if compression and compression != "Uncompressed": raise NotImplementedError( @@ -773,6 +818,10 @@ def _( ) (closed,) = options lop, rop = expr.BooleanFunction._BETWEEN_OPS[closed] + if not POLARS_VERSION_LT_139: + # IsBetween type coercion does not + # insert casts for Decimal/Float comparisons. + column, lo, hi = _align_decimal_float_for_comparison(column, lo, hi) return expr.BinOp( dtype, plc.binaryop.BinaryOperator.LOGICAL_AND, @@ -845,14 +894,9 @@ def _( dtype: DataType, schema: Schema, ) -> expr.Expr: - if isinstance( + if hasattr(plrs._expr_nodes, "RollingGroupOptions") and isinstance( node.options, plrs._expr_nodes.RollingGroupOptions - ): # pragma: no cover; polars gives Aexpr::Rolling node now - # TODO: As of polars 1.36.0, rolling is represented as a Rolling expression node - # but is currently not implemented in the python node visitor. Once we support it, - # we should move this branch to separate translator dispatch function for the new node. - - # pl.col("a").rolling(...) + ): # pragma: no cover; polars >=1.36 uses AExpr::Rolling now with set_expr_context(translator, ExecutionContext.ROLLING): agg = translator.translate_expr(n=node.function, schema=schema) name_generator = translator._internal_name_gen or unique_names(schema) @@ -953,6 +997,61 @@ def _( assert_never(node.options) +def _translate_rolling( + node: plrs._expr_nodes.Rolling, + translator: Translator, + dtype: DataType, + schema: Schema, +) -> expr.Expr: + # pl.col("a").rolling(...) + with set_expr_context(translator, ExecutionContext.ROLLING): + agg_expr = translator.translate_expr(n=node.function, schema=schema) + orderby = translator.visitor.view_expression(node.index_column).name + orderby_dtype = schema[orderby].plc_type + if plc.traits.is_integral(orderby_dtype): + orderby_dtype = plc.DataType(plc.TypeId.INT64) + name_generator = unique_names(schema) + aggs, named_post_agg = decompose_single_agg( + expr.NamedExpr(next(name_generator), agg_expr), + name_generator, + is_top=True, + context=ExecutionContext.ROLLING, + ) + named_aggs = [a for a, _ in aggs] + closed_window = node.closed_window + if isinstance(named_post_agg.value, expr.Col): + (named_agg,) = named_aggs + return expr.RollingWindow( + named_agg.value.dtype, + orderby_dtype, + node.offset, + node.period, + closed_window, + orderby, + named_agg.value, + ) + replacements: dict[expr.Expr, expr.Expr] = { + expr.Col(a.value.dtype, a.name): expr.RollingWindow( + a.value.dtype, + orderby_dtype, + node.offset, + node.period, + closed_window, + orderby, + a.value, + ) + for a in named_aggs + } + return replace([named_post_agg.value], replacements)[0] + + +# Rolling was added in polars 1.39. Use explicit registration (not annotation-based) +# so that Python 3.14's stricter get_type_hints() resolution does not fail when +# running on older polars versions that lack this node type. +if hasattr(plrs._expr_nodes, "Rolling"): + _translate_expr.register(plrs._expr_nodes.Rolling)(_translate_rolling) + + @_translate_expr.register def _( node: plrs._expr_nodes.Literal, diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/__init__.py b/python/cudf_polars/cudf_polars/engine/__init__.py similarity index 100% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/__init__.py rename to python/cudf_polars/cudf_polars/engine/__init__.py diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py b/python/cudf_polars/cudf_polars/engine/core.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py rename to python/cudf_polars/cudf_polars/engine/core.py index 2e8d14c7e01..45aeaa8d270 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/core.py +++ b/python/cudf_polars/cudf_polars/engine/core.py @@ -25,15 +25,15 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IRExecutionContext -from cudf_polars.experimental.base import StatsCollector -from cudf_polars.experimental.parallel import lower_ir_graph -from cudf_polars.experimental.rapidsmpf.collectives import ReserveOpIDs -from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id -from cudf_polars.experimental.rapidsmpf.core import generate_network -from cudf_polars.experimental.rapidsmpf.tracing import log_query_plan -from cudf_polars.experimental.rapidsmpf.utils import empty_table_chunk -from cudf_polars.experimental.statistics import collect_statistics -from cudf_polars.experimental.utils import _concat +from cudf_polars.streaming.actor_graph.collectives import ReserveOpIDs +from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id +from cudf_polars.streaming.actor_graph.core import generate_network +from cudf_polars.streaming.actor_graph.tracing import log_query_plan +from cudf_polars.streaming.actor_graph.utils import empty_table_chunk +from cudf_polars.streaming.base import StatsCollector +from cudf_polars.streaming.parallel import lower_ir_graph +from cudf_polars.streaming.statistics import collect_statistics +from cudf_polars.streaming.utils import _concat from cudf_polars.utils.config import get_total_device_memory if TYPE_CHECKING: @@ -48,8 +48,8 @@ from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.base import PartitionInfo - from cudf_polars.experimental.parallel import ConfigOptions + from cudf_polars.streaming.base import PartitionInfo + from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -176,7 +176,7 @@ def __init__( ): # Refuse to construct if a ``DefaultSingletonEngine`` is alive # (no-op for the singleton itself). - from cudf_polars.experimental.rapidsmpf.frontend.default_singleton_engine import ( + from cudf_polars.engine.default_singleton_engine import ( check_no_live_default_singleton, ) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py b/python/cudf_polars/cudf_polars/engine/dask.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py rename to python/cudf_polars/cudf_polars/engine/dask.py index 23e85074a0b..509980fc2c2 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/dask.py +++ b/python/cudf_polars/cudf_polars/engine/dask.py @@ -27,14 +27,14 @@ import rmm.mr -from cudf_polars.experimental.rapidsmpf.frontend.core import ( +from cudf_polars.engine.core import ( ClusterInfo, StreamingEngine, check_reserved_keys, evaluate_on_rank, resolve_rapidsmpf_options, ) -from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( +from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, bind_to_gpu, ) @@ -48,9 +48,9 @@ from rapidsmpf.streaming.cudf.channel_metadata import ChannelMetadata from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.parallel import ConfigOptions - from cudf_polars.experimental.rapidsmpf.frontend.core import T - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.engine.core import T + from cudf_polars.engine.options import StreamingOptions + from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -87,7 +87,7 @@ def dask_setup(nanny: distributed.Nanny) -> None: Usage:: dask worker SCHEDULER:8786 --nworkers N --nthreads 1 \ - --preload-nanny cudf_polars.experimental.rapidsmpf.frontend.dask + --preload-nanny cudf_polars.engine.dask Parameters ---------- @@ -566,9 +566,7 @@ class DaskEngine(StreamingEngine): and sets ``CUDA_VISIBLE_DEVICES`` per worker), disable some or all of the built-in binding to avoid conflicts: - >>> from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( - ... HardwareBindingPolicy, - ... ) + >>> from cudf_polars.engine.hardware_binding import HardwareBindingPolicy >>> with DaskEngine( # doctest: +SKIP ... dask_client=dc, ... engine_options={ @@ -581,7 +579,7 @@ class DaskEngine(StreamingEngine): one GPU per worker before the worker process spawns:: dask worker SCHEDULER:8786 --nworkers N --nthreads 1 \ - --preload-nanny cudf_polars.experimental.rapidsmpf.frontend.dask + --preload-nanny cudf_polars.engine.dask Then connect from the client:: @@ -784,9 +782,7 @@ def from_options( Examples -------- - >>> from cudf_polars.experimental.rapidsmpf.frontend.options import ( - ... StreamingOptions, - ... ) + >>> from cudf_polars.engine.options import StreamingOptions >>> opts = StreamingOptions(num_streaming_threads=4, fallback_mode="silent") >>> with DaskEngine.from_options(opts) as engine: # doctest: +SKIP ... result = pl.LazyFrame({"a": [1, 2, 3]}).collect(engine=engine) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/default_singleton_engine.py b/python/cudf_polars/cudf_polars/engine/default_singleton_engine.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/default_singleton_engine.py rename to python/cudf_polars/cudf_polars/engine/default_singleton_engine.py index e5433770427..73f08730df6 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/default_singleton_engine.py +++ b/python/cudf_polars/cudf_polars/engine/default_singleton_engine.py @@ -16,10 +16,10 @@ ) from rapidsmpf.progress_thread import ProgressThread -from cudf_polars.experimental.rapidsmpf.frontend.core import ( +from cudf_polars.engine.core import ( resolve_rapidsmpf_options, ) -from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine +from cudf_polars.engine.spmd import SPMDEngine if TYPE_CHECKING: from collections.abc import Callable diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/hardware_binding.py b/python/cudf_polars/cudf_polars/engine/hardware_binding.py similarity index 100% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/hardware_binding.py rename to python/cudf_polars/cudf_polars/engine/hardware_binding.py diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/options.py b/python/cudf_polars/cudf_polars/engine/options.py similarity index 99% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/options.py rename to python/cudf_polars/cudf_polars/engine/options.py index d620ab7a826..21ed6c961f0 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/options.py +++ b/python/cudf_polars/cudf_polars/engine/options.py @@ -15,7 +15,7 @@ from rapidsmpf.config import Options from rapidsmpf.utils.string import parse_boolean -from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( +from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, ) from cudf_polars.utils.config import MemoryResourceConfig @@ -273,8 +273,7 @@ class StreamingOptions: Env: ``CUDF_POLARS__CUDA_STREAM_POLICY``. Category: engine. hardware_binding - Hardware binding policy. Pass a - :class:`~cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.HardwareBindingPolicy` + Hardware binding policy. Pass a :class:`~cudf_polars.engine.hardware_binding.HardwareBindingPolicy` instance for fine-grained control. Env: ``CUDF_POLARS__HARDWARE_BINDING`` (JSON object, e.g. ``'{"enabled": false}'``). diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py b/python/cudf_polars/cudf_polars/engine/ray.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py rename to python/cudf_polars/cudf_polars/engine/ray.py index 4e74e392c8c..532e863bdf7 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/ray.py +++ b/python/cudf_polars/cudf_polars/engine/ray.py @@ -23,14 +23,14 @@ import rmm.mr -from cudf_polars.experimental.rapidsmpf.frontend.core import ( +from cudf_polars.engine.core import ( ClusterInfo, StreamingEngine, check_reserved_keys, evaluate_on_rank, resolve_rapidsmpf_options, ) -from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( +from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, bind_to_gpu, ) @@ -46,9 +46,9 @@ from ray.actor import ActorHandle from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.parallel import ConfigOptions - from cudf_polars.experimental.rapidsmpf.frontend.core import T - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.engine.core import T + from cudf_polars.engine.options import StreamingOptions + from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -163,10 +163,9 @@ class RankActor: Notes ----- - Calls :func:`~cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.bind_to_gpu` - at construction time, before RMM and communicator initialisation, so that - CPU affinity, NUMA memory policy, and ``UCX_NET_DEVICES`` are set as early - as possible. + Calls :func:`~cudf_polars.engine.hardware_binding.bind_to_gpu` at construction + time, before RMM and communicator initialisation, so that CPU affinity, NUMA + memory policy, and ``UCX_NET_DEVICES`` are set as early as possible. """ def __init__( @@ -432,11 +431,10 @@ class RayEngine(StreamingEngine): Creates a RapidsMPF Ray cluster and returns an engine that can be passed to ``LazyFrame.collect(engine=engine)``. - Prefer :meth:`from_options` for typical use — pass a - :class:`~cudf_polars.experimental.rapidsmpf.frontend.options.StreamingOptions` - instance for a unified, typed interface. The ``__init__`` parameters - (``rapidsmpf_options``, ``executor_options``, ``engine_options``) are - intended for advanced use when fine-grained control is needed. + Prefer :meth:`from_options` for typical use. Pass a :class:`~cudf_polars.engine.options.StreamingOptions` + instance for a unified, typed interface. The ``__init__`` parameters (``rapidsmpf_options``, + ``executor_options``, ``engine_options``) are intended for advanced use when + fine-grained control is needed. Prefer the context-manager form in scripts: it guarantees that actors and Ray are shut down even if an exception is raised. In interactive environments @@ -695,9 +693,7 @@ def from_options( Examples -------- - >>> from cudf_polars.experimental.rapidsmpf.frontend.options import ( - ... StreamingOptions, - ... ) + >>> from cudf_polars.engine.options import StreamingOptions >>> opts = StreamingOptions(num_streaming_threads=4, fallback_mode="silent") >>> with RayEngine.from_options(opts) as engine: # doctest: +SKIP ... result = pl.LazyFrame({"a": [1, 2, 3]}).collect(engine=engine) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py b/python/cudf_polars/cudf_polars/engine/spmd.py similarity index 93% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py rename to python/cudf_polars/cudf_polars/engine/spmd.py index e996b2de733..64e7b189350 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/frontend/spmd.py +++ b/python/cudf_polars/cudf_polars/engine/spmd.py @@ -27,8 +27,7 @@ from pylibcudf.contiguous_split import pack from cudf_polars.containers import DataFrame, DataType -from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id -from cudf_polars.experimental.rapidsmpf.frontend.core import ( +from cudf_polars.engine.core import ( ClusterInfo, StreamingEngine, all_gather_host_data, @@ -36,11 +35,12 @@ evaluate_on_rank, resolve_rapidsmpf_options, ) -from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( +from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, bind_to_gpu, ) -from cudf_polars.experimental.rapidsmpf.utils import set_memory_resource +from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id +from cudf_polars.streaming.actor_graph.utils import set_memory_resource from cudf_polars.utils.config import ( MemoryResourceConfig, SPMDContext, @@ -58,9 +58,9 @@ import polars as pl from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.parallel import ConfigOptions - from cudf_polars.experimental.rapidsmpf.frontend.core import T - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.engine.core import T + from cudf_polars.engine.options import StreamingOptions + from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -204,11 +204,10 @@ class SPMDEngine(StreamingEngine): such as shuffles, all-gathers, and joins, coordinate across ranks to produce a globally consistent result. - Prefer :meth:`from_options` for typical use — pass a - :class:`~cudf_polars.experimental.rapidsmpf.frontend.options.StreamingOptions` - instance for a unified, typed interface. The ``__init__`` parameters - (``rapidsmpf_options``, ``executor_options``, ``engine_options``) are - intended for advanced use when fine-grained control is needed. + Prefer :meth:`from_options` for typical use. Pass a :class:`~cudf_polars.engine.options.StreamingOptions` + instance for a unified, typed interface. The ``__init__`` parameters (``rapidsmpf_options``, + ``executor_options``, ``engine_options``) are intended for advanced use when + fine-grained control is needed. This class is the primary entry point for SPMD execution. It: @@ -277,7 +276,7 @@ class SPMDEngine(StreamingEngine): Every rank must issue the *same* sequence of Polars queries in the *same* order. Collective operations (shuffles, all-gathers, joins) are matched - across ranks by a monotonically increasing operation ID — if one rank calls + across ranks by a monotonically increasing operation ID; if one rank calls a collective that another rank does not, all ranks will deadlock. This means your driver script must be fully deterministic: avoid rank-conditional ``collect`` calls, early exits, or any branching that would cause different @@ -297,8 +296,7 @@ class SPMDEngine(StreamingEngine): executor_options Executor-specific options (e.g. ``max_rows_per_partition``). engine_options - Engine-specific keyword arguments (e.g. ``raise_on_fail``, - ``parquet_options``). + Engine-specific keyword arguments (e.g. ``raise_on_fail``, ``parquet_options``). Raises ------ @@ -307,13 +305,11 @@ class SPMDEngine(StreamingEngine): Notes ----- - Calls - :func:`~cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.bind_to_gpu` - at construction time, before RMM and communicator initialisation, so that - CPU affinity, NUMA memory policy, and ``UCX_NET_DEVICES`` are set as early - as possible. By default, binding is skipped under ``rrun`` (which already - performs its own binding) — see - :attr:`HardwareBindingPolicy.skip_under_rrun`. + Calls :func:`~cudf_polars.engine.hardware_binding.bind_to_gpu` at construction + time, before RMM and communicator initialisation, so that CPU affinity, NUMA + memory policy, and ``UCX_NET_DEVICES`` are set as early as possible. By default, + binding is skipped under ``rrun`` (which already performs its own binding), + see :attr:`HardwareBindingPolicy.skip_under_rrun`. Examples -------- @@ -447,9 +443,7 @@ def from_options(cls, options: StreamingOptions) -> SPMDEngine: Examples -------- - >>> from cudf_polars.experimental.rapidsmpf.frontend.options import ( - ... StreamingOptions, - ... ) + >>> from cudf_polars.engine.options import StreamingOptions >>> opts = StreamingOptions(num_streaming_threads=8, fallback_mode="silent") >>> with SPMDEngine.from_options(opts) as engine: # doctest: +SKIP ... result = df.lazy().collect(engine=engine) diff --git a/python/cudf_polars/cudf_polars/experimental/__init__.py b/python/cudf_polars/cudf_polars/experimental/__init__.py deleted file mode 100644 index 6fd93bf5157..00000000000 --- a/python/cudf_polars/cudf_polars/experimental/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -"""Experimental features, which can change without any deprecation period.""" - -from __future__ import annotations - -__all__: list[str] = [] diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py b/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py deleted file mode 100644 index 8591ed18cdd..00000000000 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils.py +++ /dev/null @@ -1,27 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Benchmark utilities.""" - -from __future__ import annotations - -from cudf_polars.experimental.benchmarks.utils_new_frontends import ( - COUNT_DTYPE, - QueryResult, - RunConfig, - build_parser, - get_data, - parse_args, - run_duckdb, - run_polars, -) - -__all__: list[str] = [ - "COUNT_DTYPE", - "QueryResult", - "RunConfig", - "build_parser", - "get_data", - "parse_args", - "run_duckdb", - "run_polars", -] diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/__init__.py b/python/cudf_polars/cudf_polars/experimental/rapidsmpf/__init__.py deleted file mode 100644 index 7eedbf12bd4..00000000000 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -"""RapidsMPF streaming-engine support.""" - -from __future__ import annotations - -# Side-effect imports: each module registers -# ``@generate_ir_sub_network.register(...)`` handlers at import time so the -# dispatch table is populated before any query is evaluated. -import cudf_polars.experimental.rapidsmpf.collectives.shuffle -import cudf_polars.experimental.rapidsmpf.collectives.sort -import cudf_polars.experimental.rapidsmpf.groupby -import cudf_polars.experimental.rapidsmpf.io -import cudf_polars.experimental.rapidsmpf.join -import cudf_polars.experimental.rapidsmpf.over -import cudf_polars.experimental.rapidsmpf.repartition -import cudf_polars.experimental.rapidsmpf.union # noqa: F401 - -__all__: list[str] = [] diff --git a/python/cudf_polars/cudf_polars/streaming/__init__.py b/python/cudf_polars/cudf_polars/streaming/__init__.py new file mode 100644 index 00000000000..c1e71e1959e --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""Streaming execution layer for cudf-polars.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py new file mode 100644 index 00000000000..a4dd049a0ea --- /dev/null +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/__init__.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""RapidsMPF streaming-engine support.""" + +from __future__ import annotations + +import cudf_polars.streaming.actor_graph.collectives.shuffle +import cudf_polars.streaming.actor_graph.collectives.sort + +# Side-effect imports: each module registers +# ``@generate_ir_sub_network.register(...)`` handlers at import time so the +# dispatch table is populated before any query is evaluated. +import cudf_polars.streaming.actor_graph.groupby +import cudf_polars.streaming.actor_graph.io +import cudf_polars.streaming.actor_graph.join +import cudf_polars.streaming.actor_graph.over +import cudf_polars.streaming.actor_graph.repartition +import cudf_polars.streaming.actor_graph.union # noqa: F401 + +__all__: list[str] = [] diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/__init__.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py similarity index 64% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/__init__.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py index b6fc25f335f..5cbf011ca35 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/__init__.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/__init__.py @@ -1,10 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -"""Collective operations for the RapidsMPF streaming runtime.""" +"""Collective operations for building a RapidsMPF actor graph.""" from __future__ import annotations -from cudf_polars.experimental.rapidsmpf.collectives.common import ( +from cudf_polars.streaming.actor_graph.collectives.common import ( ReserveOpIDs, reserve_op_id, ) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/allgather.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py similarity index 100% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/allgather.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/allgather.py diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py index fa5cb995ac5..68758f6c1ff 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/common.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/common.py @@ -12,11 +12,11 @@ from cudf_polars.dsl.ir import Distinct, GroupBy, Sort from cudf_polars.dsl.traversal import traversal -from cudf_polars.experimental.io import StreamingSink -from cudf_polars.experimental.join import Join -from cudf_polars.experimental.over import Over -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.shuffle import Shuffle +from cudf_polars.streaming.io import StreamingSink +from cudf_polars.streaming.join import Join +from cudf_polars.streaming.over import Over +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.shuffle import Shuffle if TYPE_CHECKING: from collections.abc import Iterator diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py index d3952657135..7da5259afd1 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/shuffle.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/shuffle.py @@ -29,17 +29,17 @@ import pylibcudf.partitioning from cudf_polars.dsl.expr import Col -from cudf_polars.experimental.rapidsmpf.dispatch import ( +from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) -from cudf_polars.experimental.rapidsmpf.nodes import shutdown_on_error -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.nodes import shutdown_on_error +from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, _is_already_partitioned, recv_metadata, send_metadata, ) -from cudf_polars.experimental.shuffle import Shuffle +from cudf_polars.streaming.shuffle import Shuffle from cudf_polars.utils.cuda_stream import stream_ordered_after if TYPE_CHECKING: @@ -52,7 +52,7 @@ from rmm.pylibrmm.stream import Stream from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.rapidsmpf.core import SubNetGenerator + from cudf_polars.streaming.actor_graph.core import SubNetGenerator class ShuffleManager: diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/sort.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/sort.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py index 6d63562a08a..75087141c9b 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/collectives/sort.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/collectives/sort.py @@ -25,14 +25,14 @@ from cudf_polars.dsl.expr import Col, NamedExpr from cudf_polars.dsl.ir import Empty, Sort from cudf_polars.dsl.utils.naming import names_to_indices, unique_names -from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager -from cudf_polars.experimental.rapidsmpf.collectives.shuffle import ShuffleManager -from cudf_polars.experimental.rapidsmpf.dispatch import generate_ir_sub_network -from cudf_polars.experimental.rapidsmpf.nodes import ( +from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager +from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager +from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.nodes import ( default_node_single, shutdown_on_error, ) -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, ChunkStore, NormalizedPartitioning, @@ -49,8 +49,8 @@ replay_buffered_channel, send_metadata, ) -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.sort import ( +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.sort import ( _get_final_sort_boundaries, _has_simple_zlice, _select_local_split_candidates, @@ -64,8 +64,8 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator - from cudf_polars.experimental.rapidsmpf.tracing import ActorTracer + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.typing import Schema from cudf_polars.utils.config import StreamingExecutor diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py similarity index 93% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/core.py index 8ec9c9faec5..837384a1561 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/core.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/core.py @@ -19,12 +19,12 @@ Union, ) from cudf_polars.dsl.traversal import CachingVisitor, traversal -from cudf_polars.experimental.over import Over -from cudf_polars.experimental.rapidsmpf.dispatch import FanoutInfo -from cudf_polars.experimental.rapidsmpf.nodes import ( +from cudf_polars.streaming.actor_graph.dispatch import FanoutInfo +from cudf_polars.streaming.actor_graph.nodes import ( generate_ir_sub_network_wrapper, metadata_drain_node, ) +from cudf_polars.streaming.over import Over from cudf_polars.utils.config import SPMDContext if TYPE_CHECKING: @@ -40,12 +40,12 @@ import polars as pl from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.base import PartitionInfo, StatsCollector - from cudf_polars.experimental.parallel import ConfigOptions - from cudf_polars.experimental.rapidsmpf.dispatch import ( + from cudf_polars.streaming.actor_graph.dispatch import ( GenState, SubNetGenerator, ) + from cudf_polars.streaming.base import PartitionInfo, StatsCollector + from cudf_polars.streaming.parallel import ConfigOptions from cudf_polars.utils.config import StreamingExecutor @@ -74,7 +74,7 @@ def evaluate_logical_plan( # For default_singleton, inject the process-wide DefaultSingletonEngine instance # into config_options before treating it as a regular SPMDEngine. if config_options.executor.cluster == "default_singleton": - from cudf_polars.experimental.rapidsmpf.frontend.default_singleton_engine import ( + from cudf_polars.engine.default_singleton_engine import ( DefaultSingletonEngine, ) @@ -97,7 +97,7 @@ def evaluate_logical_plan( ): match config_options.executor.cluster: case "spmd" | "default_singleton": - from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( + from cudf_polars.engine.spmd import ( evaluate_pipeline_spmd_mode, ) @@ -108,7 +108,7 @@ def evaluate_logical_plan( query_id=query_id, ) case "ray": - from cudf_polars.experimental.rapidsmpf.frontend.ray import ( + from cudf_polars.engine.ray import ( evaluate_pipeline_ray_mode, ) @@ -119,7 +119,7 @@ def evaluate_logical_plan( query_id=query_id, ) case "dask": - from cudf_polars.experimental.rapidsmpf.frontend.dask import ( + from cudf_polars.engine.dask import ( evaluate_pipeline_dask_mode, ) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py index a6c9bcbeb08..2554d95fe75 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/dispatch.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/dispatch.py @@ -16,11 +16,11 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.base import ( + from cudf_polars.streaming.actor_graph.utils import ChannelManager + from cudf_polars.streaming.base import ( PartitionInfo, StatsCollector, ) - from cudf_polars.experimental.rapidsmpf.utils import ChannelManager from cudf_polars.utils.config import ConfigOptions, StreamingExecutor diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/groupby.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/groupby.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py index 243255738e3..22be38dade4 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/groupby.py @@ -24,12 +24,11 @@ from cudf_polars.dsl.expr import Col, NamedExpr from cudf_polars.dsl.ir import IR, Distinct, GroupBy, Select from cudf_polars.dsl.utils.naming import unique_names -from cudf_polars.experimental.groupby import combine, decompose -from cudf_polars.experimental.rapidsmpf.collectives.shuffle import ShuffleManager -from cudf_polars.experimental.rapidsmpf.dispatch import ( +from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager +from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, NormalizedPartitioning, _make_hash_shuffle_metadata, @@ -45,7 +44,8 @@ send_metadata, shutdown_on_error, ) -from cudf_polars.experimental.repartition import Repartition +from cudf_polars.streaming.groupby import combine, decompose +from cudf_polars.streaming.repartition import Repartition if TYPE_CHECKING: from rapidsmpf.communicator.communicator import Communicator @@ -53,8 +53,8 @@ from rapidsmpf.streaming.core.channel import Channel from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator - from cudf_polars.experimental.rapidsmpf.tracing import ActorTracer + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.tracing import ActorTracer from cudf_polars.typing import Schema diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py similarity index 91% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/io.py index f99482fb938..78672bb0cab 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/io.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/io.py @@ -27,27 +27,15 @@ _prepare_parquet_predicate, ) from cudf_polars.dsl.to_ast import to_parquet_filter -from cudf_polars.experimental.base import ( - IOPartitionFlavor, - IOPartitionPlan, - PartitionInfo, -) -from cudf_polars.experimental.io import ( - SplitScan, - StreamingSink, - _prepare_sink_directory, - _sink_to_file, - scan_partition_plan, -) -from cudf_polars.experimental.rapidsmpf.dispatch import ( +from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) -from cudf_polars.experimental.rapidsmpf.nodes import ( +from cudf_polars.streaming.actor_graph.nodes import ( define_actor, metadata_feeder_node, shutdown_on_error, ) -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, chunk_to_frame, empty_table_chunk, @@ -56,20 +44,30 @@ recv_metadata, send_metadata, ) -from cudf_polars.experimental.utils import _dynamic_planning_on +from cudf_polars.streaming.base import ( + IOPartitionFlavor, +) +from cudf_polars.streaming.io import ( + SplitScan, + StreamingSink, + _prepare_sink_directory, + _sink_to_file, +) +from cudf_polars.streaming.utils import _dynamic_planning_on if TYPE_CHECKING: - from collections.abc import MutableMapping - from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.base import StatsCollector - from cudf_polars.experimental.dispatch import LowerIRTransformer - from cudf_polars.experimental.rapidsmpf.core import SubNetGenerator - from cudf_polars.experimental.rapidsmpf.tracing import ActorTracer + from cudf_polars.streaming.actor_graph.core import SubNetGenerator + from cudf_polars.streaming.actor_graph.tracing import ActorTracer + from cudf_polars.streaming.base import ( + IOPartitionPlan, + PartitionInfo, + StatsCollector, + ) from cudf_polars.utils.config import ParquetOptions @@ -129,24 +127,6 @@ async def drain(self) -> None: await self.ch_out.drain(self.context) -def lower_dataframescan_rapidsmpf( - ir: DataFrameScan, rec: LowerIRTransformer -) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: - """Lower a DataFrameScan node for the RapidsMPF streaming runtime.""" - config_options = rec.state["config_options"] - - # NOTE: We calculate the expected partition count - # to help trigger fallback warnings in lower_ir_graph. - # The generate_ir_sub_network logic is NOT required - # to obey this partition count. However, the count - # WILL match after an IO operation (for now). - rows_per_partition = config_options.executor.max_rows_per_partition - nrows = max(ir.df.shape()[0], 1) - count = math.ceil(nrows / rows_per_partition) - - return ir, {ir: PartitionInfo(count=count)} - - @define_actor() async def dataframescan_node( context: Context, @@ -317,37 +297,6 @@ def _( return nodes, channels -def lower_scan_rapidsmpf( - ir: Scan, rec: LowerIRTransformer -) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: - """Lower a Scan node for the RapidsMPF streaming runtime.""" - config_options = rec.state["config_options"] - if ( - ir.typ in ("csv", "parquet", "ndjson") - and ir.n_rows == -1 - and ir.skip_rows == 0 - and ir.row_index is None - ): - # NOTE: We calculate the expected partition count - # to help trigger fallback warnings in lower_ir_graph. - # The generate_ir_sub_network logic is NOT required - # to obey this partition count. However, the count - # WILL match after an IO operation (for now). - plan = scan_partition_plan(ir, rec.state["stats"], config_options) - paths = list(ir.paths) - if plan.flavor == IOPartitionFlavor.SPLIT_FILES: - count = plan.factor * len(paths) - else: - count = math.ceil(len(paths) / plan.factor) - - return ir, {ir: PartitionInfo(count=count, io_plan=plan)} - else: - plan = IOPartitionPlan( - flavor=IOPartitionFlavor.SINGLE_READ, factor=len(ir.paths) - ) - return ir, {ir: PartitionInfo(count=1, io_plan=plan)} - - async def read_chunk( context: Context, scan: IR, diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/join.py index 6fd5280b9e4..2c0efb5bf64 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/join.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/join.py @@ -30,13 +30,13 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IR, Join from cudf_polars.dsl.utils.naming import names_to_indices -from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager -from cudf_polars.experimental.rapidsmpf.collectives.shuffle import _global_shuffle -from cudf_polars.experimental.rapidsmpf.dispatch import ( +from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager +from cudf_polars.streaming.actor_graph.collectives.shuffle import _global_shuffle +from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) -from cudf_polars.experimental.rapidsmpf.nodes import default_node_multi -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.nodes import default_node_multi +from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, NormalizedPartitioning, TableSizeStats, @@ -53,8 +53,8 @@ send_metadata, shutdown_on_error, ) -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.utils import _concat +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.utils import _concat if TYPE_CHECKING: from collections.abc import Iterable, MutableMapping @@ -66,9 +66,9 @@ from rapidsmpf.streaming.cudf.bloom_filter import BloomFilterChunk from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.base import PartitionInfo - from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator - from cudf_polars.experimental.rapidsmpf.tracing import ActorTracer + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.tracing import ActorTracer + from cudf_polars.streaming.base import PartitionInfo from cudf_polars.utils.config import StreamingExecutor diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py similarity index 99% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py index 56e6c1ed5ac..9432fd7295d 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/nodes.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/nodes.py @@ -20,10 +20,10 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IR, Empty -from cudf_polars.experimental.rapidsmpf.dispatch import ( +from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, chunkwise_evaluate, empty_table_chunk, @@ -42,7 +42,7 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator @define_actor() diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/over.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/over.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/over.py index c2c72ae9dab..d1569d561b2 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/over.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/over.py @@ -59,13 +59,12 @@ from cudf_polars.dsl.expressions.base import ExecutionContext from cudf_polars.dsl.utils.naming import unique_names from cudf_polars.dsl.utils.reshape import broadcast -from cudf_polars.experimental.over import Over, _build_over_groupby_irs -from cudf_polars.experimental.rapidsmpf.collectives.shuffle import ( +from cudf_polars.streaming.actor_graph.collectives.shuffle import ( LocalRepartitioner, ShuffleManager, ) -from cudf_polars.experimental.rapidsmpf.dispatch import generate_ir_sub_network -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, ChunkStore, NormalizedPartitioning, @@ -86,6 +85,7 @@ send_metadata, shutdown_on_error, ) +from cudf_polars.streaming.over import Over, _build_over_groupby_irs if TYPE_CHECKING: from rapidsmpf.communicator.communicator import Communicator @@ -97,8 +97,8 @@ from cudf_polars.dsl.expr import Col from cudf_polars.dsl.ir import IR, GroupBy, IRExecutionContext, Select - from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator - from cudf_polars.experimental.rapidsmpf.utils import TableSizeStats + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.utils import TableSizeStats @dataclass(frozen=True) diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py index dba7dd5a0a3..fcfc897e172 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/repartition.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/repartition.py @@ -17,17 +17,17 @@ ) from cudf_polars.containers import DataFrame -from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager -from cudf_polars.experimental.rapidsmpf.dispatch import generate_ir_sub_network -from cudf_polars.experimental.rapidsmpf.nodes import shutdown_on_error -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager +from cudf_polars.streaming.actor_graph.dispatch import generate_ir_sub_network +from cudf_polars.streaming.actor_graph.nodes import shutdown_on_error +from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, empty_table_chunk, recv_metadata, send_metadata, ) -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.utils import _concat +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.utils import _concat if TYPE_CHECKING: from rapidsmpf.communicator.communicator import Communicator @@ -35,7 +35,7 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator @define_actor() diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py index c9f39cab068..6560bb8bcc2 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/tracing.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/tracing.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING from cudf_polars.dsl.tracing import LOG_TRACES, Scope -from cudf_polars.experimental.explain import SerializablePlan +from cudf_polars.streaming.explain import SerializablePlan if TYPE_CHECKING: import pylibcudf as plc diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/union.py similarity index 93% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/union.py index df4c6d5dc94..a2664fa3a79 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/union.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/union.py @@ -11,11 +11,11 @@ from rapidsmpf.streaming.cudf.table_chunk import TableChunk from cudf_polars.dsl.ir import Union -from cudf_polars.experimental.rapidsmpf.dispatch import ( +from cudf_polars.streaming.actor_graph.dispatch import ( generate_ir_sub_network, ) -from cudf_polars.experimental.rapidsmpf.nodes import define_actor, shutdown_on_error -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.nodes import define_actor, shutdown_on_error +from cudf_polars.streaming.actor_graph.utils import ( ChannelManager, empty_table_chunk, gather_in_task_group, @@ -30,7 +30,7 @@ from rapidsmpf.streaming.core.context import Context from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.rapidsmpf.core import SubNetGenerator + from cudf_polars.streaming.actor_graph.core import SubNetGenerator @define_actor() diff --git a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py similarity index 99% rename from python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py rename to python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py index 39546950223..ac87b16d587 100644 --- a/python/cudf_polars/cudf_polars/experimental/rapidsmpf/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/actor_graph/utils.py @@ -41,8 +41,9 @@ from cudf_polars.dsl.ir import Cache, Filter, GroupBy, HStack, Join, Projection, Select from cudf_polars.dsl.tracing import Scope from cudf_polars.dsl.utils.naming import names_to_indices -from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager -from cudf_polars.experimental.utils import _concat +from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager +from cudf_polars.streaming.actor_graph.tracing import ActorTracer +from cudf_polars.streaming.utils import _concat from cudf_polars.utils.dtypes import make_empty_column if TYPE_CHECKING: @@ -64,8 +65,7 @@ from rmm.pylibrmm.stream import Stream from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.rapidsmpf.dispatch import SubNetGenerator - from cudf_polars.experimental.rapidsmpf.tracing import ActorTracer + from cudf_polars.streaming.actor_graph.dispatch import SubNetGenerator from cudf_polars.typing import Schema @@ -166,7 +166,6 @@ async def shutdown_on_error( # Create tracer only if LOG_TRACES is enabled and IR is provided tracer: ActorTracer | None = None contextvars: dict[str, Any] = {} - from cudf_polars.experimental.rapidsmpf.tracing import ActorTracer ir_id = trace_ir.get_stable_id() ir_type = type(trace_ir).__name__ diff --git a/python/cudf_polars/cudf_polars/experimental/base.py b/python/cudf_polars/cudf_polars/streaming/base.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/base.py rename to python/cudf_polars/cudf_polars/streaming/base.py index 80ff0dfacbd..9415c1d245a 100644 --- a/python/cudf_polars/cudf_polars/experimental/base.py +++ b/python/cudf_polars/cudf_polars/streaming/base.py @@ -125,7 +125,7 @@ def deserialize(cls, entries: list[SerializedStatsEntry], ir: IR) -> StatsCollec ir Root of the (pre-lowered) IR graph on the local rank. """ - from cudf_polars.experimental.io import DataFrameSourceInfo, ParquetSourceInfo + from cudf_polars.streaming.io import DataFrameSourceInfo, ParquetSourceInfo _deserializers: dict[ str, type[ParquetSourceInfo] | type[DataFrameSourceInfo] diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/__init__.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/__init__.py similarity index 100% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/__init__.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/__init__.py diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/asserts.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py similarity index 100% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/asserts.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/asserts.py diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py index 9d6d7ee5f7c..bf43487bf8d 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds.py @@ -20,7 +20,7 @@ import polars as pl try: - from cudf_polars.experimental.benchmarks.utils import ( + from cudf_polars.streaming.benchmarks.utils import ( COUNT_DTYPE, build_parser, parse_args, @@ -59,7 +59,7 @@ def __getattr__(cls, name: str): # type: ignore[no-untyped-def] if valid_query(name): q_num = int(name[1:]) module: ModuleType = importlib.import_module( - f"cudf_polars.experimental.benchmarks.pdsds_queries.q{q_num}" + f"cudf_polars.streaming.benchmarks.pdsds_queries.q{q_num}" ) return getattr(module, cls.q_impl) raise AttributeError(f"{name} is not a valid query name") @@ -77,7 +77,7 @@ class PDSDSPolarsQueries(PDSDSQueries): q_impl = "polars_impl" # See comments for EXPECTED_CASTS and EXPECTED_CASTS_DECIMAL - # in cudf/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py + # in cudf/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py # for more details. EXPECTED_CASTS_DECIMAL: ClassVar[dict] = { 2: [ diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_parameters/README.md b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_parameters/README.md similarity index 99% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_parameters/README.md rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_parameters/README.md index 6dd6bf2ab36..04081d33361 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_parameters/README.md +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_parameters/README.md @@ -119,7 +119,7 @@ Pre-generated parameters for PDS-DS queries across a set of scale factors. ### Python API ```python -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters # Load randomly generated parameters for a specific query and scale factor params = load_parameters(scale_factor=100, query_id=1) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_parameters/__init__.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_parameters/__init__.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_parameters/__init__.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_parameters/__init__.py index 71c3300b078..fcd68e54120 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_parameters/__init__.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_parameters/__init__.py @@ -7,7 +7,7 @@ from typing import Any -from cudf_polars.experimental.benchmarks.pdsds_parameters.parameter_substitutions import ( +from cudf_polars.streaming.benchmarks.pdsds_parameters.parameter_substitutions import ( PARAMETERS, ) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_parameters/parameter_substitutions.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_parameters/parameter_substitutions.py similarity index 100% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_parameters/parameter_substitutions.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_parameters/parameter_substitutions.py diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/__init__.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/__init__.py similarity index 100% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/__init__.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/__init__.py diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q1.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q1.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q1.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q1.py index 81fd42ea30e..a5aebfa87dd 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q1.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q1.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q10.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q10.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q10.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q10.py index c0eaa5ed677..e238554907f 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q10.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q10.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q11.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q11.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q11.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q11.py index 5b5b17cb967..98c15bfe674 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q11.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q11.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q12.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q12.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q12.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q12.py index 76f8e47fc09..8ca5b6cbeb5 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q12.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q12.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q13.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q13.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q13.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q13.py index a31639aa833..e7c0366aab6 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q13.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q13.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q14.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q14.py similarity index 93% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q14.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q14.py index 57946e3044d..bc8010f1a61 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q14.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q14.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -335,65 +335,66 @@ def polars_impl(run_config: RunConfig) -> QueryResult: item = get_data(run_config.dataset_path, "item", run_config.suffix) date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix) + cross_items = build_cross_items( + store_sales, catalog_sales, web_sales, item, date_dim, year=year + ) + average_sales = build_average_sales( + store_sales, catalog_sales, web_sales, date_dim, year=year + ) + + # week_dates is ≤7 rows (one calendar week), computed once as a 1-partition frame. + # Push the week filter into each channel before the UNION via a semi-join so that + # ~99% of rows (everything outside the target week) are dropped before the + # expensive cross_items join and groupby. + target_week = ( + date_dim.filter( + (pl.col("d_year") == year + 1) + & (pl.col("d_moy") == 12) + & (pl.col("d_dom") == day) + ) + .select("d_week_seq") + .unique() + ) + week_dates = date_dim.join(target_week, on="d_week_seq").select("d_date_sk") + all_sales = pl.concat( [ - store_sales.select( + store_sales.join( + week_dates, left_on="ss_sold_date_sk", right_on="d_date_sk", how="semi" + ).select( [ pl.lit("store").alias("channel"), pl.col("ss_item_sk").alias("item_sk"), pl.col("ss_quantity").alias("quantity"), pl.col("ss_list_price").alias("list_price"), - pl.col("ss_sold_date_sk").alias("date_sk"), ] ), - catalog_sales.select( + catalog_sales.join( + week_dates, left_on="cs_sold_date_sk", right_on="d_date_sk", how="semi" + ).select( [ pl.lit("catalog").alias("channel"), pl.col("cs_item_sk").alias("item_sk"), pl.col("cs_quantity").alias("quantity"), pl.col("cs_list_price").alias("list_price"), - pl.col("cs_sold_date_sk").alias("date_sk"), ] ), - web_sales.select( + web_sales.join( + week_dates, left_on="ws_sold_date_sk", right_on="d_date_sk", how="semi" + ).select( [ pl.lit("web").alias("channel"), pl.col("ws_item_sk").alias("item_sk"), pl.col("ws_quantity").alias("quantity"), pl.col("ws_list_price").alias("list_price"), - pl.col("ws_sold_date_sk").alias("date_sk"), ] ), ] ) - cross_items = build_cross_items( - store_sales, catalog_sales, web_sales, item, date_dim, year=year - ) - average_sales = build_average_sales( - store_sales, catalog_sales, web_sales, date_dim, year=year - ) - - # d_week_seq target is the same for all 3 channels; compute it once. - target_week = ( - date_dim.filter( - (pl.col("d_year") == year + 1) - & (pl.col("d_moy") == 12) - & (pl.col("d_dom") == day) - ) - .select("d_week_seq") - .unique() - ) - week_dates = date_dim.join(target_week, on="d_week_seq").select("d_date_sk") - - # Build y: all 3 channels in a single pipeline. - # cross_items and average_sales each appear once — no CSE needed. - # After group_by the frame is tiny, so the cross join with the 1-row - # average_sales frame is negligible even if Polars fuses it into an IEJoin. y = ( all_sales.join(cross_items, left_on="item_sk", right_on="ss_item_sk") .join(item, left_on="item_sk", right_on="i_item_sk") - .join(week_dates, left_on="date_sk", right_on="d_date_sk") .group_by(["channel", "i_brand_id", "i_class_id", "i_category_id"]) .agg( [ diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q15.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q15.py similarity index 93% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q15.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q15.py index fa0f8938862..31440d311cb 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q15.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q15.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q16.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q16.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q16.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q16.py index 4f91f854318..81fc272a4a9 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q16.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q16.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q17.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q17.py similarity index 65% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q17.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q17.py index 997a1546d47..cb503ffdaba 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q17.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q17.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -110,38 +110,103 @@ def polars_impl(run_config: RunConfig) -> QueryResult: sort_by = {"i_item_id": False, "i_item_desc": False, "s_state": False} limit = 100 - store_sales_base = ( + q1 = f"{year}Q1" + q1_q3 = [f"{year}Q1", f"{year}Q2", f"{year}Q3"] + + # Pre-filter date_dim to only qualifying d_date_sk values. + d1_dates = date_dim.filter(pl.col("d_quarter_name") == q1).select("d_date_sk") + d_q3_dates = date_dim.filter(pl.col("d_quarter_name").is_in(q1_q3)).select( + "d_date_sk" + ) + + # store_returns has [6] partitions — at the broadcast limit. Filter it to Q1-Q3 dates + # first, then use the (customer, item) pairs it contains to pre-filter both store_sales + # and catalog_sales before those larger tables enter the expensive shuffle joins. + store_returns_filtered = store_returns.join( + d_q3_dates, left_on="sr_returned_date_sk", right_on="d_date_sk", how="semi" + ).select(["sr_customer_sk", "sr_item_sk", "sr_ticket_number", "sr_return_quantity"]) + + # (customer, item) pairs present in any qualifying store return; stays at [6] partitions + # so broadcast is free. Polars will CACHE this shared subplan. + sr_customer_item = store_returns_filtered.select(["sr_customer_sk", "sr_item_sk"]) + + store_sales_filtered = ( store_sales.join( - date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk", suffix="_d1" + d1_dates, left_on="ss_sold_date_sk", right_on="d_date_sk", how="semi" + ) + .join( + sr_customer_item, + left_on=["ss_customer_sk", "ss_item_sk"], + right_on=["sr_customer_sk", "sr_item_sk"], + how="semi", + ) + .select( + [ + "ss_customer_sk", + "ss_item_sk", + "ss_store_sk", + "ss_ticket_number", + "ss_quantity", + ] + ) + .join( + item.select(["i_item_sk", "i_item_id", "i_item_desc"]), + left_on="ss_item_sk", + right_on="i_item_sk", + ) + .join( + store.select(["s_store_sk", "s_state"]), + left_on="ss_store_sk", + right_on="s_store_sk", + ) + .select( + [ + "ss_customer_sk", + "ss_item_sk", + "ss_ticket_number", + "ss_quantity", + "i_item_id", + "i_item_desc", + "s_state", + ] ) - .join(item, left_on="ss_item_sk", right_on="i_item_sk") - .join(store, left_on="ss_store_sk", right_on="s_store_sk") - .filter(pl.col("d_quarter_name") == f"{year}Q1") ) - store_returns_base = store_returns.join( - date_dim, left_on="sr_returned_date_sk", right_on="d_date_sk", suffix="_d2" - ).filter(pl.col("d_quarter_name").is_in([f"{year}Q1", f"{year}Q2", f"{year}Q3"])) - - catalog_sales_base = catalog_sales.join( - date_dim, left_on="cs_sold_date_sk", right_on="d_date_sk", suffix="_d3" - ).filter(pl.col("d_quarter_name").is_in([f"{year}Q1", f"{year}Q2", f"{year}Q3"])) + catalog_sales_filtered = ( + catalog_sales.join( + d_q3_dates, left_on="cs_sold_date_sk", right_on="d_date_sk", how="semi" + ) + .join( + sr_customer_item, + left_on=["cs_bill_customer_sk", "cs_item_sk"], + right_on=["sr_customer_sk", "sr_item_sk"], + how="semi", + ) + .select(["cs_bill_customer_sk", "cs_item_sk", "cs_quantity"]) + ) return QueryResult( frame=( - store_sales_base.join( - store_returns_base, + store_sales_filtered.join( + store_returns_filtered, left_on=["ss_customer_sk", "ss_item_sk", "ss_ticket_number"], right_on=["sr_customer_sk", "sr_item_sk", "sr_ticket_number"], - how="inner", - suffix="_sr", + ) + .select( + [ + "ss_customer_sk", + "ss_item_sk", + "ss_quantity", + "sr_return_quantity", + "i_item_id", + "i_item_desc", + "s_state", + ] ) .join( - catalog_sales_base, + catalog_sales_filtered, left_on=["ss_customer_sk", "ss_item_sk"], right_on=["cs_bill_customer_sk", "cs_item_sk"], - how="inner", - suffix="_cs", ) .group_by(["i_item_id", "i_item_desc", "s_state"]) .agg( diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q18.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q18.py similarity index 73% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q18.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q18.py index 9c2b9f227ef..9c0784cbd05 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q18.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q18.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -121,10 +121,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: catalog_sales = get_data( run_config.dataset_path, "catalog_sales", run_config.suffix ) - customer_demographics_1 = get_data( - run_config.dataset_path, "customer_demographics", run_config.suffix - ) - customer_demographics_2 = get_data( + customer_demographics = get_data( run_config.dataset_path, "customer_demographics", run_config.suffix ) customer = get_data(run_config.dataset_path, "customer", run_config.suffix) @@ -134,30 +131,52 @@ def polars_impl(run_config: RunConfig) -> QueryResult: date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix) item = get_data(run_config.dataset_path, "item", run_config.suffix) + # Pre-filter each dimension table before joining against catalog_sales [45 partitions]. + # d_year not in GROUP BY — semi-join keeps only the date key in the pipeline. + filtered_dates = date_dim.filter(pl.col("d_year") == year).select("d_date_sk") + filtered_cd1 = customer_demographics.filter( + (pl.col("cd_gender") == gen) & (pl.col("cd_education_status") == es) + ).select(["cd_demo_sk", "cd_dep_count"]) + filtered_customer = customer.filter(pl.col("c_birth_month").is_in(month)).select( + ["c_customer_sk", "c_current_cdemo_sk", "c_current_addr_sk", "c_birth_year"] + ) + filtered_addr = customer_address.filter(pl.col("ca_state").is_in(state)).select( + ["ca_address_sk", "ca_county", "ca_state", "ca_country"] + ) + base_query = ( - catalog_sales.join(date_dim, left_on="cs_sold_date_sk", right_on="d_date_sk") - .join(item, left_on="cs_item_sk", right_on="i_item_sk") + catalog_sales.select( + [ + "cs_sold_date_sk", + "cs_item_sk", + "cs_bill_cdemo_sk", + "cs_bill_customer_sk", + "cs_quantity", + "cs_list_price", + "cs_coupon_amt", + "cs_sales_price", + "cs_net_profit", + ] + ) .join( - customer_demographics_1, - left_on="cs_bill_cdemo_sk", - right_on="cd_demo_sk", - suffix="_cd1", + filtered_dates, left_on="cs_sold_date_sk", right_on="d_date_sk", how="semi" + ) + .join( + item.select(["i_item_sk", "i_item_id"]), + left_on="cs_item_sk", + right_on="i_item_sk", ) - .join(customer, left_on="cs_bill_customer_sk", right_on="c_customer_sk") + .join(filtered_cd1, left_on="cs_bill_cdemo_sk", right_on="cd_demo_sk") .join( - customer_demographics_2, + filtered_customer, left_on="cs_bill_customer_sk", right_on="c_customer_sk" + ) + .join( + customer_demographics.select("cd_demo_sk"), left_on="c_current_cdemo_sk", right_on="cd_demo_sk", - suffix="_cd2", - ) - .join(customer_address, left_on="c_current_addr_sk", right_on="ca_address_sk") - .filter( - (pl.col("cd_gender") == gen) - & (pl.col("cd_education_status") == es) - & pl.col("c_birth_month").is_in(month) - & (pl.col("d_year") == year) - & pl.col("ca_state").is_in(state) + how="semi", ) + .join(filtered_addr, left_on="c_current_addr_sk", right_on="ca_address_sk") ) agg_exprs = [ diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q19.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q19.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q19.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q19.py index 0a971dd0dcd..334429b2a0f 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q19.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q19.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q2.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q2.py similarity index 88% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q2.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q2.py index 998cccb6017..2753c64d5c6 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q2.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q2.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -136,10 +136,6 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ), ] ) - # Step 2: Create wswscs CTE equivalent (aggregate by week and day of week) - # First join with date_dim to get day names - wscs_with_dates = wscs.join(date_dim, left_on="sold_date_sk", right_on="d_date_sk") - # Create separate aggregations for each day to better control null handling days = ( "Sunday", "Monday", @@ -158,35 +154,26 @@ def polars_impl(run_config: RunConfig) -> QueryResult: "fri_sales", "sat_sales", ) - # Start with all week sequences - all_weeks = wscs_with_dates.select("d_week_seq").unique() - wswscs = all_weeks - + # Pre-filter date_dim to 4 years ([year-1, year, year+1, year+2]) to capture + # boundary weeks that span year transitions (e.g. from Dec 28 to Jan 3). Filtering to + # only [year, year+1] incorrectly excludes Dec days whose d_week_seq also + # appears in year's date_dim, producing null sales for those boundary weeks. + date_dim_prefilter = date_dim.filter( + pl.col("d_year").is_in([year - 1, year, year + 1, year + 2]) + ).select(["d_date_sk", "d_week_seq", "d_day_name"]) wswscs = ( - wscs_with_dates.with_columns( + wscs.join(date_dim_prefilter, left_on="sold_date_sk", right_on="d_date_sk") + .group_by("d_week_seq") + .agg( [ pl.when(pl.col("d_day_name") == day) .then(pl.col("sales_price")) .otherwise(None) + .sum() .alias(name) for day, name in zip(days, day_cols, strict=True) ] ) - .group_by("d_week_seq") - .agg( - *(pl.col(name).sum().alias(name) for name in day_cols), - *(pl.col(name).count().alias(f"{name}_count") for name in day_cols), - ) - .with_columns( - [ - pl.when(pl.col(f"{name}_count") > 0) - .then(pl.col(name)) - .otherwise(None) - .alias(name) - for name in day_cols - ] - ) - .select(["d_week_seq", *day_cols]) ) # Step 3: Create year data (y subquery equivalent) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q20.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q20.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py index 08140a708e5..c8282098b01 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q20.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q20.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q21.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q21.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q21.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q21.py index b264ff92d93..f96b9f12df2 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q21.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q21.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q22.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q22.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q22.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q22.py index 5592608416c..2151ab0cd62 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q22.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q22.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q23.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q23.py similarity index 83% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q23.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q23.py index debe1b512b3..01304a7683f 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q23.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q23.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -105,10 +105,16 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ) web_sales = get_data(run_config.dataset_path, "web_sales", run_config.suffix) + # Pre-filter date_dim to the 4-year window so the inner join with store_sales + # naturally excludes out-of-window records before the expensive item join. + year_set = [year, year + 1, year + 2, year + 3] + date_dim_years = date_dim.filter(pl.col("d_year").is_in(year_set)) + frequent_ss_items = ( - store_sales.join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") + store_sales.join( + date_dim_years, left_on="ss_sold_date_sk", right_on="d_date_sk" + ) .join(item, left_on="ss_item_sk", right_on="i_item_sk") - .filter(pl.col("d_year").is_in([year, year + 1, year + 2, year + 3])) .with_columns(pl.col("i_item_desc").str.slice(0, 30).alias("itemdesc")) .group_by(["itemdesc", "ss_item_sk", "d_date"]) .agg(pl.len().alias("cnt")) @@ -121,8 +127,11 @@ def polars_impl(run_config: RunConfig) -> QueryResult: # only valid because we know that the TPC-DS includes a foreign key here, so all # customers in store_sales _must_ be entries that exist somewhere in customer. store_sales.filter(pl.col("ss_customer_sk").is_not_null()) - .join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") - .filter(pl.col("d_year").is_in([year, year + 1, year + 2, year + 3])) + .join( + date_dim_years.select("d_date_sk"), + left_on="ss_sold_date_sk", + right_on="d_date_sk", + ) .group_by("ss_customer_sk") .agg((pl.col("ss_quantity") * pl.col("ss_sales_price")).sum().alias("csales")) ) @@ -146,13 +155,12 @@ def polars_impl(run_config: RunConfig) -> QueryResult: (pl.col("d_year") == year) & (pl.col("d_moy") == month) ).select("d_date_sk") + # Join order: most selective filters first (date_target ~1.2%, frequent_ss_items, + # best_customers semi ~5%), then customer last — it's a non-filtering name lookup + # that only adds c_last_name/c_first_name, so running it on the already-reduced + # row set avoids the full catalog_sales/web_sales scan width. catalog_part = ( - catalog_sales.join( - customer.select(["c_customer_sk", "c_last_name", "c_first_name"]), - left_on="cs_bill_customer_sk", - right_on="c_customer_sk", - ) - .join(date_target, left_on="cs_sold_date_sk", right_on="d_date_sk") + catalog_sales.join(date_target, left_on="cs_sold_date_sk", right_on="d_date_sk") .join(frequent_ss_items, left_on="cs_item_sk", right_on="ss_item_sk") .join( best_customers, @@ -160,17 +168,17 @@ def polars_impl(run_config: RunConfig) -> QueryResult: right_on="ss_customer_sk", how="semi", ) + .join( + customer.select(["c_customer_sk", "c_last_name", "c_first_name"]), + left_on="cs_bill_customer_sk", + right_on="c_customer_sk", + ) .group_by(["c_last_name", "c_first_name"]) .agg((pl.col("cs_quantity") * pl.col("cs_list_price")).sum().alias("sales")) ) web_part = ( - web_sales.join( - customer.select(["c_customer_sk", "c_last_name", "c_first_name"]), - left_on="ws_bill_customer_sk", - right_on="c_customer_sk", - ) - .join(date_target, left_on="ws_sold_date_sk", right_on="d_date_sk") + web_sales.join(date_target, left_on="ws_sold_date_sk", right_on="d_date_sk") .join(frequent_ss_items, left_on="ws_item_sk", right_on="ss_item_sk") .join( best_customers, @@ -178,6 +186,11 @@ def polars_impl(run_config: RunConfig) -> QueryResult: right_on="ss_customer_sk", how="semi", ) + .join( + customer.select(["c_customer_sk", "c_last_name", "c_first_name"]), + left_on="ws_bill_customer_sk", + right_on="c_customer_sk", + ) .group_by(["c_last_name", "c_first_name"]) .agg((pl.col("ws_quantity") * pl.col("ws_list_price")).sum().alias("sales")) ) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q24.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q24.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q24.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q24.py index 6cde550aea0..fe506bf767c 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q24.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q24.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q25.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q25.py similarity index 56% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q25.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q25.py index 8006585a170..603785ef77f 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q25.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q25.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -96,17 +96,6 @@ def polars_impl(run_config: RunConfig) -> QueryResult: store = get_data(run_config.dataset_path, "store", run_config.suffix) item = get_data(run_config.dataset_path, "item", run_config.suffix) - d1, d2, d3 = [ - date_dim.clone().select( - [ - pl.col("d_date_sk").alias(f"{p}_date_sk"), - pl.col("d_moy").alias(f"{p}_moy"), - pl.col("d_year").alias(f"{p}_year"), - ] - ) - for p in ("d1", "d2", "d3") - ] - sort_by = { "i_item_id": False, "i_item_desc": False, @@ -114,31 +103,103 @@ def polars_impl(run_config: RunConfig) -> QueryResult: "s_store_name": False, } limit = 100 + + # d1: only April of the target year — very selective (~1/60 of date_dim rows). + # d2/d3: from April to October of the target year: same condition, one pre-filtered frame. + d1_dates = date_dim.filter( + (pl.col("d_moy") == 4) & (pl.col("d_year") == year) + ).select("d_date_sk") + d2_d3_dates = date_dim.filter( + pl.col("d_moy").is_between(4, 10) & (pl.col("d_year") == year) + ).select("d_date_sk") + + # store_returns [6] ≤ broadcast limit: filter to qualifying return dates first, + # then extract (customer, item) pairs to pre-filter ss and cs before shuffle joins. + store_returns_filtered = store_returns.join( + d2_d3_dates, left_on="sr_returned_date_sk", right_on="d_date_sk", how="semi" + ).select(["sr_customer_sk", "sr_item_sk", "sr_ticket_number", "sr_net_loss"]) + sr_customer_item = store_returns_filtered.select(["sr_customer_sk", "sr_item_sk"]) + + store_sales_filtered = ( + store_sales.join( + d1_dates, left_on="ss_sold_date_sk", right_on="d_date_sk", how="semi" + ) + .join( + sr_customer_item, + left_on=["ss_customer_sk", "ss_item_sk"], + right_on=["sr_customer_sk", "sr_item_sk"], + how="semi", + ) + .select( + [ + "ss_customer_sk", + "ss_item_sk", + "ss_store_sk", + "ss_ticket_number", + "ss_net_profit", + ] + ) + .join( + item.select(["i_item_sk", "i_item_id", "i_item_desc"]), + left_on="ss_item_sk", + right_on="i_item_sk", + ) + .join( + store.select(["s_store_sk", "s_store_id", "s_store_name"]), + left_on="ss_store_sk", + right_on="s_store_sk", + ) + .select( + [ + "ss_customer_sk", + "ss_item_sk", + "ss_ticket_number", + "ss_net_profit", + "i_item_id", + "i_item_desc", + "s_store_id", + "s_store_name", + ] + ) + ) + + catalog_sales_filtered = ( + catalog_sales.join( + d2_d3_dates, left_on="cs_sold_date_sk", right_on="d_date_sk", how="semi" + ) + .join( + sr_customer_item, + left_on=["cs_bill_customer_sk", "cs_item_sk"], + right_on=["sr_customer_sk", "sr_item_sk"], + how="semi", + ) + .select(["cs_bill_customer_sk", "cs_item_sk", "cs_net_profit"]) + ) + return QueryResult( frame=( - store_sales.join(d1, left_on="ss_sold_date_sk", right_on="d1_date_sk") - .join(item, left_on="ss_item_sk", right_on="i_item_sk") - .join(store, left_on="ss_store_sk", right_on="s_store_sk") - .join( - store_returns, + store_sales_filtered.join( + store_returns_filtered, left_on=["ss_customer_sk", "ss_item_sk", "ss_ticket_number"], right_on=["sr_customer_sk", "sr_item_sk", "sr_ticket_number"], ) - .join(d2, left_on="sr_returned_date_sk", right_on="d2_date_sk") + .select( + [ + "ss_customer_sk", + "ss_item_sk", + "ss_net_profit", + "sr_net_loss", + "i_item_id", + "i_item_desc", + "s_store_id", + "s_store_name", + ] + ) .join( - catalog_sales, + catalog_sales_filtered, left_on=["ss_customer_sk", "ss_item_sk"], right_on=["cs_bill_customer_sk", "cs_item_sk"], ) - .join(d3, left_on="cs_sold_date_sk", right_on="d3_date_sk") - .filter( - (pl.col("d1_moy") == 4) - & (pl.col("d1_year") == year) - & (pl.col("d2_moy").is_between(4, 10)) - & (pl.col("d2_year") == year) - & (pl.col("d3_moy").is_between(4, 10)) - & (pl.col("d3_year") == year) - ) .group_by(["i_item_id", "i_item_desc", "s_store_id", "s_store_name"]) .agg( [ diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q26.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q26.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q26.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q26.py index 4861f6fb21c..d6cad5c4171 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q26.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q26.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q27.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q27.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q27.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q27.py index 073a4c090af..fcd8dd5c2e6 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q27.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q27.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q28.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q28.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q28.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q28.py index 0b8a8f68f6a..d1938fb8644 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q28.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q28.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q29.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q29.py similarity index 56% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q29.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q29.py index a39954745bd..9abbdcd79b8 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q29.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q29.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -97,23 +97,6 @@ def polars_impl(run_config: RunConfig) -> QueryResult: store = get_data(run_config.dataset_path, "store", run_config.suffix) item = get_data(run_config.dataset_path, "item", run_config.suffix) - d1, d2 = [ - date_dim.clone().select( - [ - pl.col("d_date_sk").alias(f"{p}_date_sk"), - pl.col("d_moy").alias(f"{p}_moy"), - pl.col("d_year").alias(f"{p}_year"), - ] - ) - for p in ("d1", "d2") - ] - d3 = date_dim.clone().select( - [ - pl.col("d_date_sk").alias("d3_date_sk"), - pl.col("d_year").alias("d3_year"), - ] - ) - sort_by = { "i_item_id": False, "i_item_desc": False, @@ -121,30 +104,107 @@ def polars_impl(run_config: RunConfig) -> QueryResult: "s_store_name": False, } limit = 100 + + # d1: one specific month of the target year — most selective filter. + # d2: 4-month window of the target year. + # d3: 3-year window — less selective but still worth pushing before the cs shuffle join. + d1_dates = date_dim.filter( + (pl.col("d_moy") == month) & (pl.col("d_year") == year) + ).select("d_date_sk") + d2_dates = date_dim.filter( + pl.col("d_moy").is_between(month, month + 3) & (pl.col("d_year") == year) + ).select("d_date_sk") + d3_dates = date_dim.filter( + pl.col("d_year").is_in([year, year + 1, year + 2]) + ).select("d_date_sk") + + # store_returns [6] ≤ broadcast limit: apply d2 date filter, then use + # (customer, item) pairs to pre-filter ss and cs before shuffle joins. + store_returns_filtered = store_returns.join( + d2_dates, left_on="sr_returned_date_sk", right_on="d_date_sk", how="semi" + ).select(["sr_customer_sk", "sr_item_sk", "sr_ticket_number", "sr_return_quantity"]) + sr_customer_item = store_returns_filtered.select(["sr_customer_sk", "sr_item_sk"]) + + store_sales_filtered = ( + store_sales.join( + d1_dates, left_on="ss_sold_date_sk", right_on="d_date_sk", how="semi" + ) + .join( + sr_customer_item, + left_on=["ss_customer_sk", "ss_item_sk"], + right_on=["sr_customer_sk", "sr_item_sk"], + how="semi", + ) + .select( + [ + "ss_customer_sk", + "ss_item_sk", + "ss_store_sk", + "ss_ticket_number", + "ss_quantity", + ] + ) + .join( + item.select(["i_item_sk", "i_item_id", "i_item_desc"]), + left_on="ss_item_sk", + right_on="i_item_sk", + ) + .join( + store.select(["s_store_sk", "s_store_id", "s_store_name"]), + left_on="ss_store_sk", + right_on="s_store_sk", + ) + .select( + [ + "ss_customer_sk", + "ss_item_sk", + "ss_ticket_number", + "ss_quantity", + "i_item_id", + "i_item_desc", + "s_store_id", + "s_store_name", + ] + ) + ) + + catalog_sales_filtered = ( + catalog_sales.join( + d3_dates, left_on="cs_sold_date_sk", right_on="d_date_sk", how="semi" + ) + .join( + sr_customer_item, + left_on=["cs_bill_customer_sk", "cs_item_sk"], + right_on=["sr_customer_sk", "sr_item_sk"], + how="semi", + ) + .select(["cs_bill_customer_sk", "cs_item_sk", "cs_quantity"]) + ) + return QueryResult( frame=( - store_sales.join(d1, left_on="ss_sold_date_sk", right_on="d1_date_sk") - .join(item, left_on="ss_item_sk", right_on="i_item_sk") - .join(store, left_on="ss_store_sk", right_on="s_store_sk") - .join( - store_returns, + store_sales_filtered.join( + store_returns_filtered, left_on=["ss_customer_sk", "ss_item_sk", "ss_ticket_number"], right_on=["sr_customer_sk", "sr_item_sk", "sr_ticket_number"], ) - .join(d2, left_on="sr_returned_date_sk", right_on="d2_date_sk") + .select( + [ + "ss_customer_sk", + "ss_item_sk", + "ss_quantity", + "sr_return_quantity", + "i_item_id", + "i_item_desc", + "s_store_id", + "s_store_name", + ] + ) .join( - catalog_sales, + catalog_sales_filtered, left_on=["ss_customer_sk", "ss_item_sk"], right_on=["cs_bill_customer_sk", "cs_item_sk"], ) - .join(d3, left_on="cs_sold_date_sk", right_on="d3_date_sk") - .filter( - (pl.col("d1_moy") == month) - & (pl.col("d1_year") == year) - & (pl.col("d2_moy").is_between(month, month + 3)) - & (pl.col("d2_year") == year) - & (pl.col("d3_year").is_in([year, year + 1, year + 2])) - ) .group_by(["i_item_id", "i_item_desc", "s_store_id", "s_store_name"]) .agg( [ diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q3.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q3.py similarity index 91% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q3.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q3.py index cb2c46a4fee..5a7b7251d9a 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q3.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q3.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q30.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q30.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q30.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q30.py index c00742696fe..ab288319da9 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q30.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q30.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q31.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q31.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py index e50aa9ba1b1..95351f103b6 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q31.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q31.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q32.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q32.py similarity index 93% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q32.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q32.py index 4a804e42336..9134f3366c2 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q32.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q32.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q33.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q33.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q33.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q33.py index 2b448284852..6d9882c0a43 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q33.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q33.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q34.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q34.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q34.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q34.py index 6113775cd8d..bfaf81dff00 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q34.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q34.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q35.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q35.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q35.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q35.py index 96ef2987bb8..9cb77a9e9f2 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q35.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q35.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q36.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q36.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q36.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q36.py index 56b4ef9c92c..0ad51ccbba8 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q36.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q36.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q37.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q37.py similarity index 93% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q37.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q37.py index 3e2bcbcfd76..08c96b3b87a 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q37.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q37.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q38.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q38.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q38.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q38.py index 122cac89704..cd1a091e39b 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q38.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q38.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q39.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q39.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q39.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q39.py index 56c3348b354..db1081a565f 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q39.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q39.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q4.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q4.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q4.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q4.py index df5060fb414..33edfcd14e3 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q4.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q4.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q40.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q40.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q40.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q40.py index 31570a956c5..0399de2dc78 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q40.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q40.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q41.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q41.py similarity index 93% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q41.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q41.py index 6010ecdcd45..715bc244f6e 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q41.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q41.py @@ -11,11 +11,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig Q_NUM = 41 diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q42.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q42.py similarity index 92% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q42.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q42.py index dbcecf5fc58..0a6ba9d374e 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q42.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q42.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q43.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q43.py similarity index 85% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q43.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q43.py index b0c1023d655..4e76911e53b 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q43.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q43.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -91,10 +91,20 @@ def polars_impl(run_config: RunConfig) -> QueryResult: year = params["year"] gmt = params["gmt"] - # Load tables date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix) store_sales = get_data(run_config.dataset_path, "store_sales", run_config.suffix) store = get_data(run_config.dataset_path, "store", run_config.suffix) + + # Pre-filter lookup tables before joining against store_sales [58 partitions]. + # d_year not needed after filter; d_day_name drives the conditional agg columns. + filtered_dates = date_dim.filter(pl.col("d_year") == year).select( + ["d_date_sk", "d_day_name"] + ) + # s_gmt_offset not needed after filter; keep group-by output columns. + filtered_store = store.filter(pl.col("s_gmt_offset") == gmt).select( + ["s_store_sk", "s_store_name", "s_store_id"] + ) + sort_by = { "s_store_name": False, "s_store_id": False, @@ -107,15 +117,13 @@ def polars_impl(run_config: RunConfig) -> QueryResult: "sat_sales": False, } limit = 100 - # Main query with joins and conditional aggregations return QueryResult( frame=( - store_sales.join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") - .join(store, left_on="ss_store_sk", right_on="s_store_sk") - .filter((pl.col("s_gmt_offset") == gmt) & (pl.col("d_year") == year)) + store_sales.select(["ss_sold_date_sk", "ss_store_sk", "ss_sales_price"]) + .join(filtered_dates, left_on="ss_sold_date_sk", right_on="d_date_sk") + .join(filtered_store, left_on="ss_store_sk", right_on="s_store_sk") .with_columns( [ - # Pre-compute conditional sales amounts for each day pl.when(pl.col("d_day_name") == "Sunday") .then(pl.col("ss_sales_price")) .otherwise(0) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q44.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q44.py similarity index 67% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q44.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q44.py index 204c9c90a76..3232dffe5f2 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q44.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q44.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -90,87 +90,55 @@ def polars_impl(run_config: RunConfig) -> QueryResult: store_sk = params["store_sk"] - # Load tables store_sales = get_data(run_config.dataset_path, "store_sales", run_config.suffix) item = get_data(run_config.dataset_path, "item", run_config.suffix) - # Step 1: Calculate benchmark (average profit for store with null demographics) + # Benchmark: global mean profit for the store with null demographics — single row. + # Use a constant-key equi-join instead of how="cross" so the streaming executor + # treats it as a broadcast join (1 row ≤ broadcast_join_limit) rather than a + # ConditionalJoin that falls back from multi-partition mode. benchmark = ( store_sales.filter( - (pl.col("ss_store_sk") == store_sk) & (pl.col("ss_cdemo_sk").is_null()) + (pl.col("ss_store_sk") == store_sk) & pl.col("ss_cdemo_sk").is_null() ) - .group_by("ss_store_sk") - .agg( - [ - pl.col("ss_net_profit").mean().alias("profit_mean"), - pl.col("ss_net_profit").count().alias("profit_count"), - ] - ) - .with_columns( - [ - pl.when(pl.col("profit_count") > 0) - .then(pl.col("profit_mean")) - .otherwise(None) - .alias("benchmark_profit") - ] - ) - .select("benchmark_profit") + .select(pl.col("ss_net_profit").mean().alias("benchmark_profit")) + .with_columns(pl.lit(1, dtype=pl.Int32).alias("_key")) ) - # Step 2: Calculate item-level average profits for store + # Item-level average profits, broadcast-joined with the 1-row benchmark. item_profits = ( store_sales.filter(pl.col("ss_store_sk") == store_sk) .group_by("ss_item_sk") - .agg( - [ - pl.col("ss_net_profit").mean().alias("profit_mean"), - pl.col("ss_net_profit").count().alias("profit_count"), - ] - ) - .with_columns( - [ - pl.when(pl.col("profit_count") > 0) - .then(pl.col("profit_mean")) - .otherwise(None) - .alias("avg(ss_net_profit)") - ] - ) - .drop(["profit_mean", "profit_count"]) - .join(benchmark, how="cross") - .filter(pl.col("avg(ss_net_profit)") > (0.9 * pl.col("benchmark_profit"))) + .agg(pl.col("ss_net_profit").mean().alias("avg_profit")) + .with_columns(pl.lit(1, dtype=pl.Int32).alias("_key")) + .join(benchmark, on="_key") + .filter(pl.col("avg_profit") > 0.9 * pl.col("benchmark_profit")) + .select(["ss_item_sk", "avg_profit"]) ) - # Step 3: Create ascending ranking (worst to best) ascending_rank = ( - item_profits.with_columns( - [pl.col("avg(ss_net_profit)").rank(method="ordinal").alias("rnk")] - ) + item_profits.with_columns(pl.col("avg_profit").rank(method="min").alias("rnk")) .filter(pl.col("rnk") < 11) .select(["ss_item_sk", "rnk"]) ) - # Step 4: Create descending ranking (best to worst) descending_rank = ( item_profits.with_columns( - [ - pl.col("avg(ss_net_profit)") - .rank(method="ordinal", descending=True) - .alias("rnk") - ] + pl.col("avg_profit").rank(method="min", descending=True).alias("rnk") ) .filter(pl.col("rnk") < 11) .select(["ss_item_sk", "rnk"]) ) + item_cols = item.select(["i_item_sk", "i_product_name"]) sort_by = {"rnk": False} limit = 100 - # Step 5: Join rankings and get product names return QueryResult( frame=( ascending_rank.join(descending_rank, on="rnk", how="inner", suffix="_desc") - .join(item, left_on="ss_item_sk", right_on="i_item_sk", how="inner") + .join(item_cols, left_on="ss_item_sk", right_on="i_item_sk", how="inner") .join( - item, + item_cols, left_on="ss_item_sk_desc", right_on="i_item_sk", how="inner", diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q45.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q45.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q45.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q45.py index c02a7090db2..b556634be0a 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q45.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q45.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q46.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q46.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q46.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q46.py index aaa5dea0778..94a24031853 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q46.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q46.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q47.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q47.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q47.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q47.py index 4a1278c771a..3fae2b8b9ee 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q47.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q47.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q48.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q48.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q48.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q48.py index 604607e6811..3a600f6584e 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q48.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q48.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q49.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q49.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py index f1d6e78ccea..ee7ddfa3aae 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q49.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q49.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q5.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q5.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q5.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q5.py index 95b5dd71a9c..e8c3216d153 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q5.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q5.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q50.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q50.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q50.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q50.py index 14638535a68..2ba9c7bc51c 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q50.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q50.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q51.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q51.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py index ddcae0f5dd1..35f5814f869 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q51.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q51.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q52.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q52.py similarity index 74% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q52.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q52.py index 2a8f74151b3..2d6843744a8 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q52.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q52.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -69,15 +69,21 @@ def polars_impl(run_config: RunConfig) -> QueryResult: sort_by = {"d_year": False, "ext_price": True, "brand_id": False} limit = 100 + + # Pre-filter both lookup tables before joining against store_sales [87 partitions]. + # date_dim keeps d_year because it appears in the GROUP BY. + filtered_dates = date_dim.filter( + (pl.col("d_moy") == month) & (pl.col("d_year") == year) + ).select(["d_date_sk", "d_year"]) + filtered_item = item.filter(pl.col("i_manager_id") == manager_id).select( + ["i_item_sk", "i_brand", "i_brand_id"] + ) + return QueryResult( frame=( - store_sales.join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") - .join(item, left_on="ss_item_sk", right_on="i_item_sk") - .filter( - (pl.col("i_manager_id") == manager_id) - & (pl.col("d_moy") == month) - & (pl.col("d_year") == year) - ) + store_sales.select(["ss_sold_date_sk", "ss_item_sk", "ss_ext_sales_price"]) + .join(filtered_dates, left_on="ss_sold_date_sk", right_on="d_date_sk") + .join(filtered_item, left_on="ss_item_sk", right_on="i_item_sk") .group_by(["d_year", "i_brand", "i_brand_id"]) .agg(pl.col("ss_ext_sales_price").sum().alias("ext_price")) .select( diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q53.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q53.py similarity index 79% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q53.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q53.py index 88a9e1992ee..4a4f5ca9d85 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q53.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q53.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -101,25 +101,37 @@ def polars_impl(run_config: RunConfig) -> QueryResult: date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix) store = get_data(run_config.dataset_path, "store", run_config.suffix) month_seq_list = list(range(dms, dms + 12)) + + # Pre-filter lookup tables before joining against store_sales [87 partitions]. + # date_dim: keep d_qoy because it appears in the GROUP BY. + filtered_dates = date_dim.filter( + pl.col("d_month_seq").is_in(month_seq_list) + ).select(["d_date_sk", "d_qoy"]) + # item: apply both OR'd rule groups up front; only i_manufact_id needed after. + filtered_item = item.filter( + ( + pl.col("i_category").is_in(categories1) + & pl.col("i_class").is_in(classes1) + & pl.col("i_brand").is_in(brands1) + ) + | ( + pl.col("i_category").is_in(categories2) + & pl.col("i_class").is_in(classes2) + & pl.col("i_brand").is_in(brands2) + ) + ).select(["i_item_sk", "i_manufact_id"]) + grouped_data = ( - store_sales.join(item, left_on="ss_item_sk", right_on="i_item_sk") - .join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") - .join(store, left_on="ss_store_sk", right_on="s_store_sk") - .filter(pl.col("d_month_seq").is_in(month_seq_list)) - .filter( - # First rule group - ( - (pl.col("i_category").is_in(categories1)) - & (pl.col("i_class").is_in(classes1)) - & (pl.col("i_brand").is_in(brands1)) - ) - | - # Second rule group - ( - (pl.col("i_category").is_in(categories2)) - & (pl.col("i_class").is_in(classes2)) - & (pl.col("i_brand").is_in(brands2)) - ) + store_sales.select( + ["ss_sold_date_sk", "ss_item_sk", "ss_store_sk", "ss_sales_price"] + ) + .join(filtered_item, left_on="ss_item_sk", right_on="i_item_sk") + .join(filtered_dates, left_on="ss_sold_date_sk", right_on="d_date_sk") + .join( + store.select("s_store_sk"), + left_on="ss_store_sk", + right_on="s_store_sk", + how="semi", ) .group_by(["i_manufact_id", "d_qoy"]) .agg([pl.col("ss_sales_price").sum().alias("sum_sales_raw")]) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q54.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q54.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q54.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q54.py index afe8e3db09c..59ea271678c 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q54.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q54.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q55.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q55.py similarity index 72% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q55.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q55.py index e6cfbfef9e6..7f48a7ed6d7 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q55.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q55.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -65,15 +65,25 @@ def polars_impl(run_config: RunConfig) -> QueryResult: item = get_data(run_config.dataset_path, "item", run_config.suffix) sort_by = {"ext_price": True, "brand_id": False} limit = 100 + + # d_year not in GROUP BY so date filter can be a semi-join (no date columns needed). + filtered_dates = date_dim.filter( + (pl.col("d_moy") == month) & (pl.col("d_year") == year) + ).select("d_date_sk") + filtered_item = item.filter(pl.col("i_manager_id") == manager_id).select( + ["i_item_sk", "i_brand", "i_brand_id"] + ) + return QueryResult( frame=( - store_sales.join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") - .join(item, left_on="ss_item_sk", right_on="i_item_sk") - .filter( - (pl.col("i_manager_id") == manager_id) - & (pl.col("d_moy") == month) - & (pl.col("d_year") == year) + store_sales.select(["ss_sold_date_sk", "ss_item_sk", "ss_ext_sales_price"]) + .join( + filtered_dates, + left_on="ss_sold_date_sk", + right_on="d_date_sk", + how="semi", ) + .join(filtered_item, left_on="ss_item_sk", right_on="i_item_sk") .group_by(["i_brand", "i_brand_id"]) .agg(pl.col("ss_ext_sales_price").sum().alias("ext_price")) .select( diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q56.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q56.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q56.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q56.py index 3ce11ce2b60..646ad01c8f7 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q56.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q56.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q57.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q57.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q57.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q57.py index 1ee90911f10..6b3ab51830e 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q57.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q57.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q58.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q58.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q58.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q58.py index 5fc54970a1f..22a94063a8f 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q58.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q58.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q59.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q59.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q59.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q59.py index f2648f447e9..7e5d4bd62d1 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q59.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q59.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q6.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q6.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q6.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q6.py index bbd71270c0d..e08c5befc36 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q6.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q6.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q60.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q60.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q60.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q60.py index c295ef731a1..5209bc08c4d 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q60.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q60.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q61.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q61.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q61.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q61.py index f5d7d4644a9..6cdebf5e8ed 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q61.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q61.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q62.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q62.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q62.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q62.py index b08d51b9150..6aa9124b06e 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q62.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q62.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q63.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q63.py similarity index 69% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q63.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q63.py index 3e9b06cb553..faed1d49b28 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q63.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q63.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -89,42 +89,54 @@ def polars_impl(run_config: RunConfig) -> QueryResult: date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix) store = get_data(run_config.dataset_path, "store", run_config.suffix) - inner_query = ( - store_sales.join(item, left_on="ss_item_sk", right_on="i_item_sk") - .join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") - .join(store, left_on="ss_store_sk", right_on="s_store_sk") - .filter( - pl.col("d_month_seq").is_in([dms + i for i in range(12)]) - & ( - ( - pl.col("i_category").is_in(["Books", "Children", "Electronics"]) - & pl.col("i_class").is_in( - ["personal", "portable", "reference", "self-help"] - ) - & pl.col("i_brand").is_in( - [ - "scholaramalgamalg #14", - "scholaramalgamalg #7", - "exportiunivamalg #9", - "scholaramalgamalg #9", - ] - ) - ) - | ( - pl.col("i_category").is_in(["Women", "Music", "Men"]) - & pl.col("i_class").is_in( - ["accessories", "classical", "fragrances", "pants"] - ) - & pl.col("i_brand").is_in( - [ - "amalgimporto #1", - "edu packscholar #1", - "exportiimporto #1", - "importoamalg #1", - ] - ) - ) + # Pre-filter both lookup tables before joining against store_sales [58 partitions]. + # item: apply both OR'd rule groups up front; only i_manager_id needed after. + # date_dim: keep d_moy because it appears in the GROUP BY. + filtered_item = item.filter( + ( + pl.col("i_category").is_in(["Books", "Children", "Electronics"]) + & pl.col("i_class").is_in( + ["personal", "portable", "reference", "self-help"] + ) + & pl.col("i_brand").is_in( + [ + "scholaramalgamalg #14", + "scholaramalgamalg #7", + "exportiunivamalg #9", + "scholaramalgamalg #9", + ] + ) + ) + | ( + pl.col("i_category").is_in(["Women", "Music", "Men"]) + & pl.col("i_class").is_in( + ["accessories", "classical", "fragrances", "pants"] ) + & pl.col("i_brand").is_in( + [ + "amalgimporto #1", + "edu packscholar #1", + "exportiimporto #1", + "importoamalg #1", + ] + ) + ) + ).select(["i_item_sk", "i_manager_id"]) + filtered_dates = date_dim.filter( + pl.col("d_month_seq").is_in([dms + i for i in range(12)]) + ).select(["d_date_sk", "d_moy"]) + + inner_query = ( + store_sales.select( + ["ss_sold_date_sk", "ss_item_sk", "ss_store_sk", "ss_sales_price"] + ) + .join(filtered_item, left_on="ss_item_sk", right_on="i_item_sk") + .join(filtered_dates, left_on="ss_sold_date_sk", right_on="d_date_sk") + .join( + store.select("s_store_sk"), + left_on="ss_store_sk", + right_on="s_store_sk", + how="semi", ) .group_by(["i_manager_id", "d_moy"]) .agg([pl.col("ss_sales_price").sum().alias("sum_sales")]) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q64.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q64.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py index aaa2468498a..b7046d2e65a 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q64.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q64.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q65.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q65.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q65.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q65.py index ad4829c4d8b..39fc0af08bd 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q65.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q65.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q66.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q66.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q66.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q66.py index 158555736b4..70a495daa90 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q66.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q66.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q67.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q67.py similarity index 86% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q67.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q67.py index 0c6859ca55a..8f0c18a6e59 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q67.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q67.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -214,11 +214,35 @@ def polars_impl(run_config: RunConfig) -> QueryResult: store = get_data(run_config.dataset_path, "store", run_config.suffix) item = get_data(run_config.dataset_path, "item", run_config.suffix) + # Pre-filter date_dim to the 12-month window before joining against store_sales [58]. + # d_month_seq not needed after filter; keep group-by output columns. + filtered_dates = date_dim.filter( + pl.col("d_month_seq").is_between(dms, dms + 11) + ).select(["d_date_sk", "d_year", "d_qoy", "d_moy"]) + base_data = ( - store_sales.join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") - .join(store, left_on="ss_store_sk", right_on="s_store_sk") - .join(item, left_on="ss_item_sk", right_on="i_item_sk") - .filter(pl.col("d_month_seq").is_between(dms, dms + 11)) + store_sales.select( + [ + "ss_sold_date_sk", + "ss_item_sk", + "ss_store_sk", + "ss_sales_price", + "ss_quantity", + ] + ) + .join(filtered_dates, left_on="ss_sold_date_sk", right_on="d_date_sk") + .join( + store.select(["s_store_sk", "s_store_id"]), + left_on="ss_store_sk", + right_on="s_store_sk", + ) + .join( + item.select( + ["i_item_sk", "i_category", "i_class", "i_brand", "i_product_name"] + ), + left_on="ss_item_sk", + right_on="i_item_sk", + ) .with_columns( (pl.col("ss_sales_price") * pl.col("ss_quantity")) .fill_null(0) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q68.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q68.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q68.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q68.py index 258872a7495..b85a3cd5d9b 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q68.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q68.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q69.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q69.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q69.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q69.py index 844eb5d9b84..5a542da51b1 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q69.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q69.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q7.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q7.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q7.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q7.py index f762a64becf..023bc473a4a 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q7.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q7.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q70.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q70.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q70.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q70.py index 1491321757f..4999ae78720 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q70.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q70.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q71.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q71.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q71.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q71.py index 4ae7df5121c..6241c242ac1 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q71.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q71.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q72.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q72.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q72.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q72.py index d7de7ddaf98..f264053438c 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q72.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q72.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import COUNT_DTYPE, QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import COUNT_DTYPE, QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q73.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q73.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q73.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q73.py index 88dff01105b..0228f1dd8cc 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q73.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q73.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q74.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q74.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q74.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q74.py index 8c42981b83c..2a9aad88ab1 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q74.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q74.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q75.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q75.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q75.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q75.py index fdd5dd0a598..36ec9fd2e0b 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q75.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q75.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q76.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q76.py similarity index 82% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q76.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q76.py index e05d9b28f65..21198788592 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q76.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q76.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -106,10 +106,16 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ) item = get_data(run_config.dataset_path, "item", run_config.suffix) date_dim = get_data(run_config.dataset_path, "date_dim", run_config.suffix) + + # Project lookup tables to only the columns needed in each component. + date_cols = date_dim.select(["d_date_sk", "d_year", "d_qoy"]) + item_cols = item.select(["i_item_sk", "i_category"]) + store_component = ( store_sales.filter(pl.col(nullcol_ss).is_null()) - .join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") - .join(item, left_on="ss_item_sk", right_on="i_item_sk") + .select(["ss_sold_date_sk", "ss_item_sk", "ss_ext_sales_price"]) + .join(date_cols, left_on="ss_sold_date_sk", right_on="d_date_sk") + .join(item_cols, left_on="ss_item_sk", right_on="i_item_sk") .select( [ pl.lit("store").alias("channel"), @@ -123,8 +129,9 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ) web_component = ( web_sales.filter(pl.col(nullcol_ws).is_null()) - .join(date_dim, left_on="ws_sold_date_sk", right_on="d_date_sk") - .join(item, left_on="ws_item_sk", right_on="i_item_sk") + .select(["ws_sold_date_sk", "ws_item_sk", "ws_ext_sales_price"]) + .join(date_cols, left_on="ws_sold_date_sk", right_on="d_date_sk") + .join(item_cols, left_on="ws_item_sk", right_on="i_item_sk") .select( [ pl.lit("web").alias("channel"), @@ -138,8 +145,9 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ) catalog_component = ( catalog_sales.filter(pl.col(nullcol_cs).is_null()) - .join(date_dim, left_on="cs_sold_date_sk", right_on="d_date_sk") - .join(item, left_on="cs_item_sk", right_on="i_item_sk") + .select(["cs_sold_date_sk", "cs_item_sk", "cs_ext_sales_price"]) + .join(date_cols, left_on="cs_sold_date_sk", right_on="d_date_sk") + .join(item_cols, left_on="cs_item_sk", right_on="i_item_sk") .select( [ pl.lit("catalog").alias("channel"), @@ -166,10 +174,7 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .agg( [ pl.len().cast(pl.Int64).alias("sales_cnt"), - pl.when(pl.col("ext_sales_price").count() > 0) - .then(pl.col("ext_sales_price").sum()) - .otherwise(None) - .alias("sales_amt"), + pl.col("ext_sales_price").sum().alias("sales_amt"), ] ) .select( diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q77.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q77.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q77.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q77.py index 3404e6fd7d9..eeacd977399 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q77.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q77.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q78.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q78.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q78.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q78.py index 3af998b89fa..7e2a9019779 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q78.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q78.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q79.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q79.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q79.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q79.py index 2e44cc32dd1..741aa0acd14 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q79.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q79.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q8.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q8.py similarity index 80% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q8.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q8.py index 05c47660ea6..0405f3ed176 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q8.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q8.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -105,19 +105,36 @@ def polars_impl(run_config: RunConfig) -> QueryResult: .select(pl.col("ca_zip").str.slice(0, 2).alias("ca_zip_prefix")) ) - # Main query: join store_sales with date_dim, store, and filter by zip codes + # Pre-filter date_dim; d_year/d_qoy not needed after filter — semi-join. + filtered_dates = date_dim.filter( + (pl.col("d_year") == year) & (pl.col("d_qoy") == qoy) + ).select("d_date_sk") + return QueryResult( frame=( - store_sales.join(date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk") - .join(store, left_on="ss_store_sk", right_on="s_store_sk") - .with_columns(pl.col("s_zip").str.slice(0, 2).alias("s_zip_prefix")) + store_sales.select(["ss_sold_date_sk", "ss_store_sk", "ss_net_profit"]) + .join( + filtered_dates, + left_on="ss_sold_date_sk", + right_on="d_date_sk", + how="semi", + ) + .join( + store.select( + [ + "s_store_sk", + "s_store_name", + pl.col("s_zip").str.slice(0, 2).alias("s_zip_prefix"), + ] + ), + left_on="ss_store_sk", + right_on="s_store_sk", + ) .join( intersect_zips, left_on="s_zip_prefix", right_on="ca_zip_prefix", ) - .filter(pl.col("d_qoy") == qoy) - .filter(pl.col("d_year") == year) .group_by("s_store_name") .agg(pl.col("ss_net_profit").sum().alias("sum")) .sort("s_store_name", nulls_last=True) diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q80.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q80.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q80.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q80.py index 68157f05a82..382c381a317 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q80.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q80.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q81.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q81.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q81.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q81.py index 704035793bb..e018c59b2d7 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q81.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q81.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q82.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q82.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q82.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q82.py index 2c0112f0837..103ca91e9b4 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q82.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q82.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q83.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q83.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q83.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q83.py index 05858eef99c..7ee2a30961c 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q83.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q83.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q84.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q84.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q84.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q84.py index e8fdaba7842..62540107a07 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q84.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q84.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q85.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q85.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q85.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q85.py index 9fb980d15b1..6964fe31edf 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q85.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q85.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q86.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q86.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q86.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q86.py index 2c815317e88..cbddf727473 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q86.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q86.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q87.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q87.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py index 5dcf0054b09..1de0babdc64 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q87.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q87.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q88.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q88.py similarity index 71% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q88.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q88.py index a086bb609cf..2216d24c55e 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q88.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q88.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -142,81 +142,76 @@ def polars_impl(run_config: RunConfig) -> QueryResult: ) time_dim = get_data(run_config.dataset_path, "time_dim", run_config.suffix) store = get_data(run_config.dataset_path, "store", run_config.suffix) - hd_filter = ( + + # Pre-filter each small table before joining against store_sales [58 partitions]. + filtered_hdemo = household_demographics.filter( ((pl.col("hd_dep_count") == hd1) & (pl.col("hd_vehicle_count") <= hd1 + 2)) | ((pl.col("hd_dep_count") == hd2) & (pl.col("hd_vehicle_count") <= hd2 + 2)) | ((pl.col("hd_dep_count") == hd3) & (pl.col("hd_vehicle_count") <= hd3 + 2)) + ).select("hd_demo_sk") + filtered_store = store.filter(pl.col("s_store_name") == s_store_name).select( + "s_store_sk" ) - base_query = ( - store_sales.join( - time_dim, left_on="ss_sold_time_sk", right_on="t_time_sk", how="inner" - ) - .join( - household_demographics, - left_on="ss_hdemo_sk", - right_on="hd_demo_sk", - how="inner", - ) - .join(store, left_on="ss_store_sk", right_on="s_store_sk", how="inner") - .filter( - hd_filter - & ( - pl.col("s_store_name").is_not_null() - & (pl.col("s_store_name") == s_store_name) - ) + # Restrict time_dim to the union of all 8 slot conditions; every surviving row maps + # to exactly one bucket, so the downstream pl.when chain is exhaustive. + filtered_time = time_dim.filter( + ((pl.col("t_hour") == 8) & (pl.col("t_minute") >= 30)) + | pl.col("t_hour").is_in([9, 10, 11]) + | ((pl.col("t_hour") == 12) & (pl.col("t_minute") < 30)) + ).select(["t_time_sk", "t_hour", "t_minute"]) + + bucket_names = [ + "h8_30_to_9", + "h9_to_9_30", + "h9_30_to_10", + "h10_to_10_30", + "h10_30_to_11", + "h11_to_11_30", + "h11_30_to_12", + "h12_to_12_30", + ] + + # Collapse the 58-partition store_sales pipeline to an 8-row bucket-count table first. + # The 8 conditional sums in the final select then operate on [1] partition, so even if + # the streaming executor creates separate sub-plans for each sum, each reads only the + # tiny CACHE'd group_by output rather than re-scanning store_sales. + counts_lf = ( + store_sales.select(["ss_sold_time_sk", "ss_hdemo_sk", "ss_store_sk"]) + .join(filtered_time, left_on="ss_sold_time_sk", right_on="t_time_sk") + .join(filtered_hdemo, left_on="ss_hdemo_sk", right_on="hd_demo_sk", how="semi") + .join(filtered_store, left_on="ss_store_sk", right_on="s_store_sk", how="semi") + .select( + pl.when((pl.col("t_hour") == 8) & (pl.col("t_minute") >= 30)) + .then(pl.lit(0)) + .when((pl.col("t_hour") == 9) & (pl.col("t_minute") < 30)) + .then(pl.lit(1)) + .when((pl.col("t_hour") == 9) & (pl.col("t_minute") >= 30)) + .then(pl.lit(2)) + .when((pl.col("t_hour") == 10) & (pl.col("t_minute") < 30)) + .then(pl.lit(3)) + .when((pl.col("t_hour") == 10) & (pl.col("t_minute") >= 30)) + .then(pl.lit(4)) + .when((pl.col("t_hour") == 11) & (pl.col("t_minute") < 30)) + .then(pl.lit(5)) + .when((pl.col("t_hour") == 11) & (pl.col("t_minute") >= 30)) + .then(pl.lit(6)) + .when((pl.col("t_hour") == 12) & (pl.col("t_minute") < 30)) + .then(pl.lit(7)) + .alias("bucket") ) + .group_by("bucket") + .agg(pl.len().cast(pl.Int64).alias("cnt")) ) + return QueryResult( - frame=base_query.select( + frame=counts_lf.select( [ - pl.when((pl.col("t_hour") == 8) & (pl.col("t_minute") >= 30)) - .then(1) - .otherwise(0) - .sum() - .cast(pl.Int64) - .alias("h8_30_to_9"), - pl.when((pl.col("t_hour") == 9) & (pl.col("t_minute") < 30)) - .then(1) - .otherwise(0) - .sum() - .cast(pl.Int64) - .alias("h9_to_9_30"), - pl.when((pl.col("t_hour") == 9) & (pl.col("t_minute") >= 30)) - .then(1) - .otherwise(0) - .sum() - .cast(pl.Int64) - .alias("h9_30_to_10"), - pl.when((pl.col("t_hour") == 10) & (pl.col("t_minute") < 30)) - .then(1) - .otherwise(0) - .sum() - .cast(pl.Int64) - .alias("h10_to_10_30"), - pl.when((pl.col("t_hour") == 10) & (pl.col("t_minute") >= 30)) - .then(1) - .otherwise(0) - .sum() - .cast(pl.Int64) - .alias("h10_30_to_11"), - pl.when((pl.col("t_hour") == 11) & (pl.col("t_minute") < 30)) - .then(1) - .otherwise(0) - .sum() - .cast(pl.Int64) - .alias("h11_to_11_30"), - pl.when((pl.col("t_hour") == 11) & (pl.col("t_minute") >= 30)) - .then(1) - .otherwise(0) - .sum() - .cast(pl.Int64) - .alias("h11_30_to_12"), - pl.when((pl.col("t_hour") == 12) & (pl.col("t_minute") < 30)) - .then(1) - .otherwise(0) + pl.when(pl.col("bucket") == i) + .then(pl.col("cnt")) + .otherwise(pl.lit(0).cast(pl.Int64)) .sum() - .cast(pl.Int64) - .alias("h12_to_12_30"), + .alias(name) + for i, name in enumerate(bucket_names) ] ), sort_by=[], diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q89.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q89.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q89.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q89.py index 5f611668009..b43cd136fb6 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q89.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q89.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q9.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q9.py similarity index 65% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q9.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q9.py index d42218179ea..f845d4833c2 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q9.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q9.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -102,57 +102,56 @@ def polars_impl(run_config: RunConfig) -> QueryResult: aggcelse = params["aggcelse"] rc = params["rc"] - # Load required tables store_sales = get_data(run_config.dataset_path, "store_sales", run_config.suffix) reason = get_data(run_config.dataset_path, "reason", run_config.suffix) - # Define bucket configurations: (min_qty, max_qty, count_threshold) - buckets = [ - (1, 20, rc[0]), - (21, 40, rc[1]), - (41, 60, rc[2]), - (61, 80, rc[3]), - (81, 100, rc[4]), - ] - - bucket_expressions = [] - for i, (min_qty, max_qty, _) in enumerate(buckets, 1): - condition = pl.col("ss_quantity").is_between(min_qty, max_qty, closed="both") - bucket_expressions.extend( - [ - condition.sum().alias(f"count_{i}"), - pl.when(condition) - .then(pl.col(aggcthen)) - .otherwise(None) - .mean() - .alias(f"avg_then_{i}"), - pl.when(condition) - .then(pl.col(aggcelse)) - .otherwise(None) - .mean() - .alias(f"avg_else_{i}"), - ] + thresholds = pl.LazyFrame({"bucket": [1, 2, 3, 4, 5], "threshold": list(rc)}) + + # Single scan: the 5 ss_quantity ranges are non-overlapping, so a group_by + # computes all counts and averages in one pass over store_sales. + stats = ( + store_sales.with_columns( + pl.when(pl.col("ss_quantity").is_between(1, 20)) + .then(pl.lit(1)) + .when(pl.col("ss_quantity").is_between(21, 40)) + .then(pl.lit(2)) + .when(pl.col("ss_quantity").is_between(41, 60)) + .then(pl.lit(3)) + .when(pl.col("ss_quantity").is_between(61, 80)) + .then(pl.lit(4)) + .when(pl.col("ss_quantity").is_between(81, 100)) + .then(pl.lit(5)) + .alias("bucket") ) - - combined_stats = store_sales.select(bucket_expressions) - - # Select appropriate value per bucket based on count threshold - bucket_values = [] - for i, (_min_qty, _max_qty, threshold) in enumerate(buckets, 1): - bucket = ( - pl.when(pl.col(f"count_{i}") > threshold) - .then(pl.col(f"avg_then_{i}")) - .otherwise(pl.col(f"avg_else_{i}")) - .alias(f"bucket{i}") + .filter(pl.col("bucket").is_not_null()) + .group_by("bucket") + .agg( + pl.len().alias("count"), + pl.col(aggcthen).mean().alias("avg_then"), + pl.col(aggcelse).mean().alias("avg_else"), + ) + .join(thresholds, on="bucket") + .select( + pl.col("bucket"), + pl.when(pl.col("count") > pl.col("threshold")) + .then(pl.col("avg_then")) + .otherwise(pl.col("avg_else")) + .alias("value"), ) - bucket_values.append(bucket) + .sort("bucket") + ) + + # Pivot 5 rows → 1 row with 5 named columns (operates on 5 rows, trivially fast) + wide = stats.select( + pl.col("value").filter(pl.col("bucket") == i).first().alias(f"bucket{i}") + for i in range(1, 6) + ) - # Create result DataFrame with one row (using reason table as in SQL) return QueryResult( frame=( reason.filter(pl.col("r_reason_sk") == 1) - .join(combined_stats, how="cross") - .select(bucket_values) + .join(wide, how="cross") + .select([f"bucket{i}" for i in range(1, 6)]) .limit(1) ), sort_by=[], diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q90.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q90.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q90.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q90.py index 3b1e9f90ef3..e995b4c22da 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q90.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q90.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q91.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q91.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q91.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q91.py index 829c172cb0c..cfac30225b2 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q91.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q91.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q92.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q92.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q92.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q92.py index 934ed050ef5..d187daeb4b8 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q92.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q92.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q93.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q93.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q93.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q93.py index fc7a640566c..bcb1f486184 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q93.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q93.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q94.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q94.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q94.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q94.py index feefa3108cf..d98b113957a 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q94.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q94.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q95.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q95.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q95.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q95.py index 0b752fcc8ca..1010dc5bc08 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q95.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q95.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q96.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q96.py similarity index 92% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q96.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q96.py index a2d9a78f411..4ddb2673c88 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q96.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q96.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q97.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q97.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q97.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q97.py index b2af2b88817..4ee2ac1b90c 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q97.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q97.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q98.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q98.py similarity index 78% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q98.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q98.py index a70bf71269b..00ad3d62a71 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q98.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q98.py @@ -10,11 +10,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: @@ -76,17 +76,34 @@ def polars_impl(run_config: RunConfig) -> QueryResult: end_date_py = start_date_py + timedelta(days=30) start_date = pl.date(start_date_py.year, start_date_py.month, start_date_py.day) end_date = pl.date(end_date_py.year, end_date_py.month, end_date_py.day) + + # Pre-filter item to matching categories before joining against store_sales [58 partitions]. + filtered_item = item.filter( + pl.col("i_category").is_in(params["categories"]) + ).select( + [ + "i_item_sk", + "i_item_id", + "i_item_desc", + "i_current_price", + "i_class", + "i_category", + ] + ) + # Pre-filter date_dim to the 30-day window; d_date not needed after filter — semi-join. + filtered_dates = date_dim.filter( + pl.col("d_date").is_between(start_date, end_date, closed="both") + ).select("d_date_sk") + return QueryResult( frame=( - store_sales.join( - item, left_on="ss_item_sk", right_on="i_item_sk", how="inner" - ) + store_sales.select(["ss_sold_date_sk", "ss_item_sk", "ss_ext_sales_price"]) + .join(filtered_item, left_on="ss_item_sk", right_on="i_item_sk") .join( - date_dim, left_on="ss_sold_date_sk", right_on="d_date_sk", how="inner" - ) - .filter( - pl.col("i_category").is_in(params["categories"]) - & pl.col("d_date").is_between(start_date, end_date, closed="both") + filtered_dates, + left_on="ss_sold_date_sk", + right_on="d_date_sk", + how="semi", ) .group_by( [ diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q99.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q99.py similarity index 96% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q99.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q99.py index a7e3da8578f..5aaaa13dbc3 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsds_queries/q99.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsds_queries/q99.py @@ -9,11 +9,11 @@ import polars as pl -from cudf_polars.experimental.benchmarks.pdsds_parameters import load_parameters -from cudf_polars.experimental.benchmarks.utils import QueryResult, get_data +from cudf_polars.streaming.benchmarks.pdsds_parameters import load_parameters +from cudf_polars.streaming.benchmarks.utils import QueryResult, get_data if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig def duckdb_impl(run_config: RunConfig) -> str: diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py similarity index 99% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py index 59a71fe8d13..79260dc096a 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/pdsh.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/pdsh.py @@ -20,7 +20,7 @@ import polars as pl try: - from cudf_polars.experimental.benchmarks.utils import ( + from cudf_polars.streaming.benchmarks.utils import ( COUNT_DTYPE, QueryResult, RunConfig, @@ -36,7 +36,7 @@ COUNT_DTYPE = None # type: ignore[assignment] if TYPE_CHECKING: - from cudf_polars.experimental.benchmarks.utils import RunConfig + from cudf_polars.streaming.benchmarks.utils import RunConfig # Without this setting, the first IO task to run # on each worker takes ~15 sec extra diff --git a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py rename to python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py index 3c85971f091..fc2159d9b4f 100644 --- a/python/cudf_polars/cudf_polars/experimental/benchmarks/utils_new_frontends.py +++ b/python/cudf_polars/cudf_polars/streaming/benchmarks/utils.py @@ -63,13 +63,13 @@ from cudf_polars.dsl.ir import IRExecutionContext from cudf_polars.dsl.tracing import Scope from cudf_polars.dsl.translate import Translator - from cudf_polars.experimental.benchmarks.asserts import ( + from cudf_polars.engine.core import StreamingEngine + from cudf_polars.streaming.benchmarks.asserts import ( ValidationError, assert_tpch_result_equal, ) - from cudf_polars.experimental.explain import explain_query - from cudf_polars.experimental.parallel import evaluate_streaming - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine + from cudf_polars.streaming.explain import explain_query + from cudf_polars.streaming.parallel import evaluate_streaming from cudf_polars.utils.config import ConfigOptions CUDF_POLARS_AVAILABLE = True @@ -79,8 +79,8 @@ if TYPE_CHECKING: from collections.abc import Callable - from cudf_polars.experimental.explain import SerializablePlan - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.engine.options import StreamingOptions + from cudf_polars.streaming.explain import SerializablePlan POLARS_VALIDATION_OPTIONS = { "check_row_order": True, @@ -400,7 +400,7 @@ class RunConfig: # All streaming/rapidsmpf/engine knobs streaming_options: StreamingOptions = dataclasses.field( default_factory=lambda: __import__( - "cudf_polars.experimental.rapidsmpf.frontend.options", + "cudf_polars.engine.options", fromlist=["StreamingOptions"], ).StreamingOptions() ) @@ -447,7 +447,7 @@ def __post_init__(self) -> None: # noqa: D105 @classmethod def from_args(cls, args: argparse.Namespace) -> RunConfig: """Create a RunConfig from command line arguments.""" - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.engine.options import StreamingOptions streaming_options = StreamingOptions._from_argparse(args) @@ -921,7 +921,7 @@ def run_polars_query( print_query_plan(q_id, q, args, run_config, engine, print_plans=args.print_plans) plan = None if (args.explain or args.explain_logical) and engine is not None: - from cudf_polars.experimental.explain import serialize_query + from cudf_polars.streaming.explain import serialize_query plan = serialize_query(q, engine) @@ -1179,7 +1179,7 @@ def run_polars_spmd( validation_files: dict[int, Path] | None, ) -> None: """Run benchmark queries using SPMD execution via the ``rrun`` launcher.""" - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine executor_options = get_executor_options(run_config, benchmark=benchmark) # "cluster" is reserved — SPMDEngine sets it @@ -1193,10 +1193,10 @@ def run_polars_spmd( executor_options=executor_options, engine_options=engine_options, ) as engine: - from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id - from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( + from cudf_polars.engine.spmd import ( allgather_polars_dataframe, ) + from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id def _allgather_result(df: pl.DataFrame) -> pl.DataFrame: with reserve_op_id() as op_id: @@ -1238,7 +1238,7 @@ def run_polars_ray( validation_files: dict[int, Path] | None, ) -> None: """Run benchmark queries using Ray actor-based distributed execution.""" - from cudf_polars.experimental.rapidsmpf.frontend.ray import RayEngine + from cudf_polars.engine.ray import RayEngine executor_options = get_executor_options(run_config, benchmark=benchmark) # "cluster" is reserved — RayEngine sets it @@ -1289,7 +1289,7 @@ def run_polars_dask( """Run benchmark queries using Dask distributed execution.""" import distributed - from cudf_polars.experimental.rapidsmpf.frontend.dask import DaskEngine + from cudf_polars.engine.dask import DaskEngine executor_options = get_executor_options(run_config, benchmark=benchmark) # "cluster" is reserved — DaskEngine sets it @@ -1749,7 +1749,7 @@ def parse(query: str | int) -> list[int]: def build_parser(num_queries: int = 22) -> argparse.ArgumentParser: """Build the argument parser for PDS-H/PDS-DS benchmarks.""" - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.engine.options import StreamingOptions parser = argparse.ArgumentParser( prog="Cudf-Polars PDS-H/PDS-DS Benchmarks", diff --git a/python/cudf_polars/cudf_polars/experimental/dispatch.py b/python/cudf_polars/cudf_polars/streaming/dispatch.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/dispatch.py rename to python/cudf_polars/cudf_polars/streaming/dispatch.py index 9ff0cc3156b..81937e6ce6f 100644 --- a/python/cudf_polars/cudf_polars/experimental/dispatch.py +++ b/python/cudf_polars/cudf_polars/streaming/dispatch.py @@ -14,7 +14,7 @@ from cudf_polars.dsl import ir from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.base import ( + from cudf_polars.streaming.base import ( PartitionInfo, StatsCollector, ) diff --git a/python/cudf_polars/cudf_polars/experimental/distinct.py b/python/cudf_polars/cudf_polars/streaming/distinct.py similarity index 93% rename from python/cudf_polars/cudf_polars/experimental/distinct.py rename to python/cudf_polars/cudf_polars/streaming/distinct.py index 564fe570919..07904fcf357 100644 --- a/python/cudf_polars/cudf_polars/experimental/distinct.py +++ b/python/cudf_polars/cudf_polars/streaming/distinct.py @@ -11,11 +11,11 @@ from cudf_polars.dsl.expressions.base import Col, NamedExpr from cudf_polars.dsl.ir import Distinct -from cudf_polars.experimental.base import PartitionInfo -from cudf_polars.experimental.dispatch import lower_ir_node -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.shuffle import Shuffle -from cudf_polars.experimental.utils import ( +from cudf_polars.streaming.base import PartitionInfo +from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.shuffle import Shuffle +from cudf_polars.streaming.utils import ( _dynamic_planning_on, _lower_ir_fallback, ) @@ -24,7 +24,7 @@ from collections.abc import MutableMapping from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.dispatch import LowerIRTransformer + from cudf_polars.streaming.dispatch import LowerIRTransformer from cudf_polars.utils.config import ConfigOptions, StreamingExecutor diff --git a/python/cudf_polars/cudf_polars/experimental/explain.py b/python/cudf_polars/cudf_polars/streaming/explain.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/explain.py rename to python/cudf_polars/cudf_polars/streaming/explain.py index d50d9fae0ae..3e876f4f9e0 100644 --- a/python/cudf_polars/cudf_polars/experimental/explain.py +++ b/python/cudf_polars/cudf_polars/streaming/explain.py @@ -26,9 +26,9 @@ ) from cudf_polars.dsl.translate import Translator from cudf_polars.dsl.traversal import traversal -from cudf_polars.experimental.parallel import lower_ir_graph -from cudf_polars.experimental.shuffle import Shuffle -from cudf_polars.experimental.statistics import ( +from cudf_polars.streaming.parallel import lower_ir_graph +from cudf_polars.streaming.shuffle import Shuffle +from cudf_polars.streaming.statistics import ( collect_statistics, ) from cudf_polars.utils.config import ConfigOptions @@ -40,7 +40,7 @@ from cudf_polars.dsl.expressions.base import Expr from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.base import PartitionInfo, StatsCollector + from cudf_polars.streaming.base import PartitionInfo, StatsCollector Serializable: TypeAlias = ( diff --git a/python/cudf_polars/cudf_polars/experimental/expressions.py b/python/cudf_polars/cudf_polars/streaming/expressions.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/expressions.py rename to python/cudf_polars/cudf_polars/streaming/expressions.py index ef94da02dea..ef21cb33d17 100644 --- a/python/cudf_polars/cudf_polars/experimental/expressions.py +++ b/python/cudf_polars/cudf_polars/streaming/expressions.py @@ -51,10 +51,12 @@ from cudf_polars.dsl.traversal import ( CachingVisitor, ) -from cudf_polars.experimental.base import PartitionInfo -from cudf_polars.experimental.over import _decompose_grouped_window_node -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.utils import _dynamic_planning_on +from cudf_polars.streaming.base import PartitionInfo +from cudf_polars.streaming.distinct import lower_distinct +from cudf_polars.streaming.over import _decompose_grouped_window_node +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.shuffle import Shuffle +from cudf_polars.streaming.utils import _dynamic_planning_on if TYPE_CHECKING: from collections.abc import Generator, MutableMapping, Sequence @@ -180,8 +182,6 @@ def _decompose_unique( A mapping from unique nodes in the new graph to associated partitioning information. """ - from cudf_polars.experimental.distinct import lower_distinct - (child,) = unique.children (maintain_order,) = unique.options columns, input_ir, partition_info = select( @@ -314,8 +314,6 @@ def _decompose_agg_node( (child,) = agg.children pi = partition_info[input_ir] if pi.count > 1 and [ne.value for ne in pi.partitioned_on] != [input_ir]: - from cudf_polars.experimental.shuffle import Shuffle - children, input_ir, partition_info = select( [UnaryFunction(agg.dtype, "unique", (False,), child)], input_ir, @@ -512,6 +510,7 @@ def _decompose( input_ir = HConcat( schema, True, # noqa: FBT003 + False, # noqa: FBT003 *unique_input_irs, ) partition_info[input_ir] = PartitionInfo(count=partition_count) diff --git a/python/cudf_polars/cudf_polars/experimental/groupby.py b/python/cudf_polars/cudf_polars/streaming/groupby.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/groupby.py rename to python/cudf_polars/cudf_polars/streaming/groupby.py index 6a17b56bfc5..309b41e772b 100644 --- a/python/cudf_polars/cudf_polars/experimental/groupby.py +++ b/python/cudf_polars/cudf_polars/streaming/groupby.py @@ -30,11 +30,11 @@ from cudf_polars.dsl.ir import GroupBy, Select, Slice from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.naming import unique_names -from cudf_polars.experimental.base import PartitionInfo -from cudf_polars.experimental.dispatch import lower_ir_node -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.shuffle import Shuffle -from cudf_polars.experimental.utils import ( +from cudf_polars.streaming.base import PartitionInfo +from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.shuffle import Shuffle +from cudf_polars.streaming.utils import ( _dynamic_planning_on, _lower_ir_fallback, ) @@ -44,7 +44,7 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.parallel import LowerIRTransformer + from cudf_polars.streaming.parallel import LowerIRTransformer # Supported multi-partition aggregations diff --git a/python/cudf_polars/cudf_polars/experimental/io.py b/python/cudf_polars/cudf_polars/streaming/io.py similarity index 92% rename from python/cudf_polars/cudf_polars/experimental/io.py rename to python/cudf_polars/cudf_polars/streaming/io.py index 0c0d08d3210..4472735ac35 100644 --- a/python/cudf_polars/cudf_polars/experimental/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -24,13 +24,13 @@ Sink, ) from cudf_polars.dsl.tracing import nvtx_annotate_cudf_polars -from cudf_polars.experimental.base import ( +from cudf_polars.streaming.base import ( IOPartitionFlavor, IOPartitionPlan, PartitionInfo, SerializedDataSourceInfo, ) -from cudf_polars.experimental.dispatch import lower_ir_node +from cudf_polars.streaming.dispatch import lower_ir_node from cudf_polars.utils.config import Cluster from cudf_polars.utils.cuda_stream import get_cuda_stream from cudf_polars.utils.versions import POLARS_VERSION_LT_137 @@ -41,12 +41,12 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.experimental.base import ( + from cudf_polars.streaming.base import ( DataSourceInfo, SerializedDataSourceInfo, StatsCollector, ) - from cudf_polars.experimental.dispatch import LowerIRTransformer + from cudf_polars.streaming.dispatch import LowerIRTransformer from cudf_polars.typing import Schema from cudf_polars.utils.config import ( ConfigOptions, @@ -59,9 +59,18 @@ def _( ir: DataFrameScan, rec: LowerIRTransformer ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: - from cudf_polars.experimental.rapidsmpf.io import lower_dataframescan_rapidsmpf + config_options = rec.state["config_options"] - return lower_dataframescan_rapidsmpf(ir, rec) + # NOTE: We calculate the expected partition count + # to help trigger fallback warnings in lower_ir_graph. + # The generate_ir_sub_network logic is NOT required + # to obey this partition count. However, the count + # WILL match after an IO operation (for now). + rows_per_partition = config_options.executor.max_rows_per_partition + nrows = max(ir.df.shape()[0], 1) + count = math.ceil(nrows / rows_per_partition) + + return ir, {ir: PartitionInfo(count=count)} def scan_partition_plan( @@ -255,9 +264,31 @@ def _( def _( ir: Scan, rec: LowerIRTransformer ) -> tuple[IR, MutableMapping[IR, PartitionInfo]]: - from cudf_polars.experimental.rapidsmpf.io import lower_scan_rapidsmpf + config_options = rec.state["config_options"] + if ( + ir.typ in ("csv", "parquet", "ndjson") + and ir.n_rows == -1 + and ir.skip_rows == 0 + and ir.row_index is None + ): + # NOTE: We calculate the expected partition count + # to help trigger fallback warnings in lower_ir_graph. + # The generate_ir_sub_network logic is NOT required + # to obey this partition count. However, the count + # WILL match after an IO operation (for now). + plan = scan_partition_plan(ir, rec.state["stats"], config_options) + paths = list(ir.paths) + if plan.flavor == IOPartitionFlavor.SPLIT_FILES: + count = plan.factor * len(paths) + else: + count = math.ceil(len(paths) / plan.factor) - return lower_scan_rapidsmpf(ir, rec) + return ir, {ir: PartitionInfo(count=count, io_plan=plan)} + else: + plan = IOPartitionPlan( + flavor=IOPartitionFlavor.SINGLE_READ, factor=len(ir.paths) + ) + return ir, {ir: PartitionInfo(count=1, io_plan=plan)} class StreamingSink(IR): @@ -412,7 +443,7 @@ def _sink_to_file( # Path.open returns IO[Any] but SinkInfo needs more specific IO types sink = plc.io.types.SinkInfo([f]) # type: ignore[arg-type] Sink._write_csv(sink, use_options, df) - elif kind == "Json" if POLARS_VERSION_LT_137 else "NDJson": + elif kind == ("Json" if POLARS_VERSION_LT_137 else "NDJson"): mode = "wb" if writer_state is None else "ab" with Path.open(Path(path), mode) as f: # Path.open returns IO[Any] but SinkInfo needs more specific IO types diff --git a/python/cudf_polars/cudf_polars/experimental/join.py b/python/cudf_polars/cudf_polars/streaming/join.py similarity index 95% rename from python/cudf_polars/cudf_polars/experimental/join.py rename to python/cudf_polars/cudf_polars/streaming/join.py index 69abff7fd5b..ec31d27e79b 100644 --- a/python/cudf_polars/cudf_polars/experimental/join.py +++ b/python/cudf_polars/cudf_polars/streaming/join.py @@ -9,11 +9,11 @@ from typing import TYPE_CHECKING from cudf_polars.dsl.ir import ConditionalJoin, Join, Slice -from cudf_polars.experimental.base import PartitionInfo -from cudf_polars.experimental.dispatch import lower_ir_node -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.shuffle import Shuffle -from cudf_polars.experimental.utils import ( +from cudf_polars.streaming.base import PartitionInfo +from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.shuffle import Shuffle +from cudf_polars.streaming.utils import ( _dynamic_planning_on, _fallback_inform, _lower_ir_fallback, @@ -24,7 +24,7 @@ from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.parallel import LowerIRTransformer + from cudf_polars.streaming.parallel import LowerIRTransformer def _maybe_shuffle_frame( diff --git a/python/cudf_polars/cudf_polars/experimental/over.py b/python/cudf_polars/cudf_polars/streaming/over.py similarity index 98% rename from python/cudf_polars/cudf_polars/experimental/over.py rename to python/cudf_polars/cudf_polars/streaming/over.py index abc90d35810..f6222acd39d 100644 --- a/python/cudf_polars/cudf_polars/experimental/over.py +++ b/python/cudf_polars/cudf_polars/streaming/over.py @@ -11,7 +11,7 @@ from cudf_polars.dsl.expr import Agg, Col, Len, NamedExpr from cudf_polars.dsl.ir import IR, GroupBy, Select from cudf_polars.dsl.utils.naming import names_to_indices, unique_names -from cudf_polars.experimental.groupby import combine, decompose +from cudf_polars.streaming.groupby import combine, decompose if TYPE_CHECKING: from collections.abc import Generator, MutableMapping @@ -20,7 +20,7 @@ from cudf_polars.dsl.expr import GroupedWindow from cudf_polars.dsl.expressions.base import Expr from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.experimental.base import PartitionInfo + from cudf_polars.streaming.base import PartitionInfo from cudf_polars.typing import Schema from cudf_polars.utils.config import ConfigOptions diff --git a/python/cudf_polars/cudf_polars/experimental/parallel.py b/python/cudf_polars/cudf_polars/streaming/parallel.py similarity index 91% rename from python/cudf_polars/cudf_polars/experimental/parallel.py rename to python/cudf_polars/cudf_polars/streaming/parallel.py index da9e4e6fd6b..5396787d73d 100644 --- a/python/cudf_polars/cudf_polars/experimental/parallel.py +++ b/python/cudf_polars/cudf_polars/streaming/parallel.py @@ -10,13 +10,16 @@ import polars as pl -import cudf_polars.experimental.distinct -import cudf_polars.experimental.groupby -import cudf_polars.experimental.io -import cudf_polars.experimental.join -import cudf_polars.experimental.select -import cudf_polars.experimental.shuffle -import cudf_polars.experimental.sort # noqa: F401 +# Side-effect imports: each module registers ``@lower_ir_node.register(...)`` +# handlers at import time so the dispatch table is populated before any query +# is lowered. +import cudf_polars.streaming.distinct +import cudf_polars.streaming.groupby +import cudf_polars.streaming.io +import cudf_polars.streaming.join +import cudf_polars.streaming.select +import cudf_polars.streaming.shuffle +import cudf_polars.streaming.sort # noqa: F401 from cudf_polars.containers import DataType from cudf_polars.dsl.expr import Col, Literal, NamedExpr from cudf_polars.dsl.ir import ( @@ -33,11 +36,11 @@ ) from cudf_polars.dsl.traversal import CachingVisitor, traversal from cudf_polars.dsl.utils.naming import unique_names -from cudf_polars.experimental.base import PartitionInfo -from cudf_polars.experimental.dispatch import lower_ir_node -from cudf_polars.experimental.io import _clear_source_info_cache -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.utils import ( +from cudf_polars.streaming.base import PartitionInfo +from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.io import _clear_source_info_cache +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.utils import ( _contains_over, _dynamic_planning_on, _lower_ir_fallback, @@ -46,8 +49,8 @@ if TYPE_CHECKING: from collections.abc import MutableMapping - from cudf_polars.experimental.base import StatsCollector - from cudf_polars.experimental.dispatch import LowerIRTransformer, State + from cudf_polars.streaming.base import StatsCollector + from cudf_polars.streaming.dispatch import LowerIRTransformer, State from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -122,7 +125,7 @@ def evaluate_streaming( # Clear source info cache in case data was overwritten _clear_source_info_cache() - from cudf_polars.experimental.rapidsmpf.core import evaluate_logical_plan + from cudf_polars.streaming.actor_graph.core import evaluate_logical_plan result, _ = evaluate_logical_plan(ir, config_options, collect_metadata=False) return result diff --git a/python/cudf_polars/cudf_polars/experimental/repartition.py b/python/cudf_polars/cudf_polars/streaming/repartition.py similarity index 100% rename from python/cudf_polars/cudf_polars/experimental/repartition.py rename to python/cudf_polars/cudf_polars/streaming/repartition.py diff --git a/python/cudf_polars/cudf_polars/experimental/select.py b/python/cudf_polars/cudf_polars/streaming/select.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/select.py rename to python/cudf_polars/cudf_polars/streaming/select.py index 87c34db1b1c..4536f0f3e22 100644 --- a/python/cudf_polars/cudf_polars/experimental/select.py +++ b/python/cudf_polars/cudf_polars/streaming/select.py @@ -14,15 +14,15 @@ from cudf_polars.dsl.ir import Empty, HConcat, HStack, Projection, Scan, Select, Union from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.naming import unique_names -from cudf_polars.experimental.base import PartitionInfo -from cudf_polars.experimental.dispatch import lower_ir_node -from cudf_polars.experimental.expressions import ( +from cudf_polars.streaming.base import PartitionInfo +from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.expressions import ( decompose_expr_graph, make_expr_decomposer, ) -from cudf_polars.experimental.over import _fuse_over_nodes -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.utils import ( +from cudf_polars.streaming.over import _fuse_over_nodes +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.utils import ( _contains_unsupported_fill_strategy, _dynamic_planning_on, _lower_ir_fallback, @@ -32,7 +32,7 @@ from collections.abc import MutableMapping, Sequence from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.parallel import LowerIRTransformer + from cudf_polars.streaming.parallel import LowerIRTransformer from cudf_polars.typing import Schema from cudf_polars.utils.config import ConfigOptions @@ -89,7 +89,7 @@ def _( ) # Partition-wise default - Import here to avoid circular import - from cudf_polars.experimental.parallel import _lower_ir_pwise + from cudf_polars.streaming.parallel import _lower_ir_pwise return _lower_ir_pwise(ir, rec, preserve_partitioning=True) @@ -176,6 +176,7 @@ def decompose_select( new_ir = HConcat( select_ir.schema, True, # noqa: FBT003 + False, # noqa: FBT003 *selections, ) partition_info[new_ir] = PartitionInfo( @@ -350,6 +351,7 @@ def _fuse_simple_reductions( new_hconcat = HConcat( hconcat_schema, True, # noqa: FBT003 + False, # noqa: FBT003 *new_decomposed_select_irs, ) count = max(pi[c].count for c in new_decomposed_select_irs) diff --git a/python/cudf_polars/cudf_polars/experimental/shuffle.py b/python/cudf_polars/cudf_polars/streaming/shuffle.py similarity index 89% rename from python/cudf_polars/cudf_polars/experimental/shuffle.py rename to python/cudf_polars/cudf_polars/streaming/shuffle.py index 9381126775f..b556b217a2c 100644 --- a/python/cudf_polars/cudf_polars/experimental/shuffle.py +++ b/python/cudf_polars/cudf_polars/streaming/shuffle.py @@ -8,8 +8,8 @@ from cudf_polars.dsl.ir import IR from cudf_polars.dsl.tracing import log_do_evaluate -from cudf_polars.experimental.dispatch import lower_ir_node -from cudf_polars.experimental.utils import _dynamic_planning_on +from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.utils import _dynamic_planning_on if TYPE_CHECKING: from collections.abc import MutableMapping @@ -17,8 +17,8 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import IRExecutionContext - from cudf_polars.experimental.dispatch import LowerIRTransformer - from cudf_polars.experimental.parallel import PartitionInfo + from cudf_polars.streaming.dispatch import LowerIRTransformer + from cudf_polars.streaming.parallel import PartitionInfo from cudf_polars.typing import Schema @@ -77,7 +77,7 @@ def _( # Simple lower_ir_node handling for the default hash-based shuffle. # More-complex logic (e.g. joining and sorting) should # be handled separately. - from cudf_polars.experimental.parallel import PartitionInfo + from cudf_polars.streaming.parallel import PartitionInfo (child,) = ir.children diff --git a/python/cudf_polars/cudf_polars/experimental/sort.py b/python/cudf_polars/cudf_polars/streaming/sort.py similarity index 97% rename from python/cudf_polars/cudf_polars/experimental/sort.py rename to python/cudf_polars/cudf_polars/streaming/sort.py index fa610324c2d..652a706bad7 100644 --- a/python/cudf_polars/cudf_polars/experimental/sort.py +++ b/python/cudf_polars/cudf_polars/streaming/sort.py @@ -15,8 +15,8 @@ from cudf_polars.dsl.ir import Slice, Sort from cudf_polars.dsl.traversal import traversal from cudf_polars.dsl.utils.naming import unique_names -from cudf_polars.experimental.dispatch import lower_ir_node -from cudf_polars.experimental.utils import ( +from cudf_polars.streaming.dispatch import lower_ir_node +from cudf_polars.streaming.utils import ( _lower_ir_fallback, ) @@ -26,8 +26,8 @@ from rmm.pylibrmm.stream import Stream from cudf_polars.dsl.ir import IR - from cudf_polars.experimental.base import PartitionInfo - from cudf_polars.experimental.dispatch import LowerIRTransformer + from cudf_polars.streaming.base import PartitionInfo + from cudf_polars.streaming.dispatch import LowerIRTransformer def find_sort_splits( diff --git a/python/cudf_polars/cudf_polars/experimental/statistics.py b/python/cudf_polars/cudf_polars/streaming/statistics.py similarity index 94% rename from python/cudf_polars/cudf_polars/experimental/statistics.py rename to python/cudf_polars/cudf_polars/streaming/statistics.py index bbfd7fbad9c..096808956d3 100644 --- a/python/cudf_polars/cudf_polars/experimental/statistics.py +++ b/python/cudf_polars/cudf_polars/streaming/statistics.py @@ -9,8 +9,8 @@ from cudf_polars.dsl.ir import DataFrameScan, Scan from cudf_polars.dsl.traversal import traversal -from cudf_polars.experimental.base import StatsCollector -from cudf_polars.experimental.io import _build_source_info +from cudf_polars.streaming.base import StatsCollector +from cudf_polars.streaming.io import _build_source_info if TYPE_CHECKING: from cudf_polars.dsl.ir import IR diff --git a/python/cudf_polars/cudf_polars/experimental/utils.py b/python/cudf_polars/cudf_polars/streaming/utils.py similarity index 93% rename from python/cudf_polars/cudf_polars/experimental/utils.py rename to python/cudf_polars/cudf_polars/streaming/utils.py index 848a4d44759..dfe20ef3dbb 100644 --- a/python/cudf_polars/cudf_polars/experimental/utils.py +++ b/python/cudf_polars/cudf_polars/streaming/utils.py @@ -13,7 +13,7 @@ from cudf_polars.dsl.expr import Col, GroupedWindow, UnaryFunction from cudf_polars.dsl.ir import Union from cudf_polars.dsl.traversal import traversal -from cudf_polars.experimental.base import PartitionInfo +from cudf_polars.streaming.base import PartitionInfo if TYPE_CHECKING: from collections.abc import MutableMapping, Sequence @@ -21,7 +21,7 @@ from cudf_polars.containers import DataFrame from cudf_polars.dsl.expr import Expr from cudf_polars.dsl.ir import IR, IRExecutionContext - from cudf_polars.experimental.dispatch import LowerIRTransformer + from cudf_polars.streaming.dispatch import LowerIRTransformer from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -62,8 +62,8 @@ def _lower_ir_fallback( # Catch-all single-partition lowering logic. # If any children contain multiple partitions, # those children will be collapsed with `Repartition`. - from cudf_polars.experimental.repartition import Repartition - from cudf_polars.experimental.select import _inline_hstack_false + from cudf_polars.streaming.repartition import Repartition + from cudf_polars.streaming.select import _inline_hstack_false # Make sure we avoid mixed-length columns in intermediate TableChunks. ir = _inline_hstack_false(ir) diff --git a/python/cudf_polars/cudf_polars/testing/asserts.py b/python/cudf_polars/cudf_polars/testing/asserts.py index 69985bda6d9..fb3e078f26d 100644 --- a/python/cudf_polars/cudf_polars/testing/asserts.py +++ b/python/cudf_polars/cudf_polars/testing/asserts.py @@ -93,7 +93,7 @@ def assert_gpu_result_equal( if ( engine.config.get("executor_options", {}).get("cluster") == "spmd" ): # pragma: no cover - from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( + from cudf_polars.engine.spmd import ( SPMDEngine, allgather_polars_dataframe, ) diff --git a/python/cudf_polars/cudf_polars/testing/engine_utils.py b/python/cudf_polars/cudf_polars/testing/engine_utils.py index b80afc8ea86..f3af427ed9c 100644 --- a/python/cudf_polars/cudf_polars/testing/engine_utils.py +++ b/python/cudf_polars/cudf_polars/testing/engine_utils.py @@ -14,8 +14,8 @@ import polars as pl - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.engine.core import StreamingEngine + from cudf_polars.engine.options import StreamingOptions STREAMING_ENGINE_FIXTURE_PARAMS: list[str] = [] @@ -66,7 +66,7 @@ def __init__(self, full_name: str): def is_streaming_engine(obj: Any) -> bool: """Return ``True`` if ``obj`` is a :class:`StreamingEngine`.""" try: - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine + from cudf_polars.engine.core import StreamingEngine except ImportError: # pragma: no cover; only triggered without rapidsmpf return False return isinstance(obj, StreamingEngine) @@ -93,7 +93,7 @@ def warns_on_spmd( # pragma: no cover; rapidsmpf-only path import pytest - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine if when and isinstance(engine, SPMDEngine): return pytest.warns(*args, **kwargs) @@ -117,7 +117,7 @@ def create_streaming_options( ------- The streaming options for the given block size. """ - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.engine.options import StreamingOptions from cudf_polars.utils.config import StreamingFallbackMode # ``allow_gpu_sharing=True`` is always set so the cached multi-rank @@ -162,7 +162,7 @@ def merge_streaming_options( ------- The merged streaming options with overrides overriding any base options. """ - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.engine.options import StreamingOptions return StreamingOptions(**{**base.to_dict(), **overrides.to_dict()}) diff --git a/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py b/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py index a7a1d189df2..4b5f2a0e39c 100644 --- a/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py +++ b/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py @@ -60,7 +60,7 @@ def pytest_configure(config: pytest.Config) -> None: if variant == "in-memory": engine = polars.GPUEngine(executor="in-memory", raise_on_fail=raise_on_fail) else: - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine executor_options: dict[str, object] = {} if blocksize == "small": @@ -135,8 +135,7 @@ def pytest_report_header(config: pytest.Config) -> str: return f"injected GPU engine: {cls.__module__}.{cls.__name__}" -# TODO: This is just Mapping[str, str]? -EXPECTED_FAILURES: Mapping[str, str | tuple[str, bool]] = { +EXPECTED_FAILURES: Mapping[str, str] = { "tests/unit/io/test_csv.py::test_read_csv_only_loads_selected_columns": "Memory usage won't be correct due to GPU", "tests/unit/io/test_delta.py::test_scan_delta_version": "Need to expose hive partitioning", "tests/unit/io/test_delta.py::test_scan_delta_relative": "Need to expose hive partitioning", @@ -256,7 +255,15 @@ def pytest_report_header(config: pytest.Config) -> str: "tests/unit/operations/test_slice.py::test_schema_slice_on_literal_23999[lit1-0-len1-False]": "List literal loses nesting in slice: cudf#19610", "tests/unit/operations/test_slice.py::test_schema_slice_on_literal_23999[lit1-offset1-0-False]": "List literal loses nesting in slice: cudf#19610", "tests/unit/operations/test_slice.py::test_schema_slice_on_literal_23999[lit1-offset1-len1-False]": "List literal loses nesting in slice: cudf#19610", - "tests/unit/functions/test_concat.py::test_concat_horizontally_strict": "polars doesnt hand us the hconcat options. Fixed in 1.39.", + "tests/unit/functions/test_concat.py::test_concat_with_empty_dataframes_strict_25725": "https://github.com/rapidsai/cudf/issues/21644", + "tests/unit/sql/test_window_functions.py::test_over_with_order_by": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_over_with_cumulative_window_funcs": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_window_function_order_by_multi": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_window_cumulative_agg_with_nulls": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_window_named_window": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_window_multiple_named_windows": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_window_frame_validation": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/operations/test_window.py::test_over_literal_cum_sum_26800": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", } @@ -302,6 +309,12 @@ def pytest_report_header(config: pytest.Config) -> str: # Short term adds in the aftermath of the rapidsmpf switch to get CI passing "tests/unit/io/test_lazy_parquet.py::test_scan_parquet_local_with_async": "Flaky, otherwise TBD", "tests/unit/operations/test_join.py::test_join_where_nested_expr_21066": "Flaky, otherwise TBD", + "tests/unit/io/test_scan.py::test_scan_metrics[True-parquet]": "Checks to IO metric logs specific to Polars CPU", + "tests/unit/io/test_scan.py::test_scan_metrics[True-csv]": "Checks to IO metric logs specific to Polars CPU", + "tests/unit/io/test_scan.py::test_scan_metrics[True-ndjson]": "Checks to IO metric logs specific to Polars CPU", + "tests/unit/io/test_scan.py::test_scan_metrics[False-parquet]": "Checks to IO metric logs specific to Polars CPU", + "tests/unit/io/test_scan.py::test_scan_metrics[False-csv]": "Checks to IO metric logs specific to Polars CPU", + "tests/unit/io/test_scan.py::test_scan_metrics[False-ndjson]": "Checks to IO metric logs specific to Polars CPU", } @@ -370,12 +383,14 @@ def pytest_report_header(config: pytest.Config) -> str: # xfail for tests that produce different results than CPU Polars STREAMING_ENGINE_EXPECTED_FAILURES: Mapping[str, str] = { "tests/unit/functions/range/test_linear_space.py::test_linear_space_num_samples_expr": "https://github.com/rapidsai/cudf/issues/22072", + "tests/unit/functions/test_concat.py::test_concat_horizontal_zero_width_height_mismatch_26876": "https://github.com/rapidsai/cudf/issues/21644", + "tests/unit/functions/test_concat.py::test_concat_horizontally_strict": "Correct polars.exceptions.ShapeError raised but it's in a ExceptionGroup", + "tests/unit/interop/test_interop.py::test_0_width_df_roundtrip": "https://github.com/rapidsai/cudf/issues/21644", "tests/unit/lazyframe/test_projections.py::test_join_projection_pushdown_struct_field_as_key_24446": "https://github.com/rapidsai/cudf/issues/22105", "tests/unit/operations/test_slice.py::test_slice_pushdown_literal_projection_14349": "https://github.com/rapidsai/cudf/issues/22072", "tests/unit/operations/test_group_by.py::test_group_by_lit_series": "Incorrect broadcasting of literals in groupby-agg", "tests/unit/operations/test_group_by.py::test_group_by_series_partitioned": "https://github.com/rapidsai/cudf/issues/22072", "tests/unit/operations/test_group_by.py::test_partitioned_group_by_chunked": "https://github.com/rapidsai/cudf/issues/22072", - "tests/unit/interop/test_interop.py::test_0_width_df_roundtrip": "https://github.com/rapidsai/cudf/issues/21644", "tests/unit/operations/test_group_by.py::test_group_by_unique_parametric[n_unique-True-True]": "https://github.com/rapidsai/cudf/issues/21641", "tests/unit/operations/test_group_by.py::test_unique_head_tail_26429[1]": "https://github.com/rapidsai/cudf/issues/22075", "tests/unit/operations/test_group_by.py::test_unique_head_tail_26429[4]": "https://github.com/rapidsai/cudf/issues/22075", @@ -425,7 +440,14 @@ def pytest_report_header(config: pytest.Config) -> str: "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes43]": "https://github.com/rapidsai/cudf/issues/22085", "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_15338[False-dtypes44]": "https://github.com/rapidsai/cudf/issues/22085", "tests/unit/operations/test_join.py::test_join_numeric_key_upcast_order": "https://github.com/rapidsai/cudf/issues/22085", + "tests/unit/operations/test_window.py::test_over_literal_cum_sum_26800": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", "tests/unit/sql/test_joins.py::test_cross_join_unnest_from_cte": "https://github.com/rapidsai/cudf/issues/22073", + "tests/unit/sql/test_window_functions.py::test_over_with_cumulative_window_funcs": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_over_with_order_by": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_window_cumulative_agg_with_nulls": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_window_function_order_by_multi": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_window_frame_validation": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", + "tests/unit/sql/test_window_functions.py::test_window_multiple_named_window": "TODO: https://github.com/rapidsai/cudf/pull/22048#discussion_r3238041970", } @@ -450,15 +472,5 @@ def pytest_collection_modifyitems( is not None ): item.add_marker(pytest.mark.xfail(reason=s_reason)) - elif (entry := EXPECTED_FAILURES.get(item.nodeid, None)) is not None: - if isinstance(entry, tuple): - # the second entry in the tuple is the condition to xfail on - reason, condition = entry - item.add_marker( - pytest.mark.xfail( - condition=condition, - reason=reason, - ), - ) - else: - item.add_marker(pytest.mark.xfail(reason=entry)) + elif (reason := EXPECTED_FAILURES.get(item.nodeid, None)) is not None: + item.add_marker(pytest.mark.xfail(reason=reason)) diff --git a/python/cudf_polars/cudf_polars/testing/io.py b/python/cudf_polars/cudf_polars/testing/io.py index d43f6f3d983..98ec8d7da6e 100644 --- a/python/cudf_polars/cudf_polars/testing/io.py +++ b/python/cudf_polars/cudf_polars/testing/io.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 """IO testing utilities.""" @@ -99,7 +99,7 @@ def make_lazy_frame( n_rows : optional, int Slice to apply to the final LazyFrame before returning. """ - from cudf_polars.experimental.io import _clear_source_info_cache + from cudf_polars.streaming.io import _clear_source_info_cache _clear_source_info_cache() diff --git a/python/cudf_polars/cudf_polars/utils/config.py b/python/cudf_polars/cudf_polars/utils/config.py index 5b0f30f88c1..b1d1c0c975c 100644 --- a/python/cudf_polars/cudf_polars/utils/config.py +++ b/python/cudf_polars/cudf_polars/utils/config.py @@ -44,7 +44,7 @@ import rmm.mr - from cudf_polars.experimental.rapidsmpf.frontend.ray import RankActor + from cudf_polars.engine.ray import RankActor __all__ = [ @@ -462,9 +462,8 @@ class SPMDContext: This dataclass is **not picklable** because :class:`Communicator`, :class:`Context`, and :class:`~concurrent.futures.ThreadPoolExecutor` cannot be serialized. In SPMD mode each rank constructs its own - ``SPMDContext`` locally inside - :class:`~cudf_polars.experimental.rapidsmpf.frontend.spmd.SPMDEngine`, so - pickling is never required. Do not use this class with Dask or any other + ``SPMDContext`` locally inside :class:`~cudf_polars.engine.spmd.SPMDEngine`, + so pickling is never required. Do not use this class with Dask or any other framework that serializes executor configuration across process boundaries. Parameters @@ -490,16 +489,15 @@ class RayContext: .. note:: This dataclass holds Ray actor handles, which are only valid within the Ray session that created them. It is stripped from ``config_options`` - before pickling for remote actor calls in - :func:`~cudf_polars.experimental.rapidsmpf.frontend.ray.evaluate_pipeline_ray_mode` - by :class:`~cudf_polars.experimental.rapidsmpf.frontend.ray.RayEngine`. - Do not persist or transfer this object across Ray sessions. + before pickling for remote actor calls in :func:`~cudf_polars.engine.ray.evaluate_pipeline_ray_mode` + by :class:`~cudf_polars.engine.ray.RayEngine`. Do not persist or transfer + this object across Ray sessions. Parameters ---------- rank_actors - List of :class:`~cudf_polars.experimental.rapidsmpf.frontend.ray.RankActor` - handles, one per GPU in the cluster. + List of :class:`~cudf_polars.engine.ray.RankActor` handles, one per GPU + in the cluster. """ rank_actors: list[ActorHandle[RankActor]] @@ -514,7 +512,7 @@ class DaskContext: This dataclass holds a :class:`~distributed.Client` handle, which is only valid within the Dask session that created it. It is stripped from ``config_options`` before pickling for remote worker calls in - :func:`~cudf_polars.experimental.rapidsmpf.frontend.dask.evaluate_pipeline_dask_mode`. + :func:`~cudf_polars.engine.dask.evaluate_pipeline_dask_mode`. Do not persist or transfer this object across Dask sessions. Parameters @@ -525,7 +523,7 @@ class DaskContext: Unique identifier for this RapidsMPF bootstrap session. owned_client Client to close on shutdown, if created internally by - :class:`~cudf_polars.experimental.rapidsmpf.frontend.dask.DaskEngine`. + :class:`~cudf_polars.engine.dask.DaskEngine`. owned_cluster Cluster to close on shutdown, if created internally. """ diff --git a/python/cudf_polars/cudf_polars/utils/versions.py b/python/cudf_polars/cudf_polars/utils/versions.py index 4b8e94f479b..25e4be8ed1d 100644 --- a/python/cudf_polars/cudf_polars/utils/versions.py +++ b/python/cudf_polars/cudf_polars/utils/versions.py @@ -15,6 +15,7 @@ POLARS_VERSION_LT_136 = POLARS_VERSION < parse("1.36.0") POLARS_VERSION_LT_137 = POLARS_VERSION < parse("1.37.0") POLARS_VERSION_LT_138 = POLARS_VERSION < parse("1.38.0") +POLARS_VERSION_LT_139 = POLARS_VERSION < parse("1.39.0") def _ensure_polars_version() -> None: diff --git a/python/cudf_polars/docs/cudf-polars-mp.md b/python/cudf_polars/docs/cudf-polars-mp.md index 4b56d67dea9..f11fd24ac17 100644 --- a/python/cudf_polars/docs/cudf-polars-mp.md +++ b/python/cudf_polars/docs/cudf-polars-mp.md @@ -54,7 +54,7 @@ environment variable if set, otherwise let the underlying library apply its own built-in default. ```python -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions opts = StreamingOptions( num_streaming_threads=8, @@ -68,9 +68,9 @@ Pass the options object to `from_options()` on any engine — this is the recommended constructor for typical use: ```python -from cudf_polars.experimental.rapidsmpf.frontend.dask import DaskEngine -from cudf_polars.experimental.rapidsmpf.frontend.ray import RayEngine -from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine +from cudf_polars.engine.dask import DaskEngine +from cudf_polars.engine.ray import RayEngine +from cudf_polars.engine.spmd import SPMDEngine with RayEngine.from_options(opts) as engine: result = df.lazy().collect(engine=engine) @@ -148,7 +148,7 @@ Binding is controlled by the `hardware_binding` executor option, which accepts a `HardwareBindingPolicy` instance: ```python -from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( +from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, ) ``` @@ -242,7 +242,7 @@ broadcasts it to all actors, so every rank always executes the same query. ### Running in Ray mode -`RayEngine` is imported from `cudf_polars.experimental.rapidsmpf.frontend.ray`. On construction it: +`RayEngine` is imported from `cudf_polars.engine.ray`. On construction it: 1. Calls `ray.init()` if Ray is not already running 2. Creates one `RankActor` per GPU @@ -255,8 +255,8 @@ The recommended way to construct a `RayEngine` is via `from_options()`: ```python import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.rapidsmpf.frontend.ray import RayEngine +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.ray import RayEngine opts = StreamingOptions(num_streaming_threads=8, fallback_mode="silent") @@ -287,7 +287,7 @@ does not call `ray.shutdown()` on exit. ```python import ray import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.ray import RayEngine +from cudf_polars.engine.ray import RayEngine ray.init(address="auto") @@ -384,7 +384,7 @@ Conceptually the system looks like this: ### Running in Dask mode -`DaskEngine` is imported from `cudf_polars.experimental.rapidsmpf.frontend.dask`. On construction it: +`DaskEngine` is imported from `cudf_polars.engine.dask`. On construction it: 1. If `dask_client` is `None`, creates a `distributed.LocalCluster` (one worker per visible GPU) and a `distributed.Client` 2. Bootstraps a UCXX communicator across all workers @@ -393,7 +393,7 @@ Conceptually the system looks like this: ```python import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.dask import DaskEngine +from cudf_polars.engine.dask import DaskEngine with DaskEngine() as engine: result = ( @@ -419,7 +419,7 @@ Bring-your-own-client variant: ```python from distributed import Client import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.dask import DaskEngine +from cudf_polars.engine.dask import DaskEngine with Client("scheduler-address:8786") as dc: with DaskEngine(dask_client=dc) as engine: @@ -444,8 +444,8 @@ sets `CUDA_VISIBLE_DEVICES` per worker — disable the built-in binding to avoid conflicts: ```python -from cudf_polars.experimental.rapidsmpf.frontend.dask import DaskEngine -from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( +from cudf_polars.engine.dask import DaskEngine +from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, ) @@ -467,14 +467,14 @@ built-in nanny preload to assign one GPU per worker. The preload sets ```bash # On each node — launch one worker per GPU with a single thread each: dask worker SCHEDULER:8786 --nworkers N --nthreads 1 \ - --preload-nanny cudf_polars.experimental.rapidsmpf.frontend.dask + --preload-nanny cudf_polars.engine.dask ``` Then connect from the client: ```python from distributed import Client -from cudf_polars.experimental.rapidsmpf.frontend.dask import DaskEngine +from cudf_polars.engine.dask import DaskEngine with Client("SCHEDULER:8786") as dc: with DaskEngine(dask_client=dc) as engine: @@ -584,7 +584,7 @@ every rank, call `allgather_polars_dataframe()`. ### Running in SPMD mode `SPMDEngine` is the primary entry point for SPMD execution. It is a context -manager imported from `cudf_polars.experimental.rapidsmpf.frontend.spmd`. On construction it: +manager imported from `cudf_polars.engine.spmd`. On construction it: 1. Bootstraps a communicator: UCXX when running under `rrun`, otherwise a single-rank communicator that requires no external library. @@ -601,9 +601,9 @@ The recommended way to construct an `SPMDEngine` is via `from_options()`: # multi-GPU launch: rrun -n 4 python my_script.py # single-GPU (no rrun needed): python my_script.py import polars as pl -from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( +from cudf_polars.streaming.collectives.common import reserve_op_id +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.spmd import ( SPMDEngine, allgather_polars_dataframe, ) @@ -686,8 +686,8 @@ with SPMDEngine() as engine: `allgather_polars_dataframe()` to gather all fragments: ```python -from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id -from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( +from cudf_polars.streaming.collectives.common import reserve_op_id +from cudf_polars.engine.spmd import ( SPMDEngine, allgather_polars_dataframe, ) @@ -723,7 +723,7 @@ retains ownership and can reuse it across multiple `SPMDEngine` lifetimes. ```python from rapidsmpf import bootstrap from rapidsmpf.progress_thread import ProgressThread -from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine +from cudf_polars.engine.spmd import SPMDEngine # Bootstrap once. comm = bootstrap.create_ucxx_comm(progress_thread=ProgressThread()) diff --git a/python/cudf_polars/docs/overview.md b/python/cudf_polars/docs/overview.md index cbc2eb66b63..307f95bebec 100644 --- a/python/cudf_polars/docs/overview.md +++ b/python/cudf_polars/docs/overview.md @@ -631,19 +631,19 @@ annotated with nvtx ranges. # Query Plans -The module `cudf_polars.experimental.explain` contains functions for dumping +The module `cudf_polars.streaming.explain` contains functions for dumping the query for a given `LazyFrame`. ## Structured Output -`cudf_polars.experimental.explain.serialize_query` can be used to output +`cudf_polars.streaming.explain.serialize_query` can be used to output the query plan in a structured format. ```python >>> import dataclasses >>> import polars as pl ->>> from cudf_polars.experimental.explain import serialize_query +>>> from cudf_polars.streaming.explain import serialize_query >>> q = pl.LazyFrame({"a": ['a', 'b', 'a'], "b": [1, 2, 3]}).group_by("a").agg(pl.len()) >>> dataclasses.asdict(serialize_query(q, engine=pl.GPUEngine())) {'roots': ['526964741'], diff --git a/python/cudf_polars/pyproject.toml b/python/cudf_polars/pyproject.toml index 9827ea67c74..c881aaa9aea 100644 --- a/python/cudf_polars/pyproject.toml +++ b/python/cudf_polars/pyproject.toml @@ -24,7 +24,7 @@ dependencies = [ "cuda-python>=13.0.1,<14.0", "nvidia-ml-py>=12", "packaging", - "polars>=1.35,<1.39", + "polars>=1.35,<1.40", "pylibcudf==26.8.*,>=0.0.0a0", "rapidsmpf==26.8.*,>=0.0.0a0", "typing_extensions>=4.0.0", @@ -96,12 +96,12 @@ exclude_also = [ "class .*\\bProtocol(?:\\[[^]]+\\])?\\):", "assert_never\\(" ] -# The cudf_polars test suite does not exercise the plugin or all experimental -# tests (some require optional dependencies and others are not yet fully covered), -# so we omit them from coverage checks. +# The cudf_polars test suite are missing some coverage. +# See https://github.com/rapidsai/cudf/issues/22568. omit = [ "cudf_polars/testing/inject_gpu_engine.py", - "cudf_polars/experimental/**", + "cudf_polars/streaming/**", + "cudf_polars/engine/**", ] [tool.ruff] diff --git a/python/cudf_polars/tests/conftest.py b/python/cudf_polars/tests/conftest.py index 0e4938b775f..d02e8fef353 100644 --- a/python/cudf_polars/tests/conftest.py +++ b/python/cudf_polars/tests/conftest.py @@ -20,9 +20,9 @@ if TYPE_CHECKING: from collections.abc import Callable, Generator - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.core import StreamingEngine + from cudf_polars.engine.options import StreamingOptions + from cudf_polars.engine.spmd import SPMDEngine # Number of ranks for multi-rank streaming engines that share one GPU @@ -95,15 +95,15 @@ def _unconfigured_engine( engine: StreamingEngine match _engine_param.engine_name: case "spmd": - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine engine = SPMDEngine() case "dask": # pragma: no cover - from cudf_polars.experimental.rapidsmpf.frontend.dask import DaskEngine + from cudf_polars.engine.dask import DaskEngine engine = DaskEngine(engine_options={"allow_gpu_sharing": True}) case "ray": # pragma: no cover - from cudf_polars.experimental.rapidsmpf.frontend.ray import RayEngine + from cudf_polars.engine.ray import RayEngine # Always specify num_ranks so the engine has a fixed size # regardless of how many GPUs the host happens to have; @@ -240,7 +240,7 @@ def engine( engine, options = _unconfigured_engine if options is None: return engine - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine + from cudf_polars.engine.core import StreamingEngine assert isinstance(engine, StreamingEngine) return configure_streaming_engine(engine, options) diff --git a/python/cudf_polars/tests/experimental/__init__.py b/python/cudf_polars/tests/experimental/__init__.py deleted file mode 100644 index db71916f3c4..00000000000 --- a/python/cudf_polars/tests/experimental/__init__.py +++ /dev/null @@ -1,8 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -"""Testing experimental features.""" - -from __future__ import annotations - -__all__: list[str] = [] diff --git a/python/cudf_polars/tests/expressions/test_booleanfunction.py b/python/cudf_polars/tests/expressions/test_booleanfunction.py index 94463c62ca7..93a7d3ebeec 100644 --- a/python/cudf_polars/tests/expressions/test_booleanfunction.py +++ b/python/cudf_polars/tests/expressions/test_booleanfunction.py @@ -165,6 +165,21 @@ def test_boolean_isbetween(engine: pl.GPUEngine, closed, bounds): assert_gpu_result_equal(q, engine=engine) +@pytest.mark.parametrize("closed", ["both", "left", "right", "none"]) +def test_boolean_isbetween_decimal_float(engine: pl.GPUEngine, closed): + df = pl.LazyFrame( + { + "a": pl.Series([1, 2, 3, 4], dtype=pl.Decimal(scale=2)), + "lo": pl.Series([0.5, 1.5, 2.5, 3.5], dtype=pl.Float64), + "hi": pl.Series([1.5, 2.5, 3.5, 4.5], dtype=pl.Float64), + } + ) + + q = df.select(pl.col("a").is_between(pl.col("lo"), pl.col("hi"), closed=closed)) + + assert_gpu_result_equal(q, engine=engine) + + @pytest.mark.parametrize( "expr", [pl.any_horizontal("*"), pl.all_horizontal("*")], ids=["any", "all"] ) diff --git a/python/cudf_polars/tests/expressions/test_rolling.py b/python/cudf_polars/tests/expressions/test_rolling.py index 6c00f66347d..f3384b2af73 100644 --- a/python/cudf_polars/tests/expressions/test_rolling.py +++ b/python/cudf_polars/tests/expressions/test_rolling.py @@ -13,12 +13,21 @@ assert_gpu_result_equal, assert_ir_translation_raises, ) -from cudf_polars.utils.versions import POLARS_VERSION_LT_136 +from cudf_polars.utils.versions import POLARS_VERSION_LT_136, POLARS_VERSION_LT_139 if TYPE_CHECKING: from cudf_polars.typing import RankMethod +# In polars 1.36-1.38, rolling window expressions (pl.col(...).rolling(...)) +# are represented as an opaque Rust node that is not accessible via view_expression(). +# Support requires polars <1.36 (via Window+RollingGroupOptions) or >=1.39 (via Rolling). +skip_rolling_expr_136_to_138 = pytest.mark.skipif( + not POLARS_VERSION_LT_136 and POLARS_VERSION_LT_139, + reason="Rolling window expressions are not accessible in polars 1.36-1.38, see https://github.com/pola-rs/polars/pull/25117", +) + + @pytest.fixture def df(): return pl.LazyFrame( @@ -32,14 +41,9 @@ def df(): ) +@skip_rolling_expr_136_to_138 @pytest.mark.parametrize("time_unit", ["ns", "us", "ms"]) -def test_rolling_datetime(engine: pl.GPUEngine, request, time_unit): - request.applymarker( - pytest.mark.xfail( - condition=not POLARS_VERSION_LT_136, - reason="Polars translates this to AExpr::Rolling which is NotImplemented", - ) - ) +def test_rolling_datetime(engine: pl.GPUEngine, time_unit): dates = [ "2020-01-01 13:45:48", "2020-01-01 16:42:13", @@ -62,11 +66,8 @@ def test_rolling_datetime(engine: pl.GPUEngine, request, time_unit): assert_gpu_result_equal(q, engine=engine) -def test_rolling_date(engine: pl.GPUEngine, request): - if not POLARS_VERSION_LT_136: - request.applymarker( - pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") - ) +@skip_rolling_expr_136_to_138 +def test_rolling_date(engine: pl.GPUEngine): dates = [ "2020-01-01", "2020-01-01", @@ -87,12 +88,9 @@ def test_rolling_date(engine: pl.GPUEngine, request): assert_gpu_result_equal(q, engine=engine) +@skip_rolling_expr_136_to_138 @pytest.mark.parametrize("dtype", [pl.Int32, pl.UInt32, pl.Int64, pl.UInt64]) -def test_rolling_integral_orderby(engine: pl.GPUEngine, request, dtype): - if not POLARS_VERSION_LT_136: - request.applymarker( - pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") - ) +def test_rolling_integral_orderby(engine: pl.GPUEngine, dtype): df = pl.LazyFrame( { "orderby": pl.Series([1, 4, 8, 10, 12, 13, 14, 22], dtype=dtype), @@ -106,14 +104,11 @@ def test_rolling_integral_orderby(engine: pl.GPUEngine, request, dtype): assert_gpu_result_equal(q, engine=engine) -@pytest.mark.skipif( - POLARS_VERSION_LT_136, - reason="Rolling expression node only exists in polars >= 1.36", -) -def test_rolling_translation_raises(): +@skip_rolling_expr_136_to_138 +def test_rolling_agg_before_rolling(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3], "b": [1, 2, 3]}) q = df.with_columns(pl.col("a").sum().rolling("b", period="2i")) - assert_ir_translation_raises(q, NotImplementedError) + assert_gpu_result_equal(q, engine=engine) def test_rolling_collect_list_raises(): @@ -129,11 +124,8 @@ def test_rolling_collect_list_raises(): ) -def test_unsorted_raises(engine_raise_on_fail: pl.GPUEngine, request): - if not POLARS_VERSION_LT_136: - request.applymarker( - pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") - ) +@skip_rolling_expr_136_to_138 +def test_unsorted_raises(engine_raise_on_fail: pl.GPUEngine): df = pl.LazyFrame({"orderby": [1, 2, 4, 2], "values": [1, 2, 3, 4]}) q = df.select(pl.col("values").sum().rolling("orderby", period="2i")) with pytest.raises(pl.exceptions.InvalidOperationError): @@ -145,11 +137,8 @@ def test_unsorted_raises(engine_raise_on_fail: pl.GPUEngine, request): q.collect(engine=engine_raise_on_fail) -def test_orderby_nulls_raises_computeerror(engine_raise_on_fail: pl.GPUEngine, request): - if not POLARS_VERSION_LT_136: - request.applymarker( - pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") - ) +@skip_rolling_expr_136_to_138 +def test_orderby_nulls_raises_computeerror(engine_raise_on_fail: pl.GPUEngine): df = pl.LazyFrame({"orderby": [1, 2, 4, None], "values": [1, 2, 3, 4]}) q = df.select(pl.col("values").sum().rolling("orderby", period="2i")) with pytest.raises(pl.exceptions.InvalidOperationError): @@ -160,23 +149,15 @@ def test_orderby_nulls_raises_computeerror(engine_raise_on_fail: pl.GPUEngine, r q.collect(engine=engine_raise_on_fail) -def test_invalid_duration_spec_raises_in_translation(request): - if not POLARS_VERSION_LT_136: - request.applymarker( - pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") - ) +@skip_rolling_expr_136_to_138 +def test_invalid_duration_spec_raises_in_translation(): df = pl.LazyFrame({"orderby": [1, 2, 4, 5], "values": [1, 2, 3, 4]}) q = df.select(pl.col("values").sum().rolling("orderby", period="3d")) assert_ir_translation_raises(q, pl.exceptions.InvalidOperationError) -def test_rolling_inside_groupby_raises(request): - request.applymarker( - pytest.mark.xfail( - condition=not POLARS_VERSION_LT_136, - reason="not supported as of polars 1.36", - ) - ) +@pytest.mark.xfail(condition=not POLARS_VERSION_LT_136, reason="not supported") +def test_rolling_inside_groupby_raises(): df = pl.LazyFrame( {"keys": [1, 1, 1, 2], "orderby": [1, 2, 4, 2], "values": [1, 2, 3, 4]} ) @@ -188,11 +169,8 @@ def test_rolling_inside_groupby_raises(request): assert_ir_translation_raises(q, NotImplementedError) -def test_rolling_sum_all_null_window_returns_null(engine: pl.GPUEngine, request): - if not POLARS_VERSION_LT_136: - request.applymarker( - pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") - ) +@skip_rolling_expr_136_to_138 +def test_rolling_sum_all_null_window_returns_null(engine: pl.GPUEngine): df = pl.LazyFrame( { "orderby": [1, 2, 3, 4, 5, 6], diff --git a/python/cudf_polars/tests/expressions/test_sort.py b/python/cudf_polars/tests/expressions/test_sort.py index 3f2f0dccc66..0fd0354308e 100644 --- a/python/cudf_polars/tests/expressions/test_sort.py +++ b/python/cudf_polars/tests/expressions/test_sort.py @@ -78,6 +78,20 @@ def test_setsorted(engine: pl.GPUEngine, request, descending, nulls_last, with_n assert_gpu_result_equal(q, engine=engine) +@pytest.mark.parametrize("descending", [False, True]) +@pytest.mark.parametrize("nulls_last", [False, True]) +def test_setsorted_expr_with_nulls(engine: pl.GPUEngine, descending, nulls_last): + sorted_values = sorted([1, 2, 3, 4, 5], reverse=descending) + values: list[int | None] = [*sorted_values] + if nulls_last: + values.append(None) + else: + values.insert(0, None) + ldf = pl.LazyFrame({"a": values}) + q = ldf.select(pl.col("a").set_sorted(descending=descending)) + assert_gpu_result_equal(q, engine=engine) + + def test_sort_concat_filtered_to_empty(engine: pl.GPUEngine): df = pl.LazyFrame({"a": [1, 2, 3]}) q = pl.concat([df.filter(pl.col("a") == 0), df.filter(pl.col("a") == 4)]).sort("a") diff --git a/python/cudf_polars/tests/streaming/__init__.py b/python/cudf_polars/tests/streaming/__init__.py new file mode 100644 index 00000000000..1e0eca9b14d --- /dev/null +++ b/python/cudf_polars/tests/streaming/__init__.py @@ -0,0 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""Testing the streaming execution layer.""" + +from __future__ import annotations + +__all__: list[str] = [] diff --git a/python/cudf_polars/tests/experimental/test_agg.py b/python/cudf_polars/tests/streaming/test_agg.py similarity index 100% rename from python/cudf_polars/tests/experimental/test_agg.py rename to python/cudf_polars/tests/streaming/test_agg.py diff --git a/python/cudf_polars/tests/experimental/test_all_gather_host_data.py b/python/cudf_polars/tests/streaming/test_all_gather_host_data.py similarity index 94% rename from python/cudf_polars/tests/experimental/test_all_gather_host_data.py rename to python/cudf_polars/tests/streaming/test_all_gather_host_data.py index c85598a8c64..da53cf58b69 100644 --- a/python/cudf_polars/tests/experimental/test_all_gather_host_data.py +++ b/python/cudf_polars/tests/streaming/test_all_gather_host_data.py @@ -9,11 +9,11 @@ import pytest -from cudf_polars.experimental.rapidsmpf.frontend.core import ( +from cudf_polars.engine.core import ( ClusterInfo, all_gather_host_data, ) -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions pytestmark = pytest.mark.spmd diff --git a/python/cudf_polars/tests/experimental/test_allgather.py b/python/cudf_polars/tests/streaming/test_allgather.py similarity index 94% rename from python/cudf_polars/tests/experimental/test_allgather.py rename to python/cudf_polars/tests/streaming/test_allgather.py index 4b3a751214a..48d84afdb00 100644 --- a/python/cudf_polars/tests/experimental/test_allgather.py +++ b/python/cudf_polars/tests/streaming/test_allgather.py @@ -15,8 +15,8 @@ import pylibcudf as plc from cudf_polars.dsl.ir import IRExecutionContext -from cudf_polars.experimental.rapidsmpf.collectives.allgather import AllGatherManager -from cudf_polars.experimental.rapidsmpf.utils import allgather_reduce +from cudf_polars.streaming.actor_graph.collectives.allgather import AllGatherManager +from cudf_polars.streaming.actor_graph.utils import allgather_reduce async def _test_allgather(engine) -> None: diff --git a/python/cudf_polars/tests/experimental/test_bind_to_gpu.py b/python/cudf_polars/tests/streaming/test_bind_to_gpu.py similarity index 83% rename from python/cudf_polars/tests/experimental/test_bind_to_gpu.py rename to python/cudf_polars/tests/streaming/test_bind_to_gpu.py index aa1c439bc01..1b177c379bb 100644 --- a/python/cudf_polars/tests/experimental/test_bind_to_gpu.py +++ b/python/cudf_polars/tests/streaming/test_bind_to_gpu.py @@ -78,7 +78,7 @@ def _run_in_subprocess(target: Callable[[], None]) -> None: def _reset_bind_state() -> None: """Reset the module-level bind state so each subprocess starts clean.""" - from cudf_polars.experimental.rapidsmpf.frontend import hardware_binding + from cudf_polars.engine import hardware_binding hardware_binding._bind_done = False @@ -91,10 +91,8 @@ def _reset_bind_state() -> None: def _body_bind_called_once() -> None: _reset_bind_state() - with patch( - "cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.bind" - ) as mock_bind: - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( + with patch("cudf_polars.engine.hardware_binding.bind") as mock_bind: + from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, bind_to_gpu, ) @@ -114,10 +112,10 @@ def _body_bind_falls_back_to_gpu_0() -> None: _reset_bind_state() mock_bind = MagicMock(side_effect=[RuntimeError("no CUDA_VISIBLE_DEVICES"), None]) with patch( - "cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.bind", + "cudf_polars.engine.hardware_binding.bind", mock_bind, ): - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( + from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, bind_to_gpu, ) @@ -145,10 +143,10 @@ def _body_bind_raise_on_fail_propagates_exception() -> None: _reset_bind_state() mock_bind = MagicMock(side_effect=RuntimeError("binding failed")) with patch( - "cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.bind", + "cudf_polars.engine.hardware_binding.bind", mock_bind, ): - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( + from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, bind_to_gpu, ) @@ -166,10 +164,10 @@ def _body_bind_raise_on_fail_false_suppresses_exception() -> None: _reset_bind_state() mock_bind = MagicMock(side_effect=RuntimeError("binding failed")) with patch( - "cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.bind", + "cudf_polars.engine.hardware_binding.bind", mock_bind, ): - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( + from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, bind_to_gpu, ) @@ -186,10 +184,8 @@ def _body_bind_thread_safe() -> None: import threading _reset_bind_state() - with patch( - "cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.bind" - ) as mock_bind: - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( + with patch("cudf_polars.engine.hardware_binding.bind") as mock_bind: + from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, bind_to_gpu, ) @@ -216,11 +212,11 @@ def test_bind_thread_safe() -> None: def _body_bind_done_flag_set() -> None: - from cudf_polars.experimental.rapidsmpf.frontend import hardware_binding + from cudf_polars.engine import hardware_binding _reset_bind_state() assert not hardware_binding._bind_done - with patch("cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.bind"): + with patch("cudf_polars.engine.hardware_binding.bind"): hardware_binding.bind_to_gpu(hardware_binding.HardwareBindingPolicy()) assert hardware_binding._bind_done @@ -232,10 +228,8 @@ def test_bind_done_flag_set() -> None: def _body_bind_disabled() -> None: _reset_bind_state() - with patch( - "cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.bind" - ) as mock_bind: - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( + with patch("cudf_polars.engine.hardware_binding.bind") as mock_bind: + from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, bind_to_gpu, ) @@ -251,10 +245,8 @@ def test_bind_disabled() -> None: def _body_bind_enable_once_false() -> None: _reset_bind_state() - with patch( - "cudf_polars.experimental.rapidsmpf.frontend.hardware_binding.bind" - ) as mock_bind: - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( + with patch("cudf_polars.engine.hardware_binding.bind") as mock_bind: + from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, bind_to_gpu, ) @@ -276,7 +268,7 @@ def test_bind_enable_once_false() -> None: def test_get_visible_gpu_ids_from_env(monkeypatch: pytest.MonkeyPatch) -> None: - from cudf_polars.experimental.rapidsmpf.frontend.dask import _get_visible_gpu_ids + from cudf_polars.engine.dask import _get_visible_gpu_ids monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1,4") assert _get_visible_gpu_ids() == ["3", "1", "4"] @@ -284,7 +276,7 @@ def test_get_visible_gpu_ids_from_env(monkeypatch: pytest.MonkeyPatch) -> None: def test_dask_setup_assigns_gpus(monkeypatch: pytest.MonkeyPatch) -> None: """dask_setup assigns round-robin CUDA_VISIBLE_DEVICES to each nanny.""" - import cudf_polars.experimental.rapidsmpf.frontend.dask as dask_mod + import cudf_polars.engine.dask as dask_mod monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") monkeypatch.setattr(dask_mod, "_nanny_preload_counter", 0) @@ -302,7 +294,7 @@ def test_dask_setup_assigns_gpus(monkeypatch: pytest.MonkeyPatch) -> None: def test_dask_setup_wraps_around(monkeypatch: pytest.MonkeyPatch) -> None: """Counter wraps around when workers exceed GPUs.""" - import cudf_polars.experimental.rapidsmpf.frontend.dask as dask_mod + import cudf_polars.engine.dask as dask_mod monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "0,1") monkeypatch.setattr(dask_mod, "_nanny_preload_counter", 0) @@ -322,7 +314,7 @@ def test_dask_setup_wraps_around(monkeypatch: pytest.MonkeyPatch) -> None: def test_dask_setup_rejects_worker() -> None: """dask_setup raises TypeError when used with --preload instead of --preload-nanny.""" - import cudf_polars.experimental.rapidsmpf.frontend.dask as dask_mod + import cudf_polars.engine.dask as dask_mod worker = MagicMock(spec=distributed.Worker) with pytest.raises(TypeError, match="--preload-nanny"): diff --git a/python/cudf_polars/tests/experimental/test_dask.py b/python/cudf_polars/tests/streaming/test_dask.py similarity index 97% rename from python/cudf_polars/tests/experimental/test_dask.py rename to python/cudf_polars/tests/streaming/test_dask.py index 93ef4318490..18b3edc248f 100644 --- a/python/cudf_polars/tests/experimental/test_dask.py +++ b/python/cudf_polars/tests/streaming/test_dask.py @@ -12,13 +12,13 @@ import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.utils.config import DaskContext distributed = pytest.importorskip("distributed") -from cudf_polars.experimental.rapidsmpf.frontend.dask import DaskEngine # noqa: E402 +from cudf_polars.engine.dask import DaskEngine # noqa: E402 if TYPE_CHECKING: from collections.abc import Iterator @@ -233,7 +233,7 @@ def test_reset_rejects_construction_time_engine_options( reset_engine: DaskEngine, ) -> None: """``_reset`` rejects ``engine_options`` keys read at worker setup.""" - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( + from cudf_polars.engine.hardware_binding import ( HardwareBindingPolicy, ) diff --git a/python/cudf_polars/tests/experimental/test_dataframescan.py b/python/cudf_polars/tests/streaming/test_dataframescan.py similarity index 94% rename from python/cudf_polars/tests/experimental/test_dataframescan.py rename to python/cudf_polars/tests/streaming/test_dataframescan.py index c7a0cd3a4e7..265cbca7d07 100644 --- a/python/cudf_polars/tests/experimental/test_dataframescan.py +++ b/python/cudf_polars/tests/streaming/test_dataframescan.py @@ -11,9 +11,9 @@ from cudf_polars import Translator from cudf_polars.dsl.traversal import traversal -from cudf_polars.experimental.parallel import lower_ir_graph -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.statistics import collect_statistics +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.parallel import lower_ir_graph +from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.utils.config import ConfigOptions diff --git a/python/cudf_polars/tests/experimental/test_default_singleton_engine.py b/python/cudf_polars/tests/streaming/test_default_singleton_engine.py similarity index 90% rename from python/cudf_polars/tests/experimental/test_default_singleton_engine.py rename to python/cudf_polars/tests/streaming/test_default_singleton_engine.py index d481484d532..b2086ba8b1b 100644 --- a/python/cudf_polars/tests/experimental/test_default_singleton_engine.py +++ b/python/cudf_polars/tests/streaming/test_default_singleton_engine.py @@ -43,10 +43,10 @@ def proc_pool() -> Generator[ProcessPoolExecutor, None, None]: def _reset_singleton_module_state() -> None: """Tear down any leftover engine and reset every module-level slot.""" - from cudf_polars.experimental.rapidsmpf.frontend import ( + from cudf_polars.engine import ( default_singleton_engine as dse, ) - from cudf_polars.experimental.rapidsmpf.frontend.default_singleton_engine import ( + from cudf_polars.engine.default_singleton_engine import ( DefaultSingletonEngine, ) @@ -87,10 +87,10 @@ def _body_lifecycle() -> None: import polars as pl - from cudf_polars.experimental.rapidsmpf.frontend.default_singleton_engine import ( + from cudf_polars.engine.default_singleton_engine import ( DefaultSingletonEngine, ) - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine # Construction succeeds; isinstance check skips its parent's # ``check_no_live_default_singleton`` so SPMDEngine.__init__ runs cleanly. @@ -123,10 +123,10 @@ def _body_default_path_routing() -> None: """ import polars as pl - from cudf_polars.experimental.rapidsmpf.frontend import ( + from cudf_polars.engine import ( default_singleton_engine as dse, ) - from cudf_polars.experimental.rapidsmpf.frontend.default_singleton_engine import ( + from cudf_polars.engine.default_singleton_engine import ( DefaultSingletonEngine, ) @@ -163,7 +163,7 @@ def _body_concurrent_warm_path() -> None: """Concurrent ``get_or_create()`` calls return the same instance.""" import threading - from cudf_polars.experimental.rapidsmpf.frontend.default_singleton_engine import ( + from cudf_polars.engine.default_singleton_engine import ( DefaultSingletonEngine, ) @@ -195,10 +195,10 @@ def _body_atexit_no_op() -> None: ``DefaultSingletonEngine.shutdown`` (registered once at import as the atexit hook) is a no-op when no engine is live. """ - from cudf_polars.experimental.rapidsmpf.frontend import ( + from cudf_polars.engine import ( default_singleton_engine as dse, ) - from cudf_polars.experimental.rapidsmpf.frontend.default_singleton_engine import ( + from cudf_polars.engine.default_singleton_engine import ( DefaultSingletonEngine, ) @@ -221,11 +221,11 @@ def _body_singleton_blocked_when_explicit_alive() -> None: """ import pytest - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine - from cudf_polars.experimental.rapidsmpf.frontend.default_singleton_engine import ( + from cudf_polars.engine.core import StreamingEngine + from cudf_polars.engine.default_singleton_engine import ( DefaultSingletonEngine, ) - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine assert StreamingEngine._active_engine_count() == 0 @@ -260,13 +260,13 @@ def _body_worker_thread_isolation() -> None: import pytest - from cudf_polars.experimental.rapidsmpf.frontend import ( + from cudf_polars.engine import ( default_singleton_engine as dse, ) - from cudf_polars.experimental.rapidsmpf.frontend.default_singleton_engine import ( + from cudf_polars.engine.default_singleton_engine import ( DefaultSingletonEngine, ) - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine # 1) Construction runs on the worker thread. recorded: dict[str, threading.Thread] = {} @@ -326,13 +326,13 @@ def _body_shutdown_timeout() -> None: import pytest - from cudf_polars.experimental.rapidsmpf.frontend import ( + from cudf_polars.engine import ( default_singleton_engine as dse, ) - from cudf_polars.experimental.rapidsmpf.frontend.default_singleton_engine import ( + from cudf_polars.engine.default_singleton_engine import ( DefaultSingletonEngine, ) - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine release_worker = threading.Event() real_done = threading.Event() diff --git a/python/cudf_polars/tests/experimental/test_distinct.py b/python/cudf_polars/tests/streaming/test_distinct.py similarity index 97% rename from python/cudf_polars/tests/experimental/test_distinct.py rename to python/cudf_polars/tests/streaming/test_distinct.py index 83d7233cad5..6ee42842c1b 100644 --- a/python/cudf_polars/tests/experimental/test_distinct.py +++ b/python/cudf_polars/tests/streaming/test_distinct.py @@ -8,7 +8,7 @@ import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions from cudf_polars.testing.asserts import assert_gpu_result_equal diff --git a/python/cudf_polars/tests/experimental/test_empty_channels.py b/python/cudf_polars/tests/streaming/test_empty_channels.py similarity index 96% rename from python/cudf_polars/tests/experimental/test_empty_channels.py rename to python/cudf_polars/tests/streaming/test_empty_channels.py index 530fc8e7ffb..b4b542fe323 100644 --- a/python/cudf_polars/tests/experimental/test_empty_channels.py +++ b/python/cudf_polars/tests/streaming/test_empty_channels.py @@ -9,7 +9,7 @@ import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions from cudf_polars.testing.asserts import assert_gpu_result_equal diff --git a/python/cudf_polars/tests/experimental/test_explain.py b/python/cudf_polars/tests/streaming/test_explain.py similarity index 99% rename from python/cudf_polars/tests/experimental/test_explain.py rename to python/cudf_polars/tests/streaming/test_explain.py index 8259cced75c..df61f366179 100644 --- a/python/cudf_polars/tests/experimental/test_explain.py +++ b/python/cudf_polars/tests/streaming/test_explain.py @@ -13,12 +13,12 @@ import polars as pl -from cudf_polars.experimental.explain import ( +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.explain import ( _fmt_row_count, explain_query, serialize_query, ) -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.testing.io import make_lazy_frame, make_partitioned_source diff --git a/python/cudf_polars/tests/experimental/test_filter.py b/python/cudf_polars/tests/streaming/test_filter.py similarity index 93% rename from python/cudf_polars/tests/experimental/test_filter.py rename to python/cudf_polars/tests/streaming/test_filter.py index b8b4fb2749c..fa72da6f5bf 100644 --- a/python/cudf_polars/tests/experimental/test_filter.py +++ b/python/cudf_polars/tests/streaming/test_filter.py @@ -7,7 +7,7 @@ import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.testing.engine_utils import warns_on_spmd diff --git a/python/cudf_polars/tests/experimental/test_groupby.py b/python/cudf_polars/tests/streaming/test_groupby.py similarity index 97% rename from python/cudf_polars/tests/experimental/test_groupby.py rename to python/cudf_polars/tests/streaming/test_groupby.py index 6ca11387da0..99835e125fd 100644 --- a/python/cudf_polars/tests/experimental/test_groupby.py +++ b/python/cudf_polars/tests/streaming/test_groupby.py @@ -11,8 +11,8 @@ import polars as pl -from cudf_polars.experimental.rapidsmpf.collectives.shuffle import ShuffleManager -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.actor_graph.collectives.shuffle import ShuffleManager from cudf_polars.testing.asserts import assert_gpu_result_equal @@ -92,7 +92,7 @@ def test_dynamic_groupby_single_row(streaming_engine): # --------------------------------------------------------------------------- -# Tests migrated from tests/experimental/test_groupby.py +# Tests migrated from tests/streaming/test_groupby.py # --------------------------------------------------------------------------- @@ -262,7 +262,7 @@ def test_groupby_literal_key(df, streaming_engine): # --------------------------------------------------------------------------- -# Tests migrated from tests/experimental/test_groupby.py +# Tests migrated from tests/streaming/test_groupby.py # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/experimental/test_hstack.py b/python/cudf_polars/tests/streaming/test_hstack.py similarity index 95% rename from python/cudf_polars/tests/experimental/test_hstack.py rename to python/cudf_polars/tests/streaming/test_hstack.py index b6753d31adb..01861ea30e7 100644 --- a/python/cudf_polars/tests/experimental/test_hstack.py +++ b/python/cudf_polars/tests/streaming/test_hstack.py @@ -15,10 +15,10 @@ from cudf_polars.dsl.expressions.base import ExecutionContext from cudf_polars.dsl.ir import Filter, HStack from cudf_polars.dsl.traversal import traversal -from cudf_polars.experimental.parallel import lower_ir_graph -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.repartition import Repartition -from cudf_polars.experimental.statistics import collect_statistics +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.parallel import lower_ir_graph +from cudf_polars.streaming.repartition import Repartition +from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.testing.engine_utils import warns_on_spmd from cudf_polars.utils.config import ConfigOptions, StreamingFallbackMode diff --git a/python/cudf_polars/tests/experimental/test_io_multirank.py b/python/cudf_polars/tests/streaming/test_io_multirank.py similarity index 93% rename from python/cudf_polars/tests/experimental/test_io_multirank.py rename to python/cudf_polars/tests/streaming/test_io_multirank.py index be627f25caf..265984b6836 100644 --- a/python/cudf_polars/tests/experimental/test_io_multirank.py +++ b/python/cudf_polars/tests/streaming/test_io_multirank.py @@ -10,7 +10,7 @@ import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions from cudf_polars.testing.asserts import assert_sink_result_equal from cudf_polars.utils.config import Cluster, StreamingExecutor @@ -18,7 +18,7 @@ from collections.abc import Callable from pathlib import Path - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine + from cudf_polars.engine.core import StreamingEngine # Runs the spmd variant even under rrun with nranks > 1. The ray/dask # variants skip themselves in that environment. diff --git a/python/cudf_polars/tests/experimental/test_join.py b/python/cudf_polars/tests/streaming/test_join.py similarity index 95% rename from python/cudf_polars/tests/experimental/test_join.py rename to python/cudf_polars/tests/streaming/test_join.py index 0945e25f1d4..727de9afad7 100644 --- a/python/cudf_polars/tests/experimental/test_join.py +++ b/python/cudf_polars/tests/streaming/test_join.py @@ -12,12 +12,12 @@ from cudf_polars import Translator from cudf_polars.dsl.ir import Cache, Join from cudf_polars.dsl.traversal import traversal -from cudf_polars.experimental.base import PartitionInfo -from cudf_polars.experimental.parallel import lower_ir_graph -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.rapidsmpf.join import _use_pwise_join -from cudf_polars.experimental.shuffle import Shuffle -from cudf_polars.experimental.statistics import collect_statistics +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.actor_graph.join import _use_pwise_join +from cudf_polars.streaming.base import PartitionInfo +from cudf_polars.streaming.parallel import lower_ir_graph +from cudf_polars.streaming.shuffle import Shuffle +from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.testing.engine_utils import warns_on_spmd from cudf_polars.utils.config import ConfigOptions, StreamingExecutor @@ -72,7 +72,7 @@ def test_dynamic_join_right_full_reverse(left, right, streaming_engine_factory, # --------------------------------------------------------------------------- -# Tests migrated from tests/experimental/test_join.py +# Tests migrated from tests/streaming/test_join.py # --------------------------------------------------------------------------- @@ -113,7 +113,7 @@ def test_join_conditional(reverse, max_rows_per_partition, streaming_engine_fact # --------------------------------------------------------------------------- -# Tests migrated from tests/experimental/test_join.py +# Tests migrated from tests/streaming/test_join.py # --------------------------------------------------------------------------- @@ -268,7 +268,7 @@ def test_join_maintain_order_fallback_streaming( # --------------------------------------------------------------------------- -# Tests migrated from tests/experimental/test_join.py (IR inspection) +# Tests migrated from tests/streaming/test_join.py (IR inspection) # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/experimental/test_metadata.py b/python/cudf_polars/tests/streaming/test_metadata.py similarity index 98% rename from python/cudf_polars/tests/experimental/test_metadata.py rename to python/cudf_polars/tests/streaming/test_metadata.py index 969fca298df..5c555e3593b 100644 --- a/python/cudf_polars/tests/experimental/test_metadata.py +++ b/python/cudf_polars/tests/streaming/test_metadata.py @@ -23,13 +23,13 @@ from cudf_polars.containers import DataFrame, DataType from cudf_polars.dsl import expr from cudf_polars.dsl.ir import GroupBy, HStack, Projection, Select, Sort -from cudf_polars.experimental.rapidsmpf.collectives.sort import ( +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.actor_graph.collectives.sort import ( _is_already_sorted, _sort_to_order_keys, ) -from cudf_polars.experimental.rapidsmpf.core import evaluate_logical_plan -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.core import evaluate_logical_plan +from cudf_polars.streaming.actor_graph.utils import ( NormalizedPartitioning, maybe_remap_partitioning, ) diff --git a/python/cudf_polars/tests/experimental/test_options.py b/python/cudf_polars/tests/streaming/test_options.py similarity index 95% rename from python/cudf_polars/tests/experimental/test_options.py rename to python/cudf_polars/tests/streaming/test_options.py index 5eb1c301376..926fc505a17 100644 --- a/python/cudf_polars/tests/experimental/test_options.py +++ b/python/cudf_polars/tests/streaming/test_options.py @@ -10,10 +10,12 @@ import pytest -from cudf_polars.experimental.rapidsmpf.frontend.options import ( +from cudf_polars.engine.hardware_binding import HardwareBindingPolicy +from cudf_polars.engine.options import ( UNSPECIFIED, StreamingOptions, Unspecified, + _parse_memory_resource_config, ) from cudf_polars.utils.config import MemoryResourceConfig @@ -169,10 +171,6 @@ def test_unbounded_file_read_cache_env_var(monkeypatch: pytest.MonkeyPatch) -> N def test_parse_memory_resource_config() -> None: """_parse_memory_resource_config converts a JSON string to MemoryResourceConfig.""" - from cudf_polars.experimental.rapidsmpf.frontend.options import ( - _parse_memory_resource_config, - ) - config = _parse_memory_resource_config('{"qualname": "rmm.mr.CudaMemoryResource"}') assert isinstance(config, MemoryResourceConfig) assert config.qualname == "rmm.mr.CudaMemoryResource" @@ -368,10 +366,6 @@ def test_to_dict_roundtrip_empty() -> None: def test_hardware_binding_in_engine_options() -> None: - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( - HardwareBindingPolicy, - ) - result = StreamingOptions( hardware_binding=HardwareBindingPolicy(enabled=False) ).to_engine_options() @@ -379,20 +373,12 @@ def test_hardware_binding_in_engine_options() -> None: def test_hardware_binding_env_var_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( - HardwareBindingPolicy, - ) - monkeypatch.setenv("CUDF_POLARS__HARDWARE_BINDING", '{"enabled": false}') result = StreamingOptions().to_engine_options() assert result["hardware_binding"] == HardwareBindingPolicy(enabled=False) def test_hardware_binding_cli_json() -> None: - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( - HardwareBindingPolicy, - ) - parser = argparse.ArgumentParser() StreamingOptions._add_cli_args(parser) args = parser.parse_args(["--hardware-binding", '{"raise_on_fail": true}']) @@ -407,10 +393,6 @@ def test_hardware_binding_invalid_json(monkeypatch: pytest.MonkeyPatch) -> None: def test_hardware_binding_cli_disabled() -> None: - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( - HardwareBindingPolicy, - ) - parser = argparse.ArgumentParser() StreamingOptions._add_cli_args(parser) args = parser.parse_args(["--hardware-binding", '{"enabled": false}']) diff --git a/python/cudf_polars/tests/experimental/test_parallel.py b/python/cudf_polars/tests/streaming/test_parallel.py similarity index 95% rename from python/cudf_polars/tests/experimental/test_parallel.py rename to python/cudf_polars/tests/streaming/test_parallel.py index e057b561016..a7b95dd3c5d 100644 --- a/python/cudf_polars/tests/experimental/test_parallel.py +++ b/python/cudf_polars/tests/streaming/test_parallel.py @@ -13,7 +13,7 @@ from cudf_polars import Translator from cudf_polars.dsl.traversal import traversal -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions from cudf_polars.testing.asserts import assert_gpu_result_equal @@ -69,7 +69,7 @@ def test_fallback_on_concat_zlice(spmd_engine_factory) -> None: # --------------------------------------------------------------------------- -# Tests migrated from tests/experimental/test_parallel.py +# Tests migrated from tests/streaming/test_parallel.py # --------------------------------------------------------------------------- @@ -89,7 +89,7 @@ def test_evaluate_streaming(streaming_engine): # --------------------------------------------------------------------------- -# Tests migrated from tests/experimental/test_parallel.py (round 3) +# Tests migrated from tests/streaming/test_parallel.py (round 3) # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/experimental/test_ray.py b/python/cudf_polars/tests/streaming/test_ray.py similarity index 98% rename from python/cudf_polars/tests/experimental/test_ray.py rename to python/cudf_polars/tests/streaming/test_ray.py index f62c3e3b831..8427099fb02 100644 --- a/python/cudf_polars/tests/experimental/test_ray.py +++ b/python/cudf_polars/tests/streaming/test_ray.py @@ -13,10 +13,11 @@ import polars as pl +from cudf_polars.engine.hardware_binding import HardwareBindingPolicy from cudf_polars.utils.config import RayContext ray = pytest.importorskip("ray") -from cudf_polars.experimental.rapidsmpf.frontend.ray import RayEngine # noqa: E402 +from cudf_polars.engine.ray import RayEngine # noqa: E402 if TYPE_CHECKING: from collections.abc import Iterator @@ -304,10 +305,6 @@ def test_reset_rejects_construction_time_engine_options( reset_engine: RayEngine, ) -> None: """``_reset`` rejects ``engine_options`` keys read at actor construction.""" - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( - HardwareBindingPolicy, - ) - with pytest.raises(ValueError, match="hardware_binding"): reset_engine._reset( engine_options={ diff --git a/python/cudf_polars/tests/experimental/test_rolling.py b/python/cudf_polars/tests/streaming/test_rolling.py similarity index 95% rename from python/cudf_polars/tests/experimental/test_rolling.py rename to python/cudf_polars/tests/streaming/test_rolling.py index 6c7a71ce0b9..743fe7aebcc 100644 --- a/python/cudf_polars/tests/experimental/test_rolling.py +++ b/python/cudf_polars/tests/streaming/test_rolling.py @@ -7,11 +7,11 @@ import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.spmd import SPMDEngine from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.testing.engine_utils import warns_on_spmd -from cudf_polars.utils.versions import POLARS_VERSION_LT_136 +from cudf_polars.utils.versions import POLARS_VERSION_LT_136, POLARS_VERSION_LT_139 @pytest.fixture @@ -21,11 +21,11 @@ def engine(streaming_engine_factory): ) -def test_rolling_datetime(request, engine): - if not POLARS_VERSION_LT_136: - request.applymarker( - pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") - ) +@pytest.mark.skipif( + not POLARS_VERSION_LT_136 and POLARS_VERSION_LT_139, + reason="Rolling window expressions are not accessible in polars 1.36-1.38", +) +def test_rolling_datetime(engine): dates = [ "2020-01-01 13:45:48", "2020-01-01 16:42:13", diff --git a/python/cudf_polars/tests/experimental/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py similarity index 93% rename from python/cudf_polars/tests/experimental/test_scan.py rename to python/cudf_polars/tests/streaming/test_scan.py index 6776eee84ea..063ae53ad25 100644 --- a/python/cudf_polars/tests/experimental/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -8,9 +8,9 @@ import polars as pl from cudf_polars import Translator -from cudf_polars.experimental.parallel import lower_ir_graph -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.statistics import collect_statistics +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.parallel import lower_ir_graph +from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.testing.io import make_partitioned_source from cudf_polars.utils.config import ConfigOptions @@ -53,7 +53,7 @@ def test_scan_parquet_use_rapidsmpf_native(tmp_path, df, streaming_engine_factor # --------------------------------------------------------------------------- -# Tests migrated from tests/experimental/test_scan.py +# Tests migrated from tests/streaming/test_scan.py # --------------------------------------------------------------------------- diff --git a/python/cudf_polars/tests/experimental/test_select.py b/python/cudf_polars/tests/streaming/test_select.py similarity index 98% rename from python/cudf_polars/tests/experimental/test_select.py rename to python/cudf_polars/tests/streaming/test_select.py index 188e9e392d0..7eb5c1a93eb 100644 --- a/python/cudf_polars/tests/experimental/test_select.py +++ b/python/cudf_polars/tests/streaming/test_select.py @@ -16,8 +16,8 @@ from cudf_polars.containers import DataType from cudf_polars.dsl import expr from cudf_polars.dsl.ir import HStack, Projection, Select -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.select import _inline_hstack_false, _sub_expr +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.select import _inline_hstack_false, _sub_expr from cudf_polars.testing.asserts import ( assert_gpu_result_equal, ) diff --git a/python/cudf_polars/tests/experimental/test_shuffler.py b/python/cudf_polars/tests/streaming/test_shuffler.py similarity index 95% rename from python/cudf_polars/tests/experimental/test_shuffler.py rename to python/cudf_polars/tests/streaming/test_shuffler.py index 84818741815..7f4fd46d15e 100644 --- a/python/cudf_polars/tests/experimental/test_shuffler.py +++ b/python/cudf_polars/tests/streaming/test_shuffler.py @@ -16,14 +16,14 @@ import polars as pl from cudf_polars.containers import DataFrame, DataType -from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id -from cudf_polars.experimental.rapidsmpf.collectives.shuffle import ( +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.spmd import allgather_polars_dataframe +from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id +from cudf_polars.streaming.actor_graph.collectives.shuffle import ( LocalRepartitioner, ShuffleManager, ) -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.rapidsmpf.frontend.spmd import allgather_polars_dataframe -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.streaming.actor_graph.utils import ( _is_already_partitioned, ) from cudf_polars.testing.asserts import assert_gpu_result_equal diff --git a/python/cudf_polars/tests/experimental/test_sink.py b/python/cudf_polars/tests/streaming/test_sink.py similarity index 98% rename from python/cudf_polars/tests/experimental/test_sink.py rename to python/cudf_polars/tests/streaming/test_sink.py index 5c969728c7a..cf160d432e5 100644 --- a/python/cudf_polars/tests/experimental/test_sink.py +++ b/python/cudf_polars/tests/streaming/test_sink.py @@ -9,7 +9,7 @@ import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions from cudf_polars.testing.asserts import assert_sink_result_equal diff --git a/python/cudf_polars/tests/experimental/test_sort.py b/python/cudf_polars/tests/streaming/test_sort.py similarity index 97% rename from python/cudf_polars/tests/experimental/test_sort.py rename to python/cudf_polars/tests/streaming/test_sort.py index b147358dc4b..533b5933a02 100644 --- a/python/cudf_polars/tests/experimental/test_sort.py +++ b/python/cudf_polars/tests/streaming/test_sort.py @@ -7,7 +7,7 @@ import polars as pl -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions from cudf_polars.testing.asserts import assert_gpu_result_equal diff --git a/python/cudf_polars/tests/experimental/test_spilling.py b/python/cudf_polars/tests/streaming/test_spilling.py similarity index 97% rename from python/cudf_polars/tests/experimental/test_spilling.py rename to python/cudf_polars/tests/streaming/test_spilling.py index c543d0d3a61..e368bc3ac1f 100644 --- a/python/cudf_polars/tests/experimental/test_spilling.py +++ b/python/cudf_polars/tests/streaming/test_spilling.py @@ -19,8 +19,8 @@ import pylibcudf as plc -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.rapidsmpf.utils import ( +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.actor_graph.utils import ( make_spill_function, ) diff --git a/python/cudf_polars/tests/experimental/test_spmd.py b/python/cudf_polars/tests/streaming/test_spmd.py similarity index 98% rename from python/cudf_polars/tests/experimental/test_spmd.py rename to python/cudf_polars/tests/streaming/test_spmd.py index a61268cf387..75ca4888d08 100644 --- a/python/cudf_polars/tests/experimental/test_spmd.py +++ b/python/cudf_polars/tests/streaming/test_spmd.py @@ -16,12 +16,13 @@ import rmm.mr -from cudf_polars.experimental.rapidsmpf.collectives.common import reserve_op_id -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.rapidsmpf.frontend.spmd import ( +from cudf_polars.engine.hardware_binding import HardwareBindingPolicy +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.engine.spmd import ( SPMDEngine, allgather_polars_dataframe, ) +from cudf_polars.streaming.actor_graph.collectives.common import reserve_op_id from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.utils.config import MemoryResourceConfig @@ -380,10 +381,6 @@ def test_reset_rejects_construction_time_engine_options( comm: Communicator, ) -> None: """``_reset`` rejects ``engine_options`` keys read at engine construction.""" - from cudf_polars.experimental.rapidsmpf.frontend.hardware_binding import ( - HardwareBindingPolicy, - ) - with SPMDEngine(comm=comm) as engine: with pytest.raises(ValueError, match="hardware_binding"): engine._reset( diff --git a/python/cudf_polars/tests/experimental/test_statistics.py b/python/cudf_polars/tests/streaming/test_statistics.py similarity index 91% rename from python/cudf_polars/tests/experimental/test_statistics.py rename to python/cudf_polars/tests/streaming/test_statistics.py index 42014a02106..61d5659dc70 100644 --- a/python/cudf_polars/tests/experimental/test_statistics.py +++ b/python/cudf_polars/tests/streaming/test_statistics.py @@ -9,12 +9,12 @@ import pytest from rapidsmpf.statistics import Statistics -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions if TYPE_CHECKING: from collections.abc import Callable - from cudf_polars.experimental.rapidsmpf.frontend.core import StreamingEngine + from cudf_polars.engine.core import StreamingEngine # Runs the spmd variant even under rrun with nranks > 1. The ray/dask # variants skip themselves in that environment. diff --git a/python/cudf_polars/tests/experimental/test_stats.py b/python/cudf_polars/tests/streaming/test_stats.py similarity index 97% rename from python/cudf_polars/tests/experimental/test_stats.py rename to python/cudf_polars/tests/streaming/test_stats.py index 7d0d5dc01e4..252f7462d04 100644 --- a/python/cudf_polars/tests/experimental/test_stats.py +++ b/python/cudf_polars/tests/streaming/test_stats.py @@ -14,14 +14,14 @@ from cudf_polars import Translator from cudf_polars.containers import DataType from cudf_polars.dsl.ir import Empty, Projection -from cudf_polars.experimental.base import SerializedDataSourceInfo, StatsCollector -from cudf_polars.experimental.io import ( +from cudf_polars.engine.options import StreamingOptions +from cudf_polars.streaming.base import SerializedDataSourceInfo, StatsCollector +from cudf_polars.streaming.io import ( DataFrameSourceInfo, ParquetSourceInfo, _clear_source_info_cache, ) -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions -from cudf_polars.experimental.statistics import collect_statistics +from cudf_polars.streaming.statistics import collect_statistics from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.testing.io import make_lazy_frame, make_partitioned_source from cudf_polars.utils.config import ConfigOptions diff --git a/python/cudf_polars/tests/experimental/test_tracing.py b/python/cudf_polars/tests/streaming/test_tracing.py similarity index 93% rename from python/cudf_polars/tests/experimental/test_tracing.py rename to python/cudf_polars/tests/streaming/test_tracing.py index cf1d8a5cbf2..9a1e71a00e2 100644 --- a/python/cudf_polars/tests/experimental/test_tracing.py +++ b/python/cudf_polars/tests/streaming/test_tracing.py @@ -20,7 +20,7 @@ def test_structlog_streaming_node_events(): import polars as pl rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine df = pl.DataFrame({"x": range(100), "y": ["a", "b"] * 50}) q = df.lazy().filter(pl.col("x") > 50).group_by("y").agg(pl.col("x").sum()) @@ -54,7 +54,7 @@ def test_structlog_contains_expected_ir_types(): import polars as pl rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine df = pl.DataFrame({"x": range(100), "y": ["a", "b"] * 50}) q = df.lazy().filter(pl.col("x") > 50).group_by("y").agg(pl.col("x").sum()) @@ -86,7 +86,7 @@ def test_structlog_disabled_by_default(): import polars as pl rmm.mr.set_current_device_resource(rmm.mr.ManagedMemoryResource()) - from cudf_polars.experimental.rapidsmpf.frontend.spmd import SPMDEngine + from cudf_polars.engine.spmd import SPMDEngine df = pl.DataFrame({"x": range(10), "y": ["a", "b"] * 5}) q = df.lazy().filter(pl.col("x") > 5) diff --git a/python/cudf_polars/tests/experimental/test_union.py b/python/cudf_polars/tests/streaming/test_union.py similarity index 100% rename from python/cudf_polars/tests/experimental/test_union.py rename to python/cudf_polars/tests/streaming/test_union.py diff --git a/python/cudf_polars/tests/experimental/test_unique.py b/python/cudf_polars/tests/streaming/test_unique.py similarity index 97% rename from python/cudf_polars/tests/experimental/test_unique.py rename to python/cudf_polars/tests/streaming/test_unique.py index 1a157c3fe21..27b5454030d 100644 --- a/python/cudf_polars/tests/experimental/test_unique.py +++ b/python/cudf_polars/tests/streaming/test_unique.py @@ -8,7 +8,7 @@ import polars as pl from polars.testing import assert_frame_equal -from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions +from cudf_polars.engine.options import StreamingOptions from cudf_polars.testing.asserts import assert_gpu_result_equal from cudf_polars.testing.engine_utils import warns_on_spmd diff --git a/python/cudf_polars/tests/experimental/test_utils.py b/python/cudf_polars/tests/streaming/test_utils.py similarity index 81% rename from python/cudf_polars/tests/experimental/test_utils.py rename to python/cudf_polars/tests/streaming/test_utils.py index 4445f68ff10..28afb8eebfc 100644 --- a/python/cudf_polars/tests/experimental/test_utils.py +++ b/python/cudf_polars/tests/streaming/test_utils.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -8,7 +8,7 @@ from cudf_polars.containers import DataType from cudf_polars.dsl import expr -from cudf_polars.experimental.utils import _leaf_column_names +from cudf_polars.streaming.utils import _leaf_column_names def test_leaf_column_names(): diff --git a/python/cudf_polars/tests/test_hconcat.py b/python/cudf_polars/tests/test_hconcat.py index 761b673a682..25c277dc268 100644 --- a/python/cudf_polars/tests/test_hconcat.py +++ b/python/cudf_polars/tests/test_hconcat.py @@ -2,11 +2,14 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations +import pytest + import polars as pl from cudf_polars.containers import DataType from cudf_polars.dsl.ir import DataFrameScan, Empty, HConcat, IRExecutionContext from cudf_polars.testing.asserts import assert_gpu_result_equal +from cudf_polars.utils.versions import POLARS_VERSION_LT_139 def test_hconcat(engine: pl.GPUEngine): @@ -30,6 +33,17 @@ def test_hconcat_different_heights(engine: pl.GPUEngine): assert_gpu_result_equal(q, engine=engine) +@pytest.mark.skipif(POLARS_VERSION_LT_139, reason="strict keyword added in polars 1.39") +def test_hconcat_strict_different_heights(): + left = pl.LazyFrame({"a": [1, 2, 3]}) + right = pl.LazyFrame({"b": [4, 5]}) + q = pl.concat([left, right], how="horizontal", strict=True) + with pytest.raises(pl.exceptions.ShapeError): + q.collect() + with pytest.raises(pl.exceptions.ShapeError): + q.collect(engine=pl.GPUEngine(executor="in-memory", raise_on_fail=True)) + + def test_hconcat_should_broadcast(): # HConcat with should_broadcast=True is used by the streaming engine to # recombine decomposed expressions. Test it by constructing IR directly. @@ -42,7 +56,7 @@ def test_hconcat_should_broadcast(): child2 = DataFrameScan({"b": DataType(pl.Float64())}, df2._df, None) schema = {"a": DataType(pl.Int64()), "b": DataType(pl.Float64())} - node = HConcat(schema, True, child1, child2) # noqa: FBT003 + node = HConcat(schema, True, False, child1, child2) # noqa: FBT003 result = node.evaluate(cache={}, timer=None, context=context) polars_result = result.to_polars() diff --git a/python/cudf_polars/tests/test_parquet_filters.py b/python/cudf_polars/tests/test_parquet_filters.py index 2009800ced5..d77a627fdb7 100644 --- a/python/cudf_polars/tests/test_parquet_filters.py +++ b/python/cudf_polars/tests/test_parquet_filters.py @@ -65,6 +65,17 @@ def test_scan_by_hand(expr, selection, pq_file, chunked): ) +def test_parquet_filter_ne_missing(tmp_path): + df = pl.DataFrame({"a": [1, 2, None, 3, None]}) + df.write_parquet(tmp_path / "df.parquet", row_group_size=3) + q = pl.scan_parquet(tmp_path / "df.parquet").filter( + pl.col("a").ne_missing(pl.lit(2, dtype=pl.Int64)) + ) + assert_gpu_result_equal( + q, engine=pl.GPUEngine(executor="in-memory", raise_on_fail=True) + ) + + def test_parquet_filter_boolean_column(engine: pl.GPUEngine, tmp_path): df = pl.DataFrame({"x": [1, 2, 3], "y": [True, False, True]}) df.write_parquet(tmp_path / "df.parquet") diff --git a/python/cudf_polars/tests/test_repr.py b/python/cudf_polars/tests/test_repr.py index df234a24af6..5d50e49cf53 100644 --- a/python/cudf_polars/tests/test_repr.py +++ b/python/cudf_polars/tests/test_repr.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations @@ -10,7 +10,7 @@ import polars as pl from cudf_polars.dsl.translate import Translator -from cudf_polars.experimental.base import PartitionInfo +from cudf_polars.streaming.base import PartitionInfo def test_partition_info_rich_repr() -> None: diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index 4a23f9aaec0..dffbfefd986 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.py @@ -21,6 +21,7 @@ from cudf_polars.testing.io import make_partitioned_source from cudf_polars.utils.versions import ( POLARS_VERSION_LT_138, + POLARS_VERSION_LT_139, ) if TYPE_CHECKING: @@ -587,6 +588,10 @@ def get_handler(req: Request) -> Response: ) +@pytest.mark.xfail( + condition=not POLARS_VERSION_LT_139, + reason="polars 1.39+ ndjson remote reader requires range request support", +) def test_scan_ndjson_remote( engine: pl.GPUEngine, tmp_path: Path, @@ -656,7 +661,7 @@ def test_hits_scan_row_index_duplicate(engine: pl.GPUEngine, request, tmp_path): request.applymarker( pytest.mark.xfail( condition=not POLARS_VERSION_LT_138, - reason="polars fails ahead of time", + reason="polars >= 1.38 raises duplicate row_index name ahead of time", ) ) pl.DataFrame({"col": [1, 2, 3]}).write_parquet(tmp_path / "a.parquet") @@ -705,8 +710,7 @@ def test_scan_tiny_file_not_compressed(engine: pl.GPUEngine, tmp_path): @pytest.mark.skipif( - POLARS_VERSION_LT_138, - reason="height parameter added in Polars 1.38", + POLARS_VERSION_LT_138, reason="pl.LazyFrame(height=...) requires polars >= 1.38" ) @pytest.mark.parametrize("custom_engine", [None, NO_CHUNK_ENGINE]) def test_scan_parquet_zero_width_with_limit( diff --git a/python/cudf_polars/tests/test_window_functions.py b/python/cudf_polars/tests/test_window_functions.py index 2dbb06fc271..743aa3152e6 100644 --- a/python/cudf_polars/tests/test_window_functions.py +++ b/python/cudf_polars/tests/test_window_functions.py @@ -13,7 +13,7 @@ assert_gpu_result_equal, assert_ir_translation_raises, ) -from cudf_polars.utils.versions import POLARS_VERSION_LT_136 +from cudf_polars.utils.versions import POLARS_VERSION_LT_136, POLARS_VERSION_LT_139 @pytest.fixture @@ -126,15 +126,15 @@ def test_over_mapping_strategy( assert_ir_translation_raises(q, NotImplementedError) +@pytest.mark.skipif( + not POLARS_VERSION_LT_136 and POLARS_VERSION_LT_139, + reason="Rolling window expressions are not accessible in polars 1.36-1.38", +) @pytest.mark.parametrize("period", ["2d", "3d"]) def test_rolling( engine: pl.GPUEngine, request, df: pl.LazyFrame, agg_expr, period: str ): """Test rolling window functions over time series.""" - if not POLARS_VERSION_LT_136: - request.applymarker( - pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") - ) window_expr = agg_expr.rolling(period=period, index_column="date") result_name = f"{agg_expr!s}_rolling_{period}" window_expr = window_expr.alias(result_name) @@ -147,7 +147,7 @@ def test_rolling( def test_rolling_unsupported(df: pl.LazyFrame, unsupported_agg_expr): """Test rolling window functions over time series.""" window_expr = unsupported_agg_expr.rolling(period="2d", index_column="date") - result_name = f"{agg_expr!s}_rolling" + result_name = f"{unsupported_agg_expr!s}_rolling" window_expr = window_expr.alias(result_name) query = df.with_columns(window_expr) @@ -155,13 +155,13 @@ def test_rolling_unsupported(df: pl.LazyFrame, unsupported_agg_expr): assert_ir_translation_raises(query, NotImplementedError) +@pytest.mark.skipif( + not POLARS_VERSION_LT_136 and POLARS_VERSION_LT_139, + reason="Rolling window expressions are not accessible in polars 1.36-1.38", +) @pytest.mark.parametrize("closed", ["left", "right", "both", "none"]) def test_rolling_closed(engine: pl.GPUEngine, request, df: pl.LazyFrame, closed: str): """Test rolling window functions with different closed parameters.""" - if not POLARS_VERSION_LT_136: - request.applymarker( - pytest.mark.xfail(reason="See https://github.com/pola-rs/polars/pull/25117") - ) # ignore is for polars' ClosedInterval, which isn't publicly exported. # https://github.com/pola-rs/polars/issues/17420 query = df.with_columns( diff --git a/python/cudf_polars/tests/testing/test_asserts.py b/python/cudf_polars/tests/testing/test_asserts.py index dded6078d96..9ae4cd64cc6 100644 --- a/python/cudf_polars/tests/testing/test_asserts.py +++ b/python/cudf_polars/tests/testing/test_asserts.py @@ -11,7 +11,7 @@ import polars as pl import polars.testing -from cudf_polars.experimental.benchmarks.asserts import ( +from cudf_polars.streaming.benchmarks.asserts import ( ValidationError, assert_tpch_result_equal, ) diff --git a/python/cudf_polars/tests/testing/test_engine_utils.py b/python/cudf_polars/tests/testing/test_engine_utils.py index ff8ceeb215a..ce1440abba7 100644 --- a/python/cudf_polars/tests/testing/test_engine_utils.py +++ b/python/cudf_polars/tests/testing/test_engine_utils.py @@ -43,7 +43,7 @@ def test_create_streaming_options_small(): def test_create_streaming_options_overrides_merge(): """Overrides take precedence over the blocksize baseline.""" - from cudf_polars.experimental.rapidsmpf.frontend.options import StreamingOptions + from cudf_polars.engine.options import StreamingOptions overrides = StreamingOptions(max_rows_per_partition=999) merged = merge_streaming_options(create_streaming_options("medium"), overrides)