Github user HyukjinKwon commented on a diff in the pull request:
https://github.com/apache/spark/pull/16976#discussion_r103059540
--- Diff:
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/csv/UnivocityParser.scala
---
@@ -269,3 +273,89 @@ private[csv] class UnivocityParser(
}
}
}
+
+private[csv] object UnivocityParser {
+
+ /**
+ * Parses a stream that contains CSV strings and turns it into an
iterator of tokens.
+ */
+ def tokenizeStream(
+ inputStream: InputStream,
+ shouldDropHeader: Boolean,
+ tokenizer: CsvParser): Iterator[Array[String]] = new
Iterator[Array[String]] {
+ // Note that, here we assume `inputStream` is the whole file that
might include the header.
+ tokenizer.beginParsing(inputStream)
+ private var nextRecord = {
+ if (shouldDropHeader) {
+ tokenizer.parseNext()
+ }
+ tokenizer.parseNext()
+ }
+
+ override def hasNext: Boolean = nextRecord != null
+
+ override def next(): Array[String] = {
+ if (!hasNext) {
+ throw new NoSuchElementException("End of stream")
+ }
+ val curRecord = nextRecord
+ nextRecord = tokenizer.parseNext()
+ curRecord
+ }
+ }
+
+ /**
+ * Parses a stream that contains CSV strings and turns it into an
iterator of rows.
+ */
+ def parseStream(
+ inputStream: InputStream,
+ shouldDropHeader: Boolean,
+ parser: UnivocityParser): Iterator[InternalRow] = {
+ val options = parser.options
+ val tokenizer = parser.tokenizer
+ val corruptFieldIndex = parser.corruptFieldIndex
+ val shouldKeepCurrentInput = options.permissive &&
corruptFieldIndex.isDefined
+
+ val tokenizedIter = tokenizeStream(inputStream, shouldDropHeader,
tokenizer)
+
+ val parsedIter = new Iterator[Option[InternalRow]] {
+ private var currentInput: String = _
+
+ override def hasNext: Boolean = tokenizedIter.hasNext
+
+ override def next(): Option[InternalRow] = {
+ if (shouldKeepCurrentInput) {
+ // We should keep the current content before actually parsing in
this case.
+ // Otherwise, we lose the string.
+ currentInput =
tokenizer.getContext.currentParsedContent().stripLineEnd
--- End diff --
It expects,
> Returns a String with the input character sequence parsed to produce the
current record.
For example,
input
```csv
0,2013-111-11 12:13:14
1,1983-08-04
```
output (sequentially)
- `CsvParser.parseNext()`
```scala
Array("0", "2013-111-11 12:13:14")
```
- `CsvParser.getContext.currentParsedContent()`
```
0,2013-111-11 12:13:14\n
```
- `CsvParser.parseNext()`
```scala
Array("1", "1983-08-04")
```
- `CsvParser.getContext.currentParsedContent()`
```
1,1983-08-04\n
```
It seems re-use the internal buffer but it requires an overhead. Let me try
to avoid this extra call.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at infrastruct...@apache.org or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
