HyukjinKwon commented on code in PR #59929: URL: https://github.com/apache/spark/pull/59929#discussion_r4461914182
########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnDomainsSuite.scala: ########## @@ -0,0 +1,381 @@ +/* + * 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.spark.sql.execution.datasources.v2.columnar + +import java.sql.Date + +import scala.jdk.CollectionConverters._ +import scala.util.Random + +import org.scalatest.concurrent.TimeLimits +import org.scalatest.time.SpanSugar._ + +import org.apache.spark.sql.{Column, QueryTest, Row} +import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} +import org.apache.spark.sql.catalyst.expressions.{Literal, PredicateHelper} +import org.apache.spark.sql.catalyst.plans.logical.Filter +import org.apache.spark.sql.connector.expressions.{Expression => V2Expression, FieldReference, LiteralValue} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Strategy +import org.apache.spark.sql.functions.{col, expr, not} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.UTF8String + +/** Tests how [[ColumnDomains]] derives the domains of columns from the predicates of a query. */ +class ColumnDomainsSuite extends QueryTest with SharedSparkSession with PredicateHelper + with TimeLimits { + + private val schema = StructType.fromDDL("id BIGINT, i INT, d DOUBLE, f FLOAT, s STRING, " + + "b BINARY, flag BOOLEAN, dt DATE, ts TIMESTAMP, dec DECIMAL(10, 2), " + + "st STRUCT<x: INT, y: STRING>, arr ARRAY<INT>, c STRING COLLATE UTF8_LCASE") + + private def char(codePoint: Int): String = new String(Character.toChars(codePoint)) + + /** + * The predicates that Data Source V2 pushes for the filter, as in the analyzed plan with the + * constants folded: the optimizer would rewrite some of them, such as NOT (x > 1). + */ + private def predicates(filter: String, schema: StructType): Seq[Predicate] = { + predicates(expr(filter), schema) + } + + private def predicates(filter: Column, schema: StructType): Seq[Predicate] = { + val df = spark.createDataFrame(Seq.empty[Row].asJava, schema) + val condition = df.where(filter).queryExecution.analyzed.collectFirst { + case Filter(condition, _) => condition + }.get + splitConjunctivePredicates(condition.transformUp { + case e if e.foldable && !e.isInstanceOf[Literal] => Literal.create(e.eval(), e.dataType) + }).flatMap(DataSourceV2Strategy.translateFilterV2) + } + + private def domains(filter: String): Seq[String] = domains(expr(filter)) + + private def domains(filter: Column): Seq[String] = { + ColumnDomains.derive(predicates(filter, schema), schema).domains.map(_.toString) + } + + test("comparisons, IN and null checks") { + assert(domains("id > 5") == Seq("id: (5L, +inf)")) + assert(domains("5 < id") == Seq("id: (5L, +inf)")) + assert(domains("id >= 5 AND id < 10") == Seq("id: [5L, 10L)")) + assert(domains("id <= 5") == Seq("id: (-inf, 5L]")) + assert(domains("id = 5") == Seq("id: 5L")) + assert(domains("id <> 5") == Seq("id: (-inf, 5L), (5L, +inf)")) + assert(domains("id <=> 5") == Seq("id: 5L")) + // x <=> v is false, not null, when x is null, so NOT (x <=> v) is true for null. + assert(domains("NOT (id <=> 5)") == Seq("id: NULL, (-inf, 5L), (5L, +inf)")) + assert(domains("id <=> NULL") == Seq("id: NULL")) + assert(domains("NOT (id <=> NULL)") == Seq("id: ANY")) + assert(domains("id IS NULL") == Seq("id: NULL")) + assert(domains("id IS NOT NULL") == Seq("id: ANY")) + assert(domains("id IN (3, 1, 2, 3)") == Seq("id: 1L, 2L, 3L")) + assert(domains("id NOT IN (1, 3)") == Seq("id: (-inf, 1L), (1L, 3L), (3L, +inf)")) + // Comparisons with null are null: so is NOT (x IN (..., NULL)), which matches no row. Review Comment: Done in 9fc11f6, as suggested. ########## sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/columnar/ColumnDomainsSuite.scala: ########## @@ -0,0 +1,381 @@ +/* + * 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.spark.sql.execution.datasources.v2.columnar + +import java.sql.Date + +import scala.jdk.CollectionConverters._ +import scala.util.Random + +import org.scalatest.concurrent.TimeLimits +import org.scalatest.time.SpanSugar._ + +import org.apache.spark.sql.{Column, QueryTest, Row} +import org.apache.spark.sql.catalyst.{CatalystTypeConverters, InternalRow} +import org.apache.spark.sql.catalyst.expressions.{Literal, PredicateHelper} +import org.apache.spark.sql.catalyst.plans.logical.Filter +import org.apache.spark.sql.connector.expressions.{Expression => V2Expression, FieldReference, LiteralValue} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Strategy +import org.apache.spark.sql.functions.{col, expr, not} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.UTF8String + +/** Tests how [[ColumnDomains]] derives the domains of columns from the predicates of a query. */ +class ColumnDomainsSuite extends QueryTest with SharedSparkSession with PredicateHelper + with TimeLimits { + + private val schema = StructType.fromDDL("id BIGINT, i INT, d DOUBLE, f FLOAT, s STRING, " + + "b BINARY, flag BOOLEAN, dt DATE, ts TIMESTAMP, dec DECIMAL(10, 2), " + + "st STRUCT<x: INT, y: STRING>, arr ARRAY<INT>, c STRING COLLATE UTF8_LCASE") + + private def char(codePoint: Int): String = new String(Character.toChars(codePoint)) + + /** + * The predicates that Data Source V2 pushes for the filter, as in the analyzed plan with the + * constants folded: the optimizer would rewrite some of them, such as NOT (x > 1). + */ + private def predicates(filter: String, schema: StructType): Seq[Predicate] = { + predicates(expr(filter), schema) + } + + private def predicates(filter: Column, schema: StructType): Seq[Predicate] = { + val df = spark.createDataFrame(Seq.empty[Row].asJava, schema) + val condition = df.where(filter).queryExecution.analyzed.collectFirst { + case Filter(condition, _) => condition + }.get + splitConjunctivePredicates(condition.transformUp { + case e if e.foldable && !e.isInstanceOf[Literal] => Literal.create(e.eval(), e.dataType) + }).flatMap(DataSourceV2Strategy.translateFilterV2) + } + + private def domains(filter: String): Seq[String] = domains(expr(filter)) + + private def domains(filter: Column): Seq[String] = { + ColumnDomains.derive(predicates(filter, schema), schema).domains.map(_.toString) + } + + test("comparisons, IN and null checks") { + assert(domains("id > 5") == Seq("id: (5L, +inf)")) + assert(domains("5 < id") == Seq("id: (5L, +inf)")) + assert(domains("id >= 5 AND id < 10") == Seq("id: [5L, 10L)")) + assert(domains("id <= 5") == Seq("id: (-inf, 5L]")) + assert(domains("id = 5") == Seq("id: 5L")) + assert(domains("id <> 5") == Seq("id: (-inf, 5L), (5L, +inf)")) + assert(domains("id <=> 5") == Seq("id: 5L")) + // x <=> v is false, not null, when x is null, so NOT (x <=> v) is true for null. + assert(domains("NOT (id <=> 5)") == Seq("id: NULL, (-inf, 5L), (5L, +inf)")) + assert(domains("id <=> NULL") == Seq("id: NULL")) + assert(domains("NOT (id <=> NULL)") == Seq("id: ANY")) + assert(domains("id IS NULL") == Seq("id: NULL")) + assert(domains("id IS NOT NULL") == Seq("id: ANY")) + assert(domains("id IN (3, 1, 2, 3)") == Seq("id: 1L, 2L, 3L")) + assert(domains("id NOT IN (1, 3)") == Seq("id: (-inf, 1L), (1L, 3L), (3L, +inf)")) + // Comparisons with null are null: so is NOT (x IN (..., NULL)), which matches no row. + assert(domains("id = NULL") == Seq("id: NONE")) + assert(domains("NOT (id = NULL)") == Seq("id: NONE")) + assert(domains("id IN (1, NULL)") == Seq("id: 1L")) + assert(domains("NOT (id IN (1, NULL))") == Seq("id: NONE")) + // x IN (), which only the DataFrame API builds, is never true: it is false, and null for null + // with spark.sql.legacy.nullInEmptyListBehavior. So NOT (x IN ()) constrains no column. + assert(domains(col("id").isin()) == Seq("id: NONE")) + assert(domains(not(col("id").isin())).isEmpty) + assert(domains(not(col("id").isInCollection(Nil)) && col("id") > 5) == Seq("id: (5L, +inf)")) + assert(domains("id = 1 AND id = 2") == Seq("id: NONE")) + // The domains of all the predicates, which are combined with AND, at most one per column. + assert(domains("id > 5 AND i = 1 AND id < 8") == Seq("id: (5L, 8L)", "i: 1")) + assert(domains("true").isEmpty) + assert(ColumnDomains.derive(predicates("false", schema), schema).domains.isEmpty) + } + + test("AND, OR and NOT") { + assert(domains("id > 5 OR id < 0") == Seq("id: (-inf, 0L), (5L, +inf)")) + assert(domains("id > 5 OR id > 3") == Seq("id: (3L, +inf)")) + assert(domains("id >= 3 OR id = 3 OR id < 1") == Seq("id: (-inf, 1L), [3L, +inf)")) + // [1, 2) and [2, 3] touch, and so do (-inf, 5) and [5, +inf), which then constrain id to not + // be null only. + assert(domains("(id >= 1 AND id < 2) OR (id >= 2 AND id <= 3)") == Seq("id: [1L, 3L]")) + assert(domains("id < 5 OR id >= 5") == Seq("id: ANY")) + assert(domains("id < 5 OR id >= 5 OR id IS NULL").isEmpty) + // A column that only one side of OR constrains can have any value in the rows of the other. + assert(domains("id > 5 OR i > 5").isEmpty) + assert(domains("id > 5 AND (i > 1 OR id < 3)") == Seq("id: (5L, +inf)")) + assert(domains("(id > 5 AND i = 1) OR (id < 0 AND i = 2)") == + Seq("id: (-inf, 0L), (5L, +inf)", "i: 1, 2")) + // A side that matches no row does not widen the other. + assert(domains("id = NULL OR i > 5") == Seq("i: (5, +inf)")) + assert(domains("NOT (id > 5 OR i IS NULL)") == Seq("id: (-inf, 5L]", "i: ANY")) + assert(domains("NOT (id > 5 AND i < 2)").isEmpty) + assert(domains("NOT NOT (id > 5)") == Seq("id: (5L, +inf)")) + // Predicates on other expressions constrain no column. + assert(domains("id + 1 > 5").isEmpty) + assert(domains("id = i").isEmpty) + assert(domains("id + 1 > 5 AND id < 3") == Seq("id: (-inf, 3L)")) + assert(domains("NOT (id + 1 > 5) AND id < 3") == Seq("id: (-inf, 3L)")) + } + + test("NaN and -0.0") { + // NaN is equal to NaN and greater than any other value, including positive infinity. + assert(domains("d > 1.0") == Seq("d: NaN, (1.0D, +inf)")) + assert(domains("d < 1.0") == Seq("d: (-inf, 1.0D)")) + assert(domains("d <> 1.0") == Seq("d: NaN, (-inf, 1.0D), (1.0D, +inf)")) + assert(domains("NOT (d < 1.0)") == Seq("d: NaN, [1.0D, +inf)")) + assert(domains("d > CAST('Infinity' AS DOUBLE)") == + Seq("d: NaN, (CAST('Infinity' AS DOUBLE), +inf)")) + assert(domains("d = CAST('NaN' AS DOUBLE)") == Seq("d: NaN")) + assert(domains("d <> CAST('NaN' AS DOUBLE)") == Seq("d: ANY")) + assert(domains("d > CAST('NaN' AS DOUBLE)") == Seq("d: NONE")) + assert(domains("d >= CAST('NaN' AS DOUBLE)") == Seq("d: NaN")) + assert(domains("d < CAST('NaN' AS DOUBLE)") == Seq("d: ANY")) + assert(domains("d <= CAST('NaN' AS DOUBLE)") == Seq("d: NaN, ANY")) + assert(domains("d IN (1.0, CAST('NaN' AS DOUBLE))") == Seq("d: NaN, 1.0D")) + assert(domains("d NOT IN (1.0)") == Seq("d: NaN, (-inf, 1.0D), (1.0D, +inf)")) + assert(domains("d IS NOT NULL") == Seq("d: NaN, ANY")) + // -0.0 is equal to 0.0, which a data source may order after -0.0. + assert(domains("d = -0.0") == Seq("d: 0.0D")) Review Comment: Done in 9fc11f6: the cases now use `d = -0.0D` and `f > CAST('-0.0' AS FLOAT)`, with the same expected strings, and a comment says why. With `normalize` removed, the double case now fails with `List("d: -0.0D") did not equal List("d: 0.0D")`. ########## sql/core/src/test/resources/native-datasource/test_native_datasource.cc: ########## @@ -1075,6 +1083,108 @@ int pruneColumns(void* handle, const char* const* columns, int64_t count, bool* }); } +bool boolAt(const ArrowArray* array, int64_t i) { + int64_t bit = array->offset + i; + return (static_cast<const uint8_t*>(array->buffers[1])[bit / 8] >> (bit % 8)) & 1; +} + +// Formats the ranges of a domain, such as "[1,5) (7,+inf)", after checking their schema. +std::string formatRanges(const SparkDsDomain* domain, const std::string& format) { + const ArrowSchema* schema = domain->ranges_schema; + const ArrowArray* array = domain->ranges; + if (schema == nullptr || std::string(schema->format) != "+s" || schema->n_children != 4 || + array->n_children != 4) { + throw std::runtime_error("the ranges are not a struct array of 4 fields"); + } + const char* names[] = {"low", "high", "low_inclusive", "high_inclusive"}; + for (int i = 0; i < 4; i++) { + std::string expected = i < 2 ? format : "b"; + if (std::string(schema->children[i]->name) != names[i] || + (!expected.empty() && std::string(schema->children[i]->format) != expected)) { + throw std::runtime_error(std::string("unexpected field ") + schema->children[i]->name + + " of format " + schema->children[i]->format); + } + } + std::string ranges; + for (int64_t r = 0; r < array->length; r++) { + bool low_inclusive = boolAt(array->children[2], r); + bool high_inclusive = boolAt(array->children[3], r); + std::string low = isValid(array->children[0], r) + ? formatValue(schema->children[0], array->children[0], r) : "-inf"; + std::string high = isValid(array->children[1], r) + ? formatValue(schema->children[1], array->children[1], r) : "+inf"; + ranges += std::string(r > 0 ? " " : "") + (low_inclusive ? "[" : "(") + low + "," + high + + (high_inclusive ? "]" : ")"); + } + return ranges; +} + +// The ids [lo, hi) that the domain of id has, as "lo,hi;lo,hi". +std::string idRanges(const SparkDsDomain* domain) { + if (domain->ranges == nullptr) return std::to_string(INT64_MIN) + "," + std::to_string(INT64_MAX); + const ArrowArray* array = domain->ranges; + std::string result; + for (int64_t r = 0; r < array->length; r++) { + const ArrowArray* low = array->children[0]; + const ArrowArray* high = array->children[1]; + int64_t lo = INT64_MIN; + int64_t hi = INT64_MAX; + if (isValid(low, r)) { + lo = static_cast<const int64_t*>(low->buffers[1])[low->offset + r]; + if (!boolAt(array->children[2], r) && lo < INT64_MAX) lo++; + } + if (isValid(high, r)) { + hi = static_cast<const int64_t*>(high->buffers[1])[high->offset + r]; + if (boolAt(array->children[3], r) && hi < INT64_MAX) hi++; + } + result += (result.empty() ? "" : ";") + std::to_string(lo) + "," + std::to_string(hi); + } + return result; +} + +int pushDomains(void* handle, const SparkDsDomain* const* domains, int64_t count, bool* used) { Review Comment: Done in 9fc11f6: the test library fails `push_domains` with the option `fail` set to `push_domains`. "report the errors of native code" checks that the query fails with `NATIVE_DATA_SOURCE_ERROR` and the message of the library. It also checks that `ArrowUtils.rootAllocator` gets back all its memory after the failure, and after a scan with an `IN` filter whose domains the library takes. With the release of the ranges removed from `NativeArrow.exportDomains`, the check fails with 68202 bytes still allocated. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
