From 426b8ecfc421d7b2e28122ce0ff200014594d086 Mon Sep 17 00:00:00 2001 From: Lundomn <134036682+Lundomn@users.noreply.github.com> Date: Mon, 12 Oct 2026 02:49:15 +0000 Subject: [PATCH 1/4] docs: document Spark JDBC isolation for MatrixOne --- README.md | 1 + README_CN.md | 1 + docs/guides/spark-jdbc.md | 52 ++++ test/manual/spark_jdbc/README.md | 46 ++++ .../spark_jdbc/spark_jdbc_isolation_probe.py | 257 ++++++++++++++++++ 5 files changed, 357 insertions(+) create mode 100644 docs/guides/spark-jdbc.md create mode 100644 test/manual/spark_jdbc/README.md create mode 100644 test/manual/spark_jdbc/spark_jdbc_isolation_probe.py diff --git a/README.md b/README.md index e1dab5ae5ddaa..cb11889763c98 100644 --- a/README.md +++ b/README.md @@ -248,6 +248,7 @@ Ready to dive deeper? Explore our comprehensive collection of hands-on tutorials | Tutorial | Language/Framework | Description | |----------|-------------------|-------------| | [Java CRUD Demo](https://docs.matrixorigin.cn/mo/en/latest/MatrixOne/Tutorial/develop-java-crud-demo.html) | Java | Java application development | +| [Spark JDBC writes](docs/guides/spark-jdbc.md) | Python / Spark | Configure Spark JDBC writes for MatrixOne's supported isolation level | | [SpringBoot and JPA CRUD Demo](https://docs.matrixorigin.cn/mo/en/latest/MatrixOne/Tutorial/springboot-hibernate-crud-demo.html) | Java | SpringBoot with Hibernate/JPA | | [PyMySQL CRUD Demo](https://docs.matrixorigin.cn/mo/en/latest/MatrixOne/Tutorial/develop-python-crud-demo.html) | Python | Basic database operations with Python | | [SQLAlchemy CRUD Demo](https://docs.matrixorigin.cn/mo/en/latest/MatrixOne/Tutorial/sqlalchemy-python-crud-demo.html) | Python | Python with SQLAlchemy ORM | diff --git a/README_CN.md b/README_CN.md index e845b9c976232..cd1f44815ccd0 100644 --- a/README_CN.md +++ b/README_CN.md @@ -242,6 +242,7 @@ for row in results.rows: | 教程 | 语言/框架 | 说明 | |----------|-------------------|-------------| | [Java CRUD 示例](https://docs.matrixorigin.cn/mo/en/latest/MatrixOne/Tutorial/develop-java-crud-demo.html) | Java | Java 应用开发 | +| [Spark JDBC 写入指南](docs/guides/spark-jdbc.md) | Python / Spark | 为 MatrixOne 配置受支持的 JDBC 写入隔离级别 | | [SpringBoot 和 JPA CRUD 示例](https://docs.matrixorigin.cn/mo/en/latest/MatrixOne/Tutorial/springboot-hibernate-crud-demo.html) | Java | SpringBoot + Hibernate/JPA | | [PyMySQL CRUD 示例](https://docs.matrixorigin.cn/mo/en/latest/MatrixOne/Tutorial/develop-python-crud-demo.html) | Python | Python 基础数据库操作 | | [SQLAlchemy CRUD 示例](https://docs.matrixorigin.cn/mo/en/latest/MatrixOne/Tutorial/sqlalchemy-python-crud-demo.html) | Python | Python + SQLAlchemy ORM | diff --git a/docs/guides/spark-jdbc.md b/docs/guides/spark-jdbc.md new file mode 100644 index 0000000000000..5eec2272caa5b --- /dev/null +++ b/docs/guides/spark-jdbc.md @@ -0,0 +1,52 @@ +# Spark JDBC writes to MatrixOne + +[Spark's JDBC write option documentation](https://spark.apache.org/docs/3.5.6/sql-data-sources-jdbc.html#data-source-option) +lists `READ_UNCOMMITTED` as the default `isolationLevel`. MatrixOne rejects +that isolation level, so a Spark write that leaves the default in place fails +when Connector/J configures the writer connection. Set Spark's JDBC +`isolationLevel` option to `READ_COMMITTED` on every write: + +```python +connection_properties = { + "driver": "com.mysql.cj.jdbc.Driver", + "user": "your_user", + "password": "your_password", + "isolationLevel": "READ_COMMITTED", +} + +df.write.jdbc( + url="jdbc:mysql://127.0.0.1:6001/your_database", + table="target_table", + mode="append", + properties=connection_properties, +) +``` + +`isolationLevel` is a Spark JDBC **write** option. Spark opens JDBC connections +for its writer tasks, so setting isolation on a separate SQL connection does +not configure those writes. MatrixOne supports `READ_COMMITTED`; its session +isolation variable displays this as `READ-COMMITTED`. + +Read the committed output with a new JDBC read when checking a pipeline: + +```python +read_properties = { + "driver": "com.mysql.cj.jdbc.Driver", + "user": "your_user", + "password": "your_password", +} +readback = spark.read.jdbc( + url="jdbc:mysql://127.0.0.1:6001/your_database", + table="target_table", + properties=read_properties, +) +``` + +The repository includes a runnable compatibility probe that checks the default +failure, the explicit `READ_COMMITTED` write across two writer partitions, and +readback through both a fresh Spark JDBC read and a separate SQL connection: +[`test/manual/spark_jdbc/README.md`](../../test/manual/spark_jdbc/README.md). + +This option selects a supported isolation level for each Spark write. It does +not make a multi-partition Spark append atomic or change Spark's task retry +behavior. diff --git a/test/manual/spark_jdbc/README.md b/test/manual/spark_jdbc/README.md new file mode 100644 index 0000000000000..f1ee8afb1e953 --- /dev/null +++ b/test/manual/spark_jdbc/README.md @@ -0,0 +1,46 @@ +# Spark JDBC isolation probe + +This probe reproduces the Spark JDBC write-isolation compatibility boundary +with Spark 3.5.3 and MySQL Connector/J 8.0.33. It verifies that a writer with +Spark's default isolation fails with MatrixOne's explicit unsupported-level +error and inserts no rows, then writes the same two rows with +`isolationLevel=READ_COMMITTED` across two partitions and reads them back +through Spark and a separate JDBC connection. The probe creates uniquely named +tables and removes them when it exits. + +Use an isolated MatrixOne instance and a dedicated database. Set the connection +environment variables without placing credentials in the source: + +```sh +export MO_JDBC_URL='jdbc:mysql://127.0.0.1:6001/spark_jdbc_issue_29838' +export MO_USER='your_user' +export MO_PASSWORD='your_password' +export MYSQL_CONNECTOR_JAR='/absolute/path/to/mysql-connector-j-8.0.33.jar' +``` + +Create the database before running the probe. Install the pinned PySpark and +Connector/J dependencies, then launch it with both driver and executor +classpath entries for Connector/J: + +```sh +python3 -m venv /tmp/spark-jdbc-29838-venv +/tmp/spark-jdbc-29838-venv/bin/pip install pyspark==3.5.3 +curl -fL \ + https://repo.maven.apache.org/maven2/com/mysql/mysql-connector-j/8.0.33/mysql-connector-j-8.0.33.jar \ + -o /tmp/mysql-connector-j-8.0.33.jar +export MYSQL_CONNECTOR_JAR=/tmp/mysql-connector-j-8.0.33.jar +export PYSPARK_DRIVER_PYTHON=/tmp/spark-jdbc-29838-venv/bin/python +export PYSPARK_PYTHON=/tmp/spark-jdbc-29838-venv/bin/python + +/tmp/spark-jdbc-29838-venv/bin/spark-submit \ + --master 'local[2]' \ + --jars "$MYSQL_CONNECTOR_JAR" \ + --driver-class-path "$MYSQL_CONNECTOR_JAR" \ + test/manual/spark_jdbc/spark_jdbc_isolation_probe.py +``` + +The command can be run again to check reuse and cleanup. It prints the MatrixOne +version, Spark version, and the verified default failure and `READ_COMMITTED` +readback. The probe fails if the default write succeeds, if an unrelated error +causes the default write to fail, if any default rows were inserted, or if the +explicit-isolation write/readback differs from the expected rows. diff --git a/test/manual/spark_jdbc/spark_jdbc_isolation_probe.py b/test/manual/spark_jdbc/spark_jdbc_isolation_probe.py new file mode 100644 index 0000000000000..17429893adc92 --- /dev/null +++ b/test/manual/spark_jdbc/spark_jdbc_isolation_probe.py @@ -0,0 +1,257 @@ +#!/usr/bin/env python3 +"""Exercise Spark's JDBC write isolation against an isolated MatrixOne database.""" + +from __future__ import annotations + +import os +import sys +import uuid + +from pyspark.sql import SparkSession + + +DRIVER = "com.mysql.cj.jdbc.Driver" +EXPECTED_ROWS = [(1, 0.25), (2, 0.75)] + + +def required_env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise RuntimeError(f"{name} must be set") + return value + + +def jdbc_connection(spark: SparkSession, url: str, user: str, password: str): + return spark._jvm.java.sql.DriverManager.getConnection(url, user, password) + + +def query_rows(connection, sql: str): + statement = connection.createStatement() + result = statement.executeQuery(sql) + try: + rows = [] + while result.next(): + rows.append((int(result.getLong(1)), float(result.getDouble(2)))) + return rows + finally: + result.close() + statement.close() + + +def row_count(connection, table: str) -> int: + statement = connection.createStatement() + result = statement.executeQuery(f"SELECT COUNT(*) FROM `{table}`") + try: + if not result.next(): + raise AssertionError(f"COUNT(*) returned no row for {table}") + return int(result.getLong(1)) + finally: + result.close() + statement.close() + + +def exception_chain(exc: Exception) -> str: + messages = [str(exc)] + current = getattr(exc, "java_exception", None) + for _ in range(20): + if current is None: + break + message = str(current.toString()) + try: + error_code = int(current.getErrorCode()) + except Exception: + error_code = 0 + if error_code: + message += f" [JDBC error code {error_code}]" + messages.append(message) + current = current.getCause() + return "\n".join(messages) + + +def drop_owned_tables(spark: SparkSession, tables: list[str]) -> None: + connection = jdbc_connection( + spark, + required_env("MO_JDBC_URL"), + required_env("MO_USER"), + required_env("MO_PASSWORD"), + ) + try: + for table in tables: + statement = connection.createStatement() + try: + statement.execute(f"DROP TABLE IF EXISTS `{table}`") + result = statement.executeQuery(f"SHOW TABLES LIKE '{table}'") + try: + if result.next(): + raise AssertionError(f"probe table was not removed: {table}") + finally: + result.close() + finally: + statement.close() + finally: + connection.close() + + +def main() -> None: + url = required_env("MO_JDBC_URL") + user = required_env("MO_USER") + password = required_env("MO_PASSWORD") + connector_jar = os.path.abspath(required_env("MYSQL_CONNECTOR_JAR")) + if not os.path.isfile(connector_jar): + raise RuntimeError(f"Connector/J jar does not exist: {connector_jar}") + + suffix = uuid.uuid4().hex[:12] + default_table = f"spark_jdbc_default_{suffix}" + rc_table = f"spark_jdbc_rc_{suffix}" + tables = [default_table, rc_table] + + spark = ( + SparkSession.builder.master("local[2]") + .appName("matrixone-spark-jdbc-isolation-29838") + .config("spark.jars", connector_jar) + .getOrCreate() + ) + primary_error = None + try: + if spark.version != "3.5.3": + raise RuntimeError( + f"expected Spark 3.5.3, got {spark.version}; use the pinned probe environment" + ) + + jvm = spark._jvm + jvm.java.lang.Class.forName(DRIVER) + print(f"JDK: {spark._jvm.java.lang.System.getProperty('java.version')}") + print(f"Spark: {spark.version}") + + setup_connection = jdbc_connection(spark, url, user, password) + try: + statement = setup_connection.createStatement() + try: + version_result = statement.executeQuery("SELECT VERSION()") + try: + if not version_result.next(): + raise AssertionError("SELECT VERSION() returned no row") + print(f"MatrixOne version: {version_result.getString(1)}") + finally: + version_result.close() + isolation_result = statement.executeQuery( + "SELECT @@session.transaction_isolation" + ) + try: + if not isolation_result.next(): + raise AssertionError( + "session isolation query returned no row" + ) + print( + "Fresh JDBC session isolation: " + f"{isolation_result.getString(1)}" + ) + finally: + isolation_result.close() + for table in tables: + statement.execute( + f"CREATE TABLE `{table}` (id BIGINT NOT NULL PRIMARY KEY, value DOUBLE NOT NULL)" + ) + finally: + statement.close() + finally: + setup_connection.close() + + frame = spark.createDataFrame(EXPECTED_ROWS, "id long, value double") + + try: + frame.coalesce(1).write.jdbc( + url=url, + table=default_table, + mode="append", + properties={"driver": DRIVER, "user": user, "password": password}, + ) + except Exception as exc: + detail = exception_chain(exc) + normalized = detail.lower() + if ( + "20105" not in normalized + or "read-uncommitted" not in normalized + or "not supported" not in normalized + ): + raise AssertionError( + "default Spark JDBC write failed for a reason other than the expected " + "MatrixOne 20105 unsupported READ_UNCOMMITTED error:\n" + f"{detail}" + ) from exc + print("PASS: default Spark JDBC write rejects READ_UNCOMMITTED with MatrixOne 20105") + else: + raise AssertionError( + "default Spark JDBC write unexpectedly succeeded; MatrixOne must not silently " + "map READ_UNCOMMITTED to another isolation level" + ) + + check_connection = jdbc_connection(spark, url, user, password) + try: + if row_count(check_connection, default_table) != 0: + raise AssertionError("default write inserted rows before the isolation error") + finally: + check_connection.close() + print("PASS: default write left its pre-created table empty") + + writer_properties = { + "driver": DRIVER, + "user": user, + "password": password, + "isolationLevel": "READ_COMMITTED", + } + frame.repartition(2).write.jdbc( + url=url, + table=rc_table, + mode="append", + properties=writer_properties, + ) + + read_properties = {"driver": DRIVER, "user": user, "password": password} + spark_rows = sorted( + (int(row.id), float(row.value)) + for row in spark.read.jdbc(url=url, table=rc_table, properties=read_properties).collect() + ) + if spark_rows != EXPECTED_ROWS: + raise AssertionError(f"fresh Spark JDBC read mismatch: {spark_rows}") + + read_connection = jdbc_connection(spark, url, user, password) + try: + sql_rows = query_rows(read_connection, f"SELECT id, value FROM `{rc_table}` ORDER BY id") + if sql_rows != EXPECTED_ROWS: + raise AssertionError(f"independent JDBC SQL read mismatch: {sql_rows}") + statement = read_connection.createStatement() + result = statement.executeQuery( + f"SELECT COUNT(*), SUM(value) FROM `{rc_table}`" + ) + try: + if not result.next() or int(result.getLong(1)) != 2 or float(result.getDouble(2)) != 1.0: + raise AssertionError("independent JDBC aggregate must be count=2, sum=1.0") + finally: + result.close() + statement.close() + finally: + read_connection.close() + + print(f"PASS: READ_COMMITTED Spark writer committed both rows from two partitions: {spark_rows}") + print("PASS: fresh Spark read and independent JDBC SQL read agree; count=2, sum=1.0") + except BaseException as exc: + primary_error = exc + raise + finally: + try: + drop_owned_tables(spark, tables) + except Exception as cleanup_error: + if primary_error is None: + raise + print(f"cleanup also failed: {cleanup_error}", file=sys.stderr) + try: + spark.stop() + except Exception as stop_error: + if primary_error is None: + raise + print(f"Spark shutdown also failed: {stop_error}", file=sys.stderr) + + +if __name__ == "__main__": + main() From a52daa6790c55ebe0b4794e219a99d314019b50e Mon Sep 17 00:00:00 2001 From: Lundomn <134036682+Lundomn@users.noreply.github.com> Date: Mon, 12 Oct 2026 03:03:12 +0000 Subject: [PATCH 2/4] docs: tighten Spark JDBC writer probe coverage --- docs/guides/spark-jdbc.md | 3 +- test/manual/spark_jdbc/README.md | 4 +-- .../spark_jdbc/spark_jdbc_isolation_probe.py | 31 ++++++++++++++----- 3 files changed, 27 insertions(+), 11 deletions(-) diff --git a/docs/guides/spark-jdbc.md b/docs/guides/spark-jdbc.md index 5eec2272caa5b..c6e96f888e853 100644 --- a/docs/guides/spark-jdbc.md +++ b/docs/guides/spark-jdbc.md @@ -44,7 +44,8 @@ readback = spark.read.jdbc( The repository includes a runnable compatibility probe that checks the default failure, the explicit `READ_COMMITTED` write across two writer partitions, and -readback through both a fresh Spark JDBC read and a separate SQL connection: +readback through both a fresh Spark JDBC read and a separate SQL connection. It +places one row in each writer partition to exercise both JDBC writer tasks: [`test/manual/spark_jdbc/README.md`](../../test/manual/spark_jdbc/README.md). This option selects a supported isolation level for each Spark write. It does diff --git a/test/manual/spark_jdbc/README.md b/test/manual/spark_jdbc/README.md index f1ee8afb1e953..f6ce820164c96 100644 --- a/test/manual/spark_jdbc/README.md +++ b/test/manual/spark_jdbc/README.md @@ -4,8 +4,8 @@ This probe reproduces the Spark JDBC write-isolation compatibility boundary with Spark 3.5.3 and MySQL Connector/J 8.0.33. It verifies that a writer with Spark's default isolation fails with MatrixOne's explicit unsupported-level error and inserts no rows, then writes the same two rows with -`isolationLevel=READ_COMMITTED` across two partitions and reads them back -through Spark and a separate JDBC connection. The probe creates uniquely named +`isolationLevel=READ_COMMITTED` across two non-empty writer partitions and reads +them back through Spark and a separate JDBC connection. The probe creates uniquely named tables and removes them when it exits. Use an isolated MatrixOne instance and a dedicated database. Set the connection diff --git a/test/manual/spark_jdbc/spark_jdbc_isolation_probe.py b/test/manual/spark_jdbc/spark_jdbc_isolation_probe.py index 17429893adc92..88090d3910017 100644 --- a/test/manual/spark_jdbc/spark_jdbc_isolation_probe.py +++ b/test/manual/spark_jdbc/spark_jdbc_isolation_probe.py @@ -200,7 +200,18 @@ def main() -> None: "password": password, "isolationLevel": "READ_COMMITTED", } - frame.repartition(2).write.jdbc( + writer_rdd = spark.sparkContext.parallelize(EXPECTED_ROWS, 2) + partition_row_counts = writer_rdd.mapPartitions( + lambda rows: [sum(1 for _ in rows)] + ).collect() + if partition_row_counts != [1, 1]: + raise AssertionError( + "expected one row in each of two writer partitions, got " + f"{partition_row_counts}" + ) + writer_frame = spark.createDataFrame(writer_rdd, "id long, value double") + print(f"Spark READ_COMMITTED writer partition row counts: {partition_row_counts}") + writer_frame.write.jdbc( url=url, table=rc_table, mode="append", @@ -239,18 +250,22 @@ def main() -> None: primary_error = exc raise finally: + cleanup_error = None try: drop_owned_tables(spark, tables) - except Exception as cleanup_error: - if primary_error is None: - raise - print(f"cleanup also failed: {cleanup_error}", file=sys.stderr) + except BaseException as exc: + cleanup_error = exc try: spark.stop() - except Exception as stop_error: + except BaseException as exc: + if cleanup_error is None: + cleanup_error = exc + else: + print(f"Spark shutdown also failed: {exc}", file=sys.stderr) + if cleanup_error is not None: if primary_error is None: - raise - print(f"Spark shutdown also failed: {stop_error}", file=sys.stderr) + raise cleanup_error + print(f"cleanup also failed: {cleanup_error}", file=sys.stderr) if __name__ == "__main__": From d3b21b990e2e3edb983427b25bb72e33c82ceb45 Mon Sep 17 00:00:00 2001 From: Lundomn <134036682+Lundomn@users.noreply.github.com> Date: Mon, 12 Oct 2026 07:20:05 +0000 Subject: [PATCH 3/4] test(partition): cover TopN ANY NULL ordering --- pkg/sql/colexec/partition/topn_test.go | 89 ++++++++++++++++++++++++-- 1 file changed, 85 insertions(+), 4 deletions(-) diff --git a/pkg/sql/colexec/partition/topn_test.go b/pkg/sql/colexec/partition/topn_test.go index c78cb89b4f933..741b348a6d549 100644 --- a/pkg/sql/colexec/partition/topn_test.go +++ b/pkg/sql/colexec/partition/topn_test.go @@ -25,6 +25,7 @@ import ( "github.com/matrixorigin/matrixone/pkg/common/moerr" "github.com/matrixorigin/matrixone/pkg/common/mpool" "github.com/matrixorigin/matrixone/pkg/container/batch" + "github.com/matrixorigin/matrixone/pkg/container/nulls" "github.com/matrixorigin/matrixone/pkg/container/types" "github.com/matrixorigin/matrixone/pkg/container/vector" "github.com/matrixorigin/matrixone/pkg/pb/plan" @@ -129,15 +130,95 @@ func TestPartitionTopNLiteralOrderKeys(t *testing.T) { } func TestPartitionTopNRejectsUnsupportedOrderType(t *testing.T) { + for _, tc := range []struct { + name string + oid types.T + supported bool + }{ + {"tuple", types.T_tuple, false}, + {"interval", types.T_interval, false}, + // ANY is supported for domainless NULL values; it cannot serve as + // the unsupported-type counterexample. + {"domainless_null", types.T_any, true}, + } { + t.Run(tc.name, func(t *testing.T) { + proc := testutil.NewProcessWithOwnedMPool(t, "", mpool.MustNewZero()) + t.Cleanup(func() { require.Zero(t, proc.Mp().CurrNB()) }) + t.Cleanup(proc.Free) + arg := newTopNArgument(1) + t.Cleanup(func() { arg.Free(proc, true, nil) }) + arg.OrderBySpecs = arg.OrderBySpecs[:2] + arg.OrderBySpecs[1].Expr = topNCol(1, tc.oid) + err := arg.Prepare(proc) + if tc.supported { + require.NoError(t, err) + } else { + require.True(t, moerr.IsMoErrCode(err, moerr.ErrNotSupported), "%v", err) + } + }) + } +} + +func TestPartitionTopNNonLiteralAnyNullOrder(t *testing.T) { proc := testutil.NewProcessWithOwnedMPool(t, "", mpool.MustNewZero()) t.Cleanup(func() { require.Zero(t, proc.Mp().CurrNB()) }) t.Cleanup(proc.Free) arg := newTopNArgument(1) - t.Cleanup(func() { arg.Free(proc, true, nil) }) - arg.OrderBySpecs = arg.OrderBySpecs[:2] arg.OrderBySpecs[1].Expr = topNCol(1, types.T_any) - err := arg.Prepare(proc) - require.True(t, moerr.IsMoErrCode(err, moerr.ErrNotSupported), "%v", err) + t.Cleanup(func() { arg.Free(proc, false, nil) }) + + for generation := 0; generation < 2; generation++ { + batches := []*batch.Batch{ + makeAnyNullTopNBatch(t, proc, []int32{1, 2}, []int64{10, 20}, false), + makeAnyNullTopNBatch(t, proc, []int32{1, 2}, []int64{11, 21}, true), + } + child := colexec.NewMockOperator().WithBatchs(batches) + t.Cleanup(func() { child.Free(proc, false, nil) }) + arg.Children = nil + arg.AppendChild(child) + + require.NoError(t, arg.Prepare(proc)) + require.Len(t, arg.top.compares, 2) + require.NotNil(t, arg.top.compares[0]) + + counts := make(map[int32]int) + validPayload := map[int32][]int64{1: {10, 11}, 2: {20, 21}} + for { + result, err := arg.Call(proc) + require.NoError(t, err) + if result.Batch == nil || result.Status == vm.ExecStop { + break + } + groups := vector.MustFixedColWithTypeCheck[int32](result.Batch.Vecs[0]) + payloads := vector.MustFixedColWithTypeCheck[int64](result.Batch.Vecs[2]) + for row, group := range groups { + require.True(t, result.Batch.Vecs[1].IsNull(uint64(row))) + require.Contains(t, validPayload[group], payloads[row]) + counts[group]++ + } + } + require.Equal(t, map[int32]int{1: 1, 2: 1}, counts) + arg.Reset(proc, false, nil) + } +} + +func makeAnyNullTopNBatch(t *testing.T, proc *process.Process, groups []int32, payloads []int64, constant bool) *batch.Batch { + t.Helper() + require.Len(t, groups, len(payloads)) + bat := batch.NewWithSize(3) + bat.Vecs[0] = vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixedList(bat.Vecs[0], groups, nil, proc.Mp())) + if constant { + bat.Vecs[1] = vector.NewConstNull(types.T_any.ToType(), len(groups), proc.Mp()) + } else { + bat.Vecs[1] = vector.NewVec(types.T_any.ToType()) + bat.Vecs[1].SetLength(len(groups)) + nulls.AddRange(bat.Vecs[1].GetNulls(), 0, uint64(len(groups))) + } + bat.Vecs[2] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList(bat.Vecs[2], payloads, nil, proc.Mp())) + bat.SetRowCount(len(groups)) + return bat } func TestPartitionTopNSQLOrderFloatNaNs(t *testing.T) { From ca733a4b001b364ab467521b182fdb199d2c40ad Mon Sep 17 00:00:00 2001 From: Lundomn <134036682+Lundomn@users.noreply.github.com> Date: Mon, 12 Oct 2026 07:32:23 +0000 Subject: [PATCH 4/4] test(partition): separate ANY NULL TopN regression --- pkg/sql/colexec/partition/topn_test.go | 124 ++++++++++++------------- 1 file changed, 62 insertions(+), 62 deletions(-) diff --git a/pkg/sql/colexec/partition/topn_test.go b/pkg/sql/colexec/partition/topn_test.go index 741b348a6d549..86389fa8ec6a6 100644 --- a/pkg/sql/colexec/partition/topn_test.go +++ b/pkg/sql/colexec/partition/topn_test.go @@ -159,68 +159,6 @@ func TestPartitionTopNRejectsUnsupportedOrderType(t *testing.T) { } } -func TestPartitionTopNNonLiteralAnyNullOrder(t *testing.T) { - proc := testutil.NewProcessWithOwnedMPool(t, "", mpool.MustNewZero()) - t.Cleanup(func() { require.Zero(t, proc.Mp().CurrNB()) }) - t.Cleanup(proc.Free) - arg := newTopNArgument(1) - arg.OrderBySpecs[1].Expr = topNCol(1, types.T_any) - t.Cleanup(func() { arg.Free(proc, false, nil) }) - - for generation := 0; generation < 2; generation++ { - batches := []*batch.Batch{ - makeAnyNullTopNBatch(t, proc, []int32{1, 2}, []int64{10, 20}, false), - makeAnyNullTopNBatch(t, proc, []int32{1, 2}, []int64{11, 21}, true), - } - child := colexec.NewMockOperator().WithBatchs(batches) - t.Cleanup(func() { child.Free(proc, false, nil) }) - arg.Children = nil - arg.AppendChild(child) - - require.NoError(t, arg.Prepare(proc)) - require.Len(t, arg.top.compares, 2) - require.NotNil(t, arg.top.compares[0]) - - counts := make(map[int32]int) - validPayload := map[int32][]int64{1: {10, 11}, 2: {20, 21}} - for { - result, err := arg.Call(proc) - require.NoError(t, err) - if result.Batch == nil || result.Status == vm.ExecStop { - break - } - groups := vector.MustFixedColWithTypeCheck[int32](result.Batch.Vecs[0]) - payloads := vector.MustFixedColWithTypeCheck[int64](result.Batch.Vecs[2]) - for row, group := range groups { - require.True(t, result.Batch.Vecs[1].IsNull(uint64(row))) - require.Contains(t, validPayload[group], payloads[row]) - counts[group]++ - } - } - require.Equal(t, map[int32]int{1: 1, 2: 1}, counts) - arg.Reset(proc, false, nil) - } -} - -func makeAnyNullTopNBatch(t *testing.T, proc *process.Process, groups []int32, payloads []int64, constant bool) *batch.Batch { - t.Helper() - require.Len(t, groups, len(payloads)) - bat := batch.NewWithSize(3) - bat.Vecs[0] = vector.NewVec(types.T_int32.ToType()) - require.NoError(t, vector.AppendFixedList(bat.Vecs[0], groups, nil, proc.Mp())) - if constant { - bat.Vecs[1] = vector.NewConstNull(types.T_any.ToType(), len(groups), proc.Mp()) - } else { - bat.Vecs[1] = vector.NewVec(types.T_any.ToType()) - bat.Vecs[1].SetLength(len(groups)) - nulls.AddRange(bat.Vecs[1].GetNulls(), 0, uint64(len(groups))) - } - bat.Vecs[2] = vector.NewVec(types.T_int64.ToType()) - require.NoError(t, vector.AppendFixedList(bat.Vecs[2], payloads, nil, proc.Mp())) - bat.SetRowCount(len(groups)) - return bat -} - func TestPartitionTopNSQLOrderFloatNaNs(t *testing.T) { for _, tc := range []struct { name string @@ -823,3 +761,65 @@ func collectTopNRowsAt( } return rows } + +func TestPartitionTopNNonLiteralAnyNullOrder(t *testing.T) { + proc := testutil.NewProcessWithOwnedMPool(t, "", mpool.MustNewZero()) + t.Cleanup(func() { require.Zero(t, proc.Mp().CurrNB()) }) + t.Cleanup(proc.Free) + arg := newTopNArgument(1) + arg.OrderBySpecs[1].Expr = topNCol(1, types.T_any) + t.Cleanup(func() { arg.Free(proc, false, nil) }) + + for generation := 0; generation < 2; generation++ { + batches := []*batch.Batch{ + makeAnyNullTopNBatch(t, proc, []int32{1, 2}, []int64{10, 20}, false), + makeAnyNullTopNBatch(t, proc, []int32{1, 2}, []int64{11, 21}, true), + } + child := colexec.NewMockOperator().WithBatchs(batches) + t.Cleanup(func() { child.Free(proc, false, nil) }) + arg.Children = nil + arg.AppendChild(child) + + require.NoError(t, arg.Prepare(proc)) + require.Len(t, arg.top.compares, 2) + require.NotNil(t, arg.top.compares[0]) + + counts := make(map[int32]int) + validPayload := map[int32][]int64{1: {10, 11}, 2: {20, 21}} + for { + result, err := arg.Call(proc) + require.NoError(t, err) + if result.Batch == nil || result.Status == vm.ExecStop { + break + } + groups := vector.MustFixedColWithTypeCheck[int32](result.Batch.Vecs[0]) + payloads := vector.MustFixedColWithTypeCheck[int64](result.Batch.Vecs[2]) + for row, group := range groups { + require.True(t, result.Batch.Vecs[1].IsNull(uint64(row))) + require.Contains(t, validPayload[group], payloads[row]) + counts[group]++ + } + } + require.Equal(t, map[int32]int{1: 1, 2: 1}, counts) + arg.Reset(proc, false, nil) + } +} + +func makeAnyNullTopNBatch(t *testing.T, proc *process.Process, groups []int32, payloads []int64, constant bool) *batch.Batch { + t.Helper() + require.Len(t, groups, len(payloads)) + bat := batch.NewWithSize(3) + bat.Vecs[0] = vector.NewVec(types.T_int32.ToType()) + require.NoError(t, vector.AppendFixedList(bat.Vecs[0], groups, nil, proc.Mp())) + if constant { + bat.Vecs[1] = vector.NewConstNull(types.T_any.ToType(), len(groups), proc.Mp()) + } else { + bat.Vecs[1] = vector.NewVec(types.T_any.ToType()) + bat.Vecs[1].SetLength(len(groups)) + nulls.AddRange(bat.Vecs[1].GetNulls(), 0, uint64(len(groups))) + } + bat.Vecs[2] = vector.NewVec(types.T_int64.ToType()) + require.NoError(t, vector.AppendFixedList(bat.Vecs[2], payloads, nil, proc.Mp())) + bat.SetRowCount(len(groups)) + return bat +}