sunchao commented on code in PR #6784: URL: https://github.com/apache/datafusion-comet/pull/6784#discussion_r4458633621
########## 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) { Review Comment: The exact-byte search differs from Hadoop's custom-delimiter handling across an input-buffer boundary. I wrote the valid UTF-8 bytes `b'x' * 4092 + b'aaaaabTAIL'`, used `lineSep='aab'` and a 4096-byte Hadoop buffer, and read the same small, unsplit file through Spark's V2 text reader. Spark returned one row of length `4102`; `scan_line` and `next_row` compiled directly from this head returned two rows of lengths `4095` and `4`. The analogous case at the 64 KiB boundary has the same discrepancy. There is no BOM, compression, invalid UTF-8 or partition-column complication. Could the reader preserve Hadoop's delimiter behavior, or fall back for cases it cannot match? The current string search admits this file but changes both the records and the row count. This reproduction executes the Spark reader and the exact Rust parsing helpers separately; I did not run a complete native scan because the locked dependencies were unavailable locally. -- 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]
