From 4b513835ef86db42fc35c981ed729a5959d13448 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Sat, 26 Sep 2026 20:14:33 -0700 Subject: [PATCH 1/3] feat(operator): export the JSONL and Arrow sources, and hand Python a zoneless microsecond timestamp Co-Authored-By: Claude Opus 5.5 (1M context) --- .../core/models/schema/attribute_type.py | 12 +- .../core/models/schema/test_attribute_type.py | 57 +++-- .../python/pytexera/udf/test_udf_operator.py | 9 +- .../amber/core/tuple/AttributeTypeUtils.scala | 23 +- .../apache/texera/amber/util/ArrowUtils.scala | 75 ++++-- .../texera/amber/util/ArrowUtilsSpec.scala | 46 +++- .../source/scan/arrow/ArrowSourceOpDesc.scala | 91 ++++++- .../scan/json/JSONLScanSourceOpDesc.scala | 233 +++++++++++++++++- .../scan/arrow/ArrowSourceOpDescSpec.scala | 208 +++++++++++++++- .../scan/json/JSONLScanSourceOpDescSpec.scala | 29 +++ .../texera/amber/util/ArrowUtilsSpec.scala | 19 +- 11 files changed, 716 insertions(+), 86 deletions(-) diff --git a/amber/src/main/python/core/models/schema/attribute_type.py b/amber/src/main/python/core/models/schema/attribute_type.py index 0887816c828..1e0fdb900e3 100644 --- a/amber/src/main/python/core/models/schema/attribute_type.py +++ b/amber/src/main/python/core/models/schema/attribute_type.py @@ -98,18 +98,18 @@ def _parse_bool(v): def _parse_timestamp(v): + # A TIMESTAMP reaches a worker as a wall clock with no zone, so a value + # parsed here holds none either and compares with a row's. An offset is + # moved to the machine's zone first, as the engine's DateParserUtils does. if _is_empty_value(v): - return datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) + return datetime.datetime(1970, 1, 1) normalized_value = str(v) if normalized_value.endswith("Z"): normalized_value = normalized_value[:-1] + "+00:00" parsed_value = datetime.datetime.fromisoformat(normalized_value) - if ( - parsed_value.tzinfo is None - or parsed_value.tzinfo.utcoffset(parsed_value) is None - ): - return parsed_value.replace(tzinfo=datetime.timezone.utc) + if parsed_value.utcoffset() is not None: + return parsed_value.astimezone().replace(tzinfo=None) return parsed_value diff --git a/amber/src/test/python/core/models/schema/test_attribute_type.py b/amber/src/test/python/core/models/schema/test_attribute_type.py index 3f737dbb2c3..302b2ff7651 100644 --- a/amber/src/test/python/core/models/schema/test_attribute_type.py +++ b/amber/src/test/python/core/models/schema/test_attribute_type.py @@ -16,6 +16,7 @@ # under the License. import datetime +import time import pytest @@ -31,7 +32,19 @@ parse_bool = FROM_STRING_PARSER_MAPPING[AttributeType.BOOL] parse_timestamp = FROM_STRING_PARSER_MAPPING[AttributeType.TIMESTAMP] -EPOCH = datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) +EPOCH = datetime.datetime(1970, 1, 1) + + +@pytest.fixture +def machine_zone(monkeypatch): + # An offset is moved to the machine's zone, so the tests set one. + def set_zone(name): + monkeypatch.setenv("TZ", name) + time.tzset() + + yield set_zone + monkeypatch.undo() + time.tzset() class TestParseBool: @@ -60,35 +73,25 @@ def test_a_non_numeric_non_literal_value_is_rejected(self): class TestParseTimestamp: @pytest.mark.parametrize("empty", [None, "", " ", "\t\n"]) - def test_an_absent_value_parses_as_the_utc_epoch(self, empty): + def test_an_absent_value_parses_as_the_epoch_with_no_zone(self, empty): parsed = parse_timestamp(empty) - # Assert the exact instant *and* the tzinfo: a naive 1970-01-01 would - # compare unequal here, so dropping the timezone is caught too. assert parsed == EPOCH - assert parsed.tzinfo == datetime.timezone.utc - assert (parsed.year, parsed.month, parsed.day) == (1970, 1, 1) - - def test_a_zulu_suffix_yields_a_utc_aware_instant(self): - # Named for the observable outcome, not for the code that produces it. - # `_parse_timestamp` rewrites a trailing "Z" into "+00:00" before - # calling `fromisoformat`, but `fromisoformat` has accepted "Z" itself - # since Python 3.11 and the pyamber CI matrix is 3.11/3.12/3.13 -- so - # deleting that rewrite leaves the entire suite green, and this test - # must not be credited with pinning it. What it does pin is the - # resulting instant and its offset (and, uniquely in this file, the - # "+00:00" constant, should the rewrite ever run on an older runtime). - assert parse_timestamp("2024-05-06T07:08:09Z") == datetime.datetime( - 2024, 5, 6, 7, 8, 9, tzinfo=datetime.timezone.utc - ) + assert parsed.tzinfo is None + + def test_a_value_with_no_offset_keeps_its_wall_clock(self, machine_zone): + machine_zone("Asia/Kolkata") + parsed = parse_timestamp("2024-05-06T07:08:09") + assert parsed == datetime.datetime(2024, 5, 6, 7, 8, 9) + assert parsed.tzinfo is None - def test_a_naive_value_is_assumed_to_be_utc(self): - assert parse_timestamp("2024-05-06T07:08:09") == datetime.datetime( - 2024, 5, 6, 7, 8, 9, tzinfo=datetime.timezone.utc + def test_a_zulu_suffix_is_moved_to_the_machine_zone(self, machine_zone): + machine_zone("Asia/Kolkata") + assert parse_timestamp("2024-05-06T07:08:09Z") == datetime.datetime( + 2024, 5, 6, 12, 38, 9 ) - def test_an_explicit_offset_is_preserved(self): + def test_an_explicit_offset_is_moved_to_the_machine_zone(self, machine_zone): + machine_zone("UTC") parsed = parse_timestamp("2024-05-06T07:08:09+02:00") - assert parsed.utcoffset() == datetime.timedelta(hours=2) - assert parsed == datetime.datetime( - 2024, 5, 6, 5, 8, 9, tzinfo=datetime.timezone.utc - ) + assert parsed == datetime.datetime(2024, 5, 6, 5, 8, 9) + assert parsed.tzinfo is None diff --git a/amber/src/test/python/pytexera/udf/test_udf_operator.py b/amber/src/test/python/pytexera/udf/test_udf_operator.py index 86cd2a88c51..73ea4ca9873 100644 --- a/amber/src/test/python/pytexera/udf/test_udf_operator.py +++ b/amber/src/test/python/pytexera/udf/test_udf_operator.py @@ -206,7 +206,7 @@ def test_injected_values_are_applied_before_open(self): assert operator.count_parameter.value == 7 assert operator.enabled_parameter.value is True assert operator.created_at_parameter.value == datetime.datetime( - 2024, 1, 1, 0, 0, tzinfo=datetime.timezone.utc + 2024, 1, 1, 0, 0 ) def test_duplicate_parameter_names_with_conflicting_types_raise(self): @@ -238,12 +238,7 @@ def test_duplicate_parameter_names_with_same_type_succeed(self): ( "2024-01-01T00:00:00", AttributeType.TIMESTAMP, - datetime.datetime(2024, 1, 1, 0, 0, tzinfo=datetime.timezone.utc), - ), - ( - "2024-01-01T00:00:00Z", - AttributeType.TIMESTAMP, - datetime.datetime(2024, 1, 1, 0, 0, tzinfo=datetime.timezone.utc), + datetime.datetime(2024, 1, 1, 0, 0), ), ], ) diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/tuple/AttributeTypeUtils.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/tuple/AttributeTypeUtils.scala index 17c3dfef33c..66a955f49ad 100644 --- a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/tuple/AttributeTypeUtils.scala +++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/tuple/AttributeTypeUtils.scala @@ -144,8 +144,14 @@ object AttributeTypeUtils extends Serializable { } else { str.trim.toInt } - case int: Integer => int - case long: java.lang.Long => long.toInt + case int: Integer => int + case long: java.lang.Long => long.toInt + // An Arrow file states the width of its own integers, and a narrow + // column hands its values over as Shorts or Bytes. Without this the parse + // threw, the Arrow source caught it, and every value in the column + // arrived null. + case short: java.lang.Short => short.toInt + case byte: java.lang.Byte => byte.toInt case double: java.lang.Double => double.toInt case boolean: java.lang.Boolean => if (boolean) 1 else 0 // Timestamp and Binary are considered to be illegal here. @@ -232,10 +238,15 @@ object AttributeTypeUtils extends Serializable { def parseDouble(fieldValue: Any): java.lang.Double = { val attempt: Try[Double] = Try { fieldValue match { - case str: String => str.trim.toDouble - case int: Integer => int.toDouble - case long: java.lang.Long => long.toDouble - case double: java.lang.Double => double + case str: String => str.trim.toDouble + case int: Integer => int.toDouble + case long: java.lang.Long => long.toDouble + case double: java.lang.Double => double + // A single-precision column hands its values over as Floats, which only + // an Arrow file produces: Texera writes every double it owns as eight + // bytes. Without this the parse threw, the Arrow source caught it, and + // the whole column arrived null. + case float: java.lang.Float => float.toDouble case boolean: java.lang.Boolean => if (boolean) 1 else 0 // Timestamp and Binary are considered to be illegal here. case _ => diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala index 434519b3dea..9b6b0594d91 100644 --- a/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala +++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala @@ -84,11 +84,16 @@ object ArrowUtils extends LazyLogging { // Use the attribute type from the schema (which includes metadata) // instead of deriving it from the Arrow type val attributeType = schema.getAttributes(index).getType - // A timestamp is the one type whose field says more than the - // schema does, so it is the only one that reads the field. + // A timestamp and an unsigned integer are the types whose field + // says more than the schema does, so they are the ones that read + // the field. attributeType match { case AttributeType.TIMESTAMP => wallClockOf(value, fieldVector.getField.getType) - case _ => AttributeTypeUtils.parseField(value, attributeType) + case _ => + AttributeTypeUtils.parseField( + unsignedValueOf(value, fieldVector.getField.getType), + attributeType + ) } } catch { case e: Exception => @@ -101,6 +106,29 @@ object ArrowUtils extends LazyLogging { .build() } + /** The number an unsigned column counts to, out of the storage it counts in. + * + * Arrow's unsigned vectors hand back the raw storage, signed: `UInt1Vector` a + * Byte, `UInt4Vector` an Int, each reading the file's largest value as -1. + * `UInt2Vector` is the exception and hands back a Character, which is already + * the number but is no kind of integer to the parse below. Every other value + * passes through untouched. + * + * Nothing wider is handled because nothing wider arrives: [[toAttributeType]] + * refuses an unsigned 64-bit column, having no Texera type to hold it. + */ + private def unsignedValueOf(value: AnyRef, arrowType: ArrowType): AnyRef = + arrowType match { + case int: ArrowType.Int if !int.getIsSigned => + value match { + case byte: java.lang.Byte => Int.box(byte & 0xff) + case char: java.lang.Character => Int.box(char.toInt) + case integer: Integer => Long.box(integer.toLong & 0xffffffffL) + case other => other + } + case _ => value + } + /** The wall clock a timestamp column holds, read the way its own field states. * * A Texera TIMESTAMP carries no zone, so the wall clock is the whole of what @@ -140,13 +168,17 @@ object ArrowUtils extends LazyLogging { case TimeUnit.NANOSECOND => Instant.EPOCH.plusNanos(number) } - /** The number a field of this unit records an instant as, inverting [[instantOf]]. */ + /** The number a field of this unit records an instant as, inverting [[instantOf]]. + * Microseconds are counted from the seconds, because ChronoUnit.MICROS goes + * through nanoseconds and overflows past 2262, a year a Timestamp holds. + */ private def numberOf(instant: Instant, unit: TimeUnit): Long = unit match { case TimeUnit.SECOND => instant.getEpochSecond case TimeUnit.MILLISECOND => instant.toEpochMilli - case TimeUnit.MICROSECOND => ChronoUnit.MICROS.between(Instant.EPOCH, instant) - case TimeUnit.NANOSECOND => ChronoUnit.NANOS.between(Instant.EPOCH, instant) + case TimeUnit.MICROSECOND => + Math.addExact(Math.multiplyExact(instant.getEpochSecond, 1000000L), instant.getNano / 1000L) + case TimeUnit.NANOSECOND => ChronoUnit.NANOS.between(Instant.EPOCH, instant) } /** @@ -182,15 +214,23 @@ object ArrowUtils extends LazyLogging { @throws[AttributeTypeException] def toAttributeType(srcType: ArrowType): AttributeType = { srcType match { + // An unsigned column counts up where its storage counts down: the largest + // unsigned 32-bit value is stored as -1, so read as its storage it would + // arrive as -1 where the file means 4294967295. The next Texera integer up + // holds it, and [[unsignedValueOf]] does the reading. Past 64 bits there is + // no next one. The same widening covers the narrow widths, Texera having no + // column shorter than a 32-bit integer. case int: ArrowType.Int => - int.getBitWidth match { - case 16 | 32 => - AttributeType.INTEGER - - case 64 => - AttributeType.LONG - - case other => + (int.getBitWidth, int.getIsSigned) match { + case (8 | 16 | 32, true) => AttributeType.INTEGER + case (8 | 16, false) => AttributeType.INTEGER + case (64, true) => AttributeType.LONG + case (32, false) => AttributeType.LONG + case (64, false) => + throw new AttributeTypeUtils.AttributeTypeException( + "Unsupported unsigned 64-bit Int, which is wider than any Texera column" + ) + case (other, _) => throw new AttributeTypeUtils.AttributeTypeException( s"Unsupported Int bit width: $other" ) @@ -366,8 +406,13 @@ object ArrowUtils extends LazyLogging { case AttributeType.BOOLEAN => ArrowType.Bool.INSTANCE + // A wall clock with no zone, to the microsecond, which is what a Python + // worker maps TIMESTAMP to and the finest a datetime holds. Labelled UTC + // in milliseconds, a Python operator was handed an aware moment cut to + // the millisecond, where a JVM operator reads the same row with no zone + // and to the nanosecond. case AttributeType.TIMESTAMP => - new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC") + new ArrowType.Timestamp(TimeUnit.MICROSECOND, null) case AttributeType.BINARY => new ArrowType.Binary diff --git a/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala b/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala index e3f8014462d..e157f9892eb 100644 --- a/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala +++ b/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala @@ -50,15 +50,27 @@ class ArrowUtilsSpec extends AnyFlatSpec with Matchers { ArrowUtils.toAttributeType(new ArrowType.Int(64, true)) shouldBe AttributeType.LONG } - it should "throw AttributeTypeException for non-standard Int bit-widths" in { - // Only 16/32 (INTEGER) and 64 (LONG) are supported. Other widths used to - // be silently coerced to LONG by a `case 64 | _` catch-all; they now - // raise rather than masquerade as Int64. + it should "map every integer to the narrowest Texera type that holds it" in { + // Texera has no column shorter than a 32-bit integer, so a narrower width + // reads as one. An unsigned column needs the width above its own, counting + // up where its storage counts down: the largest unsigned 32-bit value is + // past what an INTEGER holds. + ArrowUtils.toAttributeType(new ArrowType.Int(8, true)) shouldBe AttributeType.INTEGER + ArrowUtils.toAttributeType(new ArrowType.Int(8, false)) shouldBe AttributeType.INTEGER + ArrowUtils.toAttributeType(new ArrowType.Int(16, false)) shouldBe AttributeType.INTEGER + ArrowUtils.toAttributeType(new ArrowType.Int(32, false)) shouldBe AttributeType.LONG + } + + it should "throw AttributeTypeException for an integer no Texera type holds" in { + // A width above 64 used to be silently coerced to LONG by a `case 64 | _` + // catch-all; it raises rather than masquerade as Int64. An unsigned 64-bit + // column raises for the same reason, there being no width above it to read + // it as. assertThrows[AttributeTypeException] { - ArrowUtils.toAttributeType(new ArrowType.Int(8, true)) + ArrowUtils.toAttributeType(new ArrowType.Int(128, true)) } assertThrows[AttributeTypeException] { - ArrowUtils.toAttributeType(new ArrowType.Int(128, true)) + ArrowUtils.toAttributeType(new ArrowType.Int(64, false)) } } @@ -114,9 +126,9 @@ class ArrowUtilsSpec extends AnyFlatSpec with Matchers { ArrowUtils.fromAttributeType(AttributeType.BOOLEAN) shouldBe ArrowType.Bool.INSTANCE } - it should "map TIMESTAMP to Timestamp(MILLISECOND, UTC)" in { + it should "map TIMESTAMP to a zoneless Timestamp(MICROSECOND)" in { val arrow = ArrowUtils.fromAttributeType(AttributeType.TIMESTAMP) - arrow shouldBe new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC") + arrow shouldBe new ArrowType.Timestamp(TimeUnit.MICROSECOND, null) } it should "map BINARY to ArrowType.Binary" in { @@ -357,6 +369,20 @@ class ArrowUtilsSpec extends AnyFlatSpec with Matchers { } } + // A Timestamp reaches past 2262, where counting microseconds through + // nanoseconds overflowed, and before 1970. Both come back to the microsecond. + it should "round-trip a timestamp either side of the nanosecond range" in { + val schema = Schema(List(new Attribute("t", AttributeType.TIMESTAMP))) + Seq("2500-01-01 00:00:00.123456", "1600-06-15 08:30:00.5").foreach { text => + val tuple = + Tuple.builder(schema).addSequentially(Array[Any](Timestamp.valueOf(text))).build() + withRoot(schema) { root => + ArrowUtils.appendTexeraTuple(tuple, root) + ArrowUtils.getTexeraTuple(0, root).getField[Timestamp]("t") shouldBe Timestamp.valueOf(text) + } + } + } + it should "append consecutive tuples at increasing row indices" in { val schema = Schema(List(new Attribute("s", AttributeType.STRING))) val first = Tuple.builder(schema).addSequentially(Array[Any]("first")).build() @@ -436,9 +462,9 @@ class ArrowUtilsSpec extends AnyFlatSpec with Matchers { } } - // ----- Timestamp fields that are not the UTC millisecond ones we write ----- + // ----- Timestamp fields that are not the zoneless microsecond ones we write ----- - // fromTexeraSchema only ever writes Timestamp(MILLISECOND, "UTC"), so the + // fromTexeraSchema only ever writes a zoneless Timestamp(MICROSECOND), so the // roots built above never exercise another zone or unit. An .arrow file handed // to ArrowSourceOpDesc can carry either: pandas writes a tz-aware column as // Timestamp(NANOSECOND, ). These build the field directly. diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/arrow/ArrowSourceOpDesc.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/arrow/ArrowSourceOpDesc.scala index ad1d7a34176..a4081d886cd 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/arrow/ArrowSourceOpDesc.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/arrow/ArrowSourceOpDesc.scala @@ -25,6 +25,8 @@ import org.apache.texera.amber.core.storage.DocumentFactory import org.apache.texera.amber.core.tuple.Schema import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} import org.apache.texera.amber.core.workflow.{PhysicalOp, SchemaPropagationFunc} +import org.apache.texera.amber.operator.StandaloneCodeGenerator +import org.apache.texera.amber.operator.StandaloneCodeGenerator.SourceFilePlaceholder import org.apache.texera.amber.operator.source.scan.ScanSourceOpDesc import org.apache.texera.amber.util.ArrowUtils import org.apache.texera.amber.util.JSONUtils.objectMapper @@ -38,10 +40,97 @@ import java.nio.file.{Files, StandardOpenOption} import scala.util.Using @JsonIgnoreProperties(value = Array("fileEncoding")) -class ArrowSourceOpDesc extends ScanSourceOpDesc { +class ArrowSourceOpDesc extends ScanSourceOpDesc with StandaloneCodeGenerator { fileTypeName = Option("Arrow") + // pyarrow is what reads an Arrow file, pandas going through it; the block + // below asks it for the file's own columns. + override def standaloneImports(): Seq[String] = Seq("import pyarrow as pa") + + override def standaloneSourcePath(): Option[String] = fileName + + override def generateStandaloneCode(): String = { + // Arrow says of every value whether it is there, and a numpy column has + // nowhere to put that: a missing double and a stored NaN both land on NaN. + // The nullable dtypes keep a holed integer integral too, so a long past 2^53 + // is not rounded on its way through a float. Named here because the read + // below cannot ask pd.read_feather for them, and named for the types + // ArrowUtils.toAttributeType reads; a timestamp needs none, datetime64 + // having a slot of its own for a missing value. + val dtypes = + """|_nullable = { + | pa.bool_(): pd.BooleanDtype(), + | pa.string(): pd.StringDtype(), + | pa.float32(): pd.Float32Dtype(), + | pa.float64(): pd.Float64Dtype(), + | pa.int8(): pd.Int8Dtype(), + | pa.int16(): pd.Int16Dtype(), + | pa.int32(): pd.Int32Dtype(), + | pa.int64(): pd.Int64Dtype(), + | pa.uint8(): pd.UInt8Dtype(), + | pa.uint16(): pd.UInt16Dtype(), + | pa.uint32(): pd.UInt32Dtype(), + |}""".stripMargin + // A file pandas wrote notes in its schema what the frame it came from looked + // like, and pandas reads that note back: a column the frame was keyed by + // returns as the index, and numbered labels return as numbers where the file + // says "1" and "2". The executor reads the columns the file states, so the + // note is refused rather than undone, as the Parquet source refuses it. + val read = + s"""|with pa.ipc.open_file($SourceFilePlaceholder) as _file: + | _table = _file.read_all() + |out1df = _table.to_pandas(ignore_metadata=True, types_mapper=_nullable.get)""".stripMargin + // The widths pandas keeps and Texera has no column for. A file states the + // width and the sign of each of its numbers, and pandas reads every one of + // them back, where a Texera column is a double or a 32-bit integer and + // nothing narrower. Left alone, a single-precision column summed to a + // different number on the two sides: 16777216 and 1 add to 16777217 as + // doubles and to 16777216 as floats. An unsigned column lands on the Texera + // type its own values need, which is the executor's rule too: one counting + // to 4294967295 has to be a long. See ParquetScanSourceOpDesc, which + // normalizes the same widths for the same reason. + val widths = + """|for _column, _values in out1df.items(): + | if _values.dtype == "Float32": + | out1df[_column] = _values.astype("Float64") + | elif _values.dtype in ("Int8", "Int16", "UInt8", "UInt16"): + | out1df[_column] = _values.astype("Int32") + | elif _values.dtype == "UInt32": + | out1df[_column] = _values.astype("Int64")""".stripMargin + // A timestamp column may name a zone, and pandas keeps it on the column + // where the executor keeps only the wall clock in that zone: a Texera + // TIMESTAMP has none. Left zoned, the column reached a downstream `astype` + // that refuses to drop a zone and ended the script. pandas already holds the + // wall clock in the file's own zone, so the zone is taken off and the clock + // left as it is, which is what ArrowUtils reads. CSV and JSONL parse their + // timestamps from text, carrying no types to go on, but Arrow states its own. + val zones = + """|for _column, _values in out1df.items(): + | if isinstance(_values.dtype, pd.DatetimeTZDtype): + | out1df[_column] = _values.dt.tz_localize(None)""".stripMargin + + // The executor drops `offset` rows and then takes `limit` of them. Feather has + // no row-range read, so the same window is taken once the frame is in memory. + // + // Clamped first: the property editor refuses a negative, but a plan posted to + // the API can still carry one, and `iloc` reads it from the end where `drop` + // skips nothing and `take` keeps nothing. + val window = (offset.map(_.max(0)), limit.map(_.max(0))) match { + // The end of the window is counted in Long: two Ints the operator accepts + // can add up past what an Int holds, and the slice would come out negative + // and take the wrong rows. As in ParquetScanSourceOpDesc. + case (Some(o), Some(l)) => Some(s"$o:${o.toLong + l}") + case (Some(o), None) => Some(s"$o:") + case (None, Some(l)) => Some(s":$l") + case _ => None + } + + (Seq(dtypes, read, widths, zones) ++ window.map(w => + s"out1df = out1df.iloc[$w].reset_index(drop=True)" + )).mkString("\n") + } + @throws[IOException] override def getPhysicalOp( workflowId: WorkflowIdentity, diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDesc.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDesc.scala index faccc76f882..949af677d74 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDesc.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDesc.scala @@ -24,18 +24,22 @@ import com.fasterxml.jackson.databind.JsonNode import org.apache.texera.amber.core.executor.OpExecWithClassName import org.apache.texera.amber.core.storage.DocumentFactory import org.apache.texera.amber.core.tuple.AttributeTypeUtils.inferSchemaFromRows -import org.apache.texera.amber.core.tuple.{Attribute, Schema} +import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema} import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} import org.apache.texera.amber.core.workflow.{PhysicalOp, SchemaPropagationFunc} +import org.apache.texera.amber.operator.{StandaloneCodeGenerator, StandaloneHelpers} +import org.apache.texera.amber.operator.StandaloneCodeGenerator.SourceFilePlaceholder import org.apache.texera.amber.operator.source.scan.ScanSourceOpDesc +import org.apache.texera.amber.pybuilder.PythonTemplateBuilder.pyStringLiteral import org.apache.texera.amber.util.JSONUtils.{JSONToMap, objectMapper} import java.io._ import java.net.URI import scala.collection.mutable.ArrayBuffer +import scala.util.Try import scala.jdk.CollectionConverters.IteratorHasAsScala -class JSONLScanSourceOpDesc extends ScanSourceOpDesc { +class JSONLScanSourceOpDesc extends ScanSourceOpDesc with StandaloneCodeGenerator { @JsonProperty(required = true, defaultValue = "false") @JsonPropertyDescription("flatten nested objects and arrays") @@ -43,6 +47,176 @@ class JSONLScanSourceOpDesc extends ScanSourceOpDesc { fileTypeName = Option("JSONL") + override def standaloneSourcePath(): Option[String] = fileName + + override def generateStandaloneCode(): String = { + val enc = fileEncoding.toString.replace("_", "-").toLowerCase + + // The executor drops and takes on the RAW lines, before any of them is + // parsed, so a line outside the window is never read as JSON and an + // unparseable one there costs nothing. Reading the whole file first and + // slicing the frame would end the export on a line the workflow skipped. + val windowed = offset.exists(_ > 0) || limit.isDefined + + // read_json parses a JSON number into a float before any dtype it is handed + // can apply, so a LONG past 2^53 is already rounded by the time the column + // exists: 9007199254740993 arrives as ...992, a value the file never held. + // Reading the lines a second time with Python's own parser, which keeps an + // integer exact, is the only way to put the column back. + val longColumns: Seq[String] = + Try(sourceSchema()).toOption.toSeq.flatMap( + _.getAttributes.filter(_.getType == AttributeType.LONG).map(_.getName) + ) + + // The executor reads every value as the text Jackson gives it, so a column + // this operator types STRING holds that text whatever the JSON held: 35.0 is + // "35.0", true is "true", and a null is "null", which is also why one null + // makes its column STRING. read_json keeps the values typed and the null a + // null. These columns are rebuilt from the lines the way LONG ones are. + val stringColumns: Seq[String] = + Try(sourceSchema()).toOption.toSeq.flatMap( + _.getAttributes.filter(_.getType == AttributeType.STRING).map(_.getName) + ) + + // The lines are read into the script whenever something below needs them, + // and read_json is then fed from those rather than from the file, so the + // file is opened once either way. + val readsLines = windowed || longColumns.nonEmpty || stringColumns.nonEmpty + val dropped = offset.filter(_ > 0).fold("_lines")(o => s"_lines[$o:]") + val taken = limit.fold(dropped)(l => s"$dropped[:${l.max(0)}]") + + val readArgs = scala.collection.mutable.ArrayBuffer[String]() + readArgs += (if (readsLines) s"""io.StringIO("".join($taken))""" else SourceFilePlaceholder) + readArgs += "lines=True" + // Text already decoded by the open above carries no encoding of its own. + if (!readsLines) readArgs += s"""encoding=${pyStringLiteral(enc)}""" + + // JSON has no timestamp of its own, so both readers infer from the text and + // do not infer alike: the schema below tries TIMESTAMP and parses what it + // can, while pd.read_json guesses from the COLUMN NAME (anything ending + // "_at" or "_time", anything called "date"). Every column is read as it + // was written, and the ones this operator decided were timestamps are + // parsed below, the way it parses them, once the frame holds them: under + // flattening a nested value is no column until the flattening builds it. + readArgs += "convert_dates=False" + + val readExpr = s"pd.read_json(${readArgs.mkString(", ")})" + // json_normalize opens a nested object and leaves a nested array whole, so + // an array arrived as one column holding a list where the executor had + // already given each element a column of its own. The flattening the + // executor does is done here instead, and json_normalize is left the frame + // to build out of records that are flat by then. + val baseExpr = + if (flatten) + s"pd.json_normalize([_texera_json_flatten(_r) for _r in $readExpr.to_dict('records')])" + else readExpr + + val lines = scala.collection.mutable.ArrayBuffer[String]() + if (readsLines) { + lines += s"""with open($SourceFilePlaceholder, "r", encoding=${pyStringLiteral( + enc + )}) as _f:""" + lines += " _lines = _f.readlines()" + } + lines += s"out1df = $baseExpr" + Try(sourceSchema()).toOption.foreach { schema => + lines += StandaloneCodeGenerator.typeAnEmptyRead("out1df", schema) + } + + if (longColumns.nonEmpty || stringColumns.nonEmpty) { + // Flattened or not, a column's name is a key of the record the lookup + // below reads, because a flattened record is flattened first. + val parsed = if (flatten) "_texera_json_flatten(json.loads(_l))" else "json.loads(_l)" + lines += s"_records = [$parsed for _l in $taken]" + longColumns.foreach { name => + val nameLit = pyStringLiteral(name) + lines += s"""out1df[$nameLit] = pd.array([_r.get($nameLit) for _r in _records], dtype="Int64")""" + } + // A missing key is the executor's null. A key holding null is its "null". + stringColumns.foreach { name => + val nameLit = pyStringLiteral(name) + lines += s"""out1df[$nameLit] = pd.Series([_texera_json_text(_r[$nameLit]) if $nameLit in _r else None for _r in _records], index=out1df.index, dtype="object")""" + } + } + + Try(sourceSchema()).toOption + .map(StandaloneCodeGenerator.parseTimestamps("out1df", _)) + .filter(_.nonEmpty) + .foreach(lines += _) + + // pandas has no plain boolean that carries a hole, so a record missing the + // key widens the column to floats, and a later cast to text read 1.0 and 0.0 + // where the executor has true and false. The nullable boolean carries both + // values and the hole. A column with no hole arrives as bool already and is + // left alone. + val booleanColumns: Seq[String] = + Try(sourceSchema()).toOption.toSeq.flatMap( + _.getAttributes + .filter(_.getType == AttributeType.BOOLEAN) + .map(a => pyStringLiteral(a.getName)) + ) + booleanColumns.foreach { nameLit => + lines += s"""if out1df[$nameLit].dtype == "float64":""" + lines += s""" out1df[$nameLit] = out1df[$nameLit].astype("boolean")""" + } + + // read_json infers each column's type from its values, which is not how this + // operator types it: a DOUBLE whose values are all whole, 3.0 among them, is + // read as int64, and an INTEGER missing a key is widened to float. Both reach + // the next operator as the other type, where 3 and 3.0 print differently. The + // schema's type is asked for, the INTEGER in the nullable width that keeps a + // hole. + Try(sourceSchema()).toOption.toSeq.flatMap(_.getAttributes).foreach { attribute => + val nameLit = pyStringLiteral(attribute.getName) + attribute.getType match { + case AttributeType.DOUBLE => + lines += s"""out1df[$nameLit] = out1df[$nameLit].astype("float64")""" + case AttributeType.INTEGER => + lines += s"""out1df[$nameLit] = out1df[$nameLit].astype("Int32")""" + case _ => + } + } + + // A JSONL file states no column order, so the schema this operator infers + // sorts the names it found and the rows the workflow sees follow that + // order. read_json keeps the order the first record happened to use, which + // is the same columns in a different order, and column order is what a + // positional read downstream and a file export both go by. + // + // Unflattened, the executor also passes over a nested object or array, so + // the key is no column of its own, where read_json keeps it as one holding + // the dict or the list. The schema names the columns the executor has. + val columns = + Try(sourceSchema()).toOption.map(_.getAttributes.map(a => pyStringLiteral(a.getName))) + lines += columns.fold("out1df = out1df[sorted(out1df.columns)]")(names => + s"out1df = out1df[[${names.mkString(", ")}]]" + ) + + lines.mkString("\n") + } + + override def standaloneHelpers(): Seq[String] = + (if (flatten) Seq(JSONLScanSourceOpDesc.JsonFlatten) else Seq.empty) ++ + (if (columnCount(AttributeType.STRING) > 0) Seq(JSONLScanSourceOpDesc.JsonText) + else Seq.empty) ++ + (if (columnCount(AttributeType.TIMESTAMP) > 0 || columnCount(AttributeType.STRING) > 0) + Seq(StandaloneHelpers.AttributeCasts) + else Seq.empty) + + override def standaloneImports(): Seq[String] = { + val windowed = offset.exists(_ > 0) || limit.isDefined + val longs = columnCount(AttributeType.LONG) + val strings = columnCount(AttributeType.STRING) + (if (windowed || longs > 0 || strings > 0) Seq("import io") else Seq.empty) ++ + (if (longs > 0 || strings > 0) Seq("import json") else Seq.empty) + } + + /** How many columns the schema gives this type, or none when it cannot be read. */ + private def columnCount(attributeType: AttributeType): Int = + Try(sourceSchema()).toOption + .map(_.getAttributes.count(_.getType == attributeType)) + .getOrElse(0) + @throws[IOException] override def getPhysicalOp( workflowId: WorkflowIdentity, @@ -87,8 +261,12 @@ class JSONLScanSourceOpDesc extends ScanSourceOpDesc { val allFields: ArrayBuffer[Map[String, String]] = ArrayBuffer() val startOffset = offset.getOrElse(0) + // A window of no rows is still a window on this file, and the file's columns + // do not depend on how many of its rows were asked for. Reading the sample + // through the limit left a Limit of 0 nothing to infer from, and the operator + // declared a schema of no columns at all. val endOffset = - startOffset + limit.getOrElse(INFER_READ_LIMIT).min(INFER_READ_LIMIT) + startOffset + limit.filter(_ > 0).getOrElse(INFER_READ_LIMIT).min(INFER_READ_LIMIT) reader .lines() .iterator() @@ -124,3 +302,52 @@ class JSONLScanSourceOpDesc extends ScanSourceOpDesc { } } + +object JSONLScanSourceOpDesc { + + /** + * One record flattened the way `JSONUtils.JSONToMap` flattens it, so the + * names the two paths give a nested value are the same names. + * + * A nested object joins its key to its parent's with a dot, and an array + * element takes its parent's name followed by its position counted from one: + * `{"items": [{"id": 1}, {"id": 2}]}` is `items1.id` and `items2.id` on both + * sides. A value nested no deeper keeps its own key. + */ + val JsonFlatten: String = + """def _texera_json_flatten(record): + | flat = {} + | stack = [(record, "")] + | while stack: + | node, parent = stack.pop() + | if isinstance(node, dict): + | for key, child in node.items(): + | path = parent + "." + key if parent else key + | if isinstance(child, (dict, list)): + | stack.append((child, path)) + | else: + | flat[path] = child + | elif isinstance(node, list): + | for index, child in enumerate(node): + | stack.append((child, parent + str(index + 1))) + | elif parent: + | flat[parent] = node + | return flat""".stripMargin + + /** + * One JSON value as the text Jackson's `asText` gives it, which is what the + * executor reads every value as: `null` for a null, `true` and `false` in + * lower case, and a number the way Java prints it, with `_texera_java_text` + * from [[StandaloneHelpers.AttributeCasts]]. An object or an array is skipped + * by the executor, so it is no value here either. + */ + val JsonText: String = + """def _texera_json_text(value): + | if value is None: + | return "null" + | if isinstance(value, (bool, int, float)): + | return _texera_java_text(value) + | if isinstance(value, str): + | return value + | return None""".stripMargin +} diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/arrow/ArrowSourceOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/arrow/ArrowSourceOpDescSpec.scala index 54dd677913e..71ae99280dd 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/arrow/ArrowSourceOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/arrow/ArrowSourceOpDescSpec.scala @@ -19,34 +19,52 @@ package org.apache.texera.amber.operator.source.scan.arrow +import com.typesafe.config.ConfigFactory import org.apache.arrow.memory.RootAllocator -import org.apache.arrow.vector.VectorSchemaRoot +import org.apache.arrow.vector.{UInt8Vector, VectorSchemaRoot} import org.apache.arrow.vector.ipc.ArrowFileWriter +import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema => ArrowFileSchema} 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.tags.IntegrationTest +import org.apache.texera.amber.operator.{LogicalOp, StandaloneCodeGenerator} import org.apache.texera.amber.operator.metadata.OperatorGroupConstants import org.apache.texera.amber.operator.source.scan.FileDecodingMethod import org.apache.texera.amber.util.ArrowUtils 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.io.{File, FileOutputStream} import java.nio.channels.Channels +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 ArrowSourceOpDescSpec extends AnyFlatSpec with Matchers { + private val NeedsPythonPackages = Tag(classOf[IntegrationTest].getName) + private val workflowId = WorkflowIdentity(1L) private val executionId = ExecutionIdentity(1L) - private def writeArrowFile(schema: Schema, rows: Seq[Array[Any]]): File = { + private def writeArrowFile( + schema: Schema, + rows: Seq[Array[Any]], + schemaMetadata: Map[String, String] = Map.empty + ): File = { val file = File.createTempFile("arrow-src-", ".arrow") file.deleteOnExit() val allocator = new RootAllocator() - val root = VectorSchemaRoot.create(ArrowUtils.fromTexeraSchema(schema), allocator) + val root = VectorSchemaRoot.create( + new ArrowFileSchema(ArrowUtils.fromTexeraSchema(schema).getFields, schemaMetadata.asJava), + allocator + ) val out = new FileOutputStream(file) val writer = new ArrowFileWriter(root, null, Channels.newChannel(out)) try { @@ -153,6 +171,61 @@ class ArrowSourceOpDescSpec extends AnyFlatSpec with Matchers { d.inferSchema() shouldBe schema } + // A missing double and a stored NaN are two different values to the engine, + // and a numpy column has one slot for both. A holed integer column loses its + // type the same way. + "ArrowSourceOpDesc.generateStandaloneCode" should + "keep a missing value apart from a NaN, and an integer integral" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val schema = Schema( + List(new Attribute("d", AttributeType.DOUBLE), new Attribute("i", AttributeType.INTEGER)) + ) + val file = writeArrowFile( + schema, + Seq( + Array[Any](null, null), + Array[Any](Double.box(Double.NaN), Int.box(7)) + ) + ) + + val d = new ArrowSourceOpDesc + d.fileName = Some(file.toURI.toString) + // The block reads the file from the script's own directory, so run there. + val workDir = Files.createTempDirectory("arrow-standalone-") + workDir.toFile.deleteOnExit() + val beside = workDir.resolve(file.getName) + Files.copy(file.toPath, beside) + + val script = workDir.resolve("run.py") + Files.write( + script, + s"""${scriptHeader(d)} + |${d.generateStandaloneCode()} + |print(str(out1df["d"].dtype), str(out1df["i"].dtype)) + |print(repr(out1df["d"].iloc[0]), repr(out1df["d"].iloc[1])) + |""".stripMargin.getBytes(StandardCharsets.UTF_8) + ) + + val process = new ProcessBuilder(python, script.toString) + .directory(workDir.toFile) + .redirectErrorStream(true) + .start() + val out = Source.fromInputStream(process.getInputStream).mkString + process.waitFor(120, TimeUnit.SECONDS) + + withClue(s"python said:\n$out") { + process.exitValue() shouldBe 0 + val lines = out.trim.linesIterator.toSeq + lines.head shouldBe "Float64 Int32" + // The first is the value that was missing, the second the NaN that was + // stored. Under numpy dtypes both printed as nan. + lines(1) should include("") + lines(1) should include("nan") + } + } + it should "throw a friendly error when the file is not a valid Arrow file" in { val bogus = File.createTempFile("not-arrow-", ".arrow") bogus.deleteOnExit() @@ -162,4 +235,131 @@ class ArrowSourceOpDescSpec extends AnyFlatSpec with Matchers { val ex = intercept[RuntimeException](d.inferSchema()) ex.getMessage shouldBe "Failed to read the .arrow file. Please ensure it is a valid Arrow file." } + + // pandas notes in an Arrow file's schema what type the frame's column labels + // had, and casts the names back to it. A frame whose columns were numbered is written + // under the names "1" and "2", which is what the executor reads, and comes + // back out of pandas labelled 1 and 2. + it should "keep the names a pandas frame of numbered columns was written under" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + val schema = + Schema( + List(new Attribute("1", AttributeType.INTEGER), new Attribute("2", AttributeType.INTEGER)) + ) + // What `pd.DataFrame({1: [7, 8], 2: [10, 20]}).to_feather(...)` leaves: the + // names are strings, and `column_indexes` is the whole of what turns them + // back into numbers. + val pandasMetadata = + """|{"index_columns": [{"kind": "range", "name": null, "start": 0, "stop": 2, "step": 1}], + | "column_indexes": [{"name": null, "field_name": null, + | "pandas_type": "int64", "numpy_type": "int64", "metadata": null}], + | "columns": [{"name": "1", "field_name": "1", + | "pandas_type": "int32", "numpy_type": "int32", "metadata": null}, + | {"name": "2", "field_name": "2", + | "pandas_type": "int32", "numpy_type": "int32", "metadata": null}], + | "pandas_version": "2.2.3"}""".stripMargin + val file = writeArrowFile( + schema, + Seq(Array[Any](7, 10), Array[Any](8, 20)), + Map("pandas" -> pandasMetadata) + ) + + val d = new ArrowSourceOpDesc + d.fileName = Some(file.toURI.toString) + d.inferSchema().getAttributeNames.toList shouldBe List("1", "2") + + val workDir = Files.createTempDirectory("arrow-labels-") + workDir.toFile.deleteOnExit() + Files.copy(file.toPath, workDir.resolve(file.getName)) + + val script = workDir.resolve("run.py") + Files.write( + script, + s"""${scriptHeader(d)} + |${d.generateStandaloneCode()} + |print(list(out1df.columns)) + |print([int(v) for v in out1df["2"]]) + |""".stripMargin.getBytes(StandardCharsets.UTF_8) + ) + + val process = new ProcessBuilder(python, script.toString) + .directory(workDir.toFile) + .redirectErrorStream(true) + .start() + val out = Source.fromInputStream(process.getInputStream).mkString + process.waitFor(120, TimeUnit.SECONDS) + + withClue(s"python said:\n$out") { + process.exitValue() shouldBe 0 + val lines = out.trim.linesIterator.toSeq + lines.head shouldBe "['1', '2']" + lines(1) shouldBe "[10, 20]" + } + } + + // Nothing in Texera holds a value past 2^63 - 1, so the column is refused at + // the schema rather than read as the -1 its storage counts down to. + it should "refuse an unsigned 64-bit column, having no Texera type for it" in { + val fields = List(new Field("u", FieldType.nullable(new ArrowType.Int(64, false)), null)) + val file = File.createTempFile("arrow-u64-", ".arrow") + file.deleteOnExit() + val allocator = new RootAllocator() + val root = VectorSchemaRoot.create(new ArrowFileSchema(fields.asJava), allocator) + val out = new FileOutputStream(file) + val writer = new ArrowFileWriter(root, null, Channels.newChannel(out)) + try { + writer.start() + root.allocateNew() + root.getVector("u").asInstanceOf[UInt8Vector].setSafe(0, -1L) + root.setRowCount(1) + writer.writeBatch() + writer.end() + } finally { + writer.close() + root.close() + allocator.close() + out.close() + } + + val d = new ArrowSourceOpDesc + d.fileName = Some(file.toURI.toString) + val ex = intercept[RuntimeException](d.inferSchema()) + ex.getMessage shouldBe "Failed to read the .arrow file. Please ensure it is a valid Arrow file." + } + + /** The block names its file by placeholder; the translator puts a name there. */ + private def bindSourceFile(d: ArrowSourceOpDesc): String = + s"""${StandaloneCodeGenerator.SourceFilePlaceholder} = "${d.standaloneSourceName().get}"""" + + /** What a script needs before the block: its imports, and the file it reads. */ + private def scriptHeader(d: ArrowSourceOpDesc): String = + (d.standaloneImports() :+ "import pandas as pd" :+ bindSourceFile(d)).mkString("\n") + + 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, pyarrow").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/source/scan/json/JSONLScanSourceOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDescSpec.scala index c0d8cf2e93c..6e038736dc3 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDescSpec.scala @@ -20,6 +20,8 @@ package org.apache.texera.amber.operator.source.scan.json import org.apache.texera.amber.core.executor.OpExecWithClassName +import org.apache.texera.amber.core.storage.FileResolver +import org.apache.texera.amber.core.tuple.AttributeType import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} import org.apache.texera.amber.operator.LogicalOp import org.apache.texera.amber.operator.metadata.OperatorGroupConstants @@ -28,6 +30,9 @@ import org.apache.texera.amber.util.JSONUtils.objectMapper import org.scalatest.flatspec.AnyFlatSpec import org.scalatest.matchers.should.Matchers +import java.nio.charset.StandardCharsets +import java.nio.file.Files + class JSONLScanSourceOpDescSpec extends AnyFlatSpec with Matchers { private val workflowId = WorkflowIdentity(1L) @@ -86,4 +91,28 @@ class JSONLScanSourceOpDescSpec extends AnyFlatSpec with Matchers { r.limit shouldBe Some(10) r.offset shouldBe Some(5) } + + // The limit bounded the sample the inference reads as well as the rows the + // operator emits, so a Limit of 0 had nothing to infer from and the operator + // declared a schema of no columns at all. A file's columns do not depend on + // how many of its rows were asked for. + "JSONLScanSourceOpDesc.sourceSchema" should "keep the file's columns when the window asks for no rows" in { + val data = Files.createTempFile("jsonl-zero-window-", ".jsonl") + data.toFile.deleteOnExit() + Files.write( + data, + "{\"id\":1,\"name\":\"alice\"}\n{\"id\":2,\"name\":\"bob\"}\n".getBytes( + StandardCharsets.UTF_8 + ) + ) + + val op = new JSONLScanSourceOpDesc + op.fileName = Some(data.toString) + op.setResolvedFileName(FileResolver.resolve(data.toString)) + op.limit = Some(0) + + val schema = op.sourceSchema() + schema.getAttributeNames shouldBe List("id", "name") + schema.getAttribute("id").getType shouldBe AttributeType.INTEGER + } } diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala index 94abd91b8ae..665dc4fc488 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala @@ -42,7 +42,7 @@ class ArrowUtilsSpec extends AnyFlatSpec { val float = new ArrowType.FloatingPoint(FloatingPointPrecision.SINGLE) val double = new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE) val half = new ArrowType.FloatingPoint(FloatingPointPrecision.HALF) - val timestamp = new ArrowType.Timestamp(TimeUnit.MILLISECOND, "UTC") + val timestamp = new ArrowType.Timestamp(TimeUnit.MICROSECOND, null) val string: ArrowType.Utf8 = ArrowType.Utf8.INSTANCE val texeraSchema: Schema = Schema() @@ -70,10 +70,13 @@ class ArrowUtilsSpec extends AnyFlatSpec { it should "convert to AttributeTypes correctly" in { assert(ArrowUtils.toAttributeType(unsignedShortInt) == AttributeType.INTEGER) assert(ArrowUtils.toAttributeType(signedShortInt) == AttributeType.INTEGER) - assert(ArrowUtils.toAttributeType(unsignedInt) == AttributeType.INTEGER) assert(ArrowUtils.toAttributeType(signedInt) == AttributeType.INTEGER) - assert(ArrowUtils.toAttributeType(unsignedLongInt) == AttributeType.LONG) assert(ArrowUtils.toAttributeType(signedLongInt) == AttributeType.LONG) + // An unsigned column needs the width above its own: read at its own width it + // would hand back the storage, and the largest unsigned 32-bit value is + // stored as -1. Past 64 bits there is no width above to read it as. + assert(ArrowUtils.toAttributeType(unsignedInt) == AttributeType.LONG) + assertThrows[AttributeTypeException](ArrowUtils.toAttributeType(unsignedLongInt)) assert(ArrowUtils.toAttributeType(boolean) == AttributeType.BOOLEAN) @@ -173,10 +176,12 @@ class ArrowUtilsSpec extends AnyFlatSpec { assert(vectorSchemaRoot.getVector(2).getObject(index).asInstanceOf[Boolean] == true) assert(vectorSchemaRoot.getVector(3).getObject(index).asInstanceOf[Double] == 1.1) - // The arrow storage type of timestamp is Long, and the field is labelled - // UTC, so the wall clock above is stored as the UTC instant of the same - // reading: ten seconds past the epoch, on a server anywhere. - assert(vectorSchemaRoot.getVector(4).getObject(index).asInstanceOf[Long] == 10000L) + // The field carries no zone, so the vector holds the wall clock above and + // hands it back as one: ten seconds past the epoch, on a server anywhere. + assert( + vectorSchemaRoot.getVector(4).getObject(index) == + java.time.LocalDateTime.of(1970, 1, 1, 0, 0, 10) + ) // the arrow storage type of string is Text assert( From cf5695a4bd56e6af9b9e3bd988ba6c012ca21143 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Sat, 26 Sep 2026 21:28:17 -0700 Subject: [PATCH 2/3] fix(workflow-core): read an Arrow field only for the columns an unsigned integer can arrive in Co-Authored-By: Claude Opus 5.5 (1M context) --- .../scala/org/apache/texera/amber/util/ArrowUtils.scala | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala index 9b6b0594d91..5678948bc0d 100644 --- a/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala +++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/util/ArrowUtils.scala @@ -85,15 +85,16 @@ object ArrowUtils extends LazyLogging { // instead of deriving it from the Arrow type val attributeType = schema.getAttributes(index).getType // A timestamp and an unsigned integer are the types whose field - // says more than the schema does, so they are the ones that read - // the field. + // says more than the schema does, so they are the only ones that + // read the field: an unsigned column arrives as INTEGER or LONG. attributeType match { case AttributeType.TIMESTAMP => wallClockOf(value, fieldVector.getField.getType) - case _ => + case AttributeType.INTEGER | AttributeType.LONG => AttributeTypeUtils.parseField( unsignedValueOf(value, fieldVector.getField.getType), attributeType ) + case _ => AttributeTypeUtils.parseField(value, attributeType) } } catch { case e: Exception => From 039d5f9f229b5bc76742d4c011ba05dcb6bc26b3 Mon Sep 17 00:00:00 2001 From: kary zheng Date: Tue, 29 Sep 2026 13:33:20 -0700 Subject: [PATCH 3/3] fix(operator): read a JSONL double exactly as the file spells it read_json's default float parser rounds 1.0000000000000002 to 1.0 where Jackson keeps it, so a filter on v > 1 kept different rows in the two paths. precise_float=True reads the value exactly. Co-Authored-By: Claude Opus 5.5 --- .../scan/json/JSONLScanSourceOpDesc.scala | 4 + .../scan/json/JSONLScanSourceOpDescSpec.scala | 88 ++++++++++++++++++- 2 files changed, 91 insertions(+), 1 deletion(-) diff --git a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDesc.scala b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDesc.scala index 949af677d74..2e3e420aa30 100644 --- a/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDesc.scala +++ b/common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDesc.scala @@ -99,6 +99,10 @@ class JSONLScanSourceOpDesc extends ScanSourceOpDesc with StandaloneCodeGenerato // parsed below, the way it parses them, once the frame holds them: under // flattening a nested value is no column until the flattening builds it. readArgs += "convert_dates=False" + // read_json's default float parser can land one ulp off the value the text + // spells, where Jackson's parse is exact: 1.0000000000000002 reads as 1.0, + // and a filter on v > 1 then drops a row the workflow keeps. + readArgs += "precise_float=True" val readExpr = s"pd.read_json(${readArgs.mkString(", ")})" // json_normalize opens a nested object and leaves a nested array whole, so diff --git a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDescSpec.scala b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDescSpec.scala index 6e038736dc3..6a1a5fff37a 100644 --- a/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDescSpec.scala +++ b/common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/source/scan/json/JSONLScanSourceOpDescSpec.scala @@ -19,22 +19,30 @@ package org.apache.texera.amber.operator.source.scan.json +import com.typesafe.config.ConfigFactory import org.apache.texera.amber.core.executor.OpExecWithClassName import org.apache.texera.amber.core.storage.FileResolver import org.apache.texera.amber.core.tuple.AttributeType import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity} -import org.apache.texera.amber.operator.LogicalOp +import org.apache.texera.amber.operator.{LogicalOp, StandaloneCodeGenerator} import org.apache.texera.amber.operator.metadata.OperatorGroupConstants import org.apache.texera.amber.operator.source.scan.FileDecodingMethod +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 JSONLScanSourceOpDescSpec extends AnyFlatSpec with Matchers { + private val NeedsPythonPackages = Tag(classOf[IntegrationTest].getName) + private val workflowId = WorkflowIdentity(1L) private val executionId = ExecutionIdentity(1L) @@ -115,4 +123,82 @@ class JSONLScanSourceOpDescSpec extends AnyFlatSpec with Matchers { schema.getAttributeNames shouldBe List("id", "name") schema.getAttribute("id").getType shouldBe AttributeType.INTEGER } + + // Jackson parses 1.0000000000000002 exactly; read_json's default parser + // rounds it to 1.0, so a filter on v > 1 would keep a different set of rows. + "JSONLScanSourceOpDesc.generateStandaloneCode" should + "read a double as exactly the value the file spells" taggedAs NeedsPythonPackages in { + val python = resolvePython().getOrElse(cancel("No runnable python executable")) + if (!canImportPandas(python)) cancel(s"'$python' cannot import pandas") + + // The block reads the file from the script's own directory, so run there. + val workDir = Files.createTempDirectory("jsonl-standalone-") + workDir.toFile.deleteOnExit() + val data = workDir.resolve("data.jsonl") + Files.write( + data, + "{\"v\": 1.0000000000000002}\n{\"v\": 1.0}\n".getBytes(StandardCharsets.UTF_8) + ) + + val op = new JSONLScanSourceOpDesc + op.fileName = Some(data.toString) + op.setResolvedFileName(FileResolver.resolve(data.toString)) + op.sourceSchema().getAttribute("v").getType shouldBe AttributeType.DOUBLE + + val script = workDir.resolve("run.py") + Files.write( + script, + s"""${scriptHeader(op)} + |${op.generateStandaloneCode()} + |print(repr(out1df["v"].tolist())) + |""".stripMargin.getBytes(StandardCharsets.UTF_8) + ) + + val process = new ProcessBuilder(python, script.toString) + .directory(workDir.toFile) + .redirectErrorStream(true) + .start() + val out = Source.fromInputStream(process.getInputStream).mkString + process.waitFor(120, TimeUnit.SECONDS) + + withClue(s"python said:\n$out") { + process.exitValue() shouldBe 0 + out.trim.linesIterator.toSeq.last shouldBe "[1.0000000000000002, 1.0]" + } + } + + /** The block names its file by placeholder; the translator puts a name there. */ + private def bindSourceFile(d: JSONLScanSourceOpDesc): String = + s"""${StandaloneCodeGenerator.SourceFilePlaceholder} = "${d.standaloneSourceName().get}"""" + + /** What a script needs before the block: its imports, and the file it reads. */ + private def scriptHeader(d: JSONLScanSourceOpDesc): String = + (d.standaloneImports() :+ "import pandas as pd" :+ bindSourceFile(d)).mkString("\n") + + 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").redirectErrorStream(true).start() + ).toOption + .exists { p => + if (!p.waitFor(60, TimeUnit.SECONDS)) { p.destroyForcibly(); false } + else p.exitValue() == 0 + } }