feat(operator): read a Parquet file as a source

common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/LogicalOp.scala

@@ -79,6 +79,7 @@ import org.apache.texera.amber.operator.source.apis.twitter.v2.{
 import org.apache.texera.amber.operator.source.dataset.FileListerSourceOpDesc
 import org.apache.texera.amber.operator.source.fetcher.URLFetcherOpDesc
 import org.apache.texera.amber.operator.source.scan.arrow.ArrowSourceOpDesc
+import org.apache.texera.amber.operator.source.scan.parquet.ParquetScanSourceOpDesc
 import org.apache.texera.amber.operator.source.scan.csv.CSVScanSourceOpDesc
 import org.apache.texera.amber.operator.source.scan.csvOld.CSVOldScanSourceOpDesc
 import org.apache.texera.amber.operator.source.scan.json.JSONLScanSourceOpDesc
@@ -273,6 +274,7 @@ trait StateTransferFunc
     new Type(value = classOf[RUDFOpDesc], name = "RUDF"),
     new Type(value = classOf[RUDFSourceOpDesc], name = "RUDFSource"),
     new Type(value = classOf[ArrowSourceOpDesc], name = "ArrowSource"),
+    new Type(value = classOf[ParquetScanSourceOpDesc], name = "ParquetFileScan"),
     new Type(value = classOf[MachineLearningScorerOpDesc], name = "Scorer"),
     new Type(value = classOf[SortOpDesc], name = "Sort"),
     new Type(value = classOf[StableMergeSortOpDesc], name = "StableMergeSort"),

common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/parquet/ParquetScanSourceOpDesc.scala

@@ -0,0 +1,182 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.operator.source.scan.parquet
+
+import com.fasterxml.jackson.annotation.JsonIgnoreProperties
+import org.apache.parquet.hadoop.ParquetFileReader
+import org.apache.parquet.io.LocalInputFile
+import org.apache.texera.amber.core.executor.OpExecWithClassName
+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.JSONUtils.objectMapper
+
+import java.io.IOException
+import java.net.URI
+import scala.util.Using
+
+/**
+  * Reads a Parquet file. The format states its own column types in a footer, so
+  * unlike the CSV and JSONL sources this one infers nothing: the INTEGER the
+  * file was written with is the INTEGER that arrives.
+  */
+@JsonIgnoreProperties(value = Array("fileEncoding"))
+class ParquetScanSourceOpDesc extends ScanSourceOpDesc with StandaloneCodeGenerator {
+
+  fileTypeName = Option("Parquet")
+
+  // A DECIMAL is the one column pandas does not land on the same type as the
+  // executor: it fills that column with decimal.Decimal objects, which a script
+  // cannot then multiply by a float. pyarrow is what pandas reads a Parquet file
+  // with; the footer is asked for directly here to read it the executor's way.
+  override def standaloneImports(): Seq[String] =
+    Seq("from datetime import date", "from decimal import Decimal", "import pyarrow.parquet as pq")
+
+  override def standaloneSourcePath(): Option[String] = fileName
+
+  override def generateStandaloneCode(): String = {
+    // No timestamps to parse, and no dtype map. pandas reads the types out of
+    // the same footer the executor does, which is the whole point of the format;
+    // the text formats parse theirs because they carry nothing to read.
+    // Read into the nullable dtypes, as the Arrow source is. Parquet says of
+    // every value whether it is there, and a numpy column has nowhere to put
+    // that: pandas widens a holed integer column through a float, where every
+    // value past 2^53 is rounded and 9007199254740993 came back as ...992. The
+    // executor reads the exact long, and the declared column stays integral.
+    //
+    // The footer is handed back to the reader stripped of its metadata. A file
+    // pandas wrote records there which of its columns held the frame's index,
+    // and pandas restores those as an index rather than as columns: a file
+    // written from a frame keyed by `customer_id` reads back without that
+    // column, and a projection naming it raises. The executor reads the columns
+    // the footer states and knows nothing of an index, so it hands `customer_id`
+    // on like any other. Stripped, pandas reads every column the file holds,
+    // under the name the footer gives it and in the order it states them.
+    val read =
+      s"""|out1df = pd.read_parquet(
+          |    $SourceFilePlaceholder,
+          |    dtype_backend="numpy_nullable",
+          |    schema=pq.read_schema($SourceFilePlaceholder).remove_metadata(),
+          |)""".stripMargin
+    // The columns pandas does not land on the same value as the executor. A
+    // DECIMAL arrives as decimal.Decimal objects and a FLOAT keeps the single
+    // precision the executor widens; an unsigned column, a duration and a zoned
+    // timestamp come back as the Arrow types the writer left in the file's
+    // metadata, where the executor reads what the footer alone states.
+    //
+    // A DATE arrives as datetime.date, where the executor reads midnight of that
+    // day; milliseconds reach the year 9999. Only an object column, since a
+    // timestamp's values are dates too. An 8- or 16-bit integer is widened to
+    // the INTEGER the executor reads, or 120 + 10 wraps to -126.
+    //
+    // Named in the nullable spelling, `Float32` and not `float32`, and widened
+    // into the nullable dtype in turn, so a hole stays a hole rather than
+    // becoming the NaN a numpy column would have to write it as.
+    val columns =
+      """|for _column, _values in out1df.items():
+         |    _first = next(iter(_values.dropna()), None)
+         |    if isinstance(_first, Decimal):
+         |        out1df[_column] = _values.astype("Float64")
+         |    elif _values.dtype == object and isinstance(_first, date):
+         |        out1df[_column] = _values.astype("datetime64[ms]")
+         |    elif isinstance(_values.dtype, pd.DatetimeTZDtype):
+         |        out1df[_column] = _values.dt.tz_convert("UTC").dt.tz_localize(None)
+         |    elif _values.dtype.kind in "um":
+         |        out1df[_column] = _values.astype("Int64")
+         |    elif _values.dtype in ("Int8", "Int16"):
+         |        out1df[_column] = _values.astype("Int32")
+         |    elif _values.dtype == "Float32":
+         |        out1df[_column] = _values.astype("Float64")""".stripMargin
+    // The executor drops `offset` rows and then takes `limit` of them. Parquet
+    // can skip whole row groups but not an arbitrary row range, so the same
+    // window is taken once the frame is in memory, as the Arrow source does.
+    //
+    // Clamped first: a plan from the API can carry a negative, which `iloc`
+    // counts from the end.
+    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.
+      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(read, columns) ++ window.map(w => s"out1df = out1df.iloc[$w].reset_index(drop=True)"))
+      .mkString("\n")
+  }
+
+  @throws[IOException]
+  override def getPhysicalOp(
+      workflowId: WorkflowIdentity,
+      executionId: ExecutionIdentity
+  ): PhysicalOp = {
+    PhysicalOp
+      .sourcePhysicalOp(
+        workflowId,
+        executionId,
+        operatorIdentifier,
+        OpExecWithClassName(
+          "org.apache.texera.amber.operator.source.scan.parquet.ParquetScanSourceOpExec",
+          objectMapper.writeValueAsString(this)
+        )
+      )
+      .withInputPorts(operatorInfo.inputPorts)
+      .withOutputPorts(operatorInfo.outputPorts)
+      .withPropagateSchema(
+        SchemaPropagationFunc(_ => Map(operatorInfo.outputPorts.head.id -> inferSchema()))
+      )
+  }
+
+  /** The file's own schema, read from its footer. No rows are read to get it. */
+  @Override
+  def inferSchema(): Schema = {
+    require(
+      fileResolved(),
+      "No file selected. Please select a valid .parquet file from the 'File' dropdown in the right panel."
+    )
+
+    val uri = new URI(fileName.get)
+    if (uri.getScheme == "file") {
+      require(
+        new java.io.File(uri).isFile,
+        "The selected item is a folder or does not exist. Please select an actual .parquet file from the 'File' dropdown."
+      )
+    }
+    val file = DocumentFactory.openReadonlyDocument(uri).asFile()
+
+    Using(ParquetFileReader.open(new LocalInputFile(file.toPath))) { reader =>
+      ParquetSchemaMapping.toTexeraSchema(reader.getFooter.getFileMetaData.getSchema)
+    }.recoverWith {
+      case e: UnsupportedOperationException => scala.util.Failure(e)
+      case scala.util.control.NonFatal(e) =>
+        scala.util.Failure(
+          new RuntimeException(
+            "Failed to read the .parquet file. Please ensure it is a valid Parquet file.",
+            e
+          )
+        )
+    }.get
+  }
+}

common/workflow-operator/src/main/scala/org/apache/texera/amber/operator/source/scan/parquet/ParquetScanSourceOpExec.scala

@@ -0,0 +1,184 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.operator.source.scan.parquet
+
+import org.apache.parquet.example.data.Group
+import org.apache.parquet.example.data.simple.NanoTime
+import org.apache.parquet.example.data.simple.convert.GroupRecordConverter
+import org.apache.parquet.hadoop.ParquetFileReader
+import org.apache.parquet.io.{ColumnIOFactory, LocalInputFile}
+import org.apache.parquet.schema.LogicalTypeAnnotation.{
+  DateLogicalTypeAnnotation,
+  DecimalLogicalTypeAnnotation,
+  IntLogicalTypeAnnotation,
+  JsonLogicalTypeAnnotation,
+  StringLogicalTypeAnnotation,
+  TimeUnit,
+  TimestampLogicalTypeAnnotation
+}
+import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName
+import org.apache.parquet.schema.{MessageType, PrimitiveType}
+import org.apache.texera.amber.core.executor.SourceOperatorExecutor
+import org.apache.texera.amber.core.storage.DocumentFactory
+import org.apache.texera.amber.core.tuple.TupleLike
+import org.apache.texera.amber.util.JSONUtils.objectMapper
+
+import java.net.URI
+import java.sql.Timestamp
+import java.time.{Instant, LocalDate, LocalDateTime, ZoneOffset}
+import java.time.temporal.ChronoUnit
+import scala.jdk.CollectionConverters._
+
+class ParquetScanSourceOpExec(descString: String) extends SourceOperatorExecutor {
+
+  /** The Julian day 1970-01-01 falls on, which an INT96 counts its days from. */
+  private val UnixEpochJulianDay = 2440588L
+
+  private val desc: ParquetScanSourceOpDesc =
+    objectMapper.readValue(descString, classOf[ParquetScanSourceOpDesc])
+  private var reader: Option[ParquetFileReader] = None
+
+  override def open(): Unit = {
+    val file = DocumentFactory.openReadonlyDocument(new URI(desc.fileName.get)).asFile()
+    reader = Some(ParquetFileReader.open(new LocalInputFile(file.toPath)))
+  }
+
+  override def produceTuple(): Iterator[TupleLike] = {
+    val fileReader = reader.get
+    val messageType = fileReader.getFooter.getFileMetaData.getSchema
+    val columns = messageType.getFields.asScala.toVector
+
+    // One row group at a time: the format stores rows in groups and a reader
+    // that asked for all of them at once would hold the whole file in memory,
+    // which is the thing a columnar format is chosen to avoid.
+    val rows: Iterator[TupleLike] = Iterator
+      .continually(fileReader.readNextRowGroup())
+      .takeWhile(_ != null)
+      .flatMap { pages =>
+        val recordReader = new ColumnIOFactory()
+          .getColumnIO(messageType)
+          .getRecordReader(pages, new GroupRecordConverter(messageType))
+        (0L until pages.getRowCount).iterator.map { _ =>
+          val group = recordReader.read()
+          TupleLike(columns.indices.map(i => readField(group, i, messageType)): _*)
+        }
+      }
+
+    val afterOffset = rows.drop(desc.offset.getOrElse(0))
+    desc.limit.fold(afterOffset)(afterOffset.take)
+  }
+
+  /** One cell, as the Texera type [[ParquetSchemaMapping]] said the column is. */
+  private def readField(group: Group, index: Int, messageType: MessageType): Any = {
+    // An optional column that was not written for this row repeats zero times.
+    // Parquet has no "null value": absence is the null.
+    if (group.getFieldRepetitionCount(index) == 0) return null
+    val primitive = messageType.getType(index).asPrimitiveType()
+    primitive.getLogicalTypeAnnotation match {
+      case annotation: DecimalLogicalTypeAnnotation =>
+        decimalOf(group, index, primitive, annotation)
+      // An unsigned column counts up where its storage counts down, so the bits
+      // are read again as the number the file means: -1 in an unsigned 32-bit
+      // column is 4294967295, which is what pandas reads there too. The narrower
+      // widths are stored as themselves and the same reading leaves them alone.
+      case annotation: IntLogicalTypeAnnotation if !annotation.isSigned =>
+        group.getInteger(index, 0).toLong & 0xffffffffL
+      case _ => primitiveOf(group, index, primitive)
+    }
+  }
+
+  /** The number a DECIMAL column stands for: its integer, with the point moved. */
+  private def decimalOf(
+      group: Group,
+      index: Int,
+      primitive: PrimitiveType,
+      annotation: DecimalLogicalTypeAnnotation
+  ): Double = {
+    val unscaled = primitive.getPrimitiveTypeName match {
+      case PrimitiveTypeName.INT32 => BigInt(group.getInteger(index, 0))
+      case PrimitiveTypeName.INT64 => BigInt(group.getLong(index, 0))
+      // The wider ones store the integer as its big-endian two's complement.
+      case _ => BigInt(group.getBinary(index, 0).getBytes)
+    }
+    BigDecimal(unscaled, annotation.getScale).toDouble
+  }
+
+  /** One cell of a column the file states nothing more about than its storage. */
+  private def primitiveOf(group: Group, index: Int, primitive: PrimitiveType): Any = {
+    primitive.getPrimitiveTypeName match {
+      case PrimitiveTypeName.BOOLEAN => group.getBoolean(index, 0)
+      case PrimitiveTypeName.FLOAT   => group.getFloat(index, 0).toDouble
+      case PrimitiveTypeName.DOUBLE  => group.getDouble(index, 0)
+      case PrimitiveTypeName.INT32 =>
+        val raw = group.getInteger(index, 0)
+        primitive.getLogicalTypeAnnotation match {
+          // A DATE is a count of days, and Texera's nearest column is a moment.
+          // Midnight of that day, in the same UTC the file counts from.
+          case _: DateLogicalTypeAnnotation =>
+            Timestamp.valueOf(LocalDate.ofEpochDay(raw.toLong).atStartOfDay)
+          case _ => raw
+        }
+      case PrimitiveTypeName.INT64 =>
+        val raw = group.getLong(index, 0)
+        primitive.getLogicalTypeAnnotation match {
+          case annotation: TimestampLogicalTypeAnnotation =>
+            // A Texera TIMESTAMP carries no zone, so the count from the epoch is
+            // read with UTC arithmetic and the wall clock it lands on is the
+            // whole of the value. `new Timestamp(millis)` would instead shift it
+            // by whatever zone the machine running the workflow is set to, and
+            // the exported script, which reads the same file with pandas, would
+            // disagree by exactly that offset.
+            Timestamp.valueOf(
+              LocalDateTime.ofInstant(instantOf(raw, annotation.getUnit), ZoneOffset.UTC)
+            )
+          case _ => raw
+        }
+      case PrimitiveTypeName.BINARY =>
+        primitive.getLogicalTypeAnnotation match {
+          case _: StringLogicalTypeAnnotation | _: JsonLogicalTypeAnnotation =>
+            group.getBinary(index, 0).toStringUsingUTF8
+          case _ => group.getBinary(index, 0).getBytes
+        }
+      // Twelve bytes holding a Julian day and the nanoseconds into it, which is
+      // how a timestamp was written before there was an annotation for one. The
+      // same UTC arithmetic as above: the wall clock it lands on is the value.
+      case PrimitiveTypeName.INT96 =>
+        val nanoTime = NanoTime.fromBinary(group.getInt96(index, 0))
+        val moment = Instant.EPOCH
+          .plus(nanoTime.getJulianDay - UnixEpochJulianDay, ChronoUnit.DAYS)
+          .plusNanos(nanoTime.getTimeOfDayNanos)
+        Timestamp.valueOf(LocalDateTime.ofInstant(moment, ZoneOffset.UTC))
+      case PrimitiveTypeName.FIXED_LEN_BYTE_ARRAY => group.getBinary(index, 0).getBytes
+    }
+  }
+
+  /** The moment a timestamp counts to, in whatever unit the file counts in. A
+    * `java.sql.Timestamp` holds nanoseconds, so a file written in micros keeps
+    * the digits under the millisecond that rounding would drop.
+    */
+  private def instantOf(value: Long, unit: TimeUnit): Instant =
+    unit match {
+      case TimeUnit.MILLIS => Instant.EPOCH.plus(value, ChronoUnit.MILLIS)
+      case TimeUnit.MICROS => Instant.EPOCH.plus(value, ChronoUnit.MICROS)
+      case TimeUnit.NANOS  => Instant.EPOCH.plusNanos(value)
+    }
+
+  override def close(): Unit = reader.foreach(_.close())
+}