diff --git a/python/cudf_polars/cudf_polars/dsl/ir.py b/python/cudf_polars/cudf_polars/dsl/ir.py index 8dc366275852..78fc9e6594bc 100644 --- a/python/cudf_polars/cudf_polars/dsl/ir.py +++ b/python/cudf_polars/cudf_polars/dsl/ir.py @@ -51,6 +51,7 @@ from cudf_polars.dsl.to_ast import _DECIMAL_IDS, to_ast, to_parquet_filter from cudf_polars.dsl.tracing import log_do_evaluate, nvtx_annotate_cudf_polars from cudf_polars.dsl.utils.naming import unique_names +from cudf_polars.dsl.utils.per_path import PerPathValues from cudf_polars.dsl.utils.reshape import broadcast from cudf_polars.dsl.utils.windows import ( offsets_to_windows, @@ -665,6 +666,7 @@ class Scan(IR): __slots__ = ( "cached_parquet_info", "cloud_options", + "hive_parts", "include_file_paths", "n_rows", "parquet_options", @@ -689,8 +691,9 @@ class Scan(IR): "include_file_paths", "predicate", "parquet_options", + "hive_parts", ) - _n_non_child_args = 12 + _n_non_child_args = 13 typ: str """What type of file are we reading? Parquet, CSV, etc...""" reader_options: dict[str, Any] @@ -713,6 +716,8 @@ class Scan(IR): """Mask to apply to the read dataframe.""" parquet_options: ParquetOptions """Parquet-specific options.""" + hive_parts: PerPathValues | None + """Hive partition values, one per path.""" cached_parquet_info: list[CachedParquetInfo] | None """Cached parquet file metadata.""" @@ -733,6 +738,7 @@ def __init__( include_file_paths: str | None, predicate: expr.NamedExpr | None, parquet_options: ParquetOptions, + hive_parts: PerPathValues | None = None, cached_parquet_info: list[CachedParquetInfo] | None = None, ): self.schema = schema @@ -758,10 +764,12 @@ def __init__( include_file_paths, predicate, parquet_options, + hive_parts, cached_parquet_info, ) self.children = () self.parquet_options = parquet_options + self.hive_parts = hive_parts self.cached_parquet_info = cached_parquet_info Scan._validate_cached_parquet_info(self.paths, self.cached_parquet_info) @@ -882,36 +890,50 @@ def get_hashable(self) -> Hashable: self.include_file_paths, self.predicate, self.parquet_options, + self.hive_parts, ) @staticmethod def add_file_paths( - name: str, paths: list[str], rows_per_path: list[int], df: DataFrame + name: str, + paths: list[str], + df: DataFrame, + *, + rows_per_path: Sequence[int] | None = None, + source_index: plc.Column | None = None, ) -> DataFrame: """ Add a Column of file paths to the DataFrame. - Each path is repeated according to the number of rows read from it. + Parameters + ---------- + name + Name of the column to add. + paths + The paths read, in source order. + df + Frame to add the column to. + rows_per_path + Number of rows read from each path. + source_index + Column giving the source each output row came from. Takes + precedence over ``rows_per_path`` when both are available. + + Returns + ------- + ``df`` with the file path column appended. """ - (filepaths,) = plc.filling.repeat( - plc.Table( - [ - plc.Column.from_arrow( - pl.Series(values=map(str, paths)), - stream=df.stream, - ) - ] - ), - plc.Column.from_arrow( - pl.Series(values=rows_per_path, dtype=pl.datatypes.Int32()), - stream=df.stream, - ), - stream=df.stream, - ).columns() - dtype = DataType(pl.String()) - return df.with_columns( - [Column(filepaths, name=name, dtype=dtype)], stream=df.stream + per_path = PerPathValues( + pl.DataFrame( + {name: [str(path) for path in paths]}, schema={name: pl.String()} + ) ) + if source_index is not None: + columns = per_path.gather(source_index, stream=df.stream) + else: + assert rows_per_path is not None + columns = per_path.repeat(rows_per_path, stream=df.stream) + return df.with_columns(columns, stream=df.stream) @staticmethod @nvtx_annotate_cudf_polars(message="Scan._get_parquet_row_count_from_metadata") @@ -943,6 +965,84 @@ def _get_parquet_row_count_from_metadata( num_rows = min(num_rows, n_rows) return max(num_rows, 0) + @staticmethod + @nvtx_annotate_cudf_polars(message="Scan._parquet_rows_per_path") + def _parquet_rows_per_path( + paths: list[str], + skip_rows: int, + n_rows: int, + cached_parquet_info: list[CachedParquetInfo] | None, + ) -> list[int]: + """ + Rows each path contributes, from file metadata. + + Used when no filter is pushed down. + + Parameters + ---------- + paths + The paths to read, in source order. + skip_rows + Number of leading rows to skip, counted across the paths as a + whole rather than per path. + n_rows + Maximum number of rows to read once ``skip_rows`` have been + skipped, or ``-1`` for no limit. + cached_parquet_info + Prefetched file metadata. + + Returns + ------- + Rows contributed by each path, in source order. + """ + if cached_parquet_info is not None: + Scan._validate_cached_parquet_info(paths, cached_parquet_info) + totals = [info.file_metadata.num_rows for info in cached_parquet_info] + else: + totals = [ + metadata.num_rows + for metadata in plc.io.parquet_metadata.read_parquet_footers( + plc.io.SourceInfo(paths) + ) + ] + available = max(sum(totals) - skip_rows, 0) + budget = available if n_rows == -1 else min(n_rows, available) + counts = [0] * len(totals) + for i, total in enumerate(totals): + if budget == 0: + break + skipped = min(skip_rows, total) + skip_rows -= skipped + counts[i] = min(total - skipped, budget) + budget -= counts[i] + return counts + + @staticmethod + def _pop_source_index( + table: plc.Table, names: Sequence[str], *, prepended: bool + ) -> tuple[plc.Column | None, plc.Table, list[str]]: + """ + Split off the source index column the parquet reader prepends. + + Parameters + ---------- + table + Table as returned by the reader. + names + Column names of ``table``. + prepended + Whether the reader was asked to prepend the source index. + + Returns + ------- + The source index column, or ``None`` if it was not requested, along + with the remaining table and its column names. + """ + if not prepended: + return None, table, list(names) + columns = table.columns() + return columns[0], plc.Table(columns[1:]), list(names[1:]) + @staticmethod def _apply_parquet_projection( table: plc.Table, names: Sequence[str], with_columns: Sequence[str] | None @@ -972,6 +1072,7 @@ def do_evaluate( include_file_paths: str | None, predicate: expr.NamedExpr | None, parquet_options: ParquetOptions, + hive_parts: PerPathValues | None, cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, @@ -1085,8 +1186,8 @@ def read_csv_header( df = Scan.add_file_paths( include_file_paths, seen_paths, - [t.num_rows() for t in tables], df, + rows_per_path=[t.num_rows() for t in tables], ) elif typ == "parquet": if cached_parquet_info is not None: @@ -1103,14 +1204,29 @@ def read_csv_header( parquet_metadatas = None source_info = plc.io.SourceInfo(paths) + hive_names = ( + frozenset(hive_parts.names) if hive_parts is not None else frozenset() + ) + file_columns = ( + with_columns + if with_columns is None or not hive_names + else [name for name in with_columns if name not in hive_names] + ) + rows_per_path: list[int] | None = None + if hive_parts is not None and file_columns == []: + rows_per_path = cls._parquet_rows_per_path( + paths, skip_rows, n_rows, cached_parquet_info + ) + filters = None if predicate is not None and row_index is None: # Can't apply filters during read if we have a row index. filters, residual_expr = to_parquet_filter( _prepare_parquet_predicate( - predicate.value, paths, schema, with_columns + predicate.value, paths, schema, file_columns ), stream=stream, + unreadable_columns=hive_names, ) if filters is not None: effective_predicate = ( @@ -1118,6 +1234,14 @@ def read_csv_header( if residual_expr is not None else None ) + # The reader drops its per-source row counts once it applies a + # filter, so the source of each row has to be read instead. Hive + # columns need it too. + prepend_source_index = ( + hive_parts is not None + and not hive_parts.is_uniform + and file_columns != [] + ) or (include_file_paths is not None and filters is not None) builder = plc.io.parquet.ParquetReaderOptions.builder(source_info) if filters is not None and parquet_options.use_jit_filter: builder.use_jit_filter(use_jit_filter=True) @@ -1125,8 +1249,10 @@ def read_csv_header( plc.TypeId.DECIMAL128 ).build() - if with_columns is not None: - parquet_reader_options.set_column_names(with_columns) + if file_columns is not None: + parquet_reader_options.set_column_names(file_columns) + if prepend_source_index: + parquet_reader_options.enable_prepend_source_index_column(val=True) if filters is not None: parquet_reader_options.set_filter(filters) if n_rows != -1: @@ -1153,13 +1279,18 @@ def read_csv_header( concatenated_columns[i] = plc.concatenate.concatenate( [concatenated_columns[i], columns.pop()], stream=stream ) - table, names = cls._apply_parquet_projection( - plc.Table(concatenated_columns), names, with_columns + source_index, table, names = cls._pop_source_index( + plc.Table(concatenated_columns), + names, + prepended=prepend_source_index, ) + table, names = cls._apply_parquet_projection(table, names, file_columns) if not names: table = plc.Table( table.columns(), - num_rows=cls._get_parquet_row_count_from_metadata( + num_rows=sum(rows_per_path) + if rows_per_path is not None + else cls._get_parquet_row_count_from_metadata( paths, skip_rows, n_rows, @@ -1175,7 +1306,11 @@ def read_csv_header( ) if include_file_paths is not None: df = Scan.add_file_paths( # pragma: no cover - include_file_paths, paths, chunk.num_rows_per_source, df + include_file_paths, + paths, + df, + rows_per_path=rows_per_path or chunk.num_rows_per_source, + source_index=source_index, ) else: tbl_w_meta = plc.io.parquet.read_parquet( @@ -1185,13 +1320,18 @@ def read_csv_header( ) # TODO: consider nested column names? col_names = tbl_w_meta.column_names(include_children=False) + source_index, table, col_names = cls._pop_source_index( + tbl_w_meta.tbl, col_names, prepended=prepend_source_index + ) table, col_names = cls._apply_parquet_projection( - tbl_w_meta.tbl, col_names, with_columns + table, col_names, file_columns ) if not col_names: table = plc.Table( table.columns(), - num_rows=cls._get_parquet_row_count_from_metadata( + num_rows=sum(rows_per_path) + if rows_per_path is not None + else cls._get_parquet_row_count_from_metadata( paths, skip_rows, n_rows, @@ -1207,8 +1347,27 @@ def read_csv_header( ) if include_file_paths is not None: df = Scan.add_file_paths( - include_file_paths, paths, tbl_w_meta.num_rows_per_source, df + include_file_paths, + paths, + df, + rows_per_path=rows_per_path or tbl_w_meta.num_rows_per_source, + source_index=source_index, ) + if hive_parts is not None: + if source_index is not None: + hive_columns = hive_parts.gather(source_index, stream=stream) + elif rows_per_path is not None: + hive_columns = hive_parts.repeat(rows_per_path, stream=stream) + else: + hive_columns = hive_parts.broadcast(df.num_rows, stream=stream) + df = df.with_columns(hive_columns, stream=stream) + df = df.select( + [ + name + for name in schema + if row_index is None or name != row_index[0] + ] + ) if filters is not None and effective_predicate is None: return df elif typ == "ndjson": diff --git a/python/cudf_polars/cudf_polars/dsl/to_ast.py b/python/cudf_polars/cudf_polars/dsl/to_ast.py index edd33c87efce..cc8fc9fa57c0 100644 --- a/python/cudf_polars/cudf_polars/dsl/to_ast.py +++ b/python/cudf_polars/cudf_polars/dsl/to_ast.py @@ -15,7 +15,7 @@ from cudf_polars.containers import DataType from cudf_polars.dsl import expr -from cudf_polars.dsl.traversal import CachingVisitor, reuse_if_unchanged +from cudf_polars.dsl.traversal import CachingVisitor, reuse_if_unchanged, traversal from cudf_polars.typing import GenericTransformer if TYPE_CHECKING: @@ -277,8 +277,13 @@ def _extract_conjuncts(node: expr.Expr) -> list[expr.Expr]: def _to_parquet_filter( - node: expr.Expr, mapper: Transformer + node: expr.Expr, mapper: Transformer, unreadable_columns: frozenset[str] ) -> plc_expr.Expression | None: + if unreadable_columns and any( + isinstance(child, expr.Col) and child.name in unreadable_columns + for child in traversal([node]) + ): + return None # Converts a boolean column reference (e.g., filter(pl.col("foo"))) # to an explicit comparison for parquet filters (e.g., filter(pl.col("foo") == True)). # TODO: Have polars pass us the comparison instead @@ -296,7 +301,9 @@ def _to_parquet_filter( def to_parquet_filter( - node: expr.Expr, stream: Stream + node: expr.Expr, + stream: Stream, + unreadable_columns: frozenset[str] = frozenset(), ) -> tuple[plc_expr.Expression | None, expr.Expr | None]: """ Convert an expression to libcudf AST nodes suitable for parquet filtering. @@ -307,6 +314,11 @@ def to_parquet_filter( Expression to convert. stream CUDA stream used for device memory operations and kernel launches. + unreadable_columns + Names of columns that are part of the scan's schema but are not stored + in the files, such as hive partition keys. Conjuncts referencing them + are left to the residual so they can be applied once those columns have + been materialized. Returns ------- @@ -321,13 +333,13 @@ def to_parquet_filter( mapper: Transformer = CachingVisitor( _to_ast, state={"for_parquet": True, "stream": stream} ) - whole = _to_parquet_filter(node, mapper) + whole = _to_parquet_filter(node, mapper, unreadable_columns) if whole is not None: return whole, None can_handle_filters = [] cant_handle_exprs = [] for conjunct in _extract_conjuncts(node): - f = _to_parquet_filter(conjunct, mapper) + f = _to_parquet_filter(conjunct, mapper, unreadable_columns) if f is not None: can_handle_filters.append(f) else: diff --git a/python/cudf_polars/cudf_polars/dsl/translate.py b/python/cudf_polars/cudf_polars/dsl/translate.py index 23642bf64196..2619c418e5a4 100644 --- a/python/cudf_polars/cudf_polars/dsl/translate.py +++ b/python/cudf_polars/cudf_polars/dsl/translate.py @@ -30,6 +30,7 @@ from cudf_polars.dsl.utils.aggregations import decompose_single_agg from cudf_polars.dsl.utils.groupby import rewrite_groupby from cudf_polars.dsl.utils.naming import unique_names +from cudf_polars.dsl.utils.per_path import PerPathValues from cudf_polars.dsl.utils.replace import replace from cudf_polars.dsl.utils.rolling import rewrite_rolling from cudf_polars.typing import Schema @@ -501,8 +502,11 @@ def _(node: plrs._ir_nodes.Scan, translator: Translator, schema: Schema) -> ir.I raise NotImplementedError( "Iceberg format is not supported in cudf-polars. Furthermore, row-level deletions are not supported." ) # pragma: no cover - if not POLARS_VERSION_LT_142 and node.hive_parts is not None: - raise NotImplementedError("Hive-partitioned scans are not supported") + hive_parts = ( + None + if POLARS_VERSION_LT_142 or node.hive_parts is None + else PerPathValues.from_polars(pl.DataFrame._from_pydf(node.hive_parts)) + ) config_options = translator.config_options parquet_options = config_options.parquet_options @@ -550,6 +554,7 @@ def _(node: plrs._ir_nodes.Scan, translator: Translator, schema: Schema) -> ir.I ) ), parquet_options, + hive_parts=hive_parts, cached_parquet_info=None, ) diff --git a/python/cudf_polars/cudf_polars/dsl/utils/per_path.py b/python/cudf_polars/cudf_polars/dsl/utils/per_path.py new file mode 100644 index 000000000000..fea1b1cd592a --- /dev/null +++ b/python/cudf_polars/cudf_polars/dsl/utils/per_path.py @@ -0,0 +1,215 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Values attached to each path of a file scan.""" + +from __future__ import annotations + +import dataclasses +import functools +from typing import TYPE_CHECKING, Any + +import polars as pl + +import pylibcudf as plc + +from cudf_polars.containers import Column, DataType + +if TYPE_CHECKING: + from collections.abc import Sequence + + from rmm.pylibrmm.stream import Stream + +__all__ = ["PerPathValues"] + + +@dataclasses.dataclass(frozen=True, eq=False, repr=False) +class PerPathValues: + """ + Values that vary per path in a file scan, one row per path. + + For example, a hive-partitioned scan of a dataset written with + ``partition_by=["cat", "part"]`` supplies four paths and this frame:: + + shape: (4, 2) + ┌──────┬─────┐ + │ part ┆ cat │ + │ --- ┆ --- │ + │ i64 ┆ str │ + ╞══════╪═════╡ + │ 1 ┆ u │ + │ 2 ┆ u │ + │ 3 ┆ v │ + │ 4 ┆ v │ + └──────┴─────┘ + + Row ``i`` holds the values for path ``i``. + The columns might not be in the order the keys appear in the + paths. + + Parameters + ---------- + df + One row per path, one column per value to materialize. + """ + + df: pl.DataFrame + + @classmethod + def from_polars(cls, df: pl.DataFrame) -> PerPathValues | None: + """ + Build from the ``hive_parts`` dataframe of a polars ``Scan`` node. + + Parameters + ---------- + df + One row per path in the scan, one column per hive key. + + Returns + ------- + The partition values, or ``None`` if no hive columns are needed. + """ + if df.width == 0: + # Polars filtered out all paths + return None + return cls(df) + + @functools.cached_property + def names(self) -> tuple[str, ...]: + """Names of the columns to materialize.""" + return tuple(self.df.columns) + + @functools.cached_property + def dtypes(self) -> tuple[DataType, ...]: + """Datatype of each column.""" + return tuple(DataType(dtype) for dtype in self.df.dtypes) + + @property + def num_paths(self) -> int: + """Number of paths these values describe.""" + return self.df.height + + @functools.cached_property + def is_uniform(self) -> bool: + """Whether every path shares the same values.""" + return all(series.n_unique() <= 1 for series in self.df.iter_columns()) + + def slice(self, start: int, stop: int) -> PerPathValues: + """ + Restrict to the paths in ``range(start, stop)``. + + Parameters + ---------- + start + Index of the first path to keep. + stop + Index one past the last path to keep. + + Returns + ------- + Values for the selected paths. + """ + return type(self)(self.df.slice(start, stop - start)) + + def broadcast(self, num_rows: int, *, stream: Stream) -> list[Column]: + """ + Materialize the values as columns of ``num_rows`` equal rows. + + Only valid when :attr:`is_uniform` holds, since every output row is + given the values of the first path. + + Parameters + ---------- + num_rows + Length of the returned columns. + stream + CUDA stream used for device memory operations and kernel launches. + + Returns + ------- + One column per value. + """ + return self._to_columns( + plc.filling.repeat( + plc.Table.from_arrow(self.df.head(1), stream=stream), + num_rows, + stream=stream, + ) + ) + + def repeat(self, rows_per_path: Sequence[int], *, stream: Stream) -> list[Column]: + """ + Materialize the values by repeating each path's row. + + Prefer gather when a source index is known. + + Parameters + ---------- + rows_per_path + Number of output rows contributed by each path. + stream + CUDA stream used for device memory operations and kernel launches. + + Returns + ------- + One column per value, of length ``sum(rows_per_path)``. + """ + return self._to_columns( + plc.filling.repeat( + plc.Table.from_arrow(self.df, stream=stream), + plc.Column.from_arrow( + pl.Series(values=rows_per_path, dtype=pl.Int32()), stream=stream + ), + stream=stream, + ) + ) + + def gather(self, source_index: plc.Column, *, stream: Stream) -> list[Column]: + """ + Materialize the values by indexing them with a source index. + + Parameters + ---------- + source_index + Column giving, for each output row, the index of the path it was + read from. + stream + CUDA stream used for device memory operations and kernel launches. + + Returns + ------- + One column per value, aligned with ``source_index``. + """ + return self._to_columns( + plc.copying.gather( + plc.Table.from_arrow(self.df, stream=stream), + source_index, + plc.copying.OutOfBoundsPolicy.DONT_CHECK, + stream=stream, + ) + ) + + def _to_columns(self, table: plc.Table) -> list[Column]: + return [ + Column(column, name=name, dtype=dtype) + for column, name, dtype in zip( + table.columns(), self.names, self.dtypes, strict=True + ) + ] + + def __repr__(self) -> str: + """Representation showing the values.""" + return f"{type(self).__name__}(df={self.df!r})" + + def __hash__(self) -> int: + """Hash of the schema and of every value.""" + return hash( + (tuple(self.df.schema.items()), tuple(self.df.hash_rows().to_list())) + ) + + def __eq__(self, other: Any) -> bool: + """Whether two sets of values agree.""" + return ( + isinstance(other, PerPathValues) + and self.df.schema == other.df.schema + and self.df.equals(other.df) + ) diff --git a/python/cudf_polars/cudf_polars/streaming/io.py b/python/cudf_polars/cudf_polars/streaming/io.py index 638691f30045..6409449bf26d 100644 --- a/python/cudf_polars/cudf_polars/streaming/io.py +++ b/python/cudf_polars/cudf_polars/streaming/io.py @@ -49,6 +49,7 @@ from cudf_polars.containers import DataType from cudf_polars.dsl.expr import NamedExpr from cudf_polars.dsl.ir import CachedParquetInfo, IRExecutionContext + from cudf_polars.dsl.utils.per_path import PerPathValues from cudf_polars.streaming.base import ( DataSourceInfo, SerializedDataSourceInfo, @@ -210,6 +211,7 @@ def hybrid_scan_eligible( row_index: tuple[str, int] | None, include_file_paths: str | None, predicate: NamedExpr | None, + hive_parts: PerPathValues | None, ) -> bool: """Whether a parquet split is eligible for the HybridScanReader path.""" return ( @@ -218,6 +220,8 @@ def hybrid_scan_eligible( and row_index is None and include_file_paths is None and predicate is not None + # TODO: Support hive partitioning + and hive_parts is None ) @@ -358,6 +362,7 @@ class SplitScan(IR): __slots__ = ( "base_scan", "cached_parquet_info", + "hive_parts", "parquet_options", "paths", "schema", @@ -371,8 +376,9 @@ class SplitScan(IR): "split_index", "total_splits", "parquet_options", + "hive_parts", ) - _n_non_child_args = 13 + _n_non_child_args = 14 base_scan: Scan """Scan operation this node is based on.""" paths: list[str] @@ -383,6 +389,8 @@ class SplitScan(IR): """Total number of splits.""" parquet_options: ParquetOptions """Parquet-specific options.""" + hive_parts: PerPathValues | None + """Hive partition values for this split's path.""" cached_parquet_info: list[CachedParquetInfo] | None def __init__( @@ -393,6 +401,7 @@ def __init__( split_index: int, total_splits: int, parquet_options: ParquetOptions, + hive_parts: PerPathValues | None = None, cached_parquet_info: list[CachedParquetInfo] | None = None, ): self.schema = schema @@ -414,9 +423,11 @@ def __init__( base_scan.include_file_paths, base_scan.predicate, parquet_options, + hive_parts, cached_parquet_info, ) self.parquet_options = parquet_options + self.hive_parts = hive_parts self.cached_parquet_info = cached_parquet_info self.children = () if base_scan.typ not in ("parquet",): # pragma: no cover @@ -434,6 +445,7 @@ def get_hashable(self) -> Hashable: self.split_index, self.total_splits, self.parquet_options, + self.hive_parts, ) @classmethod @@ -452,6 +464,7 @@ def do_evaluate( include_file_paths: str | None, predicate: NamedExpr | None, parquet_options: ParquetOptions, + hive_parts: PerPathValues | None, cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, @@ -510,6 +523,7 @@ def do_evaluate( row_index=row_index, include_file_paths=include_file_paths, predicate=predicate, + hive_parts=hive_parts, ): assert predicate is not None assert cached_parquet_info is not None @@ -567,7 +581,8 @@ def do_evaluate( include_file_paths, predicate, parquet_options, - cached_parquet_info, + hive_parts=hive_parts, + cached_parquet_info=cached_parquet_info, context=context, ) @@ -583,6 +598,7 @@ class FusedScan(IR): __slots__ = ( "base_scan", "cached_parquet_info", + "hive_parts", "parquet_options", "paths", "schema", @@ -592,14 +608,17 @@ class FusedScan(IR): "base_scan", "paths", "parquet_options", + "hive_parts", ) - _n_non_child_args = 11 + _n_non_child_args = 12 base_scan: Scan """Scan operation this node is based on.""" paths: list[str] """File paths assigned to this task.""" parquet_options: ParquetOptions """Parquet-specific options.""" + hive_parts: PerPathValues | None + """Hive partition values for this task's paths.""" cached_parquet_info: list[CachedParquetInfo] | None """Cached parquet metadata.""" @@ -609,12 +628,14 @@ def __init__( base_scan: Scan, paths: list[str], parquet_options: ParquetOptions, + hive_parts: PerPathValues | None = None, cached_parquet_info: list[CachedParquetInfo] | None = None, ): self.schema = schema self.base_scan = base_scan self.paths = paths self.parquet_options = parquet_options + self.hive_parts = hive_parts self.cached_parquet_info = cached_parquet_info self._non_child_args = ( base_scan.schema, @@ -628,6 +649,7 @@ def __init__( base_scan.include_file_paths, base_scan.predicate, parquet_options, + hive_parts, cached_parquet_info, ) self.children = () @@ -640,6 +662,7 @@ def get_hashable(self) -> Hashable: self.base_scan.get_hashable(), tuple(self.paths), self.parquet_options, + self.hive_parts, ) @classmethod @@ -656,6 +679,7 @@ def do_evaluate( include_file_paths: str | None, predicate: NamedExpr | None, parquet_options: ParquetOptions, + hive_parts: PerPathValues | None, cached_parquet_info: list[CachedParquetInfo] | None, *, context: IRExecutionContext, @@ -674,7 +698,8 @@ def do_evaluate( include_file_paths, predicate, parquet_options, - cached_parquet_info, + hive_parts=hive_parts, + cached_parquet_info=cached_parquet_info, context=context, ) @@ -793,7 +818,12 @@ def for_split_files( sindex = local_offset % plan.factor scans: list[SplitScan] = [] splits_created = 0 - for path in local_paths: + for path_index, path in enumerate(local_paths, start=path_offset): + hive_parts = ( + None + if base_scan.hive_parts is None + else base_scan.hive_parts.slice(path_index, path_index + 1) + ) while sindex < plan.factor and splits_created < local_count: scans.append( SplitScan( @@ -803,6 +833,7 @@ def for_split_files( sindex, plan.factor, parquet_options, + hive_parts, None, ) ) @@ -832,6 +863,9 @@ def for_fused_files( base_scan, base_scan.paths[offset : offset + plan.factor], parquet_options, + None + if base_scan.hive_parts is None + else base_scan.hive_parts.slice(offset, offset + plan.factor), None, ) for offset in range(paths_start, paths_end, plan.factor) 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 1f5543358e89..20af2ac6f0c3 100644 --- a/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py +++ b/python/cudf_polars/cudf_polars/testing/inject_gpu_engine.py @@ -139,12 +139,14 @@ def pytest_report_header(config: pytest.Config) -> str: EXPECTED_FAILURES: dict[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", - "tests/unit/io/test_delta.py::test_read_delta_version": "Need to expose hive partitioning", - "tests/unit/io/test_delta.py::test_scan_delta_schema_evolution_nested_struct_field_19915": "Need to expose hive partitioning", + "tests/unit/io/test_delta.py::test_scan_delta_version": "Delta schema evolution not yet implemented in cudf-polars: the table's files have differing column counts", + "tests/unit/io/test_delta.py::test_scan_delta_relative": "Delta schema evolution not yet implemented in cudf-polars: the table's files have differing column counts", + "tests/unit/io/test_delta.py::test_read_delta_version": "Delta schema evolution not yet implemented in cudf-polars: the table's files have differing column counts", + "tests/unit/io/test_delta.py::test_scan_delta_schema_evolution_nested_struct_field_19915": "Delta schema evolution not yet implemented in cudf-polars: the table's files have differing column counts", "tests/unit/io/test_delta.py::test_scan_delta_nanosecond_timestamp": "polars generates the wrong schema: https://github.com/pola-rs/polars/issues/23949", "tests/unit/io/test_delta.py::test_scan_delta_nanosecond_timestamp_nested": "polars generates the wrong schema: https://github.com/pola-rs/polars/issues/23949", + "tests/unit/io/test_delta.py::test_sink_delta": "Delta schema evolution not yet implemented in cudf-polars: 'All sources must have the same schema'", + "tests/unit/io/test_delta.py::test_write_delta": "Delta schema evolution not yet implemented in cudf-polars: 'All sources must have the same schema'", "tests/unit/io/test_iceberg.py::test_scan_iceberg_row_index_renamed": "Iceberg support not yet implemented in cudf-polars", "tests/unit/io/test_iceberg.py::test_scan_iceberg_extra_columns": "Iceberg support not yet implemented in cudf-polars", "tests/unit/io/test_iceberg.py::test_scan_iceberg_extra_struct_fields": "Iceberg support not yet implemented in cudf-polars", diff --git a/python/cudf_polars/tests/dsl/test_per_path.py b/python/cudf_polars/tests/dsl/test_per_path.py new file mode 100644 index 000000000000..f8c38de04f70 --- /dev/null +++ b/python/cudf_polars/tests/dsl/test_per_path.py @@ -0,0 +1,193 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import pytest + +import polars as pl +from polars.testing import assert_frame_equal + +import pylibcudf as plc + +from cudf_polars.containers import DataFrame, DataType +from cudf_polars.dsl.utils.per_path import PerPathValues +from cudf_polars.utils.cuda_stream import get_cuda_stream +from cudf_polars.utils.versions import POLARS_VERSION_LT_138 + + +@pytest.fixture +def per_path() -> PerPathValues: + return PerPathValues(pl.DataFrame({"part": [1, 2, 3], "cat": ["u", "u", "v"]})) + + +def test_from_polars() -> None: + df = pl.DataFrame( + {"part": [1, 2], "cat": ["u", "v"]}, + schema={"part": pl.Int32, "cat": pl.String}, + ) + got = PerPathValues.from_polars(df) + assert got == PerPathValues(df) + assert got is not None + assert got.names == ("part", "cat") + assert got.dtypes == (DataType(pl.Int32()), DataType(pl.String())) + + +@pytest.mark.skipif( + POLARS_VERSION_LT_138, + reason="height parameter added in Polars 1.38", +) +def test_from_polars_zero_width() -> None: + assert PerPathValues.from_polars(pl.DataFrame(height=3)) is None + + +def test_hashable(per_path: PerPathValues) -> None: + assert hash(per_path) == hash( + PerPathValues(pl.DataFrame({"part": [1, 2, 3], "cat": ["u", "u", "v"]})) + ) + + +def test_dtypes_distinguish_identical_values() -> None: + values = {"part": [1, 2]} + narrow = PerPathValues(pl.DataFrame(values, schema={"part": pl.Int32})) + wide = PerPathValues(pl.DataFrame(values, schema={"part": pl.Int64})) + assert narrow != wide + assert hash(narrow) != hash(wide) + + +def test_names_distinguish_identical_values() -> None: + # hash_rows digests the values but not the column they sit in, so the + # schema is what tells these two apart. + values = [1, 2] + part = PerPathValues(pl.DataFrame({"part": values})) + cat = PerPathValues(pl.DataFrame({"cat": values})) + assert part != cat + assert hash(part) != hash(cat) + + +def test_path_order_matters() -> None: + forwards = PerPathValues(pl.DataFrame({"part": [1, 2]})) + backwards = PerPathValues(pl.DataFrame({"part": [2, 1]})) + assert forwards != backwards + assert hash(forwards) != hash(backwards) + + +def test_tall_partitions_are_distinguished() -> None: + # Polars elides the middle of a tall frame's repr, so identity cannot + # rest on it. hash_rows sees every row. + tall = PerPathValues(pl.DataFrame({"part": range(40)})) + other = PerPathValues(pl.DataFrame({"part": [*range(20), 999, *range(21, 40)]})) + assert repr(tall.df) == repr(other.df) + assert tall != other + assert hash(tall) != hash(other) + + +def test_not_equal_to_other_types(per_path: PerPathValues) -> None: + assert per_path != per_path.df + + +def test_repr(per_path: PerPathValues) -> None: + assert repr(per_path) == f"PerPathValues(df={per_path.df!r})" + + +def test_num_paths(per_path: PerPathValues) -> None: + assert per_path.num_paths == 3 + + +@pytest.mark.parametrize( + "values,expected", + [ + ({"part": [1, 2, 3], "cat": ["u", "u", "v"]}, False), + ({"part": [1, 1, 1], "cat": ["u", "u", "v"]}, False), + ({"part": [1, 1, 1], "cat": ["u", "u", "u"]}, True), + ({"part": [1], "cat": ["u"]}, True), + ({"part": [None, None], "cat": [None, None]}, True), + ], +) +def test_is_uniform(values, expected) -> None: + assert PerPathValues(pl.DataFrame(values)).is_uniform is expected + + +def test_slice(per_path: PerPathValues) -> None: + assert per_path.slice(1, 3) == PerPathValues( + pl.DataFrame({"part": [2, 3], "cat": ["u", "v"]}) + ) + + +def test_broadcast() -> None: + stream = get_cuda_stream() + per_path = PerPathValues(pl.DataFrame({"part": [7], "cat": ["u"]})) + got = DataFrame(per_path.broadcast(3, stream=stream), stream=stream) + assert_frame_equal( + got.to_polars(), pl.DataFrame({"part": [7, 7, 7], "cat": ["u", "u", "u"]}) + ) + + +def test_broadcast_uses_first_path() -> None: + stream = get_cuda_stream() + per_path = PerPathValues(pl.DataFrame({"part": [7, 7]})) + got = DataFrame(per_path.broadcast(2, stream=stream), stream=stream) + assert_frame_equal(got.to_polars(), pl.DataFrame({"part": [7, 7]})) + + +def test_broadcast_null() -> None: + stream = get_cuda_stream() + per_path = PerPathValues(pl.DataFrame({"part": [None]}, schema={"part": pl.Int64})) + got = DataFrame(per_path.broadcast(2, stream=stream), stream=stream) + assert_frame_equal( + got.to_polars(), pl.DataFrame({"part": [None, None]}, schema={"part": pl.Int64}) + ) + + +def test_repeat(per_path: PerPathValues) -> None: + stream = get_cuda_stream() + got = DataFrame(per_path.repeat([2, 0, 1], stream=stream), stream=stream) + assert_frame_equal( + got.to_polars(), + pl.DataFrame({"part": [1, 1, 3], "cat": ["u", "u", "v"]}), + ) + + +def test_repeat_empty(per_path: PerPathValues) -> None: + stream = get_cuda_stream() + got = DataFrame(per_path.repeat([0, 0, 0], stream=stream), stream=stream) + assert got.num_rows == 0 + + +def test_gather(per_path: PerPathValues) -> None: + stream = get_cuda_stream() + source_index = plc.Column.from_arrow( + pl.Series(values=[0, 0, 2, 1], dtype=pl.Int32()), stream=stream + ) + got = DataFrame(per_path.gather(source_index, stream=stream), stream=stream) + assert_frame_equal( + got.to_polars(), + pl.DataFrame({"part": [1, 1, 3, 2], "cat": ["u", "u", "v", "u"]}), + ) + + +def test_gather_preserves_nulls() -> None: + stream = get_cuda_stream() + per_path = PerPathValues( + pl.DataFrame({"part": [1, None]}, schema={"part": pl.Int64}) + ) + source_index = plc.Column.from_arrow( + pl.Series(values=[1, 0], dtype=pl.Int32()), stream=stream + ) + got = DataFrame(per_path.gather(source_index, stream=stream), stream=stream) + assert_frame_equal(got.to_polars(), pl.DataFrame({"part": [None, 1]})) + + +@pytest.mark.parametrize( + "dtype", + [pl.Int32, pl.Int64, pl.String, pl.Float64, pl.Boolean, pl.Date, pl.Datetime("us")], +) +def test_dtypes_survive_conversion(dtype: pl.DataType) -> None: + stream = get_cuda_stream() + series = pl.Series("part", [0, 1], dtype=pl.Int64).cast(dtype, strict=False) + per_path = PerPathValues(series.to_frame()) + source_index = plc.Column.from_arrow( + pl.Series(values=[1, 0], dtype=pl.Int32()), stream=stream + ) + got = DataFrame(per_path.gather(source_index, stream=stream), stream=stream) + assert_frame_equal(got.to_polars(), series.reverse().to_frame()) diff --git a/python/cudf_polars/tests/streaming/test_scan.py b/python/cudf_polars/tests/streaming/test_scan.py index 72389d58e9fd..a6011b694c76 100644 --- a/python/cudf_polars/tests/streaming/test_scan.py +++ b/python/cudf_polars/tests/streaming/test_scan.py @@ -3,6 +3,7 @@ from __future__ import annotations +import functools import math from typing import TYPE_CHECKING, cast @@ -35,18 +36,23 @@ SplitScan, StreamingScan, expand_scan_for_rank, + hybrid_scan_eligible, scan_partition_plan, ) 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.engine_utils import SMALL_MAX_ROWS_PER_PARTITION +from cudf_polars.testing.engine_utils import ( + SMALL_MAX_ROWS_PER_PARTITION, + is_streaming_engine, +) from cudf_polars.testing.io import make_partitioned_source from cudf_polars.utils.config import ( ConfigOptions, MaxConcurrentIOTasks, ParquetOptions, ) +from cudf_polars.utils.versions import POLARS_VERSION_LT_142 if TYPE_CHECKING: import concurrent.futures @@ -178,7 +184,9 @@ def recording_prefetch( ) scan = _make_parquet_scan(paths) - fused = FusedScan(scan.schema, scan, paths, scan.parquet_options, None) + fused = FusedScan( + scan.schema, scan, paths, scan.parquet_options, cached_parquet_info=None + ) streaming_scan = StreamingScan([fused], scan, "fused") result = prefetch_parquet_file_metadata_for_ir( @@ -194,7 +202,9 @@ def test_prefetch_parquet_file_metadata_remote_only(tmp_path, df) -> None: local_path = str(next(tmp_path.glob("*.parquet"))) scan = _make_parquet_scan([local_path]) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) + fused = FusedScan( + scan.schema, scan, scan.paths, scan.parquet_options, cached_parquet_info=[] + ) streaming_scan = StreamingScan([fused], scan, "fused") # Local paths are skipped entirely when remote_only=True. @@ -410,7 +420,7 @@ def _make_parquet_scan( None, None, parquet_options, - None, + cached_parquet_info=None, ) @@ -499,7 +509,9 @@ def test_expand_scan_for_rank_split_files( def test_streaming_scan_raises() -> None: # This isn't reachable by normal cudf-polars usage. scan = _make_parquet_scan(["file.parquet"]) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) + fused = FusedScan( + scan.schema, scan, scan.paths, scan.parquet_options, cached_parquet_info=[] + ) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): StreamingScan.do_evaluate([fused], scan, context=ctx) @@ -566,6 +578,7 @@ def test_scan_path_mismatch_raises() -> None: scan.include_file_paths, scan.predicate, scan.parquet_options, + scan.hive_parts, [], context=ctx, ) @@ -576,7 +589,9 @@ def test_streaming_scan_missing_prefetch_metadata_raises() -> None: scan = _make_parquet_scan( ["file.parquet"], parquet_options=ParquetOptions(prefetch_file_metadata=True) ) - fused = FusedScan(scan.schema, scan, scan.paths, scan.parquet_options, []) + fused = FusedScan( + scan.schema, scan, scan.paths, scan.parquet_options, cached_parquet_info=[] + ) ctx = IRExecutionContext() with pytest.raises(NotImplementedError, match=r"StreamingScan.do_evaluate"): @@ -607,6 +622,7 @@ def test_split_scan_do_evaluate_missing_prefetch_metadata() -> None: None, None, parquet_options, + None, [], context=context, ) @@ -672,9 +688,15 @@ def test_fused_scan_identity_equality() -> None: paths = ["a.parquet"] info = _make_cached_parquet_info(paths) - a = FusedScan(base.schema, base, paths, base.parquet_options, info) - b = FusedScan(base.schema, base, paths, base.parquet_options, info.copy()) - c = FusedScan(base.schema, base, ["b.parquet"], base.parquet_options, info) + a = FusedScan( + base.schema, base, paths, base.parquet_options, cached_parquet_info=info + ) + b = FusedScan( + base.schema, base, paths, base.parquet_options, cached_parquet_info=info.copy() + ) + c = FusedScan( + base.schema, base, ["b.parquet"], base.parquet_options, cached_parquet_info=info + ) assert a == b assert hash(a) == hash(b) @@ -685,11 +707,33 @@ def test_split_scan_identity_equality() -> None: base = _make_parquet_scan(["a.parquet"]) info = _make_cached_parquet_info(base.paths) - a = SplitScan(base.schema, base, base.paths, 0, 4, base.parquet_options, info) + a = SplitScan( + base.schema, + base, + base.paths, + 0, + 4, + base.parquet_options, + cached_parquet_info=info, + ) b = SplitScan( - base.schema, base, base.paths, 0, 4, base.parquet_options, info.copy() + base.schema, + base, + base.paths, + 0, + 4, + base.parquet_options, + cached_parquet_info=info.copy(), + ) + c = SplitScan( + base.schema, + base, + base.paths, + 1, + 4, + base.parquet_options, + cached_parquet_info=info, ) - c = SplitScan(base.schema, base, base.paths, 1, 4, base.parquet_options, info) assert a == b assert hash(a) == hash(b) @@ -705,7 +749,7 @@ def test_streaming_scan_identity_equality() -> None: 0, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=10), + cached_parquet_info=_make_cached_parquet_info(base.paths, size=10), ) split_same = SplitScan( base.schema, @@ -714,7 +758,7 @@ def test_streaming_scan_identity_equality() -> None: 0, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=11), + cached_parquet_info=_make_cached_parquet_info(base.paths, size=11), ) split_diff = SplitScan( base.schema, @@ -723,7 +767,7 @@ def test_streaming_scan_identity_equality() -> None: 1, 2, base.parquet_options, - _make_cached_parquet_info(base.paths, size=10), + cached_parquet_info=_make_cached_parquet_info(base.paths, size=10), ) a = StreamingScan([split], base, "split") @@ -753,22 +797,38 @@ def test_cached_parquet_info_excluded_from_identity() -> None: None, None, base.parquet_options, - info, + cached_parquet_info=info, ) assert scan_without == scan_with assert hash(scan_without) == hash(scan_with) split_without = SplitScan( - base.schema, base, base.paths, 0, 4, base.parquet_options, None + base.schema, + base, + base.paths, + 0, + 4, + base.parquet_options, + cached_parquet_info=None, ) split_with = SplitScan( - base.schema, base, base.paths, 0, 4, base.parquet_options, info + base.schema, + base, + base.paths, + 0, + 4, + base.parquet_options, + cached_parquet_info=info, ) assert split_without == split_with assert hash(split_without) == hash(split_with) - fused_without = FusedScan(base.schema, base, base.paths, base.parquet_options, None) - fused_with = FusedScan(base.schema, base, base.paths, base.parquet_options, info) + fused_without = FusedScan( + base.schema, base, base.paths, base.parquet_options, cached_parquet_info=None + ) + fused_with = FusedScan( + base.schema, base, base.paths, base.parquet_options, cached_parquet_info=info + ) assert fused_without == fused_with assert hash(fused_without) == hash(fused_with) @@ -837,3 +897,157 @@ def test_scan_partition_plan_nearest( plan = scan_partition_plan(scan, FooStats(scan, file_size), _make_config(10)) assert plan.factor == expected_factor assert plan.flavor == expected_flavor + + +requires_hive_ir = pytest.mark.skipif( + POLARS_VERSION_LT_142, + reason="hive::HivePartitionedDf not exposed in the logical plan before 1.42", +) + + +@pytest.fixture +def hive_root(tmp_path: Path) -> Path: + """Hive dataset with several row groups per file, to allow file splitting.""" + root = tmp_path / "hive" + pl.DataFrame( + { + "x": range(600), + "part": [i // 200 for i in range(600)], + } + ).write_parquet(root, partition_by=["part"], row_group_size=25) + return root + + +@requires_hive_ir +@pytest.mark.parametrize( + "target_partition_size,expected_flavor", + [ + (1_000, IOPartitionFlavor.SPLIT_FILES), + (1_000_000, IOPartitionFlavor.FUSED_FILES), + ], +) +@pytest.mark.parametrize( + "query", + [ + lambda lf: lf, + lambda lf: lf.select("part"), + lambda lf: lf.filter(pl.col("x") > 400), + lambda lf: lf.filter(pl.col("part") == 1), + lambda lf: lf.filter((pl.col("part") == 1) & (pl.col("x") > 250)), + lambda lf: lf.group_by("part").agg(pl.col("x").sum()), + ], +) +def test_hive_partitioned_streaming_scan( + hive_root: Path, + streaming_engine_factory: Callable[..., StreamingEngine], + target_partition_size: int, + expected_flavor: IOPartitionFlavor, + query, +) -> None: + streaming_engine = streaming_engine_factory( + StreamingOptions(target_partition_size=target_partition_size), + ) + q = query(pl.scan_parquet(hive_root, hive_partitioning=True)) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) + + +@requires_hive_ir +def test_hive_partitioned_split_scan_slices_partitions( + hive_root: Path, engine: pl.GPUEngine +) -> None: + if not is_streaming_engine(engine): + pytest.skip("SplitScan/FusedScan are only built for streaming engines") + q = pl.scan_parquet(hive_root, hive_partitioning=True) + scan = cast("Scan", Translator(q._ldf.visit(), engine).translate_ir()) + assert scan.hive_parts is not None + + streaming = StreamingScan.for_split_files( + scan, + IOPartitionPlan(2, IOPartitionFlavor.SPLIT_FILES), + 2 * len(scan.paths), + rank=0, + nranks=1, + parquet_options=ParquetOptions(), + ) + # Both splits of a file see that file's partition values, and nothing else. + hive_parts = [] + for split in streaming.scans: + assert split.hive_parts is not None + assert split.hive_parts.num_paths == 1 + assert split.hive_parts.is_uniform + hive_parts.append(split.hive_parts) + assert [parts.df.rows() for parts in hive_parts] == [ + [(0,)], + [(0,)], + [(1,)], + [(1,)], + [(2,)], + [(2,)], + ] + + +@requires_hive_ir +def test_hive_partitioned_fused_scan_slices_partitions( + hive_root: Path, engine: pl.GPUEngine +) -> None: + if not is_streaming_engine(engine): + pytest.skip("SplitScan/FusedScan are only built for streaming engines") + q = pl.scan_parquet(hive_root, hive_partitioning=True) + scan = cast("Scan", Translator(q._ldf.visit(), engine).translate_ir()) + + streaming = StreamingScan.for_fused_files( + scan, + IOPartitionPlan(2, IOPartitionFlavor.FUSED_FILES), + 2, + rank=0, + nranks=1, + parquet_options=ParquetOptions(), + ) + hive_parts = [] + for fused in streaming.scans: + assert fused.hive_parts is not None + hive_parts.append(fused.hive_parts) + assert [parts.df.rows() for parts in hive_parts] == [[(0,), (1,)], [(2,)]] + + +@requires_hive_ir +def test_hive_partitioned_scan_skips_hybrid_scan( + hive_root: Path, engine: pl.GPUEngine +) -> None: + # The hybrid reader cannot keep hive columns out of what it asks the file + # for, so a hive scan must fall back to the regular reader. + if not is_streaming_engine(engine): + pytest.skip("SplitScan/FusedScan are only built for streaming engines") + q = pl.scan_parquet(hive_root, hive_partitioning=True).filter(pl.col("x") > 400) + scan = cast("Scan", Translator(q._ldf.visit(), engine).translate_ir()) + assert scan.hive_parts is not None + assert scan.predicate is not None + + eligibility = functools.partial( + hybrid_scan_eligible, + ParquetOptions(use_hybrid_scan=True), + cached_parquet_info=_make_cached_parquet_info(scan.paths), + row_index=None, + include_file_paths=None, + predicate=scan.predicate, + ) + assert eligibility(hive_parts=None) is True + assert eligibility(hive_parts=scan.hive_parts) is False + + +@requires_hive_ir +def test_hive_partitioned_scan_with_hybrid_scan_enabled( + hive_root: Path, + streaming_engine_factory: Callable[..., StreamingEngine], +) -> None: + streaming_engine = streaming_engine_factory( + StreamingOptions( + target_partition_size=1_000, + parquet_options={ + "prefetch_file_metadata": True, + "use_hybrid_scan": True, + }, + ), + ) + q = pl.scan_parquet(hive_root, hive_partitioning=True).filter(pl.col("x") > 400) + assert_gpu_result_equal(q, engine=streaming_engine, check_row_order=False) diff --git a/python/cudf_polars/tests/test_scan.py b/python/cudf_polars/tests/test_scan.py index b757892b455f..7f20e2138dc0 100644 --- a/python/cudf_polars/tests/test_scan.py +++ b/python/cudf_polars/tests/test_scan.py @@ -38,6 +38,27 @@ from werkzeug import Request +requires_hive_ir = pytest.mark.skipif( + POLARS_VERSION_LT_142, + reason="hive::HivePartitionedDf not exposed in the logical plan before 1.42", +) + + +@pytest.fixture +def hive_root(tmp_path: Path) -> Path: + """A dataset partitioned by two keys, with differing rows per file.""" + root = tmp_path / "hive" + pl.DataFrame( + { + "a": [1, 2, 3, 4, 5, 6, 7, 8, 9], + "b": ["x", "y", "z", "w", "p", "q", "r", "s", "t"], + "part": [1, 1, 1, 2, 2, 3, 3, 4, 4], + "cat": ["u", "u", "u", "u", "u", "v", "v", "v", "v"], + } + ).write_parquet(root, partition_by=["cat", "part"]) + return root + + NO_CHUNK_ENGINE = pl.GPUEngine( executor="in-memory", raise_on_fail=True, parquet_options={"chunked": False} ) @@ -206,6 +227,7 @@ def test_scan_do_evaluate_missing_prefetch_metadata() -> None: None, None, parquet_options, + None, [], context=context, ) @@ -452,6 +474,50 @@ def test_scan_include_file_path( assert_gpu_result_equal(q, engine=NO_CHUNK_ENGINE) +@pytest.mark.parametrize( + "predicate", + [ + None, + pl.col("a") > 20, + pl.col("a") < 0, + pl.col("a") >= 10, + ], +) +def test_scan_parquet_include_file_path_with_predicate( + engine: pl.GPUEngine, tmp_path: Path, predicate +) -> None: + # A pushed-down filter clears the reader's per-source row counts, so the + # paths have to be recovered from the source index instead. + for i, height in enumerate([3, 2, 4]): + pl.DataFrame({"a": [i * 10 + j for j in range(height)]}).write_parquet( + tmp_path / f"part-{i}.parquet" + ) + q = pl.scan_parquet(tmp_path, include_file_paths="files") + if predicate is not None: + q = q.filter(predicate) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) + + +@requires_hive_ir +def test_scan_parquet_include_file_path_with_hive( + engine: pl.GPUEngine, hive_root: Path +) -> None: + q = pl.scan_parquet( + hive_root, hive_partitioning=True, include_file_paths="files" + ).filter(pl.col("a") > 3) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) + + +@requires_hive_ir +def test_scan_parquet_include_file_path_with_hive_only_projection( + engine: pl.GPUEngine, hive_root: Path +) -> None: + q = pl.scan_parquet( + hive_root, hive_partitioning=True, include_file_paths="files" + ).select("part", "files") + assert_gpu_result_equal(q, engine=engine, check_row_order=False) + + @pytest.fixture( scope="module", params=["no_slice", "skip_to_end", "skip_partial", "partial"] ) @@ -870,14 +936,155 @@ def test_scan_parquet_is_between_literal_dtype_mismatch_22622( assert_gpu_result_equal(q, engine=engine) -@pytest.mark.skipif( - POLARS_VERSION_LT_142, - reason="hive::HivePartitionedDf not exposed in the logical plan before 1.42", +@requires_hive_ir +@pytest.mark.parametrize( + "query", + [ + lambda lf: lf, + lambda lf: lf.select("a"), + lambda lf: lf.select("part"), + lambda lf: lf.select("part", "cat"), + lambda lf: lf.select("part", "a"), + lambda lf: lf.select("cat", "b", "part"), + lambda lf: lf.filter(pl.col("part") == 2), + lambda lf: lf.filter(pl.col("part") >= 3), + lambda lf: lf.filter(pl.col("part") == 99), + lambda lf: lf.filter(pl.col("cat") == "v"), + lambda lf: lf.filter(pl.col("a") > 4), + lambda lf: lf.filter((pl.col("part") >= 3) & (pl.col("a") > 6)), + lambda lf: lf.filter((pl.col("cat") == "v") | (pl.col("a") > 6)), + lambda lf: lf.filter(pl.col("part") > pl.col("a")), + lambda lf: lf.filter(pl.col("part").is_in([1, 3])), + lambda lf: lf.filter(pl.col("part") % 2 == 0).select("part"), + lambda lf: lf.filter(pl.col("a") > 3).select("part"), + lambda lf: lf.head(4), + lambda lf: lf.select("part").head(4), + lambda lf: lf.slice(2, 5), + lambda lf: lf.with_row_index(), + lambda lf: lf.group_by("part").agg(pl.col("a").sum()), + ], ) -def test_scan_parquet_hive_partitioned_raises( +def test_scan_parquet_hive_partitioned( + engine: pl.GPUEngine, hive_root: Path, query +) -> None: + q = query(pl.scan_parquet(hive_root, hive_partitioning=True)) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) + + +@requires_hive_ir +def test_scan_parquet_hive_partitioned_schema_override( + engine: pl.GPUEngine, hive_root: Path +) -> None: + q = pl.scan_parquet( + hive_root, hive_schema={"cat": pl.String, "part": pl.Int32} + ).filter(pl.col("part") > 1) + assert_gpu_result_equal(q, engine=engine) + + +@requires_hive_ir +def test_scan_parquet_hive_partitioned_single_file( engine: pl.GPUEngine, tmp_path: Path ) -> None: (tmp_path / "part=1").mkdir() pl.DataFrame({"x": [1, 2, 3]}).write_parquet(tmp_path / "part=1" / "data.parquet") q = pl.scan_parquet(tmp_path, hive_schema={"part": pl.Int32}) - assert_ir_translation_raises(q, engine, NotImplementedError) + assert_gpu_result_equal(q, engine=engine) + + +@requires_hive_ir +def test_scan_parquet_hive_partitioned_shadowed_column( + engine: pl.GPUEngine, tmp_path: Path +) -> None: + for name, values in [("part=1", [100, 200]), ("part=2", [300, 400])]: + (tmp_path / name).mkdir() + pl.DataFrame({"a": [1, 2], "part": values}).write_parquet( + tmp_path / name / "data.parquet" + ) + q = pl.scan_parquet(tmp_path, hive_partitioning=True) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) + + +@requires_hive_ir +def test_scan_parquet_hive_partitioned_null_value( + engine: pl.GPUEngine, tmp_path: Path +) -> None: + pl.DataFrame({"a": [1, 2, 3, 4], "part": ["u", "u", None, None]}).write_parquet( + tmp_path / "hive", partition_by=["part"] + ) + q = pl.scan_parquet(tmp_path / "hive", hive_partitioning=True) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) + + +@requires_hive_ir +@pytest.mark.parametrize( + "query", + [ + lambda lf: lf, + lambda lf: lf.select("part"), + lambda lf: lf.filter(pl.col("part") > 1), + lambda lf: lf.filter(pl.col("a") > 4), + ], +) +def test_scan_parquet_hive_partitioned_chunked(hive_root: Path, query) -> None: + q = query(pl.scan_parquet(hive_root, hive_partitioning=True)) + assert_gpu_result_equal( + q, + engine=pl.GPUEngine( + executor="in-memory", + raise_on_fail=True, + parquet_options={"chunked": True, "chunk_read_limit": 1}, + ), + check_row_order=False, + ) + + +@requires_hive_ir +@pytest.mark.parametrize( + "dtype", + [pl.Date, pl.Datetime("us"), pl.Float64, pl.Boolean, pl.Int64, pl.String], +) +@pytest.mark.parametrize( + "query", [lambda lf: lf, lambda lf: lf.select("part")], ids=["all", "hive_only"] +) +def test_scan_parquet_hive_partitioned_dtypes( + engine: pl.GPUEngine, tmp_path: Path, dtype: pl.DataType, query +) -> None: + values = pl.Series([0, 1, 1, 0], dtype=pl.Int64).cast(dtype, strict=False) + root = tmp_path / "hive" + pl.DataFrame({"a": [1, 2, 3, 4], "part": values}).write_parquet( + root, partition_by=["part"] + ) + q = query(pl.scan_parquet(root, hive_schema={"part": dtype})) + assert_gpu_result_equal(q, engine=engine, check_row_order=False) + + +@requires_hive_ir +@pytest.mark.parametrize( + "offset,length", + [(0, None), (0, 4), (4, None), (3, 3), (9, None), (20, 5), (0, 0)], +) +def test_scan_parquet_hive_only_projection_sliced( + engine: pl.GPUEngine, tmp_path: Path, offset: int, length: int | None +) -> None: + root = tmp_path / "hive" + for part, height in enumerate([3, 2, 4]): + (root / f"part={part}").mkdir(parents=True) + pl.DataFrame({"a": range(height)}).write_parquet( + root / f"part={part}" / "data.parquet" + ) + q = pl.scan_parquet(root, hive_schema={"part": pl.Int64}).select("part") + q = q.slice(offset) if length is None else q.slice(offset, length) + assert_gpu_result_equal(q, engine=engine) + + +@requires_hive_ir +def test_scan_parquet_hive_partitioned_uniform_multiple_files( + engine: pl.GPUEngine, tmp_path: Path +) -> None: + (tmp_path / "part=1").mkdir() + for index in range(3): + pl.DataFrame({"a": [index, index + 1]}).write_parquet( + tmp_path / "part=1" / f"{index}.parquet" + ) + q = pl.scan_parquet(tmp_path, hive_schema={"part": pl.Int32}) + assert_gpu_result_equal(q, engine=engine, check_row_order=False)