You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2026/10/22 08:33:12 UTC
(spark) branch branch-4.x updated: [SPARK-60391][SQL] Add DataFrame API for BIN BY
This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new 198934b3b63b [SPARK-60391][SQL] Add DataFrame API for BIN BY
198934b3b63b is described below
commit 198934b3b63b8a5055c91c1d77a1cdf508426f60
Author: Nikolina Vraneš <ni...@databricks.com>
AuthorDate: Thu Oct 22 08:32:33 2026 +0000
[SPARK-60391][SQL] Add DataFrame API for BIN BY
### What changes were proposed in this pull request?
Add DataFrame APIs for the BIN BY relation operator:
- Flat Scala/Java `Dataset.binBy(...)` overloads for the required inputs, alignment, and output names.
- A PySpark `DataFrame.binBy(...)` method for classic Spark and Spark Connect.
- A Spark Connect relation with Scala/Python client and server support.
The Scala/Java and PySpark APIs use the same parameter names and are marked experimental.
Resolve column expressions through the normal analyzer before BIN BY validation so DataFrame column references retain their identity. Add API documentation, validation tests, and Connect protocol goldens.
JIRA: [SPARK-60391](https://issues.apache.org/jira/browse/SPARK-60391)
### Why are the changes needed?
DataFrame users need to split timestamp ranges into fixed-width bins and distribute measurements across the overlapping bins. These APIs let them compose BIN BY with other DataFrame operations without registering a view and writing a SQL BIN BY clause.
### Does this PR introduce _any_ user-facing change?
Yes. The new APIs expose BIN BY to Scala, Java, and PySpark users, including Spark Connect.
Scala:
```scala
import org.apache.spark.sql.functions.{col, expr}
df.binBy(
rangeStart = col("start"),
rangeEnd = col("end"),
binWidth = expr("INTERVAL 5 MINUTES"),
distributeUniform = Array(col("value")))
```
PySpark:
```python
from pyspark.sql import functions as sf
df.binBy("start", "end", sf.expr("INTERVAL 5 MINUTES"), ["value"])
```
The result scales the selected FLOAT/DOUBLE values by each bin's overlap ratio and appends `bin_start`, `bin_end`, and `bin_distribute_ratio`. Users can set `alignTo` and rename the appended columns with `binStartName`, `binEndName`, and `binDistributeRatioName`.
The range and distribution arguments must reference top-level columns. Project computed, cast, or nested values with `select` first.
The APIs use the existing internal, default-off `spark.sql.binByRelationOperator.enabled` gate. It must be enabled in the session to run these examples.
### How was this patch tested?
Added coverage for SQL/API equivalence, optional settings, invalid inputs, bound and nested column references, and Spark Connect plan serialization.
After the rebase and replacement of the builder with the flat API, all 88 selected tests passed:
- `BinBySuite`: 34 tests.
- `JavaDataFrameSuite`, including `testBinBy`: 27 tests.
- BIN BY cases in `ClientDatasetSuite`, `PlanGenerationTestSuite`, `ProtoToParsedPlanTestSuite`, and `SparkConnectProtoSuite`: 6 tests.
- PySpark BIN BY tests in classic, Connect parity, and plan-only suites: 21 tests.
`ResolveBinBySuite` and the BIN BY SQL result and analyzer goldens in `SQLQueryTestSuite` also passed before the rebase. Module-scoped Scala formatting checks, Python Ruff checks, and `git diff --check` passed for the flat API changes.
### Was this patch authored or co-authored using generative AI tooling?
Generated-by: OpenAI Codex CLI 0.163.0.
Closes #59648 from vranes/bin-by-dataframe-api.
Authored-by: Nikolina Vraneš <ni...@databricks.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
(cherry picked from commit 76a21482d6ca54d1587e0b0398bbf49a2799af6e)
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
project/MimaExcludes.scala | 4 +-
.../source/reference/pyspark.sql/dataframe.rst | 3 +-
python/pyspark/sql/classic/dataframe.py | 72 ++++-
python/pyspark/sql/connect/dataframe.py | 66 ++++
python/pyspark/sql/connect/plan.py | 48 +++
python/pyspark/sql/connect/proto/relations_pb2.py | 358 +++++++++++----------
python/pyspark/sql/connect/proto/relations_pb2.pyi | 141 ++++++++
python/pyspark/sql/dataframe.py | 79 +++++
.../pyspark/sql/tests/connect/test_connect_plan.py | 42 +++
.../sql/tests/connect/test_parity_dataframe.py | 34 ++
python/pyspark/sql/tests/test_dataframe.py | 264 ++++++++++++++-
.../main/scala/org/apache/spark/sql/Dataset.scala | 105 +++++-
.../spark/sql/catalyst/analysis/ResolveBinBy.scala | 44 +--
.../plans/logical/basicLogicalOperators.scala | 11 +-
.../sql/catalyst/analysis/ResolveBinBySuite.scala | 29 +-
.../apache/spark/sql/PlanGenerationTestSuite.scala | 33 ++
.../spark/sql/connect/ClientDatasetSuite.scala | 58 ++++
.../main/protobuf/spark/connect/relations.proto | 31 ++
.../org/apache/spark/sql/connect/Dataset.scala | 57 ++++
.../query-tests/explain-results/binBy.explain | 3 +
.../explain-results/binBy_all_options.explain | 3 +
.../test/resources/query-tests/queries/binBy.json | 179 +++++++++++
.../resources/query-tests/queries/binBy.proto.bin | Bin 0 -> 1480 bytes
.../query-tests/queries/binBy_all_options.json | 247 ++++++++++++++
.../queries/binBy_all_options.proto.bin | Bin 0 -> 2131 bytes
.../sql/connect/planner/SparkConnectPlanner.scala | 18 +-
.../sql/connect/ProtoToParsedPlanTestSuite.scala | 1 +
.../connect/planner/SparkConnectProtoSuite.scala | 54 ++++
.../org/apache/spark/sql/classic/Dataset.scala | 55 ++++
.../org/apache/spark/sql/JavaDataFrameSuite.java | 33 ++
.../sql-tests/analyzer-results/bin-by.sql.out | 44 +++
.../src/test/resources/sql-tests/inputs/bin-by.sql | 14 +
.../resources/sql-tests/results/bin-by.sql.out | 48 +++
.../scala/org/apache/spark/sql/BinBySuite.scala | 140 ++++++++
34 files changed, 2083 insertions(+), 235 deletions(-)
diff --git a/project/MimaExcludes.scala b/project/MimaExcludes.scala
index 473cc851c410..46c92fcb133d 100644
--- a/project/MimaExcludes.scala
+++ b/project/MimaExcludes.scala
@@ -43,7 +43,9 @@ object MimaExcludes {
"org.apache.spark.ml.regression.DecisionTreeRegressionModel.numLeave"),
// [SPARK-59154] Remove unused prediction variance helper after inlining its implementation.
ProblemFilters.exclude[DirectMissingMethodProblem](
- "org.apache.spark.ml.regression.DecisionTreeRegressionModel.predictVariance")
+ "org.apache.spark.ml.regression.DecisionTreeRegressionModel.predictVariance"),
+ // BIN BY DataFrame API.
+ ProblemFilters.exclude[ReversedMissingMethodProblem]("org.apache.spark.sql.Dataset.binBy")
)
// Exclude rules for 4.3.x from 4.2.0 (add 4.3-specific filters below as needed).
diff --git a/python/docs/source/reference/pyspark.sql/dataframe.rst b/python/docs/source/reference/pyspark.sql/dataframe.rst
index 22e88bb67d5e..43fca87b4e69 100644
--- a/python/docs/source/reference/pyspark.sql/dataframe.rst
+++ b/python/docs/source/reference/pyspark.sql/dataframe.rst
@@ -31,6 +31,7 @@ DataFrame
DataFrame.alias
DataFrame.approxQuantile
DataFrame.asTable
+ DataFrame.binBy
DataFrame.cache
DataFrame.checkpoint
DataFrame.coalesce
@@ -191,4 +192,4 @@ plotting functions via the ``PySparkPlotAccessor``. Users can call specific plot
PySparkPlotAccessor.scatter
PySparkPlotAccessor.box
PySparkPlotAccessor.kde
- PySparkPlotAccessor.hist
\ No newline at end of file
+ PySparkPlotAccessor.hist
diff --git a/python/pyspark/sql/classic/dataframe.py b/python/pyspark/sql/classic/dataframe.py
index ddef85e3008e..6b47a23f5135 100644
--- a/python/pyspark/sql/classic/dataframe.py
+++ b/python/pyspark/sql/classic/dataframe.py
@@ -49,7 +49,7 @@ from pyspark.errors import (
)
from pyspark.resource import ResourceProfile
from pyspark.serializers import BatchedSerializer, CPickleSerializer, UTF8Deserializer
-from pyspark.sql.classic.column import _to_java_column, _to_list, _to_seq
+from pyspark.sql.classic.column import _to_java_column, _to_java_column_opt, _to_list, _to_seq
from pyspark.sql.column import Column
from pyspark.sql.dataframe import (
DataFrame as ParentDataFrame,
@@ -1219,6 +1219,76 @@ class DataFrame(ParentDataFrame, PandasMapOpsMixin, PandasConversionMixin):
jgd = self._jdf.groupingSets(jgrouping_sets, self._jcols_ordinal(*cols))
return GroupedData(jgd, self)
+ def binBy(
+ self,
+ rangeStart: "ColumnOrName",
+ rangeEnd: "ColumnOrName",
+ binWidth: Column,
+ distributeUniform: Union[List["ColumnOrName"], Tuple["ColumnOrName", ...]],
+ *,
+ alignTo: Optional[Column] = None,
+ binStartName: Optional[str] = None,
+ binEndName: Optional[str] = None,
+ binDistributeRatioName: Optional[str] = None,
+ ) -> ParentDataFrame:
+ expressions = {"binWidth": binWidth}
+ if alignTo is not None:
+ expressions["alignTo"] = alignTo
+ for arg_name, expression in expressions.items():
+ if not isinstance(expression, Column):
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": arg_name,
+ "expected_type": "Column",
+ "arg_type": type(expression).__name__,
+ },
+ )
+
+ if not isinstance(distributeUniform, (list, tuple)):
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": "distributeUniform",
+ "expected_type": "list or tuple",
+ "arg_type": type(distributeUniform).__name__,
+ },
+ )
+ aliases = {
+ "binStartName": binStartName,
+ "binEndName": binEndName,
+ "binDistributeRatioName": binDistributeRatioName,
+ }
+ for arg_name, name in aliases.items():
+ if name is not None and not isinstance(name, str):
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": arg_name,
+ "expected_type": "str",
+ "arg_type": type(name).__name__,
+ },
+ )
+
+ assert self._sc._gateway is not None
+ assert self._sc._jvm is not None
+ columns = to_java_array(
+ self._sc._gateway,
+ self._sc._jvm.org.apache.spark.sql.Column,
+ [_to_java_column(column) for column in distributeUniform],
+ )
+ jdf = self._jdf.binBy(
+ _to_java_column(rangeStart),
+ _to_java_column(rangeEnd),
+ _to_java_column(binWidth),
+ columns,
+ _to_java_column_opt(alignTo),
+ binStartName,
+ binEndName,
+ binDistributeRatioName,
+ )
+ return DataFrame(jdf, self.sparkSession)
+
def unpivot(
self,
ids: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]],
diff --git a/python/pyspark/sql/connect/dataframe.py b/python/pyspark/sql/connect/dataframe.py
index 18b3ed2c359a..cb25a25a3e62 100644
--- a/python/pyspark/sql/connect/dataframe.py
+++ b/python/pyspark/sql/connect/dataframe.py
@@ -999,6 +999,72 @@ class DataFrame(ParentDataFrame):
session=self._session,
)
+ def binBy(
+ self,
+ rangeStart: "ColumnOrName",
+ rangeEnd: "ColumnOrName",
+ binWidth: Column,
+ distributeUniform: Union[List["ColumnOrName"], Tuple["ColumnOrName", ...]],
+ *,
+ alignTo: Optional[Column] = None,
+ binStartName: Optional[str] = None,
+ binEndName: Optional[str] = None,
+ binDistributeRatioName: Optional[str] = None,
+ ) -> ParentDataFrame:
+ expressions = {"binWidth": binWidth}
+ if alignTo is not None:
+ expressions["alignTo"] = alignTo
+ for arg_name, expression in expressions.items():
+ if not isinstance(expression, Column):
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": arg_name,
+ "expected_type": "Column",
+ "arg_type": type(expression).__name__,
+ },
+ )
+
+ if not isinstance(distributeUniform, (list, tuple)):
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": "distributeUniform",
+ "expected_type": "list or tuple",
+ "arg_type": type(distributeUniform).__name__,
+ },
+ )
+ aliases = {
+ "binStartName": binStartName,
+ "binEndName": binEndName,
+ "binDistributeRatioName": binDistributeRatioName,
+ }
+ for arg_name, name in aliases.items():
+ if name is not None and not isinstance(name, str):
+ raise PySparkTypeError(
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": arg_name,
+ "expected_type": "str",
+ "arg_type": type(name).__name__,
+ },
+ )
+
+ return DataFrame(
+ plan.BinBy(
+ child=self._plan,
+ range_start=F._to_col(rangeStart),
+ range_end=F._to_col(rangeEnd),
+ bin_width=F._to_col(binWidth),
+ distribute_uniform=[F._to_col(col) for col in distributeUniform],
+ align_to=F._to_col(alignTo) if alignTo is not None else None,
+ bin_start_name=binStartName,
+ bin_end_name=binEndName,
+ bin_distribute_ratio_name=binDistributeRatioName,
+ ),
+ session=self._session,
+ )
+
def unpivot(
self,
ids: Union["ColumnOrName", List["ColumnOrName"], Tuple["ColumnOrName", ...]],
diff --git a/python/pyspark/sql/connect/plan.py b/python/pyspark/sql/connect/plan.py
index 297f20e2de05..351f1be4f716 100644
--- a/python/pyspark/sql/connect/plan.py
+++ b/python/pyspark/sql/connect/plan.py
@@ -1742,6 +1742,54 @@ class WithColumnsRenamed(LogicalPlan):
return plan
+class BinBy(LogicalPlan):
+ """Logical plan for splitting timestamp ranges into fixed-width bins."""
+
+ def __init__(
+ self,
+ child: Optional[LogicalPlan],
+ range_start: Column,
+ range_end: Column,
+ bin_width: Column,
+ distribute_uniform: Sequence[Column],
+ align_to: Optional[Column] = None,
+ bin_start_name: Optional[str] = None,
+ bin_end_name: Optional[str] = None,
+ bin_distribute_ratio_name: Optional[str] = None,
+ ) -> None:
+ columns = [range_start, range_end, bin_width, *distribute_uniform]
+ if align_to is not None:
+ columns.append(align_to)
+ super().__init__(child, self._collect_references(columns))
+ self.range_start = range_start
+ self.range_end = range_end
+ self.bin_width = bin_width
+ self.distribute_uniform = distribute_uniform
+ self.align_to = align_to
+ self.bin_start_name = bin_start_name
+ self.bin_end_name = bin_end_name
+ self.bin_distribute_ratio_name = bin_distribute_ratio_name
+
+ def plan(self, session: "SparkConnectClient") -> proto.Relation:
+ assert self._child is not None
+ plan = self._create_proto_relation()
+ node = plan.bin_by
+ node.input.CopyFrom(self._child.plan(session))
+ node.range_start.CopyFrom(self.range_start.to_plan(session))
+ node.range_end.CopyFrom(self.range_end.to_plan(session))
+ node.bin_width.CopyFrom(self.bin_width.to_plan(session))
+ node.distribute_uniform.extend(col.to_plan(session) for col in self.distribute_uniform)
+ if self.align_to is not None:
+ node.align_to.CopyFrom(self.align_to.to_plan(session))
+ if self.bin_start_name is not None:
+ node.bin_start_name = self.bin_start_name
+ if self.bin_end_name is not None:
+ node.bin_end_name = self.bin_end_name
+ if self.bin_distribute_ratio_name is not None:
+ node.bin_distribute_ratio_name = self.bin_distribute_ratio_name
+ return self._with_relations(plan, session)
+
+
class Unpivot(LogicalPlan):
"""Logical plan object for a unpivot operation."""
diff --git a/python/pyspark/sql/connect/proto/relations_pb2.py b/python/pyspark/sql/connect/proto/relations_pb2.py
index 8b56455e69f2..18b0ec53cb69 100644
--- a/python/pyspark/sql/connect/proto/relations_pb2.py
+++ b/python/pyspark/sql/connect/proto/relations_pb2.py
@@ -44,7 +44,7 @@ from pyspark.sql.connect.proto import ml_common_pb2 as spark_dot_connect_dot_ml_
DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
- b'\n\x1dspark/connect/relations.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1fspark/connect/expressions.proto\x1a\x19spark/connect/types.proto\x1a\x1bspark/connect/catalog.proto\x1a\x1aspark/connect/common.proto\x1a\x1dspark/connect/ml_common.proto"\xc9\x1f\n\x08Relation\x12\x35\n\x06\x63ommon\x18\x01 \x01(\x0b\x32\x1d.spark.connect.RelationCommonR\x06\x63ommon\x12)\n\x04read\x18\x02 \x01(\x0b\x32\x13.spark.connect.ReadH\x00R\x04read\x12\x32\n\x07project\x18\x03 \x [...]
+ b'\n\x1dspark/connect/relations.proto\x12\rspark.connect\x1a\x19google/protobuf/any.proto\x1a\x1fspark/connect/expressions.proto\x1a\x19spark/connect/types.proto\x1a\x1bspark/connect/catalog.proto\x1a\x1aspark/connect/common.proto\x1a\x1dspark/connect/ml_common.proto"\xf8\x1f\n\x08Relation\x12\x35\n\x06\x63ommon\x18\x01 \x01(\x0b\x32\x1d.spark.connect.RelationCommonR\x06\x63ommon\x12)\n\x04read\x18\x02 \x01(\x0b\x32\x13.spark.connect.ReadH\x00R\x04read\x12\x32\n\x07project\x18\x03 \x [...]
)
_globals = globals()
@@ -82,181 +82,183 @@ if not _descriptor._USE_C_DESCRIPTORS:
_globals["_PARSE_OPTIONSENTRY"]._loaded_options = None
_globals["_PARSE_OPTIONSENTRY"]._serialized_options = b"8\001"
_globals["_RELATION"]._serialized_start = 224
- _globals["_RELATION"]._serialized_end = 4265
- _globals["_MLRELATION"]._serialized_start = 4268
- _globals["_MLRELATION"]._serialized_end = 4752
- _globals["_MLRELATION_TRANSFORM"]._serialized_start = 4480
- _globals["_MLRELATION_TRANSFORM"]._serialized_end = 4715
- _globals["_FETCH"]._serialized_start = 4755
- _globals["_FETCH"]._serialized_end = 5086
- _globals["_FETCH_METHOD"]._serialized_start = 4871
- _globals["_FETCH_METHOD"]._serialized_end = 5086
- _globals["_FETCH_METHOD_ARGS"]._serialized_start = 4959
- _globals["_FETCH_METHOD_ARGS"]._serialized_end = 5086
- _globals["_UNKNOWN"]._serialized_start = 5088
- _globals["_UNKNOWN"]._serialized_end = 5097
- _globals["_RELATIONCOMMON"]._serialized_start = 5100
- _globals["_RELATIONCOMMON"]._serialized_end = 5242
- _globals["_SQL"]._serialized_start = 5245
- _globals["_SQL"]._serialized_end = 5723
- _globals["_SQL_ARGSENTRY"]._serialized_start = 5539
- _globals["_SQL_ARGSENTRY"]._serialized_end = 5629
- _globals["_SQL_NAMEDARGUMENTSENTRY"]._serialized_start = 5631
- _globals["_SQL_NAMEDARGUMENTSENTRY"]._serialized_end = 5723
- _globals["_WITHRELATIONS"]._serialized_start = 5725
- _globals["_WITHRELATIONS"]._serialized_end = 5842
- _globals["_READ"]._serialized_start = 5845
- _globals["_READ"]._serialized_end = 6562
- _globals["_READ_NAMEDTABLE"]._serialized_start = 6023
- _globals["_READ_NAMEDTABLE"]._serialized_end = 6215
- _globals["_READ_NAMEDTABLE_OPTIONSENTRY"]._serialized_start = 6157
- _globals["_READ_NAMEDTABLE_OPTIONSENTRY"]._serialized_end = 6215
- _globals["_READ_DATASOURCE"]._serialized_start = 6218
- _globals["_READ_DATASOURCE"]._serialized_end = 6549
- _globals["_READ_DATASOURCE_OPTIONSENTRY"]._serialized_start = 6157
- _globals["_READ_DATASOURCE_OPTIONSENTRY"]._serialized_end = 6215
- _globals["_RELATIONCHANGES"]._serialized_start = 6565
- _globals["_RELATIONCHANGES"]._serialized_end = 6797
- _globals["_RELATIONCHANGES_OPTIONSENTRY"]._serialized_start = 6157
- _globals["_RELATIONCHANGES_OPTIONSENTRY"]._serialized_end = 6215
- _globals["_PROJECT"]._serialized_start = 6799
- _globals["_PROJECT"]._serialized_end = 6916
- _globals["_FILTER"]._serialized_start = 6918
- _globals["_FILTER"]._serialized_end = 7030
- _globals["_JOIN"]._serialized_start = 7033
- _globals["_JOIN"]._serialized_end = 7694
- _globals["_JOIN_JOINDATATYPE"]._serialized_start = 7372
- _globals["_JOIN_JOINDATATYPE"]._serialized_end = 7464
- _globals["_JOIN_JOINTYPE"]._serialized_start = 7467
- _globals["_JOIN_JOINTYPE"]._serialized_end = 7675
- _globals["_SETOPERATION"]._serialized_start = 7697
- _globals["_SETOPERATION"]._serialized_end = 8176
- _globals["_SETOPERATION_SETOPTYPE"]._serialized_start = 8013
- _globals["_SETOPERATION_SETOPTYPE"]._serialized_end = 8127
- _globals["_LIMIT"]._serialized_start = 8178
- _globals["_LIMIT"]._serialized_end = 8254
- _globals["_OFFSET"]._serialized_start = 8256
- _globals["_OFFSET"]._serialized_end = 8335
- _globals["_TAIL"]._serialized_start = 8337
- _globals["_TAIL"]._serialized_end = 8412
- _globals["_AGGREGATE"]._serialized_start = 8415
- _globals["_AGGREGATE"]._serialized_end = 9181
- _globals["_AGGREGATE_PIVOT"]._serialized_start = 8830
- _globals["_AGGREGATE_PIVOT"]._serialized_end = 8941
- _globals["_AGGREGATE_GROUPINGSETS"]._serialized_start = 8943
- _globals["_AGGREGATE_GROUPINGSETS"]._serialized_end = 9019
- _globals["_AGGREGATE_GROUPTYPE"]._serialized_start = 9022
- _globals["_AGGREGATE_GROUPTYPE"]._serialized_end = 9181
- _globals["_SORT"]._serialized_start = 9184
- _globals["_SORT"]._serialized_end = 9344
- _globals["_DROP"]._serialized_start = 9347
- _globals["_DROP"]._serialized_end = 9488
- _globals["_DEDUPLICATE"]._serialized_start = 9491
- _globals["_DEDUPLICATE"]._serialized_end = 9731
- _globals["_LOCALRELATION"]._serialized_start = 9733
- _globals["_LOCALRELATION"]._serialized_end = 9822
- _globals["_CACHEDLOCALRELATION"]._serialized_start = 9824
- _globals["_CACHEDLOCALRELATION"]._serialized_end = 9896
- _globals["_CHUNKEDCACHEDLOCALRELATION"]._serialized_start = 9898
- _globals["_CHUNKEDCACHEDLOCALRELATION"]._serialized_end = 10010
- _globals["_CACHEDREMOTERELATION"]._serialized_start = 10012
- _globals["_CACHEDREMOTERELATION"]._serialized_end = 10067
- _globals["_SAMPLE"]._serialized_start = 10070
- _globals["_SAMPLE"]._serialized_end = 10343
- _globals["_RANGE"]._serialized_start = 10346
- _globals["_RANGE"]._serialized_end = 10491
- _globals["_SUBQUERYALIAS"]._serialized_start = 10493
- _globals["_SUBQUERYALIAS"]._serialized_end = 10607
- _globals["_REPARTITION"]._serialized_start = 10610
- _globals["_REPARTITION"]._serialized_end = 10752
- _globals["_SHOWSTRING"]._serialized_start = 10755
- _globals["_SHOWSTRING"]._serialized_end = 10897
- _globals["_HTMLSTRING"]._serialized_start = 10899
- _globals["_HTMLSTRING"]._serialized_end = 11013
- _globals["_STATSUMMARY"]._serialized_start = 11015
- _globals["_STATSUMMARY"]._serialized_end = 11107
- _globals["_STATDESCRIBE"]._serialized_start = 11109
- _globals["_STATDESCRIBE"]._serialized_end = 11190
- _globals["_STATCROSSTAB"]._serialized_start = 11192
- _globals["_STATCROSSTAB"]._serialized_end = 11293
- _globals["_STATCOV"]._serialized_start = 11295
- _globals["_STATCOV"]._serialized_end = 11391
- _globals["_STATCORR"]._serialized_start = 11394
- _globals["_STATCORR"]._serialized_end = 11531
- _globals["_STATAPPROXQUANTILE"]._serialized_start = 11534
- _globals["_STATAPPROXQUANTILE"]._serialized_end = 11698
- _globals["_STATFREQITEMS"]._serialized_start = 11700
- _globals["_STATFREQITEMS"]._serialized_end = 11825
- _globals["_STATSAMPLEBY"]._serialized_start = 11828
- _globals["_STATSAMPLEBY"]._serialized_end = 12137
- _globals["_STATSAMPLEBY_FRACTION"]._serialized_start = 12029
- _globals["_STATSAMPLEBY_FRACTION"]._serialized_end = 12128
- _globals["_NAFILL"]._serialized_start = 12140
- _globals["_NAFILL"]._serialized_end = 12274
- _globals["_NADROP"]._serialized_start = 12277
- _globals["_NADROP"]._serialized_end = 12411
- _globals["_NAREPLACE"]._serialized_start = 12414
- _globals["_NAREPLACE"]._serialized_end = 12710
- _globals["_NAREPLACE_REPLACEMENT"]._serialized_start = 12569
- _globals["_NAREPLACE_REPLACEMENT"]._serialized_end = 12710
- _globals["_TODF"]._serialized_start = 12712
- _globals["_TODF"]._serialized_end = 12800
- _globals["_WITHCOLUMNSRENAMED"]._serialized_start = 12803
- _globals["_WITHCOLUMNSRENAMED"]._serialized_end = 13185
- _globals["_WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY"]._serialized_start = 13047
- _globals["_WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY"]._serialized_end = 13114
- _globals["_WITHCOLUMNSRENAMED_RENAME"]._serialized_start = 13116
- _globals["_WITHCOLUMNSRENAMED_RENAME"]._serialized_end = 13185
- _globals["_WITHCOLUMNS"]._serialized_start = 13187
- _globals["_WITHCOLUMNS"]._serialized_end = 13306
- _globals["_WITHWATERMARK"]._serialized_start = 13309
- _globals["_WITHWATERMARK"]._serialized_end = 13443
- _globals["_HINT"]._serialized_start = 13446
- _globals["_HINT"]._serialized_end = 13578
- _globals["_UNPIVOT"]._serialized_start = 13581
- _globals["_UNPIVOT"]._serialized_end = 13908
- _globals["_UNPIVOT_VALUES"]._serialized_start = 13838
- _globals["_UNPIVOT_VALUES"]._serialized_end = 13897
- _globals["_TRANSPOSE"]._serialized_start = 13910
- _globals["_TRANSPOSE"]._serialized_end = 14032
- _globals["_UNRESOLVEDTABLEVALUEDFUNCTION"]._serialized_start = 14034
- _globals["_UNRESOLVEDTABLEVALUEDFUNCTION"]._serialized_end = 14159
- _globals["_TOSCHEMA"]._serialized_start = 14161
- _globals["_TOSCHEMA"]._serialized_end = 14267
- _globals["_REPARTITIONBYEXPRESSION"]._serialized_start = 14270
- _globals["_REPARTITIONBYEXPRESSION"]._serialized_end = 14473
- _globals["_MAPPARTITIONS"]._serialized_start = 14476
- _globals["_MAPPARTITIONS"]._serialized_end = 14708
- _globals["_GROUPMAP"]._serialized_start = 14711
- _globals["_GROUPMAP"]._serialized_end = 15561
- _globals["_TRANSFORMWITHSTATEINFO"]._serialized_start = 15564
- _globals["_TRANSFORMWITHSTATEINFO"]._serialized_end = 15787
- _globals["_COGROUPMAP"]._serialized_start = 15790
- _globals["_COGROUPMAP"]._serialized_end = 16316
- _globals["_APPLYINPANDASWITHSTATE"]._serialized_start = 16319
- _globals["_APPLYINPANDASWITHSTATE"]._serialized_end = 16676
- _globals["_COMMONINLINEUSERDEFINEDTABLEFUNCTION"]._serialized_start = 16679
- _globals["_COMMONINLINEUSERDEFINEDTABLEFUNCTION"]._serialized_end = 16923
- _globals["_PYTHONUDTF"]._serialized_start = 16926
- _globals["_PYTHONUDTF"]._serialized_end = 17103
- _globals["_COMMONINLINEUSERDEFINEDDATASOURCE"]._serialized_start = 17106
- _globals["_COMMONINLINEUSERDEFINEDDATASOURCE"]._serialized_end = 17257
- _globals["_PYTHONDATASOURCE"]._serialized_start = 17259
- _globals["_PYTHONDATASOURCE"]._serialized_end = 17334
- _globals["_COLLECTMETRICS"]._serialized_start = 17337
- _globals["_COLLECTMETRICS"]._serialized_end = 17473
- _globals["_PARSE"]._serialized_start = 17476
- _globals["_PARSE"]._serialized_end = 17886
- _globals["_PARSE_OPTIONSENTRY"]._serialized_start = 6157
- _globals["_PARSE_OPTIONSENTRY"]._serialized_end = 6215
- _globals["_PARSE_PARSEFORMAT"]._serialized_start = 17765
- _globals["_PARSE_PARSEFORMAT"]._serialized_end = 17875
- _globals["_ASOFJOIN"]._serialized_start = 17889
- _globals["_ASOFJOIN"]._serialized_end = 18364
- _globals["_LATERALJOIN"]._serialized_start = 18367
- _globals["_LATERALJOIN"]._serialized_end = 18597
- _globals["_NEARESTBYJOIN"]._serialized_start = 18600
- _globals["_NEARESTBYJOIN"]._serialized_end = 18893
- _globals["_ZIP"]._serialized_start = 18895
- _globals["_ZIP"]._serialized_end = 18992
+ _globals["_RELATION"]._serialized_end = 4312
+ _globals["_MLRELATION"]._serialized_start = 4315
+ _globals["_MLRELATION"]._serialized_end = 4799
+ _globals["_MLRELATION_TRANSFORM"]._serialized_start = 4527
+ _globals["_MLRELATION_TRANSFORM"]._serialized_end = 4762
+ _globals["_FETCH"]._serialized_start = 4802
+ _globals["_FETCH"]._serialized_end = 5133
+ _globals["_FETCH_METHOD"]._serialized_start = 4918
+ _globals["_FETCH_METHOD"]._serialized_end = 5133
+ _globals["_FETCH_METHOD_ARGS"]._serialized_start = 5006
+ _globals["_FETCH_METHOD_ARGS"]._serialized_end = 5133
+ _globals["_UNKNOWN"]._serialized_start = 5135
+ _globals["_UNKNOWN"]._serialized_end = 5144
+ _globals["_RELATIONCOMMON"]._serialized_start = 5147
+ _globals["_RELATIONCOMMON"]._serialized_end = 5289
+ _globals["_SQL"]._serialized_start = 5292
+ _globals["_SQL"]._serialized_end = 5770
+ _globals["_SQL_ARGSENTRY"]._serialized_start = 5586
+ _globals["_SQL_ARGSENTRY"]._serialized_end = 5676
+ _globals["_SQL_NAMEDARGUMENTSENTRY"]._serialized_start = 5678
+ _globals["_SQL_NAMEDARGUMENTSENTRY"]._serialized_end = 5770
+ _globals["_WITHRELATIONS"]._serialized_start = 5772
+ _globals["_WITHRELATIONS"]._serialized_end = 5889
+ _globals["_READ"]._serialized_start = 5892
+ _globals["_READ"]._serialized_end = 6609
+ _globals["_READ_NAMEDTABLE"]._serialized_start = 6070
+ _globals["_READ_NAMEDTABLE"]._serialized_end = 6262
+ _globals["_READ_NAMEDTABLE_OPTIONSENTRY"]._serialized_start = 6204
+ _globals["_READ_NAMEDTABLE_OPTIONSENTRY"]._serialized_end = 6262
+ _globals["_READ_DATASOURCE"]._serialized_start = 6265
+ _globals["_READ_DATASOURCE"]._serialized_end = 6596
+ _globals["_READ_DATASOURCE_OPTIONSENTRY"]._serialized_start = 6204
+ _globals["_READ_DATASOURCE_OPTIONSENTRY"]._serialized_end = 6262
+ _globals["_RELATIONCHANGES"]._serialized_start = 6612
+ _globals["_RELATIONCHANGES"]._serialized_end = 6844
+ _globals["_RELATIONCHANGES_OPTIONSENTRY"]._serialized_start = 6204
+ _globals["_RELATIONCHANGES_OPTIONSENTRY"]._serialized_end = 6262
+ _globals["_PROJECT"]._serialized_start = 6846
+ _globals["_PROJECT"]._serialized_end = 6963
+ _globals["_FILTER"]._serialized_start = 6965
+ _globals["_FILTER"]._serialized_end = 7077
+ _globals["_JOIN"]._serialized_start = 7080
+ _globals["_JOIN"]._serialized_end = 7741
+ _globals["_JOIN_JOINDATATYPE"]._serialized_start = 7419
+ _globals["_JOIN_JOINDATATYPE"]._serialized_end = 7511
+ _globals["_JOIN_JOINTYPE"]._serialized_start = 7514
+ _globals["_JOIN_JOINTYPE"]._serialized_end = 7722
+ _globals["_SETOPERATION"]._serialized_start = 7744
+ _globals["_SETOPERATION"]._serialized_end = 8223
+ _globals["_SETOPERATION_SETOPTYPE"]._serialized_start = 8060
+ _globals["_SETOPERATION_SETOPTYPE"]._serialized_end = 8174
+ _globals["_LIMIT"]._serialized_start = 8225
+ _globals["_LIMIT"]._serialized_end = 8301
+ _globals["_OFFSET"]._serialized_start = 8303
+ _globals["_OFFSET"]._serialized_end = 8382
+ _globals["_TAIL"]._serialized_start = 8384
+ _globals["_TAIL"]._serialized_end = 8459
+ _globals["_AGGREGATE"]._serialized_start = 8462
+ _globals["_AGGREGATE"]._serialized_end = 9228
+ _globals["_AGGREGATE_PIVOT"]._serialized_start = 8877
+ _globals["_AGGREGATE_PIVOT"]._serialized_end = 8988
+ _globals["_AGGREGATE_GROUPINGSETS"]._serialized_start = 8990
+ _globals["_AGGREGATE_GROUPINGSETS"]._serialized_end = 9066
+ _globals["_AGGREGATE_GROUPTYPE"]._serialized_start = 9069
+ _globals["_AGGREGATE_GROUPTYPE"]._serialized_end = 9228
+ _globals["_SORT"]._serialized_start = 9231
+ _globals["_SORT"]._serialized_end = 9391
+ _globals["_DROP"]._serialized_start = 9394
+ _globals["_DROP"]._serialized_end = 9535
+ _globals["_DEDUPLICATE"]._serialized_start = 9538
+ _globals["_DEDUPLICATE"]._serialized_end = 9778
+ _globals["_LOCALRELATION"]._serialized_start = 9780
+ _globals["_LOCALRELATION"]._serialized_end = 9869
+ _globals["_CACHEDLOCALRELATION"]._serialized_start = 9871
+ _globals["_CACHEDLOCALRELATION"]._serialized_end = 9943
+ _globals["_CHUNKEDCACHEDLOCALRELATION"]._serialized_start = 9945
+ _globals["_CHUNKEDCACHEDLOCALRELATION"]._serialized_end = 10057
+ _globals["_CACHEDREMOTERELATION"]._serialized_start = 10059
+ _globals["_CACHEDREMOTERELATION"]._serialized_end = 10114
+ _globals["_SAMPLE"]._serialized_start = 10117
+ _globals["_SAMPLE"]._serialized_end = 10390
+ _globals["_RANGE"]._serialized_start = 10393
+ _globals["_RANGE"]._serialized_end = 10538
+ _globals["_SUBQUERYALIAS"]._serialized_start = 10540
+ _globals["_SUBQUERYALIAS"]._serialized_end = 10654
+ _globals["_REPARTITION"]._serialized_start = 10657
+ _globals["_REPARTITION"]._serialized_end = 10799
+ _globals["_SHOWSTRING"]._serialized_start = 10802
+ _globals["_SHOWSTRING"]._serialized_end = 10944
+ _globals["_HTMLSTRING"]._serialized_start = 10946
+ _globals["_HTMLSTRING"]._serialized_end = 11060
+ _globals["_STATSUMMARY"]._serialized_start = 11062
+ _globals["_STATSUMMARY"]._serialized_end = 11154
+ _globals["_STATDESCRIBE"]._serialized_start = 11156
+ _globals["_STATDESCRIBE"]._serialized_end = 11237
+ _globals["_STATCROSSTAB"]._serialized_start = 11239
+ _globals["_STATCROSSTAB"]._serialized_end = 11340
+ _globals["_STATCOV"]._serialized_start = 11342
+ _globals["_STATCOV"]._serialized_end = 11438
+ _globals["_STATCORR"]._serialized_start = 11441
+ _globals["_STATCORR"]._serialized_end = 11578
+ _globals["_STATAPPROXQUANTILE"]._serialized_start = 11581
+ _globals["_STATAPPROXQUANTILE"]._serialized_end = 11745
+ _globals["_STATFREQITEMS"]._serialized_start = 11747
+ _globals["_STATFREQITEMS"]._serialized_end = 11872
+ _globals["_STATSAMPLEBY"]._serialized_start = 11875
+ _globals["_STATSAMPLEBY"]._serialized_end = 12184
+ _globals["_STATSAMPLEBY_FRACTION"]._serialized_start = 12076
+ _globals["_STATSAMPLEBY_FRACTION"]._serialized_end = 12175
+ _globals["_NAFILL"]._serialized_start = 12187
+ _globals["_NAFILL"]._serialized_end = 12321
+ _globals["_NADROP"]._serialized_start = 12324
+ _globals["_NADROP"]._serialized_end = 12458
+ _globals["_NAREPLACE"]._serialized_start = 12461
+ _globals["_NAREPLACE"]._serialized_end = 12757
+ _globals["_NAREPLACE_REPLACEMENT"]._serialized_start = 12616
+ _globals["_NAREPLACE_REPLACEMENT"]._serialized_end = 12757
+ _globals["_TODF"]._serialized_start = 12759
+ _globals["_TODF"]._serialized_end = 12847
+ _globals["_WITHCOLUMNSRENAMED"]._serialized_start = 12850
+ _globals["_WITHCOLUMNSRENAMED"]._serialized_end = 13232
+ _globals["_WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY"]._serialized_start = 13094
+ _globals["_WITHCOLUMNSRENAMED_RENAMECOLUMNSMAPENTRY"]._serialized_end = 13161
+ _globals["_WITHCOLUMNSRENAMED_RENAME"]._serialized_start = 13163
+ _globals["_WITHCOLUMNSRENAMED_RENAME"]._serialized_end = 13232
+ _globals["_WITHCOLUMNS"]._serialized_start = 13234
+ _globals["_WITHCOLUMNS"]._serialized_end = 13353
+ _globals["_WITHWATERMARK"]._serialized_start = 13356
+ _globals["_WITHWATERMARK"]._serialized_end = 13490
+ _globals["_HINT"]._serialized_start = 13493
+ _globals["_HINT"]._serialized_end = 13625
+ _globals["_BINBY"]._serialized_start = 13628
+ _globals["_BINBY"]._serialized_end = 14194
+ _globals["_UNPIVOT"]._serialized_start = 14197
+ _globals["_UNPIVOT"]._serialized_end = 14524
+ _globals["_UNPIVOT_VALUES"]._serialized_start = 14454
+ _globals["_UNPIVOT_VALUES"]._serialized_end = 14513
+ _globals["_TRANSPOSE"]._serialized_start = 14526
+ _globals["_TRANSPOSE"]._serialized_end = 14648
+ _globals["_UNRESOLVEDTABLEVALUEDFUNCTION"]._serialized_start = 14650
+ _globals["_UNRESOLVEDTABLEVALUEDFUNCTION"]._serialized_end = 14775
+ _globals["_TOSCHEMA"]._serialized_start = 14777
+ _globals["_TOSCHEMA"]._serialized_end = 14883
+ _globals["_REPARTITIONBYEXPRESSION"]._serialized_start = 14886
+ _globals["_REPARTITIONBYEXPRESSION"]._serialized_end = 15089
+ _globals["_MAPPARTITIONS"]._serialized_start = 15092
+ _globals["_MAPPARTITIONS"]._serialized_end = 15324
+ _globals["_GROUPMAP"]._serialized_start = 15327
+ _globals["_GROUPMAP"]._serialized_end = 16177
+ _globals["_TRANSFORMWITHSTATEINFO"]._serialized_start = 16180
+ _globals["_TRANSFORMWITHSTATEINFO"]._serialized_end = 16403
+ _globals["_COGROUPMAP"]._serialized_start = 16406
+ _globals["_COGROUPMAP"]._serialized_end = 16932
+ _globals["_APPLYINPANDASWITHSTATE"]._serialized_start = 16935
+ _globals["_APPLYINPANDASWITHSTATE"]._serialized_end = 17292
+ _globals["_COMMONINLINEUSERDEFINEDTABLEFUNCTION"]._serialized_start = 17295
+ _globals["_COMMONINLINEUSERDEFINEDTABLEFUNCTION"]._serialized_end = 17539
+ _globals["_PYTHONUDTF"]._serialized_start = 17542
+ _globals["_PYTHONUDTF"]._serialized_end = 17719
+ _globals["_COMMONINLINEUSERDEFINEDDATASOURCE"]._serialized_start = 17722
+ _globals["_COMMONINLINEUSERDEFINEDDATASOURCE"]._serialized_end = 17873
+ _globals["_PYTHONDATASOURCE"]._serialized_start = 17875
+ _globals["_PYTHONDATASOURCE"]._serialized_end = 17950
+ _globals["_COLLECTMETRICS"]._serialized_start = 17953
+ _globals["_COLLECTMETRICS"]._serialized_end = 18089
+ _globals["_PARSE"]._serialized_start = 18092
+ _globals["_PARSE"]._serialized_end = 18502
+ _globals["_PARSE_OPTIONSENTRY"]._serialized_start = 6204
+ _globals["_PARSE_OPTIONSENTRY"]._serialized_end = 6262
+ _globals["_PARSE_PARSEFORMAT"]._serialized_start = 18381
+ _globals["_PARSE_PARSEFORMAT"]._serialized_end = 18491
+ _globals["_ASOFJOIN"]._serialized_start = 18505
+ _globals["_ASOFJOIN"]._serialized_end = 18980
+ _globals["_LATERALJOIN"]._serialized_start = 18983
+ _globals["_LATERALJOIN"]._serialized_end = 19213
+ _globals["_NEARESTBYJOIN"]._serialized_start = 19216
+ _globals["_NEARESTBYJOIN"]._serialized_end = 19509
+ _globals["_ZIP"]._serialized_start = 19511
+ _globals["_ZIP"]._serialized_end = 19608
# @@protoc_insertion_point(module_scope)
diff --git a/python/pyspark/sql/connect/proto/relations_pb2.pyi b/python/pyspark/sql/connect/proto/relations_pb2.pyi
index 2d17e88446d6..a286e560c0dc 100644
--- a/python/pyspark/sql/connect/proto/relations_pb2.pyi
+++ b/python/pyspark/sql/connect/proto/relations_pb2.pyi
@@ -113,6 +113,7 @@ class Relation(google.protobuf.message.Message):
RELATION_CHANGES_FIELD_NUMBER: builtins.int
NEAREST_BY_JOIN_FIELD_NUMBER: builtins.int
ZIP_FIELD_NUMBER: builtins.int
+ BIN_BY_FIELD_NUMBER: builtins.int
FILL_NA_FIELD_NUMBER: builtins.int
DROP_NA_FIELD_NUMBER: builtins.int
REPLACE_FIELD_NUMBER: builtins.int
@@ -229,6 +230,8 @@ class Relation(google.protobuf.message.Message):
@property
def zip(self) -> global___Zip: ...
@property
+ def bin_by(self) -> global___BinBy: ...
+ @property
def fill_na(self) -> global___NAFill:
"""NA functions"""
@property
@@ -318,6 +321,7 @@ class Relation(google.protobuf.message.Message):
relation_changes: global___RelationChanges | None = ...,
nearest_by_join: global___NearestByJoin | None = ...,
zip: global___Zip | None = ...,
+ bin_by: global___BinBy | None = ...,
fill_na: global___NAFill | None = ...,
drop_na: global___NADrop | None = ...,
replace: global___NAReplace | None = ...,
@@ -345,6 +349,8 @@ class Relation(google.protobuf.message.Message):
b"approx_quantile",
"as_of_join",
b"as_of_join",
+ "bin_by",
+ b"bin_by",
"cached_local_relation",
b"cached_local_relation",
"cached_remote_relation",
@@ -478,6 +484,8 @@ class Relation(google.protobuf.message.Message):
b"approx_quantile",
"as_of_join",
b"as_of_join",
+ "bin_by",
+ b"bin_by",
"cached_local_relation",
b"cached_local_relation",
"cached_remote_relation",
@@ -651,6 +659,7 @@ class Relation(google.protobuf.message.Message):
"relation_changes",
"nearest_by_join",
"zip",
+ "bin_by",
"fill_na",
"drop_na",
"replace",
@@ -3479,6 +3488,138 @@ class Hint(google.protobuf.message.Message):
global___Hint = Hint
+class BinBy(google.protobuf.message.Message):
+ """Splits timestamp ranges across fixed bins and scales selected FLOAT or DOUBLE columns."""
+
+ DESCRIPTOR: google.protobuf.descriptor.Descriptor
+
+ INPUT_FIELD_NUMBER: builtins.int
+ RANGE_START_FIELD_NUMBER: builtins.int
+ RANGE_END_FIELD_NUMBER: builtins.int
+ BIN_WIDTH_FIELD_NUMBER: builtins.int
+ ALIGN_TO_FIELD_NUMBER: builtins.int
+ DISTRIBUTE_UNIFORM_FIELD_NUMBER: builtins.int
+ BIN_START_NAME_FIELD_NUMBER: builtins.int
+ BIN_END_NAME_FIELD_NUMBER: builtins.int
+ BIN_DISTRIBUTE_RATIO_NAME_FIELD_NUMBER: builtins.int
+ @property
+ def input(self) -> global___Relation:
+ """(Required) The input relation."""
+ @property
+ def range_start(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression:
+ """(Required) Inclusive range start timestamp."""
+ @property
+ def range_end(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression:
+ """(Required) Exclusive range end, with the same timestamp type as range_start."""
+ @property
+ def bin_width(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression:
+ """(Required) A positive, constant day-time interval expression."""
+ @property
+ def align_to(self) -> pyspark.sql.connect.proto.expressions_pb2.Expression:
+ """(Optional) A constant origin with the same timestamp type as the range columns."""
+ @property
+ def distribute_uniform(
+ self,
+ ) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[
+ pyspark.sql.connect.proto.expressions_pb2.Expression
+ ]:
+ """(Required) At least one FLOAT or DOUBLE column to scale."""
+ bin_start_name: builtins.str
+ """(Optional) Name for the appended bin-start column. Defaults to bin_start."""
+ bin_end_name: builtins.str
+ """(Optional) Name for the appended bin-end column. Defaults to bin_end."""
+ bin_distribute_ratio_name: builtins.str
+ """(Optional) Name for the appended overlap-ratio column. Defaults to bin_distribute_ratio."""
+ def __init__(
+ self,
+ *,
+ input: global___Relation | None = ...,
+ range_start: pyspark.sql.connect.proto.expressions_pb2.Expression | None = ...,
+ range_end: pyspark.sql.connect.proto.expressions_pb2.Expression | None = ...,
+ bin_width: pyspark.sql.connect.proto.expressions_pb2.Expression | None = ...,
+ align_to: pyspark.sql.connect.proto.expressions_pb2.Expression | None = ...,
+ distribute_uniform: collections.abc.Iterable[
+ pyspark.sql.connect.proto.expressions_pb2.Expression
+ ]
+ | None = ...,
+ bin_start_name: builtins.str | None = ...,
+ bin_end_name: builtins.str | None = ...,
+ bin_distribute_ratio_name: builtins.str | None = ...,
+ ) -> None: ...
+ def HasField(
+ self,
+ field_name: typing_extensions.Literal[
+ "_bin_distribute_ratio_name",
+ b"_bin_distribute_ratio_name",
+ "_bin_end_name",
+ b"_bin_end_name",
+ "_bin_start_name",
+ b"_bin_start_name",
+ "align_to",
+ b"align_to",
+ "bin_distribute_ratio_name",
+ b"bin_distribute_ratio_name",
+ "bin_end_name",
+ b"bin_end_name",
+ "bin_start_name",
+ b"bin_start_name",
+ "bin_width",
+ b"bin_width",
+ "input",
+ b"input",
+ "range_end",
+ b"range_end",
+ "range_start",
+ b"range_start",
+ ],
+ ) -> builtins.bool: ...
+ def ClearField(
+ self,
+ field_name: typing_extensions.Literal[
+ "_bin_distribute_ratio_name",
+ b"_bin_distribute_ratio_name",
+ "_bin_end_name",
+ b"_bin_end_name",
+ "_bin_start_name",
+ b"_bin_start_name",
+ "align_to",
+ b"align_to",
+ "bin_distribute_ratio_name",
+ b"bin_distribute_ratio_name",
+ "bin_end_name",
+ b"bin_end_name",
+ "bin_start_name",
+ b"bin_start_name",
+ "bin_width",
+ b"bin_width",
+ "distribute_uniform",
+ b"distribute_uniform",
+ "input",
+ b"input",
+ "range_end",
+ b"range_end",
+ "range_start",
+ b"range_start",
+ ],
+ ) -> None: ...
+ @typing.overload
+ def WhichOneof(
+ self,
+ oneof_group: typing_extensions.Literal[
+ "_bin_distribute_ratio_name", b"_bin_distribute_ratio_name"
+ ],
+ ) -> typing_extensions.Literal["bin_distribute_ratio_name"] | None: ...
+ @typing.overload
+ def WhichOneof(
+ self, oneof_group: typing_extensions.Literal["_bin_end_name", b"_bin_end_name"]
+ ) -> typing_extensions.Literal["bin_end_name"] | None: ...
+ @typing.overload
+ def WhichOneof(
+ self, oneof_group: typing_extensions.Literal["_bin_start_name", b"_bin_start_name"]
+ ) -> typing_extensions.Literal["bin_start_name"] | None: ...
+
+global___BinBy = BinBy
+
class Unpivot(google.protobuf.message.Message):
"""Unpivot a DataFrame from wide format to long format, optionally leaving identifier columns set."""
diff --git a/python/pyspark/sql/dataframe.py b/python/pyspark/sql/dataframe.py
index d92290b8443e..908bfa5a5cd1 100644
--- a/python/pyspark/sql/dataframe.py
+++ b/python/pyspark/sql/dataframe.py
@@ -4343,6 +4343,85 @@ class DataFrame:
"""
...
+ @dispatch_df_method
+ def binBy(
+ self,
+ rangeStart: "ColumnOrName",
+ rangeEnd: "ColumnOrName",
+ binWidth: Column,
+ distributeUniform: Union[List["ColumnOrName"], Tuple["ColumnOrName", ...]],
+ *,
+ alignTo: Optional[Column] = None,
+ binStartName: Optional[str] = None,
+ binEndName: Optional[str] = None,
+ binDistributeRatioName: Optional[str] = None,
+ ) -> "DataFrame":
+ """Split timestamp ranges across fixed bins and proportionally scale selected columns.
+
+ .. versionadded:: 4.4.0
+
+ Supports Spark Connect.
+
+ Parameters
+ ----------
+ rangeStart : :class:`Column` or str
+ Top-level timestamp column for the inclusive range start.
+ rangeEnd : :class:`Column` or str
+ Top-level timestamp column for the exclusive range end.
+ Both range columns must have the same timestamp type.
+ binWidth : :class:`Column`
+ Positive, constant day-time interval expression,
+ for example ``sf.expr("INTERVAL 5 MINUTES")``.
+ distributeUniform : list or tuple
+ Nonempty collection of top-level FLOAT or DOUBLE columns to scale.
+ Each entry is a :class:`Column` or a column name.
+ alignTo : :class:`Column`, optional
+ Constant alignment origin with the same timestamp type as the range columns.
+ The default is 1970-01-01 00:00:00, in the session time zone for TIMESTAMP.
+ binStartName : str, optional
+ Name for the appended bin-start column. The default is ``bin_start``.
+ binEndName : str, optional
+ Name for the appended bin-end column. The default is ``bin_end``.
+ binDistributeRatioName : str, optional
+ Name for the appended overlap-ratio column.
+ The default is ``bin_distribute_ratio``.
+
+ Returns
+ -------
+ :class:`DataFrame`
+ Input columns with the selected values scaled, followed by the three bin columns.
+
+ Notes
+ -----
+ This API is experimental.
+
+ This operation requires ``spark.sql.binByRelationOperator.enabled`` to be true.
+ The range and distribution arguments must directly reference top-level columns.
+ Use :meth:`select` to project computed, cast, or nested values into columns first.
+
+ Each selected value must represent the change over its input range.
+ The ratio uses the real elapsed time that overlaps each bin.
+ TIMESTAMP bins use the session time zone for widths of at least one day.
+ TIMESTAMP_NTZ bins and sub-day widths use UTC arithmetic.
+
+ A zero-length range produces one row with a ratio of 1.0.
+ An inverted range raises an error.
+ A NULL range bound produces one row with NULL bin columns and NULL scaled values.
+
+ Examples
+ --------
+ >>> from pyspark.sql import functions as sf
+ >>> spark.conf.set("spark.sql.binByRelationOperator.enabled", "true")
+ >>> df = spark.sql("SELECT TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start, "
+ ... "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end, 100.0D AS value")
+ >>> df.binBy("start", "end", sf.expr("INTERVAL 5 MINUTES"), ["value"]).select(
+ ... "value", "bin_distribute_ratio"
+ ... ).collect()
+ [Row(value=50.0, bin_distribute_ratio=0.5), Row(value=50.0, bin_distribute_ratio=0.5)]
+ >>> spark.conf.unset("spark.sql.binByRelationOperator.enabled")
+ """
+ ...
+
@dispatch_df_method
def unpivot(
self,
diff --git a/python/pyspark/sql/tests/connect/test_connect_plan.py b/python/pyspark/sql/tests/connect/test_connect_plan.py
index 6ee4a8993dcb..76ada94870a1 100644
--- a/python/pyspark/sql/tests/connect/test_connect_plan.py
+++ b/python/pyspark/sql/tests/connect/test_connect_plan.py
@@ -108,6 +108,48 @@ class SparkConnectPlanTests(PlanOnlyTestFixture):
plan = self.connect.sql("SELECT 1")._plan.to_proto(self.connect)
self.assertEqual(plan.root.sql.query, "SELECT 1")
+ def test_bin_by(self):
+ df = self.connect.readTable(table_name=self.tbl_name)
+ result = df.binBy("start", "end", lit(datetime.timedelta(minutes=5)), ["value"])
+ plan = result._plan.to_proto(self.connect)
+ node = plan.root.bin_by
+ self.assertEqual(node.input.read.named_table.unparsed_identifier, self.tbl_name)
+ self.assertEqual(node.range_start.unresolved_attribute.unparsed_identifier, "start")
+ self.assertEqual(node.range_end.unresolved_attribute.unparsed_identifier, "end")
+ self.assertEqual(node.bin_width.literal.day_time_interval, 300000000)
+ self.assertEqual(len(node.distribute_uniform), 1)
+ self.assertEqual(
+ node.distribute_uniform[0].unresolved_attribute.unparsed_identifier, "value"
+ )
+ self.assertFalse(node.HasField("align_to"))
+ self.assertFalse(node.HasField("bin_start_name"))
+ self.assertFalse(node.HasField("bin_end_name"))
+ self.assertFalse(node.HasField("bin_distribute_ratio_name"))
+
+ def test_bin_by_optional_fields_and_bound_columns(self):
+ df = self.connect.readTable(table_name=self.tbl_name)
+ result = df.binBy(
+ df["start"],
+ df["end"],
+ lit(datetime.timedelta(minutes=5)),
+ (df["value"], df["value2"]),
+ alignTo=lit(None).cast("timestamp"),
+ binStartName="bin.start",
+ binEndName="bin`end",
+ binDistributeRatioName="",
+ )
+ plan = result._plan.to_proto(self.connect)
+ node = plan.root.bin_by
+ for expression in [node.range_start, node.range_end, *node.distribute_uniform]:
+ self.assertEqual(expression.unresolved_attribute.plan_id, df._plan._plan_id)
+ self.assertEqual(len(node.distribute_uniform), 2)
+ self.assertTrue(node.HasField("align_to"))
+ self.assertTrue(node.align_to.cast.expr.literal.HasField("null"))
+ self.assertEqual(node.bin_start_name, "bin.start")
+ self.assertEqual(node.bin_end_name, "bin`end")
+ self.assertTrue(node.HasField("bin_distribute_ratio_name"))
+ self.assertEqual(node.bin_distribute_ratio_name, "")
+
def test_simple_project(self):
plan = self.connect.readTable(table_name=self.tbl_name)._plan.to_proto(self.connect)
self.assertIsNotNone(plan.root, "Root relation must be set")
diff --git a/python/pyspark/sql/tests/connect/test_parity_dataframe.py b/python/pyspark/sql/tests/connect/test_parity_dataframe.py
index fe38894e0776..2da7ae77a7c5 100644
--- a/python/pyspark/sql/tests/connect/test_parity_dataframe.py
+++ b/python/pyspark/sql/tests/connect/test_parity_dataframe.py
@@ -17,11 +17,45 @@
import unittest
+from pyspark.errors import AnalysisException
+from pyspark.sql import Row, functions
from pyspark.sql.tests.test_dataframe import DataFrameTestsMixin
from pyspark.testing.connectutils import ReusedConnectTestCase
class DataFrameParityTests(DataFrameTestsMixin, ReusedConnectTestCase):
+ def test_bin_by_with_non_strict_column_resolution(self):
+ with self.sql_conf(
+ {
+ "spark.sql.binByRelationOperator.enabled": "true",
+ "spark.sql.analyzer.strictDataFrameColumnResolution": "false",
+ }
+ ):
+ df = self.spark.sql(
+ "SELECT TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start, "
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end, 100.0D AS value"
+ )
+ changed = (
+ df.withColumn("start", functions.expr("TIMESTAMP_NTZ '2024-01-01 00:02:00'"))
+ .withColumn("end", functions.expr("TIMESTAMP_NTZ '2024-01-01 00:18:00'"))
+ .withColumn("value", functions.lit(200.0))
+ )
+ result = changed.binBy(
+ df["start"], df["end"], functions.expr("INTERVAL 5 MINUTES"), [df["value"]]
+ )
+ self.assertEqual(
+ result.orderBy("bin_start").select("value").collect(),
+ [Row(37.5), Row(62.5), Row(62.5), Row(37.5)],
+ )
+ with self.assertRaises(AnalysisException) as error:
+ changed.binBy(
+ changed["start"],
+ changed["end"],
+ functions.expr("INTERVAL 5 MINUTES"),
+ [changed["missing"]],
+ ).collect()
+ self.assertEqual(error.exception.getCondition(), "CANNOT_RESOLVE_DATAFRAME_COLUMN")
+
def test_help_command(self):
df = self.spark.createDataFrame(data=[{"foo": "bar"}, {"foo": "baz"}])
super().check_help_command(df)
diff --git a/python/pyspark/sql/tests/test_dataframe.py b/python/pyspark/sql/tests/test_dataframe.py
index 0393cb1b6164..4111e732b84d 100644
--- a/python/pyspark/sql/tests/test_dataframe.py
+++ b/python/pyspark/sql/tests/test_dataframe.py
@@ -32,7 +32,7 @@ from pyspark.errors import (
PySparkValueError,
QueryContextType,
)
-from pyspark.sql import DataFrame, Row, functions
+from pyspark.sql import DataFrame, Row, functions, is_remote
from pyspark.sql.functions import (
array,
col,
@@ -65,6 +65,268 @@ from pyspark.testing.utils import (
class DataFrameTestsMixin:
+ def test_bin_by(self):
+ with self.sql_conf({"spark.sql.binByRelationOperator.enabled": "true"}):
+ df = self.spark.sql(
+ "SELECT TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start, "
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end, 100.0D AS value, 20.0F AS value2"
+ )
+ with self.temp_view("bin_by_input"):
+ df.createOrReplaceTempView("bin_by_input")
+ expected = self.spark.sql(
+ "SELECT * FROM bin_by_input BIN BY (RANGE start TO end "
+ "BIN WIDTH INTERVAL 5 MINUTES DISTRIBUTE UNIFORM (value, value2))"
+ )
+ for bound in [False, True]:
+ with self.subTest(bound_columns=bound):
+ result = df.binBy(
+ df["start"] if bound else "start",
+ df["end"] if bound else "end",
+ functions.expr("INTERVAL 5 MINUTES"),
+ (df["value"], df["value2"]) if bound else ["value", "value2"],
+ )
+ assertDataFrameEqual(result, expected)
+
+ def test_bin_by_origin_and_quoted_names(self):
+ with self.sql_conf({"spark.sql.binByRelationOperator.enabled": "true"}):
+ df = self.spark.sql(
+ "SELECT TIMESTAMP_NTZ '2024-01-01 00:00:00' AS `range.start`, "
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS `range``end`, 100.0D AS `v``.x`"
+ ).alias("metrics")
+ result = df.binBy(
+ "metrics.`range.start`",
+ col("metrics.`range``end`"),
+ functions.expr("INTERVAL 5 MINUTES"),
+ [df["`v``.x`"]],
+ alignTo=functions.expr("TIMESTAMP_NTZ '2024-01-01 00:02:00'"),
+ binStartName="bin.start",
+ binEndName="bin`end",
+ binDistributeRatioName="",
+ )
+ self.assertEqual(
+ result.columns, ["range.start", "range`end", "v`.x", "bin.start", "bin`end", ""]
+ )
+ self.assertEqual(
+ result.orderBy("`bin.start`").select("`v``.x`", "``").collect(),
+ [Row(20.0, 0.2), Row(50.0, 0.5), Row(30.0, 0.3)],
+ )
+
+ def test_bin_by_after_unpivot(self):
+ with self.sql_conf(
+ {
+ "spark.sql.binByRelationOperator.enabled": "true",
+ "spark.sql.analyzer.strictDataFrameColumnResolution": "false",
+ }
+ ):
+ df = self.spark.sql(
+ "SELECT TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start, "
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end, 100.0D AS value"
+ ).unpivot(["start", "end"], ["value"], "metric", "value")
+ result = df.binBy(
+ df["start"], df["end"], functions.expr("INTERVAL 5 MINUTES"), [df["value"]]
+ )
+ self.assertEqual(
+ result.select("metric", "value").collect(),
+ [Row("value", 50.0), Row("value", 50.0)],
+ )
+ unrelated = self.spark.sql("SELECT 200.0D AS value")
+ with self.assertRaises(AnalysisException) as error:
+ df.binBy(
+ df["start"],
+ df["end"],
+ functions.expr("INTERVAL 5 MINUTES"),
+ [unrelated["value"]],
+ ).collect()
+ self.assertEqual(
+ error.exception.getCondition(),
+ "CANNOT_RESOLVE_DATAFRAME_COLUMN"
+ if is_remote()
+ else "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_APPEAR_IN_OPERATION",
+ )
+
+ def test_bin_by_rejects_replaced_bound_columns(self):
+ with self.sql_conf(
+ {
+ "spark.sql.binByRelationOperator.enabled": "true",
+ "spark.sql.analyzer.strictDataFrameColumnResolution": "true",
+ }
+ ):
+ df = self.spark.sql(
+ "SELECT TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start, "
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end, 100.0D AS value"
+ )
+ width = functions.expr("INTERVAL 5 MINUTES")
+ condition = (
+ "CANNOT_RESOLVE_DATAFRAME_COLUMN"
+ if is_remote()
+ else "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_APPEAR_IN_OPERATION"
+ )
+ cases = [
+ ("start", functions.expr("TIMESTAMP_NTZ '2024-01-01 00:02:00'")),
+ ("end", functions.expr("TIMESTAMP_NTZ '2024-01-01 00:08:00'")),
+ ("value", lit(200.0)),
+ ]
+ changed = df.withColumns(dict(cases))
+ self.assertEqual(changed.binBy("start", "end", width, ["value"]).count(), 2)
+ self.assertEqual(
+ changed.binBy(changed["start"], changed["end"], width, [changed["value"]]).count(),
+ 2,
+ )
+ for name, replacement in cases:
+ with self.subTest(column=name):
+ changed = df.withColumn(name, replacement)
+ with self.assertRaises(AnalysisException) as error:
+ changed.binBy(df["start"], df["end"], width, [df["value"]]).collect()
+ self.assertEqual(error.exception.getCondition(), condition)
+
+ def test_bin_by_analysis_errors(self):
+ with self.sql_conf({"spark.sql.binByRelationOperator.enabled": "true"}):
+ df = self.spark.sql(
+ "SELECT TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start, "
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end, 100.0D AS value"
+ )
+ cases = [
+ ({"distributeUniform": []}, "BIN_BY_MISSING_DISTRIBUTE"),
+ (
+ {"distributeUniform": ["value", df["value"]]},
+ "BIN_BY_DUPLICATE_DISTRIBUTE_COLUMN",
+ ),
+ (
+ {"distributeUniform": [col("value") + 1]},
+ "BIN_BY_REQUIRES_TOP_LEVEL_COLUMN",
+ ),
+ (
+ {"distributeUniform": [col("missing") + 1]},
+ "UNRESOLVED_COLUMN.WITH_SUGGESTION",
+ ),
+ (
+ {"rangeStart": col("start").cast("timestamp")},
+ "BIN_BY_REQUIRES_TOP_LEVEL_COLUMN",
+ ),
+ (
+ {"alignTo": lit(None).cast("timestamp_ntz")},
+ "BIN_BY_NULL_ARGUMENT",
+ ),
+ (
+ {"binWidth": functions.expr("INTERVAL 0 MINUTES")},
+ "BIN_BY_NON_POSITIVE_BIN_WIDTH",
+ ),
+ ]
+ for overrides, condition in cases:
+ with self.subTest(condition=condition, overrides=overrides):
+ args = dict(
+ rangeStart="start",
+ rangeEnd="end",
+ binWidth=functions.expr("INTERVAL 5 MINUTES"),
+ distributeUniform=["value"],
+ )
+ args.update(overrides)
+ with self.assertRaises(AnalysisException) as error:
+ df.binBy(**args).collect()
+ self.assertEqual(error.exception.getCondition(), condition)
+
+ def test_bin_by_result_column_handles(self):
+ for strict in ["true", "false"]:
+ with (
+ self.subTest(strict=strict),
+ self.sql_conf(
+ {
+ "spark.sql.binByRelationOperator.enabled": "true",
+ "spark.sql.analyzer.strictDataFrameColumnResolution": strict,
+ }
+ ),
+ ):
+ df = self.spark.sql(
+ "SELECT TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start, "
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end, 100.0D AS value"
+ )
+ result = df.binBy("start", "end", functions.expr("INTERVAL 5 MINUTES"), ["value"])
+ assertDataFrameEqual(
+ result.select(result["bin_start"], result["value"]),
+ result.select("bin_start", "value"),
+ )
+ if is_remote() and strict == "false":
+ self.assertEqual(result.select(df["value"]).collect(), [Row(50.0), Row(50.0)])
+ else:
+ with self.assertRaises(AnalysisException) as error:
+ result.select(df["value"]).collect()
+ self.assertEqual(
+ error.exception.getCondition(),
+ "CANNOT_RESOLVE_DATAFRAME_COLUMN"
+ if is_remote()
+ else "MISSING_ATTRIBUTES.RESOLVED_ATTRIBUTE_APPEAR_IN_OPERATION",
+ )
+
+ def test_bin_by_rejects_nested_column(self):
+ with self.sql_conf(
+ {
+ "spark.sql.binByRelationOperator.enabled": "true",
+ "spark.sql.analyzer.strictDataFrameColumnResolution": "true",
+ }
+ ):
+ df = self.spark.sql(
+ "SELECT TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start, "
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end, 100.0D AS value"
+ ).select(struct("start").alias("s"), "end", "value")
+ for after_unpivot in [False, True]:
+ input_df = (
+ df.unpivot(["s", "end"], ["value"], "metric", "value") if after_unpivot else df
+ )
+ for bound in [False, True]:
+ with self.subTest(after_unpivot=after_unpivot, bound=bound):
+ with self.assertRaises(AnalysisException) as error:
+ input_df.binBy(
+ input_df["s.start"] if bound else "s.start",
+ "end",
+ functions.expr("INTERVAL 5 MINUTES"),
+ ["value"],
+ ).collect()
+ self.check_error(
+ exception=error.exception,
+ errorClass="BIN_BY_REQUIRES_TOP_LEVEL_COLUMN",
+ messageParameters={"columnName": '"s.start"'},
+ )
+
+ def test_bin_by_argument_types(self):
+ df = self.spark.range(1)
+ cases = [
+ ({"distributeUniform": "id"}, "distributeUniform", "list or tuple", "str"),
+ ({"distributeUniform": [["value"]]}, "col", "Column or str", "list"),
+ ({"distributeUniform": [("value",)]}, "col", "Column or str", "tuple"),
+ ({"binWidth": "INTERVAL 5 MINUTES"}, "binWidth", "Column", "str"),
+ ({"binWidth": None}, "binWidth", "Column", "NoneType"),
+ ({"alignTo": "origin"}, "alignTo", "Column", "str"),
+ ]
+ cases += [
+ ({name: 1}, name, "str", "int")
+ for name in ("binStartName", "binEndName", "binDistributeRatioName")
+ ]
+ for overrides, arg_name, expected_type, arg_type in cases:
+ with self.subTest(overrides=overrides):
+ args = dict(
+ binWidth=functions.expr("INTERVAL 5 MINUTES"), distributeUniform=["value"]
+ )
+ args.update(overrides)
+ with self.assertRaises(PySparkTypeError) as error:
+ df.binBy("start", "end", **args)
+ self.check_error(
+ exception=error.exception,
+ errorClass="NOT_EXPECTED_TYPE",
+ messageParameters={
+ "arg_name": arg_name,
+ "expected_type": expected_type,
+ "arg_type": arg_type,
+ },
+ )
+
+ def test_bin_by_disabled(self):
+ with self.sql_conf({"spark.sql.binByRelationOperator.enabled": "false"}):
+ with self.assertRaises(AnalysisException) as error:
+ self.spark.range(1).binBy(
+ "id", "id", functions.expr("INTERVAL 5 MINUTES"), ["id"]
+ ).collect()
+ self.assertEqual(error.exception.getCondition(), "UNSUPPORTED_FEATURE.BIN_BY")
+
def test_range(self):
self.assertEqual(self.spark.range(1, 1).count(), 0)
self.assertEqual(self.spark.range(1, 0, -1).count(), 1)
diff --git a/sql/api/src/main/scala/org/apache/spark/sql/Dataset.scala b/sql/api/src/main/scala/org/apache/spark/sql/Dataset.scala
index e5d6263995e3..372fb7fde7e2 100644
--- a/sql/api/src/main/scala/org/apache/spark/sql/Dataset.scala
+++ b/sql/api/src/main/scala/org/apache/spark/sql/Dataset.scala
@@ -21,7 +21,7 @@ import java.util
import scala.jdk.CollectionConverters._
import scala.reflect.runtime.universe.TypeTag
-import org.apache.spark.annotation.{ClassicOnly, DeveloperApi, Stable, Unstable}
+import org.apache.spark.annotation.{ClassicOnly, DeveloperApi, Experimental, Stable, Unstable}
import org.apache.spark.api.java.JavaRDD
import org.apache.spark.api.java.function._
import org.apache.spark.rdd.RDD
@@ -1624,6 +1624,109 @@ abstract class Dataset[T] extends Serializable {
groupByKey(ToScalaUDF(func))(encoder)
}
+ /**
+ * Splits timestamp ranges across fixed bins and proportionally scales selected FLOAT or DOUBLE
+ * columns. The result includes all input columns and three appended columns: `bin_start`,
+ * `bin_end`, and `bin_distribute_ratio`.
+ *
+ * The range and distribution arguments must directly reference top-level columns of this
+ * Dataset. Use `select` to project computed, cast, or nested values into columns first.
+ *
+ * The range columns must have the same timestamp type. The bin width must be a positive,
+ * constant day-time interval. Bins align to 1970-01-01 00:00:00, in the session time zone for
+ * TIMESTAMP.
+ *
+ * A NULL range bound produces one row with NULL bin columns and NULL scaled values. A
+ * zero-length range produces one row with an overlap ratio of 1.0. An inverted range raises an
+ * error.
+ *
+ * {{{
+ * df.binBy(col("start"), col("end"), expr("INTERVAL 5 MINUTES"), Array(col("value")))
+ * }}}
+ *
+ * This operation requires `spark.sql.binByRelationOperator.enabled` to be true.
+ *
+ * @param rangeStart
+ * Timestamp column for the inclusive range start.
+ * @param rangeEnd
+ * Timestamp column for the exclusive range end.
+ * @param binWidth
+ * Positive, constant day-time interval expression.
+ * @param distributeUniform
+ * Nonempty array of FLOAT or DOUBLE columns to scale.
+ * @group untypedrel
+ * @since 4.4.0
+ */
+ @Experimental
+ def binBy(
+ rangeStart: Column,
+ rangeEnd: Column,
+ binWidth: Column,
+ distributeUniform: Array[Column]): DataFrame
+
+ /**
+ * Splits timestamp ranges across fixed bins with an explicit alignment origin.
+ *
+ * @see
+ * `org.apache.spark.sql.Dataset.binBy(Column, Column, Column, Array[Column])`
+ * @param rangeStart
+ * Top-level timestamp column for the inclusive range start.
+ * @param rangeEnd
+ * Top-level timestamp column for the exclusive range end, with the same type as `rangeStart`.
+ * @param binWidth
+ * Positive, constant day-time interval expression.
+ * @param distributeUniform
+ * Nonempty array of top-level FLOAT or DOUBLE columns to scale.
+ * @param alignTo
+ * Constant alignment origin with the same timestamp type as the range columns, or null for
+ * the default origin.
+ * @group untypedrel
+ * @since 4.4.0
+ */
+ @Experimental
+ def binBy(
+ rangeStart: Column,
+ rangeEnd: Column,
+ binWidth: Column,
+ distributeUniform: Array[Column],
+ alignTo: Column): DataFrame
+
+ /**
+ * Splits timestamp ranges across fixed bins with optional alignment and output column names.
+ *
+ * @see
+ * `org.apache.spark.sql.Dataset.binBy(Column, Column, Column, Array[Column])`
+ * @param rangeStart
+ * Top-level timestamp column for the inclusive range start.
+ * @param rangeEnd
+ * Top-level timestamp column for the exclusive range end, with the same type as `rangeStart`.
+ * @param binWidth
+ * Positive, constant day-time interval expression.
+ * @param distributeUniform
+ * Nonempty array of top-level FLOAT or DOUBLE columns to scale.
+ * @param alignTo
+ * Constant alignment origin with the same timestamp type as the range columns, or null for
+ * the default origin.
+ * @param binStartName
+ * Name for the appended bin-start column, or null for `bin_start`.
+ * @param binEndName
+ * Name for the appended bin-end column, or null for `bin_end`.
+ * @param binDistributeRatioName
+ * Name for the appended overlap-ratio column, or null for `bin_distribute_ratio`.
+ * @group untypedrel
+ * @since 4.4.0
+ */
+ @Experimental
+ def binBy(
+ rangeStart: Column,
+ rangeEnd: Column,
+ binWidth: Column,
+ distributeUniform: Array[Column],
+ alignTo: Column,
+ binStartName: String,
+ binEndName: String,
+ binDistributeRatioName: String): DataFrame
+
/**
* Unpivot a DataFrame from wide format to long format, optionally leaving identifier columns
* set. This is the reverse to `groupBy(...).pivot(...).agg(...)`, except for the aggregation,
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveBinBy.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveBinBy.scala
index d90326cc3e0b..cc5a7d971548 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveBinBy.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/analysis/ResolveBinBy.scala
@@ -17,19 +17,16 @@
package org.apache.spark.sql.catalyst.analysis
-import org.apache.spark.SparkException
-import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Expression}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, Expression, ExtractValue}
import org.apache.spark.sql.catalyst.plans.logical.{BinBy, LogicalPlan, UnresolvedBinBy}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.catalyst.trees.TreePattern.UNRESOLVED_BIN_BY
-import org.apache.spark.sql.catalyst.util.StringUtils
import org.apache.spark.sql.errors.QueryCompilationErrors
import org.apache.spark.sql.internal.SQLConf
/**
- * Resolves [[UnresolvedBinBy]] into [[BinBy]]: looks up column references against the child's
- * output, validates types and foldability, and captures the session local time zone for the
- * physical execution to use.
+ * Resolves [[UnresolvedBinBy]] into [[BinBy]]: validates column inputs, types and foldability,
+ * and captures the session local time zone for the physical execution to use.
*/
object ResolveBinBy extends Rule[LogicalPlan] {
@@ -41,21 +38,16 @@ object ResolveBinBy extends Rule[LogicalPlan] {
case b: UnresolvedBinBy => resolve(b)
}
- // `binWidthExpr` must be fully resolved before validation can proceed. The optional
- // `originExpr` must also be resolved when present; an absent `ALIGN TO` clause defaults to
- // `1970-01-01 00:00:00` and bypasses user-facing foldability and type checks. The
- // range/distribute column references can stay as `UnresolvedAttribute`.
private def readyToResolve(b: UnresolvedBinBy): Boolean = {
- b.childrenResolved && b.binWidthExpr.resolved && b.originExpr.forall(_.resolved)
+ b.childrenResolved && b.expressions.forall(_.resolved)
}
private def resolve(b: UnresolvedBinBy): LogicalPlan = {
- val resolver = SQLConf.get.resolver
val child = b.child
- val rangeStart = resolveColumn(b.rangeStartCol, child, resolver)
- val rangeEnd = resolveColumn(b.rangeEndCol, child, resolver)
- val distributeAttributes = b.distributeColumns.map(c => resolveColumn(c, child, resolver))
+ val rangeStart = resolveColumn(b.rangeStartCol)
+ val rangeEnd = resolveColumn(b.rangeEndCol)
+ val distributeAttributes = b.distributeColumns.map(resolveColumn)
val parameters = BinByResolution.validateAndComputeParameters(
rangeStart = rangeStart,
@@ -80,28 +72,12 @@ object ResolveBinBy extends Rule[LogicalPlan] {
timeZoneId = parameters.timeZoneId)
}
- private def resolveColumn(
- expr: Expression,
- child: LogicalPlan,
- resolver: Resolver): Attribute = expr match {
- case u: UnresolvedAttribute =>
- // Still unresolved here means genuinely absent: ResolveReferences rewrites resolvable refs.
- child.resolve(u.nameParts, resolver) match {
- case Some(a: Attribute) => a
- case _ =>
- throw QueryCompilationErrors.unresolvedColumnError(
- u.name,
- proposal = StringUtils.orderSuggestedIdentifiersBySimilarity(
- u.name,
- candidates = child.output.map(attribute => attribute.qualifier :+ attribute.name)))
- }
+ private def resolveColumn(expr: Expression): Attribute = expr match {
case a: Attribute => a
- case alias: Alias if alias.resolved =>
+ case alias @ Alias(_: ExtractValue, _) =>
// A resolved nested ref (e.g. `struct.field`) exists but isn't a top-level column.
throw QueryCompilationErrors.binByRequiresTopLevelColumnError(alias.child)
case other =>
- // Genuinely unexpected; the grammar restricts these positions to multipartIdentifier.
- throw SparkException.internalError(
- s"Unexpected expression in BIN BY column position: $other")
+ throw QueryCompilationErrors.binByRequiresTopLevelColumnError(other)
}
}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala
index 2dc356e8fcd1..ac9b17652ed6 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala
@@ -1795,16 +1795,17 @@ object BinByOutputAliases {
}
/**
- * Unresolved counterpart of [[BinBy]] produced by the parser. `ResolveBinBy` enforces type and
- * foldability constraints and rewrites this into a resolved [[BinBy]].
+ * Unresolved counterpart of [[BinBy]] produced by the parser and DataFrame API. Column positions
+ * accept expressions; `ResolveBinBy` requires them to resolve to top-level attributes. It also
+ * enforces type and foldability constraints and rewrites this into a resolved [[BinBy]].
*
* @param binWidthExpr Bin-width expression (DAY-TIME INTERVAL).
- * @param rangeStartCol Reference to the row's measurement-window start column.
- * @param rangeEndCol Reference to the row's measurement-window end column.
+ * @param rangeStartCol Expression for the row's measurement-window start column.
+ * @param rangeEndCol Expression for the row's measurement-window end column.
* @param originExpr Optional alignment anchor expression (`ALIGN TO` clause). `None`
* when the clause is omitted; `ResolveBinBy` defaults the resolved
* plan's origin to `1970-01-01 00:00:00`.
- * @param distributeColumns Columns whose values are proportionally redistributed across sub-rows.
+ * @param distributeColumns Expressions for values to proportionally redistribute across sub-rows.
* @param outputAliases Optional renames for the three appended output columns.
* @param child Input relation.
*/
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveBinBySuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveBinBySuite.scala
index ff44ae657411..d3de045af4e7 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveBinBySuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/analysis/ResolveBinBySuite.scala
@@ -195,8 +195,9 @@ class ResolveBinBySuite extends AnalysisTest {
val tsStartU = UnresolvedAttribute("ts_start")
val tsEndU = UnresolvedAttribute("ts_end")
val valueU = UnresolvedAttribute("value")
- val result = ResolveBinBy.apply(
- unresolved(rangeStart = tsStartU, rangeEnd = tsEndU, distribute = Seq(valueU)))
+ val result = getAnalyzer.executeAndCheck(
+ unresolved(rangeStart = tsStartU, rangeEnd = tsEndU, distribute = Seq(valueU)),
+ new QueryPlanningTracker)
val bi = result.asInstanceOf[BinBy]
assert(bi.rangeStart.exprId == tsStart.exprId)
assert(bi.rangeEnd.exprId == tsEnd.exprId)
@@ -209,8 +210,9 @@ class ResolveBinBySuite extends AnalysisTest {
val md = new MetadataBuilder().putString("comment", "a measure").build()
val valueMd = AttributeReference("value", DoubleType, nullable = true, md)()
val child = SubqueryAlias("m", LocalRelation(tsStart, tsEnd, valueMd, label))
- val bi = ResolveBinBy.apply(
- unresolved(child = child, distribute = Seq(UnresolvedAttribute(Seq("m", "value")))))
+ val bi = getAnalyzer.executeAndCheck(
+ unresolved(child = child, distribute = Seq(UnresolvedAttribute(Seq("m", "value")))),
+ new QueryPlanningTracker)
.asInstanceOf[BinBy]
// The input is read (held in distributeColumns) but not forwarded by identity.
@@ -283,7 +285,7 @@ class ResolveBinBySuite extends AnalysisTest {
distributeColumns = Seq(UnresolvedAttribute(Seq("t2", "value"))),
outputAliases = BinByOutputAliases.empty,
child = join)
- val resolved1 = ResolveBinBy.apply(u1).asInstanceOf[BinBy]
+ val resolved1 = getAnalyzer.executeAndCheck(u1, new QueryPlanningTracker).asInstanceOf[BinBy]
assert(resolved1.rangeStart.exprId == t1Start.exprId)
assert(resolved1.rangeEnd.exprId == t1End.exprId)
@@ -296,25 +298,28 @@ class ResolveBinBySuite extends AnalysisTest {
distributeColumns = Seq(UnresolvedAttribute(Seq("t2", "value"))),
outputAliases = BinByOutputAliases.empty,
child = join)
- val resolved2 = ResolveBinBy.apply(u2).asInstanceOf[BinBy]
+ val resolved2 = getAnalyzer.executeAndCheck(u2, new QueryPlanningTracker).asInstanceOf[BinBy]
assert(resolved2.rangeStart.exprId == t2Start.exprId)
assert(resolved2.rangeEnd.exprId == t2End.exprId)
}
test("rejects unresolvable column references with UNRESOLVED_COLUMN") {
- expectError(
+ Seq(
unresolved(rangeStart = UnresolvedAttribute("nonexistent")),
- "UNRESOLVED_COLUMN.WITH_SUGGESTION")
- expectError(
- unresolved(distribute = Seq(UnresolvedAttribute("nonexistent"))),
- "UNRESOLVED_COLUMN.WITH_SUGGESTION")
+ unresolved(distribute = Seq(UnresolvedAttribute("nonexistent")))).foreach { plan =>
+ val ex = intercept[SparkThrowable] {
+ getAnalyzer.executeAndCheck(plan, new QueryPlanningTracker)
+ }
+ assert(ex.getCondition == "UNRESOLVED_COLUMN.WITH_SUGGESTION")
+ }
}
test("UNRESOLVED_COLUMN suggestions are ordered by similarity to the missing name") {
// `valeu` is one edit away from `value`, so `value` is suggested first and the remaining
// columns follow by edit distance.
val ex = intercept[SparkThrowable](
- ResolveBinBy.apply(unresolved(rangeStart = UnresolvedAttribute("valeu"))))
+ getAnalyzer.executeAndCheck(
+ unresolved(rangeStart = UnresolvedAttribute("valeu")), new QueryPlanningTracker))
assert(ex.getCondition == "UNRESOLVED_COLUMN.WITH_SUGGESTION")
assert(ex.getMessageParameters.get("objectName") == "`valeu`")
assert(ex.getMessageParameters.get("proposal") == "`value`, `label`, `ts_end`, `ts_start`")
diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
index 1e8267f1b0bc..f2c1529cfbca 100644
--- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
+++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/PlanGenerationTestSuite.scala
@@ -610,6 +610,39 @@ class PlanGenerationTestSuite extends ConnectFunSuite with Logging {
simple.selectExpr("rand(123) BETWEEN 0.1 AND 0.2")
}
+ test("binBy") {
+ val input = session
+ .range(1)
+ .selectExpr(
+ "TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start",
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end",
+ "100.0D AS value")
+ input.binBy(
+ rangeStart = input("start"),
+ rangeEnd = input("end"),
+ binWidth = fn.expr("INTERVAL 5 MINUTES"),
+ distributeUniform = Array(input("value")))
+ }
+
+ test("binBy all_options") {
+ val input = session
+ .range(1)
+ .selectExpr(
+ "TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start",
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end",
+ "100.0D AS value",
+ "20.0F AS value2")
+ input.binBy(
+ rangeStart = input("start"),
+ rangeEnd = input("end"),
+ binWidth = fn.expr("INTERVAL 5 MINUTES"),
+ distributeUniform = Array(input("value"), input("value2")),
+ alignTo = fn.expr("TIMESTAMP_NTZ '2024-01-01 00:02:00'"),
+ binStartName = "window_start",
+ binEndName = "window_end",
+ binDistributeRatioName = "ratio")
+ }
+
test("unpivot values") {
simple.unpivot(
ids = Array(fn.col("id"), fn.col("a")),
diff --git a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/ClientDatasetSuite.scala b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/ClientDatasetSuite.scala
index 2eed9f187456..e97695b88e18 100644
--- a/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/ClientDatasetSuite.scala
+++ b/sql/connect/client/jvm/src/test/scala/org/apache/spark/sql/connect/ClientDatasetSuite.scala
@@ -178,6 +178,64 @@ class ClientDatasetSuite extends ConnectFunSuite {
}
}
+ test("BIN BY preserves bound columns and optional fields") {
+ val df = ss.newDataFrame(_ => ())
+ val end = df("end")
+ val binWidth = expr("INTERVAL 5 MINUTES")
+ val values = Array(df("value"), df("value2"))
+ val default = df
+ .binBy(
+ rangeStart = df("start"),
+ rangeEnd = end,
+ binWidth = binWidth,
+ distributeUniform = values)
+ .plan
+ .getRoot
+ .getBinBy
+ assert(default.getInput == df.plan.getRoot)
+ assert(
+ default.getRangeStart.getUnresolvedAttribute.getPlanId == df.plan.getRoot.getCommon.getPlanId)
+ assert(default.getRangeEnd == toExpr(end))
+ assert(default.getBinWidth == toExpr(binWidth))
+ assert(default.getDistributeUniformCount == 2)
+ assert(default.getDistributeUniform(0) == toExpr(values(0)))
+ assert(default.getDistributeUniform(1) == toExpr(values(1)))
+ assert(
+ !default.hasAlignTo && !default.hasBinStartName && !default.hasBinEndName &&
+ !default.hasBinDistributeRatioName)
+
+ val renamed = df
+ .binBy(
+ rangeStart = df("start"),
+ rangeEnd = end,
+ binWidth = binWidth,
+ distributeUniform = values,
+ alignTo = null,
+ binStartName = "bin.start",
+ binEndName = "bin`end",
+ binDistributeRatioName = "")
+ .plan
+ .getRoot
+ .getBinBy
+ assert(renamed.getBinStartName == "bin.start")
+ assert(renamed.getBinEndName == "bin`end")
+ assert(renamed.hasBinDistributeRatioName && renamed.getBinDistributeRatioName.isEmpty)
+ assert(!renamed.hasAlignTo)
+
+ val origin = lit(null).cast("timestamp")
+ val aligned = df
+ .binBy(
+ rangeStart = df("start"),
+ rangeEnd = end,
+ binWidth = binWidth,
+ distributeUniform = Array(df("value")),
+ alignTo = origin)
+ .plan
+ .getRoot
+ .getBinBy
+ assert(aligned.hasAlignTo && aligned.getAlignTo == toExpr(origin))
+ }
+
test("serialize as null") {
val session = newSparkSession()
val ds = session.range(10)
diff --git a/sql/connect/common/src/main/protobuf/spark/connect/relations.proto b/sql/connect/common/src/main/protobuf/spark/connect/relations.proto
index 151719743703..ac6f9b917894 100644
--- a/sql/connect/common/src/main/protobuf/spark/connect/relations.proto
+++ b/sql/connect/common/src/main/protobuf/spark/connect/relations.proto
@@ -84,6 +84,7 @@ message Relation {
RelationChanges relation_changes = 46;
NearestByJoin nearest_by_join = 47;
Zip zip = 48;
+ BinBy bin_by = 49;
// NA functions
NAFill fill_na = 90;
@@ -969,6 +970,36 @@ message Hint {
repeated Expression parameters = 3;
}
+// Splits timestamp ranges across fixed bins and scales selected FLOAT or DOUBLE columns.
+message BinBy {
+ // (Required) The input relation.
+ Relation input = 1;
+
+ // (Required) Inclusive range start timestamp.
+ Expression range_start = 2;
+
+ // (Required) Exclusive range end, with the same timestamp type as range_start.
+ Expression range_end = 3;
+
+ // (Required) A positive, constant day-time interval expression.
+ Expression bin_width = 4;
+
+ // (Optional) A constant origin with the same timestamp type as the range columns.
+ Expression align_to = 5;
+
+ // (Required) At least one FLOAT or DOUBLE column to scale.
+ repeated Expression distribute_uniform = 6;
+
+ // (Optional) Name for the appended bin-start column. Defaults to bin_start.
+ optional string bin_start_name = 7;
+
+ // (Optional) Name for the appended bin-end column. Defaults to bin_end.
+ optional string bin_end_name = 8;
+
+ // (Optional) Name for the appended overlap-ratio column. Defaults to bin_distribute_ratio.
+ optional string bin_distribute_ratio_name = 9;
+}
+
// Unpivot a DataFrame from wide format to long format, optionally leaving identifier columns set.
message Unpivot {
// (Required) The input relation.
diff --git a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/Dataset.scala b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/Dataset.scala
index 0a18a3526d4b..36c526907e55 100644
--- a/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/Dataset.scala
+++ b/sql/connect/common/src/main/scala/org/apache/spark/sql/connect/Dataset.scala
@@ -674,6 +674,63 @@ class Dataset[T] private[sql] (
groupingSets = Some(groupingSetMsgs))
}
+ /** @inheritdoc */
+ def binBy(
+ rangeStart: Column,
+ rangeEnd: Column,
+ binWidth: Column,
+ distributeUniform: Array[Column]): DataFrame = {
+ binBy(
+ rangeStart = rangeStart,
+ rangeEnd = rangeEnd,
+ binWidth = binWidth,
+ distributeUniform = distributeUniform,
+ alignTo = null)
+ }
+
+ /** @inheritdoc */
+ def binBy(
+ rangeStart: Column,
+ rangeEnd: Column,
+ binWidth: Column,
+ distributeUniform: Array[Column],
+ alignTo: Column): DataFrame = {
+ binBy(
+ rangeStart = rangeStart,
+ rangeEnd = rangeEnd,
+ binWidth = binWidth,
+ distributeUniform = distributeUniform,
+ alignTo = alignTo,
+ binStartName = null,
+ binEndName = null,
+ binDistributeRatioName = null)
+ }
+
+ /** @inheritdoc */
+ def binBy(
+ rangeStart: Column,
+ rangeEnd: Column,
+ binWidth: Column,
+ distributeUniform: Array[Column],
+ alignTo: Column,
+ binStartName: String,
+ binEndName: String,
+ binDistributeRatioName: String): DataFrame = {
+ val columns = Seq(rangeStart, rangeEnd, binWidth) ++ distributeUniform ++ Option(alignTo)
+ sparkSession.newDataFrame(columns) { builder =>
+ val binBy = builder.getBinByBuilder
+ .setInput(plan.getRoot)
+ .setRangeStart(toExpr(rangeStart))
+ .setRangeEnd(toExpr(rangeEnd))
+ .setBinWidth(toExpr(binWidth))
+ Option(alignTo).foreach(origin => binBy.setAlignTo(toExpr(origin)))
+ distributeUniform.foreach(column => binBy.addDistributeUniform(toExpr(column)))
+ Option(binStartName).foreach(binBy.setBinStartName)
+ Option(binEndName).foreach(binBy.setBinEndName)
+ Option(binDistributeRatioName).foreach(binBy.setBinDistributeRatioName)
+ }
+ }
+
/** @inheritdoc */
def unpivot(
ids: Array[Column],
diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/binBy.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/binBy.explain
new file mode 100644
index 000000000000..39c13db1749a
--- /dev/null
+++ b/sql/connect/common/src/test/resources/query-tests/explain-results/binBy.explain
@@ -0,0 +1,3 @@
+BinBy range=[start#0: timestamp_ntz, end#0: timestamp_ntz], binWidth=INTERVAL '0 00:05:00' DAY TO SECOND, alignTo=1970-01-01 00:00:00, distribute=[value#0: double], scaledDistribute=[value#0: double], appends=[bin_start#0: timestamp_ntz, bin_end#0: timestamp_ntz, bin_distribute_ratio#0: double]
++- Project [2024-01-01 00:00:00 AS start#0, 2024-01-01 00:10:00 AS end#0, 100.0 AS value#0]
+ +- Range (0, 1, step=1, splits=Some(2))
diff --git a/sql/connect/common/src/test/resources/query-tests/explain-results/binBy_all_options.explain b/sql/connect/common/src/test/resources/query-tests/explain-results/binBy_all_options.explain
new file mode 100644
index 000000000000..8df0b513a5f9
--- /dev/null
+++ b/sql/connect/common/src/test/resources/query-tests/explain-results/binBy_all_options.explain
@@ -0,0 +1,3 @@
+BinBy range=[start#0: timestamp_ntz, end#0: timestamp_ntz], binWidth=INTERVAL '0 00:05:00' DAY TO SECOND, alignTo=2024-01-01 00:02:00, distribute=[value#0: double, value2#0: float], scaledDistribute=[value#0: double, value2#0: float], appends=[window_start#0: timestamp_ntz, window_end#0: timestamp_ntz, ratio#0: double]
++- Project [2024-01-01 00:00:00 AS start#0, 2024-01-01 00:10:00 AS end#0, 100.0 AS value#0, 20.0 AS value2#0]
+ +- Range (0, 1, step=1, splits=Some(2))
diff --git a/sql/connect/common/src/test/resources/query-tests/queries/binBy.json b/sql/connect/common/src/test/resources/query-tests/queries/binBy.json
new file mode 100644
index 000000000000..4490a8ef33e9
--- /dev/null
+++ b/sql/connect/common/src/test/resources/query-tests/queries/binBy.json
@@ -0,0 +1,179 @@
+{
+ "common": {
+ "planId": "2"
+ },
+ "binBy": {
+ "input": {
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "range": {
+ "start": "0",
+ "end": "1",
+ "step": "1"
+ }
+ },
+ "expressions": [{
+ "expressionString": {
+ "expression": "TIMESTAMP_NTZ \u00272024-01-01 00:00:00\u0027 AS start"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.connect.Dataset",
+ "methodName": "selectExpr",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "expressionString": {
+ "expression": "TIMESTAMP_NTZ \u00272024-01-01 00:10:00\u0027 AS end"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.connect.Dataset",
+ "methodName": "selectExpr",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "expressionString": {
+ "expression": "100.0D AS value"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.connect.Dataset",
+ "methodName": "selectExpr",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }]
+ }
+ },
+ "rangeStart": {
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "start",
+ "planId": "1"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.Dataset",
+ "methodName": "apply",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ },
+ "rangeEnd": {
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "end",
+ "planId": "1"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.Dataset",
+ "methodName": "apply",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ },
+ "binWidth": {
+ "expressionString": {
+ "expression": "INTERVAL 5 MINUTES"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "expr",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ },
+ "distributeUniform": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "value",
+ "planId": "1"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.Dataset",
+ "methodName": "apply",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }]
+ }
+}
\ No newline at end of file
diff --git a/sql/connect/common/src/test/resources/query-tests/queries/binBy.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/binBy.proto.bin
new file mode 100644
index 000000000000..77a6d60fe982
Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/binBy.proto.bin differ
diff --git a/sql/connect/common/src/test/resources/query-tests/queries/binBy_all_options.json b/sql/connect/common/src/test/resources/query-tests/queries/binBy_all_options.json
new file mode 100644
index 000000000000..13b8d53c3602
--- /dev/null
+++ b/sql/connect/common/src/test/resources/query-tests/queries/binBy_all_options.json
@@ -0,0 +1,247 @@
+{
+ "common": {
+ "planId": "2"
+ },
+ "binBy": {
+ "input": {
+ "common": {
+ "planId": "1"
+ },
+ "project": {
+ "input": {
+ "common": {
+ "planId": "0"
+ },
+ "range": {
+ "start": "0",
+ "end": "1",
+ "step": "1"
+ }
+ },
+ "expressions": [{
+ "expressionString": {
+ "expression": "TIMESTAMP_NTZ \u00272024-01-01 00:00:00\u0027 AS start"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.connect.Dataset",
+ "methodName": "selectExpr",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "expressionString": {
+ "expression": "TIMESTAMP_NTZ \u00272024-01-01 00:10:00\u0027 AS end"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.connect.Dataset",
+ "methodName": "selectExpr",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "expressionString": {
+ "expression": "100.0D AS value"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.connect.Dataset",
+ "methodName": "selectExpr",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "expressionString": {
+ "expression": "20.0F AS value2"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.connect.Dataset",
+ "methodName": "selectExpr",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }]
+ }
+ },
+ "rangeStart": {
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "start",
+ "planId": "1"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.Dataset",
+ "methodName": "apply",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ },
+ "rangeEnd": {
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "end",
+ "planId": "1"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.Dataset",
+ "methodName": "apply",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ },
+ "binWidth": {
+ "expressionString": {
+ "expression": "INTERVAL 5 MINUTES"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "expr",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ },
+ "alignTo": {
+ "expressionString": {
+ "expression": "TIMESTAMP_NTZ \u00272024-01-01 00:02:00\u0027"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.functions$",
+ "methodName": "expr",
+ "fileName": "functions.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ },
+ "distributeUniform": [{
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "value",
+ "planId": "1"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.Dataset",
+ "methodName": "apply",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }, {
+ "unresolvedAttribute": {
+ "unparsedIdentifier": "value2",
+ "planId": "1"
+ },
+ "common": {
+ "origin": {
+ "jvmOrigin": {
+ "stackTrace": [{
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.Dataset",
+ "methodName": "apply",
+ "fileName": "Dataset.scala"
+ }, {
+ "classLoaderName": "app",
+ "declaringClass": "org.apache.spark.sql.PlanGenerationTestSuite",
+ "methodName": "~~trimmed~anonfun~~",
+ "fileName": "PlanGenerationTestSuite.scala"
+ }]
+ }
+ }
+ }
+ }],
+ "binStartName": "window_start",
+ "binEndName": "window_end",
+ "binDistributeRatioName": "ratio"
+ }
+}
\ No newline at end of file
diff --git a/sql/connect/common/src/test/resources/query-tests/queries/binBy_all_options.proto.bin b/sql/connect/common/src/test/resources/query-tests/queries/binBy_all_options.proto.bin
new file mode 100644
index 000000000000..850dd31df9a2
Binary files /dev/null and b/sql/connect/common/src/test/resources/query-tests/queries/binBy_all_options.proto.bin differ
diff --git a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
index e821ff228e6c..0c8971a49603 100644
--- a/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
+++ b/sql/connect/server/src/main/scala/org/apache/spark/sql/connect/planner/SparkConnectPlanner.scala
@@ -52,7 +52,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
import org.apache.spark.sql.catalyst.parser.{NamedParameterContext, ParameterContext, ParseException, ParserUtils, PositionalParameterContext}
import org.apache.spark.sql.catalyst.plans.{Cross, FullOuter, Inner, JoinType, LeftAnti, LeftOuter, LeftSemi, RightOuter, UsingJoin}
import org.apache.spark.sql.catalyst.plans.logical
-import org.apache.spark.sql.catalyst.plans.logical.{AppendColumns, Assignment, CoGroup, CollectMetrics, CommandResult, DeduplicateAllColumnsAsKey, DeduplicateKeyColumns, DeleteAction, DeserializeToObject, Except, FlatMapGroupsWithState, InsertAction, InsertStarAction, Intersect, JoinWith, LocalRelation, LogicalGroupState, LogicalPlan, MapGroups, MapPartitions, MergeAction, Project, Sample, SerializeFromObject, Sort, SubqueryAlias, TimeModes, TransformWithState, TypedFilter, Union, Unpivo [...]
+import org.apache.spark.sql.catalyst.plans.logical.{AppendColumns, Assignment, BinByOutputAliases, CoGroup, CollectMetrics, CommandResult, DeduplicateAllColumnsAsKey, DeduplicateKeyColumns, DeleteAction, DeserializeToObject, Except, FlatMapGroupsWithState, InsertAction, InsertStarAction, Intersect, JoinWith, LocalRelation, LogicalGroupState, LogicalPlan, MapGroups, MapPartitions, MergeAction, Project, Sample, SerializeFromObject, Sort, SubqueryAlias, TimeModes, TransformWithState, TypedF [...]
import org.apache.spark.sql.catalyst.streaming.InternalOutputModes
import org.apache.spark.sql.catalyst.trees.{CurrentOrigin, Origin, TreePattern}
import org.apache.spark.sql.catalyst.types.DataTypeUtils
@@ -206,6 +206,7 @@ class SparkConnectPlanner(
transformChunkedCachedLocalRelation(rel.getChunkedCachedLocalRelation)
case proto.Relation.RelTypeCase.HINT => transformHint(rel.getHint)
case proto.Relation.RelTypeCase.UNPIVOT => transformUnpivot(rel.getUnpivot)
+ case proto.Relation.RelTypeCase.BIN_BY => transformBinBy(rel.getBinBy)
case proto.Relation.RelTypeCase.TRANSPOSE => transformTranspose(rel.getTranspose)
case proto.Relation.RelTypeCase.UNRESOLVED_TABLE_VALUED_FUNCTION =>
transformUnresolvedTableValuedFunction(rel.getUnresolvedTableValuedFunction)
@@ -1409,6 +1410,21 @@ class SparkConnectPlanner(
rel.getArgumentsList.asScala.map(transformExpression).toSeq)
}
+ private def transformBinBy(rel: proto.BinBy): LogicalPlan = {
+ UnresolvedBinBy(
+ binWidthExpr = transformExpression(rel.getBinWidth),
+ rangeStartCol = transformExpression(rel.getRangeStart),
+ rangeEndCol = transformExpression(rel.getRangeEnd),
+ originExpr = if (rel.hasAlignTo) Some(transformExpression(rel.getAlignTo)) else None,
+ distributeColumns = rel.getDistributeUniformList.asScala.map(transformExpression).toSeq,
+ outputAliases = BinByOutputAliases(
+ binStart = if (rel.hasBinStartName) Some(rel.getBinStartName) else None,
+ binEnd = if (rel.hasBinEndName) Some(rel.getBinEndName) else None,
+ binRatio =
+ if (rel.hasBinDistributeRatioName) Some(rel.getBinDistributeRatioName) else None),
+ child = transformRelation(rel.getInput))
+ }
+
private def transformUnpivot(rel: proto.Unpivot): LogicalPlan = {
val ids = rel.getIdsList.asScala.toArray.map { expr =>
Column(transformExpression(expr))
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ProtoToParsedPlanTestSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ProtoToParsedPlanTestSuite.scala
index 77e3730d72b9..4b5e36ce23b5 100644
--- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ProtoToParsedPlanTestSuite.scala
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/ProtoToParsedPlanTestSuite.scala
@@ -123,6 +123,7 @@ class ProtoToParsedPlanTestSuite extends SharedSparkSession with ResourceHelper
"org.apache.spark.sql.connect.plugin.ExampleExpressionPlugin")
.set(org.apache.spark.sql.internal.SQLConf.ANSI_ENABLED.key, false.toString)
.set(org.apache.spark.sql.internal.SQLConf.USE_COMMON_EXPR_ID_FOR_ALIAS.key, false.toString)
+ .set(SQLConf.BIN_BY_ENABLED.key, true.toString)
}
protected val suiteBaseResourcePath = commonResourcePath.resolve("query-tests")
diff --git a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectProtoSuite.scala b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectProtoSuite.scala
index 1a5189e24304..0f23b50e17f1 100644
--- a/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectProtoSuite.scala
+++ b/sql/connect/server/src/test/scala/org/apache/spark/sql/connect/planner/SparkConnectProtoSuite.scala
@@ -44,6 +44,7 @@ import org.apache.spark.sql.connector.catalog.{Column => ColumnV2, Identifier, I
import org.apache.spark.sql.connector.catalog.CatalogV2Implicits.CatalogHelper
import org.apache.spark.sql.execution.arrow.ArrowConverters
import org.apache.spark.sql.functions._
+import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types._
import org.apache.spark.unsafe.types.UTF8String
import org.apache.spark.util.Utils
@@ -1017,6 +1018,59 @@ class SparkConnectProtoSuite extends PlanTest with SparkConnectPlanTest {
comparePlans(connectTestRelation.hint("COALESCE", 3), sparkTestRelation.hint("COALESCE", 3))
}
+ test("Test BIN BY") {
+ withSQLConf(SQLConf.BIN_BY_ENABLED.key -> "true") {
+ val schema = new StructType()
+ .add("start", TimestampNTZType)
+ .add("end", TimestampNTZType)
+ .add("value", DoubleType)
+ .add("value2", FloatType)
+ val input = createLocalRelationProto(DataTypeUtils.toAttributes(schema), Seq.empty)
+ val classic = spark.createDataFrame(new java.util.ArrayList[Row](), schema)
+ val binBy = proto.BinBy
+ .newBuilder()
+ .setInput(input)
+ .setRangeStart("start".protoAttr)
+ .setRangeEnd("end".protoAttr)
+ .setBinWidth(
+ proto.Expression
+ .newBuilder()
+ .setExpressionString(
+ proto.Expression.ExpressionString.newBuilder().setExpression("INTERVAL 5 MINUTES")))
+ .addDistributeUniform("value".protoAttr)
+ .addDistributeUniform("value2".protoAttr)
+ comparePlans(
+ proto.Relation.newBuilder().setBinBy(binBy).build(),
+ classic.binBy(
+ rangeStart = col("start"),
+ rangeEnd = col("end"),
+ binWidth = expr("INTERVAL 5 MINUTES"),
+ distributeUniform = Array(col("value"), col("value2"))))
+
+ val origin = "TIMESTAMP_NTZ '2024-01-01 00:02:00'"
+ binBy
+ .setAlignTo(
+ proto.Expression
+ .newBuilder()
+ .setExpressionString(
+ proto.Expression.ExpressionString.newBuilder().setExpression(origin)))
+ .setBinStartName("bin.start")
+ .setBinEndName("bin`end")
+ .setBinDistributeRatioName("")
+ comparePlans(
+ proto.Relation.newBuilder().setBinBy(binBy).build(),
+ classic.binBy(
+ rangeStart = col("start"),
+ rangeEnd = col("end"),
+ binWidth = expr("INTERVAL 5 MINUTES"),
+ distributeUniform = Array(col("value"), col("value2")),
+ alignTo = expr(origin),
+ binStartName = "bin.start",
+ binEndName = "bin`end",
+ binDistributeRatioName = ""))
+ }
+ }
+
test("Test Unpivot") {
val connectPlan0 =
connectTestRelation.unpivot(Seq("id".protoAttr), Seq("name".protoAttr), "variable", "value")
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/classic/Dataset.scala b/sql/core/src/main/scala/org/apache/spark/sql/classic/Dataset.scala
index 8f03d91dc55b..7e855f37da64 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/classic/Dataset.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/classic/Dataset.scala
@@ -1072,6 +1072,61 @@ class Dataset[T] private[sql](
withGroupingKey.newColumns)
}
+ /** @inheritdoc */
+ def binBy(
+ rangeStart: Column,
+ rangeEnd: Column,
+ binWidth: Column,
+ distributeUniform: Array[Column]): DataFrame = {
+ binBy(
+ rangeStart = rangeStart,
+ rangeEnd = rangeEnd,
+ binWidth = binWidth,
+ distributeUniform = distributeUniform,
+ alignTo = null)
+ }
+
+ /** @inheritdoc */
+ def binBy(
+ rangeStart: Column,
+ rangeEnd: Column,
+ binWidth: Column,
+ distributeUniform: Array[Column],
+ alignTo: Column): DataFrame = {
+ binBy(
+ rangeStart = rangeStart,
+ rangeEnd = rangeEnd,
+ binWidth = binWidth,
+ distributeUniform = distributeUniform,
+ alignTo = alignTo,
+ binStartName = null,
+ binEndName = null,
+ binDistributeRatioName = null)
+ }
+
+ /** @inheritdoc */
+ def binBy(
+ rangeStart: Column,
+ rangeEnd: Column,
+ binWidth: Column,
+ distributeUniform: Array[Column],
+ alignTo: Column,
+ binStartName: String,
+ binEndName: String,
+ binDistributeRatioName: String): DataFrame = withPlan {
+ UnresolvedBinBy(
+ binWidthExpr = binWidth.expr,
+ rangeStartCol = rangeStart.expr,
+ rangeEndCol = rangeEnd.expr,
+ originExpr = Option(alignTo).map(_.expr),
+ distributeColumns = distributeUniform.map(_.expr).toImmutableArraySeq,
+ outputAliases = BinByOutputAliases(
+ binStart = Option(binStartName),
+ binEnd = Option(binEndName),
+ binRatio = Option(binDistributeRatioName)),
+ child = logicalPlan)
+ }
+
/** @inheritdoc */
def unpivot(
ids: Array[Column],
diff --git a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
index a780c057e903..6948a7e0c241 100644
--- a/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
+++ b/sql/core/src/test/java/test/org/apache/spark/sql/JavaDataFrameSuite.java
@@ -129,6 +129,39 @@ public class JavaDataFrameSuite {
df.show(1000);
}
+ @Test
+ public void testBinBy() {
+ spark.conf().set("spark.sql.binByRelationOperator.enabled", "true");
+ Dataset<Row> input = spark.sql(
+ "SELECT TIMESTAMP_NTZ '2024-01-01 00:00:00' AS start, " +
+ "TIMESTAMP_NTZ '2024-01-01 00:10:00' AS end, 100.0D AS value, 20.0F AS value2");
+ Column width = expr("INTERVAL 5 MINUTES");
+ Column origin = expr("TIMESTAMP_NTZ '2024-01-01 00:00:00'");
+ Dataset<Row> result = input.binBy(
+ input.col("start"), input.col("end"), width,
+ new Column[] {input.col("value"), input.col("value2")}, origin,
+ "lower", "upper", "ratio");
+ Assertions.assertArrayEquals(
+ new String[] {"start", "end", "value", "value2", "lower", "upper", "ratio"},
+ result.columns());
+ Assertions.assertEquals(
+ Arrays.asList(RowFactory.create(50.0, 10.0f, 0.5), RowFactory.create(50.0, 10.0f, 0.5)),
+ result.select("value", "value2", "ratio").collectAsList());
+
+ Column[] distributeColumns = {input.col("value")};
+ Dataset<Row> snapshotResult = input.binBy(
+ input.col("start"), input.col("end"), width, distributeColumns);
+ Dataset<Row> aligned = input.binBy(
+ input.col("start"), input.col("end"), width, distributeColumns, origin);
+
+ distributeColumns[0] = input.col("value2");
+
+ Assertions.assertEquals(
+ Arrays.asList(RowFactory.create(50.0, 20.0f), RowFactory.create(50.0, 20.0f)),
+ snapshotResult.select("value", "value2").collectAsList());
+ Assertions.assertEquals(snapshotResult.collectAsList(), aligned.collectAsList());
+ }
+
public static class Bean implements Serializable {
private double a = 0.0;
private Integer[] b = { 0, 1 };
diff --git a/sql/core/src/test/resources/sql-tests/analyzer-results/bin-by.sql.out b/sql/core/src/test/resources/sql-tests/analyzer-results/bin-by.sql.out
index e426a5d1bac5..243350306772 100644
--- a/sql/core/src/test/resources/sql-tests/analyzer-results/bin-by.sql.out
+++ b/sql/core/src/test/resources/sql-tests/analyzer-results/bin-by.sql.out
@@ -120,3 +120,47 @@ org.apache.spark.sql.AnalysisException
"columnType" : "\"DOUBLE\""
}
}
+
+
+-- !query
+SELECT * FROM metrics BIN BY (
+ RANGE missing TO ts_end
+ BIN WIDTH INTERVAL '5' MINUTE
+ DISTRIBUTE UNIFORM (value)
+)
+-- !query analysis
+org.apache.spark.sql.catalyst.ExtendedAnalysisException
+{
+ "errorClass" : "UNRESOLVED_COLUMN.WITH_SUGGESTION",
+ "sqlState" : "42703",
+ "messageParameters" : {
+ "objectName" : "`missing`",
+ "proposal" : "`ts_end`, `value`, `ts_start`"
+ },
+ "queryContext" : [ {
+ "objectType" : "",
+ "objectName" : "",
+ "startIndex" : 23,
+ "stopIndex" : 119,
+ "fragment" : "BIN BY (\n RANGE missing TO ts_end\n BIN WIDTH INTERVAL '5' MINUTE\n DISTRIBUTE UNIFORM (value)\n)"
+ } ]
+}
+
+
+-- !query
+SELECT * FROM (
+ SELECT named_struct('start', ts_start) AS s, ts_end, value FROM metrics
+) BIN BY (
+ RANGE s.start TO ts_end
+ BIN WIDTH INTERVAL '5' MINUTE
+ DISTRIBUTE UNIFORM (value)
+)
+-- !query analysis
+org.apache.spark.sql.AnalysisException
+{
+ "errorClass" : "BIN_BY_REQUIRES_TOP_LEVEL_COLUMN",
+ "sqlState" : "0A000",
+ "messageParameters" : {
+ "columnName" : "\"s.start\""
+ }
+}
diff --git a/sql/core/src/test/resources/sql-tests/inputs/bin-by.sql b/sql/core/src/test/resources/sql-tests/inputs/bin-by.sql
index 5413a148ef77..af759f512dee 100644
--- a/sql/core/src/test/resources/sql-tests/inputs/bin-by.sql
+++ b/sql/core/src/test/resources/sql-tests/inputs/bin-by.sql
@@ -56,3 +56,17 @@ SELECT * FROM metrics BIN BY (
BIN WIDTH INTERVAL '5' MINUTE
DISTRIBUTE UNIFORM (value)
);
+
+SELECT * FROM metrics BIN BY (
+ RANGE missing TO ts_end
+ BIN WIDTH INTERVAL '5' MINUTE
+ DISTRIBUTE UNIFORM (value)
+);
+
+SELECT * FROM (
+ SELECT named_struct('start', ts_start) AS s, ts_end, value FROM metrics
+) BIN BY (
+ RANGE s.start TO ts_end
+ BIN WIDTH INTERVAL '5' MINUTE
+ DISTRIBUTE UNIFORM (value)
+);
diff --git a/sql/core/src/test/resources/sql-tests/results/bin-by.sql.out b/sql/core/src/test/resources/sql-tests/results/bin-by.sql.out
index 7630f54276a2..f5e7dd965ba3 100644
--- a/sql/core/src/test/resources/sql-tests/results/bin-by.sql.out
+++ b/sql/core/src/test/resources/sql-tests/results/bin-by.sql.out
@@ -119,3 +119,51 @@ org.apache.spark.sql.AnalysisException
"columnType" : "\"DOUBLE\""
}
}
+
+
+-- !query
+SELECT * FROM metrics BIN BY (
+ RANGE missing TO ts_end
+ BIN WIDTH INTERVAL '5' MINUTE
+ DISTRIBUTE UNIFORM (value)
+)
+-- !query schema
+struct<>
+-- !query output
+org.apache.spark.sql.catalyst.ExtendedAnalysisException
+{
+ "errorClass" : "UNRESOLVED_COLUMN.WITH_SUGGESTION",
+ "sqlState" : "42703",
+ "messageParameters" : {
+ "objectName" : "`missing`",
+ "proposal" : "`ts_end`, `value`, `ts_start`"
+ },
+ "queryContext" : [ {
+ "objectType" : "",
+ "objectName" : "",
+ "startIndex" : 23,
+ "stopIndex" : 119,
+ "fragment" : "BIN BY (\n RANGE missing TO ts_end\n BIN WIDTH INTERVAL '5' MINUTE\n DISTRIBUTE UNIFORM (value)\n)"
+ } ]
+}
+
+
+-- !query
+SELECT * FROM (
+ SELECT named_struct('start', ts_start) AS s, ts_end, value FROM metrics
+) BIN BY (
+ RANGE s.start TO ts_end
+ BIN WIDTH INTERVAL '5' MINUTE
+ DISTRIBUTE UNIFORM (value)
+)
+-- !query schema
+struct<>
+-- !query output
+org.apache.spark.sql.AnalysisException
+{
+ "errorClass" : "BIN_BY_REQUIRES_TOP_LEVEL_COLUMN",
+ "sqlState" : "0A000",
+ "messageParameters" : {
+ "columnName" : "\"s.start\""
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/BinBySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/BinBySuite.scala
index e51f9955adc6..b57b1246ffc8 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/BinBySuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/BinBySuite.scala
@@ -23,6 +23,7 @@ import java.time.{LocalDateTime, ZoneId, ZoneOffset}
import org.apache.spark.{SparkArithmeticException, SparkThrowable}
import org.apache.spark.sql.catalyst.util.DateTimeUtils
import org.apache.spark.sql.execution.BinByExec
+import org.apache.spark.sql.functions.{col, expr, lit}
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.test.SharedSparkSession
@@ -63,6 +64,10 @@ class BinBySuite extends QueryTest with SharedSparkSession {
}
}
+ private def apiInput: DataFrame = spark.sql(
+ """SELECT TIMESTAMP '2024-01-01 00:00:00' AS start,
+ | TIMESTAMP '2024-01-01 00:10:00' AS end, 100.0D AS value""".stripMargin)
+
private def binByExecArgs(sqlText: String): String = {
val plan = spark.sql(sqlText).queryExecution.executedPlan
val binByExec = plan.collectFirst {
@@ -697,4 +702,139 @@ class BinBySuite extends QueryTest with SharedSparkSession {
"bin_distribute_ratio#x: double]")
}
}
+
+ test("BIN BY DataFrame API supports optional alignment and output names") {
+ withSQLConf(
+ SQLConf.BIN_BY_ENABLED.key -> "true",
+ SQLConf.SESSION_LOCAL_TIMEZONE.key -> "UTC") {
+ val input = apiInput
+ val width = expr("INTERVAL 5 MINUTES")
+ val renamed = input.binBy(
+ rangeStart = input("start"),
+ rangeEnd = input("end"),
+ binWidth = width,
+ distributeUniform = Array(input("value")),
+ alignTo = null,
+ binStartName = "bin.start",
+ binEndName = "bin`end",
+ binDistributeRatioName = "ratio")
+ assert(renamed.columns.toSeq ==
+ Seq("start", "end", "value", "bin.start", "bin`end", "ratio"))
+
+ val aligned = input.binBy(
+ rangeStart = col("start"),
+ rangeEnd = col("end"),
+ binWidth = width,
+ distributeUniform = Array(col("value")),
+ alignTo = expr("TIMESTAMP '2024-01-01 00:02:00'"))
+ checkAnswer(aligned.select("value", "bin_start", "bin_distribute_ratio"), Seq(
+ Row(20.0, tsAt("2023-12-31 23:57:00"), 0.2),
+ Row(50.0, tsAt("2024-01-01 00:02:00"), 0.5),
+ Row(30.0, tsAt("2024-01-01 00:07:00"), 0.3)))
+ val result = input.binBy(
+ rangeStart = input("start"),
+ rangeEnd = input("end"),
+ binWidth = width,
+ distributeUniform = Array(input("value")))
+ assert(result.columns.toSeq ==
+ Seq("start", "end", "value", "bin_start", "bin_end", "bin_distribute_ratio"))
+ checkAnswer(result.select("value", "bin_start", "bin_end", "bin_distribute_ratio"), Seq(
+ Row(50.0, tsAt("2024-01-01 00:00:00"), tsAt("2024-01-01 00:05:00"), 0.5),
+ Row(50.0, tsAt("2024-01-01 00:05:00"), tsAt("2024-01-01 00:10:00"), 0.5)))
+ checkAnswer(renamed, result)
+ }
+ }
+
+ test("BIN BY DataFrame API validates column inputs and an explicit NULL origin") {
+ withSQLConf(SQLConf.BIN_BY_ENABLED.key -> "true") {
+ val input = apiInput
+ val width = expr("INTERVAL 5 MINUTES")
+ checkError(
+ exception = intercept[AnalysisException] {
+ input.binBy(
+ rangeStart = col("start"),
+ rangeEnd = col("end"),
+ binWidth = width,
+ distributeUniform = Array.empty[Column])
+ },
+ condition = "BIN_BY_MISSING_DISTRIBUTE",
+ parameters = Map.empty)
+ checkError(
+ exception = intercept[AnalysisException] {
+ input.binBy(
+ rangeStart = col("start"),
+ rangeEnd = col("end"),
+ binWidth = width,
+ distributeUniform = Array(col("value"), input("value")))
+ },
+ condition = "BIN_BY_DUPLICATE_DISTRIBUTE_COLUMN",
+ parameters = Map("columnName" -> "`value`"))
+ checkError(
+ exception = intercept[AnalysisException] {
+ input.binBy(
+ rangeStart = col("start"),
+ rangeEnd = col("end"),
+ binWidth = width,
+ distributeUniform = Array(col("value") + 1))
+ },
+ condition = "BIN_BY_REQUIRES_TOP_LEVEL_COLUMN",
+ parameters = Map("columnName" -> "\"(value + 1)\""))
+ checkError(
+ exception = intercept[AnalysisException] {
+ input.binBy(
+ rangeStart = col("start"),
+ rangeEnd = col("end"),
+ binWidth = width,
+ distributeUniform = Array(col("value").as("v")))
+ },
+ condition = "BIN_BY_REQUIRES_TOP_LEVEL_COLUMN",
+ parameters = Map("columnName" -> "\"value AS v\""))
+ checkError(
+ exception = intercept[AnalysisException] {
+ input.binBy(
+ rangeStart = col("start").as("s"),
+ rangeEnd = col("end"),
+ binWidth = width,
+ distributeUniform = Array(col("value")))
+ },
+ condition = "BIN_BY_REQUIRES_TOP_LEVEL_COLUMN",
+ parameters = Map("columnName" -> "\"start AS s\""))
+ checkError(
+ exception = intercept[AnalysisException] {
+ input.binBy(
+ rangeStart = col("start").cast("timestamp_ntz"),
+ rangeEnd = col("end"),
+ binWidth = width,
+ distributeUniform = Array(col("value")))
+ },
+ condition = "BIN_BY_REQUIRES_TOP_LEVEL_COLUMN",
+ parameters = Map("columnName" -> "\"CAST(start AS TIMESTAMP_NTZ)\""))
+ checkError(
+ exception = intercept[AnalysisException] {
+ input.binBy(
+ rangeStart = col("start"),
+ rangeEnd = col("end"),
+ binWidth = width,
+ distributeUniform = Array(col("value")),
+ alignTo = lit(null).cast("timestamp"))
+ },
+ condition = "BIN_BY_NULL_ARGUMENT",
+ parameters = Map("clause" -> "ALIGN TO"))
+ }
+ }
+
+ test("BIN BY DataFrame API respects the feature gate") {
+ withSQLConf(SQLConf.BIN_BY_ENABLED.key -> "false") {
+ checkError(
+ exception = intercept[AnalysisException] {
+ apiInput.binBy(
+ rangeStart = col("start"),
+ rangeEnd = col("end"),
+ binWidth = expr("INTERVAL 5 MINUTES"),
+ distributeUniform = Array(col("value")))
+ },
+ condition = "UNSUPPORTED_FEATURE.BIN_BY",
+ parameters = Map.empty)
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org