You are viewing a plain text version of this content. The canonical link for it is here.
Posted to issues@sedona.apache.org by "jiayuasu (via GitHub)" <gi...@apache.org> on 2026/06/27 01:19:45 UTC
[PR] feat(python/sedonadb): add DataFrame.with_column and with_column_renamed [sedona-db]
jiayuasu opened a new pull request, #1010:
URL: https://github.com/apache/sedona-db/pull/1010
Adds two relational column operations to the Python `DataFrame`, in the engine-neutral (Spark/Polars) naming, continuing the Ibis/DuckDB-style relational surface.
## API
```python
df.with_column("c", df.a + df.b) # add a column
df.with_column("b", df.b * 2) # replace an existing column in place
df.with_column("k", lit(9)) # constant column
df.with_column("a_copy", "a") # copy a column by name
df.with_column_renamed("b", "c") # rename one column
```
- `with_column(name, expr)` — add a column, or replace the existing column of that name in place (order preserved). `expr` accepts an `Expr`, a column-name `str` (copied), or a `lit()` literal. For building several columns at once, `select(..., **kwargs)` is usually more convenient.
- `with_column_renamed(old, new)` — rename a single column, preserving order and all other columns.
## Footgun guard
DataFusion's `with_column_renamed` **silently no-ops** when the source column doesn't exist, which hides typos. The Python wrapper validates the source name and raises a `KeyError` listing the available columns — same pattern already used by `drop`.
## Implementation
| File | Change |
|---|---|
| `python/sedonadb/src/dataframe.rs` | `InternalDataFrame::with_column` / `with_column_renamed` — thin wrappers over DataFusion. |
| `python/sedonadb/python/sedonadb/dataframe.py` | `DataFrame.with_column` (Expr/str/Literal coercion + name-type check) and `with_column_renamed` (KeyError guard). |
## Test plan
`tests/expr/test_dataframe_with_column.py`:
- `with_column`: append; in-place replace (order preserved); from a column-name str; from a literal; bad name type; bad value type; lazy return.
- `with_column_renamed`: rename; order preservation; unknown-column `KeyError` (listing available columns); lazy return.
Local: 11 unit + 33 doctests + expr suite + `ruff` + `cargo fmt --check` + `clippy -Dwarnings` all clean.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): add DataFrame.with_column and with_column_renamed [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu commented on PR #1010:
URL: https://github.com/apache/sedona-db/pull/1010#issuecomment-4840575568
Good points — I'm convinced. The chained-`with_column` anti-pattern (my own docstring flags it), the efficiency win of a single projection, the "exactly one obvious way" principle, and the Ibis/DuckDB framing this API already advertises all line up. Let me pivot this PR to the bulk versions:
- **`mutate(**exprs)`** — keep all existing columns plus the named/replaced ones, mirroring how `select(**kwargs)` already coerces (and accepting positional aliased `Expr`s too). This fills the gap `select` doesn't: "all columns plus these" without re-listing everything.
- **`rename(...)`** — bulk rename.
One thing I want your call on before I implement `rename`: the **direction/signature**, since the popular libraries disagree and I'd rather not bake in a silent-swap footgun:
1. **`rename({"old": "new"})`** — dict, old→new. Matches **pandas and Polars**.
2. **`rename(new_name="old_name")`** — kwargs, new→old. Matches **Ibis and dplyr** (the framing this API cites), though the new name must be a valid identifier and the reversal surprises pandas users.
You linked Ibis and noted pandas rename is "not very flexible," so I'm guessing you lean toward (2) — but which would you like? (Or a dict in the Ibis `{new: old}` direction for arbitrary names?)
Also — good note on `con.col()`/`con.lit()` in the tests; I'll switch the new tests to the context constructors as part of the rework.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): add DataFrame.with_column and with_column_renamed [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #1010:
URL: https://github.com/apache/sedona-db/pull/1010#discussion_r3492380036
##########
python/sedonadb/tests/expr/test_dataframe_with_column.py:
##########
@@ -0,0 +1,106 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pandas as pd
+import pandas.testing as pdt
+import pytest
+
+from sedonadb.dataframe import DataFrame
+from sedonadb.expr import col, lit
Review Comment:
Tests should use `con.col()` and `con.lit()` to set a good example...these versions aren't user-accessible via `sedona.db.xx` and don't have am embedded context (so chaining via `.geo.xxx()` doesn't work). There are some preexisting tests that I missed before this got added but at least for new ones we should use the constructors from the context.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): add DataFrame.with_column and with_column_renamed [sedona-db]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #1010:
URL: https://github.com/apache/sedona-db/pull/1010#discussion_r3488554772
##########
python/sedonadb/python/sedonadb/dataframe.py:
##########
@@ -1067,6 +1067,89 @@ def except_distinct(self, other: "DataFrame") -> "DataFrame":
self._check_set_op_compatible(other, "except_distinct")
return DataFrame(self._ctx, self._impl.except_distinct(other._impl))
+ def with_column(self, name: str, expr: Union[Expr, str, _SedonaLit]) -> "DataFrame":
+ """Add a column, or replace an existing column of the same name.
+
+ Returns a new lazy DataFrame with all existing columns plus `name`
+ (Spark / Polars `with_column` naming). If `name` already exists it
+ is replaced in place; otherwise it is appended. To build several
+ columns at once, `select(...)` with keyword arguments is usually
+ more convenient.
+
+ Args:
+ name: The output column name.
+ expr: The column expression. An `Expr`, a column-name `str`
+ (copied as that column), or a `lit()` literal.
+
+ Examples:
+
+ >>> sd = sedona.db.connect()
+ >>> df = sd.sql("SELECT 1 AS a, 2 AS b")
+ >>> df.with_column("c", df.a + df.b).show()
+ ┌───────┬───────┬───────┐
+ │ a ┆ b ┆ c │
+ │ int64 ┆ int64 ┆ int64 │
+ ╞═══════╪═══════╪═══════╡
+ │ 1 ┆ 2 ┆ 3 │
+ └───────┴───────┴───────┘
+ """
+ if not isinstance(name, str):
+ raise TypeError(
+ f"with_column() name must be a str, got {type(name).__name__}"
+ )
+
+ if isinstance(expr, Expr):
+ inner_expr = expr._impl
+ elif isinstance(expr, str):
+ inner_expr = _col(expr)._impl
+ elif isinstance(expr, _SedonaLit):
+ inner_expr = _to_expr(expr)._impl
+ else:
+ raise TypeError(
+ f"with_column() expects an Expr, str, or Literal value, "
+ f"got {type(expr).__name__}"
+ )
+
+ return DataFrame(self._ctx, self._impl.with_column(name, inner_expr))
+
+ def with_column_renamed(self, old_name: str, new_name: str) -> "DataFrame":
+ """Rename a single column.
+
+ Returns a new lazy DataFrame with `old_name` renamed to `new_name`
+ (Spark naming); all other columns are unchanged. To rename several
+ columns, chain calls or use `select(*[col(k).alias(v) ...])`.
+
+ Args:
+ old_name: The existing column name.
+ new_name: The new name.
+
+ Raises:
+ KeyError: If `old_name` is not a column. (The message lists the
+ available columns.)
+
+ Examples:
+
+ >>> sd = sedona.db.connect()
+ >>> df = sd.sql("SELECT 1 AS a, 2 AS b")
+ >>> df.with_column_renamed("b", "c").show()
+ ┌───────┬───────┐
+ │ a ┆ c │
+ │ int64 ┆ int64 │
+ ╞═══════╪═══════╡
+ │ 1 ┆ 2 │
+ └───────┴───────┘
+ """
+ # DataFusion silently no-ops when the column is missing, which hides
+ # typos. Validate Python-side so the caller gets a clear KeyError
+ # listing the available columns (matching `drop`).
+ columns = self._impl.columns()
+ if old_name not in columns:
Review Comment:
`with_column_renamed` doesn’t validate `old_name`/`new_name` types before doing `old_name not in columns`. If a non-str is passed (e.g. `None`/`1`), this will raise a misleading `KeyError` instead of a `TypeError`, and it also diverges from `drop()` which type-checks its column-name arguments first.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): add DataFrame.mutate and rename [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu merged PR #1010:
URL: https://github.com/apache/sedona-db/pull/1010
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): add DataFrame.mutate (rename to follow) [sedona-db]
Posted by "paleolimbot (via GitHub)" <gi...@apache.org>.
paleolimbot commented on code in PR #1010:
URL: https://github.com/apache/sedona-db/pull/1010#discussion_r3500020897
##########
python/sedonadb/python/sedonadb/dataframe.py:
##########
@@ -1067,6 +1067,61 @@ def except_distinct(self, other: "DataFrame") -> "DataFrame":
self._check_set_op_compatible(other, "except_distinct")
return DataFrame(self._ctx, self._impl.except_distinct(other._impl))
+ def mutate(self, **exprs: Union[Expr, str, _SedonaLit]) -> "DataFrame":
Review Comment:
Should this accept positional expressions as well like `select()`?
##########
python/sedonadb/python/sedonadb/dataframe.py:
##########
@@ -1067,6 +1067,61 @@ def except_distinct(self, other: "DataFrame") -> "DataFrame":
self._check_set_op_compatible(other, "except_distinct")
return DataFrame(self._ctx, self._impl.except_distinct(other._impl))
+ def mutate(self, **exprs: Union[Expr, str, _SedonaLit]) -> "DataFrame":
+ """Add or replace columns, keeping all existing ones.
+
+ Each keyword adds a column named by the keyword and computed from
+ its value (Ibis / dplyr `mutate`). A keyword matching an existing
+ column replaces it in place; new columns are appended in the order
+ given. Unlike `select`, you don't re-list the columns you want to
+ keep — `mutate` keeps them all.
+
+ Args:
+ **exprs: Keyword column definitions. Each value is an `Expr`, a
+ column-name `str` (copied), or a `lit()` literal. At least
+ one is required. (Column names that aren't valid Python
+ identifiers can be produced via `select`.)
+
+ Examples:
+
+ >>> sd = sedona.db.connect()
+ >>> df = sd.sql("SELECT 1 AS a, 2 AS b")
+ >>> df.mutate(c=df.a + df.b, b=df.b * 10).show()
+ ┌───────┬───────┬───────┐
+ │ a ┆ b ┆ c │
+ │ int64 ┆ int64 ┆ int64 │
+ ╞═══════╪═══════╪═══════╡
+ │ 1 ┆ 20 ┆ 3 │
+ └───────┴───────┴───────┘
+ """
+ if not exprs:
+ raise ValueError("mutate() requires at least one keyword column")
+
+ # Coerce each value to an Expr aliased to its target name, mirroring
+ # how select(**kwargs) handles keyword columns.
+ mutations = {}
+ for name, e in exprs.items():
+ if isinstance(e, Expr):
+ mutations[name] = e.alias(name)
+ elif isinstance(e, str):
+ mutations[name] = _col(e).alias(name)
+ elif isinstance(e, _SedonaLit):
+ mutations[name] = _to_expr(e).alias(name)
+ else:
+ raise TypeError(
+ f"mutate() expects Expr, str, or Literal values, "
+ f"got {type(e).__name__} for '{name}'"
+ )
Review Comment:
Should there be a helper in the expressions file that collects named + positional expressions for `select()` and `mutate()` (which as far as I can tell share the same rules for input)?
##########
python/sedonadb/python/sedonadb/dataframe.py:
##########
@@ -1067,6 +1067,61 @@ def except_distinct(self, other: "DataFrame") -> "DataFrame":
self._check_set_op_compatible(other, "except_distinct")
return DataFrame(self._ctx, self._impl.except_distinct(other._impl))
+ def mutate(self, **exprs: Union[Expr, str, _SedonaLit]) -> "DataFrame":
+ """Add or replace columns, keeping all existing ones.
+
+ Each keyword adds a column named by the keyword and computed from
+ its value (Ibis / dplyr `mutate`). A keyword matching an existing
+ column replaces it in place; new columns are appended in the order
+ given. Unlike `select`, you don't re-list the columns you want to
+ keep — `mutate` keeps them all.
+
+ Args:
+ **exprs: Keyword column definitions. Each value is an `Expr`, a
+ column-name `str` (copied), or a `lit()` literal. At least
+ one is required. (Column names that aren't valid Python
+ identifiers can be produced via `select`.)
+
+ Examples:
+
+ >>> sd = sedona.db.connect()
+ >>> df = sd.sql("SELECT 1 AS a, 2 AS b")
+ >>> df.mutate(c=df.a + df.b, b=df.b * 10).show()
+ ┌───────┬───────┬───────┐
+ │ a ┆ b ┆ c │
+ │ int64 ┆ int64 ┆ int64 │
+ ╞═══════╪═══════╪═══════╡
+ │ 1 ┆ 20 ┆ 3 │
+ └───────┴───────┴───────┘
+ """
+ if not exprs:
+ raise ValueError("mutate() requires at least one keyword column")
+
+ # Coerce each value to an Expr aliased to its target name, mirroring
+ # how select(**kwargs) handles keyword columns.
+ mutations = {}
+ for name, e in exprs.items():
+ if isinstance(e, Expr):
+ mutations[name] = e.alias(name)
+ elif isinstance(e, str):
+ mutations[name] = _col(e).alias(name)
+ elif isinstance(e, _SedonaLit):
+ mutations[name] = _to_expr(e).alias(name)
+ else:
+ raise TypeError(
+ f"mutate() expects Expr, str, or Literal values, "
+ f"got {type(e).__name__} for '{name}'"
+ )
+
+ # Build a single projection: existing columns (replaced in place
+ # where a mutation matches the name), then the new columns appended.
+ # One projection node rather than a chain of per-column ones.
+ current = self._impl.columns()
+ projection = [mutations[c] if c in mutations else _col(c) for c in current]
+ projection += [m for name, m in mutations.items() if name not in current]
Review Comment:
The hiccup here is that existing columns should be updated in place rather than appended to the end of the projection (like `with_column()`).
It's a little cryptic how this works, but this is how we do it in R, although I forget if we specifically tested this case.
https://github.com/apache/sedona-db/blob/046fdede7e207c3d127f958af7b2612afe51a5e1/r/sdplyr/R/mutate.R#L35-L47
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): add DataFrame.mutate and rename [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu commented on PR #1010:
URL: https://github.com/apache/sedona-db/pull/1010#issuecomment-4890628481
Reworked per the discussion — this PR is now `mutate` + `rename` (the single-column `with_column`/`with_column_renamed` and their Rust bindings are gone):
- **`mutate(*exprs, **named)`** — positional + keyword, same input rules as `select` (shared `collect_exprs()` helper in `expr/expression.py`), in-place replacement by output name, one projection.
- **`rename(new_name="old_name")`** — Ibis/dplyr direction. A positionally-passed dict (the pandas/Polars `{old: new}` habit) raises a clear error that prints the corrected keyword call, e.g. `rename({"a": "x"})` → *"...Rewrite {'a': 'x'} as rename(x=\"a\")."* — actionable for hand-editing or an iterating LLM. Unknown old columns raise a `KeyError` listing available columns.
On the rename direction: went with your suggestion — Ibis-style kwargs as the one true form, dict detected-and-redirected rather than silently accepted. Easy to relax to "just do the rename" later if it proves annoying.
Local: mutate/rename tests + select regression (shared-helper refactor) + doctests + ruff + cargo fmt + clippy all clean.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): add DataFrame.mutate and rename [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu commented on code in PR #1010:
URL: https://github.com/apache/sedona-db/pull/1010#discussion_r3527379794
##########
python/sedonadb/python/sedonadb/dataframe.py:
##########
@@ -1067,6 +1067,61 @@ def except_distinct(self, other: "DataFrame") -> "DataFrame":
self._check_set_op_compatible(other, "except_distinct")
return DataFrame(self._ctx, self._impl.except_distinct(other._impl))
+ def mutate(self, **exprs: Union[Expr, str, _SedonaLit]) -> "DataFrame":
Review Comment:
Done — `mutate` now takes positional expressions too (`mutate(expr, y=...)`), with the same input rules as `select`. For a positional expr, its own output name is used to match existing columns for in-place replacement.
##########
python/sedonadb/python/sedonadb/dataframe.py:
##########
@@ -1067,6 +1067,61 @@ def except_distinct(self, other: "DataFrame") -> "DataFrame":
self._check_set_op_compatible(other, "except_distinct")
return DataFrame(self._ctx, self._impl.except_distinct(other._impl))
+ def mutate(self, **exprs: Union[Expr, str, _SedonaLit]) -> "DataFrame":
+ """Add or replace columns, keeping all existing ones.
+
+ Each keyword adds a column named by the keyword and computed from
+ its value (Ibis / dplyr `mutate`). A keyword matching an existing
+ column replaces it in place; new columns are appended in the order
+ given. Unlike `select`, you don't re-list the columns you want to
+ keep — `mutate` keeps them all.
+
+ Args:
+ **exprs: Keyword column definitions. Each value is an `Expr`, a
+ column-name `str` (copied), or a `lit()` literal. At least
+ one is required. (Column names that aren't valid Python
+ identifiers can be produced via `select`.)
+
+ Examples:
+
+ >>> sd = sedona.db.connect()
+ >>> df = sd.sql("SELECT 1 AS a, 2 AS b")
+ >>> df.mutate(c=df.a + df.b, b=df.b * 10).show()
+ ┌───────┬───────┬───────┐
+ │ a ┆ b ┆ c │
+ │ int64 ┆ int64 ┆ int64 │
+ ╞═══════╪═══════╪═══════╡
+ │ 1 ┆ 20 ┆ 3 │
+ └───────┴───────┴───────┘
+ """
+ if not exprs:
+ raise ValueError("mutate() requires at least one keyword column")
+
+ # Coerce each value to an Expr aliased to its target name, mirroring
+ # how select(**kwargs) handles keyword columns.
+ mutations = {}
+ for name, e in exprs.items():
+ if isinstance(e, Expr):
+ mutations[name] = e.alias(name)
+ elif isinstance(e, str):
+ mutations[name] = _col(e).alias(name)
+ elif isinstance(e, _SedonaLit):
+ mutations[name] = _to_expr(e).alias(name)
+ else:
+ raise TypeError(
+ f"mutate() expects Expr, str, or Literal values, "
+ f"got {type(e).__name__} for '{name}'"
+ )
Review Comment:
Added `collect_exprs()` in `expr/expression.py` — it collects named + positional expressions with one set of coercion rules, and both `select` and `mutate` now use it. (`select`'s inline coercion is gone.)
##########
python/sedonadb/python/sedonadb/dataframe.py:
##########
@@ -1067,6 +1067,61 @@ def except_distinct(self, other: "DataFrame") -> "DataFrame":
self._check_set_op_compatible(other, "except_distinct")
return DataFrame(self._ctx, self._impl.except_distinct(other._impl))
+ def mutate(self, **exprs: Union[Expr, str, _SedonaLit]) -> "DataFrame":
+ """Add or replace columns, keeping all existing ones.
+
+ Each keyword adds a column named by the keyword and computed from
+ its value (Ibis / dplyr `mutate`). A keyword matching an existing
+ column replaces it in place; new columns are appended in the order
+ given. Unlike `select`, you don't re-list the columns you want to
+ keep — `mutate` keeps them all.
+
+ Args:
+ **exprs: Keyword column definitions. Each value is an `Expr`, a
+ column-name `str` (copied), or a `lit()` literal. At least
+ one is required. (Column names that aren't valid Python
+ identifiers can be produced via `select`.)
+
+ Examples:
+
+ >>> sd = sedona.db.connect()
+ >>> df = sd.sql("SELECT 1 AS a, 2 AS b")
+ >>> df.mutate(c=df.a + df.b, b=df.b * 10).show()
+ ┌───────┬───────┬───────┐
+ │ a ┆ b ┆ c │
+ │ int64 ┆ int64 ┆ int64 │
+ ╞═══════╪═══════╪═══════╡
+ │ 1 ┆ 20 ┆ 3 │
+ └───────┴───────┴───────┘
+ """
+ if not exprs:
+ raise ValueError("mutate() requires at least one keyword column")
+
+ # Coerce each value to an Expr aliased to its target name, mirroring
+ # how select(**kwargs) handles keyword columns.
+ mutations = {}
+ for name, e in exprs.items():
+ if isinstance(e, Expr):
+ mutations[name] = e.alias(name)
+ elif isinstance(e, str):
+ mutations[name] = _col(e).alias(name)
+ elif isinstance(e, _SedonaLit):
+ mutations[name] = _to_expr(e).alias(name)
+ else:
+ raise TypeError(
+ f"mutate() expects Expr, str, or Literal values, "
+ f"got {type(e).__name__} for '{name}'"
+ )
+
+ # Build a single projection: existing columns (replaced in place
+ # where a mutation matches the name), then the new columns appended.
+ # One projection node rather than a chain of per-column ones.
+ current = self._impl.columns()
+ projection = [mutations[c] if c in mutations else _col(c) for c in current]
+ projection += [m for name, m in mutations.items() if name not in current]
Review Comment:
In-place replacement is handled: the projection walks the existing columns and substitutes a mutation where its output name matches (only genuinely-new columns append). Now that positional exprs are supported, the match uses each expr's output name (new `Expr._output_name()`, via DataFusion's `schema_name`) rather than just the kwarg key. Covered by `test_mutate_replaces_in_place`, `test_mutate_positional_expr`, and `test_mutate_preserves_column_order`.
##########
python/sedonadb/tests/expr/test_dataframe_with_column.py:
##########
@@ -0,0 +1,106 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pandas as pd
+import pandas.testing as pdt
+import pytest
+
+from sedonadb.dataframe import DataFrame
+from sedonadb.expr import col, lit
Review Comment:
Addressed — the with_column test file is gone; the new `test_dataframe_mutate.py` uses `con.col()`/`con.lit()` throughout.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] feat(python/sedonadb): add DataFrame.with_column and with_column_renamed [sedona-db]
Posted by "jiayuasu (via GitHub)" <gi...@apache.org>.
jiayuasu commented on code in PR #1010:
URL: https://github.com/apache/sedona-db/pull/1010#discussion_r3489285897
##########
python/sedonadb/python/sedonadb/dataframe.py:
##########
@@ -1067,6 +1067,89 @@ def except_distinct(self, other: "DataFrame") -> "DataFrame":
self._check_set_op_compatible(other, "except_distinct")
return DataFrame(self._ctx, self._impl.except_distinct(other._impl))
+ def with_column(self, name: str, expr: Union[Expr, str, _SedonaLit]) -> "DataFrame":
+ """Add a column, or replace an existing column of the same name.
+
+ Returns a new lazy DataFrame with all existing columns plus `name`
+ (Spark / Polars `with_column` naming). If `name` already exists it
+ is replaced in place; otherwise it is appended. To build several
+ columns at once, `select(...)` with keyword arguments is usually
+ more convenient.
+
+ Args:
+ name: The output column name.
+ expr: The column expression. An `Expr`, a column-name `str`
+ (copied as that column), or a `lit()` literal.
+
+ Examples:
+
+ >>> sd = sedona.db.connect()
+ >>> df = sd.sql("SELECT 1 AS a, 2 AS b")
+ >>> df.with_column("c", df.a + df.b).show()
+ ┌───────┬───────┬───────┐
+ │ a ┆ b ┆ c │
+ │ int64 ┆ int64 ┆ int64 │
+ ╞═══════╪═══════╪═══════╡
+ │ 1 ┆ 2 ┆ 3 │
+ └───────┴───────┴───────┘
+ """
+ if not isinstance(name, str):
+ raise TypeError(
+ f"with_column() name must be a str, got {type(name).__name__}"
+ )
+
+ if isinstance(expr, Expr):
+ inner_expr = expr._impl
+ elif isinstance(expr, str):
+ inner_expr = _col(expr)._impl
+ elif isinstance(expr, _SedonaLit):
+ inner_expr = _to_expr(expr)._impl
+ else:
+ raise TypeError(
+ f"with_column() expects an Expr, str, or Literal value, "
+ f"got {type(expr).__name__}"
+ )
+
+ return DataFrame(self._ctx, self._impl.with_column(name, inner_expr))
+
+ def with_column_renamed(self, old_name: str, new_name: str) -> "DataFrame":
+ """Rename a single column.
+
+ Returns a new lazy DataFrame with `old_name` renamed to `new_name`
+ (Spark naming); all other columns are unchanged. To rename several
+ columns, chain calls or use `select(*[col(k).alias(v) ...])`.
+
+ Args:
+ old_name: The existing column name.
+ new_name: The new name.
+
+ Raises:
+ KeyError: If `old_name` is not a column. (The message lists the
+ available columns.)
+
+ Examples:
+
+ >>> sd = sedona.db.connect()
+ >>> df = sd.sql("SELECT 1 AS a, 2 AS b")
+ >>> df.with_column_renamed("b", "c").show()
+ ┌───────┬───────┐
+ │ a ┆ c │
+ │ int64 ┆ int64 │
+ ╞═══════╪═══════╡
+ │ 1 ┆ 2 │
+ └───────┴───────┘
+ """
+ # DataFusion silently no-ops when the column is missing, which hides
+ # typos. Validate Python-side so the caller gets a clear KeyError
+ # listing the available columns (matching `drop`).
+ columns = self._impl.columns()
+ if old_name not in columns:
Review Comment:
Good catch — added type checks for both `old_name` and `new_name` (raising a clear `TypeError`) before the membership lookup, matching `drop()`'s pattern. Covered by a parametrized test (`None`/`int` for each argument).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: issues-unsubscribe@sedona.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org