Skip to content

Commit 35c0076

Browse files
mykolamelnykmlmykolamelnykml
authored
Merge pull request #86 from StabRise/82-add-chunker-transformer
82 add chunker transformer
2 parents f21e840 + ad4d63d commit 35c0076

22 files changed

Lines changed: 2050 additions & 1 deletion

‎docs/source/index.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,8 @@ Benefits of using ScaleDP
4747
detectors.md
4848
ocr.md
4949
embeddings.md
50+
splitters.md
51+
schemas.md
5052
show_utils.md
5153
release_notes.md
5254

‎docs/source/models/splitters/base_splitter.md‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
(BaseSplitter)=
2+
# BaseSplitter
3+
4+
## Overview
5+
6+
`BaseSplitter` is the abstract base class for all text splitter transformers in the ScaleDP library. It extends PySpark's `Transformer` class and provides common functionality for splitting documents into chunks. This class defines the interface and shared parameters for all splitter implementations.
7+
8+
## Inheritance
9+
10+
- Extends PySpark's `Transformer` for ML pipeline compatibility.
11+
- Mixes in the following parameter mixins:
12+
- `HasInputCol` - Input column containing documents
13+
- `HasOutputCol` - Output column for results
14+
- `HasKeepInputData` - Whether to preserve input data
15+
- `HasChunkSize` - Maximum chunk size
16+
- `HasChunkOverlap` - Overlap between chunks
17+
- `HasNumPartitions` - Partition control
18+
- `HasPartitionMap` - Enable distributed processing mode
19+
- `HasWhiteList` - Whitelist filtering support
20+
21+
## Key Features
22+
23+
- **PySpark Integration**: Full compatibility with PySpark ML pipelines
24+
- **Serialization**: Support for reading and writing model parameters
25+
- **Flexible Configuration**: Extensive parameters for customization
26+
- **Extensible Design**: Foundation for specialized splitter implementations
27+
- **Batch Processing**: Support for both local and distributed processing modes
28+
29+
## Class Hierarchy
30+
31+
```
32+
BaseSplitter
33+
├── BaseTextSplitter
34+
│ └── TextSplitter (concrete implementation)
35+
```
36+
37+
## Parameters
38+
39+
| Parameter | Type | Description | Default |
40+
|-------------------|---------|--------------------------------------------------|-----------------------------|
41+
| inputCol | str | Input column name | varies by implementation |
42+
| outputCol | str | Output column name | varies by implementation |
43+
| keepInputData | bool | Keep input columns in output | True |
44+
| chunkSize | int | Size of each chunk | 500 |
45+
| chunkOverlap | int | Overlap between consecutive chunks | 0 |
46+
| numPartitions | int | Number of partitions | 1 |
47+
| partitionMap | bool | Use partitioned mapping (pandas_udf mode) | False |
48+
| whiteList | list | Whitelist of allowed items | [] |
49+
50+
## Abstract Methods
51+
52+
Subclasses must implement the following abstract methods:
53+
54+
### transform(dataset)
55+
Transforms a Spark DataFrame by applying the splitter logic.
56+
57+
**Parameters:**
58+
- `dataset` (pyspark.sql.DataFrame): Input DataFrame
59+
60+
**Returns:**
61+
- (pyspark.sql.DataFrame): DataFrame with split results
62+
63+
## Usage Guidelines
64+
65+
`BaseSplitter` is an abstract class and should not be instantiated directly. Instead, use concrete implementations like:
66+
67+
- [`TextSplitter`](./text_splitter.md) - Semantic text splitting
68+
69+
```python
70+
# Correct: Use concrete implementation
71+
from scaledp.models.splitters.TextSplitter import TextSplitter
72+
73+
splitter = TextSplitter(chunk_size=500, chunk_overlap=50)
74+
```
75+
76+
```python
77+
# Incorrect: Do not instantiate BaseSplitter directly
78+
from scaledp.models.splitters.BaseSplitter import BaseSplitter
79+
80+
# This will raise an error
81+
splitter = BaseSplitter() # Error!
82+
```
83+
84+
## Creating Custom Splitters
85+
86+
To create a custom splitter, inherit from `BaseSplitter` or `BaseTextSplitter`:
87+
88+
```python
89+
from scaledp.models.splitters.BaseTextSplitter import BaseTextSplitter
90+
from scaledp.schemas.Document import Document
91+
from scaledp.schemas.TextChunks import TextChunks
92+
93+
class CustomSplitter(BaseTextSplitter):
94+
"""Custom splitter implementation."""
95+
96+
def split(self, document: Document) -> TextChunks:
97+
"""Implement custom splitting logic."""
98+
# Your splitting algorithm here
99+
chunks = self._split_text(document.text)
100+
return TextChunks(
101+
path=document.path,
102+
chunks=chunks,
103+
exception="",
104+
processing_time=0.0
105+
)
106+
```
107+
108+
## Pipeline Integration
109+
110+
`BaseSplitter` and its subclasses are designed to work seamlessly with PySpark pipelines:
111+
112+
```python
113+
from pyspark.ml import Pipeline
114+
from scaledp.models.splitters.TextSplitter import TextSplitter
115+
116+
# Create pipeline stages
117+
splitter = TextSplitter(chunk_size=500)
118+
119+
# Create and fit pipeline
120+
pipeline = Pipeline(stages=[splitter])
121+
model = pipeline.fit(training_data)
122+
123+
# Transform data
124+
results = model.transform(test_data)
125+
```
126+
127+
## Serialization
128+
129+
All splitters support PySpark's read/write functionality:
130+
131+
```python
132+
# Save a model
133+
splitter = TextSplitter(chunk_size=500)
134+
splitter.write().overwrite().save("path/to/splitter")
135+
136+
# Load a model
137+
loaded_splitter = TextSplitter.load("path/to/splitter")
138+
```
139+
140+
## Related Classes
141+
142+
- [`BaseTextSplitter`](./base_text_splitter.md) - Abstract base for text splitters
143+
- [`TextSplitter`](./text_splitter.md) - Concrete semantic text splitter implementation
144+
- [`Document`](#Document) - Input document schema
145+
- [`TextChunks`](#TextChunks) - Output text chunks schema
146+
147+
## See Also
148+
149+
- [PySpark ML Transformers](https://spark.apache.org/docs/latest/ml-pipeline.html)
150+
- [Transformer API](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.ml.Transformer.html)

‎docs/source/models/splitters/base_text_splitter.md‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
(BaseTextSplitter)=
2+
# BaseTextSplitter
3+
4+
## Overview
5+
6+
`BaseTextSplitter` is an abstract base class for text splitting transformers in PySpark. It provides common functionality for splitting documents into chunks while preserving metadata like file paths and document types. It is designed for extensibility and serves as the foundation for concrete text splitting implementations like [`TextSplitter`](./text_splitter.md).
7+
8+
The splitter operates on **Document struct columns**, which contain structured data including text content, file path, document type, and bounding boxes.
9+
10+
## Inheritance
11+
12+
- Inherits from [`BaseSplitter`](./base_splitter.md), which provides core Spark ML transformer functionality and schema handling.
13+
- Mixes in `HasColumnValidator` and `HasDefaultEnum` for validation and enumeration support.
14+
- Extends `DefaultParamsReadable` and `DefaultParamsWritable` for serialization support.
15+
16+
## Key Features
17+
18+
- **Document-Centric**: Works with Document struct columns containing path, text, type, and bboxes
19+
- **Metadata Preservation**: Maintains document metadata (path, document type) through the splitting process
20+
- **Flexible Chunking**: Configurable chunk size and overlap for text splitting
21+
- **Distributed Processing**: Supports both regular UDF and pandas_udf (partitionMap) modes for Spark batch processing
22+
- **Error Handling**: Captures and reports processing exceptions in output
23+
24+
## Usage Example
25+
26+
```python
27+
from scaledp.models.splitters.TextSplitter import TextSplitter
28+
from scaledp.schemas.Document import Document
29+
30+
# Create a splitter with custom parameters
31+
splitter = TextSplitter(
32+
inputCol="document", # Column containing Document structs
33+
outputCol="chunks", # Output column for TextChunks
34+
chunk_size=500, # Characters per chunk
35+
chunk_overlap=50, # Character overlap between chunks
36+
)
37+
38+
# Use in a Spark pipeline
39+
result_df = splitter.transform(input_df)
40+
```
41+
42+
## Parameters
43+
44+
| Parameter | Type | Description | Default |
45+
|-------------------|---------|--------------------------------------------------|-----------------------------|
46+
| inputCol | str | Input Document struct column | "document" |
47+
| outputCol | str | Output column for TextChunks results | "chunks" |
48+
| keepInputData | bool | Keep input document column in output | True |
49+
| chunk_size | int | Size of each chunk in characters | 500 |
50+
| chunk_overlap | int | Number of characters to overlap between chunks | 0 |
51+
| numPartitions | int | Number of partitions for coalescing | 1 |
52+
| partitionMap | bool | Use pandas_udf for distributed processing | False |
53+
54+
## Input Schema
55+
56+
The input column must contain **Document struct**. For detailed schema information, see [Document Schema Documentation](../../schemas/document.md).
57+
58+
**Key Fields:**
59+
- `path` - File path or document identifier
60+
- `text` - Text content to split
61+
- `type` - Document type (e.g., "text", "pdf")
62+
- `bboxes` - Bounding boxes (empty for text documents)
63+
- `exception` - Error message if any (optional)
64+
65+
## Output Schema
66+
67+
The output column contains **TextChunks struct**. For detailed schema information, see [TextChunks Schema Documentation](../../schemas/text_chunks.md).
68+
69+
**Key Fields:**
70+
- `path` - Original document path
71+
- `chunks` - List of text chunks
72+
- `exception` - Error message if splitting failed
73+
- `processing_time` - Time taken to split document (seconds)
74+
75+
## Notes
76+
77+
- The splitter is abstract and cannot be instantiated directly. Use concrete implementations like `TextSplitter`.
78+
- Input documents must contain text and path information in the Document struct format.
79+
- Chunk overlap can help maintain context between chunks for semantic meaning.
80+
- The `partitionMap` option enables pandas_udf mode for better performance on large datasets but requires careful configuration.
81+
- All errors during splitting are captured and reported in the `exception` field of the output.
82+
83+
## Related Classes
84+
85+
- [`TextSplitter`](./text_splitter.md) - Concrete implementation using semantic text splitting
86+
- [`BaseSplitter`](./base_splitter.md) - Base transformer for all splitter implementations
87+
- [`Document`](../../schemas/document.md) - Input schema class
88+
- [`TextChunks`](../../schemas/text_chunks.md) - Output schema class
89+
- [`Box`](../../schemas/box.md) - Bounding box schema class

0 commit comments

Comments
 (0)