You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@sedona.apache.org by ji...@apache.org on 2026/06/16 03:42:14 UTC
(sedona-db) branch main updated: feat(python/sedonadb): add DataFrame.distinct and distinct_on (#961)
This is an automated email from the ASF dual-hosted git repository.
jiayuasu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/sedona-db.git
The following commit(s) were added to refs/heads/main by this push:
new b691fb740 feat(python/sedonadb): add DataFrame.distinct and distinct_on (#961)
b691fb740 is described below
commit b691fb74004aef8676d044be4b616f107160db49
Author: Jia Yu <ji...@wherobots.com>
AuthorDate: Mon Jun 15 20:42:09 2026 -0700
feat(python/sedonadb): add DataFrame.distinct and distinct_on (#961)
---
python/sedonadb/python/sedonadb/dataframe.py | 71 +++++++++++++
python/sedonadb/src/dataframe.rs | 29 ++++++
.../sedonadb/tests/expr/test_dataframe_distinct.py | 111 +++++++++++++++++++++
3 files changed, 211 insertions(+)
diff --git a/python/sedonadb/python/sedonadb/dataframe.py b/python/sedonadb/python/sedonadb/dataframe.py
index 6d436aedb..7bace3e2f 100644
--- a/python/sedonadb/python/sedonadb/dataframe.py
+++ b/python/sedonadb/python/sedonadb/dataframe.py
@@ -834,6 +834,77 @@ class DataFrame:
)
return DataFrame(self._ctx, self._impl.cross_join(other._impl))
+ def distinct(self) -> "DataFrame":
+ """Remove duplicate rows, comparing all columns.
+
+ Returns a new lazy DataFrame keeping one row from each set of
+ rows that are equal across every column (SQL `SELECT DISTINCT`).
+
+ Examples:
+
+ >>> sd = sedona.db.connect()
+ >>> df = sd.sql(
+ ... "SELECT * FROM (VALUES (1), (1), (2)) AS t(x)"
+ ... )
+ >>> df.distinct().sort("x").show()
+ ┌───────┐
+ │ x │
+ │ int64 │
+ ╞═══════╡
+ │ 1 │
+ ├╌╌╌╌╌╌╌┤
+ │ 2 │
+ └───────┘
+ """
+ return DataFrame(self._ctx, self._impl.distinct())
+
+ def distinct_on(self, *cols: Union[str, Expr]) -> "DataFrame":
+ """Remove duplicate rows, comparing only the given key columns.
+
+ Keeps one row for each distinct combination of `cols` (SQL
+ `DISTINCT ON`) and returns all columns. Which row survives for each
+ key is unspecified and not controllable today — there is no row
+ ordering applied, and a preceding `.sort(...)` is not guaranteed to
+ be respected. For a deterministic result, ensure the non-key
+ columns are the same for all rows that share a key.
+
+ Args:
+ *cols: One or more key columns (`str` names or `Expr`). At
+ least one is required.
+
+ Examples:
+
+ >>> sd = sedona.db.connect()
+ >>> df = sd.sql(
+ ... "SELECT * FROM (VALUES (1, 'a'), (1, 'a'), (2, 'b')) AS t(k, v)"
+ ... )
+ >>> df.distinct_on("k").sort("k").show()
+ ┌───────┬──────┐
+ │ k ┆ v │
+ │ int64 ┆ utf8 │
+ ╞═══════╪══════╡
+ │ 1 ┆ a │
+ ├╌╌╌╌╌╌╌┼╌╌╌╌╌╌┤
+ │ 2 ┆ b │
+ └───────┴──────┘
+ """
+ if not cols:
+ raise ValueError("distinct_on() requires at least one column")
+
+ coerced = []
+ for c in cols:
+ if isinstance(c, Expr):
+ coerced.append(c._impl)
+ elif isinstance(c, str):
+ coerced.append(_col(c)._impl)
+ else:
+ raise TypeError(
+ f"distinct_on() expects str or Expr arguments, "
+ f"got {type(c).__name__}"
+ )
+
+ return DataFrame(self._ctx, self._impl.distinct_on(coerced))
+
def limit(self, n: Optional[int], /, *, offset: int = 0) -> "DataFrame":
"""Limit result to n rows starting at offset
diff --git a/python/sedonadb/src/dataframe.rs b/python/sedonadb/src/dataframe.rs
index cb7a71626..ddbf87cdb 100644
--- a/python/sedonadb/src/dataframe.rs
+++ b/python/sedonadb/src/dataframe.rs
@@ -296,6 +296,35 @@ impl InternalDataFrame {
Ok(InternalDataFrame::new(inner, self.runtime.clone()))
}
+ fn distinct(&self) -> Result<InternalDataFrame, PySedonaError> {
+ let inner = self.inner.clone().distinct()?;
+ Ok(InternalDataFrame::new(inner, self.runtime.clone()))
+ }
+
+ /// `DISTINCT ON (on_exprs)` keeping all columns. With no sort expression
+ /// DataFusion keeps an arbitrary row per distinct key (SQL `DISTINCT ON`
+ /// without `ORDER BY`).
+ fn distinct_on(&self, on_exprs: Vec<PyExpr>) -> Result<InternalDataFrame, PySedonaError> {
+ let on: Vec<Expr> = on_exprs.into_iter().map(|e| e.inner).collect();
+
+ // Output all columns: build qualified column refs for the full schema
+ // (mirrors `qualified_column_expr`).
+ let schema = self.inner.schema();
+ let select: Vec<Expr> = (0..schema.fields().len())
+ .map(|i| {
+ let (relation, field) = schema.qualified_field(i);
+ Expr::Column(Column {
+ relation: relation.cloned(),
+ name: field.name().to_string(),
+ spans: Default::default(),
+ })
+ })
+ .collect();
+
+ let inner = self.inner.clone().distinct_on(on, select, None)?;
+ Ok(InternalDataFrame::new(inner, self.runtime.clone()))
+ }
+
fn execute<'py>(&self, py: Python<'py>) -> Result<usize, PySedonaError> {
let df = self.inner.clone();
let count = wait_for_future(py, &self.runtime, async move {
diff --git a/python/sedonadb/tests/expr/test_dataframe_distinct.py b/python/sedonadb/tests/expr/test_dataframe_distinct.py
new file mode 100644
index 000000000..738e16e7e
--- /dev/null
+++ b/python/sedonadb/tests/expr/test_dataframe_distinct.py
@@ -0,0 +1,111 @@
+# 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
+
+
+def test_distinct_all_columns(con):
+ df = con.create_data_frame(
+ pd.DataFrame({"x": [1, 1, 2, 2, 2], "y": ["a", "a", "b", "b", "c"]})
+ )
+ out = df.distinct().sort("x", "y").to_pandas()
+ pdt.assert_frame_equal(
+ out,
+ pd.DataFrame({"x": [1, 2, 2], "y": ["a", "b", "c"]}),
+ )
+
+
+def test_distinct_no_duplicates_is_identity(con):
+ df = con.create_data_frame(pd.DataFrame({"x": [1, 2, 3]}))
+ out = df.distinct().sort("x").to_pandas()
+ pdt.assert_frame_equal(out, pd.DataFrame({"x": [1, 2, 3]}))
+
+
+def test_distinct_returns_lazy_dataframe(con):
+ df = con.create_data_frame(pd.DataFrame({"x": [1, 1]}))
+ assert isinstance(df.distinct(), DataFrame)
+
+
+def test_distinct_on_single_key_keeps_all_columns(con):
+ # Companion column is constant per key, so the arbitrary row kept per
+ # key is deterministic in value.
+ df = con.create_data_frame(
+ pd.DataFrame({"k": [1, 1, 2, 2], "v": ["a", "a", "b", "b"]})
+ )
+ out = df.distinct_on("k").sort("k").to_pandas()
+ pdt.assert_frame_equal(out, pd.DataFrame({"k": [1, 2], "v": ["a", "b"]}))
+
+
+def test_distinct_on_row_count_only(con):
+ # When the companion column varies within a key, which row survives is
+ # arbitrary (no ORDER BY), so only assert the distinct-key count.
+ df = con.create_data_frame(pd.DataFrame({"k": [1, 1, 2], "v": ["a", "b", "c"]}))
+ out = df.distinct_on("k").to_pandas()
+ assert len(out) == 2
+ assert sorted(out["k"].tolist()) == [1, 2]
+
+
+def test_distinct_on_multiple_keys(con):
+ df = con.create_data_frame(
+ pd.DataFrame(
+ {
+ "k1": [1, 1, 1, 2],
+ "k2": ["a", "a", "b", "a"],
+ "v": [10, 10, 20, 30],
+ }
+ )
+ )
+ out = df.distinct_on("k1", "k2").sort("k1", "k2").to_pandas()
+ pdt.assert_frame_equal(
+ out,
+ pd.DataFrame({"k1": [1, 1, 2], "k2": ["a", "b", "a"], "v": [10, 20, 30]}),
+ )
+
+
+def test_distinct_on_accepts_expr(con):
+ df = con.create_data_frame(pd.DataFrame({"k": [1, 1, 2], "v": ["a", "a", "b"]}))
+ out = df.distinct_on(col("k")).sort("k").to_pandas()
+ pdt.assert_frame_equal(out, pd.DataFrame({"k": [1, 2], "v": ["a", "b"]}))
+
+
+def test_distinct_on_computed_expr(con):
+ # Dedup on a derived key: rows collapse by whether x > 2.
+ df = con.create_data_frame(pd.DataFrame({"x": [1, 3, 2, 4]}))
+ out = df.distinct_on(col("x") > 2).to_pandas()
+ assert len(out) == 2
+
+
+def test_distinct_on_returns_lazy_dataframe(con):
+ df = con.create_data_frame(pd.DataFrame({"k": [1, 1]}))
+ assert isinstance(df.distinct_on("k"), DataFrame)
+
+
+def test_distinct_on_empty_raises(con):
+ df = con.create_data_frame(pd.DataFrame({"k": [1]}))
+ with pytest.raises(ValueError, match="at least one column"):
+ df.distinct_on()
+
+
+def test_distinct_on_bad_type_raises(con):
+ df = con.create_data_frame(pd.DataFrame({"k": [1]}))
+ with pytest.raises(TypeError, match="distinct_on\\(\\) expects str or Expr"):
+ df.distinct_on(123)