feat: add sampling functionality to fsspec indexers (#189)

* feat: add sampling functionality to fsspec indexers

* add e2e test

* tidy

* release

Author: ahmetmeleq

CHANGELOG.md

@@ -1,9 +1,10 @@
-## 0.1.1-dev1
+## 0.1.1
 
 ### Enhancements
 
 * **Update KDB.AI vectorstore integration to 1.4**
 * **Add sqlite and postgres source connectors**
+* **Add sampling functionality for indexers in fsspec connectors**
 
 ## 0.1.0
 

test/unit/connector/fsspec/test_sampling.py

@@ -0,0 +1,49 @@
+from unstructured_ingest.v2.processes.connectors.fsspec.fsspec import (
+    FsspecIndexer,
+)
+
+
+def test_fsspec_indexer_sampling_happy_path():
+
+    indexer = FsspecIndexer(
+        connection_config="fake_connection_config",
+        index_config="fake_index_config",
+        connector_type="fake_connector_type",
+    )
+
+    all_files = [{"name": "fake_file.txt"}, {"name": "fake_file2.txt"}, {"name": "fake_file3.txt"}]
+
+    sampled_files = indexer.sample_n_files(all_files, 2)
+    assert len(sampled_files) == 2
+    for sampled_file in sampled_files:
+        assert type(sampled_file) == dict  # noqa: E721
+        assert sampled_file["name"] in [file["name"] for file in all_files]
+
+
+def test_fsspec_indexer_sampling_no_files():
+    indexer = FsspecIndexer(
+        connection_config="fake_connection_config",
+        index_config="fake_index_config",
+        connector_type="fake_connector_type",
+    )
+
+    all_files = []
+
+    sampled_files = indexer.sample_n_files(all_files, 2)
+    assert len(sampled_files) == 0
+
+
+def test_fsspec_indexer_sampling_bigger_sample_size():
+    indexer = FsspecIndexer(
+        connection_config="fake_connection_config",
+        index_config="fake_index_config",
+        connector_type="fake_connector_type",
+    )
+
+    all_files = [{"name": "fake_file.txt"}, {"name": "fake_file2.txt"}, {"name": "fake_file3.txt"}]
+
+    sampled_files = indexer.sample_n_files(all_files, 10)
+    assert len(sampled_files) == 3
+    for sampled_file in sampled_files:
+        assert type(sampled_file) == dict  # noqa: E721
+        assert sampled_file["name"] in [file["name"] for file in all_files]

test_e2e/src/s3-sample.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+
+set -e
+
+SRC_PATH=$(dirname "$(realpath "$0")")
+SCRIPT_DIR=$(dirname "$SRC_PATH")
+cd "$SCRIPT_DIR"/.. || exit 1
+OUTPUT_FOLDER_NAME=s3-sample
+OUTPUT_ROOT=${OUTPUT_ROOT:-$SCRIPT_DIR}
+OUTPUT_DIR=$OUTPUT_ROOT/structured-output/$OUTPUT_FOLDER_NAME
+WORK_DIR=$OUTPUT_ROOT/workdir/$OUTPUT_FOLDER_NAME
+DOWNLOAD_DIR=$SCRIPT_DIR/download/$OUTPUT_FOLDER_NAME
+max_processes=${MAX_PROCESSES:=$(python3 -c "import os; print(os.cpu_count())")}
+
+# shellcheck disable=SC1091
+source "$SCRIPT_DIR"/cleanup.sh
+# shellcheck disable=SC2317
+function cleanup() {
+  cleanup_dir "$OUTPUT_DIR"
+  cleanup_dir "$WORK_DIR"
+}
+trap cleanup EXIT
+
+RUN_SCRIPT=${RUN_SCRIPT:-./unstructured_ingest/main.py}
+PYTHONPATH=${PYTHONPATH:-.} "$RUN_SCRIPT" \
+  s3 \
+  --api-key "$UNS_PAID_API_KEY" \
+  --partition-by-api \
+  --partition-endpoint "https://api.unstructuredapp.io" \
+  --num-processes "$max_processes" \
+  --download-dir "$DOWNLOAD_DIR" \
+  --metadata-exclude coordinates,filename,file_directory,metadata.data_source.date_processed,metadata.last_modified,metadata.detection_class_prob,metadata.parent_id,metadata.category_depth \
+  --strategy fast \
+  --preserve-downloads \
+  --reprocess \
+  --output-dir "$OUTPUT_DIR" \
+  --verbose \
+  --remote-url s3://utic-dev-tech-fixtures/small-pdf-set/ \
+  --anonymous \
+  --work-dir "$WORK_DIR" \
+  --sample-n-files 3
+
+NUM_FILES=$(find "$OUTPUT_DIR" -type f | wc -l)
+EXPECTED_NUM_FILES=3
+
+if [ "$NUM_FILES" -ne "$EXPECTED_NUM_FILES" ]; then
+  exit 1
+fi

unstructured_ingest/__version__.py

@@ -1 +1 @@
-__version__ = "0.1.1-dev1"  # pragma: no cover
+__version__ = "0.1.1"  # pragma: no cover

unstructured_ingest/v2/processes/connectors/fsspec/fsspec.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import random
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Generator, Optional, TypeVar
@@ -63,6 +64,7 @@ def __init__(self, **data):
 
 class FsspecIndexerConfig(FileConfig, IndexerConfig):
     recursive: bool = False
+    sample_n_files: Optional[int] = None
 
 
 class FsspecAccessConfig(AccessConfig):
@@ -128,8 +130,23 @@ def get_file_data(self) -> list[dict[str, Any]]:
         filtered_files = [
             file for file in files if file.get("size") > 0 and file.get("type") == "file"
         ]
+
+        if self.index_config.sample_n_files:
+            filtered_files = self.sample_n_files(filtered_files, self.index_config.sample_n_files)
+
         return filtered_files
 
+    def sample_n_files(self, files: list[dict[str, Any]], n) -> list[dict[str, Any]]:
+        if len(files) <= n:
+            logger.warning(
+                f"number of files to be sampled={n} is not smaller than the number"
+                f" of files found ({len(files)}). Returning all of the files as the"
+                " sample."
+            )
+            return files
+
+        return random.sample(files, n)
+
     def get_metadata(self, file_data: dict) -> FileDataSourceMetadata:
         raise NotImplementedError()