From e9239985e4b59a8ab785384027eaca53b5ca69c0 Mon Sep 17 00:00:00 2001 From: simozzy Date: Thu, 18 Jun 2026 21:24:54 +0100 Subject: [PATCH 1/2] fix: treat Compression.NONE sentinel as uncompressed in COPY INTO formats Compression.NONE is a truthy Enum member, so the `if compression:` guard in CSVFormat/TSVFormat/NDJSONFormat/ParquetFormat treated the no-compression sentinel as 'a codec was specified'. For ParquetFormat this raised TypeError('Compression should be None, ZStd, or Snappy.') outright; for the others it emitted a spurious COMPRESSION = NONE option. Guard now skips both None and Compression.NONE, so the uncompressed case renders no COMPRESSION option. Explicit codecs are unchanged, and ParquetFormat still rejects codecs its writer can't use (gzip/bz2/zip). Co-Authored-By: Claude Opus 4.8 (1M context) --- databend_sqlalchemy/dml.py | 12 +++++--- tests/test_copy_format.py | 57 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 4 deletions(-) create mode 100644 tests/test_copy_format.py diff --git a/databend_sqlalchemy/dml.py b/databend_sqlalchemy/dml.py index cb2ad2b..4856d64 100644 --- a/databend_sqlalchemy/dml.py +++ b/databend_sqlalchemy/dml.py @@ -277,6 +277,7 @@ def __repr__(self): class CSVFormat(CopyFormat): format_type = "CSV" + inherit_cache = False def __init__( self, @@ -340,12 +341,13 @@ def __init__( if binary_format not in ["HEX", "BASE64"]: raise TypeError('Binary Format should be "HEX" or "BASE64".') self.options["BINARY_FORMAT"] = binary_format - if compression: + if compression and compression is not Compression.NONE: self.options["COMPRESSION"] = compression.value class TSVFormat(CopyFormat): format_type = "TSV" + inherit_cache = False def __init__( self, @@ -366,12 +368,13 @@ def __init__( if len(str(field_delimiter).encode().decode("unicode_escape")) != 1: raise TypeError("Field Delimiter should be a single character") self.options["FIELD_DELIMITER"] = f"{repr(field_delimiter)}" - if compression: + if compression and compression is not Compression.NONE: self.options["COMPRESSION"] = compression.value class NDJSONFormat(CopyFormat): format_type = "NDJSON" + inherit_cache = False def __init__( self, @@ -396,12 +399,13 @@ def __init__( 'Missing Field As should be "ERROR", "NULL", "FIELD_DEFAULT" or "TYPE_DEFAULT".' ) self.options["MISSING_FIELD_AS"] = f"{missing_field_as}" - if compression: + if compression and compression is not Compression.NONE: self.options["COMPRESSION"] = compression.value class ParquetFormat(CopyFormat): format_type = "PARQUET" + inherit_cache = False def __init__( self, @@ -416,7 +420,7 @@ def __init__( 'Missing Field As should be "ERROR" or "FIELD_DEFAULT".' ) self.options["MISSING_FIELD_AS"] = f"{missing_field_as}" - if compression: + if compression and compression is not Compression.NONE: if compression not in [Compression.ZSTD, Compression.SNAPPY]: raise TypeError( 'Compression should be None, ZStd, or Snappy.' diff --git a/tests/test_copy_format.py b/tests/test_copy_format.py new file mode 100644 index 0000000..9144e6d --- /dev/null +++ b/tests/test_copy_format.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python + +from sqlalchemy.testing import fixtures, assert_raises_message +from sqlalchemy.testing.assertions import AssertsCompiledSQL + +from databend_sqlalchemy import ( + CSVFormat, + TSVFormat, + NDJSONFormat, + ParquetFormat, + Compression, +) + + +class CompileCopyFormatCompressionTest(fixtures.TestBase, AssertsCompiledSQL): + + __only_on__ = "databend" + + _FORMATS = [ + (CSVFormat, "CSV"), + (TSVFormat, "TSV"), + (NDJSONFormat, "NDJSON"), + (ParquetFormat, "PARQUET"), + ] + + def test_no_compression_emits_no_option(self): + # ``None`` and the ``Compression.NONE`` sentinel both mean uncompressed. + # Regression: ``Compression.NONE`` is a truthy Enum member, so the old + # ``if compression:`` guard treated it as "a codec was set" — emitting a + # spurious ``COMPRESSION = NONE`` option and, for ParquetFormat, raising. + for fmt_cls, type_ in self._FORMATS: + expected = f"FILE_FORMAT = (TYPE = {type_})" + self.assert_compile(fmt_cls(), expected) + self.assert_compile(fmt_cls(compression=Compression.NONE), expected) + + def test_explicit_codec_is_emitted(self): + for fmt_cls, type_, codec in [ + (CSVFormat, "CSV", Compression.GZIP), + (TSVFormat, "TSV", Compression.GZIP), + (NDJSONFormat, "NDJSON", Compression.GZIP), + (ParquetFormat, "PARQUET", Compression.ZSTD), + (ParquetFormat, "PARQUET", Compression.SNAPPY), + ]: + self.assert_compile( + fmt_cls(compression=codec), + f"FILE_FORMAT = (TYPE = {type_}, COMPRESSION = {codec.value})", + ) + + def test_parquet_rejects_unsupported_codec(self): + # ParquetFormat still rejects codecs its writer can't use. + for codec in [Compression.GZIP, Compression.BZ2, Compression.ZIP]: + assert_raises_message( + TypeError, + "Compression should be None, ZStd, or Snappy.", + ParquetFormat, + compression=codec, + ) From c83bc03edca6c51b5cb16613d28d18535a3ecbcf Mon Sep 17 00:00:00 2001 From: simozzy Date: Mon, 19 Oct 2026 18:37:44 +0100 Subject: [PATCH 2/2] fix: emit explicit COMPRESSION = NONE instead of dropping it Omitting COMPRESSION means the server default, and for Parquet that is ZSTD, so dropping an explicit Compression.NONE silently produced compressed output. CSV/TSV/NDJSON go back to emitting it as before, and ParquetFormat now accepts NONE alongside ZSTD and SNAPPY. The unload test checks the written file's codec via inspect_parquet: compressed and uncompressed row-group sizes match only for NONE. Co-Authored-By: Claude Opus 5.5 --- databend_sqlalchemy/dml.py | 16 +++++--- tests/test_copy_format.py | 78 ++++++++++++++++++++++++++++++++------ 2 files changed, 77 insertions(+), 17 deletions(-) diff --git a/databend_sqlalchemy/dml.py b/databend_sqlalchemy/dml.py index 4856d64..dc9efa1 100644 --- a/databend_sqlalchemy/dml.py +++ b/databend_sqlalchemy/dml.py @@ -341,7 +341,7 @@ def __init__( if binary_format not in ["HEX", "BASE64"]: raise TypeError('Binary Format should be "HEX" or "BASE64".') self.options["BINARY_FORMAT"] = binary_format - if compression and compression is not Compression.NONE: + if compression: self.options["COMPRESSION"] = compression.value @@ -368,7 +368,7 @@ def __init__( if len(str(field_delimiter).encode().decode("unicode_escape")) != 1: raise TypeError("Field Delimiter should be a single character") self.options["FIELD_DELIMITER"] = f"{repr(field_delimiter)}" - if compression and compression is not Compression.NONE: + if compression: self.options["COMPRESSION"] = compression.value @@ -399,7 +399,7 @@ def __init__( 'Missing Field As should be "ERROR", "NULL", "FIELD_DEFAULT" or "TYPE_DEFAULT".' ) self.options["MISSING_FIELD_AS"] = f"{missing_field_as}" - if compression and compression is not Compression.NONE: + if compression: self.options["COMPRESSION"] = compression.value @@ -420,10 +420,14 @@ def __init__( 'Missing Field As should be "ERROR" or "FIELD_DEFAULT".' ) self.options["MISSING_FIELD_AS"] = f"{missing_field_as}" - if compression and compression is not Compression.NONE: - if compression not in [Compression.ZSTD, Compression.SNAPPY]: + if compression: + if compression not in [ + Compression.NONE, + Compression.ZSTD, + Compression.SNAPPY, + ]: raise TypeError( - 'Compression should be None, ZStd, or Snappy.' + 'Compression should be NONE, ZSTD, or SNAPPY.' ) self.options["COMPRESSION"] = compression.value diff --git a/tests/test_copy_format.py b/tests/test_copy_format.py index 9144e6d..2a32913 100644 --- a/tests/test_copy_format.py +++ b/tests/test_copy_format.py @@ -1,15 +1,21 @@ #!/usr/bin/env python -from sqlalchemy.testing import fixtures, assert_raises_message +from sqlalchemy.testing import fixtures, assert_raises_message, eq_ from sqlalchemy.testing.assertions import AssertsCompiledSQL +from sqlalchemy import Table, Column, String, text from databend_sqlalchemy import ( + CopyIntoLocation, + CopyIntoLocationOptions, CSVFormat, TSVFormat, NDJSONFormat, ParquetFormat, Compression, + StageClause, ) +import sqlalchemy +from packaging import version class CompileCopyFormatCompressionTest(fixtures.TestBase, AssertsCompiledSQL): @@ -23,15 +29,17 @@ class CompileCopyFormatCompressionTest(fixtures.TestBase, AssertsCompiledSQL): (ParquetFormat, "PARQUET"), ] - def test_no_compression_emits_no_option(self): - # ``None`` and the ``Compression.NONE`` sentinel both mean uncompressed. - # Regression: ``Compression.NONE`` is a truthy Enum member, so the old - # ``if compression:`` guard treated it as "a codec was set" — emitting a - # spurious ``COMPRESSION = NONE`` option and, for ParquetFormat, raising. + def test_omitted_compression_emits_no_option(self): for fmt_cls, type_ in self._FORMATS: - expected = f"FILE_FORMAT = (TYPE = {type_})" - self.assert_compile(fmt_cls(), expected) - self.assert_compile(fmt_cls(compression=Compression.NONE), expected) + self.assert_compile(fmt_cls(), f"FILE_FORMAT = (TYPE = {type_})") + + def test_explicit_none_is_emitted(self): + # Omitting COMPRESSION means the server default, which for Parquet is ZSTD. + for fmt_cls, type_ in self._FORMATS: + self.assert_compile( + fmt_cls(compression=Compression.NONE), + f"FILE_FORMAT = (TYPE = {type_}, COMPRESSION = NONE)", + ) def test_explicit_codec_is_emitted(self): for fmt_cls, type_, codec in [ @@ -47,11 +55,59 @@ def test_explicit_codec_is_emitted(self): ) def test_parquet_rejects_unsupported_codec(self): - # ParquetFormat still rejects codecs its writer can't use. for codec in [Compression.GZIP, Compression.BZ2, Compression.ZIP]: assert_raises_message( TypeError, - "Compression should be None, ZStd, or Snappy.", + "Compression should be NONE, ZSTD, or SNAPPY.", ParquetFormat, compression=codec, ) + + +class ParquetUnloadCompressionTest(fixtures.TablesTest): + run_create_tables = "each" + __backend__ = True + + @classmethod + def define_tables(cls, metadata): + Table("repetitive", metadata, Column("data", String(200))) + + @classmethod + def insert_data(cls, connection): + connection.execute( + cls.tables.repetitive.insert(), + [{"data": "x" * 200} for _ in range(1000)], + ) + + def _unload(self, connection, compression): + # Stages are account-wide, so each test needs its own under xdist. + stage = f"parquet_codec_{compression.value if compression else 'default'}" + connection.execute(text(f"CREATE OR REPLACE STAGE {stage}")) + connection.execute( + CopyIntoLocation( + target=StageClause(name=stage), + from_=self.tables.repetitive.select(), + file_format=ParquetFormat(compression=compression), + options=CopyIntoLocationOptions(single=True), + ) + ) + name = connection.execute(text(f"LIST @{stage}")).scalar() + sizes = connection.execute( + text( + "SELECT max_row_groups_size_compressed, " + "max_row_groups_size_uncompressed " + f"FROM inspect_parquet('@{stage}/{name}')" + ) + ).one() + connection.execute(text(f"DROP STAGE {stage}")) + return sizes + + if version.parse(sqlalchemy.__version__) >= version.parse("2.0.0"): + + def test_explicit_none_writes_uncompressed(self, connection): + compressed, uncompressed = self._unload(connection, Compression.NONE) + eq_(compressed, uncompressed) + + def test_omitted_compression_uses_server_default(self, connection): + compressed, uncompressed = self._unload(connection, None) + assert compressed < uncompressed