diff --git a/materializationengine/blueprints/upload/api.py b/materializationengine/blueprints/upload/api.py index 75c5b114..c9c24dc7 100644 --- a/materializationengine/blueprints/upload/api.py +++ b/materializationengine/blueprints/upload/api.py @@ -42,7 +42,14 @@ from materializationengine.blueprints.upload.schema_helper import get_schema_types from materializationengine.blueprints.upload.schemas import UploadRequestSchema from materializationengine.blueprints.upload.gcs_processor import GCSCsvProcessor -from materializationengine.blueprints.upload.source import SourceAccessError +from materializationengine.blueprints.upload.parquet_source import ( + is_parquet_source, + preview_parquet, +) +from materializationengine.blueprints.upload.source import ( + SourceAccessError, + parse_source_uri, +) from materializationengine.blueprints.upload.storage import ( StorageConfig, StorageService, @@ -351,9 +358,11 @@ def generate_presigned_url(datastack_name: str): @upload_bp.route("/api/source/inspect/", methods=["POST"]) @auth_requires_permission("edit", table_arg="datastack_name") def inspect_source(datastack_name: str): - """Check that a gs:// or s3:// CSV can be read, and return its size and first rows. + """Check that a gs:// or s3:// source can be read; return its size and first rows. - The wizard calls this instead of uploading a file, when the CSV is read in place. + The source is a CSV file, or parquet: a file, a folder of files (hive-partitioned + or not) or a Delta table. The wizard calls this instead of uploading a file, when + the data is read in place. """ data = request.get_json(silent=True) or {} uri = (data.get("uri") or "").strip() @@ -363,7 +372,24 @@ def inspect_source(datastack_name: str): try: processor = GCSCsvProcessor(bucket_name) if "://" not in uri: - raise ValueError("Enter a location such as gs://bucket/path.csv or s3://bucket/path.csv") + raise ValueError( + "Enter a location such as gs://bucket/path.csv, or a parquet file or folder" + ) + if is_parquet_source(uri): + parse_source_uri(uri) + info = preview_parquet(uri) + return jsonify( + { + "status": "ok", + "uri": uri, + "filename": info["name"], + "size": info["total_bytes"], + "previewRows": info["preview_rows"], + "kind": info["kind"], + "files": info["files"], + "rows": info["total_rows"], + } + ) source = processor.source(uri) filename = uri.rstrip("/").rsplit("/", 1)[-1] if not filename.lower().endswith(".csv"): diff --git a/materializationengine/blueprints/upload/gcs_processor.py b/materializationengine/blueprints/upload/gcs_processor.py index 644e8918..9e12de2a 100644 --- a/materializationengine/blueprints/upload/gcs_processor.py +++ b/materializationengine/blueprints/upload/gcs_processor.py @@ -295,6 +295,38 @@ def process_csv_range( "lines": reader.lines(), } + def process_parquet_part( + self, + source_uri: str, + root: str, + pieces: List[List[Any]], + destination_blob_name: str, + transform_function: Callable[[pd.DataFrame], pd.DataFrame], + ) -> Dict[str, Any]: + """Transform a parquet part's row groups and write them to one CSV object.""" + from materializationengine.blueprints.upload.parquet_source import ( + iter_parquet_part, + ) + + destination_blob = self.bucket.blob(destination_blob_name) + rows_read = 0 + rows_written = 0 + with destination_blob.open( + "wb", content_type="text/csv", ignore_flush=True + ) as gcs_writer_stream: + for chunk_df in iter_parquet_part( + source_uri, root, pieces, self.pandas_chunk_size + ): + rows_read += len(chunk_df) + processed_chunk_df = transform_function(chunk_df) + chunk_bytes = processed_chunk_df.to_csv(index=False, header=False).encode( + "utf-8" + ) + if chunk_bytes: + gcs_writer_stream.write(chunk_bytes) + rows_written += len(processed_chunk_df) + return {"rows_read": rows_read, "rows_written": rows_written} + def count_csv_range_lines(self, source_blob_name: str, start: int, end: int) -> int: """The number of lines in bytes [start, end) of an object.""" with self.source(source_blob_name).open_range(start, end) as blob_stream: diff --git a/materializationengine/blueprints/upload/parquet_source.py b/materializationengine/blueprints/upload/parquet_source.py new file mode 100644 index 00000000..a67ec1fc --- /dev/null +++ b/materializationengine/blueprints/upload/parquet_source.py @@ -0,0 +1,402 @@ +"""Uploads read from parquet: one file, a folder of files, or a Delta Lake table. + +A folder is read recursively. Hive-style directories (``key=value/``) become columns +with those values, as pyarrow and Spark read them. Entries whose name starts with ``_`` +or ``.`` are skipped (``_delta_log``, ``_SUCCESS``, ``.crc`` files). + +A folder holding a ``_delta_log`` is read as a Delta table: only the files in the +table's current version, since files an optimize or overwrite replaced stay on disk +until a vacuum and would otherwise be read twice. + +The work is split into parts of whole row groups from the parquet footers, so every +part knows its row count before it runs. Generated ids then need no counting pass. +""" + +import functools +import io +import os +import posixpath +import struct +from urllib.parse import quote +from typing import Any, Dict, Iterator, List, Optional, Tuple +from urllib.parse import unquote, urlparse + +import pandas as pd +import pyarrow.fs as pafs +import pyarrow.parquet as pq + +from materializationengine.blueprints.upload.source import SourceAccessError + +# Rows per parquet part task: about the work of a 128 MB CSV part, about a minute. +DEFAULT_PARQUET_PART_ROWS = 1_500_000 +FOOTER_THREADS = 32 +# One request of a file's last bytes holds the footer of all but very wide files. +FOOTER_TAIL_BYTES = 64 * 1024 + + +def is_parquet_source(uri: str) -> bool: + """True for a .parquet file or a folder, False for a .csv file. + + A name with another extension is refused; end the location with "/" to read a + folder whose name has a dot in it. + """ + raw_path = urlparse((uri or "").strip()).path + name = posixpath.basename(raw_path.rstrip("/")).lower() + if raw_path.endswith("/") or name.endswith(".parquet") or "." not in name: + return True + if name.endswith(".csv"): + return False + raise ValueError( + f"{name!r} is not a .csv or .parquet file; to read a folder, end the location with /" + ) + + +def _split(uri: str) -> Tuple[str, str, str]: + parsed = urlparse(uri.strip()) + scheme = parsed.scheme.lower() + if scheme == "file": + # local paths, for tests; the upload API only accepts gs:// and s3:// + return "file", "", parsed.path.rstrip("/") or "/" + if scheme not in ("gs", "s3"): + raise ValueError( + f"Unsupported location {uri!r}: use gs://bucket/path or s3://bucket/path" + ) + if not parsed.netloc or not parsed.path.strip("/"): + raise ValueError(f"{uri!r} does not name a file or folder") + return scheme, parsed.netloc, parsed.path.strip("/") + + +def filesystem_for(uri: str) -> Tuple[pafs.FileSystem, str]: + """A pyarrow filesystem that can read uri, and uri's path within it.""" + scheme, bucket, key = _split(uri) + if scheme == "file": + return pafs.LocalFileSystem(), key + if scheme == "s3": + return _s3_filesystem(bucket), f"{bucket}/{key}" + return _gcs_filesystem(), f"{bucket}/{key}" + + +def _filesystem_path(uri: str) -> str: + """The path of a file URI (as delta-rs lists them) within filesystem_for.""" + if "://" not in uri: + return uri + scheme, bucket, key = _split(uri) + return key if scheme == "file" else f"{bucket}/{key}" + + +@functools.lru_cache(maxsize=32) +def _s3_filesystem(bucket: str) -> pafs.S3FileSystem: + try: + region = pafs.resolve_s3_region(bucket) + except Exception: + region = "us-east-1" + return pafs.S3FileSystem(anonymous=True, region=region) + + +def _gcs_filesystem() -> pafs.GcsFileSystem: + """GCS with the deployment's Google credentials, or anonymous without any.""" + try: + import google.auth + import google.auth.transport.requests + + credentials, _ = google.auth.default( + scopes=["https://www.googleapis.com/auth/devstorage.read_only"] + ) + credentials.refresh(google.auth.transport.requests.Request()) + return pafs.GcsFileSystem( + access_token=credentials.token, + credential_token_expiration=credentials.expiry, + ) + except Exception: + return pafs.GcsFileSystem(anonymous=True) + + +def _storage_options(uri: str) -> Dict[str, str]: + """Credentials for deltalake (delta-rs), matching filesystem_for.""" + scheme, bucket, _ = _split(uri) + if scheme == "file": + return {} + if scheme == "s3": + try: + region = pafs.resolve_s3_region(bucket) + except Exception: + region = "us-east-1" + return {"aws_skip_signature": "true", "aws_region": region} + try: + import google.auth + import google.auth.transport.requests + + credentials, _ = google.auth.default( + scopes=["https://www.googleapis.com/auth/devstorage.read_only"] + ) + credentials.refresh(google.auth.transport.requests.Request()) + return {"bearer_token": credentials.token} + except Exception: + return {} + + +def _tail_reader(uri: str): + """A function (path, n) -> the last n bytes of a file, for reading footers. + + pyarrow's GCS filesystem serialized concurrent footer reads (315 files of a Delta + export: 151 s); one suffix-range request each, 32 at a time, took 3 s. + """ + scheme, bucket, _ = _split(uri) + if scheme == "file": + def tail(path, n): + with open(path, "rb") as f: + f.seek(0, os.SEEK_END) + f.seek(max(f.tell() - n, 0)) + return f.read() + return tail + if scheme == "s3": + from materializationengine.blueprints.upload.source import anonymous_s3_client + + client = anonymous_s3_client(bucket) + + def tail(path, n): + key = path.split("/", 1)[1] + return client.get_object(Bucket=bucket, Key=key, Range=f"bytes=-{n}")["Body"].read() + return tail + + import requests + + try: + import google.auth + from google.auth.transport.requests import AuthorizedSession + + credentials, _ = google.auth.default( + scopes=["https://www.googleapis.com/auth/devstorage.read_only"] + ) + session = AuthorizedSession(credentials) + except Exception: + session = requests.Session() # public objects read anonymously + + def tail(path, n): + key = path.split("/", 1)[1] + url = ( + f"https://storage.googleapis.com/download/storage/v1/b/{bucket}/o/" + f"{quote(key, safe='')}?alt=media" + ) + response = session.get(url, headers={"Range": f"bytes=-{n}"}, timeout=60) + if response.status_code in (401, 403): + raise SourceAccessError(f"Permission denied reading gs://{path}") + response.raise_for_status() + return response.content + return tail + + +def read_footer(tail, path: str) -> pq.FileMetaData: + """A parquet file's metadata, from the end of the file only.""" + data = tail(path, FOOTER_TAIL_BYTES) + if len(data) < 12 or data[-4:] != b"PAR1": + raise ValueError(f"{path} is not a parquet file") + footer_length = struct.unpack(" len(data): + data = tail(path, footer_length + 8) + # pyarrow checks only the magic bytes at either end; offsets inside are not read + return pq.ParquetFile(io.BytesIO(b"PAR1" + data[-(footer_length + 8):])).metadata + + +def _hidden(name: str) -> bool: + return name.startswith("_") or name.startswith(".") + + +def list_parquet_files(uri: str) -> Dict[str, Any]: + """The parquet files of uri, in a stable order, and what kind of source it is. + + Returns {"kind": "file" | "folder" | "delta", "root": path, "files": [paths], + "partition_columns": [names] (delta only)}. Paths are within filesystem_for(uri). + """ + fs, root = filesystem_for(uri) + try: + info = fs.get_file_info(root) + except OSError as e: + raise SourceAccessError(f"Cannot read {uri}: {e}") from e + if info.type == pafs.FileType.File: + return {"kind": "file", "root": posixpath.dirname(root), "files": [root]} + if info.type != pafs.FileType.Directory: + raise FileNotFoundError(f"{uri} does not exist") + + if fs.get_file_info(f"{root}/_delta_log").type == pafs.FileType.Directory: + from deltalake import DeltaTable + + table_uri = root if _split(uri)[0] == "file" else uri.rstrip("/") + try: + table = DeltaTable(table_uri, storage_options=_storage_options(uri)) + except Exception as e: + raise SourceAccessError(f"Cannot read the Delta table at {uri}: {e}") from e + files = sorted(_filesystem_path(u) for u in table.file_uris()) + return { + "kind": "delta", + "root": root, + "files": files, + "partition_columns": list(table.metadata().partition_columns), + "version": table.version(), + } + + try: + entries = fs.get_file_info(pafs.FileSelector(root, recursive=True)) + except OSError as e: + raise SourceAccessError(f"Cannot list {uri}: {e}") from e + files = sorted( + e.path + for e in entries + if e.type == pafs.FileType.File + and e.path.endswith(".parquet") + and not any(_hidden(part) for part in posixpath.relpath(e.path, root).split("/")) + ) + if not files: + raise FileNotFoundError(f"No parquet files under {uri}") + return {"kind": "folder", "root": root, "files": files} + + +def hive_partition_values(path: str, root: str) -> Dict[str, Any]: + """{key: value} from the key=value directories between root and the file.""" + values = {} + for part in posixpath.relpath(posixpath.dirname(path), root).split("/"): + if "=" in part: + key, _, raw = part.partition("=") + values[unquote(key)] = _typed(unquote(raw)) + return values + + +def _typed(raw: str): + if raw == "__HIVE_DEFAULT_PARTITION__": + return None + for cast in (int, float): + try: + return cast(raw) + except ValueError: + pass + return raw + + +def plan_parquet_parts( + uri: str, target_rows: int = DEFAULT_PARQUET_PART_ROWS +) -> Dict[str, Any]: + """Split a parquet source into parts of about target_rows each. + + A part is whole row groups, or a slice of one row group bigger than target_rows + (some writers put a whole file in one row group): a part must stay a short task. + + Returns {"kind", "root", "columns", "total_rows", "total_bytes", "files", + "parts": [{"pieces": [piece, ...], "rows": n}, ...]}, where a piece is + [path, [row group, ...]] or [path, [row group], first row, row count]. + """ + listing = list_parquet_files(uri) + root = listing["root"] + tail = _tail_reader(uri) + + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=FOOTER_THREADS) as pool: + footers = list(pool.map(lambda path: read_footer(tail, path), listing["files"])) + + first_partitions = hive_partition_values(listing["files"][0], root) + file_columns = list(footers[0].schema.to_arrow_schema().names) + columns = file_columns + [c for c in first_partitions if c not in file_columns] + + target_rows = max(int(target_rows), 1) + parts: List[Dict[str, Any]] = [] + pieces: List[List[Any]] = [] + rows = 0 + for path, metadata in zip(listing["files"], footers): + groups = [] + for index in range(metadata.num_row_groups): + group_rows = metadata.row_group(index).num_rows + if group_rows == 0: + continue # writers emit an empty row group for an empty file + if group_rows > target_rows: + # close what is open, then one part per slice of this row group + if groups: + pieces.append([path, groups]) + groups = [] + if pieces: + parts.append({"pieces": pieces, "rows": rows}) + pieces, rows = [], 0 + for first in range(0, group_rows, target_rows): + count = min(target_rows, group_rows - first) + parts.append({"pieces": [[path, [index], first, count]], "rows": count}) + continue + if rows and rows + group_rows > target_rows: + if groups: + pieces.append([path, groups]) + groups = [] + parts.append({"pieces": pieces, "rows": rows}) + pieces, rows = [], 0 + groups.append(index) + rows += group_rows + if groups: + pieces.append([path, groups]) + if pieces: + parts.append({"pieces": pieces, "rows": rows}) + + return { + "kind": listing["kind"], + "root": root, + "columns": columns, + "total_rows": sum(m.num_rows for m in footers), + "total_bytes": sum( + m.row_group(i).total_byte_size + for m in footers + for i in range(m.num_row_groups) + ), + "files": len(listing["files"]), + "parts": parts, + } + + +def iter_parquet_part( + uri: str, root: str, pieces: List[List[Any]], batch_size: int +) -> Iterator[pd.DataFrame]: + """The rows of a part's row groups as DataFrames of at most batch_size rows. + + Hive partition values are added as columns, unless a file has a column by that name. + """ + fs, _ = filesystem_for(uri) + for piece in pieces: + path, row_groups = piece[0], piece[1] + first, count = (piece[2], piece[3]) if len(piece) == 4 else (0, None) + partition_values = hive_partition_values(path, root) + with fs.open_input_file(path) as f: + parquet = pq.ParquetFile(f) + position = 0 # rows of the row group before this batch + for batch in parquet.iter_batches(batch_size=batch_size, row_groups=row_groups): + start, position = position, position + batch.num_rows + if count is not None: + # a slice of one row group: skip the rows before it, stop after it + if position <= first: + continue + if start >= first + count: + break + batch = batch.slice(max(first - start, 0), first + count - max(start, first)) + df = batch.to_pandas() + for key, value in partition_values.items(): + if key not in df.columns: + df[key] = value + yield df + + +def preview_parquet(uri: str, rows: int = 1) -> Dict[str, Any]: + """Columns, the first rows as strings, and the size of a parquet source.""" + plan = plan_parquet_parts(uri, target_rows=10**12) + first: Optional[pd.DataFrame] = None + for df in iter_parquet_part(uri, plan["root"], plan["parts"][0]["pieces"][:1], rows): + first = df + break + columns = plan["columns"] + preview = [columns] + if first is not None: + # object dtype keeps each value's own type (an int column stays "1001", not "1001.0") + for _, row in first.head(rows).reindex(columns=columns).astype(object).iterrows(): + preview.append(["" if pd.isna(v) else str(v) for v in row.tolist()]) + return { + "kind": plan["kind"], + "columns": columns, + "preview_rows": preview, + "total_rows": plan["total_rows"], + "total_bytes": plan["total_bytes"], + "files": plan["files"], + "name": os.path.basename(_split(uri)[2]), + } diff --git a/materializationengine/blueprints/upload/schemas.py b/materializationengine/blueprints/upload/schemas.py index 6d6c24c0..71239e8a 100644 --- a/materializationengine/blueprints/upload/schemas.py +++ b/materializationengine/blueprints/upload/schemas.py @@ -7,6 +7,7 @@ validates_schema, ) from typing import Dict, Any +from materializationengine.blueprints.upload.parquet_source import is_parquet_source from materializationengine.blueprints.upload.source import parse_source_uri @@ -45,13 +46,16 @@ def validate_source_uri(self, value): if value: try: parse_source_uri(value) + is_parquet_source(value) # refuses a file that is neither CSV nor parquet except ValueError as e: raise ValidationError(str(e)) - @validates("filename") - def validate_filename(self, value): - if not value.endswith(".csv"): - raise ValidationError("File must be a CSV") + @validates_schema + def validate_filename(self, data, **kwargs): + # a file uploaded through the browser is a CSV; a cloud location may also be + # parquet (a file, a folder of files or a Delta table) + if not data.get("source_uri") and not data.get("filename", "").endswith(".csv"): + raise ValidationError("File must be a CSV", "filename") @validates_schema def validate_ignored_columns(self, data, **kwargs): diff --git a/materializationengine/blueprints/upload/source.py b/materializationengine/blueprints/upload/source.py index ad8c6caf..0525bfb9 100644 --- a/materializationengine/blueprints/upload/source.py +++ b/materializationengine/blueprints/upload/source.py @@ -6,6 +6,8 @@ - gs://bucket/path.csv: any GCS object this deployment's service account can read, which includes every publicly readable object. - s3://bucket/path.csv: a publicly readable S3 object, read with unsigned requests. +- A parquet file, a folder of parquet files or a Delta table at such a location: + see parquet_source.py. Every source offers its size, a ranged read and a stream over a byte range, which is all the CSV processing needs. Only processed output is written, and always to the upload @@ -35,9 +37,11 @@ def parse_source_uri(uri: str) -> Tuple[str, str, str]: raise ValueError( f"Unsupported location {uri!r}: use gs://bucket/path.csv or s3://bucket/path.csv" ) - bucket, key = parsed.netloc, parsed.path.lstrip("/") - if not bucket or not key or key.endswith("/"): - raise ValueError(f"{uri!r} does not name a file: use {scheme}://bucket/path.csv") + bucket, key = parsed.netloc, parsed.path.strip("/") + if not bucket or not key: + raise ValueError( + f"{uri!r} does not name a file or folder: use {scheme}://bucket/path" + ) return scheme, bucket, key diff --git a/materializationengine/blueprints/upload/tasks.py b/materializationengine/blueprints/upload/tasks.py index 0447ea07..993d13f7 100644 --- a/materializationengine/blueprints/upload/tasks.py +++ b/materializationengine/blueprints/upload/tasks.py @@ -19,6 +19,11 @@ from dynamicannotationdb.models import SegmentationMetadata from materializationengine.blueprints.upload.gcs_processor import GCSCsvProcessor +from materializationengine.blueprints.upload.parquet_source import ( + DEFAULT_PARQUET_PART_ROWS, + is_parquet_source, + plan_parquet_parts, +) from materializationengine.blueprints.upload.processor import SchemaProcessor from materializationengine.celery_init import celery from materializationengine.cloudsql_admin import ( @@ -409,34 +414,48 @@ def progress_callback(progress_details: Dict[str, Any]): }, ) - part_size_bytes = int( - float(current_app.config.get("UPLOAD_CSV_PART_SIZE_MB", DEFAULT_CSV_PART_SIZE_MB)) - * 1024 - * 1024 - ) - header, ranges = gcs_processor.plan_csv_parts(source_blob_name, part_size_bytes) - + part_kwargs = { + "schema_type": schema_type, + "column_mapping": column_mapping, + "reference_table": reference_table, + "ignored_columns": ignored_columns, + "chunk_size": chunk_size, + "timestamp": timestamp.isoformat(), + } parts_result = None - if len(ranges) > 1: - parts_result = _process_csv_in_parts( + if is_parquet_source(source_blob_name): + # a parquet file, a folder of them (hive partitions too), or a Delta table + parts_result = _process_parquet_in_parts( gcs_processor, bucket_name=bucket_name, - source_blob_name=source_blob_name, - header=header, - ranges=ranges, + source_uri=source_blob_name, destination_blob_name=destination_blob_name, job_id=job_id_for_status, - part_kwargs={ - "schema_type": schema_type, - "column_mapping": column_mapping, - "reference_table": reference_table, - "ignored_columns": ignored_columns, - "chunk_size": chunk_size, - "timestamp": timestamp.isoformat(), - }, + part_kwargs=part_kwargs, generate_ids=schema_processor.generate_ids, id_counter_start=id_counter_start, ) + else: + part_size_bytes = int( + float(current_app.config.get("UPLOAD_CSV_PART_SIZE_MB", DEFAULT_CSV_PART_SIZE_MB)) + * 1024 + * 1024 + ) + header, ranges = gcs_processor.plan_csv_parts(source_blob_name, part_size_bytes) + + if len(ranges) > 1: + parts_result = _process_csv_in_parts( + gcs_processor, + bucket_name=bucket_name, + source_blob_name=source_blob_name, + header=header, + ranges=ranges, + destination_blob_name=destination_blob_name, + job_id=job_id_for_status, + part_kwargs=part_kwargs, + generate_ids=schema_processor.generate_ids, + id_counter_start=id_counter_start, + ) if parts_result is None: gcs_processor.process_csv_in_chunks( @@ -572,6 +591,74 @@ def _process_csv_in_parts( } +def _process_parquet_in_parts( + gcs_processor: GCSCsvProcessor, + bucket_name: str, + source_uri: str, + destination_blob_name: str, + job_id: str, + part_kwargs: Dict[str, Any], + generate_ids: bool, + id_counter_start: int, +) -> Dict[str, Any]: + """Process a parquet source in part tasks of whole row groups, then join the outputs. + + The footers give every part's row count up front, so generated ids get exact + offsets without a counting pass. + """ + target_rows = int( + current_app.config.get("UPLOAD_PARQUET_PART_ROWS", DEFAULT_PARQUET_PART_ROWS) + ) + plan = plan_parquet_parts(source_uri, target_rows) + parts = plan["parts"] + if not parts: + raise ValueError(f"{source_uri} has no rows to upload") + celery_logger.info( + f"Upload {job_id}: {plan['kind']} source with {plan['files']} files, " + f"{plan['total_rows']:,} rows, in {len(parts)} parts" + ) + update_job_status(job_id, {"total_rows": plan["total_rows"]}) + part_names = [f"processed_parts/{job_id}/part-{i:05d}.csv" for i in range(len(parts))] + + starts = [] + offset = id_counter_start + for part in parts: + starts.append(offset) + offset += part["rows"] + + try: + results = _run_csv_part_tasks( + [ + process_parquet_part.si( + bucket_name=bucket_name, + source_uri=source_uri, + root=plan["root"], + pieces=part["pieces"], + expected_rows=part["rows"], + destination_blob_name=part_names[i], + id_counter_start=starts[i] if generate_ids else id_counter_start, + job_id_for_status=job_id, + **part_kwargs, + ) + for i, part in enumerate(parts) + ], + job_id, + "Processing Parquet Data", + ) + except BaseException: + gcs_processor.delete_blobs(part_names) + raise + + gcs_processor.compose(part_names, destination_blob_name) + rows_read = sum(r["rows_read"] for r in results) + return { + "rows_read": rows_read, + "rows_written": sum(r["rows_written"] for r in results), + "dropped_rows": sum(r["dropped_rows"] for r in results), + "last_assigned_id": id_counter_start + rows_read if generate_ids else id_counter_start, + } + + def _run_csv_part_tasks(signatures: list, job_id: str, phase: str) -> List[Any]: """Send the part tasks and wait for all of them, reporting progress on the job. @@ -673,6 +760,58 @@ def process_csv_part( return result +@celery.task( + name="process:process_parquet_part", + bind=True, + acks_late=True, + autoretry_for=(Exception,), + dont_autoretry_for=(Ignore,), + max_retries=3, + retry_backoff=True, +) +def process_parquet_part( + self, + bucket_name: str, + source_uri: str, + root: str, + pieces: List[List[Any]], + expected_rows: int, + destination_blob_name: str, + schema_type: str, + column_mapping: Dict[str, str], + timestamp: str, + reference_table: str = None, + ignored_columns: List[str] = None, + chunk_size: int = 10000, + id_counter_start: int = 0, + job_id_for_status: str = None, +) -> Dict[str, Any]: + """Transform some row groups of a parquet source into one processed CSV part.""" + raise_if_upload_cancelled(job_id_for_status, "parquet part processing") + schema_processor = SchemaProcessor( + schema_type, + reference_table, + column_mapping=column_mapping, + ignored_columns=ignored_columns, + id_counter_start=id_counter_start, + ) + created = datetime.fromisoformat(timestamp) + result = GCSCsvProcessor(bucket_name, chunk_size=chunk_size).process_parquet_part( + source_uri, + root, + pieces, + destination_blob_name, + lambda chunk: schema_processor.process_chunk(chunk, created), + ) + if result["rows_read"] != expected_rows: + # the footers promised a different count: the files changed since planning + raise ValueError( + f"Read {result['rows_read']} rows from {source_uri}, expected {expected_rows}" + ) + result["dropped_rows"] = schema_processor.dropped_rows + return result + + @celery.task( name="process:count_csv_part", acks_late=True, diff --git a/static/js/columnMatch.js b/static/js/columnMatch.js new file mode 100644 index 00000000..2ce2af53 --- /dev/null +++ b/static/js/columnMatch.js @@ -0,0 +1,92 @@ +// Suggest which file column feeds each schema field, from their names. +// +// Names are split into components (on _, -, spaces, dots and camelCase). A column +// component matches a field component when they are equal or one is a prefix of the +// other ("pos" ~ "position"), worth the shorter one's length in characters. The score +// is 2 * matched characters / (characters in both names): identical names score 1. +// +// A pair is only suggested when every component of one name matched, and it scores at +// least MIN_SCORE: "pre_x" suits "pre_pt_position_x" (0.44), but neither +// "pre_pt_root_id" (its "root" and "id" match nothing there) nor "id_partition" for +// "id" (0.31). Counting characters in order without components would not work: +// "post_x" spells out a subsequence of "pre_pt_position_x". +(function (root) { + const MIN_SCORE = 0.4; + + function components(name) { + return String(name) + .replace(/([a-z0-9])([A-Z])/g, "$1_$2") + .toLowerCase() + .split(/[^a-z0-9]+/) + .filter(Boolean); + } + + function componentMatch(a, b) { + if (a === b) return a.length; + if (a.startsWith(b) || b.startsWith(a)) return Math.min(a.length, b.length); + return 0; + } + + // { score, allColumn, allField } for one column against one field. + function scoreNames(column, field) { + const columnParts = components(column); + const fieldParts = components(field); + const columnChars = columnParts.join("").length; + const fieldChars = fieldParts.join("").length; + if (!columnChars || !fieldChars) return { score: 0, allColumn: false, allField: false }; + + // pair components one to one, best matches first + const pairs = []; + columnParts.forEach((c, i) => + fieldParts.forEach((f, j) => { + const chars = componentMatch(c, f); + if (chars) pairs.push({ i, j, chars, exact: c === f }); + }) + ); + pairs.sort((p, q) => q.chars - p.chars || q.exact - p.exact); + const usedColumn = new Set(); + const usedField = new Set(); + let matched = 0; + for (const p of pairs) { + if (usedColumn.has(p.i) || usedField.has(p.j)) continue; + usedColumn.add(p.i); + usedField.add(p.j); + matched += p.chars; + } + return { + score: (2 * matched) / (columnChars + fieldChars), + allColumn: usedColumn.size === columnParts.length, + allField: usedField.size === fieldParts.length, + }; + } + + // { field: column } for the fields that have a good enough match; each column is + // used once, the best-scoring pairs first. Fields in `skip` are left out. + function suggestColumnMapping(fieldNames, columns, skip) { + const skipFields = new Set(skip || []); + const candidates = []; + fieldNames.forEach((field) => { + if (skipFields.has(field)) return; + columns.forEach((column, order) => { + const exact = String(column).toLowerCase() === String(field).toLowerCase(); + const s = scoreNames(column, field); + if (exact || ((s.allColumn || s.allField) && s.score >= MIN_SCORE)) { + candidates.push({ field, column, score: exact ? 2 : s.score, order }); + } + }); + }); + candidates.sort((a, b) => b.score - a.score || a.order - b.order); + const mapping = {}; + const usedColumns = new Set(); + for (const c of candidates) { + if (c.field in mapping || usedColumns.has(c.column)) continue; + mapping[c.field] = c.column; + usedColumns.add(c.column); + } + return mapping; + } + + const api = { suggestColumnMapping, scoreNames, components, MIN_SCORE }; + if (typeof module !== "undefined" && module.exports) module.exports = api; + else root.columnMatch = api; +})(typeof window !== "undefined" ? window : this); diff --git a/static/js/step1.js b/static/js/step1.js index 57b0ac63..e13ed09e 100644 --- a/static/js/step1.js +++ b/static/js/step1.js @@ -21,6 +21,8 @@ document.addEventListener("alpine:init", () => { sourceMode: "file", sourceUri: "", sourceSize: 0, + // for parquet sources: "file", "folder" or "delta", and their files and rows + sourceInfo: null, speedStats: { startTime: null, lastUpdateTime: null, @@ -39,6 +41,7 @@ document.addEventListener("alpine:init", () => { this.sourceMode = state.sourceMode || "file"; this.sourceUri = state.sourceUri || ""; this.sourceSize = state.sourceSize || 0; + this.sourceInfo = state.sourceInfo || null; if (this.selectedDatastack) { const datastackSelect = document.getElementById('datastackSelect'); if (datastackSelect) { @@ -71,6 +74,7 @@ document.addEventListener("alpine:init", () => { this.selectedDatastack = ""; this.sourceUri = ""; this.sourceSize = 0; + this.sourceInfo = null; this.speedStats = { startTime: null, lastUpdateTime: null, @@ -96,6 +100,7 @@ document.addEventListener("alpine:init", () => { // only sent to the server when reading in place sourceUri: this.sourceMode === "uri" ? this.sourceUri : "", sourceSize: this.sourceSize, + sourceInfo: this.sourceInfo, }; localStorage.setItem("uploadStore", JSON.stringify(stateToSave)); }, @@ -135,6 +140,7 @@ document.addEventListener("alpine:init", () => { this.previewRows = []; this.filename = ""; this.sourceSize = 0; + this.sourceInfo = null; if (!datastackName) { this.error = "Please select a datastack first."; this.status = "error"; @@ -165,6 +171,9 @@ document.addEventListener("alpine:init", () => { } this.filename = data.filename; this.sourceSize = data.size; + this.sourceInfo = data.kind + ? { kind: data.kind, files: data.files, rows: data.rows } + : null; this.previewRows = data.previewRows || []; this.status = "completed"; this.saveState(); diff --git a/static/js/step2.js b/static/js/step2.js index f1b6e0bc..5953649b 100644 --- a/static/js/step2.js +++ b/static/js/step2.js @@ -12,6 +12,8 @@ document.addEventListener("alpine:init", () => { mappingSaved: false, useMergedPositions: false, hasPositionFields: false, + // fields whose column was filled in by suggestColumns, until the user changes it + suggestedFields: [], }, init() { @@ -106,6 +108,9 @@ document.addEventListener("alpine:init", () => { (field) => field.type.toLowerCase().includes("geometry") ); this.updateDisplayFields(); + if (!isRestoring) { + this.suggestColumns(); + } this.updateIgnoredColumns(); this.state.validationErrors = []; this.saveState(); @@ -164,12 +169,51 @@ document.addEventListener("alpine:init", () => { newMapping[field.name] = currentMapping[field.name] || ""; }); this.state.columnMapping = newMapping; + this.state.suggestedFields = this.state.suggestedFields.filter( + (name) => name in newMapping + ); + this.suggestColumns(); this.updateIgnoredColumns(); this.saveState(); }, + // Fill fields that have no column yet with the best-matching column names + // (columnMatch.js). Columns already chosen are kept, and not suggested again. + suggestColumns() { + if (!window.columnMatch || !this.state.csvColumns.length) return 0; + const mapping = this.state.columnMapping; + const fieldNames = this.state.displayFields.map((field) => field.name); + const filled = fieldNames.filter((name) => mapping[name]); + const taken = new Set(filled.map((name) => mapping[name])); + const suggestions = window.columnMatch.suggestColumnMapping( + fieldNames, + this.state.csvColumns.filter((column) => !taken.has(column)), + filled + ); + Object.entries(suggestions).forEach(([field, column]) => { + mapping[field] = column; + if (!this.state.suggestedFields.includes(field)) { + this.state.suggestedFields.push(field); + } + }); + this.updateIgnoredColumns(); + this.saveState(); + return Object.keys(suggestions).length; + }, + + isSuggested(fieldName) { + return ( + this.state.suggestedFields.includes(fieldName) && + Boolean(this.state.columnMapping[fieldName]) + ); + }, + updateColumnMapping(fieldName, columnName) { this.state.columnMapping[fieldName] = columnName; + // the user's choice now, not a suggestion + this.state.suggestedFields = this.state.suggestedFields.filter( + (name) => name !== fieldName + ); this.updateIgnoredColumns(); this.saveState(); }, @@ -249,6 +293,7 @@ document.addEventListener("alpine:init", () => { this.state.displayFields = []; this.state.columnMapping = {}; this.state.ignoredColumns = []; + this.state.suggestedFields = []; this.state.validationErrors = []; this.state.mappingSaved = false; this.saveState(); @@ -276,6 +321,7 @@ document.addEventListener("alpine:init", () => { schemaModel: null, validationErrors: [], mappingSaved: false, + suggestedFields: [], }; localStorage.removeItem("schemaStore"); this.init(); diff --git a/templates/upload/step1.html b/templates/upload/step1.html index 5dfd34b2..8a448575 100644 --- a/templates/upload/step1.html +++ b/templates/upload/step1.html @@ -42,10 +42,10 @@

Step 1: Choose CSV File

- +
@@ -58,15 +58,19 @@

Step 1: Choose CSV File

- The file is read in place, without uploading it. A gs:// file must be publicly - readable or readable by this deployment's service account; an s3:// file must be - publicly readable. + The data is read in place, without uploading it: a CSV file, or parquet as one + file, a folder of files (hive-partitioned folders like key=value/ + become columns) or a Delta table (its current version is read). A gs:// location + must be publicly readable or readable by this deployment's service account; an + s3:// location must be publicly readable.
- ( MB) + ( MB) is readable.
CSV Preview
-

Map CSV Columns to Schema Fields

+
+

Map Columns to Schema Fields

+ +
+

+ Columns whose names match a field are filled in as suggestions; check them and + change any before going on. +

Map CSV Columns to Schema Fields Schema Field Type Required - CSV Column + Column @@ -152,6 +163,9 @@

Map CSV Columns to Schema Fields

Position + suggested diff --git a/templates/upload_wizard.html b/templates/upload_wizard.html index 617ed44d..71fbde8a 100644 --- a/templates/upload_wizard.html +++ b/templates/upload_wizard.html @@ -5,6 +5,7 @@ + diff --git a/tests/test_column_match_js.py b/tests/test_column_match_js.py new file mode 100644 index 00000000..cd7f2ecc --- /dev/null +++ b/tests/test_column_match_js.py @@ -0,0 +1,50 @@ +"""static/js/columnMatch.js: the wizard's suggested column mapping (run with node).""" + +import json +import os +import shutil +import subprocess + +import pytest + +SCRIPT = os.path.join(os.path.dirname(__file__), "..", "static", "js", "columnMatch.js") +FIELDS = ["id"] + [ + f"{point}_position_{axis}" for point in ("pre_pt", "post_pt", "ctr_pt") for axis in "xyz" +] + ["size"] + +pytestmark = pytest.mark.skipif(shutil.which("node") is None, reason="needs node") + + +def suggest(fields, columns, skip=None): + code = ( + f"const m = require({json.dumps(os.path.abspath(SCRIPT))});" + f"console.log(JSON.stringify(m.suggestColumnMapping({json.dumps(fields)}," + f" {json.dumps(columns)}, {json.dumps(skip or [])})));" + ) + return json.loads(subprocess.run(["node", "-e", code], capture_output=True, text=True, check=True).stdout) + + +class TestSuggestColumnMapping: + def test_identical_names(self): + columns = ["valid", "pre_pt_root_id", "id_partition"] + FIELDS + assert suggest(FIELDS, columns) == {f: f for f in FIELDS} + + def test_short_names_match_by_component(self): + columns = ["id", "pre_x", "pre_y", "pre_z", "post_x", "post_y", "post_z", "ctr_x", "ctr_y", "ctr_z", "size"] + mapping = suggest(FIELDS, columns) + assert mapping["pre_pt_position_x"] == "pre_x" + assert mapping["post_pt_position_x"] == "post_x" # not pre_x: "post" != "pre" + assert mapping["ctr_pt_position_z"] == "ctr_z" + assert len(set(mapping.values())) == len(mapping) == len(FIELDS) + + def test_camel_case_and_abbreviations(self): + mapping = suggest(FIELDS, ["prePosX", "postPosY", "size_nm"]) + assert mapping == {"pre_pt_position_x": "prePosX", "post_pt_position_y": "postPosY", "size": "size_nm"} + + def test_poor_matches_are_not_suggested(self): + # root ids are not positions; id_partition is not a unique id; a bare x is too vague + assert suggest(FIELDS, ["id_partition", "pre_pt_root_id", "post_pt_root_id", "x", "notes"]) == {} + + def test_skipped_fields_are_left_alone(self): + mapping = suggest(FIELDS, ["pre_x", "size"], skip=["size"]) + assert mapping == {"pre_pt_position_x": "pre_x"} diff --git a/tests/test_upload_parquet_source.py b/tests/test_upload_parquet_source.py new file mode 100644 index 00000000..f8d79b7d --- /dev/null +++ b/tests/test_upload_parquet_source.py @@ -0,0 +1,258 @@ +"""Uploads read from a parquet file, a (hive-partitioned) folder of them, or a Delta table.""" + +import io +from datetime import datetime, timezone +from unittest import mock + +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.dataset as ds +import pyarrow.parquet as pq +import pytest +from deltalake import write_deltalake +from flask import Flask + +from materializationengine.blueprints.upload import tasks +from materializationengine.blueprints.upload.parquet_source import ( + hive_partition_values, + is_parquet_source, + iter_parquet_part, + list_parquet_files, + plan_parquet_parts, + preview_parquet, +) +from materializationengine.blueprints.upload.schemas import UploadRequestSchema +from tests.test_upload_csv_parts import ( + BUCKET, + JOB, + MAPPING_GENERATED_IDS, + MAPPING_WITH_IDS, + FakeBucket, + FakeRedis, + make_processor, + run_parts_inline, +) + + +def synapses(n, seed=0): + rng = np.random.default_rng(seed) + df = pd.DataFrame({"id": np.arange(1000, 1000 + n, dtype=np.int64)}) + for point in ["pre", "post", "ctr"]: + for axis in "xyz": + df[f"{point}_{axis}"] = rng.integers(0, 300000, n).astype(np.int32) + df["size"] = rng.integers(1, 10000, n).astype(float) + df["bucket"] = df["id"] % 7 # a hive partition column + return df + + +def write_hive_folder(path, df, row_group_size=40): + table = pa.Table.from_pandas(df, preserve_index=False) + ds.write_dataset( + table, + str(path), + format="parquet", + partitioning=ds.partitioning(pa.schema([("bucket", pa.int64())]), flavor="hive"), + file_options=ds.ParquetFileFormat().make_write_options(), + max_rows_per_group=row_group_size, + min_rows_per_group=row_group_size, + basename_template="part-{i}.parquet", + ) + (path / "_SUCCESS").write_text("") + (path / "bucket=0" / ".part-0.parquet.crc").write_text("x") + return path + + +@pytest.fixture +def redis_clients(): + with mock.patch.object(tasks, "REDIS_CLIENT", FakeRedis()) as fake: + yield fake + + +class TestListing: + def test_is_parquet_source(self): + assert is_parquet_source("gs://b/dir/") and is_parquet_source("s3://b/x.parquet") + assert is_parquet_source("gs://b/delta_table") + assert not is_parquet_source("gs://b/x.csv") and not is_parquet_source("gs://b/X.CSV") + assert is_parquet_source("gs://b/table.v2/") # a trailing slash means a folder + with pytest.raises(ValueError, match="not a .csv or .parquet"): + is_parquet_source("gs://b/notes.txt") + + def test_hive_folder_skips_hidden_entries(self, tmp_path): + write_hive_folder(tmp_path / "t", synapses(300)) + listing = list_parquet_files(f"file://{tmp_path}/t") + assert listing["kind"] == "folder" + assert len(listing["files"]) == 7 + assert all(f.endswith(".parquet") and "/." not in f for f in listing["files"]) + + def test_hive_partition_values(self): + assert hive_partition_values("/r/a=1/b=x%2Fy/c=2.5/f.parquet", "/r") == {"a": 1, "b": "x/y", "c": 2.5} + assert hive_partition_values("/r/plain/f.parquet", "/r") == {} + assert hive_partition_values("/r/a=__HIVE_DEFAULT_PARTITION__/f.parquet", "/r") == {"a": None} + + def test_single_file(self, tmp_path): + pq.write_table(pa.Table.from_pandas(synapses(50), preserve_index=False), tmp_path / "one.parquet") + listing = list_parquet_files(f"file://{tmp_path}/one.parquet") + assert listing["kind"] == "file" and len(listing["files"]) == 1 + + def test_delta_table_reads_only_current_files(self, tmp_path): + uri = f"file://{tmp_path}/delta" + write_deltalake(str(tmp_path / "delta"), synapses(100, seed=1), partition_by=["bucket"]) + newer = synapses(60, seed=2) + write_deltalake(str(tmp_path / "delta"), newer, partition_by=["bucket"], mode="overwrite") + on_disk = list((tmp_path / "delta").rglob("*.parquet")) + listing = list_parquet_files(uri) + assert listing["kind"] == "delta" and listing["partition_columns"] == ["bucket"] + assert len(on_disk) > len(listing["files"]) # the overwritten files are still there + plan = plan_parquet_parts(uri) + assert plan["total_rows"] == 60 + rows = pd.concat(iter_parquet_part(uri, plan["root"], plan["parts"][0]["pieces"], 1000)) + assert sorted(rows["id"]) == sorted(newer["id"]) + + def test_missing(self, tmp_path): + with pytest.raises(FileNotFoundError): + list_parquet_files(f"file://{tmp_path}/nope") + (tmp_path / "empty").mkdir() + with pytest.raises(FileNotFoundError): + list_parquet_files(f"file://{tmp_path}/empty") + + +class TestPlanAndRead: + @pytest.mark.parametrize("target_rows", [1, 40, 100, 10**9]) + def test_parts_cover_every_row_group_once(self, tmp_path, target_rows): + df = synapses(900) + write_hive_folder(tmp_path / "t", df) + uri = f"file://{tmp_path}/t" + plan = plan_parquet_parts(uri, target_rows) + assert plan["total_rows"] == 900 == sum(p["rows"] for p in plan["parts"]) + # a piece is [path, groups] or a slice [path, [group], first row, rows] + seen = [ + (piece[0], g, piece[2] if len(piece) == 4 else None) + for p in plan["parts"] for piece in p["pieces"] for g in piece[1] + ] + assert len(seen) == len(set(seen)) + assert all(p["rows"] <= target_rows for p in plan["parts"]) + rows = pd.concat( + [d for p in plan["parts"] for d in iter_parquet_part(uri, plan["root"], p["pieces"], 25)] + ) + assert len(rows) == 900 + merged = rows.sort_values("id").reset_index(drop=True) + expected = df.sort_values("id").reset_index(drop=True) + assert (merged["bucket"].to_numpy() == expected["bucket"].to_numpy()).all() # from the paths + assert (merged["pre_x"].to_numpy() == expected["pre_x"].to_numpy()).all() + + @pytest.mark.parametrize("target_rows", [1, 7, 300, 999, 1000]) + def test_a_big_row_group_is_split_into_slices(self, tmp_path, target_rows): + df = synapses(1000) + pq.write_table(pa.Table.from_pandas(df, preserve_index=False), tmp_path / "big.parquet", + row_group_size=10**6) + uri = f"file://{tmp_path}/big.parquet" + plan = plan_parquet_parts(uri, target_rows) + assert len(plan["parts"]) == -(-1000 // target_rows) + assert all(p["rows"] <= target_rows for p in plan["parts"]) + rows = pd.concat( + [d for p in plan["parts"] for d in iter_parquet_part(uri, plan["root"], p["pieces"], 64)] + ) + assert rows["id"].tolist() == df["id"].tolist() # every row once, in order + + def test_preview(self, tmp_path): + write_hive_folder(tmp_path / "t", synapses(100)) + info = preview_parquet(f"file://{tmp_path}/t") + assert info["kind"] == "folder" and info["files"] == 7 and info["total_rows"] == 100 + header, first = info["preview_rows"] + assert header[-1] == "bucket" and "pre_x" in header + assert len(first) == len(header) and first[header.index("bucket")] == "0" + + +class TestProcessCsvFromParquet: + """process_csv on a parquet source matches the same rows uploaded as a CSV.""" + + def run(self, uri, data_csv, mapping, redis_clients, part_rows): + bucket = FakeBucket({"in.csv": data_csv}) + app = Flask(__name__) + app.config.update( + MATERIALIZATION_UPLOAD_BUCKET_PATH=BUCKET, + UPLOAD_PARQUET_PART_ROWS=part_rows, + UPLOAD_CSV_PART_SIZE_MB=1000, + ) + outputs = {} + with app.app_context(), \ + mock.patch.object(tasks, "GCSCsvProcessor", lambda *a, **k: make_processor(bucket, 50)), \ + mock.patch.object(tasks, "_run_csv_part_tasks", run_parts_inline), \ + mock.patch.object(tasks, "datetime") as fake_datetime: + fake_datetime.now.return_value = datetime(2026, 10, 25, tzinfo=timezone.utc) + fake_datetime.fromisoformat = datetime.fromisoformat + for name, path in [("parquet", uri), ("csv", f"gs://{BUCKET}/in.csv")]: + result = tasks.process_csv.run( + file_path=path, schema_type="synapse", column_mapping=mapping, + job_id_for_status=JOB, + ) + outputs[name] = (bucket.objects.pop(result["output_path"].split("/", 1)[1]), result) + assert not any(k.startswith("processed_parts/") for k in bucket.objects) + return outputs + + @pytest.mark.parametrize("mapping", [MAPPING_WITH_IDS, MAPPING_GENERATED_IDS], ids=["ids", "generated"]) + def test_matches_the_csv_path(self, tmp_path, redis_clients, mapping): + df = synapses(600) + df.loc[[5, 77], "pre_x"] = np.nan # dropped: missing coordinates + # sort the CSV the way the parquet parts are read: by partition folder, then file order + write_hive_folder(tmp_path / "t", df) + uri = f"file://{tmp_path}/t" + plan = plan_parquet_parts(uri, 10**9) + ordered = pd.concat(iter_parquet_part(uri, plan["root"], plan["parts"][0]["pieces"], 10**6)) + ordered = ordered[["id"] + [c for c in df.columns if c not in ("id",)]] + if "id" not in mapping: + ordered = ordered.drop(columns=["id"]) + csv = io.StringIO() + ordered.to_csv(csv, index=False) + outputs = self.run(uri, csv.getvalue().encode(), mapping, redis_clients, part_rows=120) + parquet_out, parquet_result = outputs["parquet"] + csv_out, csv_result = outputs["csv"] + assert parquet_out == csv_out + assert parquet_result["dropped_rows"] == csv_result["dropped_rows"] == 2 + assert parquet_result["last_assigned_id"] == csv_result["last_assigned_id"] + + def test_empty_source_fails(self, tmp_path, redis_clients): + pq.write_table(pa.Table.from_pandas(synapses(0), preserve_index=False), tmp_path / "e.parquet") + app = Flask(__name__) + app.config.update(MATERIALIZATION_UPLOAD_BUCKET_PATH=BUCKET) + with app.app_context(), \ + mock.patch.object(tasks, "GCSCsvProcessor", lambda *a, **k: make_processor(FakeBucket(), 50)): + with pytest.raises(ValueError, match="no rows"): + tasks.process_csv.run( + file_path=f"file://{tmp_path}/e.parquet", schema_type="synapse", + column_mapping=MAPPING_WITH_IDS, job_id_for_status=JOB, + ) + + def test_part_with_changed_row_count_fails(self, tmp_path, redis_clients): + write_hive_folder(tmp_path / "t", synapses(100)) + uri = f"file://{tmp_path}/t" + plan = plan_parquet_parts(uri, 10**9) + bucket = FakeBucket() + with mock.patch.object(tasks, "GCSCsvProcessor", lambda *a, **k: make_processor(bucket, 50)): + with pytest.raises(ValueError, match="expected"): + tasks.process_parquet_part.run( + bucket_name=BUCKET, source_uri=uri, root=plan["root"], + pieces=plan["parts"][0]["pieces"], expected_rows=101, + destination_blob_name="p.csv", schema_type="synapse", + column_mapping=MAPPING_WITH_IDS, timestamp="2026-10-25T00:00:00", + ) + + +class TestRequestSchema: + BASE = { + "columnMapping": {}, + "metadata": { + "schema_type": "synapse", "datastack_name": "d", "table_name": "t", + "description": "x", "voxel_resolution_nm_x": 4, "voxel_resolution_nm_y": 4, + "voxel_resolution_nm_z": 40, "read_permission": "PUBLIC", "write_permission": "PRIVATE", + }, + } + + def test_a_parquet_folder_needs_no_csv_name(self): + loaded = UploadRequestSchema().load({**self.BASE, "filename": "id", "sourceUri": "gs://b/t/id/"}) + assert loaded["source_uri"] == "gs://b/t/id/" + + def test_a_browser_upload_must_be_a_csv(self): + with pytest.raises(Exception, match="CSV"): + UploadRequestSchema().load({**self.BASE, "filename": "data.parquet"}) diff --git a/tests/test_upload_source.py b/tests/test_upload_source.py index 3b43922a..2bf5755d 100644 --- a/tests/test_upload_source.py +++ b/tests/test_upload_source.py @@ -72,12 +72,15 @@ def test_valid(self, uri, expected): assert parse_source_uri(uri) == expected @pytest.mark.parametrize( - "uri", ["", "k.csv", "https://b/k.csv", "gs://b", "gs://b/", "s3:///k.csv", "gs://b/dir/"] + "uri", ["", "k.csv", "https://b/k.csv", "gs://b", "gs://b/", "s3:///k.csv"] ) def test_invalid(self, uri): with pytest.raises(ValueError): parse_source_uri(uri) + def test_a_folder_is_valid(self): + assert parse_source_uri("gs://b/dir/") == ("gs", "b", "dir") + class TestS3Source: def test_size_and_ranged_reads(self): @@ -175,7 +178,7 @@ def test_preview_of_a_large_file_drops_the_cut_off_line(self, monkeypatch): @pytest.mark.parametrize( "uri, code", - [("in.csv", 400), ("ftp://b/in.csv", 400), ("gs://uploads/in.txt", 404), + [("in.csv", 400), ("ftp://b/in.csv", 400), ("gs://uploads/in.txt", 400), ("gs://uploads/missing.csv", 404)], ) def test_errors(self, uri, code): @@ -183,10 +186,10 @@ def test_errors(self, uri, code): status, body = self.call(uri, processor) assert status == code and body["status"] == "error" - def test_not_a_csv(self): + def test_neither_csv_nor_parquet(self): processor = make_processor(FakeBucket({"in.txt": b"a\n1\n"})) status, body = self.call("gs://uploads/in.txt", processor) - assert status == 400 and "CSV" in body["message"] + assert status == 400 and ".csv or .parquet" in body["message"] def test_private_gcs_names_the_service_account(self): processor = make_processor(FakeBucket())