diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/filter/SpecializedFilterOpDesc.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/filter/SpecializedFilterOpDesc.scala index e5d47ee5232..424504bd3b9 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/filter/SpecializedFilterOpDesc.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/filter/SpecializedFilterOpDesc.scala @@ -21,13 +21,16 @@ package org.apache.texera.amber.operator.filter import com.fasterxml.jackson.annotation.{JsonProperty, JsonPropertyDescription} import org.apache.texera.amber.core.executor.OpExecWithClassName +import org.apache.texera.amber.core.tuple.{AttributeType, Schema} import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} -import org.apache.texera.amber.core.workflow.{InputPort, OutputPort, PhysicalOp} -import org.apache.texera.amber.operator.StandaloneCodeGenerator +import org.apache.texera.amber.core.workflow.{InputPort, OutputPort, PhysicalOp, PortIdentity} +import org.apache.texera.amber.operator.{StandaloneCodeGenerator, StandaloneHelpers} import org.apache.texera.amber.operator.metadata.{OperatorGroupConstants, OperatorInfo} import org.apache.texera.amber.pybuilder.PythonTemplateBuilder.pyStringLiteral import org.apache.texera.amber.util.JSONUtils.objectMapper +import scala.util.Try + class SpecializedFilterOpDesc extends FilterOpDesc with StandaloneCodeGenerator { @JsonProperty(value = "predicates", required = true) @@ -63,42 +66,159 @@ class SpecializedFilterOpDesc extends FilterOpDesc with StandaloneCodeGenerator ) } - override def generateStandaloneCode(): String = { + override def generateStandaloneCode(): String = generateStandaloneCode(Map.empty) + + override def generateStandaloneCode(inputSchemas: Map[PortIdentity, Schema]): String = { // No predicate keeps no row: the executor's filter is `predicates.exists`, // which answers false on an empty list. Passing the frame through would be // the opposite answer. if (predicates.isEmpty) return "out1df = in1df.iloc[0:0].copy()" + val schema = inputSchemas.values.headOption val conditions = predicates.map { p => val colLit = pyStringLiteral(p.attribute) p.condition match { case ComparisonType.IS_NULL => s"""(in1df[$colLit].isna())""" case ComparisonType.IS_NOT_NULL => s"""(in1df[$colLit].notna())""" - case other => - val op = other.getName // returns "=", ">=", "<", etc. (see ComparisonType.java) - val pyOp = if (op == "=") "==" else op - // notna mirrors FilterPredicate, which answers false for every condition - // but IS_NULL / IS_NOT_NULL once the field is null. Only `!=` needs it — - // pandas answers True there, where every other operator answers False — - // but guarding all of them keeps the one rule visible in one place. - s"""(in1df[$colLit].notna() & (in1df[$colLit] $pyOp ${coerceValue(p.value)}))""" + case other => + // FilterPredicate finds the field through Schema, which ignores case, + // and compares by the type declared there. Without a schema the + // column's dtype stands in for its type. + val keep = schema + .flatMap(s => Try(s.getAttribute(p.attribute)).toOption) + .map(a => SpecializedFilterOpDesc.keepFor(a.getType)) + .getOrElse("_texera_filter_keep") + val opLit = pyStringLiteral(other.getName) + val valueLit = pyStringLiteral(Option(p.value).getOrElse("")) + s"$keep(in1df[$colLit], $opLit, $valueLit)" } } s"out1df = in1df[${conditions.mkString(" | ")}].reset_index(drop=True)" } - // Try numeric coercion so generated code compares column values against the right type. - // Strings that don't parse fall through to a quoted string literal. - private def coerceValue(raw: String): String = { - try { - raw.toInt.toString - } catch { - case _: NumberFormatException => - try { - raw.toDouble.toString - } catch { - case _: NumberFormatException => - pyStringLiteral(raw) - } + override def standaloneHelpers(): Seq[String] = + Seq(StandaloneHelpers.AttributeCasts, SpecializedFilterOpDesc.FilterPredicates) +} + +object SpecializedFilterOpDesc { + + /** The function that compares a column of `attributeType`, so a script calls only the one + * its column needs. + */ + def keepFor(attributeType: AttributeType): String = + attributeType match { + case AttributeType.INTEGER | AttributeType.DOUBLE => "_texera_filter_keep_double" + case AttributeType.LONG => "_texera_filter_keep_long" + case AttributeType.BOOLEAN => "_texera_filter_keep_boolean" + case AttributeType.TIMESTAMP => "_texera_filter_keep_timestamp" + case _ => "_texera_filter_keep_string" } - } + + /** `FilterPredicate.evaluate`, one function per type the engine compares apart. A null passes + * nothing. + * + * - `_texera_filter_keep_double`: INTEGER and DOUBLE by `Double.compareTo`, NaN above all. + * - `_texera_filter_keep_long`: LONG, the value read by `Long.valueOf`. + * - `_texera_filter_keep_boolean`: BOOLEAN as lower-case text. + * - `_texera_filter_keep_timestamp`: TIMESTAMP by whole milliseconds in the machine's zone. + * - `_texera_filter_keep_string`: STRING as numbers when both parse, else by UTF-16 units. + * - `_texera_filter_keep`: the type read off the column's dtype, when no schema says it. + * + * The value is read with `_texera_java_trim`, `_texera_java_double` and `_texera_parse_int`, + * and a double is ordered by `_texera_double_key`, all from + * [[StandaloneHelpers.AttributeCasts]]. + */ + val FilterPredicates: String = + """def _texera_filter_text_key(s): + | return s.encode("utf-16-be", "surrogatepass") + | + |def _texera_filter_compared(left, right): + | return (left > right) - (left < right) + | + |def _texera_filter_rows(column, condition, order): + | passes = { + | "=": lambda c: c == 0, + | ">": lambda c: c > 0, + | ">=": lambda c: c >= 0, + | "<": lambda c: c < 0, + | "<=": lambda c: c <= 0, + | "!=": lambda c: c != 0, + | }[condition] + | return pd.Series( + | [ + | False if missing else passes(order(v)) + | for v, missing in zip(column, column.isna()) + | ], + | index=column.index, + | dtype=bool, + | ) + | + |def _texera_filter_keep_double(column, condition, value): + | target = _texera_double_key(_texera_java_double(value)) + | + | def order(v): + | return _texera_filter_compared(_texera_double_key(float(v)), target) + | + | return _texera_filter_rows(column, condition, order) + | + |def _texera_filter_keep_long(column, condition, value): + | target = _texera_parse_int(value, 64) + | + | def order(v): + | return _texera_filter_compared(int(v), target) + | + | return _texera_filter_rows(column, condition, order) + | + |def _texera_filter_keep_boolean(column, condition, value): + | target = _texera_filter_text_key(_texera_java_trim(value).lower()) + | + | def order(v): + | return _texera_filter_compared(_texera_filter_text_key("true" if v else "false"), target) + | + | return _texera_filter_rows(column, condition, order) + | + |def _texera_filter_keep_timestamp(column, condition, value): + | zone = _texera_local_zone() + | + | def wall_clock(ts): + | ts = pd.Timestamp(ts).floor("ms") + | if ts.tzinfo is None: + | return ts + | return pd.Timestamp(_texera_local_wall_clock(ts.to_pydatetime(), zone)) + | + | target = wall_clock(_texera_java_trim(value)) + | + | def order(v): + | return _texera_filter_compared(wall_clock(v), target) + | + | return _texera_filter_rows(column, condition, order) + | + |def _texera_filter_keep_string(column, condition, value): + | try: + | value_number = _texera_double_key(_texera_java_double(value)) + | except ValueError: + | value_number = None + | + | def order(v): + | text = str(v) + | if value_number is not None: + | try: + | number = _texera_double_key(_texera_java_double(text)) + | return _texera_filter_compared(number, value_number) + | except ValueError: + | pass + | return _texera_filter_compared( + | _texera_filter_text_key(text), _texera_filter_text_key(value) + | ) + | + | return _texera_filter_rows(column, condition, order) + | + |def _texera_filter_keep(column, condition, value): + | dtype = column.dtype + | if pd.api.types.is_bool_dtype(dtype): + | return _texera_filter_keep_boolean(column, condition, value) + | if pd.api.types.is_integer_dtype(dtype) or pd.api.types.is_float_dtype(dtype): + | return _texera_filter_keep_double(column, condition, value) + | if pd.api.types.is_datetime64_any_dtype(dtype): + | return _texera_filter_keep_timestamp(column, condition, value) + | return _texera_filter_keep_string(column, condition, value)""".stripMargin } diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/filter/SpecializedFilterOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/filter/SpecializedFilterOpDescSpec.scala index 39e1ec59988..fb0dfdcbe49 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/filter/SpecializedFilterOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/filter/SpecializedFilterOpDescSpec.scala @@ -19,16 +19,32 @@ package org.apache.texera.amber.operator.filter +import com.typesafe.config.ConfigFactory import org.apache.texera.amber.core.executor.OpExecWithClassName +import org.apache.texera.amber.core.tuple.{AttributeType, Schema, Tuple} import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} +import org.apache.texera.amber.core.workflow.PortIdentity import org.apache.texera.amber.operator.LogicalOp import org.apache.texera.amber.operator.metadata.OperatorGroupConstants +import org.apache.texera.amber.operator.tags.IntegrationTest import org.apache.texera.amber.util.JSONUtils.objectMapper +import org.scalatest.Tag import org.scalatest.flatspec.AnyFlatSpec import org.scalatest.matchers.should.Matchers +import java.nio.charset.StandardCharsets +import java.nio.file.Files +import java.sql.Timestamp +import java.time.{LocalDateTime, ZoneOffset} +import java.util.TimeZone +import java.util.concurrent.TimeUnit +import scala.io.Source +import scala.util.Try + class SpecializedFilterOpDescSpec extends AnyFlatSpec with Matchers { + private val NeedsPythonPackages = Tag(classOf[IntegrationTest].getName) + private val workflowId = WorkflowIdentity(1L) private val executionId = ExecutionIdentity(1L) @@ -71,8 +87,6 @@ class SpecializedFilterOpDescSpec extends AnyFlatSpec with Matchers { restored.asInstanceOf[SpecializedFilterOpDesc].predicates shouldBe empty } - // A null answers false for every condition but IS_NULL / IS_NOT_NULL, which pandas - // does not do on its own for `!=`, so the emitted condition carries the guard. "SpecializedFilterOpDesc.generateStandaloneCode" should "emit one condition per predicate" in { val d = new SpecializedFilterOpDesc d.predicates = List(new FilterPredicate("age", ComparisonType.GREATER_THAN, "18")) @@ -88,4 +102,306 @@ class SpecializedFilterOpDescSpec extends AnyFlatSpec with Matchers { "out1df = in1df.iloc[0:0].copy()" } + // FilterPredicate reads `KEY` from a column `Key`, since Schema ignores case, + // and compares it by that column's declared type. Read from the dtype, a LONG + // column would compare as a double and lose the digits past 2^53. + it should "take the declared type of a column named in a different case" in { + val d = new SpecializedFilterOpDesc + d.predicates = List(new FilterPredicate("KEY", ComparisonType.EQUAL_TO, "9007199254740993")) + val code = + d.generateStandaloneCode(Map(PortIdentity() -> Schema().add("Key", AttributeType.LONG))) + code should include("_texera_filter_keep_long(") + code should include(""""=", "9007199254740993")""") + } + + /** One column of a given type, the rows the engine reads and the frame the + * script reads, and the predicates run over it. The table spells the column + * `name`, and the predicates name it `predicateName`. + */ + private case class Column( + attrType: AttributeType, + values: Seq[Any], + pandas: String, + predicates: Seq[(ComparisonType, String)], + name: String = "v", + predicateName: String = "v" + ) + + // A column holds a wall clock, which the engine reads in the JVM's zone and + // pandas holds without one, so both sides are built from the same wall clock. + private def timestamp(micros: Long): Timestamp = + Timestamp.valueOf( + LocalDateTime.ofEpochSecond( + Math.floorDiv(micros, 1000000L), + (Math.floorMod(micros, 1000000L) * 1000).toInt, + ZoneOffset.UTC + ) + ) + + // 2024-01-01 00:00:00 plus a few microseconds, and a millisecond later. + private val midnight = 1704067200000000L + + private val columns = Seq( + Column( + AttributeType.DOUBLE, + Seq(Double.NaN, -0.0, 0.0, 1.0, null, Double.PositiveInfinity), + """pd.arrays.FloatingArray( + | np.array([np.nan, -0.0, 0.0, 1.0, 0.0, np.inf]), + | np.array([False, False, False, False, True, False]), + |)""".stripMargin, + Seq( + ComparisonType.EQUAL_TO -> "0", + ComparisonType.GREATER_THAN -> "0", + ComparisonType.GREATER_THAN_OR_EQUAL_TO -> "0", + ComparisonType.LESS_THAN -> "0", + ComparisonType.LESS_THAN_OR_EQUAL_TO -> "0", + ComparisonType.NOT_EQUAL_TO -> "0", + ComparisonType.GREATER_THAN -> "1", + ComparisonType.EQUAL_TO -> "NaN", + ComparisonType.LESS_THAN -> "Infinity", + ComparisonType.EQUAL_TO -> "-0.0", + // Hex past the largest double rounds to infinity, and below the + // smallest to a zero that keeps its sign. + ComparisonType.EQUAL_TO -> "0x1.fffffffffffff8p1023", + ComparisonType.EQUAL_TO -> "-0x1p-1080" + ) + ), + Column( + AttributeType.DOUBLE, + Seq(Double.PositiveInfinity, 1.0, null, Double.NegativeInfinity), + """pd.arrays.FloatingArray( + | np.array([np.inf, 1.0, 0.0, -np.inf]), + | np.array([False, False, True, False]), + |)""".stripMargin, + Seq( + ComparisonType.EQUAL_TO -> "0x1p1024", + ComparisonType.EQUAL_TO -> "-0x1p1024", + ComparisonType.LESS_THAN -> "0x1p1024d", + ComparisonType.EQUAL_TO -> "0x1p0" + ) + ), + Column( + AttributeType.STRING, + Seq("1", "2", "10", "abc", null, " 3 ", "nan", "2d", "0x1p2000"), + """pd.Series(["1", "2", "10", "abc", None, " 3 ", "nan", "2d", "0x1p2000"], dtype=object)""", + Seq( + ComparisonType.EQUAL_TO -> "Infinity", + ComparisonType.EQUAL_TO -> "1", + ComparisonType.GREATER_THAN -> "2", + ComparisonType.EQUAL_TO -> "1.0", + ComparisonType.LESS_THAN -> "b", + ComparisonType.EQUAL_TO -> "nan", + ComparisonType.EQUAL_TO -> "2", + ComparisonType.NOT_EQUAL_TO -> "abc" + ) + ), + Column( + AttributeType.BOOLEAN, + Seq(true, false, null), + """pd.array([True, False, None], dtype="boolean")""", + Seq( + ComparisonType.EQUAL_TO -> "True", + ComparisonType.EQUAL_TO -> "1", + ComparisonType.GREATER_THAN -> "false" + ) + ), + Column( + AttributeType.INTEGER, + Seq(1, 2, null, -3), + """pd.array([1, 2, None, -3], dtype="Int32")""", + Seq(ComparisonType.GREATER_THAN -> "1.5", ComparisonType.EQUAL_TO -> "2") + ), + Column( + AttributeType.LONG, + Seq(1L, 9007199254740993L, null), + """pd.array([1, 9007199254740993, None], dtype="Int64")""", + Seq( + ComparisonType.EQUAL_TO -> "9007199254740993", + ComparisonType.GREATER_THAN -> " 1 " + ) + ), + Column( + AttributeType.TIMESTAMP, + Seq(midnight, midnight + 900, midnight + 1000).map(timestamp), + s"""pd.Series(np.array([$midnight, ${midnight + 900}, ${midnight + 1000}], dtype="datetime64[us]"))""", + Seq( + ComparisonType.EQUAL_TO -> "2024-01-01 00:00:00", + ComparisonType.GREATER_THAN -> "2024-01-01 00:00:00" + ) + ) + ) + + // FilterPredicate compares by the column's type: a STRING field and the value + // as numbers when both parse as one, a BOOLEAN as its lower-case text, an + // INTEGER as a double, a DOUBLE by Double.compareTo, and a TIMESTAMP by + // getTime. Each predicate runs through SpecializedFilterOpExec and through + // the export, and the rows each keeps are named by their input position. + it should "keep the rows the engine keeps for each column type" taggedAs NeedsPythonPackages in { + expectTheRowsTheEngineKeeps(columns) + } + + // FilterPredicate reads `KEY` from a column `Key`, since Schema ignores case, + // and compares it as the LONG declared there. 2^53 + 1 then keeps only its own + // row, where a double holds both values as 2^53 and keeps both. + it should "compare a column named in a different case by its declared type" taggedAs + NeedsPythonPackages in { + expectTheRowsTheEngineKeeps( + Seq( + Column( + AttributeType.LONG, + Seq(9007199254740992L, 9007199254740993L), + """pd.array([9007199254740992, 9007199254740993], dtype="Int64")""", + Seq(ComparisonType.EQUAL_TO -> "9007199254740993"), + name = "Key", + predicateName = "KEY" + ) + ) + ) + } + + // The engine holds a value with an offset as the wall clock it shows in the + // JVM's zone. A zone away from UTC, and a half hour off any whole hour, so a + // reading in UTC or a dropped offset keeps a different row. + it should "read a timestamp value with an offset in the machine's zone" taggedAs + NeedsPythonPackages in { + val zone = "Asia/Kolkata" + val previous = TimeZone.getDefault + TimeZone.setDefault(TimeZone.getTimeZone(zone)) + try { + val utcMidnight = midnight + 19800000000L + expectTheRowsTheEngineKeeps( + Seq( + Column( + AttributeType.TIMESTAMP, + Seq(timestamp(utcMidnight), timestamp(midnight), null), + s"""pd.Series(np.array([$utcMidnight, $midnight, "NaT"], dtype="datetime64[us]"))""", + Seq( + ComparisonType.EQUAL_TO -> "2024-01-01T00:00:00Z", + ComparisonType.EQUAL_TO -> "2024-01-01 00:00:00", + ComparisonType.EQUAL_TO -> "2024-01-01T00:00:00+05:30", + ComparisonType.EQUAL_TO -> "2023-12-31T16:00:00-08:00", + ComparisonType.GREATER_THAN -> "2024-01-01T00:00:00Z", + ComparisonType.LESS_THAN -> "2024-01-01T00:00:00Z" + ) + ) + ), + pythonZone = Some(zone) + ) + } finally TimeZone.setDefault(previous) + } + + // The offset is the one the zone had at that instant, not today's. Mexico + // City kept summer time in 2021 and stood an hour further west in 1931, so + // today's standard offset keeps the other row of each pair. + it should "read a timestamp value with the offset its zone had at that instant" taggedAs + NeedsPythonPackages in { + val zone = "America/Mexico_City" + val previous = TimeZone.getDefault + TimeZone.setDefault(TimeZone.getTimeZone(zone)) + try { + val walls = Seq( + LocalDateTime.of(2021, 6, 30, 19, 0), + LocalDateTime.of(2021, 6, 30, 18, 0), + LocalDateTime.of(1930, 12, 31, 17, 0), + LocalDateTime.of(1930, 12, 31, 18, 0) + ).map(_.toEpochSecond(ZoneOffset.UTC) * 1000000L) + expectTheRowsTheEngineKeeps( + Seq( + Column( + AttributeType.TIMESTAMP, + walls.map(timestamp), + s"""pd.Series(np.array([${walls.mkString(", ")}], dtype="datetime64[us]"))""", + Seq( + ComparisonType.EQUAL_TO -> "2021-07-01T00:00:00Z", + ComparisonType.EQUAL_TO -> "1931-01-01T00:00:00Z" + ) + ) + ), + pythonZone = Some(zone) + ) + } finally TimeZone.setDefault(previous) + } + + /** Runs each predicate through SpecializedFilterOpExec and through the export, + * and expects the export to keep the same rows, named by input position. + */ + private def expectTheRowsTheEngineKeeps( + columns: Seq[Column], + pythonZone: Option[String] = None + ): Unit = { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val runs = for { + column <- columns + (condition, value) <- column.predicates + } yield { + val schema = Schema().add(column.name, column.attrType).add("i", AttributeType.INTEGER) + val desc = new SpecializedFilterOpDesc + desc.predicates = List(new FilterPredicate(column.predicateName, condition, value)) + val exec = new SpecializedFilterOpExec(objectMapper.writeValueAsString(desc)) + exec.open() + val kept = column.values.zipWithIndex.flatMap { + case (v, i) => exec.processTuple(Tuple(schema, Array[Any](v, i)), 0) + } + exec.close() + val label = s"${column.attrType} ${condition.getName} '$value'" + val engine = s"$label: ${kept.map(_.getFields(1)).mkString(" ")}" + val code = desc.generateStandaloneCode(Map(PortIdentity() -> schema)) + val script = + s"""in1df = pd.DataFrame({${objectMapper.writeValueAsString( + column.name + )}: ${column.pandas.linesIterator.mkString("\n ")}}) + |in1df["i"] = range(len(in1df)) + |ns = {"pd": pd, "in1df": in1df} + |exec(${objectMapper.writeValueAsString( + (desc.standaloneHelpers() :+ code).mkString("\n") + )}, ns) + |print(${objectMapper.writeValueAsString( + label + ":" + )}, " ".join(str(i) for i in ns["out1df"]["i"]))""".stripMargin + (engine, script) + } + + val driver = ("import pandas as pd, numpy as np" +: runs.map(_._2)).mkString("\n") + val file = Files.createTempFile("filter-engine-parity-", ".py") + file.toFile.deleteOnExit() + Files.write(file, driver.getBytes(StandardCharsets.UTF_8)) + val builder = new ProcessBuilder(python, file.toString).redirectErrorStream(true) + pythonZone.foreach(builder.environment().put("TZ", _)) + val process = builder.start() + val out = Source.fromInputStream(process.getInputStream).mkString + process.waitFor(120, TimeUnit.SECONDS) + + withClue(s"python said:\n$out\nscript:\n$driver") { + process.exitValue() shouldBe 0 + out.linesIterator.map(_.trim).toList shouldBe runs.map(_._1.trim).toList + } + } + + private def resolvePython(): Option[String] = { + def fromConfig: Option[String] = + Try(ConfigFactory.parseResources("udf.conf").resolve()).toOption + .orElse(Try(ConfigFactory.load()).toOption) + .flatMap(c => Try(c.getConfig("python").getString("path")).toOption) + .map(_.trim) + .filter(_.nonEmpty) + + def runnable(exe: String): Boolean = + Try(new ProcessBuilder(exe, "--version").redirectErrorStream(true).start()).toOption + .exists { p => + if (!p.waitFor(5, TimeUnit.SECONDS)) { p.destroyForcibly(); false } + else p.exitValue() == 0 + } + + (fromConfig.toList ++ List("python3", "python", "py")).distinct.find(runnable) + } + + private def canImportPandas(python: String): Boolean = + Try( + new ProcessBuilder(python, "-c", "import pandas, numpy").redirectErrorStream(true).start() + ).toOption.exists { p => + if (!p.waitFor(60, TimeUnit.SECONDS)) { p.destroyForcibly(); false } + else p.exitValue() == 0 + } }