From 996fc4304a3a0784dc99e1f36be103d1e533db65 Mon Sep 17 00:00:00 2001 From: Miguel Branco Date: Wed, 3 Jul 2024 16:06:22 +0200 Subject: [PATCH] Adding max rows. --- .../raw/client/api/ProgramEnvironment.scala | 3 +- .../client/rql2/truffle/Rql2CsvWriter.scala | 13 ++++++- .../client/rql2/truffle/Rql2JsonWriter.scala | 34 +++++++++++++------ .../truffle/Rql2TruffleCompilerService.scala | 4 +-- .../java/raw/runtime/truffle/RawContext.java | 2 +- 5 files changed, 40 insertions(+), 16 deletions(-) diff --git a/client/src/main/scala/raw/client/api/ProgramEnvironment.scala b/client/src/main/scala/raw/client/api/ProgramEnvironment.scala index c71fa72c4..edd10de1e 100644 --- a/client/src/main/scala/raw/client/api/ProgramEnvironment.scala +++ b/client/src/main/scala/raw/client/api/ProgramEnvironment.scala @@ -20,5 +20,6 @@ final case class ProgramEnvironment( scopes: Set[String], options: Map[String, String], maybeTraceId: Option[String] = None, - jdbcUrl: Option[String] = None + jdbcUrl: Option[String] = None, + maxRows: Option[Long] = None ) diff --git a/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2CsvWriter.scala b/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2CsvWriter.scala index aa5e0d7fb..f23fb71da 100644 --- a/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2CsvWriter.scala +++ b/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2CsvWriter.scala @@ -47,7 +47,7 @@ import java.util.Base64 import scala.annotation.tailrec import scala.util.control.NonFatal -final class Rql2CsvWriter(os: OutputStream, lineSeparator: String) extends Closeable { +final class Rql2CsvWriter(os: OutputStream, lineSeparator: String, maxRows: Option[Long]) extends Closeable { final private val gen = try { @@ -88,6 +88,7 @@ final class Rql2CsvWriter(os: OutputStream, lineSeparator: String) extends Close write(v, t.cloneAndRemoveProp(nullable).asInstanceOf[Rql2TypeWithProperties]) } } else { + var rowsWritten = 0L t match { case Rql2IterableType(recordType: Rql2RecordType, _) => val columnNames = recordType.atts.map(_.idn) @@ -100,6 +101,11 @@ final class Rql2CsvWriter(os: OutputStream, lineSeparator: String) extends Close while (iterator.hasIteratorNextElement) { val next = iterator.getIteratorNextElement writeColumns(next, recordType) + rowsWritten += 1 + // If maxRows is defined and we have written enough rows, stop writing. + if (maxRows.exists(rowsWritten >= _)) { + return + } } case Rql2ListType(recordType: Rql2RecordType, _) => val columnNames = recordType.atts.map(_.idn) @@ -112,6 +118,11 @@ final class Rql2CsvWriter(os: OutputStream, lineSeparator: String) extends Close for (i <- 0L until size) { val next = v.getArrayElement(i) writeColumns(next, recordType) + rowsWritten += 1 + // If maxRows is defined and we have written enough rows, stop writing. + if (maxRows.exists(rowsWritten >= _)) { + return + } } case _ => throw new IOException("unsupported type") } diff --git a/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2JsonWriter.scala b/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2JsonWriter.scala index f4601fb22..1f24b3300 100644 --- a/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2JsonWriter.scala +++ b/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2JsonWriter.scala @@ -23,7 +23,7 @@ import java.time.format.DateTimeFormatter import java.util.Base64 import scala.util.control.NonFatal -final class Rql2JsonWriter(os: OutputStream) extends Closeable { +final class Rql2JsonWriter(os: OutputStream, maxRows: Option[Long]) extends Closeable { final private val gen = try { @@ -45,15 +45,15 @@ final class Rql2JsonWriter(os: OutputStream) extends Closeable { if (v.isException) { v.throwException() } else { - writeValue(v, t.cloneAndRemoveProp(tryable).asInstanceOf[Rql2TypeWithProperties]) + writeValue(v, t.cloneAndRemoveProp(tryable).asInstanceOf[Rql2TypeWithProperties], maxRows) } } else { - writeValue(v, t.cloneAndRemoveProp(tryable).asInstanceOf[Rql2TypeWithProperties]) + writeValue(v, t.cloneAndRemoveProp(tryable).asInstanceOf[Rql2TypeWithProperties], maxRows) } } @throws[IOException] - private def writeValue(v: Value, t: Rql2TypeWithProperties): Unit = { + private def writeValue(v: Value, t: Rql2TypeWithProperties, maxRows: Option[Long]): Unit = { if (t.props.contains(tryable)) { if (v.isException) { try { @@ -61,10 +61,10 @@ final class Rql2JsonWriter(os: OutputStream) extends Closeable { } catch { case NonFatal(ex) => gen.writeString(ex.getMessage) } - } else writeValue(v, t.cloneAndRemoveProp(tryable).asInstanceOf[Rql2TypeWithProperties]) + } else writeValue(v, t.cloneAndRemoveProp(tryable).asInstanceOf[Rql2TypeWithProperties], maxRows = None) } else if (t.props.contains(nullable)) { if (v.isNull) gen.writeNull() - else writeValue(v, t.cloneAndRemoveProp(nullable).asInstanceOf[Rql2TypeWithProperties]) + else writeValue(v, t.cloneAndRemoveProp(nullable).asInstanceOf[Rql2TypeWithProperties], maxRows = None) } else t match { case _: Rql2BinaryType => val bytes = (0L until v.getBufferSize).map(v.readBufferByte) @@ -112,34 +112,46 @@ final class Rql2JsonWriter(os: OutputStream) extends Closeable { val field = distincted.get(i) gen.writeFieldName(field) val a = v.getMember(field) - writeValue(a, atts(i).tipe.asInstanceOf[Rql2TypeWithProperties]) + writeValue(a, atts(i).tipe.asInstanceOf[Rql2TypeWithProperties], maxRows = None) } gen.writeEndObject() case Rql2IterableType(innerType, _) => + var rowsWritten = 0L val iterator = v.getIterator gen.writeStartArray() while (iterator.hasIteratorNextElement) { val next = iterator.getIteratorNextElement - writeValue(next, innerType.asInstanceOf[Rql2TypeWithProperties]) + writeValue(next, innerType.asInstanceOf[Rql2TypeWithProperties], maxRows = None) + rowsWritten += 1 + // If maxRows is defined and we have written enough rows, stop writing. + if (maxRows.exists(rowsWritten >= _)) { + return + } } gen.writeEndArray() case Rql2ListType(innerType, _) => + var rowsWritten = 0L val size = v.getArraySize gen.writeStartArray() for (i <- 0L until size) { val next = v.getArrayElement(i) - writeValue(next, innerType.asInstanceOf[Rql2TypeWithProperties]) + writeValue(next, innerType.asInstanceOf[Rql2TypeWithProperties], maxRows = None) + rowsWritten += 1 + // If maxRows is defined and we have written enough rows, stop writing. + if (maxRows.exists(rowsWritten >= _)) { + return + } } gen.writeEndArray() case Rql2OrType(tipes, _) if tipes.exists(Rql2TypeUtils.getProps(_).nonEmpty) => // A trick to make sur inner types do not have properties val inners = tipes.map { case inner: Rql2TypeWithProperties => Rql2TypeUtils.resetProps(inner, Set.empty) } val orProps = tipes.flatMap { case inner: Rql2TypeWithProperties => inner.props }.toSet - writeValue(v, Rql2OrType(inners, orProps)) + writeValue(v, Rql2OrType(inners, orProps), maxRows = None) case Rql2OrType(tipes, _) => val index = v.invokeMember("getIndex").asInt() val actualValue = v.invokeMember("getValue") - writeValue(actualValue, tipes(index).asInstanceOf[Rql2TypeWithProperties]) + writeValue(actualValue, tipes(index).asInstanceOf[Rql2TypeWithProperties], maxRows = None) case _ => throw new RuntimeException("unsupported type") } diff --git a/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2TruffleCompilerService.scala b/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2TruffleCompilerService.scala index c660d5d00..12090d093 100644 --- a/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2TruffleCompilerService.scala +++ b/snapi-client/src/main/scala/raw/client/rql2/truffle/Rql2TruffleCompilerService.scala @@ -279,7 +279,7 @@ class Rql2TruffleCompilerService(engineDefinition: (Engine, Boolean))(implicit p case _ => programContext.settings.config.getBoolean("raw.compiler.windows-line-ending") } val lineSeparator = if (windowsLineEnding) "\r\n" else "\n" - val w = new Rql2CsvWriter(outputStream, lineSeparator) + val w = new Rql2CsvWriter(outputStream, lineSeparator, environment.maxRows) try { w.write(v, tipe.asInstanceOf[Rql2TypeWithProperties]) w.flush() @@ -293,7 +293,7 @@ class Rql2TruffleCompilerService(engineDefinition: (Engine, Boolean))(implicit p if (!JsonPackage.outputWriteSupport(tipe)) { return ExecutionRuntimeFailure("unsupported type") } - val w = new Rql2JsonWriter(outputStream) + val w = new Rql2JsonWriter(outputStream, environment.maxRows) try { w.write(v, tipe.asInstanceOf[Rql2TypeWithProperties]) w.flush() diff --git a/snapi-truffle/src/main/java/raw/runtime/truffle/RawContext.java b/snapi-truffle/src/main/java/raw/runtime/truffle/RawContext.java index c3e80e8f1..e4d3e7a0c 100644 --- a/snapi-truffle/src/main/java/raw/runtime/truffle/RawContext.java +++ b/snapi-truffle/src/main/java/raw/runtime/truffle/RawContext.java @@ -98,7 +98,7 @@ public RawContext(RawLanguage language, Env env) { Option[]> maybeArguments = Option.empty(); this.programEnvironment = new ProgramEnvironment( - this.user, maybeArguments, scalaScopes, scalaOptions, maybeTraceId, Option.empty()); + this.user, maybeArguments, scalaScopes, scalaOptions, maybeTraceId, Option.empty(), Option.empty()); // The function registry holds snapi methods (top level functions). It is the data // structure that is used to extract a ref to a function from a piece of execute snapi.