comphead commented on code in PR #6784: URL: https://github.com/apache/datafusion-comet/pull/6784#discussion_r4459280219
########## native/core/src/execution/operators/text_scan.rs: ########## @@ -0,0 +1,796 @@ +// 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. + +//! Native Text file scan, mirroring Spark's `text` data source. +//! +//! Spark's text reader produces a single `value: string` column. With the default options each +//! line of the file is one row (split on `\n`, `\r\n`, or `\r`, with the terminator stripped). +//! With `wholetext=true` the whole file is a single row. A custom `lineSep` splits on exactly that +//! byte sequence instead of the universal-newline set. See the `TextOptions`/`HadoopFileLinesReader` +//! logic in Spark's `org.apache.spark.sql.execution.datasources.text` package. + +use crate::execution::operators::ExecutionError; +use arrow::array::StringBuilder; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use bytes::Bytes; +use datafusion::common::Result; +use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::physical_expr::projection::ProjectionExprs; +use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder}; +use datafusion::physical_plan::DisplayFormatType; +use datafusion_comet_common::decode_utf8_spark_lossy; +use datafusion_comet_proto::spark_operator::TextOptions; +use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_compression_type::FileCompressionType; +use datafusion_datasource::file_groups::FileGroup; +use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener}; +use datafusion_datasource::projection::{ProjectionOpener, SplitProjection}; +use datafusion_datasource::source::DataSourceExec; +use datafusion_datasource::{as_file_source, PartitionedFile, TableSchema}; +use futures::StreamExt; +use object_store::{ObjectStore, ObjectStoreExt}; +use std::borrow::Cow; +use std::fmt; +use std::io::Read; +use std::sync::Arc; + +/// Batch size used when the execution framework does not set one before opening the file. +const DEFAULT_TEXT_BATCH_SIZE: usize = 8192; + +/// UTF-8 byte-order mark. Spark's line reader strips a leading one; wholetext keeps it. +const UTF8_BOM: &[u8] = &[0xEF, 0xBB, 0xBF]; + +/// Cap on a batch's decoded `value` buffer bytes, kept well below Arrow's 2GB 32-bit offset limit +/// so a batch of many rows (or a few long lines) cannot overflow. Measured against the decoded +/// string length, since a lossy UTF-8 decode can expand ill-formed bytes up to 3x. +const MAX_BATCH_VALUE_BYTES: usize = 1 << 30; + +pub fn init_text_datasource_exec( + object_store_url: ObjectStoreUrl, + file_groups: Vec<Vec<PartitionedFile>>, + data_schema: SchemaRef, + _partition_schema: SchemaRef, + projection_vector: Vec<usize>, + text_options: &TextOptions, +) -> Result<Arc<DataSourceExec>, ExecutionError> { + let text_source = TextSource::new(data_schema, text_options); + + let file_groups = file_groups.into_iter().map(FileGroup::new).collect(); + + let file_scan_config = FileScanConfigBuilder::new(object_store_url, Arc::new(text_source)) + .with_file_groups(file_groups) + .with_projection_indices(Some(projection_vector))? + .build(); + + Ok(DataSourceExec::from_data_source(file_scan_config)) +} + +/// A [`FileSource`] for Spark's text format. The file schema is always the single `value` column; +/// projection (and partition columns, if any) is handled by [`ProjectionOpener`], exactly as the +/// CSV source does. +#[derive(Debug, Clone)] +struct TextSource { + table_schema: TableSchema, + projection: SplitProjection, + batch_size: Option<usize>, + metrics: ExecutionPlanMetricsSet, + whole_text: bool, + // The raw separator bytes when a custom `lineSep` is set; None means universal newline + // (`\n`, `\r\n`, `\r`). + line_sep: Option<Vec<u8>>, +} + +impl TextSource { + fn new(file_schema: SchemaRef, options: &TextOptions) -> Self { + let table_schema: TableSchema = file_schema.into(); + // `line_sep` arrives as the exact separator bytes Spark computed (`lineSeparatorInRead`, + // already encoded with the file's charset), so we split on them directly. + let line_sep = options.line_sep.clone().filter(|sep| !sep.is_empty()); + Self { + projection: SplitProjection::unprojected(&table_schema), + table_schema, + batch_size: None, + metrics: ExecutionPlanMetricsSet::new(), + whole_text: options.whole_text, + line_sep, + } + } +} + +impl From<TextSource> for Arc<dyn FileSource> { + fn from(source: TextSource) -> Self { + as_file_source(source) + } +} + +impl FileSource for TextSource { + fn create_file_opener( + &self, + object_store: Arc<dyn ObjectStore>, + base_config: &FileScanConfig, + partition_index: usize, + ) -> Result<Arc<dyn FileOpener>> { + // The inner opener must emit batches whose schema is the file schema projected by + // `file_indices` (that is what ProjectionOpener feeds its projector). For text that is + // either the single `value` column or, for a projection-less count, an empty schema. + let file_schema = self.table_schema.file_schema(); + let projected_file_schema = Arc::new(file_schema.project(&self.projection.file_indices)?); + let opener = Arc::new(TextOpener { + projected_file_schema, + file_compression_type: base_config.file_compression_type, + object_store, + partition_index, + batch_size: self.batch_size.unwrap_or(DEFAULT_TEXT_BATCH_SIZE), + metrics: self.metrics.clone(), + whole_text: self.whole_text, + line_sep: self.line_sep.clone(), + }) as Arc<dyn FileOpener>; + ProjectionOpener::try_new(self.projection.clone(), opener, file_schema) + } + + fn table_schema(&self) -> &TableSchema { + &self.table_schema + } + + fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> { + let mut conf = self.clone(); + conf.batch_size = Some(batch_size); + Arc::new(conf) + } + + fn try_pushdown_projection( + &self, + projection: &ProjectionExprs, + ) -> Result<Option<Arc<dyn FileSource>>> { + let mut source = self.clone(); + let new_projection = self.projection.source.try_merge(projection)?; + source.projection = SplitProjection::new(self.table_schema.file_schema(), &new_projection); + Ok(Some(Arc::new(source))) + } + + fn projection(&self) -> Option<&ProjectionExprs> { + Some(&self.projection.source) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn file_type(&self) -> &str { + "text" + } + + // A custom line separator can straddle a byte-range boundary, and we read whole files rather + // than boundary-aligned ranges, so this source must not let DataFusion repartition files by + // byte range. + fn supports_repartitioning(&self) -> bool { + false + } + + fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, ", whole_text={}", self.whole_text) + } + DisplayFormatType::TreeRender => Ok(()), + } + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc<dyn datafusion::physical_plan::PhysicalExpr>, + ) -> Result<datafusion::common::tree_node::TreeNodeRecursion>, + ) -> Result<datafusion::common::tree_node::TreeNodeRecursion> { + datafusion::physical_plan::apply_expression_roots(self.projection.source.iter(), f) + } +} + +/// A [`FileOpener`] that reads a whole text file and yields the projected file schema in batches. +struct TextOpener { + // The file schema projected by `file_indices`: the single `value` column, or an empty schema + // for a projection-less count. + projected_file_schema: SchemaRef, + file_compression_type: FileCompressionType, + object_store: Arc<dyn ObjectStore>, + partition_index: usize, + batch_size: usize, + metrics: ExecutionPlanMetricsSet, + whole_text: bool, + line_sep: Option<Vec<u8>>, +} + +impl FileOpener for TextOpener { + fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> { + let store = Arc::clone(&self.object_store); + let projected_file_schema = Arc::clone(&self.projected_file_schema); + let file_compression_type = self.file_compression_type; + let batch_size = self.batch_size; + let whole_text = self.whole_text; + let line_sep = self.line_sep.clone(); + let baseline_metrics = BaselineMetrics::new(&self.metrics, self.partition_index); + // Feeds the task's input bytes in the Spark UI (see CometMetricNode.reportScanInputMetrics). + let bytes_scanned = + MetricBuilder::new(&self.metrics).counter("bytes_scanned", self.partition_index); + + Ok(Box::pin(async move { + let range = partitioned_file.range.clone(); + let location = partitioned_file.object_meta.location; + let raw = store.get(&location).await?.bytes().await?; + bytes_scanned.add(raw.len()); + + // Decompress only when the file is actually compressed. On Comet's plan config the file + // is always UNCOMPRESSED (CometScanRule falls back for compressed files), so we keep + // the fetched bytes as they are and avoid a full-file copy; the decompress arm is a + // defensive handler should compression ever be plumbed through. A byte range never + // coexists with decompression (compressed files are non-splittable). + // + // NOTE: the whole object is read into memory. CometScanRule restricts native text to + // unsplit files (it falls back when Spark splits a file into byte ranges) and caps the + // file size, so in practice this reads each small file once. Lines are found and + // batches built on demand, so only the bytes and the current batch are held. + let buf: Bytes = if file_compression_type.is_compressed() { + let mut decoder = file_compression_type.convert_read(std::io::Cursor::new(raw))?; + let mut decoded = Vec::new(); + decoder.read_to_end(&mut decoded)?; + Bytes::from(decoded) + } else { + raw + }; + + // Spark's line reader (Hadoop LineRecordReader.skipUtfByteOrderMark) strips a leading + // UTF-8 BOM from the first record at file offset 0, whatever the separator. wholetext + // keeps it. Native reads unsplit files only. + let at_file_start = range.as_ref().map(|r| r.start == 0).unwrap_or(true); + // Spark may split an uncompressed file into byte ranges (TextScan.isSplitable). A line + // belongs to the split whose [start, end) contains the line's start offset -- so each + // line is emitted by exactly one split, with no duplication or gaps across splits. See + // DataFusion's CsvOpener for the same rule. wholetext files are never split. + let byte_range = if whole_text { + None + } else { + range.map(|r| (r.start.max(0) as usize, r.end.max(0) as usize)) + }; Review Comment: Two paths here cannot run for a scan the gate admits. `init_text_datasource_exec` never sets `file_compression_type` and `CometScanRule` declines compressed files, so the decompress arm, the `std::io::Read` import and the `file_compression_type` field are dead. `CometScanRule` also declines any file with `start != 0`, so `byte_range` is always the whole file and its per-row check in `next_row` does nothing. Could the decompress arm, `byte_range`, `at_file_start` and the two `range_split_*` tests go, leaving the Scala gate as the guard? If the defensive filter stays, note that Hadoop's `LineRecordReader` gives a line starting exactly at a split boundary to the earlier split, while `[start, end)` gives it to the later one. ########## native/core/src/execution/planner.rs: ########## @@ -1922,6 +1923,55 @@ impl PhysicalPlanner { Arc::new(SparkPlan::new(spark_plan.plan_id, scan, vec![])), )) } + OpStruct::TextScan(scan) => { + let data_schema = convert_spark_types_to_arrow_schema(scan.data_schema.as_slice()); + let partition_schema = + convert_spark_types_to_arrow_schema(scan.partition_schema.as_slice()); + let projection_vector: Vec<usize> = + scan.projection_vector.iter().map(|i| *i as usize).collect(); + let object_store_options: HashMap<String, String> = + scan.object_store_options.clone(); + let partition = usize::try_from(self.partition) + .ok() + .and_then(|p| scan.file_partitions.get(p)) + .ok_or_else(|| { + GeneralError(format!( + "TextScan has no file partition {} ({} partitions)", + self.partition, + scan.file_partitions.len() + )) + })?; + // Build the store from this partition's files; other partitions may read + // from a different one. + let object_store_url = match partition.partitioned_file.first() { + Some(file) => { + prepare_object_store_with_configs( + self.session_ctx.runtime_env(), + file.file_path.clone(), + &object_store_options, + )? + .0 Review Comment: The `CsvScan` arm above goes through `self.scan_store(...)`, which returns the local file system when the planner runs `without_data()` (from `native_executors/fragment.rs`), so that checking a plan needs no storage credentials. This arm calls `prepare_object_store_with_configs` directly, so a `TextScan` in such a plan would build the real store. I have not run it. Could it use `self.scan_store(&file.file_path, &object_store_options)?.0` like the CSV arm? A helper shared by both arms for the partition lookup, store and file group would keep them from drifting again. Also, `partition_schema` is always empty here because `CometScanRule` declines partition columns, and `init_text_datasource_exec` ignores it, so field 2 of `TextScan` could go. ########## spark/src/main/scala/org/apache/spark/sql/comet/CometTextNativeScanExec.scala: ########## @@ -0,0 +1,196 @@ +/* + * 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.spark.sql.comet + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder} +import org.apache.spark.sql.catalyst.plans.QueryPlan +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.FilePartition +import org.apache.spark.sql.execution.datasources.text.TextOptions +import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.execution.datasources.v2.text.TextScan +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} + +import com.google.common.base.Objects + +import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.objectstore.NativeConfig +import org.apache.comet.serde.{CometOperatorSerde, OperatorOuterClass} +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.{partition2Proto, schema2Proto} + +/* + * Native Text scan operator that delegates file reading to datafusion. Mirrors the native CSV + * scan path; produces a single `value: string` column. + */ +case class CometTextNativeScanExec( + override val nativeOp: Operator, + override val output: Seq[Attribute], + @transient override val originalPlan: BatchScanExec, + override val serializedPlanOpt: SerializedPlan) + extends CometLeafExec { + override val supportsColumnar: Boolean = true + + override val nodeName: String = "CometTextNativeScan" + + // Show only the semantic output in EXPLAIN / the SQL tab / the event log. The default would print + // the whole `nativeOp` protobuf, which carries every input file path (large for a lookup + // directory of many files). Mirrors CometNativeScanExec. + override def stringArgs: Iterator[Any] = Iterator(output) + + // The serialized partitions, which split any Spark partition that mixes object stores. + override def outputPartitioning: Partitioning = UnknownPartitioning( + nativeOp.getTextScan.getFilePartitionsCount) + + override def outputOrdering: Seq[SortOrder] = Nil + + // The scan reports the bytes it fetched as `bytes_scanned`, which CometMetricNode turns into the + // task's input bytes and records, so they show on the stage UI like any other scan. + override lazy val metrics: Map[String, SQLMetric] = + CometMetricNode.baselineMetrics(sparkContext) + + ("bytes_scanned" -> SQLMetrics.createSizeMetric(sparkContext, "number of bytes scanned")) + + // Two scans of the same files must compare equal so Spark reuses their exchanges and stages. + // That needs normalized output attribute ids and the wrapped scan's canonical form, and no + // serialization state, which differs between otherwise identical copies. + override protected def doCanonicalize(): SparkPlan = { + CometTextNativeScanExec( + nativeOp, + output.map(QueryPlan.normalizeExpressions(_, output)), + if (originalPlan == null) null else originalPlan.doCanonicalize(), + SerializedPlan(None)) + } + + override def equals(obj: Any): Boolean = { + obj match { + case other: CometTextNativeScanExec => + this.output == other.output && + this.serializedPlanOpt == other.serializedPlanOpt && + this.originalPlan == other.originalPlan + case _ => + false + } + } + + override def hashCode(): Int = { + Objects.hashCode(output, serializedPlanOpt, originalPlan) + } +} + +object CometTextNativeScanExec extends CometOperatorSerde[CometBatchScanExec] { + + override def enabledConfig: Option[ConfigEntry[Boolean]] = Some( + CometConf.COMET_TEXT_V2_NATIVE_ENABLED) + + override def convert( + op: CometBatchScanExec, + builder: Operator.Builder, + childOp: Operator*): Option[Operator] = { + val textScanBuilder = OperatorOuterClass.TextScan.newBuilder() + val textScan = op.wrapped.scan.asInstanceOf[TextScan] + val sessionState = op.session.sessionState + val options = new TextOptions(textScan.options.asScala.toMap) + // Build the Hadoop conf from the scan's read options (case preserved), matching Spark's + // TextScan.createReaderFactory, so per-read settings like fs.s3a.endpoint reach native. + val hadoopConf = + sessionState.newHadoopConfWithOptions(textScan.options.asCaseSensitiveMap.asScala.toMap) + val s3CompliantSchemes = NativeConfig.resolveS3CompliantSchemes(hadoopConf) + // The scan gate already declined schemes that only the read options route through libhdfs. + val libhdfsSchemes = NativeConfig + .resolveScanLibhdfsSchemes( + NativeConfig.resolveLibhdfsSchemes(sessionState.newHadoopConf()), + hadoopConf) + .libhdfs + val inputPartitions = op.inputPartitions.map(_.asInstanceOf[FilePartition]) + val multiStoreReason = CometScanUtils.multiStoreFallbackReason( + "Native Text scan", + inputPartitions.view.flatMap(_.files.view.map(_.pathUri)), + s3CompliantSchemes, + libhdfsSchemes, + isBucketedScan = false) + if (multiStoreReason.nonEmpty) { + // CometExecRule falls back to the wrapped scan, so tag it too for the explain output. + withFallbackReason(op, multiStoreReason.get) + withFallbackReason(op.wrapped, multiStoreReason.get) + return None + } + // The options of every scheme the scan reads, or a decline for the ABFS layouts one + // resolution cannot serve. Checked before the split below so that its reason is reported. + val fileUris = inputPartitions.view.flatMap(_.files).map(_.pathUri) + val objectStoreOptions = + NativeConfig.extractScanObjectStoreOptions(hadoopConf, fileUris) match { + case Right(options) => options + case Left(reason) => + withFallbackReason(op, reason) + withFallbackReason(op.wrapped, reason) + return None + } + // Native planning reads each partition through one object store. + val filePartitions = CometScanUtils.splitPartitionsByStore( + inputPartitions, + file => NativeConfig.objectStoreKey(file.pathUri, s3CompliantSchemes, libhdfsSchemes)) + // Spark 3.4 reports a one-partition scan as SinglePartition and plans no exchange above it, + // so the operators above rely on the scan's partition count. + if (filePartitions.length != inputPartitions.length && + !op.wrapped.outputPartitioning.isInstanceOf[UnknownPartitioning]) { + val reason = "Native Text scan would split a partition that mixes object stores, but " + + "Spark planned the operators above it for the scan's own partitioning" + withFallbackReason(op, reason) + withFallbackReason(op.wrapped, reason) + return None + } + val textOptionsProto = textOptions2Proto(options) + val dataSchemaProto = schema2Proto(textScan.dataSchema) + val readSchemaFieldNames = textScan.readDataSchema.fieldNames + val projectionVector = textScan.dataSchema.zipWithIndex + .filter { case (field, _) => + readSchemaFieldNames.contains(field.name) + } + .map(_._2.asInstanceOf[Integer]) + val partitionSchemaProto = schema2Proto(textScan.readPartitionSchema) + val partitionsProto = filePartitions.map(partition2Proto(_, textScan.readPartitionSchema)) + + textScanBuilder.putAllObjectStoreOptions(objectStoreOptions.asJava) + textScanBuilder.setTextOptions(textOptionsProto) + textScanBuilder.addAllFilePartitions(partitionsProto.asJava) + textScanBuilder.addAllDataSchema(dataSchemaProto.asJava) + textScanBuilder.addAllProjectionVector(projectionVector.asJava) + textScanBuilder.addAllPartitionSchema(partitionSchemaProto.asJava) + Some(builder.setTextScan(textScanBuilder).build()) Review Comment: About 70 lines of this `convert` match `CometCsvNativeScanExec.convert`: the Hadoop conf from the read options, `multiStoreFallbackReason`, `extractScanObjectStoreOptions`, `splitPartitionsByStore`, the Spark 3.4 `SinglePartition` check, the projection vector and `partition2Proto`. They have already drifted, since the CSV node declines `scanOnly` libhdfs schemes in `convert` while this one relies on `rootPathFallbackReason`. `equals`, `hashCode`, `stringArgs` and `outputOrdering` repeat too. Could the shared part move into a helper in `CometScanUtils` used by both? That would also carry the new `doCanonicalize` to the CSV node, which still keeps `serializedPlanOpt` in its canonical form, so I expect it to miss the same exchange reuse. I have not run that. ########## spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala: ########## @@ -477,6 +419,186 @@ case class CometScanRule(session: SparkSession) withFallbackReasons(scanExec, fallbackReasons.toSet) } + case scan: TextScan if COMET_TEXT_V2_NATIVE_ENABLED.get() => + if (scanExec.output.exists(_.isMetadataCol)) { + return withFallbackReason( + scanExec, + "Metadata columns are not supported for Text V2 scans") + } + // A native scan is useless without native execution: with exec disabled, CometExecRule + // leaves this CometBatchScanExec unconverted and it fails at runtime. Require native exec. + if (!COMET_EXEC_ENABLED.get()) { + return withFallbackReason( + scanExec, + s"Native Text scan requires ${COMET_EXEC_ENABLED.key} to be enabled") + } + // input_file_name/_block_start/_block_length are already handled by the blanket V2 gate + // near the top of transformV2Scan, so no Text-specific check is needed here. + val fallbackReasons = new ListBuffer[String]() + val schemaSupported = + CometBatchScanExec.isSchemaSupported(scan.readDataSchema, fallbackReasons) + if (!schemaSupported) { + fallbackReasons += s"Schema ${scan.readDataSchema} is not supported" + } + // Spark's text source allows only a single data column (TextScan.verifyReadSchema). A + // user-forced multi-column read schema is invalid; fall back so Spark raises its structured + // AnalysisException rather than letting the native reader fail with an opaque error. + if (scan.readDataSchema.size > 1) { + fallbackReasons += "Text data source supports only a single column" + } + // Partition columns are not implemented for native Text scans (unlike the native CSV path, + // which forwards them). Fall back to Spark for partitioned reads rather than drop them. + if (scan.readPartitionSchema.nonEmpty) { + fallbackReasons += "Comet does not support partition columns in native Text scans" + } + // The native reader loads each file whole. Spark splits a large uncompressed text file + // into byte-range partitions (TextScan.isSplitable), and each split would then re-read and + // buffer the whole file with no memory-pool reservation. Restrict native execution to + // unsplit files (the small build-side lookups this targets); Spark handles large files. + val hasSplitFile = scanExec.inputPartitions.exists { + case fp: FilePartition => fp.files.exists(_.start != 0) + case _ => false + } + if (hasSplitFile) { + fallbackReasons += "Comet native Text scan does not support files split into byte " + + "ranges (large files)" + } + // The native reader loads each file whole and does not reserve the buffer against the + // memory pool, so bound the per-file size; larger files fall back to Spark. This keeps the + // feature to the small lookup tables it targets regardless of maxPartitionBytes. + val maxTextFileSize = COMET_SCAN_TEXT_MAX_FILE_SIZE.get() + val hasOversizedFile = scanExec.inputPartitions.exists { + case fp: FilePartition => fp.files.exists(_.length > maxTextFileSize) + case _ => false + } + if (hasOversizedFile) { + fallbackReasons += "Comet native Text scan does not read files larger than " + + s"${COMET_SCAN_TEXT_MAX_FILE_SIZE.key} ($maxTextFileSize bytes)" + } Review Comment: `scanExec.inputPartitions` is matched against `FilePartition` here, again for the oversized check, for `selectedFiles` and for the wholetext check, and `selectedFiles` is then scanned for compression, `.crc` files and paths. Could `selectedFiles` be built first, with `selectedFiles.exists(_.start != 0)` and `selectedFiles.exists(_.length > maxTextFileSize)` here? `CometTextNativeScanExec.convert` already casts the same partitions with `asInstanceOf[FilePartition]`, so I expect the `case _ => false` arms to be unreachable. ########## spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala: ########## @@ -477,6 +419,186 @@ case class CometScanRule(session: SparkSession) withFallbackReasons(scanExec, fallbackReasons.toSet) } + case scan: TextScan if COMET_TEXT_V2_NATIVE_ENABLED.get() => + if (scanExec.output.exists(_.isMetadataCol)) { + return withFallbackReason( + scanExec, + "Metadata columns are not supported for Text V2 scans") + } + // A native scan is useless without native execution: with exec disabled, CometExecRule + // leaves this CometBatchScanExec unconverted and it fails at runtime. Require native exec. + if (!COMET_EXEC_ENABLED.get()) { + return withFallbackReason( + scanExec, + s"Native Text scan requires ${COMET_EXEC_ENABLED.key} to be enabled") + } + // input_file_name/_block_start/_block_length are already handled by the blanket V2 gate + // near the top of transformV2Scan, so no Text-specific check is needed here. + val fallbackReasons = new ListBuffer[String]() + val schemaSupported = + CometBatchScanExec.isSchemaSupported(scan.readDataSchema, fallbackReasons) + if (!schemaSupported) { + fallbackReasons += s"Schema ${scan.readDataSchema} is not supported" + } + // Spark's text source allows only a single data column (TextScan.verifyReadSchema). A + // user-forced multi-column read schema is invalid; fall back so Spark raises its structured + // AnalysisException rather than letting the native reader fail with an opaque error. + if (scan.readDataSchema.size > 1) { + fallbackReasons += "Text data source supports only a single column" + } + // Partition columns are not implemented for native Text scans (unlike the native CSV path, + // which forwards them). Fall back to Spark for partitioned reads rather than drop them. + if (scan.readPartitionSchema.nonEmpty) { + fallbackReasons += "Comet does not support partition columns in native Text scans" + } + // The native reader loads each file whole. Spark splits a large uncompressed text file + // into byte-range partitions (TextScan.isSplitable), and each split would then re-read and + // buffer the whole file with no memory-pool reservation. Restrict native execution to + // unsplit files (the small build-side lookups this targets); Spark handles large files. + val hasSplitFile = scanExec.inputPartitions.exists { + case fp: FilePartition => fp.files.exists(_.start != 0) + case _ => false + } + if (hasSplitFile) { + fallbackReasons += "Comet native Text scan does not support files split into byte " + + "ranges (large files)" + } + // The native reader loads each file whole and does not reserve the buffer against the + // memory pool, so bound the per-file size; larger files fall back to Spark. This keeps the + // feature to the small lookup tables it targets regardless of maxPartitionBytes. + val maxTextFileSize = COMET_SCAN_TEXT_MAX_FILE_SIZE.get() + val hasOversizedFile = scanExec.inputPartitions.exists { + case fp: FilePartition => fp.files.exists(_.length > maxTextFileSize) + case _ => false + } + if (hasOversizedFile) { + fallbackReasons += "Comet native Text scan does not read files larger than " + + s"${COMET_SCAN_TEXT_MAX_FILE_SIZE.key} ($maxTextFileSize bytes)" + } + // Build the Hadoop conf from the scan's read options with their case preserved, like + // Spark's TextScan.createReaderFactory (Hadoop keys are case sensitive, but the scan's + // option map lowercases them). + val hadoopConf = session.sessionState.newHadoopConfWithOptions( + scan.options.asCaseSensitiveMap.asScala.toMap) + val textOptions = + new org.apache.spark.sql.execution.datasources.text.TextOptions( + scan.options.asScala.toMap) + // The files this scan reads. `fileIndex.inputFiles` also lists files that partition + // pruning dropped, so the per-file checks below use the selected partitions instead. + val selectedFiles = scanExec.inputPartitions.flatMap { + case fp: FilePartition => fp.files.toSeq + case _ => Seq.empty + } + // Comet's native plan does not carry a compression codec, so the native reader would feed + // raw compressed bytes to the line splitter. Spark auto-detects a codec by file extension, + // from the decoded path (`Path(URI)`), so a name like `a b.gz` keeps its suffix. The + // version-aware ShimFileFormat.compressedFileCheck also catches the non-standard + // .gzip/.zstd extensions on Spark 4.1+ that CompressionCodecFactory alone misses. + val isCompressed = ShimFileFormat.compressedFileCheck(hadoopConf) + val hasCompressedFile = + selectedFiles.exists(file => isCompressed(new org.apache.hadoop.fs.Path(file.pathUri))) + if (hasCompressedFile) { + fallbackReasons += "Comet does not support compressed text files" + } + // Hadoop's local filesystem verifies a `.<name>.crc` checksum file when one exists and + // raises a ChecksumException if the data no longer matches it. The native reader reads the + // file directly and would return the altered rows, so fall back for such files. + if (selectedFiles.exists(file => + CometScanRule.hasLocalChecksumFile(hadoopConf, file.pathUri))) { + fallbackReasons += "Comet native Text scan does not support local files that have a " + + "Hadoop checksum (.crc) file" + } + // Spark skips unreadable/missing files when these are set; the native reader's DataFusion + // FileStream fails the whole scan instead (OnError::Fail). Fall back to match Spark. + if (SQLConf.get.ignoreCorruptFiles || scan.options.getBoolean( + "ignoreCorruptFiles", + false)) { + fallbackReasons += "Comet native Text scan does not support ignoreCorruptFiles" + } + if (SQLConf.get.ignoreMissingFiles || scan.options.getBoolean( + "ignoreMissingFiles", + false)) { + fallbackReasons += "Comet native Text scan does not support ignoreMissingFiles" + } + // Only Spark's line reader consults the settings below. The wholetext reader does not, + // so a value it never reads must not fail or fall back a wholetext query. + if (!textOptions.wholeText) { + // Spark's line reader skips any line at or over + // mapreduce.input.linerecordreader.line.maxlength (default unlimited); the key can come + // from spark.hadoop.*, core-site.xml, or a read option. The native reader does not + // truncate, so fall back when a finite limit is set. A value that is not an integer + // fails Spark's reader, so fall back and let Spark report it. + val maxLineLength = scala.util.Try( + hadoopConf.getInt("mapreduce.input.linerecordreader.line.maxlength", Int.MaxValue)) + if (maxLineLength.map(_ != Int.MaxValue).getOrElse(true)) { + fallbackReasons += "Comet native Text scan does not support " + + "mapreduce.input.linerecordreader.line.maxlength" + } + // Spark passes `mapreduce.job.input.file.must.*` to the filesystem's open call as + // mandatory options, and the read fails if the filesystem cannot honor one. The native + // reader opens files itself and cannot enforce them, so fall back and let Spark decide. + if (hadoopConf + .iterator() + .asScala + .exists(_.getKey.startsWith("mapreduce.job.input.file.must."))) { + fallbackReasons += "Comet native Text scan does not support mandatory file-open " + + "options (mapreduce.job.input.file.must.*)" + } + // Spark reads a custom line separator with Hadoop's LineReader.readCustomLine, which on + // a mismatch backtracks only one byte instead of following the separator's own prefix + // overlap. For a separator that overlaps itself (e.g. "aab", whose "aa" prefix can + // restart a match) that reader can split differently from the native reader's exact + // substring scan, and the difference depends on where the separator falls across + // Hadoop's input buffer boundaries. Fall back so the two always agree. + textOptions.lineSeparatorInRead.foreach { sep => + if (CometScanRule.separatorSelfOverlaps(sep)) { + fallbackReasons += "Comet native Text scan does not support a self-overlapping " + + "line separator (Spark's line reader can split it differently)" + } + } + } + // A wholetext file is never split (TextScan.isSplitable is false for wholeText), so the + // split-file guard above does not bound it. Reading a very large file whole is unbounded + // memory, and the whole file becomes one Arrow string value. Because a lossy UTF-8 decode + // can expand ill-formed bytes up to 3x, cap the raw file size at Int.MaxValue/3 so the + // decoded value stays under Arrow's 2GB 32-bit string-offset limit; fall back above that. + // (The native reader also fails cleanly rather than panicking if a value still overflows.) + if (textOptions.wholeText) { + val maxWholeTextBytes = + math.min(session.sessionState.conf.filesMaxPartitionBytes, Int.MaxValue.toLong / 3) + val hasLargeWholeTextFile = scanExec.inputPartitions.exists { + case fp: FilePartition => fp.files.exists(_.length > maxWholeTextBytes) + case _ => false + } + if (hasLargeWholeTextFile) { + fallbackReasons += "Comet native Text scan does not support large wholetext files " + + s"(over ${SQLConf.FILES_MAX_PARTITION_BYTES.key}, capped below 2GB)" + } + } Review Comment: `spark.comet.scan.text.maxFileSize` (64 MiB) already bounds every file, wholetext included, so with the defaults this second cap never fires first. Tying it to `maxPartitionBytes`, a line-mode splitting setting, means a user who lowers that to 16 MiB for parallelism loses native wholetext for files well under 64 MiB. Would `math.min(maxTextFileSize, Int.MaxValue.toLong / 3)` as one check keep the 2 GiB protection without the extra pass and message? ########## native/core/src/execution/operators/text_scan.rs: ########## @@ -0,0 +1,796 @@ +// 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. + +//! Native Text file scan, mirroring Spark's `text` data source. +//! +//! Spark's text reader produces a single `value: string` column. With the default options each +//! line of the file is one row (split on `\n`, `\r\n`, or `\r`, with the terminator stripped). +//! With `wholetext=true` the whole file is a single row. A custom `lineSep` splits on exactly that +//! byte sequence instead of the universal-newline set. See the `TextOptions`/`HadoopFileLinesReader` +//! logic in Spark's `org.apache.spark.sql.execution.datasources.text` package. + +use crate::execution::operators::ExecutionError; +use arrow::array::StringBuilder; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; +use bytes::Bytes; +use datafusion::common::Result; +use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::physical_expr::projection::ProjectionExprs; +use datafusion::physical_plan::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricBuilder}; +use datafusion::physical_plan::DisplayFormatType; +use datafusion_comet_common::decode_utf8_spark_lossy; +use datafusion_comet_proto::spark_operator::TextOptions; +use datafusion_datasource::file::FileSource; +use datafusion_datasource::file_compression_type::FileCompressionType; +use datafusion_datasource::file_groups::FileGroup; +use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener}; +use datafusion_datasource::projection::{ProjectionOpener, SplitProjection}; +use datafusion_datasource::source::DataSourceExec; +use datafusion_datasource::{as_file_source, PartitionedFile, TableSchema}; +use futures::StreamExt; +use object_store::{ObjectStore, ObjectStoreExt}; +use std::borrow::Cow; +use std::fmt; +use std::io::Read; +use std::sync::Arc; + +/// Batch size used when the execution framework does not set one before opening the file. +const DEFAULT_TEXT_BATCH_SIZE: usize = 8192; + +/// UTF-8 byte-order mark. Spark's line reader strips a leading one; wholetext keeps it. +const UTF8_BOM: &[u8] = &[0xEF, 0xBB, 0xBF]; + +/// Cap on a batch's decoded `value` buffer bytes, kept well below Arrow's 2GB 32-bit offset limit +/// so a batch of many rows (or a few long lines) cannot overflow. Measured against the decoded +/// string length, since a lossy UTF-8 decode can expand ill-formed bytes up to 3x. +const MAX_BATCH_VALUE_BYTES: usize = 1 << 30; + +pub fn init_text_datasource_exec( + object_store_url: ObjectStoreUrl, + file_groups: Vec<Vec<PartitionedFile>>, + data_schema: SchemaRef, + _partition_schema: SchemaRef, + projection_vector: Vec<usize>, + text_options: &TextOptions, +) -> Result<Arc<DataSourceExec>, ExecutionError> { + let text_source = TextSource::new(data_schema, text_options); + + let file_groups = file_groups.into_iter().map(FileGroup::new).collect(); + + let file_scan_config = FileScanConfigBuilder::new(object_store_url, Arc::new(text_source)) + .with_file_groups(file_groups) + .with_projection_indices(Some(projection_vector))? + .build(); + + Ok(DataSourceExec::from_data_source(file_scan_config)) +} + +/// A [`FileSource`] for Spark's text format. The file schema is always the single `value` column; +/// projection (and partition columns, if any) is handled by [`ProjectionOpener`], exactly as the +/// CSV source does. +#[derive(Debug, Clone)] +struct TextSource { + table_schema: TableSchema, + projection: SplitProjection, + batch_size: Option<usize>, + metrics: ExecutionPlanMetricsSet, + whole_text: bool, + // The raw separator bytes when a custom `lineSep` is set; None means universal newline + // (`\n`, `\r\n`, `\r`). + line_sep: Option<Vec<u8>>, +} + +impl TextSource { + fn new(file_schema: SchemaRef, options: &TextOptions) -> Self { + let table_schema: TableSchema = file_schema.into(); + // `line_sep` arrives as the exact separator bytes Spark computed (`lineSeparatorInRead`, + // already encoded with the file's charset), so we split on them directly. + let line_sep = options.line_sep.clone().filter(|sep| !sep.is_empty()); + Self { + projection: SplitProjection::unprojected(&table_schema), + table_schema, + batch_size: None, + metrics: ExecutionPlanMetricsSet::new(), + whole_text: options.whole_text, + line_sep, + } + } +} + +impl From<TextSource> for Arc<dyn FileSource> { + fn from(source: TextSource) -> Self { + as_file_source(source) + } +} + +impl FileSource for TextSource { + fn create_file_opener( + &self, + object_store: Arc<dyn ObjectStore>, + base_config: &FileScanConfig, + partition_index: usize, + ) -> Result<Arc<dyn FileOpener>> { + // The inner opener must emit batches whose schema is the file schema projected by + // `file_indices` (that is what ProjectionOpener feeds its projector). For text that is + // either the single `value` column or, for a projection-less count, an empty schema. + let file_schema = self.table_schema.file_schema(); + let projected_file_schema = Arc::new(file_schema.project(&self.projection.file_indices)?); + let opener = Arc::new(TextOpener { + projected_file_schema, + file_compression_type: base_config.file_compression_type, + object_store, + partition_index, + batch_size: self.batch_size.unwrap_or(DEFAULT_TEXT_BATCH_SIZE), + metrics: self.metrics.clone(), + whole_text: self.whole_text, + line_sep: self.line_sep.clone(), + }) as Arc<dyn FileOpener>; + ProjectionOpener::try_new(self.projection.clone(), opener, file_schema) + } + + fn table_schema(&self) -> &TableSchema { + &self.table_schema + } + + fn with_batch_size(&self, batch_size: usize) -> Arc<dyn FileSource> { + let mut conf = self.clone(); + conf.batch_size = Some(batch_size); + Arc::new(conf) + } + + fn try_pushdown_projection( + &self, + projection: &ProjectionExprs, + ) -> Result<Option<Arc<dyn FileSource>>> { + let mut source = self.clone(); + let new_projection = self.projection.source.try_merge(projection)?; + source.projection = SplitProjection::new(self.table_schema.file_schema(), &new_projection); + Ok(Some(Arc::new(source))) + } + + fn projection(&self) -> Option<&ProjectionExprs> { + Some(&self.projection.source) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn file_type(&self) -> &str { + "text" + } + + // A custom line separator can straddle a byte-range boundary, and we read whole files rather + // than boundary-aligned ranges, so this source must not let DataFusion repartition files by + // byte range. + fn supports_repartitioning(&self) -> bool { + false + } + + fn fmt_extra(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + match t { + DisplayFormatType::Default | DisplayFormatType::Verbose => { + write!(f, ", whole_text={}", self.whole_text) + } + DisplayFormatType::TreeRender => Ok(()), + } + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc<dyn datafusion::physical_plan::PhysicalExpr>, + ) -> Result<datafusion::common::tree_node::TreeNodeRecursion>, + ) -> Result<datafusion::common::tree_node::TreeNodeRecursion> { + datafusion::physical_plan::apply_expression_roots(self.projection.source.iter(), f) + } +} + +/// A [`FileOpener`] that reads a whole text file and yields the projected file schema in batches. +struct TextOpener { + // The file schema projected by `file_indices`: the single `value` column, or an empty schema + // for a projection-less count. + projected_file_schema: SchemaRef, + file_compression_type: FileCompressionType, + object_store: Arc<dyn ObjectStore>, + partition_index: usize, + batch_size: usize, + metrics: ExecutionPlanMetricsSet, + whole_text: bool, + line_sep: Option<Vec<u8>>, +} + +impl FileOpener for TextOpener { + fn open(&self, partitioned_file: PartitionedFile) -> Result<FileOpenFuture> { + let store = Arc::clone(&self.object_store); + let projected_file_schema = Arc::clone(&self.projected_file_schema); + let file_compression_type = self.file_compression_type; + let batch_size = self.batch_size; + let whole_text = self.whole_text; + let line_sep = self.line_sep.clone(); + let baseline_metrics = BaselineMetrics::new(&self.metrics, self.partition_index); + // Feeds the task's input bytes in the Spark UI (see CometMetricNode.reportScanInputMetrics). + let bytes_scanned = + MetricBuilder::new(&self.metrics).counter("bytes_scanned", self.partition_index); + + Ok(Box::pin(async move { + let range = partitioned_file.range.clone(); + let location = partitioned_file.object_meta.location; + let raw = store.get(&location).await?.bytes().await?; + bytes_scanned.add(raw.len()); + + // Decompress only when the file is actually compressed. On Comet's plan config the file + // is always UNCOMPRESSED (CometScanRule falls back for compressed files), so we keep + // the fetched bytes as they are and avoid a full-file copy; the decompress arm is a + // defensive handler should compression ever be plumbed through. A byte range never + // coexists with decompression (compressed files are non-splittable). + // + // NOTE: the whole object is read into memory. CometScanRule restricts native text to + // unsplit files (it falls back when Spark splits a file into byte ranges) and caps the + // file size, so in practice this reads each small file once. Lines are found and + // batches built on demand, so only the bytes and the current batch are held. + let buf: Bytes = if file_compression_type.is_compressed() { + let mut decoder = file_compression_type.convert_read(std::io::Cursor::new(raw))?; + let mut decoded = Vec::new(); + decoder.read_to_end(&mut decoded)?; + Bytes::from(decoded) + } else { + raw + }; + + // Spark's line reader (Hadoop LineRecordReader.skipUtfByteOrderMark) strips a leading + // UTF-8 BOM from the first record at file offset 0, whatever the separator. wholetext + // keeps it. Native reads unsplit files only. + let at_file_start = range.as_ref().map(|r| r.start == 0).unwrap_or(true); + // Spark may split an uncompressed file into byte ranges (TextScan.isSplitable). A line + // belongs to the split whose [start, end) contains the line's start offset -- so each + // line is emitted by exactly one split, with no duplication or gaps across splits. See + // DataFusion's CsvOpener for the same rule. wholetext files are never split. + let byte_range = if whole_text { + None + } else { + range.map(|r| (r.start.max(0) as usize, r.end.max(0) as usize)) + }; + let reader = TextBatchReader::new( + buf, + projected_file_schema, + batch_size, + TextLayout { + whole_text, + line_sep, + strip_bom: !whole_text && at_file_start, + byte_range, + }, + ); + + Ok(futures::stream::unfold( + (reader, baseline_metrics), + |(mut reader, metrics)| async move { + // The timer covers only building this batch, not the wait for the next poll. + let next = { + let _timer = metrics.elapsed_compute().timer(); + reader.next_batch() + }; + match next { + Ok(Some(batch)) => Some((Ok(batch), (reader, metrics))), + Ok(None) => None, + Err(e) => { + // Report the error once, then end the stream. + reader.finished = true; + Some((Err(e), (reader, metrics))) + } + } + }, + ) + .boxed()) + })) + } +} + +/// How the bytes of one file split into rows. +struct TextLayout { + /// One row for the whole file, separators included. + whole_text: bool, + /// A custom line separator's exact bytes, or `None` for the universal newline set. + line_sep: Option<Vec<u8>>, + /// Strip a leading UTF-8 BOM from the first record (line mode, split at file offset 0). + strip_bom: bool, + /// Keep only the lines whose start offset is in `[start, end)`. + byte_range: Option<(usize, usize)>, +} + +/// Reads a file's bytes as batches of `value` rows, one batch per call. +/// +/// Lines are found by scanning forward from `pos`, so nothing is kept per line: the memory beyond +/// the file bytes is the batch being built. A file of nothing but newlines would otherwise need a +/// slice descriptor per byte. +struct TextBatchReader { + buf: Bytes, + schema: SchemaRef, + batch_rows: usize, + layout: TextLayout, + /// Offset of the next unread byte in line mode. + pos: usize, + /// wholetext: the single row has been returned. + whole_taken: bool, + /// The first-record BOM check has not run yet. + first_record: bool, + /// A row read but held back because it would have overflowed the current batch. + pending: Option<(usize, usize)>, + finished: bool, +} + +impl TextBatchReader { + fn new(buf: Bytes, schema: SchemaRef, batch_size: usize, layout: TextLayout) -> Self { + let first_record = layout.strip_bom; + Self { + buf, + schema, + batch_rows: batch_size.max(1), + layout, + pos: 0, + whole_taken: false, + first_record, + pending: None, + finished: false, + } + } + + /// Offsets `(start, end)` of the next line, terminator excluded, or `None` at the end. + /// + /// Default separators are `\n`, `\r\n` and `\r`. A custom separator is an exact byte match. A + /// trailing terminator does not start another, empty line (Hadoop LineRecordReader). + fn scan_line(&mut self) -> Option<(usize, usize)> { + let buf = &self.buf[..]; + if self.layout.whole_text { + if self.whole_taken { + return None; + } + self.whole_taken = true; + return Some((0, buf.len())); + } + if self.pos >= buf.len() { + return None; + } + let start = self.pos; + match &self.layout.line_sep { + None => match buf[start..].iter().position(|&b| b == b'\n' || b == b'\r') { + None => { + self.pos = buf.len(); + Some((start, buf.len())) + } + Some(i) => { + let end = start + i; + self.pos = if buf[end] == b'\r' && buf.get(end + 1) == Some(&b'\n') { + end + 2 + } else { + end + 1 + }; + Some((start, end)) + } + }, + Some(sep) => match find_subslice(&buf[start..], sep) { + None => { + self.pos = buf.len(); + Some((start, buf.len())) + } + Some(i) => { + let end = start + i; + self.pos = end + sep.len(); + Some((start, end)) + } + }, + } + } + + /// Offsets of the next row to emit, after the first-record BOM and the byte range apply. + fn next_row(&mut self) -> Option<(usize, usize)> { + if let Some(row) = self.pending.take() { + return Some(row); + } + loop { + let (start, end) = self.scan_line()?; + let mut row_start = start; + // Hadoop reads the first record and then drops the BOM from its text. When that + // leaves nothing and the record was not terminated (the file is only the BOM), it + // counts as end of input, so the file has no rows. A BOM followed by a separator + // leaves an empty row. A custom `lineSep` equal to the BOM never reaches the strip, + // because the BOM is then the separator and the first record does not start with it. + if self.first_record { + self.first_record = false; + if self.buf[start..end].starts_with(UTF8_BOM) { + if self.buf.len() == UTF8_BOM.len() { + return None; + } + row_start = start + UTF8_BOM.len(); + } + } + if let Some((range_start, range_end)) = self.layout.byte_range { + if start >= range_end { + // Lines come in increasing offset order, so the rest are outside the range. + self.pos = self.buf.len(); + return None; + } + if start < range_start { + continue; + } + } + return Some((row_start, end)); + } + } + + /// The next batch of up to `batch_size` rows, or `None` when the file is exhausted. + fn next_batch(&mut self) -> Result<Option<RecordBatch>> { + if self.finished { + return Ok(None); + } + if self.schema.fields().is_empty() { + // The count(*) case: rows with no columns. Count lines without decoding them. + let mut rows = 0; + while rows < self.batch_rows && self.next_row().is_some() { + rows += 1; + } + if rows == 0 { + self.finished = true; + return Ok(None); + } + let options = RecordBatchOptions::new().with_row_count(Some(rows)); + return Ok(Some(RecordBatch::try_new_with_options( + Arc::clone(&self.schema), + vec![], + &options, + )?)); + } + + // Decode each line once (decode_utf8_spark_lossy borrows for valid UTF-8, allocates only + // for ill-formed input) and accumulate against the DECODED size, since a lossy decode can + // expand ill-formed bytes up to 3x. Start a new batch before the value buffer would reach + // Arrow's 32-bit offset limit. Always take at least one line so progress is guaranteed; + // if a single decoded value alone would overflow, fail cleanly rather than panic in Arrow. + // A handle to the bytes lets the decoded values borrow them while the reader advances. + let buf = self.buf.clone(); + let mut decoded: Vec<Cow<str>> = Vec::new(); + let mut data_bytes = 0usize; + while decoded.len() < self.batch_rows { + let (start, end) = match self.next_row() { + Some(row) => row, + None => break, + }; + let value = decode_utf8_spark_lossy(&buf[start..end]); + let len = value.len(); + if len > i32::MAX as usize { + return Err(datafusion::common::DataFusionError::Execution(format!( + "Text value of {len} bytes exceeds Arrow's 2GB string-offset limit" + ))); + } + if !decoded.is_empty() && data_bytes + len > MAX_BATCH_VALUE_BYTES { + self.pending = Some((start, end)); + break; + } + data_bytes += len; + decoded.push(value); + } + if decoded.is_empty() { + self.finished = true; + return Ok(None); + } + let mut builder = StringBuilder::with_capacity(decoded.len(), data_bytes); + for value in &decoded { + builder.append_value(value); + } + Ok(Some(RecordBatch::try_new( + Arc::clone(&self.schema), + vec![Arc::new(builder.finish())], + )?)) Review Comment: `decoded` collects a `Cow<str>` per row and the loop at the end copies each one into the `StringBuilder`, so each batch is walked twice, and `Vec::new()` regrows up to `batch_rows`, which is known. Could the loop append straight into a `StringBuilder::with_capacity(self.batch_rows, ..)` and track the row count and `data_bytes` itself? The overflow check only needs the length of the next decoded value. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
