diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/LogicalOp.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/LogicalOp.scala index efa46144180..5e7be51d003 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/LogicalOp.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/LogicalOp.scala @@ -459,6 +459,18 @@ abstract class LogicalOp extends PortDescriptor with Serializable { def operatorInfo: OperatorInfo + /** + * Whether the row ORDER of this operator's output is part of its contract. + * Defaults to false: the engine runs operators across parallel workers, so + * for almost every operator the output row order is an implementation-defined + * interleaving. Only operators whose very purpose is to establish an order + * override this, which here is the sort family: Sort, Stable Merge Sort and + * Sort Partitions. Anything comparing two runs of an operator reads it to + * decide whether the rows have to arrive in the same order or only be the + * same rows. + */ + def orderSensitive: Boolean = false + private def getOperatorVersion: String = { val path = "amber/src/main/scala/" val operatorPath = path + this.getClass.getPackage.getName.replace(".", "/") diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/aggregate/AggregateOpDesc.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/aggregate/AggregateOpDesc.scala index 1bace652c8b..53d28ebc3ea 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/aggregate/AggregateOpDesc.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/aggregate/AggregateOpDesc.scala @@ -22,21 +22,24 @@ package org.apache.texera.amber.operator.aggregate import com.fasterxml.jackson.annotation.{JsonProperty, JsonPropertyDescription} import com.kjetland.jackson.jsonSchema.annotations.JsonSchemaTitle import org.apache.texera.amber.core.executor.OpExecWithClassName -import org.apache.texera.amber.core.tuple.Schema +import org.apache.texera.amber.core.tuple.{AttributeType, Schema} import org.apache.texera.amber.core.virtualidentity.{ ExecutionIdentity, PhysicalOpIdentity, WorkflowIdentity } import org.apache.texera.amber.core.workflow._ -import org.apache.texera.amber.operator.LogicalOp +import org.apache.texera.amber.operator.{LogicalOp, StandaloneCodeGenerator, StandaloneHelpers} import org.apache.texera.amber.operator.metadata.annotations.AutofillAttributeNameList 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 javax.validation.constraints.{NotNull, Size} +import scala.util.Try + +class AggregateOpDesc extends LogicalOp with StandaloneCodeGenerator { -class AggregateOpDesc extends LogicalOp { @JsonProperty(value = "aggregations", required = true) @JsonPropertyDescription("multiple aggregation functions") @NotNull(message = "aggregation cannot be null") @@ -138,4 +141,354 @@ class AggregateOpDesc extends LogicalOp { inputPorts = List(InputPort()), outputPorts = List(OutputPort()) ) + + /** The engine aggregates in two phases across partitions; one process needs + * only the one groupby, or a single-row reduction when no key is grouped on. + * + * Must run before `getPhysicalPlan`, which rewrites `aggregations` in place: + * it turns COUNT into SUM for the final phase, and this reads them as + * written. + * + * SUM and AVERAGE follow the column's DECLARED type: a holed INTEGER column + * arrives as a float, and an INTEGER and a LONG arrive alike. + */ + override def generateStandaloneCode(inputSchemas: Map[PortIdentity, Schema]): String = { + val schema = inputSchemas.get(operatorInfo.inputPorts.head.id) + build(name => schema.flatMap(s => Try(s.getAttribute(name).getType).toOption)) + } + + override def generateStandaloneCode(): String = build(_ => None) + + override def standaloneHelpers(): Seq[String] = Seq(StandaloneHelpers.AttributeCasts) + + private def build(declaredType: String => Option[AttributeType]): String = { + val keys = Option(groupByKeys).getOrElse(List()) + val aggs = Option(aggregations).getOrElse(List()) + + // Identical helper definition each call — keeps the standalone module + // self-contained without relying on a shared prelude. + val concatHelper = + """def _texera_agg_concat(series): + | # The accumulator starts empty and only earns a separator once it + | # holds something, so a leading empty value adds neither text nor + | # comma: "", "a", "" concatenates to "a," and not ",a,". A null is + | # read as the empty string, which is what makes the two the same + | # here. This is concatAgg's fold, written out. + | # Each value is written as Java's toString writes it. The column's + | # own mask says which are null: pd.isna on the value would take a + | # stored NaN for one too. + | partial = "" + | for v, missing in zip(series, series.isna()): + | text = "" if missing else _texera_java_text(v) + | partial = text if partial == "" else partial + "," + text + | return partial + | + |def _texera_agg_int_sum(series): + | # The engine adds an INTEGER column as Java ints, which wrap. + | total = int(series.sum()) + | return ((total + (1 << 31)) % (1 << 32)) - (1 << 31) + | + |def _texera_agg_ts_epoch_ms(series): + | # A timestamp reaches the engine as its epoch milliseconds whatever + | # resolution the column carries, and reading the integers out of a + | # microsecond column asks for a different number than a nanosecond + | # one, so cast to milliseconds before reading them. + | # + | # A column holds a wall clock, and `Timestamp.getTime` answers for + | # the instant that wall clock names locally, so localize before + | # reading the integers out: left as UTC every one of them is a whole + | # offset away. The two flags are java.time's own reading of an hour + | # daylight saving repeats (the later one) or skips (shifted past the + | # gap). + | # + | # The engine's zone is the machine's, with its whole history; before + | # 1900 java.util.Date reads every wall clock at the standard offset. + | from dateutil.tz import tzlocal + | + | kept = series.dropna().astype("datetime64[ms]") + | zone = _texera_local_zone() + | offset_ms = int(_texera_standard_offset(zone).total_seconds()) * 1000 + | standard = kept.astype("int64") - offset_ms + | # 1900-01-01T00:00Z in epoch milliseconds. + | before = standard < -2208988800000 + | + | def localized(s): + | return s.dt.tz_localize( + | zone or tzlocal(), ambiguous=False, nonexistent=pd.Timedelta("1h") + | ).astype("int64") + | + | try: + | return localized(kept).where(~before, standard) + | except (OSError, OverflowError): + | pass + | + | # With no zone named, tzlocal() asks localtime, and Windows refuses an + | # instant before 1970. Read those at the standard offset too. + | def one(i): + | if before.iloc[i]: + | return int(standard.iloc[i]) + | try: + | return int(localized(kept.iloc[i : i + 1]).iloc[0]) + | except (OSError, OverflowError): + | return int(standard.iloc[i]) + | + | return pd.Series([one(i) for i in range(len(kept))], index=kept.index, dtype="int64") + | + |def _texera_agg_ts_sum(series): + | # SUM keeps the column's own type, so the engine adds the epoch + | # milliseconds as Java longs, which wrap, and builds a timestamp + | # from the total. Python integers carry the sum exactly, so the + | # wrap is the only place a total loses anything. + | total = int(_texera_agg_ts_epoch_ms(series).astype(object).sum()) + | total = ((total + (1 << 63)) % (1 << 64)) - (1 << 63) + | # Through datetime and not Timestamp.tz_convert, which swallows a + | # refusal from localtime and answers in UTC. + | import datetime + | + | utc = pd.Timestamp(total, unit="ms") + | zone = _texera_local_zone() + | try: + | moment = utc.to_pydatetime().replace(tzinfo=datetime.timezone.utc) + | shift = _texera_local_wall_clock(moment, zone) - moment.replace(tzinfo=None) + | except (OSError, OverflowError): + | shift = _texera_standard_offset(zone) + | return utc + shift + | + |def _texera_agg_ts_mean(series): + | # AVERAGE is declared DOUBLE whatever column it reads, so this is + | # the mean of the epoch milliseconds and not a timestamp. + | kept = _texera_agg_ts_epoch_ms(series) + | if len(kept) == 0: + | return None + | return float(kept.astype(object).sum()) / len(kept) + | + |def _texera_agg_ts_pick(series, better): + | # AttributeTypeUtils.compare reads a timestamp by getTime, so two in + | # the same millisecond tie, and the fold keeps the one it holds, + | # digits past the millisecond and all. floor and not round, as + | # getTime does. + | kept = series.dropna() + | if len(kept) == 0: + | return pd.NaT + | millis = kept.dt.floor("ms") + | pick = 0 + | for i in range(1, len(kept)): + | if better(millis.iloc[i], millis.iloc[pick]): + | pick = i + | return kept.iloc[pick] + | + |def _texera_agg_ts_min(series): + | return _texera_agg_ts_pick(series, lambda a, b: a < b) + | + |def _texera_agg_ts_max(series): + | return _texera_agg_ts_pick(series, lambda a, b: a > b) + | + |# A DOUBLE column can hold a stored NaN apart from null, and pandas' + |# reductions answer null for it. The engine skips only the nulls and folds + |# the rest as Java doubles, so these folds are written out. Each answers + |# pd.NA for a null result, which pandas keeps apart from NaN. + |def _texera_agg_double_values(series): + | return [float(v) for v in series.dropna()] + | + |def _texera_agg_double_key(v): + | # Java's Double.compare: NaN sorts above everything, -0.0 below 0.0. + | import math + | + | return (math.isnan(v), v, math.copysign(1.0, v)) + | + |def _texera_agg_double_sum(series): + | partial = 0.0 + | for v in _texera_agg_double_values(series): + | partial += v + | return partial + | + |def _texera_agg_double_mean(series): + | kept = _texera_agg_double_values(series) + | if not kept: + | return pd.NA + | partial = 0.0 + | for v in kept: + | partial += v + | return partial / len(kept) + | + |def _texera_agg_double_min(series): + | # The engine starts from Double.MAX_VALUE and reads it back as null. + | start = 1.7976931348623157e308 + | partial = start + | for v in _texera_agg_double_values(series): + | if _texera_agg_double_key(v) < _texera_agg_double_key(partial): + | partial = v + | return pd.NA if partial == start else partial + | + |def _texera_agg_double_max(series): + | # The engine starts from negative infinity and reads it back as null. + | start = float("-inf") + | partial = start + | for v in _texera_agg_double_values(series): + | if _texera_agg_double_key(v) > _texera_agg_double_key(partial): + | partial = v + | return pd.NA if partial == start else partial + | + |def _texera_agg_double_by_group(column, fold): + | # A plain float column cannot hold NaN and null apart, so the folded + | # groups go into a nullable one. + | results = column.apply(fold) + | mask = results.map(lambda v: v is pd.NA).to_numpy(bool) + | values = results.where(~mask, 0.0).to_numpy(float) + | return pd.Series(pd.arrays.FloatingArray(values, mask), index=results.index) + | + |def _texera_agg_group_ids(keys): + | # The engine keys each group on the List of the row's key fields in a + | # HashMap, under Scala's ==: NaN equals nothing, so every NaN row opens + | # a group of its own, while 0.0 meets -0.0 and a null meets a null. + | # groupby puts every NaN in one group, so each row gets its group's + | # id here. A null is read off the column's mask, since a stored NaN is + | # not one. + | import math + | + | ids = {} + | out = [] + | for values, missing in zip( + | keys.itertuples(index=False, name=None), keys.isna().to_numpy() + | ): + | key = tuple( + | None if m else object() if isinstance(v, float) and math.isnan(v) else v + | for v, m in zip(values, missing) + | ) + | out.append(ids.setdefault(key, len(ids))) + | return pd.Series(out, index=keys.index, dtype="int64")""".stripMargin + + if (keys.isEmpty) { + val rowEntries = aggs + .map(agg => + s" ${pyStringLiteral(agg.resultAttribute)}: ${aggExprScalar(agg, declaredType)}," + ) + .mkString("\n") + val resultsLit = + aggs.map(agg => pyStringLiteral(agg.resultAttribute)).mkString("[", ", ", "]") + // The engine opens a group when a row arrives, so with no rows there is no + // group to finish and it emits nothing, not one row of COUNT 0. + s"""$concatHelper + |if len(in1df): + | out1df = pd.DataFrame([{ + |${rowEntries.linesIterator.map(" " + _).mkString("\n")} + | }]) + |else: + | out1df = pd.DataFrame(columns=$resultsLit)""".stripMargin + } else { + val keysLit = keys.map(pyStringLiteral).mkString("[", ", ", "]") + val aggLines = aggs.zipWithIndex + .map { + case (agg, i) => + s"_texera_agg_s$i = ${aggExprGroupby(agg, "_texera_agg_groups", declaredType)}" + } + .mkString("\n") + // Grouped and joined back by the id of each row's group, never by the + // key values, which would pair NaN keys again. + val joinLines = aggs.indices + .map(i => s"out1df = out1df.join(_texera_agg_s$i)") + .mkString("\n") + s"""$concatHelper + |_texera_agg_ids = _texera_agg_group_ids(in1df[$keysLit]) + |_texera_agg_groups = in1df.groupby(_texera_agg_ids, sort=False) + |_texera_agg_first = ~_texera_agg_ids.duplicated() + |out1df = in1df.loc[_texera_agg_first, $keysLit].set_axis(_texera_agg_ids[_texera_agg_first]) + |$aggLines + |$joinLines + |out1df = out1df.reset_index(drop=True)""".stripMargin + } + } + + private def aggExprScalar( + agg: AggregationOperation, + declaredType: String => Option[AttributeType] + ): String = { + val attrLit = + if (agg.attribute == null || agg.attribute.isEmpty) "None" + else pyStringLiteral(agg.attribute) + val declared = Option(agg.attribute).filter(_.nonEmpty).flatMap(declaredType) + agg.aggFunction match { + case AggregationFunction.SUM => + declared match { + case Some(AttributeType.INTEGER) => s"_texera_agg_int_sum(in1df[$attrLit])" + case Some(AttributeType.TIMESTAMP) => s"_texera_agg_ts_sum(in1df[$attrLit])" + case Some(AttributeType.DOUBLE) => s"_texera_agg_double_sum(in1df[$attrLit])" + case _ => s"in1df[$attrLit].sum()" + } + case AggregationFunction.AVERAGE => + declared match { + case Some(AttributeType.TIMESTAMP) => s"_texera_agg_ts_mean(in1df[$attrLit])" + case Some(AttributeType.DOUBLE) => s"_texera_agg_double_mean(in1df[$attrLit])" + case _ => s"in1df[$attrLit].mean()" + } + case AggregationFunction.MIN => + declared match { + case Some(AttributeType.DOUBLE) => s"_texera_agg_double_min(in1df[$attrLit])" + case Some(AttributeType.TIMESTAMP) => s"_texera_agg_ts_min(in1df[$attrLit])" + case _ => s"in1df[$attrLit].min()" + } + case AggregationFunction.MAX => + declared match { + case Some(AttributeType.DOUBLE) => s"_texera_agg_double_max(in1df[$attrLit])" + case Some(AttributeType.TIMESTAMP) => s"_texera_agg_ts_max(in1df[$attrLit])" + case _ => s"in1df[$attrLit].max()" + } + case AggregationFunction.COUNT => + if (agg.attribute == null || agg.attribute.isEmpty) "int(len(in1df))" + else s"int(in1df[$attrLit].count())" + case AggregationFunction.CONCAT => s"_texera_agg_concat(in1df[$attrLit])" + } + } + + private def aggExprGroupby( + agg: AggregationOperation, + groups: String, + declaredType: String => Option[AttributeType] + ): String = { + val attrLit = + if (agg.attribute == null || agg.attribute.isEmpty) "None" + else pyStringLiteral(agg.attribute) + val resultLit = pyStringLiteral(agg.resultAttribute) + val declared = Option(agg.attribute).filter(_.nonEmpty).flatMap(declaredType) + def byDoubleFold(fold: String): String = + s"_texera_agg_double_by_group($groups[$attrLit], $fold).rename($resultLit)" + agg.aggFunction match { + case AggregationFunction.SUM => + declared match { + case Some(AttributeType.INTEGER) => + s"$groups[$attrLit].apply(_texera_agg_int_sum).rename($resultLit)" + case Some(AttributeType.TIMESTAMP) => + s"$groups[$attrLit].apply(_texera_agg_ts_sum).rename($resultLit)" + case Some(AttributeType.DOUBLE) => byDoubleFold("_texera_agg_double_sum") + case _ => s"$groups[$attrLit].sum().rename($resultLit)" + } + case AggregationFunction.AVERAGE => + declared match { + case Some(AttributeType.TIMESTAMP) => + s"$groups[$attrLit].apply(_texera_agg_ts_mean).rename($resultLit)" + case Some(AttributeType.DOUBLE) => byDoubleFold("_texera_agg_double_mean") + case _ => s"$groups[$attrLit].mean().rename($resultLit)" + } + case AggregationFunction.MIN => + declared match { + case Some(AttributeType.DOUBLE) => byDoubleFold("_texera_agg_double_min") + case Some(AttributeType.TIMESTAMP) => + s"$groups[$attrLit].apply(_texera_agg_ts_min).rename($resultLit)" + case _ => s"$groups[$attrLit].min().rename($resultLit)" + } + case AggregationFunction.MAX => + declared match { + case Some(AttributeType.DOUBLE) => byDoubleFold("_texera_agg_double_max") + case Some(AttributeType.TIMESTAMP) => + s"$groups[$attrLit].apply(_texera_agg_ts_max).rename($resultLit)" + case _ => s"$groups[$attrLit].max().rename($resultLit)" + } + case AggregationFunction.COUNT => + if (agg.attribute == null || agg.attribute.isEmpty) + s"$groups.size().rename($resultLit)" + else s"$groups[$attrLit].count().rename($resultLit)" + case AggregationFunction.CONCAT => + s"$groups[$attrLit].apply(_texera_agg_concat).rename($resultLit)" + } + } } diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/SortCriteriaUnit.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/SortCriteriaUnit.scala index f0e851110b8..9b34fe44b1a 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/SortCriteriaUnit.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/SortCriteriaUnit.scala @@ -21,13 +21,17 @@ package org.apache.texera.amber.operator.sort import com.fasterxml.jackson.annotation.{JsonProperty, JsonPropertyDescription} import org.apache.texera.amber.pybuilder.PyStringTypes.EncodableString -import org.apache.texera.amber.operator.metadata.annotations.AutofillAttributeName +import org.apache.texera.amber.operator.metadata.annotations.{AutofillAttributeName, SampleColumn} class SortCriteriaUnit { + // Verification picks a column with TIE GROUPS. Sorting on a unique key never + // puts two equal keys next to each other, so the stability this operator is + // named for goes unchecked. @JsonProperty(value = "attribute", required = true) @JsonPropertyDescription("Attribute name to sort by") @AutofillAttributeName + @SampleColumn("name") var attributeName: EncodableString = _ @JsonProperty(value = "sortPreference", required = true) diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/SortOpDesc.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/SortOpDesc.scala index db2ae21e7f6..662fd37377d 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/SortOpDesc.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/SortOpDesc.scala @@ -21,11 +21,18 @@ package org.apache.texera.amber.operator.sort import com.fasterxml.jackson.annotation.{JsonProperty, JsonPropertyDescription} import org.apache.texera.amber.core.tuple.Schema -import org.apache.texera.amber.pybuilder.PythonTemplateBuilder.PythonTemplateBuilderStringContext +import org.apache.texera.amber.pybuilder.PythonTemplateBuilder.{ + PythonTemplateBuilderStringContext, + pyStringLiteral +} import org.apache.texera.amber.core.workflow.{InputPort, OutputPort, PortIdentity} -import org.apache.texera.amber.operator.PythonOperatorDescriptor +import org.apache.texera.amber.operator.{PythonOperatorDescriptor, StandaloneCodeGenerator} import org.apache.texera.amber.operator.metadata.{OperatorGroupConstants, OperatorInfo} -class SortOpDesc extends PythonOperatorDescriptor { +class SortOpDesc extends PythonOperatorDescriptor with StandaloneCodeGenerator { + + // Sorting is this operator's contract: output row order is meaningful. + override def orderSensitive: Boolean = true + @JsonProperty(required = true) @JsonPropertyDescription("column to perform sorting on") var attributes: List[SortCriteriaUnit] = List.empty @@ -78,4 +85,11 @@ class SortOpDesc extends PythonOperatorDescriptor { outputPorts = List(OutputPort(blocking = true)) ) + override def generateStandaloneCode(): String = { + val cols = attributes.map(c => pyStringLiteral(c.attributeName)).mkString("[", ", ", "]") + val ascending = attributes + .map(c => if (c.sortPreference == SortPreference.ASC) "True" else "False") + .mkString("[", ", ", "]") + s"out1df = in1df.sort_values(by=$cols, ascending=$ascending)" + } } diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/StableMergeSortOpDesc.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/StableMergeSortOpDesc.scala index 318089da5e3..53581633154 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/StableMergeSortOpDesc.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sort/StableMergeSortOpDesc.scala @@ -24,8 +24,9 @@ import com.kjetland.jackson.jsonSchema.annotations.JsonSchemaTitle import org.apache.texera.amber.core.executor.OpExecWithClassName 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.LogicalOp +import org.apache.texera.amber.operator.{LogicalOp, StandaloneCodeGenerator} 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.collection.mutable.ListBuffer @@ -36,7 +37,10 @@ import scala.collection.mutable.ListBuffer * the lexicographic order and per-key direction (ASC/DESC). */ //TODO(#3922): disallowing sorting on binary type -class StableMergeSortOpDesc extends LogicalOp { +class StableMergeSortOpDesc extends LogicalOp with StandaloneCodeGenerator { + + // Sorting is this operator's contract: output row order is meaningful. + override def orderSensitive: Boolean = true @JsonProperty(value = "keys", required = true) @JsonSchemaTitle("Sort Keys") @@ -69,4 +73,69 @@ class StableMergeSortOpDesc extends LogicalOp { List(InputPort()), List(OutputPort(blocking = true)) ) + + // The engine runs an incremental stable merge sort with nulls last whichever + // way a key points. pandas' mergesort is stable too, so the ordering below is + // the same one. + override def generateStandaloneCode(): String = { + val criteria = Option(keys).getOrElse(ListBuffer.empty) + if (criteria.isEmpty) return "out1df = in1df.copy()" + val cols = criteria + .map(c => pyStringLiteral(c.attributeName)) + .mkString("[", ", ", "]") + val ascending = criteria + .map(c => if (c.sortPreference == SortPreference.ASC) "True" else "False") + .mkString("[", ", ", "]") + // Sort each key in three tiers, because the engine treats a null and a NaN + // differently: a null goes last whichever way the key points, while a NaN + // compares above every number, so it goes last ascending and first + // descending. A column read into a numpy dtype has one slot for both, and + // there both land in the null tier, which is where they were before. + // + // The tiers are built in a frame of their own and named by position, so + // the input keeps every column it arrived with: a helper named after the + // key would overwrite an input column that already answers to that name, + // and dropping the helper afterwards would take the payload with it. + // + // A string sorts by its UTF-16 big-endian bytes. The engine compares UTF-16 + // code units and Python compares code points, and the two disagree once a + // character above U+FFFF meets one between U+E000 and U+FFFF. + // + // A double's sign is a fourth tier. Double.compare puts -0.0 below 0.0, + // where pandas reads the two as equal and keeps them in input order. + s"""import math + |_texera_sorted = in1df.reset_index(drop=True) + |_texera_keys = pd.DataFrame(index=_texera_sorted.index) + |_texera_by = [] + |_texera_asc = [] + |for _texera_i, (_texera_col, _texera_a) in enumerate(zip($cols, $ascending)): + | _texera_null = "null_" + str(_texera_i) + | _texera_nan = "nan_" + str(_texera_i) + | _texera_val = "val_" + str(_texera_i) + | _texera_keys[_texera_null] = _texera_sorted[_texera_col].isna() + | _texera_keys[_texera_nan] = ( + | _texera_sorted[_texera_col] != _texera_sorted[_texera_col] + | ).fillna(False) + | _texera_v = _texera_sorted[_texera_col] + | if _texera_v.dtype == object or isinstance(_texera_v.dtype, pd.StringDtype): + | _texera_v = _texera_v.map( + | lambda v: v.encode("utf-16-be") if isinstance(v, str) else v + | ) + | _texera_keys[_texera_val] = _texera_v + | _texera_by += [_texera_null, _texera_nan, _texera_val] + | _texera_asc += [True, _texera_a, _texera_a] + | if pd.api.types.is_float_dtype(_texera_v.dtype): + | _texera_sign = "sign_" + str(_texera_i) + | _texera_keys[_texera_sign] = [ + | 1 if missing or v != v or math.copysign(1.0, v) > 0 else 0 + | for v, missing in zip(_texera_v, _texera_v.isna()) + | ] + | _texera_by.append(_texera_sign) + | _texera_asc.append(_texera_a) + |out1df = _texera_sorted.loc[ + | _texera_keys.sort_values( + | by=_texera_by, ascending=_texera_asc, kind="mergesort" + | ).index + |].reset_index(drop=True)""".stripMargin + } } diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sortPartitions/SortPartitionsOpDesc.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sortPartitions/SortPartitionsOpDesc.scala index 9290d549f2e..e7b962e81ec 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sortPartitions/SortPartitionsOpDesc.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/sortPartitions/SortPartitionsOpDesc.scala @@ -24,9 +24,10 @@ import com.kjetland.jackson.jsonSchema.annotations.{JsonSchemaInject, JsonSchema import org.apache.texera.amber.core.executor.OpExecWithClassName import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} import org.apache.texera.amber.core.workflow.{InputPort, OutputPort, PhysicalOp, RangePartition} -import org.apache.texera.amber.operator.LogicalOp +import org.apache.texera.amber.operator.{LogicalOp, StandaloneCodeGenerator} import org.apache.texera.amber.operator.metadata.annotations.AutofillAttributeName 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 @JsonSchemaInject(json = """ @@ -38,7 +39,10 @@ import org.apache.texera.amber.util.JSONUtils.objectMapper } } """) -class SortPartitionsOpDesc extends LogicalOp { +class SortPartitionsOpDesc extends LogicalOp with StandaloneCodeGenerator { + + // Sorting is this operator's contract: output row order is meaningful. + override def orderSensitive: Boolean = true @JsonProperty(required = true) @JsonSchemaTitle("Attribute") @@ -84,4 +88,48 @@ class SortPartitionsOpDesc extends LogicalOp { inputPorts = List(InputPort()), outputPorts = List(OutputPort(blocking = true)) ) + + // Row order is this operator's answer, so the sort matches the engine's: + // stable, ascending, nulls first. The domain bounds are partitioning hints + // that the sort itself never reads. + // + // The engine orders a NaN after positive infinity. In a nullable column, + // which is what an Arrow file is read into, pandas keeps a NaN apart from a + // null and sorts it last too. A numpy float column has one slot for both, so + // there a NaN goes first with the nulls. + // + // The rest of AttributeTypeUtils.compare is written out as tiers of a key + // frame, as in StableMergeSortOpDesc. Double.compare puts -0.0 below 0.0, + // where pandas reads the two as equal. A timestamp is compared by getTime, + // so two inside one millisecond keep their input order. A string compares + // UTF-16 code units, where pandas compares code points, and the two disagree + // once a character above U+FFFF meets one between U+E000 and U+FFFF. + override def generateStandaloneCode(): String = { + val col = pyStringLiteral(Option(sortAttributeName).getOrElse("")) + s"""import math + |out1df = in1df.reset_index(drop=True) + |_texera_v = out1df[$col] + |_texera_keys = pd.DataFrame(index=out1df.index) + |_texera_keys["held"] = _texera_v.notna() + |_texera_keys["nan"] = (_texera_v != _texera_v).fillna(False) + |if pd.api.types.is_datetime64_any_dtype(_texera_v.dtype): + | _texera_keys["val"] = _texera_v.dt.floor("ms") + |elif _texera_v.dtype == object or isinstance(_texera_v.dtype, pd.StringDtype): + | _texera_keys["val"] = _texera_v.map( + | lambda v: v.encode("utf-16-be") if isinstance(v, str) else v + | ) + |else: + | _texera_keys["val"] = _texera_v + |_texera_by = ["held", "nan", "val"] + |if pd.api.types.is_float_dtype(_texera_v.dtype): + | _texera_keys["sign"] = [ + | 1 if missing or v != v or math.copysign(1.0, v) > 0 else 0 + | for v, missing in zip(_texera_v, _texera_v.isna()) + | ] + | _texera_by.append("sign") + |out1df = out1df.loc[ + | _texera_keys.sort_values(by=_texera_by, kind="mergesort").index + |].reset_index(drop=True)""".stripMargin + } + } diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/aggregate/AggregateOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/aggregate/AggregateOpDescSpec.scala index 0f0e372ad66..8d344fd73c0 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/aggregate/AggregateOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/aggregate/AggregateOpDescSpec.scala @@ -19,17 +19,28 @@ package org.apache.texera.amber.operator.aggregate -import org.apache.texera.amber.core.tuple.{AttributeType, Schema} +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.metadata.{OperatorGroupConstants, OperatorMetadataGenerator} +import com.typesafe.config.ConfigFactory +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.util.concurrent.TimeUnit +import scala.io.Source import scala.jdk.CollectionConverters._ +import scala.util.Try class AggregateOpDescSpec extends AnyFlatSpec with Matchers { + private val NeedsPythonPackages = Tag(classOf[IntegrationTest].getName) + private val workflowId = WorkflowIdentity(1L) private val executionId = ExecutionIdentity(1L) @@ -108,6 +119,465 @@ class AggregateOpDescSpec extends AnyFlatSpec with Matchers { } } + // The integers behind a timestamp column mean microseconds or nanoseconds + // depending on the resolution it was read at, and a nanosecond total leaves + // the range of a 64-bit integer after a handful of modern dates. + it should "read a timestamp as epoch milliseconds at either resolution" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val input = Schema().add("t", AttributeType.TIMESTAMP) + val desc = descWith( + List.empty, + aggOp(AggregationFunction.SUM, "t", "ts_total"), + aggOp(AggregationFunction.AVERAGE, "t", "ts_avg") + ) + val block = desc.generateStandaloneCode(Map(PortIdentity() -> input)) + + def report(unit: String): String = + s"""in1df = pd.DataFrame({ + | "t": pd.to_datetime(["2024-01-01"] * 6).astype("datetime64[$unit]"), + |}) + |$block + |print(pd.Timestamp(out1df.iloc[0]["ts_total"]).isoformat()) + |print(repr(float(out1df.iloc[0]["ts_avg"])))""".stripMargin + + val driver = + s"""import pandas as pd + |${desc.standaloneHelpers().mkString("\n")} + |${report("ns")} + |${report("us")} + |""".stripMargin + + val out = runPython(python, driver, "aggregate-timestamp-units-", "UTC") + + withClue(s"python said:\n$out\nscript:\n$driver") { + // Six times 1704067200000 milliseconds, read back as a timestamp: a date + // no nanosecond count can hold, and the same one at either resolution. + out.trim.linesIterator.toSeq shouldBe Seq( + "2293-12-31T00:00:00", + "1704067200000.0", + "2293-12-31T00:00:00", + "1704067200000.0" + ) + } + } + + /** `Timestamp.getTime` answers for the instant a wall clock names in the + * JVM's default zone, and `new Timestamp(long)` renders one back there, so a + * script that read the column as UTC was a whole offset out on every value + * it added. Checked against Java 17 in the same zone: the sum below is + * 2078-01-01 06:00:00 and the average 1704132000000. + */ + it should "add timestamps in the zone the script runs in" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val input = Schema().add("t", AttributeType.TIMESTAMP) + val desc = descWith( + List.empty, + aggOp(AggregationFunction.SUM, "t", "ts_total"), + aggOp(AggregationFunction.AVERAGE, "t", "ts_avg") + ) + val block = desc.generateStandaloneCode(Map(PortIdentity() -> input)) + + val driver = + s"""import pandas as pd + |${desc.standaloneHelpers().mkString("\n")} + |in1df = pd.DataFrame({ + | "t": pd.to_datetime(["2024-01-01 00:00:00", "2024-01-02 00:00:00"]), + |}) + |$block + |print(pd.Timestamp(out1df.iloc[0]["ts_total"]).isoformat()) + |print(repr(float(out1df.iloc[0]["ts_avg"]))) + |""".stripMargin + + // A zone six hours behind UTC and with no daylight saving of its own, so + // one offset covers both dates. + val out = runPython(python, driver, "aggregate-timestamp-zone-", "America/Mexico_City") + + withClue(s"python said:\n$out\nscript:\n$driver") { + out.trim.linesIterator.toSeq shouldBe Seq("2078-01-01T06:00:00", "1704132000000.0") + } + } + + /** With no zone named, as on Windows without tzlocal, the script asks + * tzlocal(), which asks localtime, and Windows refuses an instant before + * 1970: the column raised and the sum came back in UTC. The driver stands in + * for both. Checked against Java 17 in the same zone: the sum below is + * 1653-12-31 18:00:00 and the average -4.9859712E12. + */ + it should "add timestamps before 1970 where localtime refuses them" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val input = Schema().add("t", AttributeType.TIMESTAMP) + val desc = descWith( + List.empty, + aggOp(AggregationFunction.SUM, "t", "ts_total"), + aggOp(AggregationFunction.AVERAGE, "t", "ts_avg") + ) + val block = desc.generateStandaloneCode(Map(PortIdentity() -> input)) + + val driver = + s"""import time + |import numpy as np + |import pandas as pd + |${desc.standaloneHelpers().mkString("\n")} + |_localtime = time.localtime + |def _refusing(seconds=None): + | if seconds is not None and seconds < 0: + | raise OSError(22, "Invalid argument") + | return _localtime() if seconds is None else _localtime(seconds) + |time.localtime = _refusing + |_texera_local_zone = lambda: None + |in1df = pd.DataFrame({ + | "t": np.array(["1600-01-01T12:00", "2024-01-01T00:00"], dtype="datetime64[ms]"), + |}) + |$block + |print(pd.Timestamp(out1df.iloc[0]["ts_total"]).isoformat()) + |print(repr(float(out1df.iloc[0]["ts_avg"]))) + |""".stripMargin + + val out = runPython(python, driver, "aggregate-timestamp-refused-", "America/Mexico_City") + + withClue(s"python said:\n$out\nscript:\n$driver") { + out.trim.linesIterator.toSeq shouldBe Seq("1653-12-31T18:00:00", "-4985971200000.0") + } + } + + // A zone whose offset has moved: Mexico City was seven hours behind in 1931, + // kept summer time until 2022 and its local mean time until 1922, and before + // 1900 java.util.Date reads it at today's standard offset. One group per row, + // so each AVERAGE is that wall clock's epoch milliseconds. + private val mexicoCityWallClocks = Seq( + "1930-12-31 17:00:00", + "2021-06-30 19:00:00", + "1920-12-31 17:23:24", + "1850-01-01 00:00:00" + ) + + it should "read each timestamp with the offset its zone had then" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val input = Schema().add("k", AttributeType.INTEGER).add("t", AttributeType.TIMESTAMP) + val desc = descWith(List("k"), aggOp(AggregationFunction.AVERAGE, "t", "ts_avg")) + val block = desc.generateStandaloneCode(Map(PortIdentity() -> input)) + + val values = mexicoCityWallClocks.map("\"" + _ + "\"").mkString("[", ", ", "]") + val driver = + s"""import pandas as pd + |${desc.standaloneHelpers().mkString("\n")} + |in1df = pd.DataFrame({ + | "k": list(range(${mexicoCityWallClocks.size})), + | "t": pd.to_datetime($values), + |}) + |$block + |for avg in out1df.sort_values("k")["ts_avg"]: + | print(int(avg)) + |""".stripMargin + + val zone = "America/Mexico_City" + val out = runPython(python, driver, "aggregate-timestamp-history-", zone) + + val machineZone = java.util.TimeZone.getDefault + java.util.TimeZone.setDefault(java.util.TimeZone.getTimeZone(zone)) + val fromEngine = + try mexicoCityWallClocks.map(w => java.sql.Timestamp.valueOf(w).getTime.toString) + finally java.util.TimeZone.setDefault(machineZone) + + withClue(s"python said:\n$out\nscript:\n$driver") { + out.trim.linesIterator.toSeq shouldBe fromEngine + } + fromEngine shouldBe Seq("-1230768000000", "1625097600000", "-1546300800000", "-3786804000000") + } + + /** Runs the generated block with the zone named to the child, because a + * timestamp's arithmetic reads the one in force where it runs. + */ + private def runPython(python: String, driver: String, prefix: String, zone: String): String = { + val script = Files.createTempFile(prefix, ".py") + script.toFile.deleteOnExit() + Files.write(script, driver.getBytes(StandardCharsets.UTF_8)) + val builder = new ProcessBuilder(python, script.toString).redirectErrorStream(true) + builder.environment().put("TZ", zone) + 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 + } + + // A group keyed on a missing value and one keyed on a NaN are two groups to + // the engine. The two are only distinct in a nullable dtype, which is what an + // Arrow file is read into. + it should "group a missing key apart from a NaN one" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val desc = descWith(List("k"), aggOp(AggregationFunction.SUM, "v", "total")) + + val driver = + s"""import pandas as pd, numpy as np + | + |in1df = pd.DataFrame({ + | "k": pd.arrays.FloatingArray(np.array([0.0, np.nan]), np.array([True, False])), + | "v": [1, 2], + |}) + |${desc.generateStandaloneCode()} + |print(len(out1df), sorted(int(x) for x in out1df["total"])) + |""".stripMargin + + val script = Files.createTempFile("aggregate-null-nan-", ".py") + script.toFile.deleteOnExit() + Files.write(script, driver.getBytes(StandardCharsets.UTF_8)) + val process = new ProcessBuilder(python, script.toString).redirectErrorStream(true).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 + // Two groups of one row each, not one group totalling 3. + out.trim shouldBe "2 [1, 2]" + } + } + + // The engine keys a group on a List of the key fields in a HashMap, under + // Scala's ==: NaN equals nothing, so each NaN row is a group of its own, + // while 0.0 and -0.0 are one group and so are two nulls. Both stages run, as + // the plan runs them. + it should "group NaN and signed-zero keys as the engine does" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val count = aggOp(AggregationFunction.COUNT, "v", "n") + val desc = descWith(List("k"), count) + val keys: Seq[java.lang.Double] = Seq(Double.NaN, Double.NaN, 0.0, -0.0, null, null, 1.0) + + val input = Schema().add("k", AttributeType.DOUBLE).add("v", AttributeType.INTEGER) + val local = new AggregateOpExec(objectMapper.writeValueAsString(desc)) + local.open() + keys.foreach(k => local.processTuple(Tuple(input, Array[Any](k, 1)), 0)) + val partials = local.onFinish(0).toList + val partialSchema = Schema().add("k", AttributeType.DOUBLE).add("n", AttributeType.INTEGER) + val global = + new AggregateOpExec(objectMapper.writeValueAsString(descWith(List("k"), count.getFinal))) + global.open() + partials.foreach(t => global.processTuple(Tuple(partialSchema, t.getFields.toArray[Any]), 0)) + val engine = global + .onFinish(0) + .map { row => + val key = row.getFields.head match { + case null => "None" + case d: java.lang.Double if d.isNaN => "nan" + case d => d.toString + } + s"$key ${row.getFields(1)}" + } + .toList + .sorted + + val pyKeys = keys.map(k => if (k == null) "None" else if (k.isNaN) "np.nan" else k.toString) + val driver = + s"""import pandas as pd, numpy as np + | + |keys = [${pyKeys.mkString(", ")}] + |in1df = pd.DataFrame({ + | "k": pd.arrays.FloatingArray( + | np.array([0.0 if k is None else k for k in keys]), + | np.array([k is None for k in keys]), + | ), + | "v": [1] * len(keys), + |}) + |${desc.generateStandaloneCode()} + |rows = [ + | ("None" if k is pd.NA else repr(float(k))) + " " + str(int(n)) + | for k, n in zip(out1df["k"], out1df["n"]) + |] + |print("\\n".join(sorted(rows))) + |""".stripMargin + + val script = Files.createTempFile("aggregate-nan-zero-", ".py") + script.toFile.deleteOnExit() + Files.write(script, driver.getBytes(StandardCharsets.UTF_8)) + val process = new ProcessBuilder(python, script.toString).redirectErrorStream(true).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.trim.linesIterator.toList shouldBe engine + } + } + + it should "emit no row for an empty input with no group keys, as the engine does" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val desc = descWith(List.empty, aggOp(AggregationFunction.COUNT, "", "cnt")) + + val exec = new AggregateOpExec(objectMapper.writeValueAsString(desc)) + exec.open() + exec.onFinish(0).toList shouldBe empty + + val driver = + s"""import pandas as pd + | + |in1df = pd.DataFrame({"v": pd.Series([], dtype="int64")}) + |${desc.generateStandaloneCode()} + |print(len(out1df), list(out1df.columns)) + |""".stripMargin + val out = runPython(python, driver, "aggregate-empty-", "UTC") + withClue(s"script:\n$driver")(out.trim shouldBe "0 ['cnt']") + } + + // A DOUBLE column read from Arrow keeps a stored NaN apart from null. The + // engine skips the null and folds the NaN in, so SUM, AVERAGE and MAX answer + // NaN while MIN, which Java's Double.compare sends NaN past, answers 1.0. + // CONCAT writes the NaN out and only a null as the empty string. + it should "fold a stored NaN in a DOUBLE column as the engine does" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val input = Schema().add("k", AttributeType.STRING).add("v", AttributeType.DOUBLE) + // Group "a" holds a NaN between two numbers and group "b" only nulls. + val rows: Seq[(String, java.lang.Double)] = + Seq(("a", 1.0), ("a", Double.NaN), ("a", 3.0), ("b", null), ("b", null)) + val functions = + Seq( + AggregationFunction.SUM, + AggregationFunction.AVERAGE, + AggregationFunction.MIN, + AggregationFunction.MAX, + AggregationFunction.CONCAT + ) + + def render(value: Any): String = + value match { + case null => "None" + case d: java.lang.Double if d.isNaN => "nan" + case d => d.toString + } + + def check(keys: List[String]): Unit = { + val desc = + () => descWith(keys, functions.map(fn => aggOp(fn, "v", fn.toString.toLowerCase)): _*) + + val exec = new AggregateOpExec(objectMapper.writeValueAsString(desc())) + exec.open() + rows.foreach { case (k, v) => exec.processTuple(Tuple(input, Array[Any](k, v)), 0) } + val engine = exec + .onFinish(0) + .map(_.getFields.toSeq.map(render).mkString(" ")) + .toSeq + .sorted + + val driver = + s"""import pandas as pd, numpy as np + | + |in1df = pd.DataFrame({ + | "k": ["a", "a", "a", "b", "b"], + | "v": pd.arrays.FloatingArray( + | np.array([1.0, np.nan, 3.0, 0.0, 0.0]), + | np.array([False, False, False, True, True]), + | ), + |}) + |${desc().standaloneHelpers().mkString("\n")} + |${desc().generateStandaloneCode(Map(PortIdentity() -> input))} + |def render(v): + | if v is pd.NA or v is None: + | return "None" + | return v if isinstance(v, str) else repr(float(v)) + |for row in sorted(" ".join(render(v) for v in r) for r in out1df.itertuples(index=False)): + | print(row) + |""".stripMargin + val out = runPython(python, driver, "aggregate-double-nan-", "UTC") + // No trim: an all-null group's CONCAT is "", which ends its line in a space. + withClue(s"script:\n$driver")(out.linesIterator.toSeq shouldBe engine) + } + + check(List("k")) + check(List.empty) + } + + // The engine compares timestamps by getTime, so two in the same millisecond + // tie and MIN and MAX keep the one they met first, with its own digits past + // the millisecond. Group "a" makes MIN keep the later digits, "b" makes MAX + // keep the earlier ones, "c" holds only a null, and "d" ties before 1970. + it should "pick a timestamp MIN and MAX by the millisecond as the engine does" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val input = Schema().add("k", AttributeType.STRING).add("t", AttributeType.TIMESTAMP) + val rows: Seq[(String, String)] = Seq( + ("a", "2024-01-01 12:00:00.0005"), + ("a", "2024-01-01 12:00:00.0001"), + ("b", "2024-01-01 12:00:00.0001"), + ("b", "2024-01-01 12:00:00.0005"), + ("c", null), + ("d", "1969-12-31 23:59:59.9991"), + ("d", "1969-12-31 23:59:59.9995") + ) + + def render(value: Any): String = + value match { + case null => "None" + case t: java.sql.Timestamp => t.toLocalDateTime.toString + case v => v.toString + } + + def check(keys: List[String]): Unit = { + val desc = () => + descWith( + keys, + aggOp(AggregationFunction.MIN, "t", "lo"), + aggOp(AggregationFunction.MAX, "t", "hi") + ) + + val exec = new AggregateOpExec(objectMapper.writeValueAsString(desc())) + exec.open() + rows.foreach { + case (k, t) => + val ts = Option(t).map(java.sql.Timestamp.valueOf).orNull + exec.processTuple(Tuple(input, Array[Any](k, ts)), 0) + } + val engine = exec + .onFinish(0) + .map(_.getFields.toSeq.map(render).mkString(" ")) + .toSeq + .sorted + + val ks = rows.map(r => "\"" + r._1 + "\"").mkString("[", ", ", "]") + val ts = rows + .map(r => Option(r._2).map("\"" + _.replace(' ', 'T') + "\"").getOrElse("\"NaT\"")) + .mkString("[", ", ", "]") + val driver = + s"""import pandas as pd, numpy as np + | + |in1df = pd.DataFrame({"k": $ks, "t": np.array($ts, dtype="datetime64[us]")}) + |${desc().standaloneHelpers().mkString("\n")} + |${desc().generateStandaloneCode(Map(PortIdentity() -> input))} + |def render(v): + | if v is None or pd.isna(v): + | return "None" + | return v if isinstance(v, str) else pd.Timestamp(v).isoformat() + |for row in sorted(" ".join(render(v) for v in r) for r in out1df.itertuples(index=False)): + | print(row) + |""".stripMargin + val out = runPython(python, driver, "aggregate-timestamp-ties-", "UTC") + withClue(s"script:\n$driver")(out.trim.linesIterator.toSeq shouldBe engine) + } + + check(List("k")) + check(List.empty) + } + + // Without a schema the type-led branches cannot be chosen. + it should "fall back to pandas' own answers when no schema is given" in { + val desc = descWith(List.empty, aggOp(AggregationFunction.SUM, "i", "int_total")) + desc.generateStandaloneCode() should include("in1df[\"i\"].sum()") + desc.generateStandaloneCode() should not include "_texera_agg_int_sum(in1df" + } + "AggregateOpDesc JSON schema" should "make the attribute optional only for count and required for every other function" in { val aggDef = OperatorMetadataGenerator @@ -131,4 +601,30 @@ class AggregateOpDescSpec extends AnyFlatSpec with Matchers { val elseRequired = rule.get("else").get("required").elements().asScala.map(_.asText()).toList elseRequired should contain("attribute") } + + 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 + } } diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/sort/StableMergeSortOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/sort/StableMergeSortOpDescSpec.scala index 933549a3b39..8458332d16e 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/sort/StableMergeSortOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/sort/StableMergeSortOpDescSpec.scala @@ -19,16 +19,29 @@ package org.apache.texera.amber.operator.sort +import com.typesafe.config.ConfigFactory import org.apache.texera.amber.core.executor.OpExecWithClassName +import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema, Tuple} import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} 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.util.concurrent.TimeUnit +import scala.collection.mutable.ListBuffer +import scala.io.Source +import scala.util.Try + class StableMergeSortOpDescSpec extends AnyFlatSpec with Matchers { + private val NeedsPythonPackages = Tag(classOf[IntegrationTest].getName) + private val workflowId = WorkflowIdentity(1L) private val executionId = ExecutionIdentity(1L) @@ -67,4 +80,260 @@ class StableMergeSortOpDescSpec extends AnyFlatSpec with Matchers { desc.keys.head.attributeName shouldBe "age" desc.keys.head.sortPreference shouldBe SortPreference.DESC } + + // A null goes last whichever way the key points; a NaN compares above every + // number, so it goes last ascending and first descending. Only a column read + // into a nullable dtype can hold both, which is where this is run. + "StableMergeSortOpDesc.generateStandaloneCode" should + "put a null last both ways and a NaN where the key points" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + def blockFor(pref: SortPreference): String = { + val unit = new SortCriteriaUnit + unit.attributeName = "v" + unit.sortPreference = pref + val d = new StableMergeSortOpDesc + d.keys = ListBuffer(unit) + d.generateStandaloneCode().linesIterator.map(" " + _).mkString("\n") + } + + val driver = + s"""import pandas as pd, numpy as np + | + |def frame(): + | # 3.0, a null, a stored NaN, 1.0, +inf + | col = pd.arrays.FloatingArray( + | np.array([3.0, 0.0, np.nan, 1.0, np.inf]), + | np.array([False, True, False, False, False]), + | ) + | return pd.DataFrame({"v": col}) + | + |def show(df): + | return [("NULL" if x is pd.NA else repr(float(x))) for x in df["v"]] + | + |if True: + | in1df = frame() + |${blockFor(SortPreference.ASC)} + | print("asc", show(out1df)) + |if True: + | in1df = frame() + |${blockFor(SortPreference.DESC)} + | print("desc", show(out1df)) + |""".stripMargin + + val script = Files.createTempFile("sort-nan-", ".py") + script.toFile.deleteOnExit() + Files.write(script, driver.getBytes(StandardCharsets.UTF_8)) + val process = new ProcessBuilder(python, script.toString).redirectErrorStream(true).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 + val lines = out.trim.linesIterator.toSeq + lines.head shouldBe "asc ['1.0', '3.0', 'inf', 'nan', 'NULL']" + lines(1) shouldBe "desc ['nan', 'inf', '3.0', '1.0', 'NULL']" + } + } + + // Double.compare puts -0.0 below 0.0, where pandas reads the two as equal and + // keeps them in input order. Each row carries its input position, and the + // engine runs the same rows both ways. + it should "order signed zeros as the engine does both ways" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val values: Seq[java.lang.Double] = Seq(0.0, -0.0, Double.NaN, 1.0, null, -0.0) + val schema = Schema().add("v", AttributeType.DOUBLE).add("i", AttributeType.INTEGER) + def descFor(pref: SortPreference): StableMergeSortOpDesc = { + val unit = new SortCriteriaUnit + unit.attributeName = "v" + unit.sortPreference = pref + val d = new StableMergeSortOpDesc + d.keys = ListBuffer(unit) + d + } + val prefs = Seq(SortPreference.ASC, SortPreference.DESC) + val engine = prefs.map { pref => + val exec = new StableMergeSortOpExec(objectMapper.writeValueAsString(descFor(pref))) + exec.open() + val rows = values.zipWithIndex.flatMap { + case (v, i) => exec.processTuple(Tuple(schema, Array[Any](v, i)), 0) + } ++ exec.onFinish(0) + rows.map(_.getFields(1).toString).mkString(" ") + } + + val calls = prefs + .map(p => s"order(${objectMapper.writeValueAsString(descFor(p).generateStandaloneCode())})") + .mkString("\n") + val driver = + s"""import pandas as pd, numpy as np + | + |def order(code): + | in1df = pd.DataFrame({ + | "v": pd.arrays.FloatingArray( + | np.array([0.0, -0.0, np.nan, 1.0, 0.0, -0.0]), + | np.array([False, False, False, False, True, False]), + | ), + | "i": range(6), + | }) + | ns = {"pd": pd, "in1df": in1df} + | exec(code, ns) + | print(" ".join(str(i) for i in ns["out1df"]["i"])) + | + |$calls + |""".stripMargin + + val script = Files.createTempFile("stable-merge-sort-zero-", ".py") + script.toFile.deleteOnExit() + Files.write(script, driver.getBytes(StandardCharsets.UTF_8)) + val process = new ProcessBuilder(python, script.toString).redirectErrorStream(true).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.trim.linesIterator.toList shouldBe engine + } + } + + // The sort tiers are named after nothing the input can be carrying, so a + // column that happens to answer to a helper's old name keeps its values. + it should "keep an input column named after a sort helper" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val unit = new SortCriteriaUnit + unit.attributeName = "x" + unit.sortPreference = SortPreference.ASC + val desc = new StableMergeSortOpDesc + desc.keys = ListBuffer(unit) + + val driver = + s"""import pandas as pd + | + |in1df = pd.DataFrame({ + | "x": [2, 1], + | "_texera_null_x": ["a", "b"], + | "_texera_nan_x": ["c", "d"], + |}) + |${desc.generateStandaloneCode()} + |print(list(out1df.columns)) + |print(list(out1df["_texera_null_x"]), list(out1df["_texera_nan_x"])) + |""".stripMargin + + val script = Files.createTempFile("sort-helper-name-", ".py") + script.toFile.deleteOnExit() + Files.write(script, driver.getBytes(StandardCharsets.UTF_8)) + val process = new ProcessBuilder(python, script.toString).redirectErrorStream(true).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 + val lines = out.trim.linesIterator.toSeq + lines.head shouldBe "['x', '_texera_null_x', '_texera_nan_x']" + lines(1) shouldBe "['b', 'a'] ['d', 'c']" + } + } + + // U+1F600 is stored as the surrogate pair D83D DE00, so the engine puts it + // before U+FF01, while Python's code-point order puts it after. + it should "order strings the way the engine does above U+FFFF" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + def cp(c: Int): String = new String(Character.toChars(c)) + val values = Seq(cp(0xff01), cp(0x1f600), "a", null, cp(0xffff)) + def show(v: String): String = + if (v == null) "NULL" else v.codePoints().toArray.map(c => f"$c%X").mkString("+") + + val schema = Schema().add(new Attribute("s", AttributeType.STRING)) + def descFor(pref: SortPreference): StableMergeSortOpDesc = { + val unit = new SortCriteriaUnit + unit.attributeName = "s" + unit.sortPreference = pref + val d = new StableMergeSortOpDesc + d.keys = ListBuffer(unit) + d + } + def engine(pref: SortPreference): String = { + val exec = new StableMergeSortOpExec(objectMapper.writeValueAsString(descFor(pref))) + exec.open() + values.foreach { v => + exec.processTuple(Tuple.builder(schema).add(schema.getAttribute("s"), v).build(), 0) + } + val out = exec.onFinish(0).map(_.asInstanceOf[Tuple].getField[String]("s")).toList + exec.close() + out.map(show).mkString(" ") + } + def block(pref: SortPreference): String = + descFor(pref).generateStandaloneCode().linesIterator.map(" " + _).mkString("\n") + + val literal = values + .map(v => + if (v == null) "None" + else v.codePoints().toArray.map(c => f"chr(0x$c%X)").mkString("+") + ) + .mkString("[", ", ", "]") + val driver = + s"""import pandas as pd + | + |def show(df): + | return " ".join( + | "NULL" if v is None else "+".join("%X" % ord(c) for c in v) for v in df["s"] + | ) + | + |if True: + | in1df = pd.DataFrame({"s": $literal}) + |${block(SortPreference.ASC)} + | print(show(out1df)) + |if True: + | in1df = pd.DataFrame({"s": $literal}) + |${block(SortPreference.DESC)} + | print(show(out1df)) + |""".stripMargin + + val script = Files.createTempFile("sort-utf16-", ".py") + script.toFile.deleteOnExit() + Files.write(script, driver.getBytes(StandardCharsets.UTF_8)) + val process = new ProcessBuilder(python, script.toString).redirectErrorStream(true).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 + val lines = out.trim.linesIterator.toSeq + engine(SortPreference.ASC) shouldBe "61 1F600 FF01 FFFF NULL" + lines.head shouldBe engine(SortPreference.ASC) + lines(1) shouldBe engine(SortPreference.DESC) + } + } + + 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 + } } diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/sortPartitions/SortPartitionsOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/sortPartitions/SortPartitionsOpDescSpec.scala index d9ae20b14a0..65d6a29e080 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/sortPartitions/SortPartitionsOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/sortPartitions/SortPartitionsOpDescSpec.scala @@ -19,17 +19,29 @@ package org.apache.texera.amber.operator.sortPartitions +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.RangePartition 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.util.concurrent.TimeUnit +import scala.io.Source +import scala.util.Try + class SortPartitionsOpDescSpec extends AnyFlatSpec with Matchers { + private val NeedsPythonPackages = Tag(classOf[IntegrationTest].getName) + private val workflowId = WorkflowIdentity(1L) private val executionId = ExecutionIdentity(1L) @@ -77,4 +89,124 @@ class SortPartitionsOpDescSpec extends AnyFlatSpec with Matchers { sp.domainMin shouldBe 1L sp.domainMax shouldBe 99L } + + // The engine puts a null first and a NaN after positive infinity. Only a + // column read into a nullable dtype can hold both, which is what an Arrow + // file is read into, so that is where this is run. + "SortPartitionsOpDesc.generateStandaloneCode" should + "put a null first and a NaN after infinity" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val driver = + s"""import pandas as pd, numpy as np + | + |# 3.0, a null, a stored NaN, 1.0, +inf + |in1df = pd.DataFrame({"v": pd.arrays.FloatingArray( + | np.array([3.0, 0.0, np.nan, 1.0, np.inf]), + | np.array([False, True, False, False, False]), + |)}) + |${newDesc("v", 0L, 10L).generateStandaloneCode()} + |print([("NULL" if x is pd.NA else repr(float(x))) for x in out1df["v"]]) + |""".stripMargin + + val script = Files.createTempFile("sort-partitions-nan-", ".py") + script.toFile.deleteOnExit() + Files.write(script, driver.getBytes(StandardCharsets.UTF_8)) + val process = new ProcessBuilder(python, script.toString).redirectErrorStream(true).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.trim shouldBe "['NULL', '1.0', '3.0', 'inf', 'nan']" + } + } + + // AttributeTypeUtils.compare: Double.compare puts -0.0 below 0.0, and a + // timestamp is compared by getTime, so two inside one millisecond are a tie + // the stable sort leaves in input order. Each row carries its input position. + it should "order signed zeros and same-millisecond timestamps as the engine does" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + def engineOrder(attrType: AttributeType, values: Seq[Any]): List[Int] = { + val schema = Schema().add("v", attrType).add("i", AttributeType.INTEGER) + val exec = new SortPartitionsOpExec(objectMapper.writeValueAsString(newDesc("v", 0L, 10L))) + exec.open() + values.zipWithIndex.foreach { + case (v, i) => exec.processTuple(Tuple(schema, Array[Any](v, i)), 0) + } + exec.onFinish(0).map(_.getFields(1).asInstanceOf[Int]).toList + } + def timestamp(micros: Long): java.sql.Timestamp = { + val t = new java.sql.Timestamp(micros / 1000) + t.setNanos(((micros % 1000000) * 1000).toInt) + t + } + + val doubles: Seq[java.lang.Double] = Seq(0.0, -0.0, Double.NaN, 1.0, null, -0.0) + // 2024-01-01 00:00:00.000900, .000100 and one millisecond before both. + val base = 1704067200000000L + val micros = Seq(base + 900, base + 100, base - 1000) + val engine = Seq( + engineOrder(AttributeType.DOUBLE, doubles), + engineOrder(AttributeType.TIMESTAMP, micros.map(timestamp)) + ).map(_.mkString(" ")) + + val code = newDesc("v", 0L, 10L).generateStandaloneCode() + val driver = + s"""import pandas as pd, numpy as np + | + |def order(column): + | in1df = pd.DataFrame({"v": column, "i": range(len(column))}) + | ns = {"pd": pd, "in1df": in1df} + | exec(${objectMapper.writeValueAsString(code)}, ns) + | print(" ".join(str(i) for i in ns["out1df"]["i"])) + | + |order(pd.arrays.FloatingArray( + | np.array([0.0, -0.0, np.nan, 1.0, 0.0, -0.0]), + | np.array([False, False, False, False, True, False]), + |)) + |order(pd.Series(np.array([${micros.mkString(", ")}], dtype="datetime64[us]"))) + |""".stripMargin + + val script = Files.createTempFile("sort-partitions-zero-ms-", ".py") + script.toFile.deleteOnExit() + Files.write(script, driver.getBytes(StandardCharsets.UTF_8)) + val process = new ProcessBuilder(python, script.toString).redirectErrorStream(true).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.trim.linesIterator.toList shouldBe engine + } + } + + 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 + } }