diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/utils/SparkPartitionPredicate.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/utils/SparkPartitionPredicate.scala index 6a92b309fd..dd545907be 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/utils/SparkPartitionPredicate.scala +++ b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/utils/SparkPartitionPredicate.scala @@ -76,6 +76,14 @@ object SparkPartitionPredicate { /** * Tests whether a partition (described by its ordered partition values) matches the given * predicate. Returns true when no predicate is provided. + * + * Callers pass the partition values reported by a lake split, and + * [[org.apache.fluss.lake.source.LakeSplit#partition]] requires one value per partition column + * for a partitioned table — a null/empty list is only legal for a non-partitioned table. An arity + * mismatch therefore means the lake plugin broke that contract, and it is rejected rather than + * silently admitted: the scan builder drops the partition predicate from the post-scan filters it + * hands back to Spark, so a split admitted here is never re-filtered and would leak rows from + * non-matching partitions into the result. */ def matchesPartition( tableInfo: TableInfo, @@ -83,9 +91,15 @@ object SparkPartitionPredicate { partitionPredicate: Option[FlussPredicate]): Boolean = partitionPredicate match { case None => true - case Some(_) if partitionValues.isEmpty => true case Some(predicate) => val rowType = PartitionUtils.partitionRowType(tableInfo) + if (partitionValues.size != rowType.getFieldCount) { + throw new IllegalArgumentException( + s"Cannot evaluate partition filter for table ${tableInfo.getTablePath}: " + + s"expected ${rowType.getFieldCount} partition value(s) for partition key(s) " + + s"${rowType.getFieldNames.asScala.mkString("[", ", ", "]")}, but the lake split " + + s"reported ${partitionValues.size}: ${partitionValues.mkString("[", ", ", "]")}.") + } predicate.test(PartitionUtils.toPartitionRow(partitionValues.asJava, rowType)) } } diff --git a/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/utils/SparkPredicateConverterTest.scala b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/utils/SparkPredicateConverterTest.scala deleted file mode 100644 index 154a2f9382..0000000000 --- a/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/utils/SparkPredicateConverterTest.scala +++ /dev/null @@ -1,278 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.apache.fluss.spark.utils - -import org.apache.fluss.predicate.{CompoundPredicate, LeafPredicate, Predicate => FlussPredicate, PredicateBuilder} -import org.apache.fluss.row.{BinaryString, Decimal, TimestampLtz, TimestampNtz} -import org.apache.fluss.types.{DataField, DataType => FlussDataType, DataTypes, RowType} - -import org.apache.spark.sql.connector.expressions.{Expression, Expressions, Literal, NamedReference} -import org.apache.spark.sql.connector.expressions.filter.{AlwaysFalse, AlwaysTrue, And, Not, Or, Predicate} -import org.apache.spark.sql.types.{BooleanType, DataType, DateType, Decimal => SparkDecimal, DecimalType, IntegerType, StringType, TimestampNTZType, TimestampType} -import org.apache.spark.unsafe.types.UTF8String -import org.assertj.core.api.Assertions.assertThat -import org.scalatest.funsuite.AnyFunSuite - -import java.lang.{Boolean => JBoolean, Byte => JByte, Long => JLong} -import java.time.LocalDate -import java.util.Arrays - -class SparkPredicateConverterTest extends AnyFunSuite { - - private val rowType: RowType = rowTypeOf( - ("id", DataTypes.INT()), - ("name", DataTypes.STRING()), - ("score", DataTypes.DOUBLE()), - ("active", DataTypes.BOOLEAN()), - ("age", DataTypes.TINYINT()), - ("balance", DataTypes.DECIMAL(10, 2)), - ("dt", DataTypes.DATE()), - ("ts", DataTypes.TIMESTAMP()), - ("tsLtz", DataTypes.TIMESTAMP_LTZ()) - ) - - test("EqualTo on integer column converts to equal predicate") { - val predicate = convert(pred("=", ref("id"), lit(Integer.valueOf(42), IntegerType))) - val expected = new PredicateBuilder(rowType).equal(0, Integer.valueOf(42)) - assertThat(predicate).isEqualTo(expected) - } - - test("EqualTo on string column wraps UTF8String as BinaryString") { - val predicate = convert(pred("=", ref("name"), lit(UTF8String.fromString("alice"), StringType))) - val expected = new PredicateBuilder(rowType).equal(1, BinaryString.fromString("alice")) - assertThat(predicate).isEqualTo(expected) - } - - test("EqualTo on null literal returns null predicate value") { - val predicate = convert(pred("=", ref("id"), lit(null, IntegerType))) - val expected = new PredicateBuilder(rowType).equal(0, null.asInstanceOf[Object]) - assertThat(predicate).isEqualTo(expected) - } - - test("EqualNullSafe with null maps to isNull") { - val predicate = convert(pred("<=>", ref("id"), lit(null, IntegerType))) - val expected = new PredicateBuilder(rowType).isNull(0) - assertThat(predicate).isEqualTo(expected) - } - - test("EqualNullSafe with value maps to equal") { - val predicate = convert(pred("<=>", ref("id"), lit(Integer.valueOf(7), IntegerType))) - val expected = new PredicateBuilder(rowType).equal(0, Integer.valueOf(7)) - assertThat(predicate).isEqualTo(expected) - } - - test("GreaterThan, GreaterThanOrEqual, LessThan, LessThanOrEqual") { - val builder = new PredicateBuilder(rowType) - assertThat(convert(pred(">", ref("id"), lit(Integer.valueOf(10), IntegerType)))) - .isEqualTo(builder.greaterThan(0, Integer.valueOf(10))) - assertThat(convert(pred(">=", ref("id"), lit(Integer.valueOf(10), IntegerType)))) - .isEqualTo(builder.greaterOrEqual(0, Integer.valueOf(10))) - assertThat(convert(pred("<", ref("id"), lit(Integer.valueOf(10), IntegerType)))) - .isEqualTo(builder.lessThan(0, Integer.valueOf(10))) - assertThat(convert(pred("<=", ref("id"), lit(Integer.valueOf(10), IntegerType)))) - .isEqualTo(builder.lessOrEqual(0, Integer.valueOf(10))) - } - - test("IsNull and IsNotNull") { - val builder = new PredicateBuilder(rowType) - assertThat(convert(pred("IS_NULL", ref("name")))) - .isEqualTo(builder.isNull(1)) - assertThat(convert(pred("IS_NOT_NULL", ref("name")))) - .isEqualTo(builder.isNotNull(1)) - } - - test("In with multiple string literals") { - val in = pred( - "IN", - ref("name"), - lit(UTF8String.fromString("a"), StringType), - lit(UTF8String.fromString("b"), StringType), - lit(UTF8String.fromString("c"), StringType)) - val expected = new PredicateBuilder(rowType).in( - 1, - Arrays.asList[Object]( - BinaryString.fromString("a"), - BinaryString.fromString("b"), - BinaryString.fromString("c"))) - assertThat(convert(in)).isEqualTo(expected) - } - - test("And composes children") { - val and = new And( - pred("=", ref("id"), lit(Integer.valueOf(1), IntegerType)), - pred("=", ref("name"), lit(UTF8String.fromString("x"), StringType))) - val builder = new PredicateBuilder(rowType) - val expected = PredicateBuilder.and( - builder.equal(0, Integer.valueOf(1)), - builder.equal(1, BinaryString.fromString("x"))) - assertThat(convert(and)).isEqualTo(expected) - } - - test("Or composes children") { - val or = new Or( - pred("=", ref("id"), lit(Integer.valueOf(1), IntegerType)), - pred("=", ref("id"), lit(Integer.valueOf(2), IntegerType))) - val builder = new PredicateBuilder(rowType) - val expected = PredicateBuilder.or( - builder.equal(0, Integer.valueOf(1)), - builder.equal(0, Integer.valueOf(2))) - assertThat(convert(or)).isEqualTo(expected) - } - - test("Not wrapping IsNull maps to isNotNull") { - val not = new Not(pred("IS_NULL", ref("name"))) - assertThat(convert(not)).isEqualTo(new PredicateBuilder(rowType).isNotNull(1)) - } - - test("Not wrapping IsNotNull maps to isNull") { - val not = new Not(pred("IS_NOT_NULL", ref("name"))) - assertThat(convert(not)).isEqualTo(new PredicateBuilder(rowType).isNull(1)) - } - - test("Not wrapping a non-null-check is not supported") { - val not = new Not(pred("=", ref("id"), lit(Integer.valueOf(1), IntegerType))) - assert(SparkPredicateConverter.convert(rowType, not).isEmpty) - } - - test("Not wrapping equality on a boolean column rewrites to its complement") { - val not = new Not(pred("=", ref("active"), lit(JBoolean.TRUE, BooleanType))) - assertThat(convert(not)) - .isEqualTo(new PredicateBuilder(rowType).equal(3, JBoolean.FALSE)) - } - - test("StringStartsWith / EndsWith / Contains on string column") { - val builder = new PredicateBuilder(rowType) - assertThat( - convert(pred("STARTS_WITH", ref("name"), lit(UTF8String.fromString("ali"), StringType)))) - .isEqualTo(builder.startsWith(1, BinaryString.fromString("ali"))) - assertThat( - convert(pred("ENDS_WITH", ref("name"), lit(UTF8String.fromString("son"), StringType)))) - .isEqualTo(builder.endsWith(1, BinaryString.fromString("son"))) - assertThat(convert(pred("CONTAINS", ref("name"), lit(UTF8String.fromString("li"), StringType)))) - .isEqualTo(builder.contains(1, BinaryString.fromString("li"))) - } - - test("StringStartsWith rejected on non-string column") { - val p = pred("STARTS_WITH", ref("id"), lit(UTF8String.fromString("1"), StringType)) - assert(SparkPredicateConverter.convert(rowType, p).isEmpty) - } - - test("Boolean literal") { - val predicate = convert(pred("=", ref("active"), lit(JBoolean.TRUE, BooleanType))) - val expected = new PredicateBuilder(rowType).equal(3, JBoolean.TRUE) - assertThat(predicate).isEqualTo(expected) - } - - test("Tinyint literal from Integer auto-narrows") { - val predicate = convert(pred("=", ref("age"), lit(Integer.valueOf(25), IntegerType))) - val expected = new PredicateBuilder(rowType).equal(4, JByte.valueOf(25.toByte)) - assertThat(predicate).isEqualTo(expected) - } - - test("Decimal literal from Spark Decimal") { - val bd = new java.math.BigDecimal("123.45") - val predicate = - convert(pred("=", ref("balance"), lit(SparkDecimal(bd, 10, 2), DecimalType(10, 2)))) - val expected = new PredicateBuilder(rowType).equal(5, Decimal.fromBigDecimal(bd, 10, 2)) - assertThat(predicate).isEqualTo(expected) - } - - test("Date literal from epoch days") { - val days = Integer.valueOf(LocalDate.of(2025, 1, 15).toEpochDay.toInt) - val predicate = convert(pred("=", ref("dt"), lit(days, DateType))) - val expected = new PredicateBuilder(rowType).equal(6, days) - assertThat(predicate).isEqualTo(expected) - } - - test("Timestamp NTZ literal from epoch micros") { - val micros = 1735639200000000L - val predicate = convert(pred("=", ref("ts"), lit(JLong.valueOf(micros), TimestampNTZType))) - val expected = new PredicateBuilder(rowType).equal(7, TimestampNtz.fromMicros(micros)) - assertThat(predicate).isEqualTo(expected) - } - - test("Timestamp LTZ literal from epoch micros") { - val micros = 1735639200000000L - val predicate = convert(pred("=", ref("tsLtz"), lit(JLong.valueOf(micros), TimestampType))) - val expected = new PredicateBuilder(rowType).equal(8, TimestampLtz.fromEpochMicros(micros)) - assertThat(predicate).isEqualTo(expected) - } - - test("Unknown column returns None") { - val p = pred("=", ref("missing"), lit(Integer.valueOf(1), IntegerType)) - assert(SparkPredicateConverter.convert(rowType, p).isEmpty) - } - - test("AlwaysTrue / AlwaysFalse return None (unsupported)") { - assert(SparkPredicateConverter.convert(rowType, new AlwaysTrue()).isEmpty) - assert(SparkPredicateConverter.convert(rowType, new AlwaysFalse()).isEmpty) - } - - test("convertPredicates returns AND of all convertible and the accepted list") { - val predicates: Seq[Predicate] = Seq( - pred("=", ref("id"), lit(Integer.valueOf(1), IntegerType)), - pred("IS_NOT_NULL", ref("name")), - pred("=", ref("unknown"), lit(Integer.valueOf(1), IntegerType)) - ) - val (predicate, accepted) = SparkPredicateConverter.convertPredicates(rowType, predicates) - assert(predicate.isDefined) - assert(predicate.get.isInstanceOf[CompoundPredicate]) - assert(accepted == Seq(predicates(0), predicates(1))) - } - - test("convertPredicates collapses a single predicate without wrapping in AND") { - val predicates: Seq[Predicate] = Seq(pred("=", ref("id"), lit(Integer.valueOf(1), IntegerType))) - val (predicate, accepted) = SparkPredicateConverter.convertPredicates(rowType, predicates) - assert(predicate.isDefined) - assert(predicate.get.isInstanceOf[LeafPredicate]) - assert(accepted == Seq(predicates.head)) - } - - test("convertPredicates with no convertible predicates returns empty") { - val predicates: Seq[Predicate] = - Seq(pred("=", ref("unknown"), lit(Integer.valueOf(1), IntegerType)), new AlwaysTrue()) - val (predicate, accepted) = SparkPredicateConverter.convertPredicates(rowType, predicates) - assert(predicate.isEmpty) - assert(accepted.isEmpty) - } - - private def convert(predicate: Predicate): FlussPredicate = - SparkPredicateConverter - .convert(rowType, predicate) - .getOrElse(fail(s"Expected predicate $predicate to be convertible")) - - private def ref(name: String): NamedReference = Expressions.column(name) - - private def lit[T](v: T, dt: DataType): Literal[T] = new Literal[T] { - override def value(): T = v - override def dataType(): DataType = dt - override def children(): Array[Expression] = Array.empty - } - - private def pred(name: String, children: Expression*): Predicate = - new Predicate(name, children.toArray) - - private def rowTypeOf(fields: (String, FlussDataType)*): RowType = { - val list = new java.util.ArrayList[DataField]() - fields.foreach { - case (name, tpe) => - list.add(new DataField(name, tpe)) - } - new RowType(list) - } -} diff --git a/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/utils/SparkPredicateUtilsTest.scala b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/utils/SparkPredicateUtilsTest.scala new file mode 100644 index 0000000000..5653510fe1 --- /dev/null +++ b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/utils/SparkPredicateUtilsTest.scala @@ -0,0 +1,561 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.fluss.spark.utils + +import org.apache.fluss.metadata.{PartitionInfo, ResolvedPartitionSpec, Schema, TableDescriptor, TableInfo, TablePath} +import org.apache.fluss.predicate.{CompoundPredicate, LeafPredicate, Predicate => FlussPredicate, PredicateBuilder} +import org.apache.fluss.row.{BinaryString, Decimal, TimestampLtz, TimestampNtz} +import org.apache.fluss.types.{DataField, DataType => FlussDataType, DataTypes, RowType} +import org.apache.fluss.utils.PartitionUtils + +import org.apache.spark.sql.connector.expressions.{Expression, Expressions, Literal, NamedReference} +import org.apache.spark.sql.connector.expressions.filter.{AlwaysFalse, AlwaysTrue, And, Not, Or, Predicate} +import org.apache.spark.sql.types.{BooleanType, DataType, DateType, Decimal => SparkDecimal, DecimalType, IntegerType, StringType, TimestampNTZType, TimestampType} +import org.apache.spark.unsafe.types.UTF8String +import org.assertj.core.api.Assertions.{assertThat, assertThatThrownBy} +import org.scalatest.funsuite.AnyFunSuite + +import java.lang.{Boolean => JBoolean, Byte => JByte, Long => JLong} +import java.time.LocalDate +import java.util.Arrays + +import scala.collection.JavaConverters._ + +class SparkPredicateUtilsTest extends AnyFunSuite { + + private val rowType: RowType = rowTypeOf( + ("id", DataTypes.INT()), + ("name", DataTypes.STRING()), + ("score", DataTypes.DOUBLE()), + ("active", DataTypes.BOOLEAN()), + ("age", DataTypes.TINYINT()), + ("balance", DataTypes.DECIMAL(10, 2)), + ("dt", DataTypes.DATE()), + ("ts", DataTypes.TIMESTAMP()), + ("tsLtz", DataTypes.TIMESTAMP_LTZ()) + ) + + // Log table partitioned by a single STRING column. + private val singleKeyTable: TableInfo = tableInfo( + columns = + Seq(("orderId", DataTypes.BIGINT()), ("amount", DataTypes.INT()), ("dt", DataTypes.STRING())), + partitionKeys = Seq("dt") + ) + + // Log table partitioned by two STRING columns, in declaration order dt, region. + private val multiKeyTable: TableInfo = tableInfo( + columns = Seq( + ("orderId", DataTypes.BIGINT()), + ("dt", DataTypes.STRING()), + ("region", DataTypes.STRING())), + partitionKeys = Seq("dt", "region") + ) + + // Primary key table partitioned by a non-STRING (INT) column. Fluss requires the partition key to + // be a subset of the primary key for a partitioned PK table. + private val pkIntKeyTable: TableInfo = tableInfo( + columns = + Seq(("id", DataTypes.BIGINT()), ("amount", DataTypes.INT()), ("day", DataTypes.INT())), + partitionKeys = Seq("day"), + primaryKeys = Seq("id", "day") + ) + + private val nonPartitionedTable: TableInfo = tableInfo( + columns = Seq(("orderId", DataTypes.BIGINT()), ("amount", DataTypes.INT())), + partitionKeys = Seq.empty + ) + + // --------------------------------------------------------------------------------------------- + // SparkPredicateConverter + // --------------------------------------------------------------------------------------------- + + test("EqualTo on integer column converts to equal predicate") { + val predicate = convert(pred("=", ref("id"), lit(Integer.valueOf(42), IntegerType))) + val expected = new PredicateBuilder(rowType).equal(0, Integer.valueOf(42)) + assertThat(predicate).isEqualTo(expected) + } + + test("EqualTo on string column wraps UTF8String as BinaryString") { + val predicate = convert(pred("=", ref("name"), lit(UTF8String.fromString("alice"), StringType))) + val expected = new PredicateBuilder(rowType).equal(1, BinaryString.fromString("alice")) + assertThat(predicate).isEqualTo(expected) + } + + test("EqualTo on null literal returns null predicate value") { + val predicate = convert(pred("=", ref("id"), lit(null, IntegerType))) + val expected = new PredicateBuilder(rowType).equal(0, null.asInstanceOf[Object]) + assertThat(predicate).isEqualTo(expected) + } + + test("EqualNullSafe with null maps to isNull") { + val predicate = convert(pred("<=>", ref("id"), lit(null, IntegerType))) + val expected = new PredicateBuilder(rowType).isNull(0) + assertThat(predicate).isEqualTo(expected) + } + + test("EqualNullSafe with value maps to equal") { + val predicate = convert(pred("<=>", ref("id"), lit(Integer.valueOf(7), IntegerType))) + val expected = new PredicateBuilder(rowType).equal(0, Integer.valueOf(7)) + assertThat(predicate).isEqualTo(expected) + } + + test("GreaterThan, GreaterThanOrEqual, LessThan, LessThanOrEqual") { + val builder = new PredicateBuilder(rowType) + assertThat(convert(pred(">", ref("id"), lit(Integer.valueOf(10), IntegerType)))) + .isEqualTo(builder.greaterThan(0, Integer.valueOf(10))) + assertThat(convert(pred(">=", ref("id"), lit(Integer.valueOf(10), IntegerType)))) + .isEqualTo(builder.greaterOrEqual(0, Integer.valueOf(10))) + assertThat(convert(pred("<", ref("id"), lit(Integer.valueOf(10), IntegerType)))) + .isEqualTo(builder.lessThan(0, Integer.valueOf(10))) + assertThat(convert(pred("<=", ref("id"), lit(Integer.valueOf(10), IntegerType)))) + .isEqualTo(builder.lessOrEqual(0, Integer.valueOf(10))) + } + + test("IsNull and IsNotNull") { + val builder = new PredicateBuilder(rowType) + assertThat(convert(pred("IS_NULL", ref("name")))) + .isEqualTo(builder.isNull(1)) + assertThat(convert(pred("IS_NOT_NULL", ref("name")))) + .isEqualTo(builder.isNotNull(1)) + } + + test("In with multiple string literals") { + val in = pred( + "IN", + ref("name"), + lit(UTF8String.fromString("a"), StringType), + lit(UTF8String.fromString("b"), StringType), + lit(UTF8String.fromString("c"), StringType)) + val expected = new PredicateBuilder(rowType).in( + 1, + Arrays.asList[Object]( + BinaryString.fromString("a"), + BinaryString.fromString("b"), + BinaryString.fromString("c"))) + assertThat(convert(in)).isEqualTo(expected) + } + + test("And composes children") { + val and = new And( + pred("=", ref("id"), lit(Integer.valueOf(1), IntegerType)), + pred("=", ref("name"), lit(UTF8String.fromString("x"), StringType))) + val builder = new PredicateBuilder(rowType) + val expected = PredicateBuilder.and( + builder.equal(0, Integer.valueOf(1)), + builder.equal(1, BinaryString.fromString("x"))) + assertThat(convert(and)).isEqualTo(expected) + } + + test("Or composes children") { + val or = new Or( + pred("=", ref("id"), lit(Integer.valueOf(1), IntegerType)), + pred("=", ref("id"), lit(Integer.valueOf(2), IntegerType))) + val builder = new PredicateBuilder(rowType) + val expected = PredicateBuilder.or( + builder.equal(0, Integer.valueOf(1)), + builder.equal(0, Integer.valueOf(2))) + assertThat(convert(or)).isEqualTo(expected) + } + + test("Not wrapping IsNull maps to isNotNull") { + val not = new Not(pred("IS_NULL", ref("name"))) + assertThat(convert(not)).isEqualTo(new PredicateBuilder(rowType).isNotNull(1)) + } + + test("Not wrapping IsNotNull maps to isNull") { + val not = new Not(pred("IS_NOT_NULL", ref("name"))) + assertThat(convert(not)).isEqualTo(new PredicateBuilder(rowType).isNull(1)) + } + + test("Not wrapping a non-null-check is not supported") { + val not = new Not(pred("=", ref("id"), lit(Integer.valueOf(1), IntegerType))) + assert(SparkPredicateConverter.convert(rowType, not).isEmpty) + } + + test("Not wrapping equality on a boolean column rewrites to its complement") { + val not = new Not(pred("=", ref("active"), lit(JBoolean.TRUE, BooleanType))) + assertThat(convert(not)) + .isEqualTo(new PredicateBuilder(rowType).equal(3, JBoolean.FALSE)) + } + + test("StringStartsWith / EndsWith / Contains on string column") { + val builder = new PredicateBuilder(rowType) + assertThat( + convert(pred("STARTS_WITH", ref("name"), lit(UTF8String.fromString("ali"), StringType)))) + .isEqualTo(builder.startsWith(1, BinaryString.fromString("ali"))) + assertThat( + convert(pred("ENDS_WITH", ref("name"), lit(UTF8String.fromString("son"), StringType)))) + .isEqualTo(builder.endsWith(1, BinaryString.fromString("son"))) + assertThat(convert(pred("CONTAINS", ref("name"), lit(UTF8String.fromString("li"), StringType)))) + .isEqualTo(builder.contains(1, BinaryString.fromString("li"))) + } + + test("StringStartsWith rejected on non-string column") { + val p = pred("STARTS_WITH", ref("id"), lit(UTF8String.fromString("1"), StringType)) + assert(SparkPredicateConverter.convert(rowType, p).isEmpty) + } + + test("Boolean literal") { + val predicate = convert(pred("=", ref("active"), lit(JBoolean.TRUE, BooleanType))) + val expected = new PredicateBuilder(rowType).equal(3, JBoolean.TRUE) + assertThat(predicate).isEqualTo(expected) + } + + test("Tinyint literal from Integer auto-narrows") { + val predicate = convert(pred("=", ref("age"), lit(Integer.valueOf(25), IntegerType))) + val expected = new PredicateBuilder(rowType).equal(4, JByte.valueOf(25.toByte)) + assertThat(predicate).isEqualTo(expected) + } + + test("Decimal literal from Spark Decimal") { + val bd = new java.math.BigDecimal("123.45") + val predicate = + convert(pred("=", ref("balance"), lit(SparkDecimal(bd, 10, 2), DecimalType(10, 2)))) + val expected = new PredicateBuilder(rowType).equal(5, Decimal.fromBigDecimal(bd, 10, 2)) + assertThat(predicate).isEqualTo(expected) + } + + test("Date literal from epoch days") { + val days = Integer.valueOf(LocalDate.of(2025, 1, 15).toEpochDay.toInt) + val predicate = convert(pred("=", ref("dt"), lit(days, DateType))) + val expected = new PredicateBuilder(rowType).equal(6, days) + assertThat(predicate).isEqualTo(expected) + } + + test("Timestamp NTZ literal from epoch micros") { + val micros = 1735639200000000L + val predicate = convert(pred("=", ref("ts"), lit(JLong.valueOf(micros), TimestampNTZType))) + val expected = new PredicateBuilder(rowType).equal(7, TimestampNtz.fromMicros(micros)) + assertThat(predicate).isEqualTo(expected) + } + + test("Timestamp LTZ literal from epoch micros") { + val micros = 1735639200000000L + val predicate = convert(pred("=", ref("tsLtz"), lit(JLong.valueOf(micros), TimestampType))) + val expected = new PredicateBuilder(rowType).equal(8, TimestampLtz.fromEpochMicros(micros)) + assertThat(predicate).isEqualTo(expected) + } + + test("Unknown column returns None") { + val p = pred("=", ref("missing"), lit(Integer.valueOf(1), IntegerType)) + assert(SparkPredicateConverter.convert(rowType, p).isEmpty) + } + + test("AlwaysTrue / AlwaysFalse return None (unsupported)") { + assert(SparkPredicateConverter.convert(rowType, new AlwaysTrue()).isEmpty) + assert(SparkPredicateConverter.convert(rowType, new AlwaysFalse()).isEmpty) + } + + test("convertPredicates returns AND of all convertible and the accepted list") { + val predicates: Seq[Predicate] = Seq( + pred("=", ref("id"), lit(Integer.valueOf(1), IntegerType)), + pred("IS_NOT_NULL", ref("name")), + pred("=", ref("unknown"), lit(Integer.valueOf(1), IntegerType)) + ) + val (predicate, accepted) = SparkPredicateConverter.convertPredicates(rowType, predicates) + assert(predicate.isDefined) + assert(predicate.get.isInstanceOf[CompoundPredicate]) + assert(accepted == Seq(predicates(0), predicates(1))) + } + + test("convertPredicates collapses a single predicate without wrapping in AND") { + val predicates: Seq[Predicate] = Seq(pred("=", ref("id"), lit(Integer.valueOf(1), IntegerType))) + val (predicate, accepted) = SparkPredicateConverter.convertPredicates(rowType, predicates) + assert(predicate.isDefined) + assert(predicate.get.isInstanceOf[LeafPredicate]) + assert(accepted == Seq(predicates.head)) + } + + test("convertPredicates with no convertible predicates returns empty") { + val predicates: Seq[Predicate] = + Seq(pred("=", ref("unknown"), lit(Integer.valueOf(1), IntegerType)), new AlwaysTrue()) + val (predicate, accepted) = SparkPredicateConverter.convertPredicates(rowType, predicates) + assert(predicate.isEmpty) + assert(accepted.isEmpty) + } + + // --------------------------------------------------------------------------------------------- + // SparkPartitionPredicate.matchesPartition + // --------------------------------------------------------------------------------------------- + + test("matchesPartition accepts a partition whose values satisfy the predicate") { + val predicate = partitionPredicateOf(singleKeyTable, dtEquals("2026-01-01")) + assertThat( + SparkPartitionPredicate + .matchesPartition(singleKeyTable, Seq("2026-01-01"), predicate)).isTrue + } + + test("matchesPartition rejects a partition whose values violate the predicate") { + val predicate = partitionPredicateOf(singleKeyTable, dtEquals("2026-01-01")) + assertThat( + SparkPartitionPredicate + .matchesPartition(singleKeyTable, Seq("2026-01-02"), predicate)).isFalse + } + + // Regression guard for the lake-only split branch of AppendPlanner.planLakePartitionedTable / + // UpsertPlanner.planLakePartitionedTable: a LakeSplit that reports no partition for a partitioned + // table used to be admitted unconditionally. Because the scan builder removes the partition + // predicate from the post-scan filters returned to Spark, such a split is never re-filtered, so + // admitting it silently leaks rows from non-matching partitions. + test("matchesPartition rejects empty partition values when a predicate is present") { + val predicate = partitionPredicateOf(singleKeyTable, dtEquals("2026-01-01")) + assertThatThrownBy( + () => + SparkPartitionPredicate + .matchesPartition(singleKeyTable, Seq.empty, predicate)) + .isInstanceOf(classOf[IllegalArgumentException]) + .hasMessageContaining("expected 1 partition value(s)") + .hasMessageContaining("[dt]") + .hasMessageContaining("reported 0") + } + + test("matchesPartition rejects a partial partition value tuple on a multi-key table") { + val predicate = partitionPredicateOf(multiKeyTable, dtEquals("2026-01-01")) + assertThatThrownBy( + () => + SparkPartitionPredicate + .matchesPartition(multiKeyTable, Seq("2026-01-01"), predicate)) + .isInstanceOf(classOf[IllegalArgumentException]) + .hasMessageContaining("expected 2 partition value(s)") + .hasMessageContaining("[dt, region]") + .hasMessageContaining("reported 1") + } + + test("matchesPartition rejects a partition value tuple longer than the partition keys") { + val predicate = partitionPredicateOf(singleKeyTable, dtEquals("2026-01-01")) + assertThatThrownBy( + () => + SparkPartitionPredicate + .matchesPartition(singleKeyTable, Seq("2026-01-01", "cn"), predicate)) + .isInstanceOf(classOf[IllegalArgumentException]) + .hasMessageContaining("reported 2") + } + + // Without a predicate there is nothing to prune, so arity is irrelevant and the split must still + // be admitted. This pins down that the guard above did not become over-strict. + test("matchesPartition admits any partition when no predicate is present") { + assertThat(SparkPartitionPredicate.matchesPartition(singleKeyTable, Seq.empty, None)).isTrue + assertThat( + SparkPartitionPredicate.matchesPartition(singleKeyTable, Seq("2026-01-01"), None)).isTrue + assertThat( + SparkPartitionPredicate.matchesPartition(multiKeyTable, Seq("only-one"), None)).isTrue + } + + test("matchesPartition evaluates all keys of a multi-key partition") { + val predicate = partitionPredicateOf( + multiKeyTable, + pred("=", ref("dt"), lit(UTF8String.fromString("2026-01-01"), StringType)), + pred("=", ref("region"), lit(UTF8String.fromString("cn"), StringType)) + ) + assertThat( + SparkPartitionPredicate + .matchesPartition(multiKeyTable, Seq("2026-01-01", "cn"), predicate)).isTrue + // second key differs + assertThat( + SparkPartitionPredicate + .matchesPartition(multiKeyTable, Seq("2026-01-01", "us"), predicate)).isFalse + } + + // Partition values always arrive as strings and are parsed against the partition row type, so a + // non-STRING partition key must still compare correctly against the Spark literal. + test("matchesPartition parses non-string partition values against the partition row type") { + val predicate = partitionPredicateOf( + pkIntKeyTable, + pred(">", ref("day"), lit(Integer.valueOf(20260101), IntegerType))) + assertThat( + SparkPartitionPredicate.matchesPartition(pkIntKeyTable, Seq("20260102"), predicate)).isTrue + assertThat( + SparkPartitionPredicate.matchesPartition(pkIntKeyTable, Seq("20260101"), predicate)).isFalse + } + + // --------------------------------------------------------------------------------------------- + // SparkPartitionPredicate.extract + // --------------------------------------------------------------------------------------------- + + test("extract returns all predicates as non-partition for a non-partitioned table") { + val predicates = Seq(pred("=", ref("amount"), lit(Integer.valueOf(1), IntegerType))) + val (nonPartition, partition) = + SparkPartitionPredicate.extract(nonPartitionedTable, predicates) + assert(nonPartition == predicates) + assertThat(partition.isDefined).isFalse + } + + test("extract separates a partition-key predicate from a data predicate") { + val dtPred = dtEquals("2026-01-01") + val amountPred = pred(">", ref("amount"), lit(Integer.valueOf(600), IntegerType)) + val (nonPartition, partition) = + SparkPartitionPredicate.extract(singleKeyTable, Seq(dtPred, amountPred)) + + assert(nonPartition == Seq(amountPred)) + assertThat(partition.isDefined).isTrue + assertThat(partition.get).isInstanceOf(classOf[LeafPredicate]) + // Field index is relative to the partition row type, where dt is the only column. + assertThat(partition.get) + .isEqualTo( + new PredicateBuilder(PartitionUtils.partitionRowType(singleKeyTable)) + .equal(0, BinaryString.fromString("2026-01-01"))) + } + + test("extract AND-combines multiple partition-key predicates") { + val dtPred = dtEquals("2026-01-01") + val regionPred = pred("=", ref("region"), lit(UTF8String.fromString("cn"), StringType)) + val (nonPartition, partition) = + SparkPartitionPredicate.extract(multiKeyTable, Seq(dtPred, regionPred)) + + assertThat(nonPartition.isEmpty).isTrue + assertThat(partition.isDefined).isTrue + assertThat(partition.get).isInstanceOf(classOf[CompoundPredicate]) + } + + // An OR spanning a partition key and a data column cannot be answered by partition pruning alone, + // so it must stay in the non-partition list for Spark to re-apply. + test("extract keeps an OR mixing partition and non-partition columns as non-partition") { + val mixed: Predicate = new Or( + dtEquals("2026-01-01"), + pred(">", ref("amount"), lit(Integer.valueOf(600), IntegerType))) + val (nonPartition, partition) = SparkPartitionPredicate.extract(singleKeyTable, Seq(mixed)) + + assert(nonPartition == Seq(mixed)) + assertThat(partition.isDefined).isFalse + } + + test("extract keeps an unconvertible predicate as non-partition") { + // NOT on a string equality is not invertible by SparkPredicateConverter. + val notEq: Predicate = new Not(dtEquals("2026-01-01")) + val (nonPartition, partition) = SparkPartitionPredicate.extract(singleKeyTable, Seq(notEq)) + + assert(nonPartition == Seq(notEq)) + assertThat(partition.isDefined).isFalse + } + + // --------------------------------------------------------------------------------------------- + // SparkPartitionPredicate.filterPartitions + // --------------------------------------------------------------------------------------------- + + test("filterPartitions keeps only the partitions matching the predicate") { + val partitions = Seq( + partitionInfo(1L, singleKeyTable, Seq("2026-01-01")), + partitionInfo(2L, singleKeyTable, Seq("2026-01-02")), + partitionInfo(3L, singleKeyTable, Seq("2026-01-03")) + ) + val predicate = partitionPredicateOf(singleKeyTable, dtEquals("2026-01-02")) + + val kept = SparkPartitionPredicate.filterPartitions(singleKeyTable, partitions, predicate) + assert(kept.map(_.getPartitionId) == Seq(2L)) + } + + test("filterPartitions returns every partition when no predicate is present") { + val partitions = Seq( + partitionInfo(1L, singleKeyTable, Seq("2026-01-01")), + partitionInfo(2L, singleKeyTable, Seq("2026-01-02"))) + val kept = SparkPartitionPredicate.filterPartitions(singleKeyTable, partitions, None) + assert(kept.map(_.getPartitionId) == Seq(1L, 2L)) + } + + test("filterPartitions can prune everything away") { + val partitions = Seq(partitionInfo(1L, singleKeyTable, Seq("2026-01-01"))) + val predicate = partitionPredicateOf(singleKeyTable, dtEquals("2099-01-01")) + val kept = SparkPartitionPredicate.filterPartitions(singleKeyTable, partitions, predicate) + assertThat(kept.isEmpty).isTrue + } + + test("filterPartitions prunes a multi-key partition list") { + val partitions = Seq( + partitionInfo(1L, multiKeyTable, Seq("2026-01-01", "cn")), + partitionInfo(2L, multiKeyTable, Seq("2026-01-01", "us")), + partitionInfo(3L, multiKeyTable, Seq("2026-01-02", "cn")) + ) + val predicate = partitionPredicateOf( + multiKeyTable, + dtEquals("2026-01-01"), + pred("=", ref("region"), lit(UTF8String.fromString("cn"), StringType))) + + val kept = SparkPartitionPredicate.filterPartitions(multiKeyTable, partitions, predicate) + assert(kept.map(_.getPartitionId) == Seq(1L)) + } + + // --------------------------------------------------------------------------------------------- + // helpers + // --------------------------------------------------------------------------------------------- + + private def convert(predicate: Predicate): FlussPredicate = + SparkPredicateConverter + .convert(rowType, predicate) + .getOrElse(fail(s"Expected predicate $predicate to be convertible")) + + /** Runs the predicates through `extract` and returns the partition predicate it produced. */ + private def partitionPredicateOf( + table: TableInfo, + predicates: Predicate*): Option[org.apache.fluss.predicate.Predicate] = { + val (_, partition) = SparkPartitionPredicate.extract(table, predicates) + assert(partition.isDefined, s"Expected $predicates to yield a partition predicate") + partition + } + + private def dtEquals(value: String): Predicate = + pred("=", ref("dt"), lit(UTF8String.fromString(value), StringType)) + + private def partitionInfo( + partitionId: Long, + table: TableInfo, + partitionValues: Seq[String]): PartitionInfo = + new PartitionInfo( + partitionId, + new ResolvedPartitionSpec(table.getPartitionKeys, partitionValues.asJava), + null) + + private def tableInfo( + columns: Seq[(String, org.apache.fluss.types.DataType)], + partitionKeys: Seq[String], + primaryKeys: Seq[String] = Seq.empty): TableInfo = { + val schemaBuilder = Schema.newBuilder() + columns.foreach { case (name, tpe) => schemaBuilder.column(name, tpe) } + if (primaryKeys.nonEmpty) { + schemaBuilder.primaryKey(primaryKeys.asJava) + } + val descriptorBuilder = TableDescriptor + .builder() + .schema(schemaBuilder.build()) + .distributedBy(1) + if (partitionKeys.nonEmpty) { + descriptorBuilder.partitionedBy(partitionKeys.asJava) + } + val now = System.currentTimeMillis() + TableInfo.of(TablePath.of("db", "t"), 1L, 1, descriptorBuilder.build(), null, now, now) + } + + private def ref(name: String): NamedReference = Expressions.column(name) + + private def lit[T](v: T, dt: DataType): Literal[T] = new Literal[T] { + override def value(): T = v + override def dataType(): DataType = dt + override def children(): Array[Expression] = Array.empty + } + + private def pred(name: String, children: Expression*): Predicate = + new Predicate(name, children.toArray) + + private def rowTypeOf(fields: (String, FlussDataType)*): RowType = { + val list = new java.util.ArrayList[DataField]() + fields.foreach { + case (name, tpe) => + list.add(new DataField(name, tpe)) + } + new RowType(list) + } +}