chore: import upstream snapshot with attribution

python/ray/data/_internal/arrow_ops/transform_polars.py

@@ -0,0 +1,59 @@
+from typing import TYPE_CHECKING, List
+
+from packaging.version import parse as parse_version
+
+try:
+    import pyarrow
+except ImportError:
+    pyarrow = None
+
+
+if TYPE_CHECKING:
+    from ray.data._internal.planner.exchange.sort_task_spec import SortKey
+
+pl = None
+
+# Polars 0.16.8 introduced the `descending` parameter for the `sort` method,
+# replacing `reverse`.
+# See https://github.com/pola-rs/polars/issues/5429 for more details.
+_POLARS_SORT_DESCENDING_MIN_VERSION = parse_version("0.16.8")
+
+
+def check_polars_installed():
+    try:
+        global pl
+        import polars as pl
+    except ImportError:
+        raise ImportError(
+            "polars not installed. Install with `pip install polars` or set "
+            "`DataContext.use_polars_sort = False` to fall back to pyarrow"
+        )
+
+
+def sort(table: "pyarrow.Table", sort_key: "SortKey") -> "pyarrow.Table":
+    check_polars_installed()
+    df = pl.from_arrow(table)
+    if parse_version(pl.__version__) >= _POLARS_SORT_DESCENDING_MIN_VERSION:
+        return df.sort(
+            sort_key.get_columns(), descending=sort_key.get_descending()
+        ).to_arrow()
+    else:
+        return df.sort(
+            sort_key.get_columns(), reverse=sort_key.get_descending()
+        ).to_arrow()
+
+
+def concat_and_sort(
+    blocks: List["pyarrow.Table"], sort_key: "SortKey", *, promote_types: bool = False
+) -> "pyarrow.Table":
+    check_polars_installed()
+    blocks = [pl.from_arrow(block) for block in blocks]
+    if parse_version(pl.__version__) >= _POLARS_SORT_DESCENDING_MIN_VERSION:
+        df = pl.concat(blocks).sort(
+            sort_key.get_columns(), descending=sort_key.get_descending()
+        )
+    else:
+        df = pl.concat(blocks).sort(
+            sort_key.get_columns(), reverse=sort_key.get_descending()
+        )
+    return df.to_arrow()